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 MEnum, Quantity, SubSection, JSON
 from nomad.units import ureg
 
 from nomad_simulations.schema_packages.atoms_state import AtomsState
@@ -190,16 +190,122 @@ class AtomCenteredFunction(ArchiveSection):
     Specifies a single function (term) in an atom-centered basis set.
     """
 
-    pass
+    function_type = Quantity(
+        type=MEnum('s', 'p', 'd', 'f', 'g', 'h', 'i', 'j'),
+        description="""
+        the angular momentum of the shell to be added.
+        """,
+    )
+
+    n_primitive = Quantity(
+        type=int,
+        description="""
+        Number of primitives.
+        """,
+    )
+
+    exponents = Quantity(
+        type=np.float32,
+        shape=['n_primitive'],
+        description="""
+        List of exponents for the basis function.
+        """,
+    )
+
+    contraction_coefficients = Quantity(
+        type=np.float32,
+        shape=['n_primitive'],
+        description="""
+        List of contraction coefficients corresponding to the exponents.
+        """,
+    )
+
+    def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
+        super().normalize(archive, logger)
+        # self.name = self.m_def.name
+
+class AtomCenteredBasisSet(BasisSetComponent):
+    """
+    Defines an atom-centered basis set, using a single JSON quantity for all basis set types,
+    while allowing different AtomState references for each basis set type.
+    """
+
+    basis_set_data = Quantity(
+        type=JSON,  # Use JSON to store basis set information, including atom references
+        description="""
+        JSON object containing all the basis set information along with atom references. Example:
+        {
+            "main_basis_set": {
+                "name": "cc-pVTZ",
+                "atoms_ref": [ref_to_atoms_1, ref_to_atoms_2]
+            },
+            "aux_c_basis_set": {
+                "name": "cc-pVTZ/C",
+                "atoms_ref": [ref_to_atoms_3]
+            },
+            "aux_j_basis_set": {
+                "name": "RIJ",
+                "atoms_ref": [ref_to_atoms_4]
+            },
+            "aux_jk_basis_set": {
+                "name": "aug-cc-pVTZ/JK",
+                "atoms_ref": [ref_to_atoms_1, ref_to_atoms_5]
+            }
+        }
+        """,
+    )
+
+    functional_composition = SubSection(
+        sub_section=AtomCenteredFunction.m_def, repeats=True
+    )
+
+    def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
+        super().normalize(archive, logger)
 
-    # TODO: design system for writing basis functions like gaussian or slater orbitals
 
 
+'''
 class AtomCenteredBasisSet(BasisSetComponent):
     """
     Defines an atom-centered basis set.
     """
 
+    main_basis_set = Quantity(
+        type=str,
+        description="""
+        Name of the main basis set.
+        """,
+    )
+
+    aux_c_basis_set = Quantity(
+        type=str,
+        description="""
+        AuxC type of basis set.
+        """,
+    )
+
+    aux_j_basis_set = Quantity(
+        type=str,
+        description="""
+        AuxJ type of basis set.
+        """,
+    )
+
+    aux_jk_basis_set = Quantity(
+        type=str,
+        description="""
+        AuxJK type of basis set.
+        """,
+    )
+
+    atoms_ref = Quantity(
+        type=AtomsState,
+        shape=['*'],
+        description="""
+        References to the `AtomsState` sections that define the atoms this basis set applies to.
+        """,
+    )
+
     functional_composition = SubSection(
         sub_section=AtomCenteredFunction.m_def, repeats=True
     )  # TODO change name
@@ -209,7 +315,7 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
         # self.name = self.m_def.name
         # TODO: set name based on basis functions
         # ? use basis set names from Basis Set Exchange
-
+'''
 
 class APWBaseOrbital(ArchiveSection):
     """