562 lines
29 KiB
Python
562 lines
29 KiB
Python
'''
|
||
Common snippets for argument processing in command line scripts.
|
||
'''
|
||
|
||
import argparse
|
||
import sys
|
||
import warnings
|
||
|
||
def flatten(S):
|
||
if S == []:
|
||
return S
|
||
if isinstance(S[0], list):
|
||
return flatten(S[0]) + flatten(S[1:])
|
||
return S[:1] + flatten(S[1:])
|
||
|
||
def make_action_sharedlist(opname, listname):
|
||
class opAction(argparse.Action):
|
||
def __call__(self, parser, args, values, option_string=None):
|
||
if (not hasattr(args, listname)) or getattr(args, listname) is None:
|
||
setattr(args, listname, list())
|
||
getattr(args, listname).append((opname, values))
|
||
return opAction
|
||
|
||
def make_dict_action(argtype=None, postaction='store', first_is_key=True):
|
||
class DictAction(argparse.Action):
|
||
#def __init__(self, option_strings, dest, nargs=None, **kwargs):
|
||
# if nargs is not None:
|
||
# raise ValueError("nargs not allowed")
|
||
# super(DictAction, self).__init__(option_strings, dest, **kwargs)
|
||
def __call__(self, parser, namespace, values, option_string=None):
|
||
if first_is_key: # For the labeled versions
|
||
key = values[0]
|
||
vals = values[1:]
|
||
else: # For the default values
|
||
key = None
|
||
vals = values
|
||
if argtype is not None:
|
||
if (first_is_key and self.nargs == 2) or (not first_is_key and self.nargs == 1):
|
||
vals = argtype(vals[0]) # avoid having lists in this case
|
||
else:
|
||
vals = [argtype(val) for val in vals]
|
||
ledict = getattr(namespace, self.dest, {})
|
||
if ledict is None:
|
||
ledict = {}
|
||
if postaction=='store':
|
||
ledict[key] = vals
|
||
elif postaction=='append':
|
||
lelist = ledict.get(key, [])
|
||
lelist.append(vals)
|
||
ledict[key] = lelist
|
||
setattr(namespace, self.dest, ledict)
|
||
return DictAction
|
||
|
||
|
||
class ArgumentProcessingError(Exception):
|
||
pass
|
||
|
||
class AppendTupleAction(argparse.Action):
|
||
''' A variation on the 'append' builtin action from argparse, but uses tuples for the internal groupings instead of lists '''
|
||
def __call__(self, parser, args, values, option_string=None):
|
||
if (not hasattr(args, self.dest)) or getattr(args, self.dest) is None:
|
||
setattr(args, self.dest, list())
|
||
getattr(args, self.dest).append(tuple(values))
|
||
|
||
def float_range(string):
|
||
"""Tries to parse a string either as one individual float value
|
||
or one of the following patterns:
|
||
|
||
first:last:increment
|
||
first:last|steps
|
||
first:last
|
||
|
||
(The last one is equivalent to first:last|50.)
|
||
Returns either float or numpy array.
|
||
"""
|
||
try:
|
||
res = float(string)
|
||
return res
|
||
except ValueError:
|
||
import re
|
||
steps = None
|
||
match = re.match(r's?([^:]+):([^|]+)\|(.+)', string)
|
||
if match:
|
||
steps = int(match.group(3))
|
||
else:
|
||
match = re.match(r's?([^:]+):([^:]+):(.+)', string)
|
||
if match:
|
||
increment = float(match.group(3))
|
||
else:
|
||
match = re.match(r's?([^:]+):(.+)', string)
|
||
if match:
|
||
steps = 50
|
||
else:
|
||
argparse.ArgumentTypeError('Invalid float/sequence format: "%s"' % string)
|
||
first = float(match.group(1))
|
||
last = float(match.group(2))
|
||
import numpy as np
|
||
if steps is not None:
|
||
return np.linspace(first, last, num=steps)
|
||
else:
|
||
return np.arange(first, last, increment)
|
||
|
||
def int_or_None(string):
|
||
"""Tries to parse a string either as an int or None (if it contains only whitespaces)"""
|
||
try:
|
||
return int(string)
|
||
except ValueError as ve:
|
||
if string.strip() == '':
|
||
return None
|
||
else:
|
||
raise ve
|
||
|
||
def sslice(string):
|
||
"""Tries to parse a string either as one individual int value
|
||
or one of the following patterns:
|
||
|
||
first:last:increment
|
||
first:last
|
||
|
||
first, last and increment must be parseable as ints
|
||
or be empty (then
|
||
|
||
In each case, 's' letter can be prepended to the whole string to avoid
|
||
argparse interpreting this as a new option (if the argument contains
|
||
'-' or '+').
|
||
|
||
Returns either int or slice containing ints or Nones.
|
||
"""
|
||
if string[0] == 's':
|
||
string = string[1:]
|
||
try:
|
||
res = int(string)
|
||
return res
|
||
except ValueError:
|
||
import re
|
||
match = re.match(r'([^:]*):([^:]*):(.*)', string)
|
||
if match:
|
||
step = int_or_None(match.group(3))
|
||
else:
|
||
match = re.match(r'([^:]*):(.*)', string)
|
||
if match:
|
||
step = None
|
||
else:
|
||
argparse.ArgumentTypeError('Invalid int/slice format: "%s"' % string)
|
||
start = int_or_None(match.group(1))
|
||
stop = int_or_None(match.group(2))
|
||
return slice(start, stop, step)
|
||
|
||
def sfloat(string):
|
||
'''Tries to match a float, or a float with prepended 's'
|
||
|
||
Used as a workaraound for argparse's negative number matcher, which does not recognize
|
||
scientific notation.
|
||
'''
|
||
try:
|
||
res = float(string)
|
||
except ValueError as exc:
|
||
if string[0] == 's':
|
||
res = float(string[1:])
|
||
else: raise exc
|
||
return res
|
||
|
||
def sint(string):
|
||
'''Tries to match an int, or an int with prepended 's'
|
||
|
||
Used as a workaraound for argparse's negative number matcher if '+' is used as a
|
||
prefix
|
||
'''
|
||
try:
|
||
res = int(string)
|
||
except ValueError as exc:
|
||
if string[0] == 's':
|
||
res = int(string[1:])
|
||
else: raise exc
|
||
return res
|
||
|
||
def material_spec(string):
|
||
"""Tries to parse a string as a material specification, i.e. a
|
||
real or complex number or one of the string in built-in Lorentz-Drude models.
|
||
|
||
Tries to interpret the string as 1) float, 2) complex, 3) Lorentz-Drude key.
|
||
Raises argparse.ArgumentTypeError on failure.
|
||
"""
|
||
from .cymaterials import lorentz_drude
|
||
if string in lorentz_drude.keys():
|
||
return string
|
||
else:
|
||
try: lemat = float(string)
|
||
except ValueError:
|
||
try: lemat = complex(string)
|
||
except ValueError as ve:
|
||
raise argparse.ArgumentTypeError("Material specification must be a supported material name %s, or a number" % (str(lorentz_drude.keys()),)) from ve
|
||
return lemat
|
||
|
||
class ArgParser:
|
||
''' Common argument parsing engine for QPMS python CLI scripts. '''
|
||
|
||
def __add_planewave_argparse_group(ap):
|
||
pwgrp = ap.add_argument_group('Incident wave specification', """
|
||
Incident wave direction is given in terms of ISO polar and azimuthal angles θ, φ,
|
||
which translate into cartesian coordinates as r̂ = (x, y, z) = (sin(θ) cos(φ), sin(θ) sin(φ), cos(θ)).
|
||
|
||
Wave polarisation is given in terms of parameters ψ, χ, where ψ is the angle between a polarisation
|
||
ellipse axis and meridian tangent θ̂, and tg χ determines axes ratio;
|
||
the electric field in the origin is then
|
||
|
||
E⃗ = cos(χ) (cos(ψ) θ̂ + sin(ψ) φ̂) + i sin(χ) (sin(ψ) θ̂ + cos(ψ) φ̂).
|
||
|
||
All the angles are given as multiples of π/2.
|
||
""" # TODO EXAMPLES
|
||
)
|
||
pwgrp.add_argument("-φ", "--phi", type=float, default=0,
|
||
help='Incident wave asimuth in multiples of π/2.')
|
||
pwgrp.add_argument("-θ", "--theta", type=float_range, default=0,
|
||
help='Incident wave polar angle in multiples of π/2. This might be a sequence in format FIRST:LAST:INCREMENT.')
|
||
pwgrp.add_argument("-ψ", "--psi", type=float, default=0,
|
||
help='Angle between polarisation ellipse axis and meridian tangent θ̂ in multiples of π/2.')
|
||
pwgrp.add_argument("-χ", "--chi", type=float, default=0,
|
||
help='Polarisation parameter χ in multiples of π/2. 0 for linear, 0.5 for circular pol.')
|
||
|
||
def __add_manyparticle_argparse_group(ap):
|
||
mpgrp = ap.add_argument_group('Many particle specification', "TODO DOC")
|
||
mpgrp.add_argument("-p", "--position", nargs='+', action=make_dict_action(argtype=sfloat, postaction='append',
|
||
first_is_key=False), help="Particle positions, cartesion coordinates (default particle properties)")
|
||
mpgrp.add_argument("+p", "++position", nargs='+', action=make_dict_action(argtype=sfloat, postaction='append',
|
||
first_is_key=True), help="Particle positions, cartesian coordinates (labeled)")
|
||
mpgrp.add_argument("-L", "--lMax", nargs=1, default={},
|
||
action=make_dict_action(argtype=int, postaction='store', first_is_key=False,),
|
||
help="Cutoff multipole degree (default)")
|
||
mpgrp.add_argument("+L", "++lMax", nargs=2,
|
||
action=make_dict_action(argtype=int, postaction='store', first_is_key=True,),
|
||
help="Cutoff multipole degree (labeled)")
|
||
mpgrp.add_argument("-m", "--material", nargs=1, default={},
|
||
action=make_dict_action(argtype=material_spec, postaction='store', first_is_key=False,),
|
||
help='particle material (Au, Ag, ... for Lorentz-Drude or number for constant refractive index) (default)')
|
||
mpgrp.add_argument("+m", "++material", nargs=2,
|
||
action=make_dict_action(argtype=material_spec, postaction='store', first_is_key=True,),
|
||
help='particle material (Au, Ag, ... for Lorentz-Drude or number for constant refractive index) (labeled)')
|
||
mpgrp.add_argument("-r", "--radius", nargs=1, default={},
|
||
action=make_dict_action(argtype=float, postaction='store', first_is_key=False,),
|
||
help='particle radius (sphere or cylinder; default)')
|
||
mpgrp.add_argument("+r", "++radius", nargs=2,
|
||
action=make_dict_action(argtype=float, postaction='store', first_is_key=True,),
|
||
help='particle radius (sphere or cylinder; labeled)')
|
||
mpgrp.add_argument("-H", "--height", nargs=1, default={},
|
||
action=make_dict_action(argtype=float, postaction='store', first_is_key=False,),
|
||
help='particle radius (cylinder; default)')
|
||
mpgrp.add_argument("+H", "++height", nargs=2,
|
||
action=make_dict_action(argtype=float, postaction='store', first_is_key=True,),
|
||
help='particle radius (cylinder; labeled)')
|
||
|
||
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_counts': lambda ap: ap.add_argument("--size", type=int, nargs=2, required=True, help='rectangular array size (particle column, row count)', metavar=('NCOLS', 'NROWS')),
|
||
'single_frequency_eV': lambda ap: ap.add_argument("-f", "--eV", type=float, required=True, help='radiation angular frequency in eV'),
|
||
'multiple_frequency_eV_optional': lambda ap: ap.add_argument("-f", "--eV", type=float, nargs='*', help='radiation angular frequency in eV (additional)'),
|
||
'seq_frequency_eV': lambda ap: ap.add_argument("-F", "--eV-seq", type=float, nargs=3, required=True, help='uniform radiation angular frequency sequence in eV', metavar=('FIRST', 'INCREMENT', 'LAST')),
|
||
'real_frequencies_eV_ng': lambda ap: ap.add_argument("-f", "--eV", type=float_range, nargs=1, action='append', required=True, help='Angular frequency (or angular frequency range) in eV'), # nargs='+', action='extend' would be better, but action='extend' requires python>=3.8
|
||
'single_material': lambda ap: ap.add_argument("-m", "--material", help='particle material (Au, Ag, ... for Lorentz-Drude or number for constant refractive index)', type=material_spec, 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=sfloat, 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_real_refractive_index': lambda ap: ap.add_argument("-n", "--refractive-index", type=float, default=1., help='background medium strictly real refractive index'),
|
||
'bg_analytical': lambda ap: ap.add_argument("-B", "--background", type=material_spec, default=1., help="Background medium specification (constant real or complex refractive index, or supported material label)"),
|
||
'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)'), # TODO consider type=argparse.FileType('w')
|
||
'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"),
|
||
'lattice2d_basis': lambda ap: ap.add_argument("-b", "--basis-vector", nargs='+', action=AppendTupleAction, help="basis vector in xy-cartesian coordinates (two required)", required=True, type=sfloat, dest='basis_vectors', metavar=('X', 'Y')),
|
||
'planewave_pol_angles': __add_planewave_argparse_group,
|
||
'multi_particle': __add_manyparticle_argparse_group,
|
||
}
|
||
|
||
feature_sets_available = { # name : (description, dependencies, atoms not in other dependencies, methods called after parsing, "virtual" features provided)
|
||
'const_real_background': ("Background medium with constant real refractive index", (), ('bg_real_refractive_index',), ('_eval_const_background_epsmu',), ('background', 'background_analytical')),
|
||
'background' : ("Most general background medium specification currently supported", ('background_analytical',), (), (), ()),
|
||
'background_analytical' : ("Background medium model holomorphic for 'reasonably large' complex frequency areas", (), ('bg_analytical',), ('_eval_analytical_background_epsmugen',), ('background',)),
|
||
'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',), ()),
|
||
'multi_particle': ("One or more particle definition (shape [curently spherical or cylindrical]), materials, and positions)", ('background',), ('multi_particle',), ('_process_multi_particle',), ()),
|
||
'single_lMax': ("Single particle lMax definition", (), ('single_lMax',), (), ()),
|
||
'single_omega': ("Single angular frequency", (), ('single_frequency_eV',), ('_eval_single_omega',), ()),
|
||
'omega_seq': ("Equidistant real frequency range with possibility of adding individual frequencies", (), ('seq_frequency_eV', 'multiple_frequency_eV_optional',), ('_eval_omega_seq',), ()),
|
||
'omega_seq_real_ng': ("Equidistant real frequency ranges or individual frequencies (new syntax)", (), ('real_frequencies_eV_ng',), ('_eval_omega_seq_real_ng',), ()),
|
||
'lattice2d': ("Specification of a generic 2d lattice (spanned by the x,y axes)", (), ('lattice2d_basis',), ('_eval_lattice2d',), ()),
|
||
'rectlattice2d': ("Specification of a rectangular 2d lattice; conflicts with lattice2d", (), ('rectlattice2d_periods',), ('_eval_rectlattice2d',), ()),
|
||
'rectlattice2d_finite': ("Specification of a rectangular 2d lattice; conflicts with lattice2d", ('rectlattice2d',), ('rectlattice2d_counts',), (), ()),
|
||
'planewave': ("Specification of a normalised plane wave (typically used for scattering) with a full polarisation state", (), ('planewave_pol_angles',), ("_process_planewave_angles",), ()),
|
||
}
|
||
|
||
|
||
def __init__(self, features=[]):
|
||
prefix_chars = '+-' if 'multi_particle' in features else '-'
|
||
self.ap = argparse.ArgumentParser(prefix_chars=prefix_chars)
|
||
self.features_enabled = set()
|
||
self.call_at_parse_list = []
|
||
self.parsed = False
|
||
for feat in features:
|
||
self.add_feature(feat)
|
||
self._emg_register = {} # EpsMuGenerator dictionary to avoid recreating equivalent instances; filled by _add_emg()
|
||
self._tmg_register = {} # TMatrixGenerator dictionary to avoid recreating equivalent instances; filled by _add_tmg()
|
||
self._bspec_register = {} # Dictionary of used BaseSpecs to keep the equivalent instances unique; filled by _add_bspec()
|
||
|
||
def _add_emg(self, emgspec):
|
||
"""Looks up whether if an EpsMuGenerator from given material_spec has been already registered, and if not, creates a new one"""
|
||
from .cymaterials import EpsMu, EpsMuGenerator, lorentz_drude
|
||
if emgspec in self._emg_register.keys():
|
||
return self._emg_register[emgspec]
|
||
else:
|
||
if isinstance(emgspec, (float, complex)):
|
||
emg = EpsMuGenerator(EpsMu(emgspec**2))
|
||
else:
|
||
emg = EpsMuGenerator(lorentz_drude[emgspec])
|
||
self._emg_register[emgspec] = emg
|
||
return emg
|
||
|
||
def _add_tmg(self, tmgspec):
|
||
"""Looks up whether if a T-matrix from given T-matrix specification tuple has been already registered, and if not, creates a new one
|
||
|
||
T-matrix specification shall be of the form
|
||
(bg_material_spec, fg_material_spec, shape_spec) where shape_spec is
|
||
(radius, height, lMax_extend)
|
||
"""
|
||
if tmgspec in self._tmg_register.keys():
|
||
return self._tmg_register[tmgspec]
|
||
else:
|
||
from .cytmatrices import TMatrixGenerator
|
||
bgspec, fgspec, (radius, height, lMax_extend) = tmgspec
|
||
bg = self._add_emg(bgspec)
|
||
fg = self._add_emg(fgspec)
|
||
if height is None:
|
||
tmgen = TMatrixGenerator.sphere(bg, fg, radius)
|
||
else:
|
||
tmgen = TMatrixGenerator.cylinder(bg, fg, radius, height, lMax_extend=lMax_extend)
|
||
self._tmg_register[tmgspec] = tmgen
|
||
return tmgen
|
||
|
||
def _add_bspec(self, key):
|
||
if key in self._bspec_register.keys():
|
||
return self._bspec_register[key]
|
||
else:
|
||
from .cybspec import BaseSpec
|
||
if isinstance(key, BaseSpec):
|
||
bspec = key
|
||
elif isinstance(key, int):
|
||
bspec = self._add_bspec(BaseSpec(lMax=key))
|
||
else: raise TypeError("Can't register this as a BaseSpec")
|
||
self._bspec_register[key] = bspec
|
||
return bspec
|
||
|
||
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, provides_virtual = 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)
|
||
for feat_virt in provides_virtual:
|
||
self.features_enabled.add(feat_virt)
|
||
|
||
def add_argument(self, *args, **kwargs):
|
||
'''Add a custom argument directly to the standard library ArgParser object'''
|
||
return self.ap.add_argument(*args, **kwargs)
|
||
|
||
def add_argument_group(self, *args, **kwargs):
|
||
'''Add a custom argument group directly to the standard library ArgParser object'''
|
||
return self.ap.add_argument_group(*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:
|
||
try:
|
||
getattr(self, method)()
|
||
except ArgumentProcessingError:
|
||
err = sys.exc_info()[1]
|
||
self.ap.error(str(err))
|
||
return self.args
|
||
|
||
def __getattr__(self, name):
|
||
return getattr(self.args, name)
|
||
|
||
|
||
# Methods to initialise the related data structures:
|
||
|
||
def _eval_const_background_epsmu(self): # feature: const_real_background
|
||
self.args.background = self.args.refractive_index
|
||
self._eval_analytical_background_epsmugen()
|
||
|
||
def _eval_analytical_background_epsmugen(self): # feature: background_analytical
|
||
a = self.args
|
||
from .cymaterials import EpsMu
|
||
if isinstance(a.background, (float, complex)):
|
||
self.background_epsmu = EpsMu(a.background**2)
|
||
self.background_emg = self._add_emg(a.background)
|
||
|
||
def _eval_single_tmgen(self): # feature: single_particle
|
||
a = self.args
|
||
from .cymaterials import EpsMuGenerator, lorentz_drude
|
||
from .cytmatrices import TMatrixGenerator
|
||
self.foreground_emg = self._add_emg(a.material)
|
||
self.tmgen = self._add_tmg((a.background, a.material, (a.radius, a.height, a.lMax_extend)))
|
||
self.bspec = self._add_bspec(a.lMax)
|
||
|
||
def _eval_single_omega(self): # feature: single_omega
|
||
from .constants import eV, hbar
|
||
self.omega = self.args.eV * eV / hbar
|
||
|
||
def _eval_omega_seq(self): # feature: omega_seq
|
||
import numpy as np
|
||
from .constants import eV, hbar
|
||
start, step, stop = self.args.eV_seq
|
||
self.omegas = np.arange(start, stop, step)
|
||
if self.args.eV:
|
||
self.omegas = np.concatenate((self.omegas, np.array(self.args.eV)))
|
||
self.omegas.sort()
|
||
self.omegas *= eV/hbar
|
||
|
||
def _eval_omega_seq_real_ng(self): # feature: omega_seq_real_ng
|
||
import numpy as np
|
||
from .constants import eV, hbar
|
||
eh = eV / hbar
|
||
self.omegas = [omega_eV * eh for omega_eV in flatten(self.args.eV)]
|
||
self.omega_max = max(om if isinstance(om, float) else max(om) for om in self.omegas)
|
||
self.omega_min = min(om if isinstance(om, float) else min(om) for om in self.omegas)
|
||
self.omega_singles = [om for om in self.omegas if isinstance(om, float)]
|
||
self.omega_ranges = [om for om in self.omegas if not isinstance(om, float)]
|
||
self.omega_descr = ("%geV" % (self.omega_max / eh)) if (self.omega_max == self.omega_min) else (
|
||
"%g–%geV" % (self.omega_min / eh, self.omega_max / eh))
|
||
self.allomegas = []
|
||
for om in self.omegas:
|
||
if isinstance(om, float):
|
||
self.allomegas.append(om)
|
||
else:
|
||
self.allomegas.extend(om)
|
||
self.allomegas = np.unique(self.allomegas)
|
||
|
||
|
||
def _eval_lattice2d(self): # feature: lattice2d
|
||
l = len(self.args.basis_vectors)
|
||
if l != 2: raise ValueError('Two basis vectors must be specified (have %d)' % l)
|
||
from .qpms_c import lll_reduce
|
||
self.direct_basis = lll_reduce(self.args.basis_vectors, delta=1.)
|
||
import numpy as np
|
||
self.reciprocal_basis1 = np.linalg.inv(self.direct_basis.T)
|
||
self.reciprocal_basis2pi = 2 * np.pi * self.reciprocal_basis1
|
||
|
||
def _eval_rectlattice2d(self): # feature: rectlattice2d
|
||
a = self.args
|
||
l = len(a.period)
|
||
if (l == 1): # square lattice
|
||
a.period = (a.period[0], a.period[0])
|
||
else:
|
||
a.period = (a.period[0], a.period[1])
|
||
if (l > 2):
|
||
raise ValueError("At most two lattice periods allowed for a rectangular lattice (got %d)" % l)
|
||
|
||
import numpy as np
|
||
a.basis_vectors = [(a.period[0], 0.), (0., a.period[1])]
|
||
self.direct_basis = np.array(a.basis_vectors)
|
||
self.reciprocal_basis1 = np.linalg.inv(self.direct_basis.T)
|
||
self.reciprocal_basis2pi = 2 * np.pi * self.reciprocal_basis1
|
||
|
||
def _process_planewave_angles(self): #feature: planewave
|
||
import math
|
||
pi2 = math.pi/2
|
||
a = self.args
|
||
a.chi = a.chi * pi2
|
||
a.psi = a.psi * pi2
|
||
a.theta = a.theta * pi2
|
||
a.phi = a.phi * pi2
|
||
|
||
def _process_multi_particle(self): # feature: multi_particle
|
||
a = self.args
|
||
self.tmspecs = {}
|
||
self.tmgens = {}
|
||
self.bspecs = {}
|
||
self.positions = {}
|
||
pos13, pos23, pos33 = False, False, False # used to
|
||
if len(a.position.keys()) == 0:
|
||
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():
|
||
try:
|
||
lMax = a.lMax.get(poslabel, 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)
|
||
if hasattr(a, 'lMax_extend'):
|
||
lMax_extend = a.lMax_extend.get(poslabel, False) or a.lMax_extend.get(None, False) or None
|
||
else:
|
||
lMax_extend = None
|
||
material = a.material.get(poslabel, False) or a.material[None]
|
||
except (TypeError, KeyError) as exc:
|
||
if poslabel is None:
|
||
raise ArgumentProcessingError("Unlabeled particles' positions (-p) specified, but some default particle properties are missing (--lMax, --radius, and --material have to be specified)") from exc
|
||
else:
|
||
raise ArgumentProcessingError(("Incomplete specification of '%s'-labeled particles: you must"
|
||
"provide at least ++lMax, ++radius, ++material arguments with the label, or the fallback arguments"
|
||
"--lMax, --radius, --material.")%(str(poslabel),)) from exc
|
||
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)
|
||
poslist_cured = []
|
||
for pos in a.position[poslabel]:
|
||
if len(pos) == 1:
|
||
pos_cured = (0., 0., pos[0])
|
||
pos13 = True
|
||
elif len(pos) == 2:
|
||
pos_cured = (pos[0], pos[1], 0.)
|
||
pos23 = True
|
||
elif len(pos) == 3:
|
||
pos_cured = pos
|
||
pos33 = True
|
||
else:
|
||
raise argparse.ArgumentTypeError("Each -p / +p argument requires 1 to 3 cartesian coordinates")
|
||
poslist_cured.append(pos_cured)
|
||
self.positions[poslabel] = poslist_cured
|
||
if pos13 and pos23:
|
||
warnings.warn("Both 1D and 2D position specifications used. The former are interpreted as z coordinates while the latter as x, y coordinates")
|
||
|
||
def get_particles(self):
|
||
"""Creates a list of Particle instances that can be directly used in ScatteringSystem.create().
|
||
|
||
Assumes that self._process_multi_particle() has been already called.
|
||
"""
|
||
from .qpms_c import Particle
|
||
plist = []
|
||
for poslabel, poss in self.positions.items():
|
||
t = self.tmgens[poslabel]
|
||
bspec = self.bspecs[poslabel]
|
||
plist.extend([Particle(pos, t, bspec=bspec) for pos in poss])
|
||
return plist
|
||
|
||
#TODO perhaps move into another module
|
||
def annotate_pdf_metadata(pdfPages, scriptname=None, keywords=None, author=None, title=None, subject=None, **kwargs):
|
||
"""Adds QPMS version-related metadata to a matplotlib PdfPages object
|
||
|
||
Use before closing the PDF file.
|
||
"""
|
||
from .qpms_c import qpms_library_version
|
||
d = pdfPages.infodict()
|
||
d['Creator'] = "QPMS%s (lib rev. %s), https://qpms.necada.org" % (
|
||
"" if scriptname is None else (" "+scriptname), qpms_library_version())
|
||
if author is not None:
|
||
d['Author'] = author
|
||
if title is not None:
|
||
d['Title'] = title
|
||
if subject is not None:
|
||
d['Subject'] = subject
|
||
if keywords is not None:
|
||
d['Keywords'] = ' '.join(keywords)
|
||
d.update(kwargs)
|