New CLI argument processing

Former-commit-id: d8fba975ccf08a11e0a4515e5af92edb7856f643

misc/finitesqlat-scatter.py

@@ -1,45 +1,34 @@
 #!/usr/bin/env python3
 
-import argparse
 import math
+from qpms.argproc import ArgParser
 
-ap = argparse.ArgumentParser()
+
+ap = ArgParser(['single_particle', 'single_omega', 'single_lMax'])
 ap.add_argument("-p", "--period", type=float, required=True, help='square lattice period')
 ap.add_argument("--Nx", type=int, required=True, help='Array size x')
 ap.add_argument("--Ny", type=int, required=True, help='Array size y')
-ap.add_argument("-f", "--eV", type=float, required=True, help='radiation angular frequency in eV')
-ap.add_argument("-m", "--material", help='particle material (Au, Ag for Lorentz-Drue or number for constant refractive index)', default='Au', required=True)
-ap.add_argument("-r", "--radius", type=float, required=True, help='particle radius (sphere or cylinder)')
-ap.add_argument("-H", "--height", type=float, help='cylindrical particle height; if not provided, particle is assumed to be spherical')
 ap.add_argument("-k", '--kx-lim', nargs=2, type=float, required=True, help='k vector', metavar=('KX_MIN', 'KX_MAX'))
 # ap.add_argument("--kpi", action='store_true', help="Indicates that the k vector is given in natural units instead of SI, i.e. the arguments given by -k shall be automatically multiplied by pi / period (given by -p argument)")
 ap.add_argument("--rank-tol", type=float, required=False)
-ap.add_argument("-n", "--refractive-index", type=float, default=1.52, help='background medium refractive index')
-ap.add_argument("-L", "--lMax", type=int, required=True, default=3, help='multipole degree cutoff')
-ap.add_argument("--lMax-extend", type=int, required=False, default=6, help='multipole degree cutoff for T-matrix calculation (cylindrical particles only')
 ap.add_argument("-o", "--output", type=str, required=False, help='output path (if not provided, will be generated automatically)')
 ap.add_argument("-N", type=int, default="151", help="Number of angles")
 ap.add_argument("-O", "--plot-out", type=str, required=False, help="path to plot output (optional)")
 ap.add_argument("-P", "--plot", action='store_true', help="if -p not given, plot to a default path")
+ap.add_argument("-g", "--save-gradually", action='store_true', help="saves the partial result after computing each irrep")
 
 
 a=ap.parse_args()
 
-if a.material in ['Ag', 'Au']:
+import logging
-    pass
+logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
-else:
+
-    try: lemat = float(a.material)
-    except ValueError:
-        try: lemat = complex(a.material)
-        except ValueError:
-            raise ValueError("--material must be either one of 'Ag', 'Au' or a number")
-    a.material = lemat
 
 particlestr = ("sph" if a.height is None else "cyl") + ("_r%gnm" % (a.radius*1e9))
-if a.height is not None: particlestr += "_h%gnm" % (a.height * 1e6)
+if a.height is not None: particlestr += "_h%gnm" % (a.height * 1e9)
 defaultprefix = "%s_p%gnm_%dx%d_m%s_n%g_angles(%g_%g)_Ey_f%geV_L%d_cn%d" % (
     particlestr, a.period*1e9, a.Nx, a.Ny, str(a.material), a.refractive_index, a.kx_lim[0], a.kx_lim[1], a.eV, a.lMax, a.N)
-print("Dafault file prefix: %s" % defaultprefix, flush=True)
+logging.info("Dafault file prefix: %s" % defaultprefix)
 
 
 import numpy as np
@@ -64,28 +53,17 @@ orig_y = (np.arange(a.Ny/2) + (0 if (a.Ny % 2) else .5)) * py
 
 orig_xy = np.stack(np.meshgrid(orig_x, orig_y), axis = -1)
 
-medium = EpsMu(a.refractive_index**2)
 
-if a.material in lorentz_drude:
+omega = ap.omega
-	emg = EpsMuGenerator(lorentz_drude[a.material])
-else: # constant refractive index
-        emg = EpsMuGenerator(EpsMu(a.material**2))
-
-if a.height is None:
-    tmgen = TMatrixGenerator.sphere(medium, emg, a.radius)
-else:
-    tmgen = TMatrixGenerator.cylinder(medium, emg, a.radius, a.height, lMax_extend=a.lMax_extend)
-
-omega = a.eV * eh
 
 bspec = BaseSpec(lMax = a.lMax)
-Tmatrix = tmgen(bspec, omega)
+Tmatrix = ap.tmgen(bspec, ap.omega)
 particles= [Particle(orig_xy[i], Tmatrix) for i in np.ndindex(orig_xy.shape[:-1])]
 
 sym = FinitePointGroup(point_group_info['D2h'])
 ss = ScatteringSystem(particles, sym)
 
-wavenumber = medium.k(omega).real # Currently, ScatteringSystem does not "remember" frequency nor wavenumber
+wavenumber = ap.background_epsmu.k(omega).real # Currently, ScatteringSystem does not "remember" frequency nor wavenumber
 
 sinalpha_list = np.linspace(a.kx_lim[0],a.kx_lim[1],a.N)
 
@@ -101,9 +79,16 @@ k_cart_list *= wavenumber
 σ_ext_list_ir = np.empty((a.N, ss.nirreps), dtype=float)
 σ_scat_list_ir = np.empty((a.N, ss.nirreps), dtype=float)
 
+outfile_tmp = defaultprefix + ".tmp" if a.output is None else a.output + ".tmp"
+
 for iri in range(ss.nirreps):
+    logging.info("processing irrep %d/%d" % (iri, ss.nirreps))
+    LU = None # to trigger garbage collection before the next call
+    translation_matrix = None
     LU = ss.scatter_solver(wavenumber,iri)
+    logging.info("LU solver created")
     translation_matrix = ss.translation_matrix_packed(wavenumber, iri, BesselType.REGULAR) + np.eye(ss.saecv_sizes[iri]) 
+    logging.info("auxillary translation matrix created")
 
     for j in range(a.N):
         # the following two could be calculated only once, but probably not a big deal
@@ -115,6 +100,11 @@ for iri in range(ss.nirreps):
         fi = LU(Tãi)
         σ_ext_list_ir[j, iri] = -np.vdot(ãi, fi).real/wavenumber**2
         σ_scat_list_ir[j, iri] = np.vdot(fi,np.dot(translation_matrix, fi)).real/wavenumber**2
+    if a.save_gradually:
+        iriout = outfile_tmp + ".%d" % iri
+        np.savez(iriout, iri=iri, meta=vars(a), sinalpha=sinalpha_list, k_cart = k_cart_list, E_cart=E_cart_list,
+		 omega=omega, wavenumber=wavenumber, σ_ext_list_ir=σ_ext_list_ir[:,iri], σ_scat_list_ir=σ_scat_list_ir[:,iri])
+        logging.info("partial results saved to %s"%iriout)
 
 σ_abs_list_ir = σ_ext_list_ir - σ_scat_list_ir
 σ_abs= np.sum(σ_abs_list_ir, axis=-1)
@@ -123,14 +113,17 @@ for iri in range(ss.nirreps):
 
 
 outfile = defaultprefix + ".npz" if a.output is None else a.output
-np.savez(outfile, meta=vars(a), k_cart = k_cart_list, E_cart=E_cart_list, σ_ext=σ_ext,σ_abs=σ_abs,σ_scat=σ_scat,
+np.savez(outfile, meta=vars(a), sinalpha=sinalpha_list, k_cart = k_cart_list, E_cart=E_cart_list, σ_ext=σ_ext,σ_abs=σ_abs,σ_scat=σ_scat,
  σ_ext_ir=σ_ext_list_ir,σ_abs_ir=σ_abs_list_ir,σ_scat_ir=σ_scat_list_ir, omega=omega, wavenumber=wavenumber
        )
-print("Saved to %s" % outfile)
+logging.info("Saved to %s" % outfile)
 
 
 if a.plot or (a.plot_out is not None):
+    import matplotlib
+    matplotlib.use('pdf')
     from matplotlib import pyplot as plt
+
     fig = plt.figure()
     ax = fig.add_subplot(111)
     ax.plot(sinalpha_list, σ_ext*1e12,label='$\sigma_\mathrm{ext}$')

qpms/argproc.py

@@ -0,0 +1,102 @@
+'''
+Common snippets for argument processing in command line scripts; legacy scripts use scripts_common.py instead.
+'''
+
+import argparse
+
+class ArgParser:
+    ''' Common argument parsing engine for QPMS python CLI scripts. '''
+    atomic_arguments = {
+            'sqlat_period': lambda ap: ap.add_argument("-p", "--period", type=float, required=True, help='square lattice period'),
+            'rectlat_Nx': lambda ap: ap.add_argument("--Nx", type=int, required=True, help='array size x'),
+            'rectlat_Ny': lambda ap: ap.add_argument("--Ny", type=int, required=True, help='array size y'),
+            'single_frequency_eV': lambda ap: ap.add_argument("-f", "--eV", type=float, required=True, help='radiation angular frequency in eV'),
+            'single_material': lambda ap: ap.add_argument("-m", "--material", help='particle material (Au, Ag, ... for Lorentz-Drude or number for constant refractive index)', default='Au', required=True),
+            'single_radius': lambda ap: ap.add_argument("-r", "--radius", type=float, required=True, help='particle radius (sphere or cylinder)'),
+            'single_height': lambda ap: ap.add_argument("-H", "--height", type=float, help='cylindrical particle height; if not provided, particle is assumed to be spherical'),
+            'single_kvec2': lambda ap: ap.add_argument("-k", '--kx-lim', nargs=2, type=float, required=True, help='k vector', metavar=('KX_MIN', 'KX_MAX')),
+            'kpi': lambda ap: ap.add_argument("--kpi", action='store_true', help="Indicates that the k vector is given in natural units instead of SI, i.e. the arguments given by -k shall be automatically multiplied by pi / period (given by -p argument)"),
+            'bg_refractive_index': lambda ap: ap.add_argument("-n", "--refractive-index", type=float, default=1.52, help='background medium refractive index'),
+            'single_lMax': lambda ap: ap.add_argument("-L", "--lMax", type=int, required=True, default=3, help='multipole degree cutoff'),
+            'single_lMax_extend': lambda ap: ap.add_argument("--lMax-extend", type=int, required=False, default=6, help='multipole degree cutoff for T-matrix calculation (cylindrical particles only'),
+            'outfile': lambda ap: ap.add_argument("-o", "--output", type=str, required=False, help='output path (if not provided, will be generated automatically)'),
+            'plot_out': lambda ap: ap.add_argument("-O", "--plot-out", type=str, required=False, help="path to plot output (optional)"),
+            'plot_do': lambda ap: ap.add_argument("-P", "--plot", action='store_true', help="if -p not given, plot to a default path"),
+    }
+
+    feature_sets_available = { # name : (description, dependencies, atoms not in other dependencies, methods called after parsing) 
+            'background': ("Background medium definition (currently only constant epsilon supported)", (), ('bg_refractive_index',), ('_eval_background_epsmu',)),
+            'single_particle': ("Single particle definition (shape [currently spherical or cylindrical]) and materials, incl. background)", ('background',), ('single_material', 'single_radius', 'single_height', 'single_lMax_extend'), ('_eval_single_tmgen',)),
+            'single_lMax': ("Single particle lMax definition", (), ('single_lMax',), ()),
+            'single_omega': ("Single angular frequency", (), ('single_frequency_eV',), ('_eval_single_omega',)),
+    }
+
+
+    def __init__(self, features=[]):
+        self.ap = argparse.ArgumentParser()
+        self.features_enabled = set()
+        self.call_at_parse_list = []
+        self.parsed = False
+        for feat in features:
+            self.add_feature(feat)
+
+    def add_feature(self, feat):
+        if feat not in self.features_enabled:
+            if feat not in ArgParser.feature_sets_available:
+                raise ValueError("Unknown ArgParser feature: %s", feat)
+            #resolve dependencies
+            _, deps, atoms, atparse = ArgParser.feature_sets_available[feat]
+            for dep in deps:
+                self.add_feature(dep)
+            for atom in atoms: # maybe check whether that atom has already been added sometimes in the future?
+                ArgParser.atomic_arguments[atom](self.ap)
+            for methodname in atparse:
+                self.call_at_parse_list.append(methodname)
+            self.features_enabled.add(feat)
+
+    def add_argument(self, *args, **kwargs):
+        '''Add a custom argument directly to the standard library ArgParser object'''
+        self.ap.add_argument(*args, **kwargs)
+
+    def parse_args(self, process_data = True,  *args, **kwargs):
+        self.args = self.ap.parse_args(*args, **kwargs)
+        if process_data:
+            for method in self.call_at_parse_list:
+                getattr(self, method)()
+        return self.args
+    
+    def __getattr__(self, name):
+        return getattr(self.args, name)
+        
+
+    # Methods to initialise the related data structures:
+
+    def _eval_background_epsmu(self): # feature: background
+        from .cymaterials import EpsMu
+        self.background_epsmu = EpsMu(self.args.refractive_index**2)
+
+    def _eval_single_tmgen(self): # feature: single_particle
+        a = self.args
+        from .cymaterials import EpsMuGenerator, lorentz_drude
+        from .cytmatrices import TMatrixGenerator
+        if a.material in lorentz_drude.keys():
+            self.foreground_emg = EpsMuGenerator(lorentz_drude[a.material])
+        else:
+            try: lemat = float(a.material)
+            except ValueError:
+                try: lemat = complex(a.material)
+                except ValueError as ve:
+                    raise ValueError("--material must be either a label such as 'Ag', 'Au', or a number") from ve
+            a.material = lemat
+            self.foreground_emg = EpsMuGenerator(EpsMu(a.material**2))
+
+        if a.height is None:
+            self.tmgen = TMatrixGenerator.sphere(self.background_epsmu, self.foreground_emg, a.radius)
+        else:
+            self.tmgen = TMatrixGenerator.cylinder(self.background_epsmu, self.foreground_emg, a.radius, a.height, lMax_extend = a.lMax_extend)
+
+    def _eval_single_omega(self): # feature: single_omega
+        from .constants import eV, hbar
+        self.omega = self.args.eV * eV / hbar
+
+

qpms/constants.py

@@ -1,6 +1,6 @@
 # unit conversions, mostly for standalone usage
 # TODO avoid importing the "heavy" qpms parts
-from scipy.constants import epsilon_0 as ε_0, c, pi as π, e as eV, hbar as ℏ, mu_0 as μ_0
+from scipy.constants import epsilon_0 as ε_0, c, pi as π, e as eV, hbar, hbar as ℏ, mu_0 as μ_0
 pi = π
 μm = 1e-6
 nm = 1e-9

qpms/scripts_common.py

@@ -1,3 +1,6 @@
+'''
+Mostly legacy code; new scripts use mostly argproc.py
+'''
 import warnings
 import argparse
 #import sys # for debugging purpose, TODO remove in production