Merge pull request #561 from ISISNeutronMuon/Replace-atomic-trajector…

…y-cython-version-with-numpy-implementation

Replace atomic trajectory cython version with numpy implementation

.gitignore

@@ -8,6 +8,8 @@
 .settings
 .vscode
 *.pyc
+mdanse_env/
+.venv/
 sandbox/
 
 # Distribution / packaging

MDANSE/Src/MDANSE/MolecularDynamics/MockTrajectory.py

@@ -23,12 +23,13 @@
 from MDANSE.Chemistry import ATOMS_DATABASE
 from MDANSE.Framework.Units import measure
 from MDANSE.Chemistry.ChemicalEntity import Atom, ChemicalSystem
-from MDANSE.Extensions import atomic_trajectory, com_trajectory
+from MDANSE.Extensions import com_trajectory
 from MDANSE.MolecularDynamics.Configuration import (
     PeriodicRealConfiguration,
     RealConfiguration,
     _Configuration,
 )
+from MDANSE.MolecularDynamics.TrajectoryUtils import atomic_trajectory
 from MDANSE.MolecularDynamics.UnitCell import UnitCell
 
 
@@ -471,7 +472,7 @@ def read_atomic_trajectory(
                     for fnum in range(first, last, step)
                 ]
             )
-            atomic_traj = atomic_trajectory.atomic_trajectory(
+            atomic_traj = atomic_trajectory(
                 coords, direct_cells, inverse_cells, box_coordinates
             )
             return atomic_traj

MDANSE/Src/MDANSE/MolecularDynamics/TrajectoryUtils.py

@@ -487,3 +487,32 @@ def sorted_atoms(
         return atoms
     else:
         return [getattr(at, attribute) for at in atoms]
+
+
+def atomic_trajectory(config, cell, rcell, box_coordinates=False):
+    """For the coordinates of a specific atom, remove all unit cell
+    jumps.
+
+    Parameters
+    ----------
+    config : np.ndarray
+        The coordinates for a specific atoms.
+    cell : np.ndarray
+        The direct matrices.
+    rcell : np.ndarray
+        The inverse matrices.
+    box_coordinates : bool
+        Returns the coordinates in fractional coordinates if true.
+
+    Returns
+    -------
+    np.ndarray
+        The input config but the unit cell jumps removed.
+    """
+    trajectory = np.einsum("ij,ikj->ik", config, rcell)
+    sdxyz = trajectory[1:, :] - trajectory[:-1, :]
+    sdxyz -= np.cumsum(np.round(sdxyz), axis=0)
+    trajectory[1:, :] = trajectory[:-1, :] + sdxyz
+    if not box_coordinates:
+        trajectory = np.einsum("ij,ikj->ik", trajectory, cell)
+    return trajectory

MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py

@@ -23,13 +23,14 @@
 from MDANSE.Framework.Units import measure
 from MDANSE.Chemistry import ATOMS_DATABASE
 from MDANSE.Chemistry.ChemicalEntity import ChemicalSystem
-from MDANSE.Extensions import atomic_trajectory, com_trajectory
+from MDANSE.Extensions import com_trajectory
 from MDANSE.MolecularDynamics.Configuration import (
     PeriodicRealConfiguration,
     RealConfiguration,
 )
 from MDANSE.MolecularDynamics.TrajectoryUtils import (
     resolve_undefined_molecules_name,
+    atomic_trajectory,
 )
 from MDANSE.MolecularDynamics.UnitCell import UnitCell
 
@@ -493,7 +494,7 @@ def read_atomic_trajectory(
                     for nf in range(first, last, step)
                 ]
             )
-            atomic_traj = atomic_trajectory.atomic_trajectory(
+            atomic_traj = atomic_trajectory(
                 coords, direct_cells, inverse_cells, box_coordinates
             )
             return atomic_traj

MDANSE/Src/MDANSE/Trajectory/MdanseTrajectory.py

@@ -21,12 +21,15 @@
 from MDANSE.MLogging import LOG
 from MDANSE.Chemistry import ATOMS_DATABASE
 from MDANSE.Chemistry.ChemicalEntity import ChemicalSystem
-from MDANSE.Extensions import atomic_trajectory, com_trajectory
+from MDANSE.Extensions import com_trajectory
 from MDANSE.MolecularDynamics.Configuration import (
     PeriodicRealConfiguration,
     RealConfiguration,
 )
-from MDANSE.MolecularDynamics.TrajectoryUtils import resolve_undefined_molecules_name
+from MDANSE.MolecularDynamics.TrajectoryUtils import (
+    resolve_undefined_molecules_name,
+    atomic_trajectory,
+)
 from MDANSE.MolecularDynamics.UnitCell import UnitCell
 
 
@@ -364,9 +367,19 @@ def read_atomic_trajectory(
         coords = grp["coordinates"][first:last:step, index, :].astype(np.float64)
 
         if self._unit_cells is not None:
-            direct_cells = np.array([uc.transposed_direct for uc in self._unit_cells])
-            inverse_cells = np.array([uc.transposed_inverse for uc in self._unit_cells])
-            atomic_traj = atomic_trajectory.atomic_trajectory(
+            direct_cells = np.array(
+                [
+                    self._unit_cells[nf].transposed_direct
+                    for nf in range(first, last, step)
+                ]
+            )
+            inverse_cells = np.array(
+                [
+                    self._unit_cells[nf].transposed_inverse
+                    for nf in range(first, last, step)
+                ]
+            )
+            atomic_traj = atomic_trajectory(
                 coords, direct_cells, inverse_cells, box_coordinates
             )
             return atomic_traj

MDANSE/Tests/UnitTests/Analysis/test_average_structure.py

@@ -22,14 +22,16 @@ def trajectory():
     trajectory = HDFTrajectoryInputData(short_traj)
     yield trajectory
 
-units = ['Angstrom', 'Bohr', 'nm', 'pm']
-formats = ['vasp', 'xyz', 'turbomole', 'abinit-in']
+
+units = ["Angstrom", "Bohr", "nm", "pm"]
+formats = ["vasp", "xyz", "turbomole", "abinit-in"]
 inputs = []
 for u in units:
     for f in formats:
-        inputs.append((u,f))
+        inputs.append((u, f))
+
 
-@pytest.mark.parametrize('output_unit,output_format', inputs)
+@pytest.mark.parametrize("output_unit,output_format", inputs)
 def test_temperature(trajectory, output_unit, output_format):
     temp_name = tempfile.mktemp()
     parameters = {}
@@ -44,5 +46,3 @@ def test_temperature(trajectory, output_unit, output_format):
     assert path.exists(temp_name)
     assert path.isfile(temp_name)
     os.remove(temp_name)
-
-

MDANSE/Tests/UnitTests/Analysis/test_dynamics.py

@@ -59,6 +59,7 @@ def test_vacf(trajectory, interp_order, normalise):
     assert path.isfile(temp_name + ".log")
     os.remove(temp_name + ".log")
 
+
 def test_pps(trajectory):
     temp_name = tempfile.mktemp()
     parameters = {}
@@ -75,6 +76,7 @@ def test_pps(trajectory):
     assert path.isfile(temp_name + ".log")
     os.remove(temp_name + ".log")
 
+
 ################################################################
 # Job parameters                                               #
 ################################################################
@@ -124,11 +126,13 @@ def parameters():
     ]:
         for rm in [("single-core", 1), ("multicore", -4)]:
             for of in ["MDAFormat", "TextFormat"]:
-                total_list.append((tp,jt, rm, of))
+                total_list.append((tp, jt, rm, of))
 
 
 @pytest.mark.parametrize("traj_path,job_type,running_mode,output_format", total_list)
-def test_dynamics_analysis(parameters, traj_path, job_type, running_mode, output_format):
+def test_dynamics_analysis(
+    parameters, traj_path, job_type, running_mode, output_format
+):
     temp_name = tempfile.mktemp()
     parameters["trajectory"] = traj_path
     parameters["running_mode"] = running_mode

MDANSE/Tests/UnitTests/Analysis/test_mdmc_h5md.py

@@ -44,4 +44,3 @@ def test_temperature_nonzero(trajectory, interp_order):
         temperature = np.array(results["/temperature"])
     os.remove(temp_name + ".mda")
     assert np.all(temperature > 0.0)
-

MDANSE/Tests/UnitTests/Analysis/test_structure.py

@@ -75,7 +75,9 @@ def parameters():
 
 
 @pytest.mark.parametrize("traj_path,job_type,running_mode,output_format", total_list)
-def test_structure_analysis(parameters, traj_path, job_type, running_mode, output_format):
+def test_structure_analysis(
+    parameters, traj_path, job_type, running_mode, output_format
+):
     temp_name = tempfile.mktemp()
     parameters["trajectory"] = traj_path
     parameters["running_mode"] = running_mode
@@ -92,4 +94,4 @@ def test_structure_analysis(parameters, traj_path, job_type, running_mode, outpu
         os.remove(temp_name + "_text.tar")
     assert path.exists(temp_name + ".log")
     assert path.isfile(temp_name + ".log")
-    os.remove(temp_name + ".log")
+    os.remove(temp_name + ".log")