From 63f7f216b022a385ebcd795f8feb71a0ebc573f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Ne=C4=8Dada?= <mmn@ivor.necada.org>
Date: Mon, 3 Jan 2022 20:58:01 +0200
Subject: [PATCH] Analysis script also for finiterectlat constant driving
 results

---
 qpms/analysis.py | 39 ++++++++++++++++++++++++++-------------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/qpms/analysis.py b/qpms/analysis.py
index 2b85b37..b473f80 100644
--- a/qpms/analysis.py
+++ b/qpms/analysis.py
@@ -42,11 +42,16 @@ def _meta_matspec2emg(matspec):
         return EpsMuGenerator(lorentz_drude[matspec])
 
 
-class FiniteRectLatModesAnalysis:
-    '''Recreates the scattering system structure from finiterectlat_modes.py'''
+class FiniteRectLatAnalysis:
+    '''Recreates the scattering system structure from finiterectlat-modes.py or 
+    finiterectlat-constant-driving.py'''
     def __init__(self, data):
         meta = data['meta'][()]
-        thegroup = 'D2h' if meta['D2'] else 'D4h'
+        scatter = 'symmetry_adapted' in meta.keys()# finiterectlat-constant-driving ; TODO unify
+        if scatter: 
+            thegroup = meta['symmetry_adapted']
+        else: # finiterectlat-modes
+            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
@@ -62,8 +67,9 @@ class FiniteRectLatModesAnalysis:
             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)
+        omega = data['omega'][()] if scatter else meta['centre'] * eh 
+        self.ss, self.ssw = ScatteringSystem.create(particles, bg, omega, sym=sym)
+        ss1, ssw1 = ScatteringSystem.create([Particle([0,0,0], tmgen, bspec)], bg, omega, sym=sym)
         self.ss1, self.ssw1 = ss1, ssw1
         
         # per-particle projectors/transformation to symmetry-adapted bases
@@ -97,19 +103,26 @@ class FiniteRectLatModesAnalysis:
 
         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.fvcs1 = fvcs1
+        self.iris1 = np.array(iris1)
         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)
+        if scatter:
+            self.typ = 'scatter'
+            self.fvcs1_fromdata = data['fvcs1']
+            self.iris1_fromdata = data['iris1']
+            self.positions_fromdata = data['positions']
+        else: # modes
+            self.typ = 'modes'
+            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'][()]
 
 
     def fullvec2grid(self, fullvec, swapxy=False):