From 19cadb86bf06e41b1843debc23416910be6ec1a0 Mon Sep 17 00:00:00 2001
From: Manuel Huber <info@manuelhu.de>
Date: Thu, 26 Sep 2024 00:19:19 +0200
Subject: [PATCH] add store api to manipulate properties

---
 src/legendoptics/copper.py    |  2 +
 src/legendoptics/fibers.py    |  9 ++++
 src/legendoptics/germanium.py |  2 +
 src/legendoptics/lar.py       | 12 ++++++
 src/legendoptics/nylon.py     |  3 ++
 src/legendoptics/pen.py       |  9 ++++
 src/legendoptics/silicon.py   |  2 +
 src/legendoptics/store.py     | 78 +++++++++++++++++++++++++++++++++++
 src/legendoptics/tetratex.py  |  2 +
 src/legendoptics/tpb.py       |  7 ++++
 tests/test_store.py           | 43 +++++++++++++++++++
 11 files changed, 169 insertions(+)
 create mode 100644 src/legendoptics/store.py
 create mode 100644 tests/test_store.py

diff --git a/src/legendoptics/copper.py b/src/legendoptics/copper.py
index 301ddb5..44df0db 100644
--- a/src/legendoptics/copper.py
+++ b/src/legendoptics/copper.py
@@ -7,12 +7,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def copper_reflectivity() -> tuple[Quantity, Quantity]:
     """Reflectivity of copper surfaces.
 
diff --git a/src/legendoptics/fibers.py b/src/legendoptics/fibers.py
index ff77dce..47436d7 100644
--- a/src/legendoptics/fibers.py
+++ b/src/legendoptics/fibers.py
@@ -12,12 +12,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import InterpolatingGraph, readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def fiber_cladding2_refractive_index() -> float:
     """Refractive index of second fiber cladding material [SaintGobainDataSheet]_.
 
@@ -26,6 +28,7 @@ def fiber_cladding2_refractive_index() -> float:
     return 1.42
 
 
+@store.register_pluggable
 def fiber_cladding1_refractive_index() -> float:
     """Refractive index of first fiber cladding material [SaintGobainDataSheet]_.
 
@@ -34,6 +37,7 @@ def fiber_cladding1_refractive_index() -> float:
     return 1.49
 
 
+@store.register_pluggable
 def fiber_core_refractive_index() -> float:
     """Refractive index of fiber core material [SaintGobainDataSheet]_.
 
@@ -42,6 +46,7 @@ def fiber_core_refractive_index() -> float:
     return 1.6
 
 
+@store.register_pluggable
 def fiber_wls_absorption(
     abs_at_400nm: Quantity = 0.7 * u.mm,
 ) -> tuple[Quantity, Quantity]:
@@ -76,6 +81,7 @@ def fiber_wls_absorption(
     return wvl, absorp
 
 
+@store.register_pluggable
 def fiber_wls_emission() -> tuple[Quantity, Quantity]:
     """[SaintGobainDataSheet]_ reports the emission spectrum for BCF-91A.
 
@@ -84,6 +90,7 @@ def fiber_wls_emission() -> tuple[Quantity, Quantity]:
     return readdatafile("psfibers_wlscomponent.dat")
 
 
+@store.register_pluggable
 def fiber_wls_timeconstant() -> Quantity:
     """WLS time constant [SaintGobainDataSheet]_.
 
@@ -92,6 +99,7 @@ def fiber_wls_timeconstant() -> Quantity:
     return 12 * u.ns
 
 
+@store.register_pluggable
 def fiber_absorption_length() -> Quantity:
     """Absorption length of fiber [SaintGobainDataSheet]_. Note this is a macroscopical value for a 1 mm fiber.
 
@@ -105,6 +113,7 @@ def fiber_absorption_length() -> Quantity:
     return 3.5 * u.m
 
 
+@store.register_pluggable
 def fiber_absorption_path_length() -> Quantity:
     """Absorption length of fiber [SaintGobainDataSheet]_, corrected for the geometry of a 1 mm square fiber.
 
diff --git a/src/legendoptics/germanium.py b/src/legendoptics/germanium.py
index 0abc454..b407725 100644
--- a/src/legendoptics/germanium.py
+++ b/src/legendoptics/germanium.py
@@ -14,12 +14,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def germanium_reflectivity() -> tuple[Quantity, Quantity]:
     """Reflectivity of germanium surfaces.
 
diff --git a/src/legendoptics/lar.py b/src/legendoptics/lar.py
index 15e824d..5c719a8 100644
--- a/src/legendoptics/lar.py
+++ b/src/legendoptics/lar.py
@@ -34,6 +34,7 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.scintillate import ScintConfig, ScintParticle
 from legendoptics.utils import (
     InterpolatingGraph,
@@ -113,6 +114,7 @@ def lar_dielectric_constant_cern2020(
     return (3 + 2 * x) / (3 - x)
 
 
+@store.register_pluggable
 def lar_dielectric_constant(
     λ: Quantity, method: ArDielectricMethods = "cern2020"
 ) -> Quantity:
@@ -130,6 +132,7 @@ def lar_dielectric_constant(
     raise ValueError(msg)
 
 
+@store.register_pluggable
 def lar_refractive_index(
     λ: Quantity, method: ArDielectricMethods = "cern2020"
 ) -> Quantity:
@@ -145,6 +148,7 @@ def lar_refractive_index(
     return np.sqrt(lar_dielectric_constant(λ, method))
 
 
+@store.register_pluggable
 def lar_emission_spectrum(λ: Quantity) -> Quantity:
     """Return the LAr emission spectrum, adapted from [Heindl2010]_.
 
@@ -160,6 +164,7 @@ def lar_emission_spectrum(λ: Quantity) -> Quantity:
     )(λ)
 
 
