Add atom centered basis set

src/nomad_simulations/schema_packages/basis_set.py

@@ -13,7 +13,7 @@
 from nomad import utils
 from nomad.datamodel.data import ArchiveSection
 from nomad.datamodel.metainfo.annotations import ELNAnnotation
-from nomad.metainfo import MEnum, Quantity, SubSection
+from nomad.metainfo import JSON, MEnum, Quantity, SubSection
 from nomad.units import ureg
 
 from nomad_simulations.schema_packages.atoms_state import AtomsState
@@ -190,25 +190,170 @@ class AtomCenteredFunction(ArchiveSection):
     Specifies a single function (term) in an atom-centered basis set.
     """
 
-    pass
+    harmonic_type = Quantity(
+        type=MEnum(
+            'spherical',
+            'cartesian',
+        ),
+        default='spherical',
+        description="""
+        Specifies whether the basis functions are spherical-harmonic or cartesian functions.
+        """,
+    )
+
+    function_type = Quantity(
+        type=MEnum(
+            's',
+            'p',
+            'd',
+            'f',
+            'g',
+            'h',
+            'i',
+            'j',
+            'k',
+            'l',
+            'sp',
+            'spd',
+            'spdf',
+        ),
+        description="""
+        L=a+b+c
+        The angular momentum of GTO to be added.
+        """,
+    )
+
+    n_primitive = Quantity(
+        type=np.int32,
+        description="""
+        Number of primitives.
+        Linear combinations of the primitive Gaussians are formed to approximate the radial extent of an STO.
+        """,
+    )
+
+    exponents = Quantity(
+        type=np.float32,
+        shape=['n_primitive'],
+        description="""
+        List of exponents for the basis function.
+        """,
+    )
+
+    contraction_coefficients = Quantity(
+        type=np.float32,
+        shape=['*'],  # Flexible shape to handle combined types (e.g. SP, SPD..)
+        description="""
+        List of contraction coefficients corresponding to the exponents.
+        """,
+    )
+
+    point_charge = Quantity(
+        type=np.float32,
+        description="""
+        the value of the point charge.
+        """,
+    )
 
-    # TODO: design system for writing basis functions like gaussian or slater orbitals
+    def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
+        """
+        Validates the input data
+        and resolves combined types like SP, SPD, SPDF, etc.
+
+        Raises ValueError: If the data is inconsistent (e.g., mismatch in exponents and coefficients).
+        """
+        super().normalize(archive, logger)
+
+        # Validate number of primitives
+        if self.n_primitive is not None:
+            if self.exponents is not None and len(self.exponents) != self.n_primitive:
+                raise ValueError(
+                    f'Mismatch in number of exponents: expected {self.n_primitive}, '
+                    f'found {len(self.exponents)}.'
+                )
+
+        # Resolve combined types
+        if self.function_type and len(self.function_type) > 1:
+            num_types = len(self.function_type)  # For SP: 2, SPD: 3, etc.
+            if self.contraction_coefficients is not None:
+                expected_coeffs = num_types * self.n_primitive
+                if len(self.contraction_coefficients) != expected_coeffs:
+                    raise ValueError(
+                        f'Mismatch in contraction coefficients for {self.function_type} type: '
+                        f'expected {expected_coeffs}, found {len(self.contraction_coefficients)}.'
+                    )
+
+                # Split coefficients into separate lists for each type
+                self.coefficient_sets = {
+                    t: self.contraction_coefficients[i::num_types]
+                    for i, t in enumerate(self.function_type)
+                }
+
+                # Debug: Log split coefficients
+                for t, coeffs in self.coefficient_sets.items():
+                    logger.info(f'{t}-type coefficients: {coeffs}')
+            else:
+                logger.warning(
+                    f'No contraction coefficients provided for {self.function_type} type.'
+                )
+
+        # For single types, ensure coefficients match primitives
+        elif self.contraction_coefficients is not None:
+            if len(self.contraction_coefficients) != self.n_primitive:
+                raise ValueError(
+                    f'Mismatch in contraction coefficients: expected {self.n_primitive}, '
+                    f'found {len(self.contraction_coefficients)}.'
+                )
 
 
 class AtomCenteredBasisSet(BasisSetComponent):
     """
     Defines an atom-centered basis set.
     """
 
+    basis_set = Quantity(
+        type=str,
+        description="""
+        name of the basis set.
+        """,
+    )
+
+    type = Quantity(
+        type=MEnum(
+            'STO',  # Slater-type orbitals
+            'GTO',  # Gaussian-type orbitals
+            'NAO',  # Numerical atomic orbitals
+            'cECP',  # Capped effective core potentials
+            'PC',  # Point charges
+        ),
+        description="""
+        Type of the basis set, e.g. STO or GTO.
+        """,
+    )
+
+    role = Quantity(
+        type=MEnum(
+            'orbital',
+            'auxiliary_scf',
+            'auxiliary_post_hf',
+            'cabs',  # complementary auxiliary basis set
+        ),
+        description="""
+        The role of the basis set.
+        """,
+    )
+
+    total_number_of_basis_functions = Quantity(
+        type=np.int32,
+        description='',
+    )
+
     functional_composition = SubSection(
         sub_section=AtomCenteredFunction.m_def, repeats=True
-    )  # TODO change name
+    )
 
     def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
         super().normalize(archive, logger)
         # self.name = self.m_def.name
-        # TODO: set name based on basis functions
-        # ? use basis set names from Basis Set Exchange
 
 
 class APWBaseOrbital(ArchiveSection):