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 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:

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_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