argproc.py: enable manual specification of BaseSpec

qpms/argproc.py

@@ -5,6 +5,7 @@ Common snippets for argument processing in command line scripts.
 import argparse
 import sys
 import warnings
+import ast
 
 def flatten(S):
     if S == []:
@@ -192,6 +193,13 @@ def material_spec(string):
                 raise argparse.ArgumentTypeError("Material specification must be a supported material name %s, or a number" % (str(lorentz_drude.keys()),)) from ve
         return lemat
 
+def string2bspec(string):
+    """
+    Converts string representation of list to BaseSpec.
+    """
+    from .cybspec import BaseSpec
+    return BaseSpec(ast.literal_eval(string))
+
 class ArgParser:
     ''' Common argument parsing engine for QPMS python CLI scripts. '''
     
@@ -248,6 +256,13 @@ class ArgParser:
         mpgrp.add_argument("+H", "++height", nargs=2,
                 action=make_dict_action(argtype=float, postaction='store', first_is_key=True,),
                 help='particle radius (cylinder; labeled)')
+        # Alternatively, add a constant T-matrix
+        mpgrp.add_argmuent("-w", "--vswf-set", nargs=1, default={},
+                action=make_dict_action(argtype=string2basespec, postaction='store', first_is_key=False),
+                help='Manual specification of VSWF set codes (format as a python list of integers); see docs on qpms_uvswfi_t for valid codes or simply use --lMax instead. Overrides --lMax.')
+        mpgrp.add_argmuent("+w", "++vswf-set", nargs=2, default={},
+                action=make_dict_action(argtype=string2basespec, postaction='store', first_is_key=True),
+                help='Manual specification of VSWF set codes (format as a python list of integers); see docs on qpms_uvswfi_t for valid codes or simply use ++lMax instead. Overrides ++lMax and --lMax.')
 
     atomic_arguments = {
             'rectlattice2d_periods': lambda ap: ap.add_argument("-p", "--period", type=float, nargs='+', required=True, help='square/rectangular lattice periods', metavar=('px','[py]')),
@@ -337,6 +352,14 @@ class ArgParser:
             return tmgen
 
     def _add_bspec(self, key):
+        """
+        Transforms a given key into a BaseSpec and registers
+        the BaseSpec with the key.
+        Always returns a BaseSpec (unless an exception occurs).
+        If an equivalent instance of BaseSpec is registered, returns that one
+        (therefore, all equivalent BaseSpecs in the register values should
+        be the same instance).
+        """
         if key in self._bspec_register.keys():
             return self._bspec_register[key]
         else:
@@ -345,7 +368,9 @@ class ArgParser:
                 bspec = key
             elif isinstance(key, int):
                 bspec = self._add_bspec(BaseSpec(lMax=key))
-            else: raise TypeError("Can't register this as a BaseSpec")
+            #else: raise TypeError("Can't register this as a BaseSpec")
+            else: 
+                bspec = self._add_bspec(BaseSpec(key))
             self._bspec_register[key] = bspec
             return bspec
 
@@ -488,8 +513,12 @@ class ArgParser:
             warnings.warn("No particle position (-p or +p) specified, assuming single particle in the origin / single particle per unit cell!")
             a.position[None] = [(0.,0.,0.)]
         for poslabel in a.position.keys():
+            # TODO HERE GOES THE CODE TRYING TO LOAD CONSTANT T-MATRIX
             try:
-                lMax = a.lMax.get(poslabel, False) or a.lMax[None]
+                lMax_or_bspec = ( a.vswf_set.get(poslabel, False) 
+                        or a.lMax.get(poslabel, False)
+                        or a.vswf_set.get(None, False)
+                        or a.lMax[None] )
                 radius = a.radius.get(poslabel, False) or a.radius[None]
                 # Height is "inherited" only together with radius
                 height = a.height.get(poslabel, None) if poslabel in a.radius.keys() else a.height.get(None, None)
@@ -508,7 +537,7 @@ class ArgParser:
             tmspec = (a.background, material, (radius, height, lMax_extend)) 
             self.tmspecs[poslabel] = tmspec
             self.tmgens[poslabel] = self._add_tmg(tmspec)
-            self.bspecs[poslabel] = self._add_bspec(lMax)
+            self.bspecs[poslabel] = self._add_bspec(lMax_or_bspec)
             poslist_cured = []
             for pos in a.position[poslabel]:
                 if len(pos) == 1:

qpms/qpms_cdefs.pxd

@@ -195,6 +195,8 @@ cdef extern from "indexing.h":
     qpms_errno_t qpms_uvswfi2tmn(qpms_uvswfi_t u, qpms_vswf_type_t* t, qpms_m_t* m, qpms_l_t* n) nogil
     qpms_m_t qpms_uvswfi2m(qpms_uvswfi_t u) nogil
     qpms_y_t qpms_lMax2nelem_sc(qpms_l_t lmax) nogil
+    void qpms_y2mn_p(qpms_y_t y, qpms_m_t *m, qpms_l_t *n)
+    qpms_l_t qpms_nelem2lMax(qpms_y_t nelem)
     # maybe more if needed
 
 # Point generators from lattices.h