abstract Hamiltonian class for "Domcke" simulations

WrightSim/experiment/_scan.py

@@ -214,7 +214,7 @@ def run(self, mp="cpu", chunk=False):
                 units=units_dict[axis_obj.parameter],
             )
         self.sig.create_variable(
-            name="time", values=self.t.reshape([1] * len(self.shape) + [-1]), units="fs"
+            name="time", values=self.t[-self.iprime:].reshape([1] * len(self.shape) + [-1]), units="fs"
         )
 
         for idx, rec_idx in enumerate(self.ham.recorded_indices):

WrightSim/hamiltonian/__init__.py

@@ -1,2 +1,3 @@
 from .default import Hamiltonian
+from ._base import AbstractRKHamiltonian
 from .TRSF_default import Hamiltonian as Hamiltonian_TRSF

WrightSim/hamiltonian/_base.py

@@ -0,0 +1,90 @@
+import numpy as np
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from ..mixed import propagate
+
+
+class AbstractRKHamiltonian(ABC):
+    """boilerplate for Runge Kutta propagator
+    """
+
+    # these need to be explicitly available to the propagator
+    # each subclass needs to define these (e.g. declare in __init__)
+    recorded_indices: List[int]
+    rho: np.array
+    omega: np.array  # density matrix frequencies, in rad/fs
+    labels: List[str]
+    tau: np.array  # relaxation rates, in fs
+
+    def __init__(self):
+        self._gamma_matrix = None
+        self._omega_matrix = None
+        self.propagator = propagate.runge_kutta
+
+    def matrix(self, efields, time) -> np.ndarray:
+        """Generate the time dependent Hamiltonian Coupling Matrix
+
+        Parameters
+        ----------
+        efields : ndarray<Complex>
+            Contains the time dependent electric fields.
+            Shape (M x T) where M is number of electric fields, and T is number of timesteps.
+        time : 1-D array <float64>
+            The time step values
+        """
+        out = 0.5j * self.rabi_matrix(efields)
+        out *= np.exp(1j * self.omega_matrix * time[None, None, :])
+        out -= self.gamma_matrix
+        return out.transpose(2, 0, 1)  # move time axis first
+
+    @property
+    def gamma_matrix(self):
+        if self._gamma_matrix is None:
+            self._gamma_matrix = np.zeros((self.omega.size,)*2)
+            np.fill_diagonal(self._gamma_matrix, 1 / self.tau)
+        return self._gamma_matrix[:, :, None]
+
+    @property
+    def omega_matrix(self):
+        if self._omega_matrix is None:
+            self._omega_matrix = self.omega[:, None] - self.omega[None, :]
+        return self._omega_matrix[:, :, None]
+
+    @abstractmethod
+    def rabi_matrix(self, efields:List[np.array]):
+        """
+        define the coupling matrix here--dipoles * efields
+        return matrix of shape [to, from, time]
+
+        Example implementation
+        -----
+            E1, E2, E3 = efields
+            out = np.zeros((self.rho.size, self.rho.size, efields[0].size), dtype=complex)
+            out[to_index, from_index] = -E1 * mu_fi
+            return out
+        """
+        pass
+
+    @property
+    def attrs(self) -> dict:
+        """define quantities to send to output attrs dict"""
+        return {
+            "rho": self.rho,
+            "omega": self.omega,
+            "propagator": "runge_kutta",
+            "tau": self.tau
+        }
+
+    def matshow(self, ax, efield:float=1.0, fontsize=10):
+        """wrapper for matplotlib.pyplot.matshow to quickly show the matrix
+        """
+        mat = self.matrix([np.array([efield])] * 3, np.array([0])).squeeze()
+        art = ax.matshow(np.abs(mat))
+        ax.get_figure().colorbar(art, ax=ax, label="amplitude")
+        labels = [f"{i}: {label}" for i, label in enumerate(self.labels)]
+        ax.set_yticks(np.arange(mat.shape[0]), labels=labels, fontsize=fontsize)
+        ax.set_xticks(np.arange(mat.shape[0]), labels=labels, rotation=45, ha="left", fontsize=fontsize, rotation_mode="anchor")
+        ax.grid(c="w", ls=":")
+

WrightSim/hamiltonian/default.py

@@ -1,11 +1,13 @@
 import numpy as np
 
 from ..mixed import propagate
+from ._base import AbstractRKHamiltonian
 
 wn_to_omega = 2 * np.pi * 3 * 10 ** -5
 
 
-class Hamiltonian:
+class Hamiltonian(AbstractRKHamiltonian):
+
     cuda_struct = """
     #include <pycuda-complex.hpp>
     #define I pycuda::complex<double>(0,1)
@@ -61,8 +63,8 @@
             The lifetime of the states in femptoseconds.
             For coherences, this is the pure dephasing time.
             For populations, this is the population lifetime.
-            If tau is scalar, the same dephasing time is used for all coherences,
-                              Populations do not decay by default (inf lifetime).
+            If tau is scalar, the same dephasing time is used for all coherences.
+            Populations do not decay by default (inf lifetime).
             This value is converted into a rate of decay, Gamma, by the multiplicative inverse.
             Default is 50.0 fs dephasing for all coherences, infinite for populations.
         mu : 1-D array <Complex> (optional)
@@ -97,6 +99,7 @@
             Default is [7, 8], the output of a TRIVE Hamiltonian.
 
         """
+        super().__init__()
         if rho is None:
             self.rho = np.zeros(len(labels), dtype=np.complex128)
             self.rho[0] = 1.0
@@ -127,10 +130,9 @@
         else:
             self.omega = omega
 
