lattices.py using now the c module

Former-commit-id: 52cecb23a8757726bcac796417ebade8242d687a

qpms/lattices.py

@@ -6,7 +6,7 @@ nx = np.newaxis
 import time
 import scipy
 import sys
-from qpms_c import get_mn_y # TODO be explicit about what is imported
+from qpms_c import get_mn_y, trans_calculator # TODO be explicit about what is imported
 from .qpms_p import cart2sph, nelem2lMax, Ã, B̃ # TODO be explicit about what is imported
 
 def _time_b(active = True, name = None, step = None):
@@ -83,6 +83,7 @@ class Scattering(object):
         self.N = positions.shape[0]
         self.k_0 = k_0
         self.lMax = lMax if lMax else nelem2lMax(TMatrices.shape[-1])
+        self.tc = trans_calculator(lMax)
         nelem = lMax * (lMax + 2) #!
         self.nelem = nelem #!
         self.prepared = False
@@ -91,7 +92,7 @@ class Scattering(object):
     def prepare(self, keep_interaction_matrix = False, verbose=False):
         btime = _time_b(verbose)
         if not self.prepared:
-            if not self.interaction_matrix:
+            if self.interaction_matrix is None:
                 self.build_interaction_matrix(verbose=verbose)
             self.lupiv = scipy.linalg.lu_factor(self.interaction_matrix,overwrite_a = not keep_interaction_matrix)
             if not keep_interaction_matrix:
@@ -106,6 +107,7 @@ class Scattering(object):
         nelem = len(my)
         leftmatrix = np.zeros((N,2,nelem,N,2,nelem), dtype=complex)
         sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients')
+        """
         for i in range(N):
             for j in range(N):
                 for yi in range(nelem):
@@ -118,6 +120,20 @@ class Scattering(object):
                             leftmatrix[j,1,yj,i,1,yi] = a
                             leftmatrix[j,0,yj,i,1,yi] = b
                             leftmatrix[j,1,yj,i,0,yi] = b
+        """
+        kdji = cart2sph(self.positions[:,nx,:] - self.positions[nx,:,:])
+        kdji[:,:,0] *= self.k_0
+        # get_AB array structure: [j,yj,i,yi]
+        a, b = self.tc.get_AB(my[nx,:,nx,nx],ny[nx,:,nx,nx],my[nx,nx,nx,:],ny[nx,nx,nx,:],
+                (kdji[:,:,0])[:,nx,:,nx], (kdji[:,:,1])[:,nx,:,nx], (kdji[:,:,2])[:,nx,:,nx],
+                False,self.J_scat)
+        mask = np.broadcast_to(np.eye(N,dtype=bool)[:,nx,:,nx],(N,nelem,N,nelem))
+        a[mask] = 0 # no self-translations
+        b[mask] = 0 
+        leftmatrix[:,0,:,:,0,:] = a
+        leftmatrix[:,1,:,:,1,:] = a
+        leftmatrix[:,0,:,:,1,:] = b
+        leftmatrix[:,1,:,:,0,:] = b
         _time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients')
         # at this point, leftmatrix is the translation matrix
         n2id = np.identity(2*nelem)
@@ -199,6 +215,7 @@ class Scattering_2D_zsym(Scattering):
         self.my, self.ny = get_mn_y(self.lMax)
         self.TE_NMz = (self.my + self.ny) % 2
         self.TM_NMz = 1 - self.TE_NMz
+        self.tc = trans_calculator(lMax)
         # TODO možnost zadávat T-matice rovnou ve zhuštěné podobě
         TMatrices_TE = TMatrices[...,self.TE_NMz[:,nx],self.TE_yz[:,nx],self.TE_NMz[nx,:],self.TE_yz[nx,:]]
         TMatrices_TM = TMatrices[...,self.TM_NMz[:,nx],self.TM_yz[:,nx],self.TM_NMz[nx,:],self.TM_yz[nx,:]]
@@ -216,13 +233,13 @@ class Scattering_2D_zsym(Scattering):
         btime = _time_b(verbose)
         if (TE_or_TM == 0): #TE
             if not self.prepared_TE:
-                if not self.interaction_matrix_TE:
+                if self.interaction_matrix_TE is None:
                     self.build_interaction_matrix(0, verbose)
                 self.lupiv_TE = scipy.linalg.lu_factor(self.interaction_matrix_TE, overwrite_a = not keep_interaction_matrix)
                 self.prepared_TE = True
         if (TE_or_TM == 1): #TM
             if not self.prepared_TM:
-                if not self.interaction_matrix_TM:
+                if self.interaction_matrix_TM is None:
                     self.build_interaction_matrix(1, verbose)
                 self.lupiv_TM = scipy.linalg.lu_factor(self.interaction_matrix_TM, overwrite_a = not keep_interaction_matrix)
                 self.prepared_TM = True
