charged-shells/charged_shells/expansion.py

from __future__ import annotations
import numpy as np
from dataclasses import dataclass
import charged_shells.functions as fn
import quaternionic
import spherical
import copy

Array = np.ndarray
Quaternion = quaternionic.array


class InvalidExpansion(Exception):
    pass


@dataclass
class Expansion:
    """Generic class for storing surface charge expansion coefficients."""

    l_array: Array
    coefs: Array
    _starting_coefs: Array = None  # initialized with the __post_init__ method
    _rotations: Quaternion = None

    def __post_init__(self):
        if self.coefs.shape[-1] != np.sum(2 * self.l_array + 1):
            raise InvalidExpansion('Number of expansion coefficients does not match the provided l_array.')
        if np.all(np.sort(self.l_array) != self.l_array) or np.all(np.unique(self.l_array) != self.l_array):
            raise InvalidExpansion('Array of l values should be unique and sorted.')
        self.coefs = self.coefs.astype(np.complex128)
        self._starting_coefs = np.copy(self.coefs)
        self._rotations = Quaternion([1., 0., 0., 0.])

    def __getitem__(self, item):
        return Expansion(self.l_array, self.coefs[item])

    @property
    def max_l(self) -> int:
        return max(self.l_array)

    @property
    def shape(self):
        return self.coefs.shape[:-1]

    def flatten(self) -> Expansion:
        new_expansion = self.clone()  # np.ndarray.flatten() also copies the array
        new_expansion.coefs = new_expansion.coefs.reshape(-1, new_expansion.coefs.shape[-1])
        new_expansion._rotations = new_expansion._rotations.reshape(-1, 4)
        return new_expansion

    def reshape(self, shape: tuple):
        self.coefs = self.coefs.reshape(shape + (self.coefs.shape[-1],))
        self._rotations = self._rotations.reshape(shape + (4,))

    @property
    def lm_arrays(self) -> (Array, Array):
        """Return l and m arrays containing all (l, m) pairs."""
        return full_lm_arrays(self.l_array)

    def repeat_over_m(self, arr: Array, axis=0) -> Array:
        if not arr.shape[axis] == len(self.l_array):
            raise ValueError('Array length should be equal to the number of l in the expansion.')
        return np.repeat(arr, 2 * self.l_array + 1, axis=axis)

    def rotate(self, rotations: Quaternion, rotate_existing=False):
        if rotate_existing:
            raise NotImplementedError("Rotation of possibly already rotated coefficients is not yet supported.")
        self._rotations = rotations
        self.coefs = expansion_rotation(rotations, self._starting_coefs, self.l_array)

    def rotate_euler(self, alpha: Array = 0, beta: Array = 0, gamma: Array = 0, rotate_existing=False):
        R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
        self.rotate(R_euler, rotate_existing=rotate_existing)

    def inverse_sign(self, exclude_00: bool = False):
        if self.l_array[0] == 0 and exclude_00:
            self.coefs[..., 1:] = -self.coefs[..., 1:]
            self._starting_coefs[..., 1:] = -self._starting_coefs[..., 1:]
            return self
        self.coefs = -self.coefs
        self._starting_coefs = -self._starting_coefs
        return self

    def charge_value(self, theta: Array | float, phi: Array | float):
        if not isinstance(theta, Array):
            theta = np.array([theta])
        if not isinstance(phi, Array):
            phi = np.array([phi])
        theta, phi = np.broadcast_arrays(theta, phi)
        full_l_array, full_m_array = self.lm_arrays
        return np.squeeze(np.real(np.sum(self.coefs[..., None] * fn.sph_harm(full_l_array[:, None],
                                                                             full_m_array[:, None],
                                                                             theta[None, :], phi[None, :]), axis=-2)))

    def clone(self) -> Expansion:
        return copy.deepcopy(self)


def m_indices_at_l(l_arr: Array, l_idx: int):
    """
    For a given l_array and index l_idx for some ell in this array, get indices of all (ell, m) coefficients
    in coefficients array.
    """
    return np.arange(np.sum(2 * l_arr[:l_idx] + 1), np.sum(2 * l_arr[:l_idx+1] + 1))


def full_lm_arrays(l_array: Array) -> (Array, Array):
    """From an array of l_values get arrays of ell and m that give you all pairs (ell, m)."""
    all_m_list = []
    for l in l_array:
        for i in range(2 * l + 1):
            all_m_list.append(-l + i)
    return np.repeat(l_array, 2 * l_array + 1), np.array(all_m_list)


def coefs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expansion:
    """Explicitly add missing expansion coefficients so that expansion includes all ell values from the target array."""
    missing_l = np.setdiff1d(target_l_array, expansion.l_array, assume_unique=True)
    fill = np.zeros(np.sum(2 * missing_l + 1))
    full_l_array1, _ = expansion.lm_arrays
    # we search for where to place missing coefs with the help of a boolean array and argmax function
    comparison_bool = (full_l_array1[:, None] - missing_l[None, :]) > 0
    indices = np.where(np.any(comparison_bool, axis=0), np.argmax(comparison_bool, axis=0), full_l_array1.shape[0])
    new_coefs = np.insert(expansion.coefs, np.repeat(indices, 2 * missing_l + 1), fill, axis=-1)
    return Expansion(target_l_array, new_coefs)


def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
    """Explicitly add zero expansion coefficients so that both expansions include coefficients for the same ell."""
    common_l_array = np.union1d(ex1.l_array, ex2.l_array)
    return coefs_fill_missing_l(ex1, common_l_array), coefs_fill_missing_l(ex2, common_l_array)


def expansion_rotation(rotations: Quaternion, coefs: Array, l_array: Array):
    """
    General function for rotations of expansion coefficients using WignerD matrices. Combines all rotations
    with each expansion given in coefs array.
    :param rotations: Quaternion array, last dimension is 4
    :param coefs: array of expansion coefficients
    :param l_array: array of all ell values of the expansion
    :return rotated coefficients, output shape is rotations.shape[:-1] + coefs.shape
    """
    rot_arrays = rotations.ndarray.reshape((-1, 4))
    coefs_reshaped = coefs.reshape((-1, coefs.shape[-1]))
    wigner_matrices = spherical.Wigner(np.max(l_array)).D(rot_arrays)
    new_coefs = np.zeros((rot_arrays.shape[0],) + coefs_reshaped.shape, dtype=np.complex128)
    for i, l in enumerate(l_array):
        Dlmn_slice = np.arange(l * (2 * l - 1) * (2 * l + 1) / 3, (l + 1) * (2 * l + 1) * (2 * l + 3) / 3).astype(int)
        all_m_indices = m_indices_at_l(l_array, i)
        wm = wigner_matrices[:, Dlmn_slice].reshape((-1, 2*l+1, 2*l+1))
        new_coefs[..., all_m_indices] = np.einsum('rnm, qm -> rqn',
                                                   wm, coefs_reshaped[:, all_m_indices])
    return new_coefs.reshape(rotations.ndarray.shape[:-1] + coefs.shape)