-        if propagator is None:
-            self.propagator = propagate.runge_kutta
-        else:
+        if propagator is not None:
             self.propagator = propagator
+
         self.phase_cycle = phase_cycle
         self.labels = labels
         self.recorded_indices = recorded_indices
@@ -178,28 +180,7 @@
         cuda.memcpy_htod(int(pointer) + 48, np.intp(int(time_orderings)))
         cuda.memcpy_htod(int(pointer) + 56, np.intp(int(recorded_indices)))
 
-    def matrix(self, efields, time):
-        """Generate the time dependant Hamiltonian Coupling Matrix.
-
-        Parameters
-        ----------
-        efields : ndarray<Complex>
-            Contains the time dependent electric fields.
-            Shape (M x T) where M is number of electric fields, and T is number of timesteps.
-        time : 1-D array <float64>
-            The time step values
-
-        Returns
-        -------
-        ndarray <Complex> 
-            Shape T x N x N array with the full Hamiltonian at each time step.
-            N is the number of states in the Density vector.
-        """
-        # TODO: Just put the body of this method in here, rather than calling _gen_matrix
-        E1, E2, E3 = efields[0:3]
-        return self._gen_matrix(E1, E2, E3, time, self.omega)
-
-    def _gen_matrix(self, E1, E2, E3, time, energies):
+    def rabi_matrix(self, efields):
         """
         creates the coupling array given the input e-fields values for a specific time, t
 
@@ -209,57 +190,51 @@
         Matrix formulated such that dephasing/relaxation is accounted for 
         outside of the matrix
         """
-        # Define transition energies
-        wag  = energies[2]
-        w2aa = energies[-1]
+        E1, E2, E3 = efields
 
         # Define dipole moments
         mu_ag = self.mu[0]
         mu_2aa = self.mu[-1]
 
         # Define helpful variables
-        A_1 = 0.5j * mu_ag * E1 * np.exp(1j * wag * time)
-        A_2 = 0.5j * mu_ag * E2 * np.exp(-1j * wag * time)
-        A_2prime = 0.5j * mu_ag * E3 * np.exp(1j * wag * time)
-        B_1 = 0.5j * mu_2aa * E1 * np.exp(1j * w2aa * time)
-        B_2 = 0.5j * mu_2aa * E2 * np.exp(-1j * w2aa * time)
-        B_2prime = 0.5j * mu_2aa * E3 * np.exp(1j * w2aa * time)
+        A_1 = 0.5j * mu_ag * E1
+        A_2 = 0.5j * mu_ag * E2
+        A_2prime = 0.5j * mu_ag * E3
+        B_1 = 0.5j * mu_2aa * E1
+        B_2 = 0.5j * mu_2aa * E2
+        B_2prime = 0.5j * mu_2aa * E3
 
         # Initailze the full array of all hamiltonians to zero
-        out = np.zeros((len(time), len(energies), len(energies)), dtype=np.complex128)
+        out = np.zeros((self.rho.size, self.rho.size, efields[0].size), dtype=np.complex128)
 
         # Add appropriate array elements, according to the time orderings
         if 3 in self.time_orderings or 5 in self.time_orderings:
-            out[:, 1, 0] = -A_2
+            out[1, 0] = -A_2
         if 4 in self.time_orderings or 6 in self.time_orderings:
-            out[:, 2, 0] = A_2prime
+            out[2, 0] = A_2prime
         if 1 in self.time_orderings or 2 in self.time_orderings:
-            out[:, 3, 0] = A_1
+            out[3, 0] = A_1
         if 3 in self.time_orderings:
-            out[:, 5, 1] = A_1
+            out[5, 1] = A_1
         if 5 in self.time_orderings:
-            out[:, 6, 1] = A_2prime
+            out[6, 1] = A_2prime
         if 4 in self.time_orderings:
-            out[:, 4, 2] = B_1
+            out[4, 2] = B_1
         if 6 in self.time_orderings:
-            out[:, 6, 2] = -A_2
+            out[6, 2] = -A_2
         if 2 in self.time_orderings:
-            out[:, 4, 3] = B_2prime
+            out[4, 3] = B_2prime
         if 1 in self.time_orderings:
-            out[:, 5, 3] = -A_2
+            out[5, 3] = -A_2
         if 2 in self.time_orderings or 4 in self.time_orderings:
-            out[:, 7, 4] = B_2
-            out[:, 8, 4] = -A_2
+            out[7, 4] = B_2
+            out[8, 4] = -A_2
         if 1 in self.time_orderings or 3 in self.time_orderings:
-            out[:, 7, 5] = -2 * A_2prime
-            out[:, 8, 5] = B_2prime
+            out[7, 5] = -2 * A_2prime
+            out[8, 5] = B_2prime
         if 5 in self.time_orderings or 6 in self.time_orderings:
-            out[:, 7, 6] = -2 * A_1
-            out[:, 8, 6] = B_1
-
-        # Add Gamma along the diagonal
-        for i in range(len(self.Gamma)):
-            out[:, i, i] = -1 * self.Gamma[i]
+            out[7, 6] = -2 * A_1
+            out[8, 6] = B_1
 
         # NOTE: NISE multiplied outputs by the approriate mu in here
         #      This mu factors out, remember to use it where needed later

tests/mixed/default.py

@@ -17,8 +17,6 @@
 )
 ham.recorded_elements = [7, 8]
 
-
-@pytest.mark.skip("this test currently fails; bugfix needed")
 def test_windowed():
     exp = ws.experiment.builtin('trive')
     exp.w1.points = w_central  # wn
@@ -84,5 +82,5 @@ def test_frequency():
 
 
 if __name__ == "__main__":
-    test_windowed()  # fails atm
+    test_windowed()
     test_frequency()