Neighbour analysis for the entire trajectory

pyiron_atomistics/atomistics/job/atomistic.py

@@ -4,12 +4,10 @@
 
 import copy
 import warnings
-from functools import wraps
-
 import numpy as np
 from ase.io import write as ase_write
-
 from pyiron_atomistics.atomistics.structure.atoms import Atoms
+from pyiron_atomistics.atomistics.structure.neighbors import NeighborsTraj
 from pyiron_atomistics.atomistics.structure.has_structure import HasStructure
 from pyiron_base import GenericParameters, GenericMaster, GenericJob as GenericJobCore, deprecate
 
@@ -464,6 +462,8 @@ def trajectory(
             snapshot_indices=None, overwrite_positions=None, overwrite_cells=None
     ):
         """
+        Returns a `Trajectory` instance containing the necessary information to describe the evolution of the atomic
+        structure during the atomistic simulation.
 
         Args:
             stride (int): The trajectories are generated with every 'stride' steps
@@ -477,7 +477,7 @@ def trajectory(
                                                  have the same length of `overwrite_positions`
 
         Returns:
-            pyiron.atomistics.job.atomistic.Trajectory: Trajectory instance
+            pyiron_atomistics.atomistics.job.atomistic.Trajectory: Trajectory instance
 
         """
         cells = self.output.cells
@@ -598,6 +598,43 @@ def write_traj(
             **kwargs
         )
 
+    def get_neighbors_snapshots(self, snapshot_indices=None, num_neighbors=12, **kwargs):
+        """
+        Get the neighbors only for the required snapshots from the trajectory
+
+        Args:
+            snapshot_indices (list/numpy.ndarray): Snapshots for which the the neighbors need to be computed
+                                                   (eg. [1, 5, 10,..., 100]
+            num_neighbors (int): The cutoff for the number of neighbors
+            **kwargs (dict): Additional arguments to be passed to the `get_neighbors()` routine
+                             (eg. cutoff_radius, norm_order , etc.)
+
+        Returns:
+            pyiron_atomistics.atomistics.structure.neighbors.NeighborsTraj: `NeighborsTraj` instances
+                                                                             containing the neighbor information.
+        """
+        return self.trajectory().get_neighbors_snapshots(snapshot_indices=snapshot_indices,
+                                                         num_neighbors=num_neighbors, **kwargs)
+
+    def get_neighbors(self, start=0, stop=-1, stride=1, num_neighbors=12, **kwargs):
+        """
+        Get the neighbors for a given section of the trajectory
+
+        Args:
+            start (int): Start point of the slice of the trajectory to be sampled
+            stop (int): End point of of the slice of the trajectory to be sampled
+            stride (int): Samples the snapshots evert `stride` steps
+            num_neighbors (int): The cutoff for the number of neighbors
+            **kwargs (dict): Additional arguments to be passed to the `get_neighbors()` routine
+                             (eg. cutoff_radius, norm_order , etc.)
+
+        Returns:
+            pyiron_atomistics.atomistics.structure.neighbors.NeighborsTraj: `NeighborsTraj` instances
+                                                                             containing the neighbor information.
+        """
+        return self.trajectory().get_neighbors(start=start, stop=stop, stride=stride,
+                                               num_neighbors=num_neighbors, **kwargs)
+
     # Compatibility functions
     @deprecate("Use get_structure()")
     def get_final_structure(self):
@@ -744,6 +781,50 @@ def __getitem__(self, item):
     def __len__(self):
         return len(self._positions)
 
+    def get_neighbors_snapshots(self, snapshot_indices=None, num_neighbors=12, **kwargs):
+        """
+        Get the neighbors only for the required snapshots from the trajectory
+
+        Args:
+            snapshot_indices (list/numpy.ndarray): Snapshots for which the the neighbors need to be computed
+                                                   (eg. [1, 5, 10,..., 100]
+            num_neighbors (int): The cutoff for the number of neighbors
+            **kwargs (dict): Additional arguments to be passed to the `get_neighbors()` routine
+                             (eg. cutoff_radius, norm_order , etc.)
+
+        Returns:
+            pyiron_atomistics.atomistics.structure.neighbors.NeighborsTraj: `NeighborsTraj` instances
+                                                                             containing the neighbor information.
+        """
+        if snapshot_indices is None:
+            snapshot_indices = np.arange(len(self), dtype=int)
+
+        n_obj = NeighborsTraj(init_structure=self._structure,
+                              positions=self._positions[snapshot_indices],
+                              cells=self._cells[snapshot_indices],
+                              num_neighbors=num_neighbors, **kwargs)
+        n_obj.compute_neighbors()
+        return n_obj
+
+    def get_neighbors(self, start=0, stop=-1, stride=1, num_neighbors=12, **kwargs):
+        """
+        Get the neighbors for a given section of the trajectory
+
+        Args:
+            start (int): Start point of the slice of the trajectory to be sampled
+            stop (int): End point of of the slice of the trajectory to be sampled
+            stride (int): Samples the snapshots evert `stride` steps
+            num_neighbors (int): The cutoff for the number of neighbors
+            **kwargs (dict): Additional arguments to be passed to the `get_neighbors()` routine
+                             (eg. cutoff_radius, norm_order , etc.)
+
+        Returns:
+            pyiron_atomistics.atomistics.structure.neighbors.NeighborsTraj: `NeighborsTraj` instances
+                                                                             containing the neighbor information.
+        """
+        snapshot_indices = np.arange(len(self))[start:stop:stride]
+        return self.get_neighbors_snapshots(snapshot_indices=snapshot_indices, num_neighbors=num_neighbors, **kwargs)
+
 
 class GenericInput(GenericParameters):
     def __init__(self, input_file_name=None, table_name="generic"):