@@ -249,9 +266,29 @@ class Scattering_2D_zsym(Scattering):
             EoMl = (1,)
         elif (TE_or_TM is None):
             EoMl = (0,1)
+        sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients')
+        kdji = cart2sph(self.positions[:,nx,:] - self.positions[nx,:,:])
+        kdji[:,:,0] *= self.k_0
+        # get_AB array structure: [j,yj,i,yi]
+        # FIXME I could save some memory by calculating only half of these coefficients
+        a, b = self.tc.get_AB(my[nx,:,nx,nx],ny[nx,:,nx,nx],my[nx,nx,nx,:],ny[nx,nx,nx,:],
+                (kdji[:,:,0])[:,nx,:,nx], (kdji[:,:,1])[:,nx,:,nx], (kdji[:,:,2])[:,nx,:,nx],
+                False,self.J_scat)
+        mask = np.broadcast_to(np.eye(N,dtype=bool)[:,nx,:,nx],(N,nelem,N,nelem))
+        a[mask] = 0 # no self-translations
+        b[mask] = 0
+        print(np.isnan(np.min(a)))
+        _time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients')
         for EoM in EoMl:
             leftmatrix = np.zeros((N,nelem,N,nelem), dtype=complex)
-            sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients, T%s part' % ('M' if EoM else 'E'))
+            y = np.arange(nelem)
+            yi = y[nx,nx,nx,:]
+            yj = y[nx,:,nx,nx]
+            mask = np.broadcast_to((((yi - yj) % 2) == 0),(N,nelem,N,nelem))
+            leftmatrix[mask] = a[mask]
+            mask = np.broadcast_to((((yi - yj) % 2) != 0),(N,nelem,N,nelem))
+            leftmatrix[mask] = b[mask]
+            """ # we use to calculate the AB coefficients here
             for i in range(N):
                 for j in range(i):
                     for yi in range(nelem):
@@ -264,6 +301,7 @@ class Scattering_2D_zsym(Scattering):
                             leftmatrix[j,yj,i,yi] = tr
                             leftmatrix[i,yi,j,yj] = tr if (0 == (my[yj]+my[yi]) % 2) else -tr
             _time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients, T%s part' % ('M' if EoM else 'E'))
+            """
             for j in range(N):
                 leftmatrix[j] = - np.tensordot(self.TMatrices_TM[j] if EoM else self.TMatrices_TE[j],leftmatrix[j],
                         axes = ([-1],[0]))
@@ -273,6 +311,8 @@ class Scattering_2D_zsym(Scattering):
                 self.interaction_matrix_TE = leftmatrix
             if EoM == 1:
                 self.interaction_matrix_TM = leftmatrix
+        a = None
+        b = None
         _time_e(btime, verbose)
 
     def scatter_partial(self, TE_or_TM, pq_0, verbose = False):

qpms/qpms_c.pyx

@@ -5,6 +5,8 @@ import numpy as np
 cimport numpy as np
 cimport cython
 from cython.parallel cimport parallel, prange
+#cimport openmp
+#openmp.omp_set_dynamic(1)
 
 ## Auxillary function for retrieving the "meshgrid-like" indices; inc. nmax
 @cython.boundscheck(False)
