[WIP] 🐛 TrajectoryData: Add pbc

Author: mbercx

src/aiida/orm/nodes/data/array/trajectory.py

@@ -8,6 +8,8 @@
 ###########################################################################
 """AiiDA class to deal with crystal structure trajectories."""
 
+from __future__ import annotations
+
 import collections.abc
 from typing import List
 
@@ -33,7 +35,7 @@ def __init__(self, structurelist=None, **kwargs):
         if structurelist is not None:
             self.set_structurelist(structurelist)
 
-    def _internal_validate(self, stepids, cells, symbols, positions, times, velocities):
+    def _internal_validate(self, stepids, cells, symbols, positions, times, velocities, pbc):
         """Internal function to validate the type and shape of the arrays. See
         the documentation of py:meth:`.set_trajectory` for a description of the
         valid shape and type of the parameters.
@@ -82,8 +84,12 @@ def _internal_validate(self, stepids, cells, symbols, positions, times, velociti
                     'have shape (s,n,3), '
                     'with s=number of steps and n=number of symbols'
                 )
+        if not (isinstance(pbc, (list, tuple)) and len(pbc) == 3 and all(isinstance(val, bool) for val in pbc)):
+            raise ValueError('`pbc` must be a list/tuple of length three with boolean values.')
 
-    def set_trajectory(self, symbols, positions, stepids=None, cells=None, times=None, velocities=None):
+    def set_trajectory(
+        self, symbols, positions, stepids=None, cells=None, times=None, velocities=None, pbc: None | list | tuple = None
+    ):
         r"""Store the whole trajectory, after checking that types and dimensions
         are correct.
 
@@ -131,14 +137,20 @@ def set_trajectory(self, symbols, positions, stepids=None, cells=None, times=Non
         :param velocities: if specified, must be a float array with the same
                       dimensions of the ``positions`` array.
                       The array contains the velocities in the atoms.
+        :param pbc: periodic boundary conditions of the structure. Should be a list of
+            length three with booleans indicating if the structure is periodic in that
+            direction. The same periodic boundary conditions are set for each step.
 
         .. todo :: Choose suitable units for velocities
         """
         import numpy
 
-        self._internal_validate(stepids, cells, symbols, positions, times, velocities)
-        # set symbols as attribute for easier querying
+        pbc = pbc or [True, True, True]
+
+        self._internal_validate(stepids, cells, symbols, positions, times, velocities, pbc)
+        # set symbols/pbc as attributes for easier querying
         self.base.attributes.set('symbols', list(symbols))
+        self.base.attributes.set('pbc', list(pbc))
         self.set_array('positions', positions)
         if stepids is not None:  # use input stepids
             self.set_array('steps', stepids)
@@ -189,7 +201,12 @@ def set_structurelist(self, structurelist):
                 raise ValueError('Symbol lists have to be the same for all of the supplied structures')
         symbols = list(symbols_first)
         positions = numpy.array([[list(s.position) for s in x.sites] for x in structurelist])
-        self.set_trajectory(stepids=stepids, cells=cells, symbols=symbols, positions=positions)
+        pbc_set = {structure.pbc for structure in structurelist}
+        if len(pbc_set) == 1:
+            pbc = pbc_set.pop()
+        else:
+            raise ValueError(f'All structures should have the same `pbc`, found: {pbc_set}')
+        self.set_trajectory(stepids=stepids, cells=cells, symbols=symbols, positions=positions, pbc=pbc)
 
     def _validate(self):
         """Verify that the required arrays are present and that their type and
@@ -206,6 +223,7 @@ def _validate(self):
                 self.get_positions(),
                 self.get_times(),
                 self.get_velocities(),
+                self.pbc,
             )
         # Should catch TypeErrors, ValueErrors, and KeyErrors for missing arrays
         except Exception as exception:
@@ -264,6 +282,11 @@ def symbols(self) -> List[str]:
         """
         return self.base.attributes.get('symbols')
 
+    @property
+    def pbc(self) -> list[bool]:
+        """Return the list of periodic boundary conditions."""
+        return self.base.attributes.get('pbc')
+
     def get_positions(self):
         """Return the array of positions, if it has already been set.
 
@@ -338,7 +361,7 @@ def get_step_data(self, index):
             cell = cells[index, :, :]
         else:
             cell = None
-        return (self.get_stepids()[index], time, cell, self.symbols, self.get_positions()[index, :, :], vel)
+        return (self.get_stepids()[index], time, cell, self.symbols, self.get_positions()[index, :, :], vel, self.pbc)
 
     def get_step_structure(self, index, custom_kinds=None):
         """Return an AiiDA :py:class:`aiida.orm.nodes.data.structure.StructureData` node
@@ -364,7 +387,7 @@ def get_step_structure(self, index, custom_kinds=None):
         from aiida.orm.nodes.data.structure import Kind, Site, StructureData
 
         # ignore step, time, and velocities
-        _, _, cell, symbols, positions, _ = self.get_step_data(index)
+        _, _, cell, symbols, positions, _, pbc = self.get_step_data(index)
 
         if custom_kinds is not None:
             kind_names = []
@@ -384,7 +407,7 @@ def get_step_structure(self, index, custom_kinds=None):
                     'passed {}, but the symbols are {}'.format(sorted(kind_names), sorted(symbols))
                 )
 
-        struc = StructureData(cell=cell)
+        struc = StructureData(cell=cell, pbc=pbc)
         if custom_kinds is not None:
             for _k in custom_kinds:
                 struc.append_kind(_k)

tests/orm/nodes/data/test_trajectory.py

@@ -10,6 +10,7 @@
 def trajectory_data():
     """Return a dictionary of data to create a ``TrajectoryData``."""
     symbols = ['H'] * 5 + ['Cl'] * 5
+    pbc = [True, True, True]
     stepids = np.arange(1000, 3000, 10)
     times = stepids * 0.01
     positions = np.arange(6000, dtype=float).reshape((200, 10, 3))
@@ -23,6 +24,7 @@ def trajectory_data():
         'cells': cells,
         'times': times,
         'velocities': velocities,
+        'pbc': pbc,
     }
 
 
@@ -104,23 +106,25 @@ def test_trajectory_get_step_data(self, trajectory_data):
         """Test the ``get_step_data`` method."""
         trajectory = TrajectoryData()
         trajectory.set_trajectory(**trajectory_data)
-        stepid, time, cell, symbols, positions, velocities = trajectory.get_step_data(-2)
+        stepid, time, cell, symbols, positions, velocities, pbc = trajectory.get_step_data(-2)
         assert stepid == trajectory_data['stepids'][-2]
         assert time == trajectory_data['times'][-2]
-        np.array_equal(cell, trajectory_data['cells'][-2, :, :])
-        np.array_equal(symbols, trajectory_data['symbols'])
-        np.array_equal(positions, trajectory_data['positions'][-2, :, :])
-        np.array_equal(velocities, trajectory_data['velocities'][-2, :, :])
+        assert np.array_equal(cell, trajectory_data['cells'][-2, :, :])
+        assert np.array_equal(symbols, trajectory_data['symbols'])
+        assert np.array_equal(pbc, trajectory_data['pbc'])
+        assert np.array_equal(positions, trajectory_data['positions'][-2, :, :])
+        assert np.array_equal(velocities, trajectory_data['velocities'][-2, :, :])
 
     def test_trajectory_get_step_data_empty(self, trajectory_data):
         """Test the `get_step_data` method when some arrays are not defined."""
         trajectory = TrajectoryData()
         trajectory.set_trajectory(symbols=trajectory_data['symbols'], positions=trajectory_data['positions'])
-        stepid, time, cell, symbols, positions, velocities = trajectory.get_step_data(3)
+        stepid, time, cell, symbols, positions, velocities, pbc = trajectory.get_step_data(3)
         assert stepid == 3
         assert time is None
         assert cell is None
         assert np.array_equal(symbols, trajectory_data['symbols'])
+        assert np.array_equal(pbc, trajectory_data['pbc'])
         assert np.array_equal(positions, trajectory_data['positions'][3, :, :])
         assert velocities is None
 
@@ -141,3 +145,14 @@ def test_trajectory_get_step_structure(self, trajectory_data):
 
         with pytest.raises(IndexError):
             trajectory.get_step_structure(500)
+
+    def test_trajectory_pbc(self):
+        """Test the `pbc` for the `TrajectoryData` ."""
+        trajectory = TrajectoryData()
+        structure = StructureData(cell=None, pbc=[False, False, False])
+        structure.append_atom(position=[0.0, 0.0, 0.0], symbols='H')
+
+        trajectory = TrajectoryData(structurelist=(structure,))
+
+        trajectory.get_step_structure(0).store()  # Verify that the `StructureData` can be stored
+        assert trajectory.get_step_structure(0).pbc == structure.pbc