Get wavelengths to plot

Author: donato-ax

src/axiomatic/pic_helpers.py

@@ -1,3 +1,5 @@
+import re
+import numpy as np
 import iklayout  # type: ignore
 import matplotlib.pyplot as plt  # type: ignore
 from ipywidgets import interactive, IntSlider  # type: ignore
@@ -301,3 +303,53 @@ def print_statements(statements: StatementDictionary, validation: Optional[State
         print("Statement:", unf_stmt.text)
         print("Formalization: UNFORMALIZABLE")
         print("\n-----------------------------------\n")
+
+
+def _str_units_to_float(str_units: str) -> float:
+    unit_conversions = {
+        "nm": 1e-3,
+        "um": 1,
+        "mm": 1e3,
+        "m": 1e6,
+    }
+    match = re.match(r"([\d\.]+)\s*([a-zA-Z]+)", str_units)
+    numeric_value = float(match.group(1))
+    unit = match.group(2)
+    return float(numeric_value * unit_conversions[unit])
+
+
+def get_wavelengths_to_plot(
+    statements: StatementDictionary, num_samples: int = 100
+) -> tuple[list[float], list[float]]:
+    """
+    Get the wavelengths to plot based on the statements.
+
+    Returns a list of wavelengths to plot the spectra and a list of vertical lines to plot on top the spectra.
+    """
+
+    min_wl = np.inf
+    max_wl = -np.inf
+    vlines = set()
+    for stmt in statements.parameter_constraints + statements.cost_functions:
+        if stmt.formalization is not None and stmt.formalization.mapping is not None:
+            for comp in stmt.formalization.mapping.values():
+                if "wavelengths" in comp.arguments:
+                    vlines = vlines | {
+                        _str_units_to_float(wl) for wl in comp.arguments["wavelengths"]
+                    }
+                if "wavelength_range" in comp.arguments:
+                    min_wl = min(min_wl, _str_units_to_float(comp.arguments["wavelength_range"][0]))
+                    max_wl = max(max_wl, _str_units_to_float(comp.arguments["wavelength_range"][1]))
+    if vlines:
+        min_wl = min(min_wl, min(vlines))
+        max_wl = max(max_wl, max(vlines))
+    if min_wl >= max_wl:
+        avg_wl = sum(vlines) / len(vlines) if vlines else 1550
+        min_wl, max_wl = avg_wl - 0.1, avg_wl + 0.1
+    else:
+        range_size = max_wl - min_wl
+        min_wl -= 0.2 * range_size
+        max_wl += 0.2 * range_size
+
+    wls = np.linspace(min_wl, max_wl, num_samples)
+    return wls.tolist(), list(vlines)