Loading Ws files

Former-commit-id: 6d72c394d5b5817efc8ecbdcb5add071111bbf2d

qpms/__init__.py

@@ -5,3 +5,4 @@ from qpms_c import *
 from .qpms_p import *
 from .lattices2d import *
 from .hexpoints import *
+from .tmatrices import *

qpms/qpms_p.py

@@ -740,9 +740,11 @@ def G0_sum_1_slow(source_cart, dest_cart, k, nmax):
     return G_Mie_scat_precalc_cart(source_cart, dest_cart, RH, RV, a=0.001, nmax=nmax, k_i=1, k_e=k, μ_i=1, μ_e=1, J_ext=1, J_scat=3)
 
 
-def loadWfile(fileName, lMax = None):
+def loadWfile(fileName, lMax = None, fatForm = True):
     '''
     Load the translation operator lattice sums generated by hexlattice_ewald.c
+    fatForm has different indices (so that it can be easily multiplied with T-Matrix
+    and twice the size.
     '''
     nparticles = 2 # TODO generalize
     um = 1e-6 # micrometer in SI units
@@ -758,11 +760,21 @@ def loadWfile(fileName, lMax = None):
     # indices: destparticle, srcparticle, wavetype, desty, srcy
     Ws.shape = (nsamples, nparticles, nparticles, 2, nelem, nelem) 
     # maybe TODO copy the following so the original data gets freed to save memory
-    freqs_weirdunits = data[:,0]
-    k0s = data[:,2]
-    ks = data[:,3:4]
+    freqs_weirdunits = np.array(data[:,0], copy=True)
+    k0s = np.array(data[:,2], copy=True)
+    ks = np.array(data[:,3:4], copy=True)
     freqs = freqs_weirdunits * c / um
+    if fatForm: #indices: (...,) destparticle, desttype, desty, srcparticle, srctype, srcy
+        Ws2 = np.moveaxis(Ws, [-5,-4,-3,-2,-1], [-4,-2,-5,-3,-1] )
+        fatWs = np.empty(Ws2.shape[:-5]+(nparticles, 2,nelem, nparticles, 2, nelem),dtype=complex)
+        fatWs[...,:,0,:,:,0,:] = Ws2[...,0,:,:,:,:] #A
+        fatWs[...,:,1,:,:,1,:] = Ws2[...,0,:,:,:,:] #A
+        fatWs[...,:,1,:,:,0,:] = Ws2[...,1,:,:,:,:] #B
+        fatWs[...,:,0,:,:,1,:] = Ws2[...,1,:,:,:,:] #B
+        Ws = fatWs
     return {
+            'lMax' : lMax,
+            'nelem' : nelem,
             'freqs' : freqs,
             'ks' : ks,
             'k0_effs' : k0s,
@@ -770,6 +782,3 @@ def loadWfile(fileName, lMax = None):
             'freqs_weirdunits' : freqs_weirdunits,
             }
 
-
-
-