- Clean eigenvalues and band structure base sections: the latter is the same as the former but with a k-points level

ndaelman-hu · ndaelman-hu · commit 52182fd30b22 · 2025-06-19T16:07:37.000+02:00
- Add various generic utilities: `quicksearch_first_value`, `inner_copy`, `check_not_none`
diff --git a/src/nomad_simulations/schema_packages/properties/band_structure.py b/src/nomad_simulations/schema_packages/properties/band_structure.py
@@ -9,13 +9,14 @@
 
 if TYPE_CHECKING:
     from nomad.datamodel.datamodel import EntryArchive
-    from nomad.metainfo import Context, Section
     from structlog.stdlib import BoundLogger
 
 from nomad_simulations.schema_packages.atoms_state import AtomsState, OrbitalsState
 from nomad_simulations.schema_packages.data_types import unit_float
 from nomad_simulations.schema_packages.numerical_settings import KSpace
 from nomad_simulations.schema_packages.physical_property import PhysicalProperty
+from nomad_simulations.schema_packages.utils.utils import check_not_none
+from nomad_simulations.schema_packages.utils.electronic import quicksearch_first_value, inner_copy
 
 configuration = config.get_plugin_entry_point(
     'nomad_simulations.schema_packages:nomad_simulations_plugin'
@@ -27,60 +28,56 @@ class BaseElectronicEigenvalues(PhysicalProperty):
     A base section used to define basic quantities for the `ElectronicEigenvalues`  and `ElectronicBandStructure` properties.
     """
 
-    iri = ''
-
     n_bands = Quantity(
         type=np.int32,
         description="""
         Number of bands / eigenvalues.
         """,
     )  # TODO: remove
 
-    value = Quantity(
-        type=np.float64,
-        unit='joule',
-        shape=['*', '*'],
-        description="""
-        Value of the electronic eigenvalues.
-        """,
-    )
-
-
 class ElectronicEigenvalues(BaseElectronicEigenvalues):
     """ """
 
     iri = 'http://fairmat-nfdi.eu/taxonomy/ElectronicEigenvalues'
 
-    spin_channel = Quantity(
-        type=np.int32,
+    # TODO: add spin annotation from @EBB2675
+
+    value = Quantity(
+        type=np.float64,
+        unit='joule',
+        shape=['level', 'spin'],
         description="""
-        Spin channel of the corresponding electronic eigenvalues. It can take values of 0 or 1.
+        Value of the electronic eigenvalues.
+        Rows correspond to the energy levels, and columns correspond to the spin channels.
         """,
     )
 
     occupation = Quantity(
         type=unit_float(),
-        shape=['*'],
+        shape=['level', 'spin'],
         description="""
-        Occupation of the electronic eigenvalues. 
+        Occupation of the electronic eigenvalues, ranging from 0 to 1.
+        Rows correspond to the energy levels, and columns correspond to the spin channels.
         """,
     )  # restructure spin for plotting?
 
     highest_occupied = Quantity(
         type=np.float64,
+        shape=['spin'],
         unit='joule',
         description="""
-        Highest occupied electronic eigenvalue. Together with `lowest_unoccupied`, it defines the
-        electronic band gap.
+        Highest occupied electronic eigenvalue for each spin channel. Together with `lowest_unoccupied`, it defines the
+        electronic band gap. Automatically resolved using binary search on sorted eigenvalues.
         """,
     )
 
     lowest_unoccupied = Quantity(
         type=np.float64,
+        shape=['spin'],
         unit='joule',
         description="""
-        Lowest unoccupied electronic eigenvalue. Together with `highest_occupied`, it defines the
-        electronic band gap.
+        Lowest unoccupied electronic eigenvalue for each spin channel. Together with `highest_occupied`, it defines the
+        electronic band gap. Automatically resolved using binary search on sorted eigenvalues.
         """,
     )
 
@@ -98,17 +95,135 @@ class ElectronicEigenvalues(BaseElectronicEigenvalues):
         """,
     )
 
+    @check_not_none('self.value', 'self.occupation')
+    def resolve_homo_lumo(self) -> None:
+        """
+        Resolve HOMO and LUMO eigenvalues using binary search on sorted eigenvalues.
+        """        
+        def process_spin_channel(spin_data):
+            """Process a single spin channel to find HOMO/LUMO."""
+            spin_values, spin_occupations = spin_data
+            lumo_idx = quicksearch_first_value(
+                spin_occupations, 0.0, tolerance=1e-6
+            )
+
+            return [
+                spin_values[lumo_idx] if lumo_idx is not None and lumo_idx >= 0 else None,
+                spin_values[lumo_idx + 1] if lumo_idx is not None and lumo_idx > 0 else None
+            ]
+
+        # Stack value and occupation arrays along last axis for apply_along_axis
+        combined_data = np.stack([self.value, self.occupation.magnitude], axis=-1)
+        results = np.apply_along_axis(process_spin_channel, axis=0, arr=combined_data.T)
+        
+        self.highest_occupied = results[:, 0] * self.value.u
+        self.lowest_unoccupied = results[:, 1] * self.value.u
+
+    def pad_out(self) -> None:
+        """
+        Pad out the value and occupation arrays along the spin channel dimension.
+        """
+        spin_index = 2
+        if np.array(self.value).shape[spin_index] == 1:  # TODO: add model_method spin_polarized
+            self.value = inner_copy(self.value, 0)  # TODO: dynamically set repetition
+        if np.array(self.occupation).shape[spin_index] == 1:  # TODO: add model_method spin_polarized
+            self.occupation = inner_copy(self.occupation, 0)  # TODO: dynamically set repetition
+
+    def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
+        super().normalize(archive, logger)
+        self.pad_out()
+        self.resolve_homo_lumo()
+
 
