Move mary

examples/cidnp.py

@@ -8,12 +8,14 @@
 
 import radicalpy as rp
 from radicalpy import kinetics, relaxation
+from radicalpy.experiments import mary_lfe_hfe, mary_loop
 from radicalpy.shared import constants as C
 from radicalpy.simulation import LiouvilleSimulation as Liouv
 from radicalpy.simulation import State
+from radicalpy.utils import is_fast_run
 
 
-def main():
+def main(tmax=5e-6, dt=5e-9, Bmax=20000, dB=100):
     # Create a radical pair with a 13C nucleus coupled to radical 1
     r1 = rp.simulation.Molecule.fromisotopes(isotopes=["13C"], hfcs=[0.0])
     r2 = rp.simulation.Molecule("radical 2")
@@ -36,8 +38,8 @@ def main():
     xi = 0.98
 
     init_state = State.SINGLET
-    time = np.arange(0, 5e-6, 5e-9)
-    B = np.arange(0, 20000, 100)
+    time = np.arange(0, tmax, dt)
+    B = np.arange(0, Bmax, dB)
     kinetic = [
         kinetics.Haberkorn(ks, State.SINGLET),
         kinetics.Haberkorn(kt, State.TRIPLET),
@@ -71,7 +73,7 @@ def main():
 
     # Run the magnetic field loop
     sim.apply_liouville_hamiltonian_modifiers(HL, kinetic + relaxations)
-    rhos = sim.mary_loop(init_state, time, B, HL, theta=None, phi=None)
+    rhos = mary_loop(sim, init_state, time, B, HL, theta=None, phi=None)
 
     # Calculate CIDNP of both singlet and triplet yields --> see Eqs. 6 and 7
     product_probabilities1 = np.real(np.trace(obs_state1 @ rhos, axis1=-2, axis2=-1))
@@ -85,7 +87,7 @@ def main():
     )
 
     dt = time[1] - time[0]
-    CIDNP, LFE, HFE = sim.mary_lfe_hfe(init_state, B, product_probabilities, dt, k)
+    CIDNP, LFE, HFE = mary_lfe_hfe(init_state, B, product_probabilities, dt, k)
 
     # Normalise the CIDNP intensity and plot and save figure
     norm_CIDNP = (CIDNP - CIDNP.max()) / (CIDNP.max() - CIDNP.min())
@@ -99,4 +101,7 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    if is_fast_run():
+        main(tmax=5e-6, dt=5e-7, Bmax=20000, dB=1000)
+    else:
+        main()

radicalpy/experiments.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 import itertools
+from typing import Optional
 
 import numpy as np
 import scipy as sp
@@ -17,6 +18,101 @@
 from .utils import anisotropy_check, mary_lorentzian, modulated_signal, reference_signal
 
 
