Cg particle support

pyproject.toml

@@ -28,6 +28,11 @@ maintainers = [
     { name = "Joseph F. Rudzinski", email = "[email protected]" }
 ]
 license = { file = "LICENSE" }
+# dependencies = [
+#     "nomad-lab>=1.3.0",
+#     "matid>=2.0.0.dev2",
+#     "nomad-simulations@file:///home/bmohr/software/nomad-simulations",
+# ]
 dependencies = [
     "nomad-lab>=1.3.0",
     "matid>=2.0.0.dev2",

src/nomad_simulations/schema_packages/model_system.py

@@ -27,6 +27,7 @@
     from structlog.stdlib import BoundLogger
 
 from nomad_simulations.schema_packages.atoms_state import AtomsState
+from nomad_simulations.schema_packages.particles_state import Particles, ParticlesState
 from nomad_simulations.schema_packages.utils import (
     get_sibling_section,
     is_not_representative,
@@ -492,6 +493,168 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
         self.name = self.m_def.name if self.name is None else self.name
 
 
+class ParticleCell(Cell):
+    """
+    A base section used to specify the atomic cell information of a system.
+    """
+
+    particles_state = SubSection(sub_section=ParticlesState.m_def, repeats=True)
+
+    n_particles = Quantity(
+        type=np.int32,
+        description="""
+        Number of atoms in the atomic cell.
+        """,
+    )
+
+    equivalent_particles = Quantity(
+        type=np.int32,
+        shape=['n_atoms'],
+        description="""
+        List of equivalent atoms as defined in `atoms`. If no equivalent atoms are found,
+        then the list is simply the index of each element, e.g.:
+            - [0, 1, 2, 3] all four atoms are non-equivalent.
+            - [0, 0, 0, 3] three equivalent atoms and one non-equivalent.
+        """,
+    )
+
+    def __init__(self, m_def: 'Section' = None, m_context: 'Context' = None, **kwargs):
+        super().__init__(m_def, m_context, **kwargs)
+        # Set the name of the section
+        self.name = self.m_def.name  #! self.name here is 'ParticleCell'
+
+    def is_equal_cell(self, other) -> bool:
+        """
+        Check if the atomic cell is equal to an`other` atomic cell by comparing the `positions` and
+        the `AtomsState[*].chemical_symbol`.
+        Args:
+            other: The other atomic cell to compare with.
+        Returns:
+            bool: True if the atomic cells are equal, False otherwise.
+        """
+        if not isinstance(other, ParticleCell):
+            return False
+
+        # Compare positions using the parent sections's `__eq__` method
+        if not super().is_equal_cell(other=other):
+            return False
+
+        # Check that the `chemical_symbol` of the atoms in `cell_1` match with the ones in `cell_2`
+        check_positions = self._check_positions(
+            positions_1=self.positions, positions_2=other.positions
+        )
+        try:
+            for particle in check_positions:
+                type_1 = self.particles_state[particle[0]].particle_type
+                type_2 = other.particles_state[particle[1]].particle_type
+                if type_1 != type_2:
+                    return False
+        except Exception:
+            return False
+        return True
+
+    def get_particle_types(self, logger: 'BoundLogger') -> list[str]:
+        """
+        Get the chemical symbols of the atoms in the atomic cell. These are defined on `atoms_state[*].chemical_symbol`.
+
+        Args:
+            logger (BoundLogger): The logger to log messages.
+
+        Returns:
+            list: The list of chemical symbols of the atoms in the atomic cell.
+        """
+        if not self.particles_state:
+            return []
+
+        particle_labels = []
+        for particle_state in self.particles_state:
+            if not particle_state.particle_type:
+                logger.warning('Could not find `ParticlesState[*].particle_type`.')
+                return []
+            particle_labels.append(particle_state.particle_type)
+        return particle_labels
+
+    def to_particles(self, logger: 'BoundLogger') -> Optional[Particles]:
+        """
+        Generates a Particles object with the most basic information from the parsed `ParticleCell`
+        section (labels, periodic_boundary_conditions, positions, and lattice_vectors).
+
+        Args:
+            logger (BoundLogger): The logger to log messages.
+
+        Returns:
+            (Optional[Particles]): The Partilces object with the basic information from the `ParticleCell`.
+        """
+        # Initialize Partilces object with labels
+        particle_labels = self.get_particle_types(logger=logger)
+        particles = Particles(symbols=particle_labels)
+
+        # PBC
+        if self.periodic_boundary_conditions is None:
+            logger.info(
+                'Could not find `ParticleCell.periodic_boundary_conditions`. They will be set to [False, False, False].'
+            )
+            self.periodic_boundary_conditions = [False, False, False]
+        particles.set_pbc(pbc=self.periodic_boundary_conditions)
+
+        # Lattice vectors
+        if self.lattice_vectors is not None:
+            particles.set_cell(cell=self.lattice_vectors.to('angstrom').magnitude)
+            print(particles.get_cell())
+        else:
+            logger.info('Could not find `ParticleCell.lattice_vectors`.')
+
+        # Positions
+        if self.positions is not None:
+            if len(self.positions) != len(self.particles_state):
+                logger.error(
+                    'Length of `ParticleCell.positions` does not coincide with the length of the `ParticleCell.particles_state`.'
+                )
+                return None
+            particles.set_positions(
+                newpositions=self.positions.to('angstrom').magnitude
+            )
+        else:
+            logger.warning('Could not find `ParticleCell.positions`.')
+            return None
+
+        return particles
+
+    def from_particles(self, particles: Particles, logger: 'BoundLogger') -> None:
+        """
+        Parses the information from a Particles object to the `ParticlesCell` section.
+
+        Args:
+            particles (Particles): The Particles object to parse.
+            logger (BoundLogger): The logger to log messages.
+        """
+        # `ParticlesState[*].particles_type`
+        for label in particles.get_particle_types():
+            particle_state = ParticlesState(particle_type=label)
+            self.particles_state.append(particle_state)
+
+        # `periodic_boundary_conditions`
+        self.periodic_boundary_conditions = particles.get_pbc()
+
+        # `lattice_vectors`
+        cell = particles.get_cell()
+        self.lattice_vectors = ase.geometry.complete_cell(cell) * ureg('angstrom')
+
+        # `positions`
+        positions = particles.get_positions()
+        if (
+            not positions.tolist()
+        ):  # ASE assigns a shape=(0, 3) array if no positions are found
+            return None
+        self.positions = positions * ureg('angstrom')
+
+    def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
+        super().normalize(archive, logger)
+
+        # Set the name of the section
+        self.name = self.m_def.name if self.name is None else self.name
+
+
 class Symmetry(ArchiveSection):
     """
     A base section used to specify the symmetry of the `AtomicCell`.
@@ -983,6 +1146,7 @@ class ModelSystem(System):
     )
 
     cell = SubSection(sub_section=Cell.m_def, repeats=True)
+    print(f'cell: {Cell.m_def}')
 
     symmetry = SubSection(sub_section=Symmetry.m_def, repeats=True)
 
@@ -1017,6 +1181,7 @@ class ModelSystem(System):
     # TODO improve description and add an example using the case in atom_indices
     bond_list = Quantity(
         type=np.int32,
+        shape=['*', 2],
         description="""
         List of pairs of atom indices corresponding to bonds (e.g., as defined by a force field)
         within this atoms_group.
@@ -1115,6 +1280,8 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
                 'Could not find the originally parsed atomic system. `Symmetry` and `ChemicalFormula` extraction is thus not run.'
             )
             return
+
+        #! self.cell[0].name is ParticleCell
         if self.cell[0].name == 'AtomicCell':
             self.cell[0].type = 'original'
             ase_atoms = self.cell[0].to_ase_atoms(logger=logger)