Refactoring model system (#170)

* work with basesections.v2

* remove the redundant child model_system subsection

* move positions Quantity to ModelSystem

* AtomsState(Element), remove helper functions, modify normalize() to check consistency between chemical_symbol and atomic_number

* remove positions and velocities from Cell

* move ASE Atoms object creation to ModelSystem

* add sub_systems to ModelSystem

* add total_charge and total_spin to ModelSystem

* add ModelSystem.sub_systems instead of ModelSystem.model_system

* move comparison functions from AtomicCell to ModelSystem

* add ParticleState and CoarseGrainedState

* Refactor symmetry resolution: use dict mapping for cell types

* remove atom_types and related tests

* remove too many warnings in get_chemical_symbols

* ModelSystem.particle_states instead of atom_states

* remove redundant positions length checks

* n_particles and particle_indices (instead of n_atoms and atom_indices)

* add a note in GeometricSpace.get_geometric_space_for_atomic_cell

* modify to_ase_atoms, test_numerical_settings.py 

* modify conftest so that model_system.positions are populated

* import the v2 Entity in physical_property.py

* Fix DOSProfile normalization factor resolution to use model_system.particle_states

* Remove unused variable from SlaterKosterBond normalization test

* fix APW-refs fixture in conftest.py

* fix spectral profile tests

* Remove redundant try/except around get_geometric_space_for_atomic_cell

* add a test for composition formula

Author: EBB2675

src/nomad_simulations/schema_packages/atoms_state.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pint
 from nomad.datamodel.data import ArchiveSection
-from nomad.datamodel.metainfo.basesections import Entity
+from nomad.datamodel.metainfo.basesections.v2 import Entity
 from nomad.metainfo import MEnum, Quantity, SubSection
 from nomad.units import ureg
 
@@ -545,12 +545,20 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
                 )
 
 
-class AtomsState(Entity):
+class ParticleState(Entity):
+    """
+    Generic base section representing the state of a particle in a simulation.
+    This can be extended to include any common quantities in the future.
+    """
+
+    pass
+
+
+class AtomsState(ParticleState):
     """
     A base section to define each atom state information.
     """
 
-    # TODO check what happens with ghost atoms that can have `chemical_symbol='X'`
     chemical_symbol = Quantity(
         type=MEnum(ase.data.chemical_symbols[1:]),
         description="""
@@ -565,8 +573,6 @@ class AtomsState(Entity):
         """,
     )
 