+@store.register_pluggable
 def lar_fano_factor() -> float:
     """Fano factor.
 
@@ -171,6 +176,7 @@ def lar_fano_factor() -> float:
     return 0.11
 
 
+@store.register_pluggable
 def lar_rayleigh(
     λ: Quantity,
     temperature: Quantity = 90 * u.K,
@@ -211,6 +217,7 @@ def lar_rayleigh(
     return (1 / inv_l).to("cm")  # simplify units
 
 
+@store.register_pluggable
 def lar_abs_length(λ: Quantity) -> Quantity:
     """Absorption length (not correctly scaled).
 
@@ -227,6 +234,7 @@ def lar_abs_length(λ: Quantity) -> Quantity:
     return np.minimum(absl, 100000 * u.cm)  # avoid large numbers
 
 
+@store.register_pluggable
 def lar_peak_attenuation_length(
     attenuation_method: ArLifetimeMethods | Quantity = "legend200-llama",
 ) -> Quantity:
@@ -244,6 +252,7 @@ def lar_peak_attenuation_length(
     return attenuation_method
 
 
+@store.register_pluggable
 def lar_lifetimes(
     triplet_lifetime_method: float | ArLifetimeMethods = "legend200-llama",
 ) -> ArScintLiftime:
@@ -265,6 +274,7 @@ def lar_lifetimes(
     return ArScintLiftime(singlet=5.95 * u.ns, triplet=triplet)
 
 
+@store.register_pluggable
 def lar_scintillation_params(
     flat_top_yield: Quantity = 31250 / u.MeV,
 ) -> ScintConfig:
@@ -308,6 +318,7 @@ def lar_scintillation_params(
     )
 
 
+@store.register_pluggable
 def pyg4_lar_attach_rindex(
     lar_mat, reg, lar_dielectric_method: ArDielectricMethods = "cern2020"
 ) -> None:
@@ -332,6 +343,7 @@ def pyg4_lar_attach_rindex(
         lar_mat.addVecPropertyPint("RINDEX", λ_full.to("eV"), rindex)
 
 
+@store.register_pluggable
 def pyg4_lar_attach_attenuation(
     lar_mat,
     reg,
diff --git a/src/legendoptics/nylon.py b/src/legendoptics/nylon.py
index cccf89e..7e1abdb 100644
--- a/src/legendoptics/nylon.py
+++ b/src/legendoptics/nylon.py
@@ -17,12 +17,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def nylon_refractive_index() -> float:
     """Refractive index in near-UV range, from [Benziger2007]_.
 
@@ -31,6 +33,7 @@ def nylon_refractive_index() -> float:
     return 1.53
 
 
+@store.register_pluggable
 def nylon_absorption() -> tuple[Quantity, Quantity]:
     """Values reported in [Agostini2018]_.
 
diff --git a/src/legendoptics/pen.py b/src/legendoptics/pen.py
index 39fd6f7..87bf9b4 100644
--- a/src/legendoptics/pen.py
+++ b/src/legendoptics/pen.py
@@ -19,6 +19,7 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.scintillate import ScintConfig, ScintParticle
 from legendoptics.utils import (
     InterpolatingGraph,
@@ -30,6 +31,7 @@
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def pen_quantum_efficiency() -> float:
     """Quantum efficiency, from [Araujo2022]_ at LAr temperature.
 
@@ -38,6 +40,7 @@ def pen_quantum_efficiency() -> float:
     return 0.69
 
 
+@store.register_pluggable
 def pen_refractive_index() -> float:
     """Refractive index from [Hong2017]_.
 
@@ -46,6 +49,7 @@ def pen_refractive_index() -> float:
     return 1.51
 
 
+@store.register_pluggable
 def pen_scint_timeconstant() -> Quantity:
     """Time constant, from [Manzanillas2022]_.
 
@@ -54,6 +58,7 @@ def pen_scint_timeconstant() -> Quantity:
     return 25.3 * u.ns
 
 
+@store.register_pluggable
 def pen_scint_light_yield() -> Quantity:
     """PEN scintillation yield for electrons, from [Manzanillas2022]_.
 
@@ -62,6 +67,7 @@ def pen_scint_light_yield() -> Quantity:
     return 5440 / u.MeV
 
 
+@store.register_pluggable
 def pen_wls_emission() -> tuple[Quantity, Quantity]:
     """WLS Emission spectrum, from [Mary1997]_.
 
@@ -70,6 +76,7 @@ def pen_wls_emission() -> tuple[Quantity, Quantity]:
     return readdatafile("pen_wlscomponent.dat")
 
 
+@store.register_pluggable
 def pen_absorption() -> tuple[Quantity, Quantity]:
     """Bulk absorption reported in [Manzanillas2022]_.
 
@@ -80,6 +87,7 @@ def pen_absorption() -> tuple[Quantity, Quantity]:
     return wvl, absorp
 
 
+@store.register_pluggable
 def pen_wls_absorption() -> tuple[Quantity, Quantity]:
     """WLS absorption of PEN.
 
@@ -99,6 +107,7 @@ def pen_wls_absorption() -> tuple[Quantity, Quantity]:
     return wvl, absorp
 
 
+@store.register_pluggable
 def pen_scintillation_params() -> ScintConfig:
     """Get a :class:`ScintConfig` object for PEN.
 
diff --git a/src/legendoptics/silicon.py b/src/legendoptics/silicon.py
index a626ac3..db2aa42 100644
--- a/src/legendoptics/silicon.py
+++ b/src/legendoptics/silicon.py
@@ -12,12 +12,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def silicon_complex_rindex() -> tuple[Quantity, Quantity, Quantity]:
     """Real and imaginary parts as tuple(wavelength, Re, Im). Measurements from [Phillip1960]_.
 
diff --git a/src/legendoptics/store.py b/src/legendoptics/store.py
new file mode 100644
index 0000000..90facb8
--- /dev/null
+++ b/src/legendoptics/store.py
@@ -0,0 +1,78 @@
+"""
+A store that allows to manipulate and swap individual material properties with custom
+implementations.
+
+:meth:`register_pluggable` is a decorator to use above all functions defining material
+properties.
+
+Users can then replace the original implementation by their own implementations using
+``replace_implementation(new_impl: Callable)``, and also switch back to the original
+implementation using ``reset_implementation()`` on the decorated function object. The
+original implementation is always available as ``original_impl()``.
+
+Apart from the decorator, this store provides functions to get and reset the status of
+all registered pluggable functions.
+"""
+
+from __future__ import annotations
+
+from functools import wraps
+from types import MethodType
+from typing import Callable
+
+_optical_property_store = []
+
+
+def register_pluggable(fn: Callable) -> Callable:
+    """Decorator that registers this function as a pluggable property function."""
+
+    # create the new wrapper object.
+    @wraps(fn)
+    def wrap(*args, **kwargs):
+        return wrap._impl(*args, **kwargs)
+
+    wrap._impl = fn
+    wrap._orig_impl = fn
+
+    # add "instance methods" of the new wrapper.
+    def reset_implementation(self) -> None:
+        """Reset to the original function implementation."""
+        self._impl = self._orig_impl
+
+    def replace_implementation(self, new_impl: Callable) -> None:
+        """Replace the underlying function implementation."""
+        self._impl = new_impl
+
+    def is_original(self) -> bool:
+        """Is the underlying function the original implementation"""
+        return self._impl == self._orig_impl
+
+    def original_impl(self) -> Callable:
+        """The original function implementation."""
+        return self._orig_impl
+
+    wrap.reset_implementation = MethodType(reset_implementation, wrap)
+    wrap.replace_implementation = MethodType(replace_implementation, wrap)
+    wrap.is_original = MethodType(is_original, wrap)
+    wrap.original_impl = MethodType(original_impl, wrap)
+
+    # store the wrapper to use it in the functions of this module.
+    _optical_property_store.append(wrap)
+
+    return wrap
+
+
+def reset_all_to_original() -> None:
+    """Reset all pluggable material property functions to their original implementations."""
+    for entry in _optical_property_store:
+        entry.reset_implementation()
+
+
+def is_all_original() -> bool:
+    """Get whether all pluggable material property functions use their original implementation."""
+    return all(p.is_original() for p in _optical_property_store)
+
+
+def get_replaced() -> list[str]:
+    """Get the names of all replaced pluggable material property functions."""
+    return [p.__name__ for p in _optical_property_store if not p.is_original()]
diff --git a/src/legendoptics/tetratex.py b/src/legendoptics/tetratex.py
index b87ad7e..e43f3a0 100644
--- a/src/legendoptics/tetratex.py
+++ b/src/legendoptics/tetratex.py
@@ -11,12 +11,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def tetratex_reflectivity() -> tuple[Quantity, Quantity]:
     """Tetratex reflectivity from [Janacek2012]_.
 
diff --git a/src/legendoptics/tpb.py b/src/legendoptics/tpb.py
index a8d34cf..d3d57d4 100644
--- a/src/legendoptics/tpb.py
+++ b/src/legendoptics/tpb.py
@@ -23,12 +23,14 @@
 import pint
 from pint import Quantity
 
+from legendoptics import store
 from legendoptics.utils import InterpolatingGraph, readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
 
 
+@store.register_pluggable
 def tpb_quantum_efficiency() -> float:
     """Quantum efficiency.
 
@@ -40,6 +42,7 @@ def tpb_quantum_efficiency() -> float:
     return 0.85
 
 
+@store.register_pluggable
 def tpb_refractive_index() -> float:
     """Refractive index from [MolbaseTPB]_.
 
@@ -48,6 +51,7 @@ def tpb_refractive_index() -> float:
     return 1.635
 
 
+@store.register_pluggable
 def tpb_wls_timeconstant() -> Quantity:
     """Time constant: arbitrary small.
 
@@ -56,6 +60,7 @@ def tpb_wls_timeconstant() -> Quantity:
     return 0.01 * u.ns
 
 
+@store.register_pluggable
 def tpb_wls_emission() -> tuple[Quantity, Quantity]:
     """WLS Emission spectrum.
 
@@ -69,6 +74,7 @@ def tpb_wls_emission() -> tuple[Quantity, Quantity]:
     return readdatafile("tpb_vm2000_wlscomponent.dat")
 
 
+@store.register_pluggable
 def tpb_polystyrene_wls_emission() -> tuple[Quantity, Quantity]:
     """WLS Emission spectrum for TPB in a polystyrene matrix.
 
@@ -82,6 +88,7 @@ def tpb_polystyrene_wls_emission() -> tuple[Quantity, Quantity]:
     return readdatafile("tpb_polystyrene_wlscomponent.dat")
 
 
+@store.register_pluggable
 def tpb_wls_absorption() -> tuple[Quantity, Quantity]:
     """Values reported in [Benson2018]_ for TPB evaporated on utraviolet-transmitting acrylic substrate.
 
diff --git a/tests/test_store.py b/tests/test_store.py
new file mode 100644
index 0000000..c68e633
--- /dev/null
+++ b/tests/test_store.py
@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from legendoptics import store
+from legendoptics.fibers import (
+    fiber_cladding1_refractive_index,
+    fiber_core_refractive_index,
+)
+
+
+def test_store():
+    assert fiber_core_refractive_index() == 1.6
+    assert fiber_core_refractive_index.is_original()
+
+    # test replacing the implementation.
+    fiber_core_refractive_index.replace_implementation(lambda: 1234)
+    assert fiber_core_refractive_index() == 1234
+    assert not fiber_core_refractive_index.is_original()
+
+    fiber_cladding1_refractive_index.replace_implementation(lambda: 1)
+
+    # test single reset.
+    fiber_core_refractive_index.reset_implementation()
+    assert fiber_core_refractive_index() == 1.6
+    assert fiber_core_refractive_index.is_original()
+    # ... other properties are not reset:
+    assert not store.is_all_original()
+    assert fiber_cladding1_refractive_index() == 1
+
+    # test global reset.
+    fiber_core_refractive_index.replace_implementation(lambda: 5678)
+    assert not fiber_core_refractive_index.is_original()
+
+    assert store.get_replaced() == [
+        "fiber_cladding1_refractive_index",
+        "fiber_core_refractive_index",
+    ]
+
+    store.reset_all_to_original()
+    assert fiber_core_refractive_index() == 1.6
+    assert fiber_core_refractive_index.is_original()
+    assert fiber_cladding1_refractive_index() == 1.49  # now also reset.
+    assert store.get_replaced() == []
+    assert store.is_all_original()