From 726e5615af14a6a8c3f999b66688fb37f8fcd7ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Ne=C4=8Dada?= <marek@necada.org>
Date: Wed, 19 Sep 2018 18:58:41 +0300
Subject: [PATCH] Wdata python processing and loading

Former-commit-id: 7c93e51efc4df0410e34c1a1770192581bf17190
---
 misc/processWfiles.py |  10 ++++
 qpms/qpms_p.py        | 109 +++++++++++++++++++++++++++++++++++++++++-
 setup.py              |   2 +-
 3 files changed, 119 insertions(+), 2 deletions(-)
 create mode 100755 misc/processWfiles.py

diff --git a/misc/processWfiles.py b/misc/processWfiles.py
new file mode 100755
index 0000000..f549768
--- /dev/null
+++ b/misc/processWfiles.py
@@ -0,0 +1,10 @@
+#!/usr/bin/env python3
+
+import sys
+from qpms import processWfiles_sameKs
+
+dest = sys.argv[1]
+srcs = sys.argv[2:]
+
+processWfiles_sameKs(srcs, dest)
+
diff --git a/qpms/qpms_p.py b/qpms/qpms_p.py
index 1cbedf4..f067117 100644
--- a/qpms/qpms_p.py
+++ b/qpms/qpms_p.py
@@ -180,7 +180,8 @@ def nelem2lMax(nelem):
     else:
         return lMax
 
-
+def lMax2nelem(lMax):
+    return lMax * (lMax+2)
 
 def lpy(nmax, z):
     # TODO TRUEVECTORIZE
@@ -782,3 +783,109 @@ def loadWfile(fileName, lMax = None, fatForm = True):
             'freqs_weirdunits' : freqs_weirdunits,
             }
 
+def processWfiles_sameKs(freqfilenames, destfilename, lMax = None):
+    '''Processes translation operator data in different files; each file is supposed to contain one frequency.
+    The Ks in the different files are expected to be exactly the same and in the same order; 
+    the result is sorted by frequencies and saved to npz file;
+    The representation is always "thin", i.e. the elements which are zero due to z-symmetry are skipped'''
+    nparticles = 2 #TODO generalize
+    
+    um = 1e-6 # micrometre in SI units
+    
+    nfiles = len(freqfilenames)
+    
+    if (nfiles < 1): return ValueError("At least one filename has to be passed")
+    
+    #Check the first file to get the "constants" set
+    filename = freqfilenames[0]
+    data = np.loadtxt(filename)
+    nk_muster = data.shape[0]
+    Wdata = data[:,5:]
+    if (lMax is None):
+        nelem = int((Wdata.shape[1]/nparticles**2)**.5/2)
+        lMax = nelem2lMax(nelem)
+    else:
+        nelem = lMax2nelem(lMax)
+    ks_muster = np.array(data[:,3:5], copy=True)
+    
+    allWs = np.empty((nfiles, nk_muster, nparticles, nparticles, 2, nelem, nelem,), dtype=complex)
+    k0s = np.empty((nfiles,))
+    freqs_weirdunits = np.empty((nfiles,))
+    EeVs = np.empty((nfiles,))
+    succread = 0
+    
+    for filename in freqfilenames:
+        data = np.loadtxt(filename)
+        if data.shape[0] != nk_muster:
+            raise ValueError("%s contains different number of lines than %s"%(filename, freqfilenames[0]))
+        ks_current = np.array(data[:,3:5])
+        if not np.all(ks_current == ks_muster):
+            raise ValueError("%s contains different k-vectors than %s"%(filename, freqfilenames[0]))
+        curWdata = data[:,5:]
+        curWs = curWdata[:,::2] + 1j*Wdata[:,1::2]
+        curWs.shape = (nk_muster, nparticles, nparticles, 2, nelem, nelem)
+        allWs[succread] = curWs
+        cur_freqs_weirdunits = np.array(data[:,0])
+        if not np.all(cur_freqs_weirdunits == cur_freqs_weirdunits[0]):
+            raise ValueError("%s contains various frequencies; we require only one frequency per file"%(filename))
+        freqs_weirdunits[succread] = cur_freqs_weirdunits[0]
+        k0s[succread] = data[0,2] # TODO check this as well?
+        EeVs[succread] = data[0,1]
+        succread += 1
+    freqs = freqs_weirdunits * c / um
+
+
+    #TODO TODO TODO sort everything by omega
+    
+    np.savez(destfilename,
+             lMax = lMax,
+             nelem = nelem,
+             npart = nparticles,
+             nfreqs = succread,
+             nk = nk_muster,
+             freqs = freqs[:succread],
+             freqs_weirdunits = freqs_weirdunits[:succread],
+             EeVs_orig = EeVs[:succread],
+             k0s = k0s[:succread],
+             ks = ks_muster,
+             Wdata = allWs[:succread]
+            )
+    return
+
+
+
+def loadWfile_processed(fileName, lMax = None, fatForm = True):
+    data = np.load(fileName);
+    nk = data['nk'][()]
+    nfreqs = data['nfreqs'][()]
+    nparticles = data['npart'][()]
+    Ws = data['Wdata']
+    if lMax is not None:
+        nelem = lMax2nelem(lMax)
+        Ws = Ws[...,:nelem,:nelem]
+    else:
+        lMax = data['lMax'][()]
+        nelem = lMax2nelem(lMax)
+    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,
+        'npart': nparticles,
+        'nfreqs': nfreqs,
+        'nk' : nk,
+        'freqs': data['freqs'],
+        'freqs_weirdunits': data['freqs_weirdunits'],
+        'EeVs_orig': data['EeVs_orig'],
+        'k0s': data['k0s'],
+        'ks': data['ks'],
+        'Ws': Ws,
+    }
+    
+
diff --git a/setup.py b/setup.py
index 4ebef4d..5016770 100644
--- a/setup.py
+++ b/setup.py
@@ -35,7 +35,7 @@ qpms_c = Extension('qpms_c',
         )
 
 setup(name='qpms',
-        version = "0.2.99",
+        version = "0.2.991",
         packages=['qpms'],
 #        setup_requires=['setuptools_cython'],
         install_requires=['cython>=0.21','quaternion','spherical_functions','py_gmm'],