argproc.py: enable manual specification of BaseSpec

This commit is contained in:
Marek Nečada 2022-05-27 11:21:59 +03:00
parent 4c0453a865
commit cd1406d23f
2 changed files with 34 additions and 3 deletions

View File

@ -5,6 +5,7 @@ Common snippets for argument processing in command line scripts.
import argparse import argparse
import sys import sys
import warnings import warnings
import ast
def flatten(S): def flatten(S):
if 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 raise argparse.ArgumentTypeError("Material specification must be a supported material name %s, or a number" % (str(lorentz_drude.keys()),)) from ve
return lemat return lemat
def string2bspec(string):
"""
Converts string representation of list to BaseSpec.
"""
from .cybspec import BaseSpec
return BaseSpec(ast.literal_eval(string))
class ArgParser: class ArgParser:
''' Common argument parsing engine for QPMS python CLI scripts. ''' ''' Common argument parsing engine for QPMS python CLI scripts. '''
@ -248,6 +256,13 @@ class ArgParser:
mpgrp.add_argument("+H", "++height", nargs=2, mpgrp.add_argument("+H", "++height", nargs=2,
action=make_dict_action(argtype=float, postaction='store', first_is_key=True,), action=make_dict_action(argtype=float, postaction='store', first_is_key=True,),
help='particle radius (cylinder; labeled)') 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 = { atomic_arguments = {
'rectlattice2d_periods': lambda ap: ap.add_argument("-p", "--period", type=float, nargs='+', required=True, help='square/rectangular lattice periods', metavar=('px','[py]')), '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 return tmgen
def _add_bspec(self, key): 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(): if key in self._bspec_register.keys():
return self._bspec_register[key] return self._bspec_register[key]
else: else:
@ -345,7 +368,9 @@ class ArgParser:
bspec = key bspec = key
elif isinstance(key, int): elif isinstance(key, int):
bspec = self._add_bspec(BaseSpec(lMax=key)) 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 self._bspec_register[key] = bspec
return 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!") 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.)] a.position[None] = [(0.,0.,0.)]
for poslabel in a.position.keys(): for poslabel in a.position.keys():
# TODO HERE GOES THE CODE TRYING TO LOAD CONSTANT T-MATRIX
try: 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] radius = a.radius.get(poslabel, False) or a.radius[None]
# Height is "inherited" only together with radius # 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) 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)) tmspec = (a.background, material, (radius, height, lMax_extend))
self.tmspecs[poslabel] = tmspec self.tmspecs[poslabel] = tmspec
self.tmgens[poslabel] = self._add_tmg(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 = [] poslist_cured = []
for pos in a.position[poslabel]: for pos in a.position[poslabel]:
if len(pos) == 1: if len(pos) == 1:

View File

@ -198,6 +198,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_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_m_t qpms_uvswfi2m(qpms_uvswfi_t u) nogil
qpms_y_t qpms_lMax2nelem_sc(qpms_l_t lmax) 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 # maybe more if needed
# Point generators from lattices.h # Point generators from lattices.h