-    orbitals_state = SubSection(sub_section=OrbitalsState.m_def, repeats=True)
-
     charge = Quantity(
         type=np.int32,
         default=0,
@@ -581,6 +587,16 @@ class AtomsState(Entity):
         """,
     )
 
+    spin = Quantity(
+        type=np.int32,
+        default=0,
+        description="""
+        Total spin quantum number, S.
+        """,
+    )
+
+    orbitals_state = SubSection(sub_section=OrbitalsState.m_def, repeats=True)
+
     core_hole = SubSection(sub_section=CoreHole.m_def, repeats=False)
 
     hubbard_interactions = SubSection(
@@ -633,3 +649,7 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
             self.chemical_symbol = self.resolve_chemical_symbol(logger=logger)
         if self.atomic_number is None:
             self.atomic_number = self.resolve_atomic_number(logger=logger)
+
+
+class CoarseGrainedState(ParticleState):
+    pass

src/nomad_simulations/schema_packages/general.py

@@ -12,6 +12,7 @@
 from nomad.datamodel.metainfo.basesections import Activity, Entity
 from nomad.metainfo import Datetime, Quantity, SchemaPackage, Section, SubSection
 
+from nomad_simulations.schema_packages.atoms_state import AtomsState
 from nomad_simulations.schema_packages.model_method import ModelMethod
 from nomad_simulations.schema_packages.model_system import ModelSystem
 from nomad_simulations.schema_packages.outputs import Outputs
@@ -195,7 +196,7 @@ class Simulation(BaseSimulation, Schema):
     def _set_system_branch_depth(
         self, system_parent: ModelSystem, branch_depth: int = 0
     ):
-        for system_child in system_parent.model_system:
+        for system_child in system_parent.sub_systems:
             system_child.branch_depth = branch_depth + 1
             self._set_system_branch_depth(
                 system_parent=system_child, branch_depth=branch_depth + 1
@@ -222,10 +223,10 @@ def set_composition_formula(
                 to the atom indices stored in system.
             """
             if not subsystems:
-                if system.atom_indices is not None and atom_labels:
-                    subsystem_labels = [atom_labels[i] for i in system.atom_indices]
-                elif system.atom_indices is not None:
-                    subsystem_labels = ['Unknown'] * len(system.atom_indices)
+                if system.particle_indices is not None and atom_labels:
+                    subsystem_labels = [atom_labels[i] for i in system.particle_indices]
+                elif system.particle_indices is not None:
+                    subsystem_labels = ['Unknown'] * len(system.particle_indices)
                 else:
                     subsystem_labels = []
             else:
@@ -249,22 +250,25 @@ def get_composition_recurs(system: ModelSystem, atom_labels: list[str]) -> None:
                 atom_labels (list[str]): The global list of atom labels corresponding
                 to the atom indices stored in system.
             """
-            subsystems = system.model_system
+            subsystems = system.sub_systems
             set_composition_formula(
                 system=system, subsystems=subsystems, atom_labels=atom_labels
             )
             if subsystems:
                 for subsystem in subsystems:
                     get_composition_recurs(system=subsystem, atom_labels=atom_labels)
 
-        atoms_state = (
-            system_parent.cell[0].atoms_state if system_parent.cell is not None else []
-        )
+        # Pull chemical symbols straight from AtomsState.chemical_symbol
         atom_labels = (
-            [atom.chemical_symbol for atom in atoms_state]
-            if atoms_state is not None
+            [
+                atom.chemical_symbol
+                for atom in system_parent.particle_states
+                if isinstance(atom, AtomsState) and atom.chemical_symbol is not None
+            ]
+            if system_parent.particle_states
             else []
         )
+
         get_composition_recurs(system=system_parent, atom_labels=atom_labels)
 
     def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
@@ -284,7 +288,7 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
         # Setting up the `branch_depth` in the parent-child tree
         for system_parent in self.model_system:
             system_parent.branch_depth = 0
-            if len(system_parent.model_system) == 0:
+            if len(system_parent.sub_systems) == 0:
                 continue
             self._set_system_branch_depth(system_parent=system_parent)
 

src/nomad_simulations/schema_packages/model_method.py

@@ -483,59 +483,74 @@ def resolve_orbital_references(
         model_index: int = -1,
     ) -> Optional[list[OrbitalsState]]:
         """
-        Resolves the references to the `OrbitalsState` sections from the child `ModelSystem` section.
+        Resolves references to the `OrbitalsState` sections from the top-level `ModelSystem`
+        that has child system(s) typed 'active_atom'. This uses the new design:
+
+        - The parent ModelSystem stores per-atom data in `particle_states`.
+        - The child system(s) typed 'active_atom' list indices in `particle_indices`.
+        - We gather OrbitalsState from each relevant particle_states entry.
 
         Args:
             model_systems (list[ModelSystem]): The list of `ModelSystem` sections.
             logger (BoundLogger): The logger to log messages.
-            model_index (int, optional): The `ModelSystem` section index from which resolve the references. Defaults to -1.
+            model_index (int, optional): The ModelSystem index to use. Defaults to -1 (the last).
 
         Returns:
-            Optional[list[OrbitalsState]]: The resolved references to the `OrbitalsState` sections.
+            Optional[list[OrbitalsState]]: The resolved references to the OrbitalsState sections.
         """
+        # Check that the requested ModelSystem exists
         try:
             model_system = model_systems[model_index]
         except IndexError:
-            logger.warning(
-                f'The `ModelSystem` section with index {model_index} was not found.'
-            )
+            logger.warning(f'No ModelSystem at index {model_index}.')
             return None
 
-        # If `ModelSystem` is not representative, the normalization will not run
+        # If the system is not representative, bail out of normalization
         if is_not_representative(model_system=model_system, logger=logger):
             return None
 
-        # If `AtomicCell` is not found, the normalization will not run
-        if not model_system.cell:
-            logger.warning('`AtomicCell` section was not found.')
+        # If no child ModelSystem sections exist, bail out of normalization
+        if not model_system.sub_systems:
+            logger.warning(
+                'No child ModelSystem found; cannot find active_atom references.'
+            )
             return None
-        atomic_cell = model_system.cell[0]
 
-        # If there is no child `ModelSystem`, the normalization will not run
-        atoms_state = atomic_cell.atoms_state
-        model_system_child = model_system.model_system
-        if not atoms_state or not model_system_child:
-            logger.warning('No `AtomsState` or child `ModelSystem` section were found.')
+        #  If no particle_states are present at the top level, we have no orbitals
+        if not model_system.particle_states:
+            logger.warning('No particle_states in the parent ModelSystem.')
             return None
 
-        # We flatten the `OrbitalsState` sections from the `ModelSystem` section
-        orbitals_ref = []
-        for active_atom in model_system_child:
-            # If the child is not an "active_atom", the normalization will not run
-            if active_atom.type != 'active_atom':
+        orbitals_ref: list[OrbitalsState] = []
+
+        # For each child in sub_systems, if type='active_atom', gather orbitals
+        for child_sys in model_system.sub_systems:
+            if child_sys.type != 'active_atom':
+                continue
+            # if no particle_indices => skip
+            if not child_sys.particle_indices:
+                logger.warning('Child system is active_atom but no particle_indices.')
                 continue
-            indices = active_atom.atom_indices
-            for index in indices:
-                try:
-                    active_atoms_state = atoms_state[index]
-                except IndexError:
+
+            # For each index in child_sys.particle_indices => fetch from parent’s particle_states
+            for idx in child_sys.particle_indices:
+                if idx < 0 or idx >= len(model_system.particle_states):
                     logger.warning(
-                        f'The `AtomsState` section with index {index} was not found.'
+                        f'Particle index {idx} out of range for particle_states.'
                     )
                     continue
-                orbitals_state = active_atoms_state.orbitals_state
-                for orbital in orbitals_state:
-                    orbitals_ref.append(orbital)
+                active_atom_state = model_system.particle_states[idx]
+
+                # If no orbitals_state => skip
+                if not active_atom_state.orbitals_state:
+                    logger.warning(
+                        f'No orbitals_state found in particle_states[{idx}].'
+                    )
+                    continue
+
+                orbitals_ref.extend(active_atom_state.orbitals_state)
+
+        # Return the collected orbitals
         return orbitals_ref
 
     def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None: