diff --git a/CodeEntropy/trajectory/mda.py b/CodeEntropy/trajectory/mda.py index 1cdb6e75..56118a3d 100644 --- a/CodeEntropy/trajectory/mda.py +++ b/CodeEntropy/trajectory/mda.py @@ -354,7 +354,10 @@ def merge_forces( select_atom_force = u_force.select_atoms("all") coordinates = self._extract_timeseries(select_atom, kind="positions") - dimensions = self._extract_timeseries(select_atom, kind="dimensions") + if u.dimensions is None: + dimensions = self._build_non_periodic_dimensions(coordinates) + else: + dimensions = self._extract_timeseries(select_atom, kind="dimensions") forces = self._extract_force_timeseries_with_fallback( select_atom_force, @@ -374,6 +377,30 @@ def merge_forces( return new_universe + @staticmethod + def _build_non_periodic_dimensions(coordinates: np.ndarray) -> np.ndarray: + """Build a box for a trajectory without box dimensions. + + The axes and covariance code needs a box for the minimum image + convention and unwrapping. The box is cubic and more than twice the + largest interatomic distance, so neither step changes the coordinates. + + Args: + coordinates: Positions with shape ``(n_frames, n_atoms, 3)``. + + Returns: + Array of shape ``(n_frames, 6)`` holding the box lengths and 90 + degree angles for every frame. + """ + span = float(np.max(np.ptp(coordinates, axis=1))) + length = max(10.0 * span, 100.0) + box = np.array([length, length, length, 90.0, 90.0, 90.0], dtype=np.float32) + logger.warning( + "Trajectory has no box dimensions; treating it as non-periodic " + f"using a {length:.1f} A cubic box." + ) + return np.tile(box, (coordinates.shape[0], 1)) + def _extract_timeseries(self, atomgroup, *, kind: str) -> np.ndarray: """Extract a time series array using explicit frame indexing. diff --git a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py index d298c486..cbe2804a 100644 --- a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py +++ b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py @@ -424,3 +424,39 @@ def test_extract_fragment_atomgroup_returns_lightweight_range_selection() -> Non "index 10:13", updating=False, ) + + +def test_build_non_periodic_dimensions_box_exceeds_twice_largest_extent(): + coordinates = np.zeros((3, 2, 3)) + coordinates[:, 1] = [50.0, 0.0, 0.0] + + out = UniverseOperations._build_non_periodic_dimensions(coordinates) + + assert out.shape == (3, 6) + assert np.all(out[:, :3] > 2 * 50.0) + assert np.allclose(out[:, 3:], 90.0) + + +def test_merge_forces_builds_box_for_non_periodic_trajectory(monkeypatch): + ops = UniverseOperations() + + u = MagicMock() + u.dimensions = None + u_force = MagicMock() + monkeypatch.setattr( + "CodeEntropy.trajectory.mda.mda.Universe", MagicMock(side_effect=[u, u_force]) + ) + ops._extract_timeseries = MagicMock(return_value=np.zeros((2, 2, 3))) + ops._extract_force_timeseries_with_fallback = MagicMock( + return_value=np.ones((2, 2, 3)) + ) + + merged = MagicMock() + monkeypatch.setattr("CodeEntropy.trajectory.mda.mda.Merge", lambda ag: merged) + + ops.merge_forces( + tprfile="tpr", trrfile="trr", forcefile="f.trr", fileformat=None, kcal=False + ) + + dimensions = merged.load_new.call_args.kwargs["dimensions"] + assert dimensions.shape == (2, 6)