From 6049c8425f9d57917e1dc74318474207a01b4cd3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Ne=C4=8Dada?= <marek@necada.org>
Date: Thu, 17 May 2018 13:00:44 +0300
Subject: [PATCH] Loading Ws files

Former-commit-id: 6d72c394d5b5817efc8ecbdcb5add071111bbf2d
---
 qpms/__init__.py |  1 +
 qpms/qpms_p.py   | 23 ++++++++++++++++-------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/qpms/__init__.py b/qpms/__init__.py
index fc034d4..3e52d07 100644
--- a/qpms/__init__.py
+++ b/qpms/__init__.py
@@ -5,3 +5,4 @@ from qpms_c import *
 from .qpms_p import *
 from .lattices2d import *
 from .hexpoints import *
+from .tmatrices import *
diff --git a/qpms/qpms_p.py b/qpms/qpms_p.py
index 44f5294..63d61ae 100644
--- a/qpms/qpms_p.py
+++ b/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,
             }
 
-
-
-