-class ElectronicBandStructure(ElectronicEigenvalues):
+class ElectronicBandStructure(BaseElectronicEigenvalues):
     """
     Accessible energies by the charges (electrons and holes) in the reciprocal space.
     """
 
     iri = 'http://fairmat-nfdi.eu/taxonomy/ElectronicBandStructure'
 
-    k_path = SubSection(sub_section=KMesh.m_def)
+    kpoint = SubSection(sub_section=KMesh.m_def)
 
-    def resolve_reciprocal_cell(self) -> Optional[pint.Quantity]:
+    value = Quantity(
+        type=np.float64,
+        unit='joule',
+        shape=['level', 'kpoint', 'spin'],
+        description="""
+        Value of the electronic eigenvalues in the reciprocal space.
+        Dimensions: [energy level, k-point, spin channel].
+        """,
+    )
+
+    occupation = Quantity(
+        type=unit_float(),
+        shape=['level', 'kpoint', 'spin'],
+        description="""
+        Occupation of the electronic eigenvalues, ranging from 0 to 1.
+        Dimensions: [energy level, k-point, spin channel].
+        """,
+    )
+
+    highest_occupied = Quantity(
+        type=np.float64,
+        shape=['kpoint', 'spin'],
+        unit='joule',
+        description="""
+        Highest occupied electronic eigenvalue for each k-point and spin channel. Together with `lowest_unoccupied`, it defines the
+        electronic band gap. Automatically resolved using binary search on sorted eigenvalues.
+        """,
+    )
+
+    lowest_unoccupied = Quantity(
+        type=np.float64,
+        shape=['kpoint', 'spin'],
+        unit='joule',
+        description="""
+        Lowest unoccupied electronic eigenvalue for each k-point and spin channel. Together with `highest_occupied`, it defines the
+        electronic band gap. Automatically resolved using binary search on sorted eigenvalues.
+        """,
+    )
+
+    @check_not_none('self.value', 'self.occupation')
+    def resolve_homo_lumo(self) -> None:
+        """
+        Resolve HOMO and LUMO eigenvalues using binary search on sorted eigenvalues for band structure.
+        """        
+        def process_kpoint_spin(kpoint_spin_data):
+            """Process a single k-point and spin channel to find HOMO/LUMO."""
+            spin_values, spin_occupations = kpoint_spin_data
+            lumo_idx = quicksearch_first_value(
+                spin_occupations, 0.0, tolerance=1e-6
+            )
+
+            return [
+                spin_values[lumo_idx] if lumo_idx is not None and lumo_idx >= 0 else None,
+                spin_values[lumo_idx + 1] if lumo_idx is not None and lumo_idx > 0 else None
+            ]
+
+        # Stack value and occupation arrays - shape: [level, kpoint, spin, 2]
+        # Apply along level axis (axis=0) for each k-point and spin combination
+        # Reshape to combine kpoint and spin dimensions for processing
+        combined_data = np.stack([self.value, self.occupation.magnitude], axis=-1)
+        n_levels, n_kpoints, n_spins, _ = combined_data.shape
+        reshaped_data = combined_data.transpose(1, 2, 0, 3).reshape(n_kpoints * n_spins, n_levels, 2)
+        results = np.apply_along_axis(process_kpoint_spin, axis=1, arr=reshaped_data)
+        
+        # Reshape back to [kpoint, spin, 2] then extract homo/lumo
+        results = results.reshape(n_kpoints, n_spins, 2)
+        self.highest_occupied = results[:, :, 0] * self.value.u
+        self.lowest_unoccupied = results[:, :, 1] * self.value.u
+
+    def pad_out(self) -> None:
+        """
+        Pad out the value and occupation arrays along the spin channel dimension.
+        """
+        spin_index = 2
+        if np.array(self.value).shape[spin_index] == 1:  # TODO: add model_method spin_polarized
+            self.value = inner_copy(self.value, 0)  # TODO: dynamically set repetition
+        if np.array(self.occupation).shape[spin_index] == 1:  # TODO: add model_method spin_polarized
+            self.occupation = inner_copy(self.occupation, 0)  # TODO: dynamically set repetition
+
+    def resolve_reciprocal_cell(self) -> pint.Quantity | None:  # ? remove
         """
         Resolve the reciprocal cell from the `KSpace` numerical settings section.
 
@@ -120,18 +235,15 @@ def resolve_reciprocal_cell(self) -> Optional[pint.Quantity]:
         )
         if numerical_settings is None:
             return None
-        k_space = None
+
         for setting in numerical_settings:
             if isinstance(setting, KSpace):
-                k_space = setting
-                break
-        if k_space is None:
-            return None
-        return k_space
+                return setting
 
     def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
         super().normalize(archive, logger)
-        self.reciprocal_cell = self.resolve_reciprocal_cell()
+        self.pad_out()
+        self.resolve_homo_lumo()
 
 
 # defunct
diff --git a/src/nomad_simulations/schema_packages/utils/electronic.py b/src/nomad_simulations/schema_packages/utils/electronic.py
@@ -0,0 +1,74 @@
+"""
+Electronic structure utility functions.
+"""
+
+import numpy as np
+
+
+def quicksearch_first_value(
+    arr: np.ndarray, target: float, tolerance: float = 1e-10
+) -> int | None:
+    """
+    Find the index of the first occurrence of a target value in a sorted array using binary search.
+
+    This function assumes the array is sorted and uses binary search to efficiently
+    locate the first element that equals the target (within tolerance).
+
+    Args:
+        arr: Sorted numpy array to search
+        target: Target value to search for
+        tolerance: Tolerance for considering values equal
+
+    Returns:
+        Index of first occurrence of target value, or None if not found
+    """
+    if len(arr) == 0:
+        return None
+
+    left, right = 0, len(arr) - 1
+    result = None
+
+    while left <= right:
+        mid = (left + right) // 2
+
+        if abs(arr[mid] - target) <= tolerance:
+            result = mid
+            right = mid - 1
+        elif arr[mid] < target - tolerance:
+            left = mid + 1
+        else:
+            right = mid - 1
+
+    return result
+
+
+def inner_copy(
+    tensor: np.ndarray, rank_selection: int | tuple[int] | slice, repeat: int = 1
+) -> np.ndarray:
+    """
+    Take a chunk of a high-ranked array and extend it with exact copies of the selection.
+
+    This function selects a portion of a tensor along its first axis and repeats it
+    the specified number of times, effectively extending the tensor.
+
+    Args:
+        tensor: Input `numpy` array to copy from
+        rank_selection: `int`, `tuple`, `slice` specifying which elements to select
+        repeat: Number of times to repeat the selection (default: 1)
+
+    Example:
+        >>> arr = np.array([[1, 2], [3, 4], [5, 6]])
+        >>> inner_copy(arr, slice(0, None), repeat=2)
+        array([[1, 2], [1, 2], [1, 2]])
+    """
+    if tensor.size == 0:
+        return tensor
+
+    selected_chunk = tensor[rank_selection]
+
+    # If selection results in 1D array, ensure it maintains proper shape
+    if selected_chunk.ndim == tensor.ndim - 1:
+        selected_chunk = np.expand_dims(selected_chunk, axis=0)
+
+    repeated_chunks = np.tile(selected_chunk, (repeat + 1, *([1] * (tensor.ndim - 1))))
+    return np.concatenate([tensor, repeated_chunks], axis=0)
diff --git a/src/nomad_simulations/schema_packages/utils/utils.py b/src/nomad_simulations/schema_packages/utils/utils.py
@@ -147,3 +147,40 @@ def wrapper(self, other) -> bool:
             return False
 
     return wrapper
+
+
+def check_not_none(*attributes: str) -> 'Callable':
+    """
+    Decorator that checks if specified object or class attributes are not `None`.
+    Returns `None` if any of the specified attributes are `None`, otherwise executes the function.
+    
+    Args:
+        *attributes: Names of attributes to check for None values
+            Use 'input.<attribute>' for input attributes
+            Use 'self.<attribute>' for object attributes
+            Use 'class.<attribute>' for class attributes
+            Use '<attribute>' for global attributes
+    """
+    def decorator(func: 'Callable') -> 'Callable':
+        def wrapper(self, *args, **kwargs):
+            for attr in attributes:
+                # Remove prefix and set source
+                if attr.startswith('input.'):
+                    attr_name = attr[6:]
+                    source = args[0]
+                elif attr.startswith('self.'):
+                    attr_name = attr[5:]
+                    source = self
+                elif attr.startswith('class.'):
+                    attr_name = attr[6:]
+                    source = self.__class__
+                else:
+                    attr_name = attr
+                    source = globals()
+                if not hasattr(source, attr_name) or getattr(source, attr_name) is None:
+                    return None
+
+            return func(self, *args, **kwargs)
+        
+        return wrapper
+    return decorator
