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