Source code for mumott.methods.utilities.fiber_fit

import sys
from tqdm import tqdm
from typing import Tuple
import numpy as np
from numpy.typing import NDArray
from mumott.core.wigner_d_utilities import load_d_matrices, calculate_sph_coefficients_rotated_by_euler_angles


[docs]def find_approximate_symmetry_axis(coefficients: NDArray[float], ell_max: int, resolution: int = 10, filter: str = None) -> Tuple[NDArray[float]]: r""" Find the axis of highest apparent symmetry voxel-by-voxel for a voxel map of sperical harmonics coefficients. As a default, the measure of degree of symmetry is a power of the function rotationally averaged around the given axis. Parameters ---------- coefficients Voxel map of spherical harmonic coefficients. ell_max Largest order :math:`\ell` used in the fitting. If :attr:`coefficients` contains higher orders, it will be truncated. resolution Number or angular steps along a half-circle, used in the search for the optimal axis. filter Weighing of different orders used to calculate the degree of symmetry. Possible values are "ramp" and "square". By default no filter is applied (`None`). Returns --------- optimal_zonal_coeffs Zonal harmonics coefficients in the frame-of-reference corresponding to the axis of highest degree of symmetry. optimal_theta Voxel-by-voxel polar angles of the axis with highest degree of symmetry. optimal_phi Voxel-by-voxel azimuthal angles of the axis with highest degree of symmetry. """ # Use only the coefficients up until ell_max truncated_coefficients = coefficients[..., :(ell_max+1)*(ell_max+2)//2] volume_shape = coefficients.shape[:-1] # Find the zonal coefficients zonal_indexes = np.zeros((ell_max+1)*(ell_max+2)//2, dtype=bool) inc = 0 f = np.ones(ell_max//2+1) for ell in range(0, ell_max+1, 2): zonal_indexes[inc + ell] = True if filter == 'ramp': f[ell//2] = ell elif filter == 'square': f[ell//2] = ell**2 inc += 2*ell + 1 # Load d matrices d_matrices = load_d_matrices(ell_max) optimal_theta = np.zeros(volume_shape) optimal_phi = np.zeros(volume_shape) maximum_degree_of_symmetry = np.zeros(volume_shape) optimal_zonal_coeffs = np.zeros((*volume_shape, ell_max//2 + 1)) # Loop through grid of directions theta_points = np.linspace(0, np.pi/2, resolution//2, endpoint=True) phi_points = np.linspace(0, 2*np.pi, 2*resolution, endpoint=False) for theta in tqdm(theta_points, file=sys.stdout): for phi in phi_points: rotated_coefficients = calculate_sph_coefficients_rotated_by_euler_angles( truncated_coefficients, Psi=-phi, Theta=-theta, Phi=None, d_matrices=d_matrices) # Check if degree of symmetry (DOS) is higher than maximum so far zonal_coeffs = rotated_coefficients[..., zonal_indexes] degree_of_symmetry = np.sum(zonal_coeffs**2*f[np.newaxis, np.newaxis, np.newaxis, :], axis=-1) indices = degree_of_symmetry > maximum_degree_of_symmetry maximum_degree_of_symmetry[indices] = degree_of_symmetry[indices] optimal_theta[indices] = theta optimal_phi[indices] = phi optimal_zonal_coeffs[indices, :] = zonal_coeffs[indices, :] return optimal_zonal_coeffs, optimal_theta, optimal_phi
[docs]def degree_of_symmetry_map(coefficients: NDArray[float], ell_max: int, resolution: int = 10, filter: str = None) -> Tuple[NDArray[float]]: r""" Make a longitude-latitude map of the degree of symmetry. Can be used to make illustrating plots and to decide which filter is more appropriate. Parameters ---------- coefficients Spherical harmonic coefficients of a single voxel. ell_max Largest order :math:`\ell` used in the fitting. If :attr:`coefficients` contains higher orders, it will be truncated. resolution Number or angular steps along a half-circle used in the search for the optimal axis. filter Weighing of different orders used to calculate the degree of symmetry. Possible values are "ramp" and "square". By default no filter is applied (`None`). Returns --------- dos Map of the calculated degree of symmetry. theta Polar angle coordinates of the map. phi Azimuthal angle coordinates of the map. """ # Use only the coefficients up until ell_max truncated_coefficients = coefficients[..., :(ell_max+1)*(ell_max+2)//2] # Find the zonal coefficients zonal_indexes = np.zeros((ell_max+1)*(ell_max+2)//2, dtype=bool) inc = 0 f = np.ones(ell_max//2+1) for ell in range(0, ell_max+1, 2): zonal_indexes[inc + ell] = True if filter == 'ramp': f[ell//2] = ell elif filter == 'square': f[ell//2] = ell**2 inc += 2*ell + 1 # Load d matrices d_matrices = load_d_matrices(ell_max) # Loop through grid of directions theta_points = np.linspace(0, np.pi, resolution, endpoint=True) phi_points = np.linspace(0, 2*np.pi, 2*resolution, endpoint=True) dos = np.zeros((resolution, 2*resolution)) for ii, theta in enumerate(theta_points): for jj, phi in enumerate(phi_points): rotated_coefficients = calculate_sph_coefficients_rotated_by_euler_angles( truncated_coefficients, Psi=-phi, Theta=-theta, Phi=None, d_matrices=d_matrices) # Check if degree of symmetry (DOS) is higher than maximum so far zonal_coeffs = rotated_coefficients[..., zonal_indexes] degree_of_symmetry = np.sum(zonal_coeffs**2*f, axis=-1) dos[ii, jj] = degree_of_symmetry theta, phi = np.meshgrid(theta_points, phi_points, indexing='ij') dos = dos / np.sum(truncated_coefficients**2) return dos, theta, phi
[docs]def symmetric_part_along_given_direction(coefficients: NDArray[float], theta: NDArray[float], phi: NDArray[float], ell_max: int) -> NDArray[float]: r""" Find the zonal harmonic coefficients along the given input directions. This can be used if eigenvector analysis is used to generate the symmetry directions used for a further zonal-harmonics refinement step. Parameters ---------- coefficients Voxel map of spherical harmonic coefficients. theta Voxel map of polar angles. phi Voxel map of azimuthal angles. ell_max Largest order :math:`\ell` used in the fitting. If :attr:`coefficients` contains higher orders, it will be truncated. Returns --------- Zonal harmonics coefficients in the frame of reference corresponding to the axis of highest degree of symmetry. """ # Use only the coefficients up until ell_max truncated_coefficients = coefficients[..., :(ell_max+1)*(ell_max+2)//2] # Find the zonal coefficients zonal_indexes = np.zeros((ell_max+1)*(ell_max+2)//2, dtype=bool) inc = 0 for ell in range(0, ell_max+1, 2): zonal_indexes[inc + ell] = True inc += 2*ell + 1 # Load d matrices d_matrices = load_d_matrices(ell_max) rotated_coefficients = calculate_sph_coefficients_rotated_by_euler_angles( truncated_coefficients, Psi=-phi, Theta=-theta, Phi=None, d_matrices=d_matrices) # Pick out symmetric part zonal_coeffs = rotated_coefficients[..., zonal_indexes] return zonal_coeffs