@@ -889,3 +970,4 @@ def __dir__(self):
             return hdf5_path.list_nodes()
         else:
             return []
+

pyiron_atomistics/atomistics/structure/neighbors.py

@@ -1047,3 +1047,147 @@ def get_cluster(dist_vec, ind_vec, prec=prec):
             ind_shell.append(ia_shells_dict)
         return ind_shell
 
+
+class NeighborsTraj:
+
+    """
+    This class generates the neighbors for a given atomistic trajectory. The assumption here is that the trajectory is
+    canonical (no change in the number and types of species throughout the trajectory). The resulting indices,
+    distances, and vectors are stored as numpy arrays.
+
+    """
+
+    def __init__(self, init_structure=None, positions=None, cells=None, num_neighbors=12, **kwargs):
+        """
+
+        Args:
+            init_structure (pyiron_atomistics.atomistics.structure.atoms.Atoms): Any given structure of the trajectory
+            positions (numpy.ndarray): The cartesian positions of the trajectories
+            cells (numpy.ndarray/None): The varying cell shapes
+            num_neighbors (int): The cutoff for the number of neighbors
+            **kwargs (dict): Additional arguments to be passed to the `get_neighbors()` routine
+                             (eg. cutoff_radius, norm_order , etc.)
+        """
+        self._init_structure = init_structure
+        self._neighbor_indices = None
+        self._neighbor_distances = None
+        self._neighbor_vectors = None
+        self._positions = positions
+        self._cells = cells
+        self._num_neighbors = num_neighbors
+        self._get_neighbors_kwargs = kwargs
+
+    @property
+    def neighbor_indices(self):
+        """
+        Neighbour indices (excluding itself) of each atom computed using the get_neighbors_traj() method
+
+        Returns:
+
+            numpy.ndarray: An array of dimension N_steps / stride x N_atoms x N_neighbors
+
+        """
+        return self._neighbor_indices
+
+    @property
+    def neighbor_distances(self):
+        """
+        Neighbour distances (excluding itself) of each atom computed using the get_neighbors_traj() method
+
+        Returns:
+
+            numpy.ndarray: An array of dimension N_steps / stride x N_atoms x N_neighbors
+
+        """
+        return self._neighbor_distances
+
+    @property
+    def neighbor_vectors(self):
+        """
+        Neighbour vectors (excluding itself) of each atom computed using the get_neighbors_traj() method
+
+        Returns:
+
+            numpy.ndarray: An array of dimension N_steps / stride x N_atoms x N_neighbors x 3
+
+        """
+        return self._neighbor_vectors
+
+    @property
+    def num_neighbors(self):
+        """
+        The maximum number of neighbors to be computed
+
+        Returns:
+
+            int: The max number of neighbors
+        """
+        return self._num_neighbors
+
+    def compute_neighbors(self):
+        """
+        Compute the neighbors across the trajectory
+        """
+        self._neighbor_indices, self._neighbor_distances, self._neighbor_vectors = \
+            _get_neighbors(self._init_structure,
+                           positions=self._positions,
+                           cells=self._cells,
+                           num_neighbors=self._num_neighbors, **self._get_neighbors_kwargs)
+
+    def to_hdf(self, hdf, group_name="neighbors_traj"):
+        """
+        Store this `NeighborTraj` instance in a .h5 file.
+
+        Args:
+            hdf (pyiron_base.generic.hdfio.FileHDFio): HDF5 group or file
+            group_name (str): Name of the HDF5 group
+
+        """
+        with hdf.open(group_name) as h_g:
+            h_g["TYPE"] = str(type(self))
+            h_g["neighbor_vectors"] = self._neighbor_vectors
+            h_g["neighbor_distances"] = self._neighbor_distances
+            h_g["neighbor_indices"] = self._neighbor_indices
+            h_g["positions"] = self._positions
+            h_g["cells"] = self._cells
+            h_g["num_neighbors"] = self._num_neighbors
+            h_g["get_neighbors_kwargs"] = self._get_neighbors_kwargs
+            self._init_structure.to_hdf(h_g, group_name="structure")
+
+    def from_hdf(self, hdf, group_name="neighbors_traj"):
+        """
+        Recreate the `NeighborTraj` instance stored in a .h5 file.
+
+        Args:
+            hdf (pyiron_base.generic.hdfio.FileHDFio): HDF5 group or file
+            group_name (str): Name of the HDF5 group
+
+        """
+        with hdf.open(group_name) as h_g:
+            self._neighbor_vectors = h_g["neighbor_vectors"]
+            self._neighbor_distances = h_g["neighbor_distances"]
+            self._neighbor_indices = h_g["neighbor_indices"]
+            self._positions = h_g["positions"]
+            self._cells = h_g["cells"]
+            self._num_neighbors = h_g["num_neighbors"]
+            self._get_neighbors_kwargs = h_g["get_neighbors_kwargs"]
+            self._init_structure = h_g["structure"].to_object()
+
+
+def _get_neighbors(struct, positions, cells=None, num_neighbors=20, **kwargs):
+    [n_steps, n_atoms, _] = positions.shape
+    indices, distances = [np.zeros((n_steps, n_atoms, num_neighbors)) for _ in range(2)]
+    vecs = np.zeros((n_steps, n_atoms, num_neighbors, 3))
+    struct = struct.copy()
+    for t, pos in enumerate(positions):
+        struct.positions = pos
+        # Do cell lookup only if its shape changes
+        if cells is not None:
+            struct.cell = cells[t]
+        # Change the `allow_ragged` based on the changes in get_neighbors()
+        neigh = struct.get_neighbors(num_neighbors=num_neighbors, allow_ragged=False, **kwargs)
+        indices[t, :, :] = neigh.indices
+        distances[t, :, :] = neigh.distances
+        vecs[t, :, :, :] = neigh.vecs
+    indices = np.array(indices, dtype=int)
+    return indices, distances, vecs

