Files
pmsco-public/pmsco/graphics/scattering.py

211 lines
6.6 KiB
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import logging
import math
import numpy as np
import scipy.interpolate
import scipy.special
logger = logging.getLogger(__name__)
try:
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
except ImportError:
Figure = None
FigureCanvas = None
logger.warning("error importing matplotlib. graphics rendering disabled.")
class TMatrix(object):
def __init__(self):
"""
self.en.shape = (n_e,)
self.tl.shape = (n_e, n_l)
"""
self.en = None
self.tl = None
def load_test_data(self):
self.en = np.array([100.])
raw = [-0.052845, -0.003238, 0.478705, 0.672581, 0.137932, 0.981700, 0.323890, 0.805299, 0.291814, 0.776792,
0.369416, 0.351845, 0.199775, 0.113314, 0.062479, 0.025691, 0.013699, 0.005283]
re_tl = np.array(raw[0::2])
im_tl = np.array(raw[1::2])
self.tl = re_tl + 1j * im_tl
def load_edac_scattering(self, f, energy=math.nan):
"""
load T matrix from EDAC scattering file
currently, only the 'tl' format is supported.
@param f: file path
@param energy: kinetic energy in eV if none is defined in the file
@return: None
"""
with open(f, "r") as fi:
h = fi.readline().rstrip().split(' ')
ne = int(h[0])
if ne > 1:
assert h[1] == 'E(eV)'
del h[1]
lmax = int(h[1])
assert h[2] == 'regular'
assert h[3] == 'tl'
self.load_edac_tl(f, ne, lmax, energy=energy)
def load_edac_tl(self, f, ne, lmax, energy=math.nan):
"""
load T matrix from EDAC scattering file in 'tl' format
@param f: file path
@param ne: number of energies (rows)
@param lmax: maximum l number (columns = 2 * (lmax + 1))
@param energy: kinetic energy in eV if none is defined in the file
@return: None
"""
if ne > 1:
self.en = np.atleast_1d(np.genfromtxt(f, skip_header=1, usecols=[0]))
start_col = 1
else:
self.en = np.asarray(energy)
start_col = 0
re_cols = range(start_col, start_col + (lmax + 1) * 2, 2)
im_cols = range(start_col + 1, start_col + (lmax + 1) * 2, 2)
re_tl = np.atleast_1d(np.genfromtxt(f, skip_header=1, usecols=re_cols))
im_tl = np.atleast_1d(np.genfromtxt(f, skip_header=1, usecols=im_cols))
self.tl = re_tl + 1j * im_tl
assert self.tl.shape == (ne, lmax + 1), "array shape mismatch"
def planewave_amplitude(self, energy, angle):
"""
total, complex plane wave scattering amplitude for given energy and angle
@param energy: kinetic energy in eV.
this can be a numeric value, a 1-dimensional numpy.ndarray,
or any value accepted by the numpy.asarray function.
@param angle: scattering angle in degrees (0..180).
this can be a numeric value, a 1-dimensional numpy.ndarray,
or any value accepted by the numpy.asarray function.
@return: 3 numpy arrays (amp, magnitude, phase) representing the scattering amplitude
versus energy and angle.
the shape of the three arrays is (n_energies, n_angles).
@arg amp: complex scattering amplitude.
@arg magnitude: magnitude (absolute value) of the scattering amplitude.
@arg phase: phase angle in radians of the scattering amplitude.
"""
if not isinstance(energy, np.ndarray):
energy = np.atleast_1d(np.asarray(energy))
ne = len(energy)
if not isinstance(angle, np.ndarray):
angle = np.atleast_1d(np.array(angle))
na = len(angle)
kinv = 1. / (0.513019932 * np.sqrt(energy))
f_tl = scipy.interpolate.interp1d(self.en, self.tl, axis=0, copy=False)
tl = f_tl(energy)
cos_angle = np.cos(np.radians(angle))
lmax = self.tl.shape[1] - 1
l = np.arange(0, lmax + 1)
amp = np.zeros((ne, na), dtype=complex)
for ia, ca in enumerate(cos_angle):
lpmn, __ = scipy.special.lpmn(0, lmax, ca)
fpart = np.outer(kinv, (2 * l + 1) * lpmn[0]) * tl
ftot = np.sum(fpart, axis=-1)
amp[:, ia] = ftot
mag = np.abs(amp)
pha = np.angle(amp)
return amp, mag, pha
def render_scattering_1d(filename, tmatrix, energy=None):
if energy is None:
en = tmatrix.en[0]
else:
en = energy
an = np.arange(0, 181, 2)
__, mag, pha = tmatrix.planewave_amplitude(en, an)
pha = pha / math.pi
canvas = FigureCanvas
fig = Figure()
canvas(fig)
ax = fig.add_subplot(211)
ax.plot(an, mag[0])
ax.set_xlabel('th (deg)')
ax.set_ylabel('mag (arb)')
ax = fig.add_subplot(212)
ax.plot(an, pha[0])
ax.set_xlabel('th (deg)')
ax.set_ylabel('pha (1/pi)')
out_filename = "{0}.{1}".format(filename, canvas.get_default_filetype())
fig.savefig(out_filename)
return out_filename
def render_scattering_2d(filename, tmatrix):
en = tmatrix.en
an = np.arange(0, 181, 2)
__, mag, pha = tmatrix.planewave_amplitude(en, an)
pha = pha / math.pi
canvas = FigureCanvas
fig = Figure()
canvas(fig)
ax = fig.add_subplot(211)
im = ax.imshow(mag, origin='lower', aspect='auto', interpolation='none')
im.set_extent((an[0], an[-1], en[0], en[-1]))
im.set_cmap("magma")
ax.set_xlabel('th (deg)')
ax.set_ylabel('E (eV)')
# cb = ax.colorbar(im, shrink=0.4, pad=0.1)
# ti = cb.get_ticks()
# ti = [0., max(ti)]
# cb.set_ticks(ti)
ax = fig.add_subplot(212)
im = ax.imshow(pha, origin='lower', aspect='auto', interpolation='none')
im.set_extent((an[0], an[-1], en[0], en[-1]))
im.set_cmap("RdBu_r")
ax.set_xlabel('th (deg)')
ax.set_ylabel('E (eV)')
# cb = ax.colorbar(im, shrink=0.4, pad=0.1)
dlo = np.nanpercentile(mag, 2)
dhi = np.nanpercentile(mag, 98)
dhi = max(abs(dlo), abs(dhi))
dlo = -dhi
im.set_clim((dlo, dhi))
# ti = cb.get_ticks()
# ti = [min(ti), 0., max(ti)]
# cb.set_ticks(ti)
out_filename = "{0}.{1}".format(filename, canvas.get_default_filetype())
fig.savefig(out_filename)
return out_filename
def render_scattering_map(filename, energy):
tmatrix = TMatrix()
tmatrix.load_edac_scattering(filename, energy)
if tmatrix.tl.shape[0] == 1:
out_filename = render_scattering_1d(filename, tmatrix)
else:
out_filename = render_scattering_2d(filename, tmatrix)
return out_filename