Rotate velocities and forces besides positions during trajectory rotation

package/CHANGELOG

@@ -23,6 +23,12 @@ The rules for this file:
  * 2.11.0
 
 Fixes
+ * `AtomGroup.rotate()` and the `rotateby` trajectory transformation now
+   also rotate velocities and forces besides positions. This also affects
+   `MDAnalysis.analysis.align.alignto()` and `AlignTraj`, since they
+   apply their fit via `AtomGroup.rotate()`. Note: this changes existing
+   behavior, as velocities/forces were previously left untouched by
+   rotation. (Issue #5421, PR #5452)
  * Fix FileLock tests for XTC and TRR: lock file is no longer removed (#5382)
  * InterRDF now correctly returns bins in parallel (PR #5344)
  * `Merge()` no longer raises a TypeError on Universes that have a `cmaps`

package/MDAnalysis/core/groups.py

@@ -1601,6 +1601,8 @@ def rotate(self, R, point=(0, 0, 0)):
         rotateby : rotate around given axis and angle
         MDAnalysis.lib.transformations : module of all coordinate transforms
 
+        .. versionchanged: 2.11.0
+        Also rotate velocities and forces if present in Timestep.
         """
         R = np.asarray(R)
         point = np.asarray(point)
@@ -1610,11 +1612,19 @@ def rotate(self, R, point=(0, 0, 0)):
         require_translation = bool(np.count_nonzero(point))
         if require_translation:
             atomgroup.translate(-point)
-        x = atomgroup.universe.trajectory.ts.positions
+        ts = atomgroup.universe.trajectory.ts
+        R_T = R.T
+        x = ts.positions
         idx = atomgroup.indices
-        x[idx] = np.dot(x[idx], R.T)
+        x[idx] = np.dot(x[idx], R_T)
         if require_translation:
             atomgroup.translate(point)
+        if ts.has_velocities:
+            v = ts.velocities
+            v[idx] = np.dot(v[idx], R_T)
+        if ts.has_forces:
+            f = ts.forces
+            f[idx] = np.dot(f[idx], R_T)
 
         return self
 

package/MDAnalysis/transformations/rotate.py

@@ -200,4 +200,8 @@ def _transform(self, ts):
         translation = matrix[:3, 3]
         ts.positions = np.dot(ts.positions, rotation)
         ts.positions += translation
+        if ts.has_velocities:
+            ts.velocities = np.dot(ts.velocities, rotation)
+        if ts.has_forces:
+            ts.forces = np.dot(ts.forces, rotation)
         return ts

testsuite/MDAnalysisTests/analysis/test_align.py

@@ -31,6 +31,7 @@
 import pytest
 from MDAnalysis import SelectionError, SelectionWarning
 from MDAnalysisTests import executable_not_found
+from MDAnalysis.lib import transformations
 from MDAnalysisTests.datafiles import (
     PSF,
     DCD,
@@ -821,3 +822,45 @@ def test_alignto_reorder_atomgroups():
     ref = u.atoms[[3, 2, 1, 0]]
     rmsd = align.alignto(mobile, ref, select="bynum 1-4")
     assert_allclose(rmsd, (0.0, 0.0))
+
+
+def test_alignto_rotates_velocities_and_forces():
+    mobile = mda.Universe.empty(
+        4, trajectory=True, velocities=True, forces=True
+    )
+    reference = mda.Universe.empty(4, trajectory=True)
+
+    mobile.add_TopologyAttr("masses", [1.0, 1.0, 1.0, 1.0])
+    reference.add_TopologyAttr("masses", [1.0, 1.0, 1.0, 1.0])
+
+    ref_pos = np.array(
+        [[1, 0, 0], [0, 1, 0], [-1, 0, 0], [0, -1, 0]], dtype=np.float64
+    )
+    reference.atoms.positions = ref_pos
+
+    angle = 23
+    known_R = transformations.rotation_matrix(np.deg2rad(angle), [-1, 2, -3])[
+        :3, :3
+    ]
+    mobile.atoms.positions = np.dot(ref_pos, known_R.T)
+
+    rng = np.random.RandomState(0)
+    orig_v = rng.random((4, 3))
+    orig_f = rng.random((4, 3))
+    mobile.atoms.velocities = orig_v.copy()
+    mobile.atoms.forces = orig_f.copy()
+
+    mobile_centered = (
+        mobile.atoms.positions - mobile.atoms.center_of_geometry()
+    )
+    ref_centered = ref_pos - ref_pos.mean(axis=0)
+    expected_R, _ = align.rotation_matrix(mobile_centered, ref_centered)
+
+    align.alignto(mobile, reference)
+
+    assert_allclose(
+        mobile.atoms.velocities, np.dot(orig_v, expected_R.T), atol=1e-6
+    )
+    assert_allclose(
+        mobile.atoms.forces, np.dot(orig_f, expected_R.T), atol=1e-6
+    )

testsuite/MDAnalysisTests/core/test_atomgroup.py

@@ -417,6 +417,34 @@ def test_rotateby(self, u, coords):
                 [-2 * np.cos(angle) + 1, -2 * np.sin(angle), 0],
             )
 
+    def test_rotate_velocities_forces(self):
+        u = mda.Universe.empty(
+            2, trajectory=True, velocities=True, forces=True
+        )
+        u.atoms.positions = np.array([[1, 0, 0], [-1, 0, 0]])
+        u.atoms.velocities = np.array([[1, 0, 0], [0, 1, 0]])
+        u.atoms.forces = np.array([[0, 0, 1], [1, 1, 0]])
+
+        orig_v = u.atoms.velocities.copy()
+        orig_f = u.atoms.forces.copy()
+
+        axis = np.array([0, 0, 1])
+        for angle in np.linspace(0, np.pi):
+            R = transformations.rotation_matrix(angle, axis)[:3, :3]
+            u.atoms.velocities = orig_v.copy()
+            u.atoms.forces = orig_f.copy()
+            u.atoms.rotate(R)
+            assert_almost_equal(u.atoms.velocities, np.dot(orig_v, R.T))
+            assert_almost_equal(u.atoms.forces, np.dot(orig_f, R.T))
+
+    def test_rotate_no_velocities_forces_does_not_raise(self):
+        u = mda.Universe.empty(2, trajectory=True)
+        orig_pos = np.array([[1, 0, 0], [-1, 0, 0]])
+        u.atoms.positions = orig_pos.copy()
+        R = transformations.rotation_matrix(1, [-1, 2, -3])[:3, :3]
+        u.atoms.rotate(R)
+        assert_almost_equal(u.atoms.positions, np.dot(orig_pos, R.T))
+
     def test_transform_rotation_only(self, u, coords):
         R = np.eye(3)
         u.atoms.rotate(R)

testsuite/MDAnalysisTests/transformations/test_rotate.py

@@ -182,6 +182,46 @@ def test_rotateby_atomgroup_com_pbc(rotate_universes):
     assert_array_almost_equal(transformed.positions, ref.positions, decimal=6)
 
 
+def test_rotateby_velocities_forces():
+    u = mda.Universe.empty(2, trajectory=True, velocities=True, forces=True)
+    u.atoms.positions = np.array([[1, 0, 0], [-1, 0, 0]])
+    u.atoms.velocities = np.array([[1, 0, 0], [0, 1, 0]])
+    u.atoms.forces = np.array([[0, 0, 1], [1, 1, 0]])
+    ts = u.trajectory.ts
+
+    orig_v = ts.velocities.copy()
+    orig_f = ts.forces.copy()
+
+    axis = [-1, 2, -3]
+    point = [0, 0, 0]
+    angle = 23
+    matrix = rotation_matrix(np.deg2rad(angle), axis, point)
+    rotation = matrix[:3, :3].T
+
+    transformed_ts = rotateby(angle, axis, point=point)(ts)
+
+    assert_array_almost_equal(
+        transformed_ts.velocities, np.dot(orig_v, rotation), decimal=6
+    )
+    assert_array_almost_equal(
+        transformed_ts.forces, np.dot(orig_f, rotation), decimal=6
+    )
+
+
+def test_rotateby_no_velocities_forces_does_not_raise():
+    u = mda.Universe.empty(2, trajectory=True)
+    orig_pos = np.array([[1, 0, 0], [-1, 0, 0]])
+    u.atoms.positions = orig_pos.copy()
+    ts = u.trajectory.ts
+    angle = 5
+    matrix = rotation_matrix(np.deg2rad(angle), [-1, 2, -3], [0, 0, 0])
+    rotation = matrix[:3, :3].T
+    transformed_ts = rotateby(angle, [-1, 2, -3], point=[0, 0, 0])(ts)
+    assert_array_almost_equal(
+        transformed_ts.positions, np.dot(orig_pos, rotation), decimal=6
+    )
+
+
 @pytest.mark.parametrize(
     "ag",
     (