public release 4.2.0 - see README.md and CHANGES.md for details
This commit is contained in:
371
pmsco/transforms/multipoles.py
Normal file
371
pmsco/transforms/multipoles.py
Normal file
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
@package pmsco.transforms.multipoles
|
||||
Transform holoscan data to multipoles expansion.
|
||||
|
||||
@author Matthias Muntwiler
|
||||
|
||||
@copyright (c) 2015-24 by Paul Scherrer Institut @n
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); @n
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
"""
|
||||
|
||||
import logging
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import os
|
||||
import scipy.special
|
||||
from typing import Any, Callable, Dict, Generator, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Union
|
||||
import h5py
|
||||
|
||||
from pmsco.data import analyse_holoscan_steps, create_data, DATATYPES, DTYPE_TPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MultipoleExpansion:
|
||||
"""
|
||||
Multipole expansion (spherical harmonics) of hologram scans
|
||||
|
||||
This class provides functions to calculate the multipole expansion of hologram scans.
|
||||
|
||||
Usage
|
||||
-----
|
||||
|
||||
1. Instantiate the object
|
||||
2. Assign a scan to `holoscan`.
|
||||
If you will use the `generate` method, the scan must include an intensity column.
|
||||
Else, the theta and phi coordinates are sufficient.
|
||||
3. Set `lmax`.
|
||||
4. Call `generate` to calculate the alm coefficients, or assign the coefficients directly.
|
||||
5. Call `expand` to calculate the filtered pattern.
|
||||
|
||||
The object keeps intermediate calculation data, e.g. spherical harmonics, to speed up further calculations.
|
||||
You have to keep the object alive in order to benefit from these performance improvements.
|
||||
Reuse of the spherical harmonics depends on the same theta-phi scan grid and same lmax.
|
||||
The class re-calculates the internal data automatically if these conditions are not met.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
|
||||
_sph_harm dimensions: l, m, tp
|
||||
_d_omege dimensions: l, m, tp
|
||||
_alm dimensions: l, m
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""
|
||||
Initialize the object
|
||||
|
||||
The scan and arrays are set to shape 0.
|
||||
The default lmax is 30 and the alm array is initialized with zeros to the corresponding shape.
|
||||
The spherical harmonics and volume element are not initialized.
|
||||
"""
|
||||
|
||||
super().__init__()
|
||||
self._scan: np.ndarray = create_data((0,), datatype='TP')
|
||||
self._original: np.ndarray = np.zeros(self._scan.shape)
|
||||
self._expansion: np.ndarray = np.zeros(self._scan.shape)
|
||||
self._lmax: int = 60
|
||||
self._alm: np.ndarray = np.zeros((self._lmax // 2 + 1, 2 * self._lmax + 1))
|
||||
self._sph_harm: Optional[np.ndarray] = None
|
||||
self._d_omega: Optional[np.ndarray] = None
|
||||
|
||||
@property
|
||||
def holoscan(self) -> np.ndarray:
|
||||
"""
|
||||
Holoscan grid and original intensity pattern.
|
||||
|
||||
The holoscan array can be created, e.g., from holo_grid or is taken from measured data.
|
||||
't' and 'p' contain the polar and azimuthal coordinates of the scan.
|
||||
The 'i' column contains the intensity value.
|
||||
|
||||
The array must contain at least the 't' and 'p' columns if you want to expand the pattern from alm.
|
||||
If you want to generate the coefficients, the 'i' (intensity) column must be set as well.
|
||||
|
||||
@attention: The property returns a copy of internal data.
|
||||
Do not change individual values or columns directly on the array.
|
||||
"""
|
||||
|
||||
tpi = np.zeros(self._scan.shape, dtype=DTYPE_TPI)
|
||||
tpi['t'] = self._scan['t']
|
||||
tpi['p'] = self._scan['p']
|
||||
tpi['i'] = self._original
|
||||
|
||||
return tpi
|
||||
|
||||
@holoscan.setter
|
||||
def holoscan(self, arr: np.ndarray):
|
||||
if arr.shape != self._scan.shape or \
|
||||
self._original.shape[0] == 0 or \
|
||||
np.all(np.abs(arr['i'] - self._original) < 0.01):
|
||||
self._scan = create_data(arr.shape, datatype='TP')
|
||||
self._discard_basis()
|
||||
|
||||
self._original = arr['i']
|
||||
self._scan['t'] = arr['t']
|
||||
self._scan['p'] = arr['p']
|
||||
|
||||
@property
|
||||
def expansion(self) -> np.ndarray:
|
||||
"""
|
||||
Return the expansion in TPI format.
|
||||
|
||||
The expansion must have been calculated using the `expand` method.
|
||||
|
||||
@return: Structured numpy array in TPI format (dtype=DTYPE_TPI).
|
||||
"""
|
||||
|
||||
tpi = create_data(self._expansion.shape, dtype=DTYPE_TPI)
|
||||
tpi['t'] = self._scan['t']
|
||||
tpi['p'] = self._scan['p']
|
||||
tpi['i'] = self._expansion
|
||||
|
||||
return tpi
|
||||
|
||||
def get_expansion(self, datatype: Optional[str] = None, dtype: Optional[npt.DTypeLike] = None) -> np.ndarray:
|
||||
"""
|
||||
Return the expansion in the requested format.
|
||||
|
||||
The expansion must have been calculated using the `expand` method.
|
||||
Either datatype or dtype must be specified, dtype takes precedence.
|
||||
|
||||
@param datatype see DATATYPES.
|
||||
@param dtype see DTYPES.
|
||||
Any structured array type will work.
|
||||
The dtype must specify at least the `i`, `t` and `p` columns.
|
||||
|
||||
@return: structured numpy array
|
||||
"""
|
||||
|
||||
if not datatype or datatype not in DATATYPES:
|
||||
datatype = "TPI"
|
||||
|
||||
tpi = create_data(self._expansion.shape, datatype=datatype, dtype=dtype)
|
||||
tpi['t'] = self._scan['t']
|
||||
tpi['p'] = self._scan['p']
|
||||
tpi['i'] = self._expansion
|
||||
|
||||
return tpi
|
||||
|
||||
@property
|
||||
def lmax(self):
|
||||
"""
|
||||
Maximum l angular number used
|
||||
|
||||
Determines the size of the alm array and number of coefficients calculated by `generate`.
|
||||
|
||||
The useful range for XPD scans is from 30 to 90.
|
||||
The default is 60, which is sufficient to reproduce fine backscattering details.
|
||||
"""
|
||||
|
||||
return self._lmax
|
||||
|
||||
@lmax.setter
|
||||
def lmax(self, value):
|
||||
if value != self._lmax:
|
||||
self._lmax = value
|
||||
self._alm = np.zeros((self._lmax // 2 + 1, 2 * self._lmax + 1))
|
||||
self._sph_harm = None
|
||||
self._d_omega = None
|
||||
|
||||
@property
|
||||
def alm(self):
|
||||
"""
|
||||
Multipole coefficients a_lm
|
||||
|
||||
Multipole coefficients a_lm are stored in a two-dimensional ndarray.
|
||||
The shape is (lmax / 2 + 1, 2 * lmax + 1),
|
||||
where the first dimension is l (0...lmax, step 2),
|
||||
and the second dimension is m (-lmax...+lmax, step 1).
|
||||
|
||||
There are two ways to modify this array:
|
||||
by assignment of a whole array or modification of individual values.
|
||||
|
||||
On assignment of a whole array, lmax is calculated from the shape of the array.
|
||||
The array is copied, NaNs are converted to zeros.
|
||||
Items |m| > l are ignored and should be zero.
|
||||
"""
|
||||
|
||||
return self._alm
|
||||
|
||||
@alm.setter
|
||||
def alm(self, value):
|
||||
lmax = (value.shape[0] - 1) * 2
|
||||
assert value.shape == (lmax // 2 + 1, 2 * lmax + 1)
|
||||
if lmax != self._lmax:
|
||||
self._lmax = lmax
|
||||
self._sph_harm = None
|
||||
self._d_omega = None
|
||||
self._alm = np.nan_to_num(value, copy=True, nan=0.0)
|
||||
|
||||
def generate(self):
|
||||
"""
|
||||
Calculate the multipole coefficients for the stored holoscan
|
||||
|
||||
Requires: lmax and holoscan
|
||||
Generates: alm
|
||||
|
||||
@return alm: Multipole coefficients in a two-dimensional ndarray.
|
||||
The shape is (lmax / 2 + 1, 2 * lmax + 1),
|
||||
where the first dimension is l (0...lmax, step 2),
|
||||
and the second dimension is m (-lmax...+lmax, step 1).
|
||||
Items |m| > l are 0.
|
||||
"""
|
||||
|
||||
if not self._check_basis():
|
||||
self._calc_basis()
|
||||
|
||||
self._alm = np.sum(self._original * np.conj(self._sph_harm) * self._d_omega, axis=-1)
|
||||
|
||||
return self._alm
|
||||
|
||||
def expand(self, lmin: Optional[int] = None, lmax: Optional[int] = None) -> np.ndarray:
|
||||
"""
|
||||
Calculate the expanded pattern from coefficients
|
||||
|
||||
The alm property and the `t` and `p` columns of `holoscan` must be valid.
|
||||
The result is available from `expansion` and as the function result.
|
||||
|
||||
The l range used in the expansion can be constricted by the lmin and lmax options.
|
||||
By default, the full alm matrix is used.
|
||||
|
||||
@param lmin: minimum l to include in expansion.
|
||||
Allowed range: 0...self.lmax.
|
||||
@param lmax: maximum l to include in expansion.
|
||||
Allowed range: 0...self.lmax.
|
||||
By default, all coefficients are included.
|
||||
@return: intensity array versus `holoscan`.
|
||||
This array contains just the intensity.
|
||||
The `expansion` property provides the full scan object in TPI structure.
|
||||
Other structures can be retrieved from the `get_expansion` method.
|
||||
"""
|
||||
|
||||
if not self._check_basis():
|
||||
self._calc_basis()
|
||||
|
||||
if lmax or lmin:
|
||||
alm = self._alm.copy()
|
||||
if lmin is not None and 0 <= lmin <= self._lmax:
|
||||
alm[0:lmin, :] = 0
|
||||
if lmax is not None and 0 <= lmax <= self._lmax:
|
||||
alm[lmax+1:self._lmax+1, :] = 0
|
||||
else:
|
||||
alm = self._alm
|
||||
|
||||
self._expansion = np.real(np.sum(np.sum(alm * self._sph_harm.transpose((2, 0, 1)), axis=-1), axis=-1))
|
||||
|
||||
return self._expansion
|
||||
|
||||
def export_alm_h5(self, filename: os.PathLike):
|
||||
"""
|
||||
Export the alm coefficients to an HDF5 file.
|
||||
|
||||
The HDF5 file contains 3 datasets: alm, l and m.
|
||||
The alm matrix is written directly as a 2D complex array
|
||||
(compound type with members r and i containing the real and imaginary parts, respectively).
|
||||
l and m ranges are linked as dimension scales.
|
||||
|
||||
@param filename: path-like
|
||||
@return: None
|
||||
"""
|
||||
|
||||
range_l = np.arange(0, self._lmax + 1, 2)
|
||||
range_m = np.arange(-self._lmax, self._lmax + 1, 1)
|
||||
with h5py.File(filename, "w") as h5:
|
||||
h5["l"] = range_l
|
||||
h5["m"] = range_m
|
||||
h5["l"].make_scale()
|
||||
h5["m"].make_scale()
|
||||
h5["alm"] = self._alm
|
||||
h5["alm"].dims[0].label = "l"
|
||||
h5["alm"].dims[1].label = "m"
|
||||
h5["alm"].dims[0].attach_scale(h5["l"])
|
||||
h5["alm"].dims[1].attach_scale(h5["m"])
|
||||
|
||||
def import_alm_h5(self, filename: os.PathLike):
|
||||
"""
|
||||
Import the alm coefficients from an HDF5 file.
|
||||
|
||||
The file must use the same structure as created by `export_alm_h5`.
|
||||
Alternatively, the real and imaginary parts can be in separate datasets "alm_r" and "alm_i", respectively.
|
||||
|
||||
@param filename: path-like
|
||||
@return:
|
||||
"""
|
||||
|
||||
with h5py.File(filename, "r") as h5:
|
||||
try:
|
||||
self.alm = h5["alm"][()]
|
||||
except KeyError:
|
||||
self.alm = complex(h5["alm_r"][()], h5["alm_i"][()])
|
||||
|
||||
def _check_basis(self):
|
||||
"""
|
||||
Check whether the internal basis functions are consistent with scan data.
|
||||
|
||||
# sph dimensions: l, m, tp
|
||||
|
||||
@return:
|
||||
"""
|
||||
|
||||
nl = self._lmax // 2 + 1
|
||||
nm = 2 * self._lmax + 1
|
||||
if self._sph_harm is None or self._d_omega is None:
|
||||
return False
|
||||
if self._sph_harm.shape[0] != nl or self._d_omega.shape[0] != nl:
|
||||
return False
|
||||
if self._sph_harm.shape[1] != nm or self._d_omega.shape[1] != nm:
|
||||
return False
|
||||
if self._sph_harm.shape[2] != self._scan.shape[0] or self._d_omega.shape[2] != self._scan.shape[0]:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _discard_basis(self):
|
||||
"""
|
||||
Discard the basis functions (spherical harmonics and volume element).
|
||||
|
||||
This is necessary when the scan positions, lmax or shape of alm coefficients have changed.
|
||||
It is not necessary after changes to the intensity column or alm values.
|
||||
The basis functions will be recalculated automatically.
|
||||
|
||||
@return: None
|
||||
"""
|
||||
|
||||
self._sph_harm = None
|
||||
self._d_omega = None
|
||||
|
||||
def _calc_basis(self):
|
||||
"""
|
||||
Calculate the basis functions.
|
||||
|
||||
The basis functions include the spherical harmonics and the volume element.
|
||||
They depend on the scan grid and lmax.
|
||||
For better performance, the values of the basis functions are evaluated on a particular scan grid
|
||||
and stored for further use.
|
||||
|
||||
They have to be re-calculated whenever lmax or the scan angles have changed.
|
||||
It is not necessary after changes to the intensity column or alm values.
|
||||
|
||||
@return: None
|
||||
"""
|
||||
|
||||
range_l = np.arange(0, self._lmax + 1, 2)
|
||||
range_m = np.arange(-self._lmax, self._lmax + 1, 1)
|
||||
tpi = np.arange(self._scan.shape[0])
|
||||
# order of dimensions is l, m, tp
|
||||
mesh_m, mesh_l, mesh_tpi = np.meshgrid(range_m, range_l, tpi)
|
||||
mesh_theta = np.deg2rad(self._scan['t'][mesh_tpi])
|
||||
mesh_phi = np.deg2rad(self._scan['p'][mesh_tpi])
|
||||
|
||||
# sph_harm returns nan for bad l/m combinations
|
||||
self._sph_harm = scipy.special.sph_harm(mesh_m, mesh_l, mesh_phi, mesh_theta)
|
||||
self._sph_harm = np.nan_to_num(self._sph_harm, copy=False, nan=0.0)
|
||||
|
||||
unique_theta, d_theta, d_phi = analyse_holoscan_steps(self._scan)
|
||||
mesh_sin = np.interp(mesh_theta, unique_theta, np.sin(np.deg2rad(unique_theta)))
|
||||
mesh_dtheta = np.interp(mesh_theta, unique_theta, np.deg2rad(d_theta))
|
||||
mesh_dphi = np.interp(mesh_theta, unique_theta, np.deg2rad(d_phi))
|
||||
self._d_omega = mesh_sin * mesh_dtheta * mesh_dphi
|
||||
Reference in New Issue
Block a user