From dd0ee9e2d4a9bd12cea38eb1f0225181fda91f7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Ne=C4=8Dada?= <marek@necada.org>
Date: Thu, 9 Dec 2021 11:19:19 +0200
Subject: [PATCH] Helper class to analyse results of finiterectlat-modes.py

---
 qpms/analysis.py | 132 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 132 insertions(+)
 create mode 100644 qpms/analysis.py

diff --git a/qpms/analysis.py b/qpms/analysis.py
new file mode 100644
index 0000000..4d79e4b
--- /dev/null
+++ b/qpms/analysis.py
@@ -0,0 +1,132 @@
+'''
+Auxiliary functions to analyse results obtained using scripts in
+the ../misc directory.
+
+This needs to be kept in sync with the scripts and qpms/argproc.py.
+'''
+
+import numpy as np
+import qpms
+from qpms.qpms_c import Particle
+from qpms.cytmatrices import CTMatrix, TMatrixGenerator
+from qpms.cybspec import BaseSpec
+from qpms.cymaterials import EpsMu, EpsMuGenerator, LorentzDrudeModel, lorentz_drude
+from qpms.symmetries import point_group_info
+from qpms import FinitePointGroup, ScatteringSystem, eV, hbar
+eh = eV / hbar
+
+def cleanarray(a, atol=1e-8, copy=True, fullNaN=True):
+    a = np.array(a, copy=copy)
+    sieve = abs(a.real) < atol
+    a[sieve] = 1j * a[sieve].imag
+    sieve = abs(a.imag) < atol
+    a[sieve] = a[sieve].real
+    if fullNaN:
+        if np.all(a == 0):
+            a[...] = 0
+    return a
+
+def nicerot(a, atol=1e-10, copy=True): 
+    '''
+    Gives an array a "nice" phase.
+    '''
+    a = np.array(a, copy=copy)
+    i = np.argmax(abs(a))
+    a = a / a[i] * abs(a[i])
+    return a
+
+def _meta_matspec2emg(matspec):
+    if isinstance(matspec, (float, complex)): # number is interpreted as refractive index
+        return EpsMuGenerator(EpsMu(matspec**2)) 
+    else:
+        return EpsMuGenerator(lorentz_drude[matspec])
+
+
+class FiniteRectLatModesAnalysis:
+    '''Recreates the scattering system structure from finiterectlat_modes.py'''
+    def __init__(self, data):
+        meta = data['meta'][()]
+        thegroup = 'D2h' if meta['D2'] else 'D4h'
+        Nx, Ny = meta['size']
+        px, py = meta['period']
+        orig_x = (np.arange(Nx/2) + (0 if (Nx % 2) else .5)) * px
+        orig_y = (np.arange(Ny/2) + (0 if (Ny % 2) else .5)) * py
+        orig_xy = np.stack(np.meshgrid(orig_x, orig_y), axis = -1)
+        bspec = BaseSpec(lMax = meta['lMax'])
+        nelem = len(bspec)
+        bg = _meta_matspec2emg(meta['background'])
+        fg = _meta_matspec2emg(meta['material'])
+        if meta.get('height', None) is None:
+            tmgen = TMatrixGenerator.sphere(bg, fg, meta['radius'])
+        else:
+            tmgen = TMatrixGenerator.cylinder(bg, fg, meta['radius'], meta['height'], lMax_extend=meta['lMax_extend'])
+        particles = [Particle(orig_xy[i], tmgen, bspec) for i in np.ndindex(orig_xy.shape[:-1])]
+        sym = FinitePointGroup(point_group_info[thegroup])
+        self.ss, self.ssw = ScatteringSystem.create(particles, bg, meta['centre'] * eh, sym=sym)
+        ss1, ssw1 = ScatteringSystem.create([Particle([0,0,0], tmgen, bspec)], bg, meta['centre']*eh, sym=sym)
+        self.ss1, self.ssw1 = ss1, ssw1
+        
+        # per-particle projectors/transformation to symmetry-adapted bases
+        fvcs1 = np.empty((nelem, nelem), dtype=complex)
+        y = 0
+        iris1 = []
+        for iri1 in range(ss1.nirreps):
+            for j in range(ss1.saecv_sizes[iri1]):
+                pvc1 = np.zeros((ss1.saecv_sizes[iri1],), dtype=complex)
+                pvc1[j] = 1
+                fvcs1[y] = ss1.unpack_vector(pvc1, iri1)
+                fvcs1[y] = cleanarray(nicerot(fvcs1[y], copy=False), copy=False)
+                y += 1
+                iris1.append(iri1)
+        
+        # Mapping between ss particles and grid positions; 
+        positions = self.ss.positions.astype(np.float32) # cast to eliminate rounding errors
+        xpositions = np.unique(positions[:,0]) 
+        assert(len(xpositions) == Nx)
+        ypositions = np.unique(positions[:,1])
+        assert(len(ypositions == Ny))
+        # particle positions as integer indices
+        posmap = np.empty((positions.shape[0],2), dtype=int)
+        #invposmap = np.empty((Nx, Ny), dtype=int)
+        for i, pos in enumerate(positions):
+            posmap[i,0] = np.searchsorted(xpositions, positions[i,0])
+            posmap[i,1] = np.searchsorted(ypositions, positions[i,1])
+            #invposmap[posmap[i,0], posmap[i, 1]] = i
+
+        self.meta = meta
+        self.npzfile = data
+        self.eigval = data['eigval']
+        self.neig = len(self.eigval)
+        self.eigvec = data['eigvec']
+        self.residuals = data['residuals']
+        self.ranktest_SV = data['ranktest_SV']
+        self.iri = data['iri'][()]
+        self.bspec = bspec
+        self.nelem = nelem
+        self.Nx, self.Ny, self.px, self.py = Nx, Ny, px, py
+        self.posmap = posmap
+        self.fullvec_poffsets = nelem * np.arange(Nx*Ny)
+        self.fvcs1 = fvcs1
+        self.iris1 = np.array(iris1)
+
+
+    def fullvec2grid(self, fullvec, swapxy=False):
+        arr = np.empty((self.Nx, self.Ny, self.nelem), dtype=complex)
+        for pi, offset in enumerate(self.fullvec_poffsets):
+            ix, iy = self.posmap[pi]
+            arr[ix, iy] = fullvec[offset:offset+self.nelem]
+        return np.swapaxes(arr, 0, 1) if swapxy else arr
+    
+    def saecv2grid(self, saecv, swapxy=False):
+        fullvec = self.ss.unpack_vector(saecv, iri=self.iri)
+        return self.fullvec2grid(fullvec, swapxy)
+    
+    def eigvec_asgrid(self, i, swapxy = False):
+        '''
+        Returns i-th eigenmode as a (Nx, Ny, nelem)-shaped grid of excitation
+        coefficients.
+        '''
+        if self.iri is not None:
+            return self.saecv2grid(self.eigvec[i], swapxy)
+        else:
+            return self.fullvec2grid(self.eigvec[i], swapxy)