Refactoring

-  Use a helper function to extract hemoglobin fraction which is used for both sO2 and bvf
-  Remove unnecessary if conditions by using 0.0 as default
- improve type hinting

simpa/utils/calculate.py

@@ -3,72 +3,63 @@
 # SPDX-License-Identifier: MIT
 
 
-from typing import Union
+from typing import Union, List, Dict, Optional
 import numpy as np
 import torch
 from scipy.interpolate import interp1d
 
 
-def calculate_oxygenation(molecule_list: list) -> Union[float, int, torch.Tensor]:
+def extract_hemoglobin_fractions(molecule_list: List) -> Dict[str, float]:
     """
-    Calculate the oxygenation level based on the volume fractions of deoxyhaemoglobin and oxyhaemoglobin.
-
-    This function takes a list of molecules and returns an oxygenation value between 0 and 1 if computable,
-    otherwise returns None.
-
+    Extract hemoglobin volume fractions from a list of molecules.
+
     :param molecule_list: List of molecules with their spectrum information and volume fractions.
-    :return: An oxygenation value between 0 and 1 if possible, or None if not computable.
+    :return: A dictionary with hemoglobin types as keys and their volume fractions as values.
     """
-    hb = None    # Volume fraction of deoxyhaemoglobin
-    hbO2 = None  # Volume fraction of oxyhaemoglobin
+
+    # Put 0.0 as default value for both hemoglobin types in case they are not present in the molecule list.
+    hemoglobin = {
+        "Deoxyhemoglobin": 0.0,
+        "Oxyhemoglobin": 0.0
+    }
 
     for molecule in molecule_list:
-        if molecule.spectrum.spectrum_name == "Deoxyhemoglobin":
-            hb = molecule.volume_fraction
-        if molecule.spectrum.spectrum_name == "Oxyhemoglobin":
-            hbO2 = molecule.volume_fraction
+        spectrum_name = molecule.spectrum.spectrum_name
+        if spectrum_name in hemoglobin:
+            hemoglobin[spectrum_name] = molecule.volume_fraction
 
-    if hb is None and hbO2 is None:
-        return None
+    return hemoglobin
 
-    if hb is None:
-        hb = 0
-    elif hbO2 is None:
-        hbO2 = 0
 
-    if hb + hbO2 < 1e-10:  # negative values are not allowed and division by (approx) zero
-        return None        # will lead to negative side effects.
+def calculate_oxygenation(molecule_list: List) -> Optional[float]:
+    """
+    Calculate the oxygenation level based on the volume fractions of deoxyhemoglobin and oxyhemoglobin.
+
+    :param molecule_list: List of molecules with their spectrum information and volume fractions.
+    :return: An oxygenation value between 0 and 1 if possible, or None if not computable.
+    """
+    hemoglobin = extract_hemoglobin_fractions(molecule_list)
+    hb, hbO2 = hemoglobin["Deoxyhemoglobin"], hemoglobin["Oxyhemoglobin"]
 
-    return hbO2 / (hb + hbO2)
+    total = hb + hbO2
 
+    # Avoid division by zero. If none of the hemoglobin types are present, the oxygenation level is not computable.
+    if total < 1e-10:
+        return None
 
-def calculate_bvf(molecule_list: list) -> Union[float, int]:
-    """
-    Calculate the blood volume fraction based on the volume fractions of deoxyhaemoglobin and oxyhaemoglobin.
+    return hbO2 / total
 
-    This function takes a list of molecules and returns a blood volume fraction value between 0 and 1.
 
+def calculate_bvf(molecule_list: List) -> Union[float, int]:
+    """
+    Calculate the blood volume fraction based on the volume fractions of deoxyhemoglobin and oxyhemoglobin.
+
     :param molecule_list: List of molecules with their spectrum information and volume fractions.
     :return: The blood volume fraction value between 0 and 1, or 0, if oxy and deoxy not present.
     """
-    hb = None # Volume fraction of deoxyhaemoglobin
-    hbO2 = None # Volume fraction of oxyhaemoglobin
-
-    for molecule in molecule_list:
-        if molecule.spectrum.spectrum_name == "Deoxyhemoglobin":
-            hb = molecule.volume_fraction
-        elif molecule.spectrum.spectrum_name == "Oxyhemoglobin":
-            hbO2 = molecule.volume_fraction
-
-    if hb is None and hbO2 is None:
-        return 0
-
-    if hb is None:
-        return hbO2
-
-    if hbO2 is None:
-        return hb
-
+    hemoglobin = extract_hemoglobin_fractions(molecule_list)
+    hb, hbO2 = hemoglobin["Deoxyhemoglobin"], hemoglobin["Oxyhemoglobin"]
+    # We can use the sum of hb and hb02 to compute blood volume fraction as the volume fraction of all molecules is 1.
     return hb + hbO2