tests/lammps/test_base.py

@@ -404,11 +404,29 @@ def test_dump_parser_water(self):
                                            snapshot_indices=snap_indices)._positions,
             orig_pos[snap_indices][:, atom_indices, :]))
         with self.job_water_dump.project_hdf5.open("output/generic") as h_out:
-            h_out["cells"] = np.array([np.zeros((3, 3))] * len(h_out["positions"]))
+            h_out["cells"] = np.repeat([np.array(water.cell)], len(h_out["positions"]), axis=0)
         self.assertTrue(np.array_equal(
             self.job_water_dump.trajectory(atom_indices=atom_indices,
                                            snapshot_indices=snap_indices)._positions,
             orig_pos[snap_indices][:, atom_indices, :]))
+        neigh_traj_obj = self.job_water_dump.get_neighbors()
+        self.assertTrue(np.allclose(np.linalg.norm(neigh_traj_obj.neighbor_vectors, axis=-1),
+                                    neigh_traj_obj.neighbor_distances))
+        h_indices = self.job_water_dump.structure.select_index("H")
+        o_indices = self.job_water_dump.structure.select_index("O")
+        self.assertLessEqual(neigh_traj_obj.neighbor_distances[:, o_indices, :2].max(), 1.2)
+        self.assertGreaterEqual(neigh_traj_obj.neighbor_distances[:, o_indices, :2].min(), 0.8)
+        self.assertTrue(np.alltrue([np.in1d(np.unique(ind_mat.flatten()), h_indices) for ind_mat in
+                                    neigh_traj_obj.neighbor_indices[:, o_indices, :2]]))
+        neigh_traj_obj_snaps = self.job_water_dump.get_neighbors_snapshots(snapshot_indices=[2, 3, 4])
+        self.assertTrue(np.allclose(neigh_traj_obj.neighbor_vectors[2:], neigh_traj_obj_snaps.neighbor_vectors))
+        neigh_traj_obj.to_hdf(self.job_water_dump.project_hdf5)
+        neigh_traj_obj_loaded = self.job_water_dump["neighbors_traj"].to_object()
+        self.assertEqual(neigh_traj_obj._init_structure, neigh_traj_obj_loaded._init_structure)
+        self.assertEqual(neigh_traj_obj._num_neighbors, neigh_traj_obj_loaded._num_neighbors)
+        self.assertTrue(np.allclose(neigh_traj_obj.neighbor_indices, neigh_traj_obj_loaded.neighbor_indices))
+        self.assertTrue(np.allclose(neigh_traj_obj.neighbor_distances, neigh_traj_obj_loaded.neighbor_distances))
+        self.assertTrue(np.allclose(neigh_traj_obj.neighbor_vectors, neigh_traj_obj_loaded.neighbor_vectors))
 
     def test_dump_parser(self):
         structure = Atoms(