tmatrices.py: get TMatrix by CL argument specification

Former-commit-id: e28bd7f80f64a810a1f2a22534e3fc1a40341dc9

TODOs.rst

@@ -6,3 +6,5 @@ LMAXVAR - besides to int, the function should support an iterable for the lMax a
           to support systems where different nanoparticles have different lMax.
 
 VECTORIZE - add support to process more than one element at once, in general.
+
+FEATURE - non-urgent general feature to implement

qpms/qpms_c.pyx

@@ -42,6 +42,9 @@ def get_mn_y(int nmax):
             i = i + 1
     return (m_arr, n_arr)
 
+def get_nelem(unsigned int lMax):
+    return lMax * (lMax + 2)
+
 def get_y_mn_unsigned(int nmax): 
     """
     Auxillary function for mapping 'unsigned m', n indices to the flat y-indexing.

qpms/scripts_common.py

@@ -1,7 +1,11 @@
-#import argparse # Do I need it when calling just parser methods?
 import warnings
+import argparse
+#import sys # for debugging purpose, TODO remove in production
+import os # because of path
 
 __TODOs__ = '''
+    - Checking validity of T-matrix ops (the arguments of --tr, --sym or similar) according to what is implemented
+      in tmatrices.py.
     - Implement a more user-friendly way to define the lattice base vectors and positions of the particles.
       cf. https://stackoverflow.com/questions/2371436/evaluating-a-mathematical-expression-in-a-string/2371789
     - low priority: allow to perform some more custom operations on T-Matrix, using some kind of parsing from the previous point
@@ -35,7 +39,7 @@ def add_argparse_unitcell_definitions(parser):
     parser.add_argument('--particle', '-p', nargs='+', action=make_action_sharedlist('particle', 'particlespec'), help='Particle label, coordinates x,y, and (optionally) path to the T-Matrix.')
     parser.add_argument('--TMatrix', '-t', nargs='+', action=make_action_sharedlist('TMatrix_path', 'particlespec'), help='Path to TMatrix file')
     parser.add_argument('--background_permittivity', action='store', type=float, default=1., help='Background medium relative permittivity (default 1)')
-    parser.add_argument('--lMax', action=make_action_sharedlist('lMax', 'particlespec'), nargs=+, help='Override lMax from the TMatrix file')
+    parser.add_argument('--lMax', action=make_action_sharedlist('lMax', 'particlespec'), nargs='+', help='Override lMax from the TMatrix file')
     popgrp=parser.add_argument_group(title='Operations')
     popgrp.add_argument('--tr', dest='ops', nargs='+', action=make_action_sharedlist('tr', 'ops'), default=list()) # the default value for dest can be set once
     popgrp.add_argument('--sym', dest='ops', nargs='+', action=make_action_sharedlist('sym', 'ops'))
@@ -62,7 +66,7 @@ def add_argparse_common_options(parser):
 
 
 
-def arg_preprocess_particles(parser, d=None, return_tuple=False):
+def arg_preprocess_particles(pargs, d=None, return_tuple=False):
     ''' 
     Nanoparticle position and T-matrix path parsing 
 
@@ -91,20 +95,20 @@ def arg_preprocess_particles(parser, d=None, return_tuple=False):
     lMax_overrides = dict()
     default_TMatrix_path = None
     default_lMax_override = None
-    if not any((arg_type == 'particle') in (arg_type, arg_content) for in pargs.particlespec):
+    if not any((arg_type == 'particle') for (arg_type, arg_content) in pargs.particlespec):
         # no particles positions given: suppose only one per unit cell, in the cell origin
         positions = {None: (0.0)}
     else:
         positions = dict()
     for arg_type, arg_content in pargs.particlespec:
-        if arg_type == 'particle' # --particle option
+        if arg_type == 'particle': # --particle option
             if  3 <= len(arg_content) <= 4:
                 try:
                     positions[arg_content[0]] = (float(arg_content[1]), float(arg_content[2]))
                 except ValueError as e:
                     e.args += ("second and third argument of --particle must be valid floats, given: ", arg_content)
                     raise
-                if len(arg_content == 4):
+                if len(arg_content) == 4:
                     if arg_content[0] in TMatrix_paths:
                         warnings.warn('T-matrix path for particle \'%s\' already specified.' 
                         'Overriding with the last value.' % arg_content[0], SyntaxWarning)
@@ -145,7 +149,10 @@ def arg_preprocess_particles(parser, d=None, return_tuple=False):
     if (set(TMatrix_paths.keys()) != set(positions.keys())) and default_TMatrix_path is None:
         raise ValueError("Position(s) of particles(s) labeled %s was given without their T-matrix"
             " and no default T-matrix was specified" 
-            % str(set(positions.keys()) - set(TMatrix_paths_keys())))
+            % str(set(positions.keys()) - set(TMatrix_paths.keys())))
+    # Fill default_TMatrix_path to those that don't have its own
+    for label in (set(positions.keys()) - set(TMatrix_paths.keys())):
+        TMatrix_paths[label] = default_TMatrix_path
     for path in TMatrix_paths.values():
         if not os.path.exists(path):
             raise ValueError("Cannot access T-matrix file %s. Does it exist?" % path)
@@ -162,8 +169,6 @@ def arg_preprocess_particles(parser, d=None, return_tuple=False):
                 e.args += 'Specified operation on undefined particle labeled \'%s\'' % label
                 raise
 
-    print(sys.stderr, "ops: ", ops) #DEBUG
-
     #### Collect all the info about the particles / their T-matrices into one list ####
     # Enumerate and assign all the _different_ T-matrices (without any intelligent group-theory checking, though)
     TMatrix_specs = dict((spec, number) 
@@ -174,7 +179,7 @@ def arg_preprocess_particles(parser, d=None, return_tuple=False):
                 for label in positions.keys()
             )))
     # particles_specs contains (label, (xpos, ypos), tmspec_index per element)
-    particles_specs = [(label, positions(label), 
+    particles_specs = [(label, positions[label], 
         TMatrix_specs[(lMax_overrides[label] if label in lMax_overrides.keys() else None, 
                        TMatrix_paths[label], 
                        tuple(ops[label]))]
@@ -186,7 +191,8 @@ def arg_preprocess_particles(parser, d=None, return_tuple=False):
     
     if d is None: 
         d = dict()
-    d['particle_specs'] = particle_specs
+    # TODO what if I used classes instead?
+    d['particle_specs'] = particles_specs
     d['TMatrix_specs'] = TMatrix_specs
     if return_tuple:
         return (particles_specs, TMatrix_specs)

qpms/tmatrices.py

@@ -1,5 +1,10 @@
 import numpy as np
 import quaternion, spherical_functions as sf # because of the Wigner matrices. These imports are SLOW.
+import re
+from scipy import interpolate
+from scipy.constants import hbar, e as eV, pi, c
+from qpms_c import get_mn_y#, get_nelem # TODO IMPORT get_nelem INTO THE FINAL MODULE
+ň = np.newaxis
 
 # Transformations of spherical bases
 def WignerD_mm(l, quat):
@@ -239,3 +244,94 @@ def symz_indexarrays(lMax, npart = 1):
     TEč = č[(mč+nč+tč) % 2 == 0]
     TMč = č[(mč+nč+tč) % 2 == 1]
     return (TEč, TMč)
+
+
+"""
+Processing T-matrix related operations from scripts
+===================================================
+
+see also scripts_common.py
+
+The unit cell is defined by a dict particle_specs and a list TMatrix_specs.
+This particular module has to provide the T-matrices according to what is defined
+in TMatrix_specs
+
+TMatrix_specs is a list of tuples (lMax_override, TMatrix_path, ops)
+where 
+  - TMatrix_path is path to the file generated by scuff-tmatrix
+  - lMax_override is int or None; if it is int and less than the lMax found from the T-matrix file,
+    lMax_override is used as the order cutoff for the output T-matrix.
+  - ops is an iterable of tuples (optype, opargs) where currently optype can be 'sym' or 'tr'
+    for symmetrization operation or some other transformation.
+"""
+
+#TODO FEATURE more basic group symmetry operations, cf. http://symmetry.otterbein.edu/tutorial/index.html
+
+# This is for finite „fractional“ rotations along the z-axis (mCN means rotation of 2π*(m/N))
+reCN = re.compile('(\d*)C(\d+)') # TODO STYLE make this regexp also accept the 3*C_5-type input, eqiv. to 3C5
+
+def get_TMatrix_fromspec(tmatrix_spec):
+    ''' TODO doc 
+    returns (TMatrices, freqs, lMax)
+    '''
+    lMax_override, tmpath, ops = tmatrix_spec
+    TMatrices, freqs, freqs_weirdunits, lMax = loadScuffTMatrices(tmpath)
+    if lMax_override is not None and (lMax_override < lMax_orig):
+        nelem = get_nelem(lMax_override)
+        TMatrices = TMatrices[...,0:nelem,:,0:nelem]
+        lMax = lMax_override
+    
+    for (optype, opargs) in ops:
+        if optype == 'sym':
+            mCN = reCN.match(opargs)
+            if opargs == 'C2' or opargs == 'C_2':
+                opmat = apply_matrix_left(yflip_yy(lMax), xflip_yy(lMax), -1)
+                TMatrices = (TMatrices + apply_matrix_left(opmat, apply_matrix_left(opmat, TMatrices, -3), -1))/2
+            elif opargs == 'σ_x' or opargs == 'σx':
+                opmat = xflip_tyty(lMax)
+                TMatrices = (TMatrices + apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1)))/2
+            elif opargs == 'σ_y' or opargs == 'σy':
+                opmat = yflip_tyty(lMax)
+                TMatrices = (TMatrices + apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1)))/2
+            elif opargs == 'σ_z' or opargs == 'σz':
+                opmat = zflip_tyty(lMax)
+                TMatrices = (TMatrices + apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1)))/2
+            elif mCN:
+                rotN = int(mCN.group(2)) # the possible m is ignored
+                TMatrix_contribs = np.empty((rotN,)+TMatrices.shape, dtype=np.complex_)
+                for i in range(rotN):
+                    rotangle = 2*np.pi*i / rotN
+                    rot = WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
+                    rotinv = WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
+                    TMatrix_contribs[i] = apply_matrix_left(rot, apply_matrix_left(rotinv, TMatrices, -3), -1)
+                TMatrices = np.sum(TMatrix_contribs, axis=0) / rotN
+            elif opargs == 'C_inf' or opargs == 'Cinf' or opargs == 'C_∞' or opargs == 'C∞':
+                raise ValueError('not implemented: ', opargs)  # TODO FEATURE
+            else:
+                raise ValueError('not implemented: ', opargs)
+        elif optype == 'tr':
+            mCN = reCN.match(opargs)
+            if opargs == 'C2' or opargs == 'C_2':
+                opmat = apply_matrix_left(yflip_yy(lMax), xflip_yy(lMax), -1)
+                TMatrices = apply_matrix_left(opmat, apply_matrix_left(opmat, TMatrices, -3), -1)/2
+            elif opargs == 'σ_x' or opargs == 'σx':
+                opmat = xflip_tyty(lMax)
+                TMatrices = apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1))/2
+            elif opargs == 'σ_y' or opargs == 'σy':
+                opmat = yflip_tyty(lMax)
+                TMatrices = apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1))/2
+            elif opargs == 'σ_z' or opargs == 'σz':
+                opmat = zflip_tyty(lMax)
+                TMatrices = apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1))/2
+            elif mCN:
+                rotN = int(mCN.group(2))
+                power = int(mCN.group(1)) if mCN.group(1) else 1
+                rotangle = 2*np.pi*power / rotN
+                rot = WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
+                rotinv = WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
+                TMatrices = apply_matrix_left(rot, apply_matrix_left(rotinv, TMatrices, -3), -1)
+            else:
+                raise ValueError('not implemented: ', opargs)
+        else:
+            raise ValueError('not implemented: ', optype)
+    return (TMatrices, freqs, lMax)