feat(experiments): Add new R1rho experiments (WIP)

- Implement two R1rho calculation methods:
  1. `wip.relaxation_15n_r1rho_eig`: Uses smallest real eigenvalue
  2. `wip.relaxation_15n_r1rho`: Simulates magnetization tilt and spin lock

- Place new experiments in `experiments.catalog.wip`
- Update experiment loader to register all plugins

Note: Work in progress, further refinements expected

chemex/experiments/catalog/wip/relaxation_15n_r1rho.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+
+from chemex.configuration.base import ExperimentConfiguration, ToBeFitted
+from chemex.configuration.conditions import ConditionsWithValidations
+from chemex.configuration.data import RelaxationDataSettings
+from chemex.configuration.experiment import ExperimentSettings
+from chemex.containers.data import Data
+from chemex.containers.dataset import load_relaxation_dataset
+from chemex.experiments.factories import Creators, factories
+from chemex.filterers import PlanesFilterer
+from chemex.nmr.basis import Basis
+from chemex.nmr.liouvillian import LiouvillianIS
+from chemex.nmr.spectrometer import Spectrometer
+from chemex.parameters.spin_system import SpinSystem
+from chemex.plotters.relaxation import RelaxationPlotter
+from chemex.printers.data import RelaxationPrinter
+from chemex.typing import ArrayBool, ArrayFloat
+
+EXPERIMENT_NAME = "wip.relaxation_15n_r1rho"
+
+
+class Relaxation15NR1RhoSettings(ExperimentSettings):
+    name: Literal["wip.relaxation_15n_r1rho"]
+
+    carrier: float
+    b1_frq: float
+    b1_inh_scale: float = np.inf
+    b1_inh_res: int = 11
+    observed_state: Literal["a", "b", "c", "d"] = "a"
+
+    @property
+    def detection(self) -> str:
+        return f"[iz_{self.observed_state}]"
+
+
+class Relaxation15NR1RhoConfig(
+    ExperimentConfiguration[
+        Relaxation15NR1RhoSettings,
+        ConditionsWithValidations,
+        RelaxationDataSettings,
+    ],
+):
+    @property
+    def to_be_fitted(self) -> ToBeFitted:
+        state = self.experiment.observed_state
+        return ToBeFitted(
+            rates=[f"r2_i_{state}"],
+            model_free=[f"tauc_{state}", f"s2_{state}"],
+        )
+
+
+def build_spectrometer(
+    config: Relaxation15NR1RhoConfig,
+    spin_system: SpinSystem,
+) -> Spectrometer:
+    settings = config.experiment
+    conditions = config.conditions
+
+    basis = Basis(type="ixyz", spin_system="nh")
+    liouvillian = LiouvillianIS(spin_system, basis, conditions)
+    spectrometer = Spectrometer(liouvillian)
+
+    spectrometer.carrier_i = settings.carrier
+    spectrometer.b1_i = settings.b1_frq
+    spectrometer.b1_i_inh_scale = settings.b1_inh_scale
+    spectrometer.b1_i_inh_res = settings.b1_inh_res
+
+    spectrometer.detection = settings.detection
+
+    return spectrometer
+
+
+@dataclass
+class Relaxation15NR1RhoSequence:
+    settings: Relaxation15NR1RhoSettings
+
+    def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
+        times = data.metadata
+
+        # Getting the starting magnetization
+        start = spectrometer.get_equilibrium()
+
+        magnetization = spectrometer.tilt_mag_along_weff_i(start)
+        magnetization = spectrometer.pulse_i(times, phase=0.0) @ magnetization
+        magnetization = spectrometer.tilt_mag_along_weff_i(magnetization, back=True)
+
+        # Return profile
+        return np.array([spectrometer.detect(mag) for mag in magnetization])
+
+    def is_reference(self, metadata: ArrayFloat) -> ArrayBool:
+        return np.full_like(metadata, fill_value=False, dtype=np.bool_)
+
+
+def register() -> None:
+    creators = Creators(
+        config_creator=Relaxation15NR1RhoConfig,
+        spectrometer_creator=build_spectrometer,
+        sequence_creator=Relaxation15NR1RhoSequence,
+        dataset_creator=load_relaxation_dataset,
+        filterer_creator=PlanesFilterer,
+        printer_creator=RelaxationPrinter,
+        plotter_creator=RelaxationPlotter,
+    )
+    factories.register(name=EXPERIMENT_NAME, creators=creators)

chemex/experiments/catalog/wip/relaxation_15n_r1rho_eig.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+
+from chemex.configuration.base import ExperimentConfiguration, ToBeFitted
+from chemex.configuration.conditions import ConditionsWithValidations
+from chemex.configuration.data import RelaxationDataSettings
+from chemex.configuration.experiment import ExperimentSettings
+from chemex.containers.data import Data
+from chemex.containers.dataset import load_relaxation_dataset
+from chemex.experiments.factories import Creators, factories
+from chemex.filterers import PlanesFilterer
+from chemex.nmr.basis import Basis
+from chemex.nmr.liouvillian import LiouvillianIS
+from chemex.nmr.spectrometer import Spectrometer
+from chemex.parameters.spin_system import SpinSystem
+from chemex.plotters.relaxation import RelaxationPlotter
+from chemex.printers.data import RelaxationPrinter
+from chemex.typing import ArrayBool, ArrayFloat
+
+EXPERIMENT_NAME = "wip.relaxation_15n_r1rho_eig"
+
+
+class Relaxation15NR1RhoSettings(ExperimentSettings):
+    name: Literal["wip.relaxation_15n_r1rho_eig"]
+
+    carrier: float
+    b1_frq: float
+    b1_inh_scale: float = np.inf
+    b1_inh_res: int = 11
+    observed_state: Literal["a", "b", "c", "d"] = "a"
+
+    @property
+    def detection(self) -> str:
+        return f"[iz_{self.observed_state}]"
+
+
+class Relaxation15NR1RhoConfig(
+    ExperimentConfiguration[
+        Relaxation15NR1RhoSettings,
+        ConditionsWithValidations,
+        RelaxationDataSettings,
+    ],
+):
+    @property
+    def to_be_fitted(self) -> ToBeFitted:
+        state = self.experiment.observed_state
+        return ToBeFitted(
+            rates=[f"r2_i_{state}"],
+            model_free=[f"tauc_{state}", f"s2_{state}"],
+        )
+
+
+def build_spectrometer(
+    config: Relaxation15NR1RhoConfig,
+    spin_system: SpinSystem,
+) -> Spectrometer:
+    settings = config.experiment
+    conditions = config.conditions
+
+    basis = Basis(type="ixyz", spin_system="nh")
+    liouvillian = LiouvillianIS(spin_system, basis, conditions)
+    spectrometer = Spectrometer(liouvillian)
+
+    spectrometer.carrier_i = settings.carrier
+    spectrometer.b1_i = settings.b1_frq
+    spectrometer.b1_i_inh_scale = settings.b1_inh_scale
+    spectrometer.b1_i_inh_res = settings.b1_inh_res
+
+    spectrometer.detection = settings.detection
+
+    return spectrometer
+
+
+@dataclass
+class Relaxation15NR1RhoSequence:
+    settings: Relaxation15NR1RhoSettings
+
+    def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
+        times = data.metadata
+
+        r1rho = spectrometer.liouvillian.calculate_r1rho()
+
+        # Return profile
+        return np.exp(-r1rho * times)
+
+    def is_reference(self, metadata: ArrayFloat) -> ArrayBool:
+        return np.full_like(metadata, fill_value=False, dtype=np.bool_)
+
+
+def register() -> None:
+    creators = Creators(
+        config_creator=Relaxation15NR1RhoConfig,
+        spectrometer_creator=build_spectrometer,
+        sequence_creator=Relaxation15NR1RhoSequence,
+        dataset_creator=load_relaxation_dataset,
+        filterer_creator=PlanesFilterer,
+        printer_creator=RelaxationPrinter,
+        plotter_creator=RelaxationPlotter,
+    )
+    factories.register(name=EXPERIMENT_NAME, creators=creators)

chemex/experiments/loader.py

@@ -1,32 +1,54 @@
-"""A simple plugin loader."""
+"""A simple plugin loader for experiment modules."""
+
 from __future__ import annotations
 
 import importlib
+from collections.abc import Iterator
+from importlib.util import find_spec
 from pkgutil import iter_modules
-from typing import Protocol, cast
+from types import ModuleType
+from typing import Protocol
 
 from chemex.experiments import catalog
+from chemex.experiments.catalog import wip
 
 
-class ExperimentModuleInterface(Protocol):
-    """Represents an experiment module interface.
-    An experiment module plugin has a single register function.
-    """
+class ExperimentModule(Protocol):
+    """Represents an experiment module interface."""
 
     @staticmethod
     def register() -> None:
         """Register the necessary items related to the experiment."""
