In [1]:
import quantumpropagator as qp
%load_ext Cython
Vm = np.ones((7,2))
Km = np.ones((7,9,3))*10e-8
theL = 7
nstates = 2
xl = theL*nstates
dthe = 0.003
In [48]:
%%cython
import numpy as np
cimport numpy as np
def wrapper_Checkaddress1D_prime(Vm,Km,theL,nstates,dthe):
a,b,c,d,k = Checkaddress1D(Vm,Km,theL,nstates,dthe,100000)
return (k)
def wrapper_Checkaddress1D(Vm,Km,theL,nstates,dthe, k):
a,b,c,d,k = Checkaddress1D(Vm,Km,theL,nstates,dthe, k)
return (np.asarray(a), np.asarray(b), np.asarray(c), np.asarray(d))
def transform(t,s):
nstates = 2
if t < 0 or t > 7-1:
#print(t,s)
res = -1
else:
res = t*nstates+s
return res
def quick_multiply(a,b,d,xl,x):
'''
a = values
b = columns
d = rows counting
i = rows
'''
I = -1j
mom = np.empty_like(x)
for i in range(xl):
for j in range(d[i],d[i+1]):
mom[i] = mom[i] + a[j] * x[b[j]]
return mom
cdef Checkaddress1D(double [:,:] Vm, double [:,:,:] Km, int theL, int nstates, double dthe, int kL):
cdef:
int s, t, k, add1
double value1
double [:] values_big
int [:] i_big, j_big, d_big
values_big = np.zeros(kL) # real
i_big = np.zeros(kL, dtype=np.intc) #integers
j_big = np.zeros(kL, dtype=np.intc) #integers
d_big = np.zeros(theL*nstates+1, dtype=np.intc) #integers
k = 0
d_big[0] = k
for t in range(theL):
for s in range(nstates):
i_add = transform(t,s)
# derivatives in theta
# dG_dt = ((1.0/12)*GRID[t-2,s]+(-2.0/3)*GRID[t-1,s]+(2.0/3)*GRID[t+1,s]+(-1.0/12)*GRID[t+2,s]) / dthe
# d2G_dt2 = (-GRID[t+2,s]+16*GRID[t+1,s]-30*GRID[t,s]+16*GRID[t-1,s]-GRID[t-2,s]) / (12 * dthe**2)
value1 = ((1.0/12) * Km[t,8,1] / dthe) + (-(1.0/12) * Km[t,8,2] / (dthe**2))
add1 = transform(t-2,s)
if add1 >= 0:
values_big[k] = value1
i_big[k] = i_add
j_big[k] = add1
k = k + 1
value1 = (-2.0/3) * Km[t,8,1] / dthe + (+(16.0/12) * Km[t,8,2] / (dthe**2))
add1 = transform(t-1,s)
if add1 >= 0:
values_big[k] = value1
i_big[k] = i_add
j_big[k] = add1
k = k + 1
value1 = (-(30.0/12) * Km[t,8,2] / (dthe**2)) + Vm[t,s]
add1 = transform(t,s)
if add1 >= 0:
values_big[k] = value1
i_big[k] = i_add
j_big[k] = add1
k = k + 1
value1 = (2.0/3) * Km[t,8,1] / dthe + (+(16.0/12) * Km[t,8,2] / (dthe**2))
add1 = transform(t+1,s)
if add1 >= 0:
values_big[k] = value1
i_big[k] = i_add
j_big[k] = add1
k = k + 1
value1 = (-(1.0/12) * Km[t,8,1] / dthe) + (-(1.0/12) * Km[t,8,2] / (dthe**2))
add1 = transform(t+2,s)
if add1 >= 0:
values_big[k] = value1
i_big[k] = i_add
j_big[k] = add1
k = k + 1
# Mtot = 0
# Ntot = 0
# for d in range(nstates): # state s is where the outer loop is, d is where the inner loop is.
# for carte in range(3): # carte is 'cartesian', meaning 0,1,2 -> x,y,z
# Mtot = Mtot - ((pulseV[carte] * Dm[t,carte,s,d] ) * GRID[t,d])
# # NAC calculation
# if s < d:
# dG_dt_oth = ((1.0/12)*GRID[t-2,d] + (-2.0/3)*GRID[t-1,d] + (2.0/3)*GRID[t+1,d] + (-1.0/12)*GRID[t+2,d])*Nm[t,s,d,2] / dthe
# elif s > d:
# dG_dt_oth = ((1.0/12)*GRID[t-2,d]*Nm[t-2,s,d,2] + (-2.0/3)*GRID[t-1,d]*Nm[t-1,s,d,2] + (2.0/3)*GRID[t+1,d]*Nm[t+1,s,d,2] + (-1.0/12)*GRID[t+2,d]*Nm[t+2,s,d,2]) / dthe
# else:
# dG_dt_oth = 0
# Ntot = Ntot - dG_dt_oth
d_big[i_add+1] = k
return(values_big, i_big, j_big, d_big, k)
In [27]:
k = wrapper_Checkaddress1D_prime(Vm,Km,theL,nstates,dthe)
print(k)
In [28]:
a,b,c,d = wrapper_Checkaddress1D(Vm,Km,theL,nstates,dthe,k)
In [29]:
a.shape, b.shape, c.shape, d.shape
Out[29]:
In [30]:
a,b,c
Out[30]:
In [31]:
lol = np.zeros((14,14))
In [45]:
lol[b,c] = a
In [46]:
qp.printMatrix2D(lol)
In [34]:
d
Out[34]:
In [35]:
x = np.ones(14, dtype=complex)
xl = 14
In [40]:
%%time
quick_multiply(a,b,d,xl,x)
Out[40]:
In [42]:
%%time
lol@x
Out[42]:
In [ ]: