1D vectorization


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)


58

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]:
((58,), (58,), (58,), (15,))

In [30]:
a,b,c


Out[30]:
(array([  9.72222222e-01,   1.48370370e-02,  -9.28703704e-04,
          9.72222222e-01,   1.48370370e-02,  -9.28703704e-04,
          1.47925926e-02,   9.72222222e-01,   1.48370370e-02,
         -9.28703704e-04,   1.47925926e-02,   9.72222222e-01,
          1.48370370e-02,  -9.28703704e-04,  -9.23148148e-04,
          1.47925926e-02,   9.72222222e-01,   1.48370370e-02,
         -9.28703704e-04,  -9.23148148e-04,   1.47925926e-02,
          9.72222222e-01,   1.48370370e-02,  -9.28703704e-04,
         -9.23148148e-04,   1.47925926e-02,   9.72222222e-01,
          1.48370370e-02,  -9.28703704e-04,  -9.23148148e-04,
          1.47925926e-02,   9.72222222e-01,   1.48370370e-02,
         -9.28703704e-04,  -9.23148148e-04,   1.47925926e-02,
          9.72222222e-01,   1.48370370e-02,  -9.28703704e-04,
         -9.23148148e-04,   1.47925926e-02,   9.72222222e-01,
          1.48370370e-02,  -9.28703704e-04,  -9.23148148e-04,
          1.47925926e-02,   9.72222222e-01,   1.48370370e-02,
         -9.23148148e-04,   1.47925926e-02,   9.72222222e-01,
          1.48370370e-02,  -9.23148148e-04,   1.47925926e-02,
          9.72222222e-01,  -9.23148148e-04,   1.47925926e-02,
          9.72222222e-01]),
 array([ 0,  0,  0,  1,  1,  1,  2,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,
         4,  4,  5,  5,  5,  5,  5,  6,  6,  6,  6,  6,  7,  7,  7,  7,  7,
         8,  8,  8,  8,  8,  9,  9,  9,  9,  9, 10, 10, 10, 10, 11, 11, 11,
        11, 12, 12, 12, 13, 13, 13], dtype=int32),
 array([ 0,  2,  4,  1,  3,  5,  0,  2,  4,  6,  1,  3,  5,  7,  0,  2,  4,
         6,  8,  1,  3,  5,  7,  9,  2,  4,  6,  8, 10,  3,  5,  7,  9, 11,
         4,  6,  8, 10, 12,  5,  7,  9, 11, 13,  6,  8, 10, 12,  7,  9, 11,
        13,  8, 10, 12,  9, 11, 13], dtype=int32))

In [31]:
lol = np.zeros((14,14))

In [45]:
lol[b,c] = a

In [46]:
qp.printMatrix2D(lol)


          1         2         3         4         5         6         7   \
1   0.972222  0.000000  0.014837  0.000000 -0.000929  0.000000  0.000000   
2   0.000000  0.972222  0.000000  0.014837  0.000000 -0.000929  0.000000   
3   0.014793  0.000000  0.972222  0.000000  0.014837  0.000000 -0.000929   
4   0.000000  0.014793  0.000000  0.972222  0.000000  0.014837  0.000000   
5  -0.000923  0.000000  0.014793  0.000000  0.972222  0.000000  0.014837   
6   0.000000 -0.000923  0.000000  0.014793  0.000000  0.972222  0.000000   
7   0.000000  0.000000 -0.000923  0.000000  0.014793  0.000000  0.972222   
8   0.000000  0.000000  0.000000 -0.000923  0.000000  0.014793  0.000000   
9   0.000000  0.000000  0.000000  0.000000 -0.000923  0.000000  0.014793   
10  0.000000  0.000000  0.000000  0.000000  0.000000 -0.000923  0.000000   
11  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000 -0.000923   
12  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000   
13  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000   
14  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000   

          8         9         10        11        12        13        14  
1   0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  
2   0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  
3   0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  
4  -0.000929  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000  
5   0.000000 -0.000929  0.000000  0.000000  0.000000  0.000000  0.000000  
6   0.014837  0.000000 -0.000929  0.000000  0.000000  0.000000  0.000000  
7   0.000000  0.014837  0.000000 -0.000929  0.000000  0.000000  0.000000  
8   0.972222  0.000000  0.014837  0.000000 -0.000929  0.000000  0.000000  
9   0.000000  0.972222  0.000000  0.014837  0.000000 -0.000929  0.000000  
10  0.014793  0.000000  0.972222  0.000000  0.014837  0.000000 -0.000929  
11  0.000000  0.014793  0.000000  0.972222  0.000000  0.014837  0.000000  
12 -0.000923  0.000000  0.014793  0.000000  0.972222  0.000000  0.014837  
13  0.000000 -0.000923  0.000000  0.014793  0.000000  0.972222  0.000000  
14  0.000000  0.000000 -0.000923  0.000000  0.014793  0.000000  0.972222  

In [34]:
d


Out[34]:
array([ 0,  3,  6, 10, 14, 19, 24, 29, 34, 39, 44, 48, 52, 55, 58], dtype=int32)

In [35]:
x = np.ones(14, dtype=complex)
xl = 14

In [40]:
%%time
quick_multiply(a,b,d,xl,x)


CPU times: user 339 µs, sys: 130 µs, total: 469 µs
Wall time: 477 µs
Out[40]:
array([  9.86130556e-001 +3.95252517e-323j,
         9.86130556e-001 +1.97626258e-323j,
         1.00092315e+000 +6.93151265e-310j,
         1.00092315e+000 -2.10723071e-121j,
         1.00000000e+000 +0.00000000e+000j,
         8.16341772e+292 +0.00000000e+000j,
         1.00000000e+000 -2.89903644e-311j,
         1.00000000e+000 +0.00000000e+000j,
        -2.78776631e+259 +0.00000000e+000j,
         1.00000000e+000 -1.10340057e+221j,
         1.00092870e+000 +0.00000000e+000j,
         6.63084096e+040 +0.00000000e+000j,
         9.86091667e-001 -2.20458621e+231j,
         9.86091667e-001 +0.00000000e+000j])

In [42]:
%%time
lol@x


CPU times: user 53 µs, sys: 20 µs, total: 73 µs
Wall time: 80.6 µs
Out[42]:
array([ 0.98613056+0.j,  0.98613056+0.j,  1.00092315+0.j,  1.00092315+0.j,
        1.00000000+0.j,  1.00000000+0.j,  1.00000000+0.j,  1.00000000+0.j,
        1.00000000+0.j,  1.00000000+0.j,  1.00092870+0.j,  1.00092870+0.j,
        0.98609167+0.j,  0.98609167+0.j])

In [ ]: