diff --git a/WrightSim/experiment/_scan.py b/WrightSim/experiment/_scan.py index c50e9e4..32fd4d4 100644 --- a/WrightSim/experiment/_scan.py +++ b/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): diff --git a/WrightSim/hamiltonian/__init__.py b/WrightSim/hamiltonian/__init__.py index 35823af..3f01407 100644 --- a/WrightSim/hamiltonian/__init__.py +++ b/WrightSim/hamiltonian/__init__.py @@ -1,2 +1,3 @@ from .default import Hamiltonian +from ._base import AbstractRKHamiltonian from .TRSF_default import Hamiltonian as Hamiltonian_TRSF diff --git a/WrightSim/hamiltonian/_base.py b/WrightSim/hamiltonian/_base.py new file mode 100644 index 0000000..99415a0 --- /dev/null +++ b/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 + 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 + 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=":") + diff --git a/WrightSim/hamiltonian/default.py b/WrightSim/hamiltonian/default.py index 2abc7a9..8ac3e51 100644 --- a/WrightSim/hamiltonian/default.py +++ b/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 #define I pycuda::complex(0,1) @@ -61,8 +63,8 @@ def __init__( 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 (optional) @@ -97,6 +99,7 @@ def __init__( 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 @@ def __init__( 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 @@ def to_device(self, pointer): 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 - 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 - The time step values - - Returns - ------- - ndarray - 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 @@ def _gen_matrix(self, E1, E2, E3, time, energies): 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 diff --git a/tests/mixed/default.py b/tests/mixed/default.py index 819015e..84513e3 100644 --- a/tests/mixed/default.py +++ b/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()