1-
21import numpy as np
32
43from ..mixed import propagate
54
5+ wn_to_omega = 2 * np .pi * 3 * 10 ** - 5
6+
67
78class Hamiltonian :
89 cuda_struct = """
@@ -23,14 +24,32 @@ class Hamiltonian:
2324 int* recorded_indices;
2425 };
2526 """
26- cuda_mem_size = 4 * 4 + np .intp (0 ).nbytes * 6
27-
28-
29- def __init__ (self , rho = None , tau = None , mu = None ,
30- omega = None , w_central = 7000. , coupling = 0 ,
31- propagator = None , phase_cycle = False ,
32- 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\' ' ],
33- time_orderings = list (range (1 ,7 )), recorded_indices = [7 , 8 ]):
27+ cuda_mem_size = 4 * 4 + np .intp (0 ).nbytes * 6
28+
29+ def __init__ (
30+ self ,
31+ rho = None ,
32+ tau = None ,
33+ mu = None ,
34+ omega = None ,
35+ w_central = 7000.0 ,
36+ coupling = 0 ,
37+ propagator = None ,
38+ phase_cycle = False ,
39+ labels = [
40+ "00" ,
41+ "01 -2" ,
42+ "10 2'" ,
43+ "10 1" ,
44+ "20 1+2'" ,
45+ "11 1-2" ,
46+ "11 2'-2" ,
47+ "10 1-2+2'" ,
48+ "21 1-2+2'" ,
49+ ],
50+ time_orderings = list (range (1 , 7 )),
51+ recorded_indices = [7 , 8 ],
52+ ):
3453 """Create a Hamiltonian object.
3554
3655 Parameters
@@ -80,32 +99,33 @@ def __init__(self, rho=None, tau=None, mu=None,
8099 """
81100 if rho is None :
82101 self .rho = np .zeros (len (labels ), dtype = np .complex128 )
83- self .rho [0 ] = 1.
102+ self .rho [0 ] = 1.0
84103 else :
85104 self .rho = np .array (rho , dtype = np .complex128 )
86105
87106 if tau is None :
88- tau = 50. # fs
107+ tau = 50.0 # fs
89108 if np .isscalar (tau ):
90109 self .tau = np .array ([np .inf , tau , tau , tau , tau , np .inf , np .inf , tau , tau ])
91110 else :
92111 self .tau = tau
93112
94- #TODO: Think about dictionaries or some other way of labeling Mu values
113+ # TODO: Think about dictionaries or some other way of labeling Mu values
95114 if mu is None :
96- self .mu = np .array ([1. , 1. ], dtype = np .complex128 )
115+ self .mu = np .array ([1.0 , 1.0 ], dtype = np .complex128 )
97116 else :
98117 self .mu = np .array (mu , dtype = np .complex128 )
99118
100119 if omega is None :
101120 w_ag = w_central
102121 w_2aa = w_ag - coupling
103- w_2ag = 2 * w_ag - coupling
104- w_gg = 0.
122+ w_2ag = 2 * w_ag - coupling
123+ w_gg = 0.0
105124 w_aa = w_gg
106125 self .omega = np .array ([w_gg , - w_ag , w_ag , w_ag , w_2ag , w_aa , w_aa , w_ag , w_2aa ])
126+ self .omega *= wn_to_omega
107127 else :
108- self .omega = w_0
128+ self .omega = omega
109129
110130 if propagator is None :
111131 self .propagator = propagate .runge_kutta
@@ -116,23 +136,28 @@ def __init__(self, rho=None, tau=None, mu=None,
116136 self .recorded_indices = recorded_indices
117137
118138 self .time_orderings = time_orderings
119- self .Gamma = 1. / self .tau
139+ self .Gamma = 1.0 / self .tau
140+
141+ @property
142+ def omega_wn (self ):
143+ return self .omega / wn_to_omega
120144
121145 def to_device (self , pointer ):
122146 """Transfer the Hamiltonian to a C struct in CUDA device memory.
123147
124148 Currently expects a pointer to an already allocated chunk of memory.
125149 """
126150 import pycuda .driver as cuda
127- #TODO: Reorganize to allocate here and return the pointer, this is more friendly
151+
152+ # TODO: Reorganize to allocate here and return the pointer, this is more friendly
128153 # Transfer the arrays which make up the hamiltonian
129154 rho = cuda .to_device (self .rho )
130155 mu = cuda .to_device (self .mu )
131156 omega = cuda .to_device (self .omega )
132157 Gamma = cuda .to_device (self .Gamma )
133158
134159 # Convert time orderings into a C boolean array of 1 and 0, offset by one
135- tos = [1 if i in self .time_orderings else 0 for i in range (1 ,7 )]
160+ tos = [1 if i in self .time_orderings else 0 for i in range (1 , 7 )]
136161
137162 # Transfer time orderings and recorded indices
138163 time_orderings = cuda .to_device (np .array (tos , dtype = np .int8 ))
@@ -141,7 +166,7 @@ def to_device(self, pointer):
141166 # Transfer metadata about the lengths of feilds
142167 cuda .memcpy_htod (pointer , np .array ([len (self .rho )], dtype = np .int32 ))
143168 cuda .memcpy_htod (int (pointer ) + 4 , np .array ([len (self .mu )], dtype = np .int32 ))
144- #TODO: generalize nTimeOrderings
169+ # TODO: generalize nTimeOrderings
145170 cuda .memcpy_htod (int (pointer ) + 8 , np .array ([6 ], dtype = np .int32 ))
146171 cuda .memcpy_htod (int (pointer ) + 12 , np .array ([len (self .recorded_indices )], dtype = np .int32 ))
147172
@@ -153,8 +178,6 @@ def to_device(self, pointer):
153178 cuda .memcpy_htod (int (pointer ) + 48 , np .intp (int (time_orderings )))
154179 cuda .memcpy_htod (int (pointer ) + 56 , np .intp (int (recorded_indices )))
155180
156-
157-
158181 def matrix (self , efields , time ):
159182 """Generate the time dependant Hamiltonian Coupling Matrix.
160183
@@ -172,9 +195,9 @@ def matrix(self, efields, time):
172195 Shape T x N x N array with the full Hamiltonian at each time step.
173196 N is the number of states in the Density vector.
174197 """
175- #TODO: Just put the body of this method in here, rather than calling _gen_matrix
176- E1 ,E2 ,E3 = efields [0 :3 ]
177- return self ._gen_matrix (E1 , E2 , E3 , time , self .omega )
198+ # TODO: Just put the body of this method in here, rather than calling _gen_matrix
199+ E1 , E2 , E3 = efields [0 :3 ]
200+ return self ._gen_matrix (E1 , E2 , E3 , time , self .omega )
178201
179202 def _gen_matrix (self , E1 , E2 , E3 , time , energies ):
180203 """
@@ -187,13 +210,13 @@ def _gen_matrix(self, E1, E2, E3, time, energies):
187210 outside of the matrix
188211 """
189212 # Define transition energies
190- wag = energies [1 ]
213+ wag = energies [1 ]
191214 w2aa = energies [- 1 ]
192-
215+
193216 # Define dipole moments
194217 mu_ag = self .mu [0 ]
195218 mu_2aa = self .mu [- 1 ]
196-
219+
197220 # Define helpful variables
198221 A_1 = 0.5j * mu_ag * E1 * np .exp (- 1j * wag * time )
199222 A_2 = 0.5j * mu_ag * E2 * np .exp (1j * wag * time )
@@ -207,46 +230,46 @@ def _gen_matrix(self, E1, E2, E3, time, energies):
207230
208231 # Add appropriate array elements, according to the time orderings
209232 if 3 in self .time_orderings or 5 in self .time_orderings :
210- out [:,1 , 0 ] = - A_2
233+ out [:, 1 , 0 ] = - A_2
211234 if 4 in self .time_orderings or 6 in self .time_orderings :
212- out [:,2 , 0 ] = A_2prime
235+ out [:, 2 , 0 ] = A_2prime
213236 if 1 in self .time_orderings or 2 in self .time_orderings :
214- out [:,3 , 0 ] = A_1
237+ out [:, 3 , 0 ] = A_1
215238 if 3 in self .time_orderings :
216- out [:,5 , 1 ] = A_1
239+ out [:, 5 , 1 ] = A_1
217240 if 5 in self .time_orderings :
218- out [:,6 , 1 ] = A_2prime
241+ out [:, 6 , 1 ] = A_2prime
219242 if 4 in self .time_orderings :
220- out [:,4 , 2 ] = B_1
243+ out [:, 4 , 2 ] = B_1
221244 if 6 in self .time_orderings :
222- out [:,6 , 2 ] = - A_2
245+ out [:, 6 , 2 ] = - A_2
223246 if 2 in self .time_orderings :
224- out [:,4 , 3 ] = B_2prime
247+ out [:, 4 , 3 ] = B_2prime
225248 if 1 in self .time_orderings :
226- out [:,5 , 3 ] = - A_2
249+ out [:, 5 , 3 ] = - A_2
227250 if 2 in self .time_orderings or 4 in self .time_orderings :
228- out [:,7 , 4 ] = B_2
229- out [:,8 , 4 ] = - A_2
251+ out [:, 7 , 4 ] = B_2
252+ out [:, 8 , 4 ] = - A_2
230253 if 1 in self .time_orderings or 3 in self .time_orderings :
231- out [:,7 , 5 ] = - 2 * A_2prime
232- out [:,8 , 5 ] = B_2prime
254+ out [:, 7 , 5 ] = - 2 * A_2prime
255+ out [:, 8 , 5 ] = B_2prime
233256 if 5 in self .time_orderings or 6 in self .time_orderings :
234- out [:,7 , 6 ] = - 2 * A_1
235- out [:,8 , 6 ] = B_1
257+ out [:, 7 , 6 ] = - 2 * A_1
258+ out [:, 8 , 6 ] = B_1
236259
237260 # Add Gamma along the diagonal
238261 for i in range (len (self .Gamma )):
239- out [:,i , i ] = - 1 * self .Gamma [i ]
262+ out [:, i , i ] = - 1 * self .Gamma [i ]
240263
241- #NOTE: NISE multiplied outputs by the approriate mu in here
264+ # NOTE: NISE multiplied outputs by the approriate mu in here
242265 # This mu factors out, remember to use it where needed later
243266 # Removed for clarity and aligning with Equation S15 of Kohler2017
244267
245268 return out
246269
247270 cuda_matrix_source = """
248271 /**
249- * Hamiltonian_matrix: Computes the Hamiltonian matrix for an indevidual time step.
272+ * Hamiltonian_matrix: Computes the Hamiltonian matrix for an individual time step.
250273 * NOTE: This differs from the Python implementation, which computes the full time
251274 * dependant hamiltonian, this only computes for a single time step
252275 * (to conserve memory).
@@ -330,5 +353,4 @@ def _gen_matrix(self, E1, E2, E3, time, energies):
330353 // Put Gamma along the diagonal
331354 for(int i=0; i<ham.nStates; i++) out[i*ham.nStates + i] = -1. * ham.Gamma[i];
332355 }
333- """
334-
356+ """
0 commit comments