examples/api/custom_tmatrix/custom_tmatrices_simple.py

   1 #!/usr/bin/env python3
   2 from qpms import TMatrixGenerator, BaseSpec, eV, hbar
   3 import numpy as np
   4 import sys
   5
   6 errors = 0
   7
   8 def tmg_diagonal_fun(tmatrix, omega):
   9     '''
  10     Example of a python function used as a custom T-matrix generator
  11
  12     It receives a CTMatrix argument with pre-filled BaseSpec
  13     (in tmatrix.spec) and angular frequency.
  14
  15     It has to fill in the T-matrix elements tmatrix[...]
  16     (a numpy array of shape (len(tmatrix.spec),len(tmatrix.spec)))
  17     and return zero (on success) or other integral value on error.
  18
  19     Note that this in justa an example of using the API,
  20     not supposed to be anything physical.
  21     '''
  22     l = tmatrix.spec.l()
  23     tmatrix[...] = np.diag(1./l**2)
  24     return 0
  25
  26 # Wrap the function as an actual TMatrixGenerator
  27 tmg_diagonal = TMatrixGenerator(tmg_diagonal_fun)
  28
  29 bspec = BaseSpec(lMax=2)
  30
  31 tmatrix = tmg_diagonal(bspec, (2.0+.01j) * eV/hbar)
  32
  33 errors += np.sum(tmatrix[...] != np.diag(1./bspec.l()**2))
  34
  35 sys.exit(errors)
  36