Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 28 additions & 1 deletion CodeEntropy/trajectory/mda.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.

Expand Down
36 changes: 36 additions & 0 deletions tests/unit/CodeEntropy/levels/test_mda_universe_operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading