[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,19 @@ 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)
+        self._internal_validate(stepids, cells, symbols, positions, times, velocities, pbc)
         # set symbols as attribute for easier querying
         self.base.attributes.set('symbols', list(symbols))
+        pbc = pbc or [True, True, True]
+        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 +200,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
@@ -264,6 +280,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 +359,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 +385,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 +405,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)