Merge pull request #387 from pyiron/flattraj

Use FlattenedStorage in NeighborsTrajectory

pyiron_atomistics/atomistics/job/atomistic.py

@@ -829,7 +829,7 @@ def get_neighbors_snapshots(self, snapshot_indices=None, num_neighbors=12, **kwa
         if snapshot_indices is None:
             snapshot_indices = np.arange(len(self), dtype=int)
 
-        n_obj = NeighborsTrajectory(self[snapshot_indices], num_neighbors=num_neighbors, **kwargs)
+        n_obj = NeighborsTrajectory(has_structure=self[snapshot_indices], num_neighbors=num_neighbors, **kwargs)
         n_obj.compute_neighbors()
         return n_obj
 

pyiron_atomistics/atomistics/structure/neighbors.py

@@ -6,7 +6,7 @@
 from sklearn.cluster import AgglomerativeClustering
 from scipy.sparse import coo_matrix
 from scipy.special import gamma
-from pyiron_base import DataContainer
+from pyiron_base import DataContainer, FlattenedStorage
 from pyiron_atomistics.atomistics.structure.analyse import get_average_of_unique_labels
 from scipy.spatial.transform import Rotation
 from scipy.special import sph_harm
@@ -1006,32 +1006,36 @@ def get_volume_of_n_sphere_in_p_norm(n=3, p=2):
 
 
 class NeighborsTrajectory(DataContainer):
-
     """
-    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.
-
+    This class generates the neighbors for a given atomistic trajectory. The resulting indices, distances, and vectors
+    are stored as numpy arrays.
     """
 
     def __new__(cls, *args, **kwargs):
         instance = super().__new__(cls, *args, **kwargs)
         object.__setattr__(instance, "_has_structure", None)
         return instance
 
-    def __init__(self, has_structure=None, num_neighbors=12, table_name="neighbors_traj", **kwargs):
+    def __init__(self, init=None, has_structure=None, num_neighbors=12, table_name="neighbors_traj", store=None, **kwargs):
         """
 
         Args:
             has_structure (:class:`.HasStructure`): object containing the structures to compute the neighbors on
             num_neighbors (int): The cutoff for the number of neighbors
-            table_name (str): Table name for the base `DataContainer` (stores this object as a group in a
-                              HDF5 file with this name)
+            table_name (str): Table name for the base `DataContainer` (stores this object as a group in a HDF5 file with
+                              this name)
+            store (FlattenedStorage): internal storage that should be used to store the neighborhood information,
+                                      creates a new one if not provided; if provided and not empty it must be compatible
+                                      with the lengths of the structures in `has_structure`, but this is *not* checked
             **kwargs (dict): Additional arguments to be passed to the `get_neighbors()` routine
                              (eg. cutoff_radius, norm_order , etc.)
         """
-        super().__init__(table_name=table_name)
+        super().__init__(init=init, table_name=table_name)
         self._has_structure = has_structure
+        self._flat_store = store if store is not None else FlattenedStorage()
+        self._flat_store.add_array("indices", dtype=np.int64, shape=(num_neighbors,), per="element")
+        self._flat_store.add_array("distances", dtype=np.float64, shape=(num_neighbors,), per="element")
+        self._flat_store.add_array("vecs", dtype=np.float64, shape=(num_neighbors, 3), per="element")
         self._neighbor_indices = None
         self._neighbor_distances = None
         self._neighbor_vectors = None
@@ -1043,36 +1047,36 @@ def indices(self):
         """
         Neighbour indices (excluding itself) of each atom computed using the get_neighbors_traj() method
 
-        Returns:
-
-            numpy.ndarray/list: An array of dimension N_steps / stride x N_atoms x N_neighbors
+        If the structures have different number of atoms, the array will have -1 on indices that are invalid.
 
+        Returns:
+            numpy.ndarray: An int array of dimension N_steps / stride x N_atoms x N_neighbors
         """
-        return self._neighbor_indices
+        return self._flat_store.get_array_filled("indices")
 
     @property
     def distances(self):
         """
         Neighbour distances (excluding itself) of each atom computed using the get_neighbors_traj() method
 
-        Returns:
-
-            numpy.ndarray/list: An array of dimension N_steps / stride x N_atoms x N_neighbors
+        If the structures have different number of atoms, the array will have NaN on indices that are invalid.
 
+        Returns:
+            numpy.ndarray: A float array of dimension N_steps / stride x N_atoms x N_neighbors
         """
-        return self._neighbor_distances
+        return self._flat_store.get_array_filled("distances")
 
     @property
     def vecs(self):
         """
         Neighbour vectors (excluding itself) of each atom computed using the get_neighbors_traj() method
 
-        Returns:
-
-            numpy.ndarray/list: An array of dimension N_steps / stride x N_atoms x N_neighbors x 3
+        If the structures have different number of atoms, the array will have NaN on indices that are invalid.
 
+        Returns:
+            numpy.ndarray: A float array of dimension N_steps / stride x N_atoms x N_neighbors x 3
         """
-        return self._neighbor_vectors
+        return self._flat_store.get_array_filled("vecs")
 
     @property
     def num_neighbors(self):
@@ -1089,21 +1093,18 @@ def compute_neighbors(self):
         """
         Compute the neighbors across the trajectory
         """
-        self._neighbor_indices, self._neighbor_distances, self._neighbor_vectors = \
-            _get_neighbors(self._has_structure,
-                           num_neighbors=self._num_neighbors, **self._get_neighbors_kwargs)
+        _get_neighbors(self._flat_store, self._has_structure,
+                       num_neighbors=self._num_neighbors, **self._get_neighbors_kwargs)
 
 
-def _get_neighbors(has_structure, num_neighbors=20, **kwargs):
-    n_steps = has_structure.number_of_structures
-    n_atoms = len(has_structure.get_structure(frame=0))
-    indices = np.zeros((n_steps, n_atoms, num_neighbors), dtype=np.int64)
-    distances = np.zeros((n_steps, n_atoms, num_neighbors))
-    vecs = np.zeros((n_steps, n_atoms, num_neighbors, 3))
-    for t, struct in enumerate(has_structure.iter_structures()):
+def _get_neighbors(store, has_structure, num_neighbors=20, **kwargs):
+    for i, struct in enumerate(has_structure.iter_structures()):
         # 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
-    return indices, distances, vecs
+        if i >= len(store):
+            store.add_chunk(len(struct), indices=neigh.indices, distances=neigh.distances, vecs=neigh.vecs)
+        else:
+            store.set_array("indices", i, neigh.indices)
+            store.set_array("distances", i, neigh.distances)
+            store.set_array("vecs", i, neigh.vecs)
+    return store.get_array_filled('indices'), store.get_array_filled('distances'), store.get_array_filled('vecs')

tests/atomistics/structure/test_neighbors_trajectory.py

@@ -0,0 +1,58 @@
+# coding: utf-8
+# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
+# Distributed under the terms of "New BSD License", see the LICENSE file.
+
+import unittest
+import numpy as np
+from pyiron_atomistics.atomistics.structure.neighbors import NeighborsTrajectory
+from pyiron_atomistics.atomistics.structure.structurestorage import StructureStorage
+from pyiron_atomistics.atomistics.structure.factory import StructureFactory
+
+class TestNeighborsTrajectory(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.canonical = StructureStorage()
+        cls.grandcanonical = StructureStorage()
+
+        bulk = StructureFactory().bulk("Fe")
+        for n in range(3):
+            bulk.positions += np.random.normal(scale=.1, size=bulk.positions.shape)
+            cls.canonical.add_structure(bulk)
+            cls.grandcanonical.add_structure(bulk.repeat( (n+1, 1, 1) ))
+
+    def test_canonical(self):
+        traj = NeighborsTrajectory(has_structure=self.canonical)
+        traj.compute_neighbors()
+        self.assertEqual(traj.indices.shape, (3, 1, 12),
+                         "indices from trajectory have wrong shape.")
+        self.assertEqual(traj.distances.shape, (3, 1, 12),
+                         "distances from trajectory have wrong shape.")
+        self.assertEqual(traj.vecs.shape, (3, 1, 12, 3),
+                         "vecs from trajectory have wrong shape.")
+        for i in range(3):
+            neigh = self.canonical.get_structure(i).get_neighbors()
+            self.assertTrue(np.array_equal(traj.indices[i], neigh.indices),
+                            "indices from trajectory and get_neighbors not equal")
+            self.assertTrue(np.array_equal(traj.distances[i], neigh.distances),
+                            "distances from trajectory and get_neighbors not equal")
+            self.assertTrue(np.array_equal(traj.vecs[i], neigh.vecs),
+                            "vecs from trajectory and get_neighbors not equal")
+
+    def test_grandcanonical(self):
+        traj = NeighborsTrajectory(has_structure=self.grandcanonical)
+        traj.compute_neighbors()
+        self.assertEqual(traj.indices.shape, (3, 3, 12),
+                         "indices from trajectory have wrong shape.")
+        self.assertEqual(traj.distances.shape, (3, 3, 12),
+                         "distances from trajectory have wrong shape.")
+        self.assertEqual(traj.vecs.shape, (3, 3, 12, 3),
+                         "vecs from trajectory have wrong shape.")
+        for i in range(3):
+            neigh = self.grandcanonical.get_structure(i).get_neighbors()
+            self.assertTrue(np.array_equal(traj.indices[i][:i+1], neigh.indices),
+                            "indices from trajectory and get_neighbors not equal")
+            self.assertTrue(np.array_equal(traj.distances[i][:i+1], neigh.distances),
+                            "distances from trajectory and get_neighbors not equal")
+            self.assertTrue(np.array_equal(traj.vecs[i][:i+1], neigh.vecs),
+                            "vecs from trajectory and get_neighbors not equal")