Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Comments and docstrings #20

Merged
merged 4 commits into from
Jan 3, 2018
Merged

Comments and docstrings #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c0a17ef
Line length, comments, remove unused mu_0
ksunden Dec 29, 2017
8a5e69f
Comments and docstrings for default hamiltonian
ksunden Dec 29, 2017
8c23c24
comments in propagation code
ksunden Dec 31, 2017
4d9f505
Address comments, spelling, wording for clarity
ksunden Jan 1, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 13 additions & 5 deletions WrightSim/experiment/_scan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,10 +80,13 @@ def _gen_efp(self, indices=None):
return efp

kernel_cuda_source = """
__global__ void kernel(double time_start, double time_end, double dt, int nEFields, double* efparams, int* phase_matching, int n_recorded, Hamiltonian* ham, pycuda::complex<double>* out)
__global__ void kernel(double time_start, double time_end, double dt, int nEFields,
double* efparams, int* phase_matching, int n_recorded, Hamiltonian* ham,
pycuda::complex<double>* out)
{
int idx = threadIdx.x + blockIdx.x * blockDim.x;
runge_kutta(time_start, time_end, dt, nEFields, efparams + (idx*5*nEFields), *(efparams + 2), phase_matching, n_recorded, *ham, out + (idx*ham->nRecorded*n_recorded));
runge_kutta(time_start, time_end, dt, nEFields, efparams + (idx * 5 * nEFields),
phase_matching, n_recorded, *ham, out + (idx * ham->nRecorded * n_recorded));
}
"""

Expand All @@ -95,7 +98,7 @@ def run(self, mp='cpu', chunk=False):
mp : {False, 'cpu', 'gpu'} (optional)
Select multiprocessing: False (or '' or None) means single-threaded.
'gpu' indicates to use the CUDA implementation
Any other value which evaluates to `True` indicates cpu multiprocessed.
Any other value which evaluates to ``True`` indicates cpu multiprocessed.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is nice

with the code as it is now, I'm fairly sure that windows machines will require an if __name__ == '__main__' to prevent infinite spawning---we might want to mention this in the docstring

I think that the "if name" windows oddity can be fixed by using the new futures library, as opposed to multiprocessing, we should probably open an issue to figure that out

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am going to refrain from tackling this particular issue for the moment. We have an ongoing discussion in #18 regarding this windows forkbomb.
I'd like to test it out and decide on the desired course of action prior to doing this particular bit of documentation.
I have not attempted running on Windows, and would like to see the problem first.

Default is 'cpu'.

Returns
Expand Down Expand Up @@ -124,10 +127,15 @@ def run(self, mp='cpu', chunk=False):
start = np.min(self.efp[..., d_ind]) - self.early_buffer
stop = np.max(self.efp[..., d_ind]) + self.late_buffer

mod = SourceModule(self.ham.cuda_struct + self.ham.cuda_matrix_source + propagate.muladd_cuda_source + propagate.dot_cuda_source + propagate.pulse_cuda_source + propagate.runge_kutta_cuda_source + Scan.kernel_cuda_source)
mod = SourceModule(self.ham.cuda_struct + self.ham.cuda_matrix_source +
propagate.muladd_cuda_source + propagate.dot_cuda_source +
propagate.pulse_cuda_source + propagate.runge_kutta_cuda_source +
Scan.kernel_cuda_source)

