Merge pull request #89 from Harry-Rich/master

Added improved method for Centre of mass calculation
bjmorgan · Jan 24, 2025 · a402f08 · a402f08
2 parents 4a8692c + 17785c9
commit a402f08
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/kinisi/parser.py b/kinisi/parser.py
@@ -207,7 +207,8 @@ def get_disps(self,
             elif self.sampling == 'multi-origin':
                 disp = np.concatenate([
                     drift_corrected[self.indices, np.newaxis, time_interval - 1],
-                    np.subtract(drift_corrected[self.indices, time_interval:], drift_corrected[self.indices, :-time_interval])
+                    np.subtract(drift_corrected[self.indices, time_interval:],
+                                drift_corrected[self.indices, :-time_interval])
                 ],
                                       axis=1)
                 if np.multiply(*disp[:, ::time_interval].shape[:2]) <= 1:
@@ -568,6 +569,7 @@ def __init__(self,
 
         structure, coords, latt, volume = self.get_structure_coords_latt(universe, sub_sample_atoms, sub_sample_traj,
                                                                          progress)
+
         if specie is not None:
             indices = self.get_indices(structure, specie, framework_indices)
         elif isinstance(specie_indices, (list, tuple)):
@@ -674,7 +676,8 @@ def _get_molecules(structure: "ase.atoms.Atoms" or "pymatgen.core.structure.Stru
                    framework_indices) -> Tuple[np.ndarray, np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
     Determine framework and non-framework indices for an :py:mod:`ase` or :py:mod:`pymatgen` or :py:mod:`MDAnalysis` compatible file when specie_indices are provided and contain multiple molecules. Warning: This function changes the structure without changing the object.
-        
+    Calculates the centre of mass of provided particle groups using the pseudo-centre of mass approach (paper to come).
+
     :param structure: Initial structure.
     :param coords: fractional coordinates for all atoms.
     :param indices: indices for the atoms in the molecules in the trajectory used in the calculation 
@@ -717,10 +720,15 @@ def _get_molecules(structure: "ase.atoms.Atoms" or "pymatgen.core.structure.Stru
     zeta = np.sin(theta)
     xi_bar = np.average(xi, axis=-2, weights=weights)
     zeta_bar = np.average(zeta, axis=-2, weights=weights)
-    theta_bar = np.arctan2(zeta_bar, xi_bar)
+    theta_bar = np.arctan2(-zeta_bar, -xi_bar) + np.pi
     new_s_coords = theta_bar / (2 * np.pi)
 
-    new_coords = np.concatenate((new_s_coords, sq_coords[:, drift_indices]), axis=1)
+    # Implementation of pseudo-centre of mass approach to centre of mass calculation (paper to come).
+    pseudo_com_recentering = ((s_coords - (new_s_coords + 0.5)[:, :, np.newaxis]) % 1)
+    com_pseudo_space = np.average(pseudo_com_recentering, weights=masses, axis=2)
+    corrected_com = (com_pseudo_space + (new_s_coords + 0.5)) % 1
+
+    new_coords = np.concatenate((corrected_com, sq_coords[:, drift_indices]), axis=1)
     new_indices = list(range(n_molecules))
     new_drift_indices = list(range(n_molecules, n_molecules + len(drift_indices)))
 

diff --git a/kinisi/tests/test_parser.py b/kinisi/tests/test_parser.py
@@ -131,7 +131,7 @@ def test_pymatgen_init_with_COG(self):
         assert_almost_equal(data.time_step, 1)
         assert_almost_equal(data.step_skip, 1)
         assert_equal(data.indices, [0])
-        assert_almost_equal(data.coords_check, [[[0.269634905, 0.262183827, 0.2]]])
+        assert_almost_equal(data.coords_check, [[[0.2733333, 0.2666667, 0.2      ]]])
 
     def test_pymatgen_init_with_COM(self):
         xd = Xdatcar(os.path.join(os.path.dirname(kinisi.__file__), 'tests/inputs/example_center.XDATCAR'))
@@ -146,7 +146,7 @@ def test_pymatgen_init_with_COM(self):
         assert_almost_equal(data.time_step, 1)
         assert_almost_equal(data.step_skip, 1)
         assert_equal(data.indices, [0])
-        assert_almost_equal(data.coords_check, [[[0.382421597, 0.2087361, 0.2]]])
+        assert_almost_equal(data.coords_check, [[[0.3788889, 0.2111111, 0.2      ]]])
 
     def test_pymatgen_init_with_framwork_indices(self):
         xd = Xdatcar(os.path.join(os.path.dirname(kinisi.__file__), 'tests/inputs/example_drift.XDATCAR'))
@@ -200,7 +200,7 @@ def test_ase_init_with_COG(self):
         assert_almost_equal(data.time_step, 1)
         assert_almost_equal(data.step_skip, 1)
         assert_equal(data.indices, [0])
-        assert_almost_equal(data.coords_check, [[[0.269634905, 0.262183827, 0.2]]])
+        assert_almost_equal(data.coords_check, [[[0.2733333, 0.2666667, 0.2]]])
 
     def test_ase_init_with_COM(self):
         traj = Trajectory(os.path.join(os.path.dirname(kinisi.__file__), 'tests/inputs/example_ase_center.traj'))
@@ -215,7 +215,7 @@ def test_ase_init_with_COM(self):
         assert_almost_equal(data.time_step, 1)
         assert_almost_equal(data.step_skip, 1)
         assert_equal(data.indices, [0])
-        assert_almost_equal(data.coords_check, [[[0.382421597, 0.2087361, 0.2]]])
+        assert_almost_equal(data.coords_check, [[[0.3788889, 0.2111111, 0.2]]])
 
     def test_ase_init_with_framwork_indices(self):
         traj = Trajectory(os.path.join(os.path.dirname(kinisi.__file__), 'tests/inputs/example_ase_drift.traj'))
@@ -266,7 +266,7 @@ def test_mda_init_with_molecules(self):
         data = parser.MDAnalysisParser(xd, **da_params)
         assert_almost_equal(data.time_step, 0.005)
         assert_almost_equal(data.step_skip, 250)
-        assert_almost_equal(data.coords_check[0, 0], [0.46465184, 0.03090944, 0.4023758])
+        assert_almost_equal(data.coords_check[0, 0], [0.5169497, 0.1174514, 0.3637794])
         assert_equal(data.indices, list(range(len(molecules))))
 
     def test_mda_init_with_COG(self):
@@ -279,7 +279,7 @@ def test_mda_init_with_COG(self):
         assert_almost_equal(data.time_step, 1)
         assert_almost_equal(data.step_skip, 1)
         assert_equal(data.indices, [0])
-        assert_almost_equal(data.coords_check, [[[0.269634905, 0.262183827, 0.2]]])
+        assert_almost_equal(data.coords_check, [[[0.2733333, 0.2666667, 0.1999999]]])
 
     def test_mda_init_with_COM(self):
         xd = mda.Universe(os.path.join(os.path.dirname(kinisi.__file__), 'tests/inputs/example_LAMMPS_center.data'),
@@ -297,7 +297,7 @@ def test_mda_init_with_COM(self):
         assert_almost_equal(data.time_step, 1)
         assert_almost_equal(data.step_skip, 1)
         assert_equal(data.indices, [0])
-        assert_almost_equal(data.coords_check, [[[0.382421597, 0.2087361, 0.2]]])
+        assert_almost_equal(data.coords_check, [[[0.3788889, 0.2111111, 0.2]]])
 
     def test_mda_init_with_framwork_indices(self):
         xd = mda.Universe(os.path.join(os.path.dirname(kinisi.__file__), 'tests/inputs/example_LAMMPS_drift.data'),
@@ -319,4 +319,4 @@ def test_mda_init_with_framwork_indices(self):
         assert_equal(data_2.drift_indices, list(range(1, 9)))
         print(data_2.dc[0])
         disp_array = [[0.0, 0.0, 0.0], [0.2, 0.2, 0.2], [0.4, 0.4, 0.4], [1, 1, 1]]
-        assert_almost_equal(data_2.dc[0], disp_array, decimal=6)
+        assert_almost_equal(data_2.dc[0], disp_array, decimal=6)