From d547cb37a10cd270b41c9ea33a33ed043b703c2f Mon Sep 17 00:00:00 2001
From: Emil VATAI <emil.vatai@gmail.com>
Date: Thu, 21 Mar 2024 22:41:38 +0900
Subject: [PATCH] Cleanup

---
 radicalpy/experiments.py | 12 ++++++------
 radicalpy/simulation.py  | 27 ++++++---------------------
 2 files changed, 12 insertions(+), 27 deletions(-)

diff --git a/radicalpy/experiments.py b/radicalpy/experiments.py
index 0998ec6..0d85d50 100644
--- a/radicalpy/experiments.py
+++ b/radicalpy/experiments.py
@@ -83,7 +83,7 @@ def semiclassical_mary(
     sim: SemiclassicalSimulation,
     num_samples: int,
     init_state: State,
-    ts: NDArray[float],
+    ts: ArrayLike,
     Bs: ArrayLike,
     D: float,
     J: float,
@@ -98,12 +98,12 @@ def semiclassical_mary(
     M = 16  # number of spin states
     trace = np.zeros((num_samples, len(ts)))
     mary = np.zeros((len(ts), len(Bs)))
-    gen = sim.semiclassical_gen2(num_samples)
+    HHs = sim.semiclassical_HHs(num_samples)
 
     for i, B0 in enumerate(tqdm(Bs)):
         Hz = sim.zeeman_hamiltonian(B0)
-        for j, Hnuc in enumerate(gen):
-            Ht = Hz + Hnuc
+        for j, HH in enumerate(HHs):
+            Ht = Hz + HH
             L = sim.convert(Ht)
             sim.apply_liouville_hamiltonian_modifiers(L, kinetics + relaxations)
             propagator = sp.sparse.linalg.expm(L * dt)
@@ -161,11 +161,11 @@ def semiclassical_kinetics_mary(
     rho_triplet = np.zeros(len(ts), dtype=complex)
     radical_pair_yield = np.zeros((1, len(ts)), dtype=complex)
     triplet_yield = np.zeros((1, len(ts)), dtype=complex)
-    gen = sim.semiclassical_gen2(num_samples)
+    HHs = sim.semiclassical_HHs(num_samples)
 
     for i, B0 in enumerate(tqdm(Bs)):
         Hz = sim.zeeman_hamiltonian(B0)
-        for j, Hnuc in enumerate(gen):
+        for j, Hnuc in enumerate(HHs):
             Ht = Hz + Hnuc
             L = sim.convert(Ht)
             kinetic_matrix[5:21, 5:21] -= L
diff --git a/radicalpy/simulation.py b/radicalpy/simulation.py
index b489033..f9d3bdb 100644
--- a/radicalpy/simulation.py
+++ b/radicalpy/simulation.py
@@ -1231,26 +1231,7 @@ def adjust_hamiltonian(self, H: np.ndarray):
 
 
 class SemiclassicalSimulation(LiouvilleSimulation):
-    def semiclassical_gen(
-        self,
-        num_samples: int,
-    ) -> np.ndarray:
-        num_particles = len(self.radicals)
-        spinops = [
-            [self.spin_operator(ri, ax) for ax in "xyz"] for ri in range(num_particles)
-        ]
-        result = np.zeros((num_samples, 4, 4), dtype=complex)
-        for i in range(num_samples):
-            for ri, m in enumerate(self.molecules):
-                std = m.semiclassical_std
-                Is = np.random.normal(0, std, size=1)
-                gamma = m.radical.gamma_mT
-                for ax in range(3):
-                    spinop = spinops[ri][ax]
-                    result[i, :, :] += gamma * spinop * Is
-        return result
-
-    def semiclassical_gen2(
+    def semiclassical_HHs(
         self,
         num_samples: int,
     ) -> np.ndarray:
@@ -1260,7 +1241,11 @@ def semiclassical_gen2(
 
         spinops = np.array([self.spin_operator(0, ax) for ax in "xyz"])
         cov = np.diag([m.semiclassical_std for m in self.molecules])
-        samples = np.random.multivariate_normal([0, 0], cov, size=(num_samples, 3))
+        samples = np.random.multivariate_normal(
+            mean=[0, 0],
+            cov=cov,
+            size=(num_samples, 3),
+        )
         result = np.einsum("nam,axy->nxy", samples, spinops)
         return result * self.radicals[0].gamma_mT