From 1643440bafc4ea6a18d93018fb11e37bf3114546 Mon Sep 17 00:00:00 2001
From: JosePizarro3 <jose.pizarro@physik.hu-berlin.de>
Date: Mon, 15 Apr 2024 10:59:50 +0200
Subject: [PATCH] Added extract_spin_polarized_property to Outputs and testing

---
 src/nomad_simulations/outputs.py | 23 ++++++++++++++++++++---
 tests/test_outputs.py            | 25 +++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/nomad_simulations/outputs.py b/src/nomad_simulations/outputs.py
index edc21331..d0ea0868 100644
--- a/src/nomad_simulations/outputs.py
+++ b/src/nomad_simulations/outputs.py
@@ -64,12 +64,29 @@ class Outputs(ArchiveSection):
     # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
     electronic_band_gap = SubSection(sub_section=ElectronicBandGap.m_def, repeats=True)
 
+    def extract_spin_polarized_property(self, property_name: str) -> list:
+        """
+        Extracts the spin-polarized properties if present from the property name and returns them as a list of two elements in
+        which each element refers to each `spin_channel`. If the return list is empty, it means that the simulation is not
+        spin-polarized (i.e., `spin_channel` is not defined).
+
+        Args:
+            property_name (str): The name of the property to be extracted.
+
+        Returns:
+            (list): The list of spin-polarized properties.
+        """
+        spin_polarized_properties = []
+        properties = getattr(self, property_name)
+        for prop in properties:
+            if prop.spin_channel is None:
+                continue
+            spin_polarized_properties.append(prop)
+        return spin_polarized_properties
+
     def normalize(self, archive, logger) -> None:
         super().normalize(archive, logger)
 
-        # Resolve if the output property `is_derived` or not.
-        # self.is_derived = self.resolve_is_derived(self.outputs_ref)
-
 
 class SCFOutputs(Outputs):
     """
diff --git a/tests/test_outputs.py b/tests/test_outputs.py
index fac4489c..081718e1 100644
--- a/tests/test_outputs.py
+++ b/tests/test_outputs.py
@@ -67,6 +67,31 @@ def test_is_scf_converged(self, threshold_change: float, result: bool):
         )
         assert is_scf_converged == result
 
+    def test_extract_spin_polarized_properties(self):
+        """
+        Test the `extract_spin_polarized_property` method.
+        """
+        outputs = Outputs()
+
+        # No spin-polarized band gap
+        band_gap_non_spin_polarized = ElectronicBandGap(variables=[])
+        band_gap_non_spin_polarized.value = 2.0 * ureg.joule
+        outputs.electronic_band_gap.append(band_gap_non_spin_polarized)
+        band_gaps = outputs.extract_spin_polarized_property('electronic_band_gap')
+        assert band_gaps == []
+
+        # Spin-polarized band gaps
+        band_gap_spin_1 = ElectronicBandGap(variables=[], spin_channel=0)
+        band_gap_spin_1.value = 1.0 * ureg.joule
+        outputs.electronic_band_gap.append(band_gap_spin_1)
+        band_gap_spin_2 = ElectronicBandGap(variables=[], spin_channel=1)
+        band_gap_spin_2.value = 1.5 * ureg.joule
+        outputs.electronic_band_gap.append(band_gap_spin_2)
+        band_gaps = outputs.extract_spin_polarized_property('electronic_band_gap')
+        assert len(band_gaps) == 2
+        assert band_gaps[0].value.magnitude == 1.0
+        assert band_gaps[1].value.magnitude == 1.5
+
     @pytest.mark.parametrize(
         'threshold_change, result',
         [(1e-3, True), (1e-5, False)],