initial investigation into harmonizing particles and atoms

src/nomad_simulations/schema_packages/atoms_state.py

@@ -552,7 +552,16 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
                 )
 
 
-class AtomsState(Entity):
+class State(Entity):
+    """
+    A base section to define the state information of the system.
+    """
+
+    def __init__(self, m_def: 'Section' = None, m_context: 'Context' = None, **kwargs):
+        super().__init__(m_def, m_context, **kwargs)
+
+
+class AtomsState(State):
     """
     A base section to define each atom state information.
     """

src/nomad_simulations/schema_packages/general.py

@@ -1,10 +1,17 @@
-from typing import TYPE_CHECKING
+#! TODO: Why is TYPE_CHECKING False?
+from typing import TYPE_CHECKING, List, Iterable, Union
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
-
-    from nomad.datamodel.datamodel import EntryArchive
-    from structlog.stdlib import BoundLogger
+if not TYPE_CHECKING:
+    from nomad.datamodel.datamodel import (
+        EntryArchive,
+    )
+    from nomad.metainfo import (
+        Context,
+        Section,
+    )
+    from structlog.stdlib import (
+        BoundLogger,
+    )
 
 import numpy as np
 from nomad.config import config
@@ -218,7 +225,9 @@ def _set_system_branch_depth(
             )
 
     #! Generalize from checks for atomic systems, error with CG input
-    def resolve_composition_formula(self, system_parent: ModelSystem) -> None:
+    def resolve_composition_formula(
+        self, system_parent: ModelSystem, logger: 'BoundLogger'
+    ) -> None:
         """Determine and set the composition formula for `system_parent` and all of its
         descendants.
 
@@ -227,7 +236,7 @@ def resolve_composition_formula(self, system_parent: ModelSystem) -> None:
         """
 
         def set_composition_formula(
-            system: ModelSystem, subsystems: list[ModelSystem], atom_labels: list[str]
+            system: ModelSystem, subsystems: list[ModelSystem], labels: list[str]
         ) -> None:
             """Determine the composition formula for `system` based on its `subsystems`.
             If `system` has no children, the atom_labels are used to determine the formula.
@@ -239,13 +248,15 @@ def set_composition_formula(
                 to the atom indices stored in system.
             """
             if not subsystems:
-                atom_indices = (
-                    system.atom_indices if system.atom_indices is not None else []
+                particle_indices = (
+                    system.particle_indices
+                    if system.particle_indices is not None
+                    else []
                 )
                 subsystem_labels = (
-                    [np.array(atom_labels)[atom_indices]]
-                    if atom_labels
-                    else ['Unknown' for atom in range(len(atom_indices))]
+                    [np.array(labels)[particle_indices]]
+                    if labels
+                    else ['Unknown' for atom in range(len(particle_indices))]
                 )
             else:
                 subsystem_labels = [
@@ -259,7 +270,7 @@ def set_composition_formula(
                     children_names=subsystem_labels
                 )
 
-        def get_composition_recurs(system: ModelSystem, atom_labels: list[str]) -> None:
+        def get_composition_recurs(system: ModelSystem, labels: list[str]) -> None:
             """Traverse the system hierarchy downward and set the branch composition for
             all (sub)systems at each level.
 
@@ -269,23 +280,17 @@ def get_composition_recurs(system: ModelSystem, atom_labels: list[str]) -> None:
                 to the atom indices stored in system.
             """
             subsystems = system.model_system
-            set_composition_formula(
-                system=system, subsystems=subsystems, atom_labels=atom_labels
-            )
+            set_composition_formula(system=system, subsystems=subsystems, labels=labels)
             if subsystems:
                 for subsystem in subsystems:
-                    get_composition_recurs(system=subsystem, atom_labels=atom_labels)
+                    get_composition_recurs(system=subsystem, labels=labels)
 
         # ! CG: system_parent.cell[0].particles_state instead of atoms_state!
-        atoms_state = (
-            system_parent.cell[0].atoms_state if system_parent.cell is not None else []
-        )
-        atom_labels = (
-            [atom.chemical_symbol for atom in atoms_state]
-            if atoms_state is not None
-            else []
-        )
-        get_composition_recurs(system=system_parent, atom_labels=atom_labels)
+        labels = []
+        if system_parent.cell is not None:
+            labels = system_parent.cell[0].get('labels', logger=logger)
+
+        get_composition_recurs(system=system_parent, labels=labels)
 
     def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
         super(Schema, self).normalize(archive, logger)
@@ -310,7 +315,7 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
 
             if is_not_representative(model_system=system_parent, logger=logger):
                 continue
-            self.resolve_composition_formula(system_parent=system_parent)
+            self.resolve_composition_formula(system_parent=system_parent, logger=logger)
 
 
 m_package.__init_metainfo__()

src/nomad_simulations/schema_packages/model_method.py

@@ -544,7 +544,7 @@ def resolve_orbital_references(
             # If the child is not an "active_atom", the normalization will not run
             if active_atom.type != 'active_atom':
                 continue
-            indices = active_atom.atom_indices
+            indices = active_atom.particle_indices
             for index in indices:
                 try:
                     active_atoms_state = atoms_state[index]