-        ...
 
 
-def import_module(name: str) -> ExperimentModuleInterface:
+def import_module(name: str) -> ExperimentModule:
     """Imports a module given a name."""
-    return cast(ExperimentModuleInterface, importlib.import_module(name))
+    return importlib.import_module(name)  # type: ignore
+
+
+def iter_experiment_modules(package: ModuleType) -> Iterator[str]:
+    """Yields experiment module names from a given package."""
+    return (
+        f"{package.__name__}.{module.name}"
+        for module in iter_modules(package.__path__)
+        if not module.ispkg
+    )
 
 
 def register_experiments() -> None:
-    """Loads the plugins defined in the plugins list."""
-    for module in iter_modules(catalog.__path__):
-        module_name = f"{catalog.__name__}.{module.name}"
-        experiment_module = import_module(module_name)
-        experiment_module.register()
+    """Loads and registers all experiment plugins."""
+    modules_to_register = [
+        *iter_experiment_modules(catalog),
+        *iter_experiment_modules(wip),
+    ]
+
+    for module_name in modules_to_register:
+        if find_spec(module_name):
+            experiment_module = import_module(module_name)
+            experiment_module.register()
+        else:
+            print(f"Warning: Module {module_name} not found")
+
+
+if __name__ == "__main__":
+    register_experiments()

chemex/nmr/basis.py

@@ -342,6 +342,12 @@ def name(self) -> str:
     def components(self) -> list[str]:
         return _BASES[self.type]
 
+    @property
+    def components_states(self) -> list[str]:
+        return [
+            f"{comp}_{state}" for state, comp in product(model.states, self.components)
+        ]
+
     @property
     def atoms(self) -> dict[str, str]:
         return {

chemex/nmr/liouvillian.py

@@ -30,6 +30,9 @@
 
 _Q_ORDER_I = {"sq": 1.0, "dq": 2.0, "tq": 3.0}
 
+# A small value used for numerical stability
+SMALL_VALUE = 1e-6
+
 
 def _make_gaussian(
     value: float,
@@ -271,6 +274,37 @@ def get_equilibrium(self) -> ArrayFloat:
             mag += self.vectors.get(f"{name}z_{state}", 0.0) * scale
         return mag
 
+    def tilt_mag_along_weff_i(
+        self, magnetization: ArrayFloat, *, back: bool = False
+    ) -> ArrayFloat:
+        basis = self.basis
+
+        w1 = self.b1_i * 2.0 * np.pi
+
+        for state in model.states:
+            index_x, index_z = (
+                basis.components_states.index(f"ix_{state}"),
+                basis.components_states.index(f"iz_{state}"),
+            )
+            components = magnetization[..., [index_x, index_z], :]
+
+            cs = self.par_values[f"cs_i_{state}"]
+            wi = -(
+                cs * self.ppm_i
+                - self.carrier_i * self.ppm_i
+                - self.offset_i * 2.0 * np.pi * np.sign(self.ppm_i)
+            )
+            theta = np.arctan2(w1, wi)
+            if back:
+                theta = -theta
+            rotation_matrix = np.array(
+                [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]
+            )
+            components_tilted = rotation_matrix @ components
+            magnetization[..., [index_x, index_z], :] = components_tilted
+
+        return magnetization
+
     def get_start_magnetization(
         self,
         terms: Iterable[str],
@@ -300,3 +334,10 @@ def offsets_to_ppms(self, offsets: ArrayFloat) -> ArrayFloat:
 
     def ppms_to_offsets(self, ppms: ArrayFloat | float) -> ArrayFloat | float:
         return (ppms - self.carrier_i) * abs(self.ppm_i) / (2.0 * np.pi)
+
+    def calculate_r1rho(self) -> float:
+        liouv = self.l_free + self.l_b1x_i
+        liouv = liouv.reshape((self.size, self.size))
+        eigenvalues = np.linalg.eigvals(liouv)
+        real_eigenvalues = eigenvalues[abs(eigenvalues.imag) < SMALL_VALUE].real
+        return -np.max(real_eigenvalues)

chemex/nmr/spectrometer.py

@@ -263,6 +263,11 @@ def get_start_magnetization(
     ) -> ArrayFloat:
         return self.liouvillian.get_start_magnetization(terms=terms, atom=atom)
 
+    def tilt_mag_along_weff_i(
+        self, magnetization: ArrayFloat, *, back: bool = False
+    ) -> ArrayFloat:
+        return self.liouvillian.tilt_mag_along_weff_i(magnetization, back=back)
+
     @property
     def detection(self) -> str:
         return self.liouvillian.detection