@@ -371,10 +373,65 @@ cdef void trans_calculator_loop_E_C_DD_iiiidddii_As_lllldddbl_DD(char **args, np
 # FIXME ERROR HANDLING!!! requires correct import and different data passed (see scipy's generated ufuncs)
 #    sf_error.check_fpe(func_name)
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef void trans_calculator_parallel_loop_E_C_DD_iiiidddii_As_lllldddbl_DD(char **args, np.npy_intp *dims, np.npy_intp *steps, void *data) nogil:
+    # E stands for error value (int), C for qpms_trans_calculator*
+    cdef np.npy_intp i, n = dims[0]
+    cdef void *func = (<trans_calculator_get_X_data_t*>data)[0].cmethod
+    #cdef complex double (*func)(qpms_trans_calculator*, double complex *, double complex *, int, int, int, int, double, double, double, int, int) nogil = (<trans_calculator_get_X_data_t*>data)[0].cmethod
+    cdef qpms_trans_calculator* c = (<trans_calculator_get_X_data_t*>data)[0].c
+    #cdef char *func_name= <char*>(<void**>data)[1] # i am not using this, nor have I saved func_name to data
+
+    cdef char *ip0
+    cdef char *ip1
+    cdef char *ip2
+    cdef char *ip3
+    cdef char *ip4
+    cdef char *ip5
+    cdef char *ip6
+    cdef char *ip7
+    cdef char *ip8
+    cdef char *op0
+    cdef char *op1
+    cdef int errval
+    for i in prange(n): # iterating over dimensions
+        ip0 = args[0] + i * steps[0]
+        ip1 = args[1] + i * steps[1]
+        ip2 = args[2] + i * steps[2]
+        ip3 = args[3] + i * steps[3]
+        ip4 = args[4] + i * steps[4]
+        ip5 = args[5] + i * steps[5]
+        ip6 = args[6] + i * steps[6]
+        ip7 = args[7] + i * steps[7]
+        ip8 = args[8] + i * steps[8]
+        op0 = args[9] + i * steps[9]
+        op1 = args[10] + i * steps[10]
+        #errval = func( 
+        errval = (<int(*)(qpms_trans_calculator*, double complex *, double complex *, int, int, int, int, double, double, double, int, int) nogil>func)(
+            c,
+            <cdouble *> op0,
+            <cdouble *> op1,
+            <int>(<np.npy_long*>ip0)[0],
+            <int>(<np.npy_long*>ip1)[0],
+            <int>(<np.npy_long*>ip2)[0],
+            <int>(<np.npy_long*>ip3)[0],
+            <double>(<np.npy_double*>ip4)[0],
+            <double>(<np.npy_double*>ip5)[0],
+            <double>(<np.npy_double*>ip6)[0],
+            <int>(<np.npy_bool*>ip7)[0],
+            <int>(<np.npy_long*>ip8)[0],
+        )
+       # TODO if (errval != 0): ...
+# FIXME ERROR HANDLING!!! requires correct import and different data passed (see scipy's generated ufuncs)
+#    sf_error.check_fpe(func_name)
+
+
 cdef np.PyUFuncGenericFunction trans_calculator_get_X_loop_funcs[1]
 trans_calculator_get_X_loop_funcs[0] = trans_calculator_loop_D_Ciiiidddii_As_D_lllldddbl
 
 cdef np.PyUFuncGenericFunction trans_calculator_get_AB_loop_funcs[1]
+#trans_calculator_get_AB_loop_funcs[0] = trans_calculator_parallel_loop_E_C_DD_iiiidddii_As_lllldddbl_DD
 trans_calculator_get_AB_loop_funcs[0] = trans_calculator_loop_E_C_DD_iiiidddii_As_lllldddbl_DD
 cdef void *trans_calculator_get_AB_elementwise_funcs[1]
 trans_calculator_get_AB_elementwise_funcs[0] = <void *>qpms_trans_calculator_get_AB_p_ext

qpms/translations.c

@@ -419,6 +419,10 @@ complex double qpms_trans_calculator_get_A_buf(const qpms_trans_calculator *c,
 		bool r_ge_d, qpms_bessel_t J,
 		complex double *bessel_buf, double *legendre_buf) {
 	if (r_ge_d) J = QPMS_BESSEL_REGULAR;
+	if (0 == kdlj.r && J != QPMS_BESSEL_REGULAR) 
+		// TODO warn? 
+		return NAN+I*NAN;
+	
 	switch(c->normalization) {
 		case QPMS_NORMALIZATION_TAYLOR:
 			{
@@ -451,6 +455,9 @@ complex double qpms_trans_calculator_get_B_buf(const qpms_trans_calculator *c,
 		bool r_ge_d, qpms_bessel_t J,
 		complex double *bessel_buf, double *legendre_buf) {
 	if (r_ge_d) J = QPMS_BESSEL_REGULAR;
+	if (0 == kdlj.r && J != QPMS_BESSEL_REGULAR) 
+		// TODO warn? 
+		return NAN+I*NAN;
 	switch(c->normalization) {
 		case QPMS_NORMALIZATION_TAYLOR:
 			{
@@ -486,6 +493,12 @@ int qpms_trans_calculator_get_AB_buf_p(const qpms_trans_calculator *c,
 		bool r_ge_d, qpms_bessel_t J,
 		complex double *bessel_buf, double *legendre_buf) {
 	if (r_ge_d) J = QPMS_BESSEL_REGULAR;
+	if (0 == kdlj.r && J != QPMS_BESSEL_REGULAR) {
+		*Adest = NAN+I*NAN;
+		*Bdest = NAN+I*NAN;
+		// TODO warn? different return value?
+		return 0;
+	}
 	switch(c->normalization) {
 		case QPMS_NORMALIZATION_TAYLOR:
 			{

setup.py

@@ -23,7 +23,7 @@ qpms_c = Extension('qpms_c',
         extra_compile_args=['-std=c99','-ggdb','-O3',
             '-DDISABLE_NDEBUG', # uncomment to enable assertions in the modules
             ],
-        libraries=['gsl', 'blas'],
+        libraries=['gsl', 'blas', 'omp'],
         runtime_library_dirs=os.environ['LD_LIBRARY_PATH'].split(':')
         )