+def mary_lfe_hfe(
+    init_state: State,
+    B: np.ndarray,
+    product_probability_seq: np.ndarray,
+    dt: float,
+    k: float,
+) -> (np.ndarray, np.ndarray, np.ndarray):
+    """Calculate MARY, LFE, HFE."""
+    MARY = np.sum(product_probability_seq, axis=1) * dt * k
+    idx = int(len(MARY) / 2) if B[0] != 0 else 0
+    minmax = max if init_state == State.SINGLET else min
+    HFE = (MARY[-1] - MARY[idx]) / MARY[idx] * 100
+    LFE = (minmax(MARY) - MARY[idx]) / MARY[idx] * 100
+    MARY = (MARY - MARY[idx]) / MARY[idx] * 100
+    return MARY, LFE, HFE
+
+
+def mary_loop(
+    sim: HilbertSimulation,
+    init_state: State,
+    time: np.ndarray,
+    B: np.ndarray,
+    H_base: np.ndarray,
+    theta: Optional[float] = None,
+    phi: Optional[float] = None,
+    hfc_anisotropy: bool = False,
+) -> np.ndarray:
+    """Generate density matrices (rhos) for MARY.
+
+    Args:
+
+        init_state (State): initial state.
+
+    Returns:
+        np.ndarray:
+
+            Density matrices.
+
+    .. todo:: Write proper docs.
+    """
+    H_zee = sim.convert(sim.zeeman_hamiltonian(1, theta, phi))
+    shape = sim._get_rho_shape(H_zee.shape[0])
+    rhos = np.zeros([len(B), len(time), *shape], dtype=complex)
+    for i, B0 in enumerate(tqdm(B)):
+        H = H_base + B0 * H_zee
+        H_sparse = sp.sparse.csc_matrix(H)
+        rhos[i] = sim.time_evolution(init_state, time, H_sparse)
+    return rhos
+
+
+def mary(
+    sim: HilbertSimulation,
+    init_state: State,
+    obs_state: State,
+    time: np.ndarray,
+    B: np.ndarray,
+    D: float,
+    J: float,
+    kinetics: list[HilbertIncoherentProcessBase] = [],
+    relaxations: list[HilbertIncoherentProcessBase] = [],
+    theta: Optional[float] = None,
+    phi: Optional[float] = None,
+    hfc_anisotropy: bool = False,
+) -> dict:
+    H = sim.total_hamiltonian(B0=0, D=D, J=J, hfc_anisotropy=hfc_anisotropy)
+
+    sim.apply_liouville_hamiltonian_modifiers(H, kinetics + relaxations)
+    rhos = mary_loop(sim, init_state, time, B, H, theta=theta, phi=phi)
+    product_probabilities = sim.product_probability(obs_state, rhos)
+
+    sim.apply_hilbert_kinetics(time, product_probabilities, kinetics)
+    k = kinetics[0].rate_constant if kinetics else 1.0
+    product_yields, product_yield_sums = sim.product_yield(
+        product_probabilities, time, k
+    )
+
+    dt = time[1] - time[0]
+    MARY, LFE, HFE = mary_lfe_hfe(init_state, B, product_probabilities, dt, k)
+    rhos = sim._square_liouville_rhos(rhos)
+
+    return dict(
+        time=time,
+        B=B,
+        theta=theta,
+        phi=phi,
+        rhos=rhos,
+        time_evolutions=product_probabilities,
+        product_yields=product_yields,
+        product_yield_sums=product_yield_sums,
+        MARY=MARY,
+        LFE=LFE,
+        HFE=HFE,
+    )
+
+
 def modulated_mary_brute_force(
     Bs: np.ndarray,
     modulation_depths: list,

radicalpy/simulation.py

@@ -683,55 +683,6 @@ def apply_hilbert_kinetics(time, product_probabilities, kinetics):
         for K in kinetics:  # skip in liouville
             K.adjust_product_probabilities(product_probabilities, time)
 
-    def mary_loop(
-        self,
-        init_state: State,
-        time: np.ndarray,
-        B: np.ndarray,
-        H_base: np.ndarray,
-        theta: Optional[float] = None,
-        phi: Optional[float] = None,
-        hfc_anisotropy: bool = False,
-    ) -> np.ndarray:
-        """Generate density matrices (rhos) for MARY.
-
-        Args:
-
-            init_state (State): initial state.
-
-        Returns:
-            np.ndarray:
-
-                Density matrices.
-
-        .. todo:: Write proper docs.
-        """
-        H_zee = self.convert(self.zeeman_hamiltonian(1, theta, phi))
-        shape = self._get_rho_shape(H_zee.shape[0])
-        rhos = np.zeros([len(B), len(time), *shape], dtype=complex)
-        for i, B0 in enumerate(tqdm(B)):
-            H = H_base + B0 * H_zee
-            H_sparse = sp.sparse.csc_matrix(H)
-            rhos[i] = self.time_evolution(init_state, time, H_sparse)
-        return rhos
-
-    @staticmethod
-    def mary_lfe_hfe(
-        init_state: State,
-        B: np.ndarray,
-        product_probability_seq: np.ndarray,
-        dt: float,
-        k: float,
-    ) -> (np.ndarray, np.ndarray, np.ndarray):
-        """Calculate MARY, LFE, HFE."""
-        MARY = np.sum(product_probability_seq, axis=1) * dt * k
-        idx = int(len(MARY) / 2) if B[0] != 0 else 0
-        minmax = max if init_state == State.SINGLET else min
-        HFE = (MARY[-1] - MARY[idx]) / MARY[idx] * 100
-        LFE = (minmax(MARY) - MARY[idx]) / MARY[idx] * 100
-        MARY = (MARY - MARY[idx]) / MARY[idx] * 100
-        return MARY, LFE, HFE
-
     @staticmethod
     def _square_liouville_rhos(rhos):
         return rhos
@@ -740,50 +691,6 @@ def _square_liouville_rhos(rhos):
     def _get_rho_shape(dim):
         return dim, dim
 
-    def MARY(
-        self,
-        init_state: State,
-        obs_state: State,
-        time: np.ndarray,
-        B: np.ndarray,
-        D: float,
-        J: float,
-        kinetics: list[HilbertIncoherentProcessBase] = [],
-        relaxations: list[HilbertIncoherentProcessBase] = [],
-        theta: Optional[float] = None,
-        phi: Optional[float] = None,
-        hfc_anisotropy: bool = False,
-    ) -> dict:
-        H = self.total_hamiltonian(B0=0, D=D, J=J, hfc_anisotropy=hfc_anisotropy)
-
-        self.apply_liouville_hamiltonian_modifiers(H, kinetics + relaxations)
-        rhos = self.mary_loop(init_state, time, B, H, theta=theta, phi=phi)
-        product_probabilities = self.product_probability(obs_state, rhos)
-
-        self.apply_hilbert_kinetics(time, product_probabilities, kinetics)
-        k = kinetics[0].rate_constant if kinetics else 1.0
-        product_yields, product_yield_sums = self.product_yield(
-            product_probabilities, time, k
-        )
-
-        dt = time[1] - time[0]
-        MARY, LFE, HFE = self.mary_lfe_hfe(init_state, B, product_probabilities, dt, k)
-        rhos = self._square_liouville_rhos(rhos)
-
-        return dict(
-            time=time,
-            B=B,
-            theta=theta,
-            phi=phi,
-            rhos=rhos,
-            time_evolutions=product_probabilities,
-            product_yields=product_yields,
-            product_yield_sums=product_yield_sums,
-            MARY=MARY,
-            LFE=LFE,
-            HFE=HFE,
-        )
-
     @staticmethod
     def convert(H: np.ndarray) -> np.ndarray:
         return H

tests/test_simulation.py

@@ -10,6 +10,7 @@
 
 import radicalpy as rp
 from radicalpy import estimations, kinetics, relaxation
+from radicalpy.experiments import mary
 from radicalpy.simulation import Basis
 
 # quick test:
@@ -361,7 +362,8 @@ def test_mary(self):
             for obs_state in rp.simulation.State:
                 if obs_state == rp.simulation.State.EQUILIBRIUM:
                     continue
-                rslt = self.sim.MARY(
+                rslt = mary(
+                    self.sim,
                     init_state,
                     obs_state,
                     self.time,
@@ -421,7 +423,8 @@ def test_ST_vs_Zeeman_basis(self):
         plt.show()
 
     def test_hyperfine_3d(self):
-        results = self.sim.MARY(
+        results = mary(
+            self.sim,
             rp.simulation.State.SINGLET,
             rp.simulation.State.TRIPLET,
             self.time,
@@ -565,7 +568,8 @@ def test_relaxation(self):
             J=0,
         )
         k = 1e6
-        results = self.sim.MARY(
+        results = mary(
+            self.sim,
             kinetics=[],
             relaxations=[
                 # relaxation.SingletTripletDephasing( k),