kernel = mod.get_function('kernel')
kernel(start, stop, np.float64(self.timestep), np.intp(3), efpPtr, pmPtr, np.intp(self.iprime), hamPtr, sigPtr, grid=(self.array.size//256,1), block=(256,1,1))
kernel(start, stop, np.float64(self.timestep), np.intp(3), efpPtr,
pmPtr, np.intp(self.iprime), hamPtr, sigPtr,
grid=(self.array.size//256,1), block=(256,1,1))

cuda.memcpy_dtoh(self.sig, sigPtr)
elif mp:
Expand Down
116 changes: 111 additions & 5 deletions WrightSim/hamiltonian/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,53 @@ def __init__(self, rho=None, tau=None, mu=None,
propagator=None, phase_cycle=False,
labels=['00','01 -2','10 2\'','10 1','20 1+2\'','11 1-2','11 2\'-2', '10 1-2+2\'', '21 1-2+2\''],
time_orderings=list(range(1,7)), recorded_indices = [7, 8]):
"""Create a Hamiltonian object.

Parameters
----------
rho : 1-D array <Complex> (optional)
The initial density vector, length defines N.
Default is 1. in the ground state (index 0), and 0. everywhere else.
tau : scalar or 1-D array (optional)
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).
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)
The dipole moments used either in computing state densities or output intensities
Array like structures are converted to the proper data type.
Order matters, and meaning is dependent on the individual Hamiltonian.
Default is two values, both initially 1.0.
omega : 1-D array <float64> (optional)
The energies of various transitions.
The default uses w_central and coupling parameters to compute the appropriate
values for a TRIVE Hamiltonian
w_central : float (optional)
The cetral frequency of a resonance for a TRIVE Hamiltonian.
Used only when ``omega`` is ``None``.
coupling : float (optional)
The copuling of states for a TRIVE Hamiltonian.
Used only when ``omega`` is ``None``.
propagator : function (optional)
The evolution function to use to evolve the density matrix over time.
Default is ``runge_kutta``.
phase_cycle : bool (optional)
Whether or not to use phase cycling.
Not applicable to all Hamiltonians.
Default is ``False``
labels : list of string (optional)
Human readable labels for the states. Not used in computation, only to keep track.
time_orderings : list of int (optional)
Which time orderings to use in the simulation.
Default is all for a three interaction process: ``[1,2,3,4,5,6]``.
recorded_indices : list of int (optional)
Which density vector states to record through the simulation.
Default is [7, 8], the output of a TRIVE Hamiltonian.

"""
if rho is None:
self.rho = np.zeros(len(labels), dtype=np.complex128)
self.rho[0] = 1.
Expand All @@ -44,6 +91,7 @@ def __init__(self, rho=None, tau=None, mu=None,
else:
self.tau = tau

#TODO: Think about dictionaries or some other way of labeling Mu values
if mu is None:
self.mu = np.array([1., 1.], dtype=np.complex128)
else:
Expand All @@ -55,7 +103,7 @@ def __init__(self, rho=None, tau=None, mu=None,
w_2ag = 2*w_ag - coupling
w_gg = 0.
w_aa = w_gg
self.omega = np.array( [w_gg, -w_ag, w_ag, w_ag, w_2ag, w_aa, w_aa, w_ag, w_2aa] )
self.omega = np.array([w_gg, -w_ag, w_ag, w_ag, w_2ag, w_aa, w_aa, w_ag, w_2aa])
else:
self.omega = w_0

Expand All @@ -71,23 +119,33 @@ def __init__(self, rho=None, tau=None, mu=None,
self.Gamma = 1./self.tau

def to_device(self, pointer):
"""Transfer the Hamiltonian to a C struct in CUDA device memory.

Currently expects a pointer to an already allocated chunk of memory.
"""
import pycuda.driver as cuda
#TODO: Reorganize to allocate here and return the pointer, this is more friendly
# Transfer the arrays which make up the hamiltonian
rho = cuda.to_device(self.rho)
mu = cuda.to_device(self.mu)
omega = cuda.to_device(self.omega)
Gamma = cuda.to_device(self.Gamma)

# Convert time orderings into a C boolean array of 1 and 0, offset by one
tos = [1 if i in self.time_orderings else 0 for i in range(1,7)]

# Transfer time orderings and recorded indices
time_orderings = cuda.to_device(np.array(tos, dtype=np.int8))
recorded_indices = cuda.to_device(np.array(self.recorded_indices, dtype=np.int32))

# Transfer metadata about the lengths of feilds
cuda.memcpy_htod(pointer, np.array([len(self.rho)], dtype=np.int32))
cuda.memcpy_htod(int(pointer) + 4, np.array([len(self.mu)], dtype=np.int32))
#TODO: generalize nTimeOrderings
cuda.memcpy_htod(int(pointer) + 8, np.array([6], dtype=np.int32))
cuda.memcpy_htod(int(pointer) + 12, np.array([len(self.recorded_indices)], dtype=np.int32))

# Set pointers in the struct
cuda.memcpy_htod(int(pointer) + 16, np.intp(int(rho)))
cuda.memcpy_htod(int(pointer) + 24, np.intp(int(mu)))
cuda.memcpy_htod(int(pointer) + 32, np.intp(int(omega)))
Expand All @@ -98,35 +156,56 @@ def to_device(self, pointer):


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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

time TODO?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm going to hold off for this PR, because I want this to be as much as possible limited to non-code, i.e. comments and docstrings. This is one of the first things to be done when I touch the code again though. This is issue #26.

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):
"""
creates the coupling array given the input e-fields values for a specific time, t
w1first selects whether w1 or w2p is the first interacting positive field

Currently neglecting pathways where w2 and w3 require different frequencies
(all TRIVE space, or DOVE on diagonal)

Matrix formulated such that dephasing/relaxation is accounted for
outside of the matrix
"""
# Define transition energies
wag = energies[1]
w2aa = energies[-1]

# 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)

# Initailze the full array of all hamiltonians to zero
out = np.zeros((len(time), len(energies), len(energies)), 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
if 4 in self.time_orderings or 6 in self.time_orderings:
Expand Down Expand Up @@ -155,39 +234,65 @@ def _gen_matrix(self, E1, E2, E3, time, energies):
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]

#NOTE: NISE multiplied outputs by the approriate mu in here
# This mu factors out, remember to use it where needed later
# Removed for clarity and aligning with Equation S15 of Kohler2016
# Removed for clarity and aligning with Equation S15 of Kohler2017

return out

cuda_matrix_source = """
/**
* Hamiltonian_matrix: Computes the Hamiltonian matrix for an indevidual time step.
* NOTE: This differs from the Python implementation, which computes the full time
* dependant hamiltonian, this only computes for a single time step
* (to conserve memory).
*
* Parameters
* ----------
* Hamiltonian ham: A struct which represents a hamiltonian,
* containing orrays omega, mu, and Gamma
* cmplx* efields: A pointer to an array containg the complex valued
* electric fields to use for evaluation
* double time: the current time step counter
*
* Output
* -------
* cmplx* out: an N x N matrix containing the transition probabilities
*
*/
__device__ void Hamiltonian_matrix(Hamiltonian ham, pycuda::complex<double>* efields,
double time, pycuda::complex<double>* out)
double time, pycuda::complex<double>* out)
{
// Define state energies
double wag = ham.omega[1];
double w2aa = ham.omega[8];

// Define dipoles
//TODO: don't assume one, generalize
pycuda::complex<double> mu_ag = 1.;//ham.mu[0];
pycuda::complex<double> mu_2aa = 1.;//ham.mu[1];

// Define the electric field values
pycuda::complex<double> E1 = efields[0];
pycuda::complex<double> E2 = efields[1];
pycuda::complex<double> E3 = efields[2];


// Define helpful variables
pycuda::complex<double> A_1 = 0.5 * I * mu_ag * E1 * pycuda::exp(-1. * I * wag * time);
pycuda::complex<double> A_2 = 0.5 * I * mu_ag * E2 * pycuda::exp(I * wag * time);
pycuda::complex<double> A_2prime = 0.5 * I * mu_ag * E3 * pycuda::exp(-1. * I * wag * time);
pycuda::complex<double> B_1 = 0.5 * I * mu_2aa * E1 * pycuda::exp(-1. * I * w2aa * time);
pycuda::complex<double> B_2 = 0.5 * I * mu_2aa * E2 * pycuda::exp(I * w2aa * time);
pycuda::complex<double> B_2prime = 0.5 * I * mu_2aa * E3 * pycuda::exp(-1. * I * w2aa * time);

//TODO: zero once, take this loop out of the inner most loop
for (int i=0; i<ham.nStates * ham.nStates; i++) out[i] = pycuda::complex<double>();

// Fill in appropriate matrix elements
if(ham.time_orderings[2] || ham.time_orderings[4])
out[1*ham.nStates + 0] = -1. * A_2;
if(ham.time_orderings[3] || ham.time_orderings[5])
Expand Down Expand Up @@ -222,6 +327,7 @@ def _gen_matrix(self, E1, E2, E3, time, energies):
out[8*ham.nStates + 6] = B_1;
}

// Put Gamma along the diagonal
for(int i=0; i<ham.nStates; i++) out[i*ham.nStates + i] = -1. * ham.Gamma[i];
}
"""
Expand Down
Loading