charged-shells/charged_shells/expansion.py

from __future__ import annotations
import numpy as np
from dataclasses import dataclass
from functools import cached_property
import charged_shells.functions as fn
import quaternionic
import spherical
import copy
from scipy.special import eval_legendre


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 = Quaternion([1., 0., 0., 0.])

    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)

    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,))

    @cached_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):
        # TODO: rotations are currently saved wrong if we start form existing coefficients not the og ones
        self._rotations = rotations
        coefs = self.coefs if rotate_existing else self._starting_coefs
        self.coefs = expansion_rotation(rotations, coefs, self.l_array)

    def rotate_euler(self, alpha: Array, beta: Array, gamma: Array, rotate_existing=False):
        # TODO: additional care required on the convention used to transform euler angles to quaternions
        # TODO: might be off for a minus sign for each? angle !!
        R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
        self.rotate(R_euler, rotate_existing=rotate_existing)

    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)


class Expansion24(Expansion):

    def __init__(self, sigma2: float, sigma4: float, sigma0: float = 0.):
        l_array = np.array([0, 2, 4])
        coefs = rot_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
        super().__init__(l_array, coefs)


class MappedExpansionQuad(Expansion):
    """Expansion that matches the outside potential of a quadrupolar impermeable particle with point charges inside."""

    def __init__(self,
                 a_bar: Array | float,
                 kappaR: Array | float,
                 sigma_tilde: float,
                 l_max: int = 20,
                 sigma0: float | Array = 0):
        """
        :param a_bar: distance between the center and off center charges
        :param kappaR: screening parameter
        :param sigma_tilde: magnitude of off-center charges / 4pi R^2
        :param l_max: maximal ell value for the expansion
        :param sigma0: total (mean) charge density
        """
        a_bar, kappaR = np.broadcast_arrays(a_bar, kappaR)

        l_array = np.array([l for l in range(l_max + 1) if l % 2 == 0])
        a_bar, kappaR, l_array_expanded = np.broadcast_arrays(a_bar[..., None],
                                                              kappaR[..., None],
                                                              l_array[None, :])

        coefs = (2 * sigma_tilde * fn.coef_C_diff(l_array_expanded, kappaR)
                 * np.sqrt(4 * np.pi * (2 * l_array_expanded + 1)) * np.power(a_bar, l_array_expanded))
        coefs = np.squeeze(rot_sym_expansion(l_array, coefs))
        coefs = expansion_total_charge(coefs, sigma0)
        super().__init__(l_array, coefs)


class GaussianCharges(Expansion):
    """Expansion for a collection of smeared charges on the sphere."""

    def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int,
                 sigma0: float | Array = 0, equal_charges: bool = True):
        """
        :param omega_k: array of positions (theta, phi) of all charges
        :param lambda_k: smear parameter for each charge or smear for different cases (if equal_charges = True)
        :param sigma1: scaling
        :param l_max: maximal ell value for the expansion
        :param sigma0: total (mean) charge density
        :param equal_charges: if this is False, length of lambda_k should be N. If True, theta0_k array will be treated
            as different expansion cases
        """
        omega_k = omega_k.reshape(-1, 2)
        if not isinstance(lambda_k, Array):
            lambda_k = np.array([lambda_k])
        if equal_charges:
            if lambda_k.ndim > 1:
                raise ValueError(f'If equal_charges=True, lambda_k should be a 1D array, got shape {lambda_k.shape}')
            lambda_k = np.full((omega_k.shape[0], lambda_k.shape[0]), lambda_k).T
        if lambda_k.shape[-1] != omega_k.shape[0]:
            raise ValueError("Number of charges (length of omega_k) should match the last dimension of lambda_k array.")
        lambda_k = lambda_k.reshape(-1, omega_k.shape[0])
        l_array = np.arange(l_max + 1)
        full_l_array, full_m_array = full_lm_arrays(l_array)
        theta_k = omega_k[:, 0]
        phi_k = omega_k[:, 1]
        summands = (lambda_k[:, None, :] / np.sinh(lambda_k[:, None, :])
                    * fn.sph_bessel_i(full_l_array[None, :, None], lambda_k[:, None, :])
                    * np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
                                          theta_k[None, None, :], phi_k[None, None, :])))
        coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
        coefs = expansion_total_charge(coefs, sigma0)
        l_array, coefs = purge_unneeded_l(l_array, coefs)
        super().__init__(l_array, coefs)


class SphericalCap(Expansion):
    """Expansion for a collection of spherical caps."""

    def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float | Array = 0,
                 equal_sizes: bool = True):
        """
        :param omega_k: array of positions (theta, phi) of all spherical caps
        :param theta0_k: sizes of each spherical caps or cap sizes for different cases (if equal_sizes = True)
        :param sigma1: charge magnitude for the single cap, currently assumes that this is equal for all caps
        :param l_max: maximal ell value for the expansion
        :param sigma0: total (mean) charge density
        :param equal_sizes: if this is False, length of theta0_k should be N. If True, theta0_k array will be treated as
            different expansion cases
        """
        omega_k = omega_k.reshape(-1, 2)
        if not isinstance(theta0_k, Array):
            theta0_k = np.array(theta0_k)
        if equal_sizes:
            if theta0_k.ndim == 0:
                theta0_k = np.full(omega_k.shape[0], theta0_k)
            elif theta0_k.ndim == 1:
                theta0_k = np.full((omega_k.shape[0], theta0_k.shape[0]), theta0_k)
            else:
                raise ValueError(f'If equal_charges=True, theta0_k should be a 1D array, got shape {theta0_k.shape}')
        if theta0_k.shape[0] != omega_k.shape[0]:
            raise ValueError("Number of charges (length of omega_k) should match the last dimension of theta0_k array.")
        rotations = Quaternion(np.stack((np.cos(omega_k[..., 0] / 2),
                                         np.sin(omega_k[..., 1]) * np.sin(omega_k[..., 0] / 2),
                                         np.cos(omega_k[..., 1]) * np.sin(omega_k[..., 0] / 2),
                                         np.zeros_like(omega_k[..., 0]))).T)
        l_array = np.arange(l_max + 1)
        coefs_l0 = -sigma1 * (np.sqrt(np.pi / (2 * l_array[None, :] + 1)) *
                              (eval_legendre(l_array[None, :] + 1, np.cos(theta0_k[..., None]))
                               - eval_legendre(l_array[None, :] - 1, np.cos(theta0_k[..., None]))))
        coefs = rot_sym_expansion(l_array, coefs_l0)
        coefs_all_single_caps = expansion_rotation(rotations, coefs, l_array)
        # Rotating is implemented in such a way that it rotates every patch to every position,
        # hence the redundancy of out of diagonal elements.
        coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
        coefs_all = expansion_total_charge(coefs_all, sigma0)
        super().__init__(l_array, np.squeeze(coefs_all))


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 rot_sym_expansion(l_array: Array, coefs: Array) -> Array:
    """Create full expansion array for rotationally symmetric distributions with only m=0 terms different form 0."""
    full_coefs = np.zeros(coefs.shape[:-1] + (np.sum(2 * l_array + 1),))
    full_coefs[..., np.cumsum(2 * l_array + 1) - l_array - 1] = coefs
    return full_coefs


def expansion_total_charge(coefs: Array, sigma0: float | Array):
    """Adds a new axis to the expansion coefficients that modifies expansion based on given net charge density."""
    if sigma0 is None:
        return coefs
    if not isinstance(sigma0, Array):
        x = copy.deepcopy(coefs)
        x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
        return x
    sigma0 = sigma0.flatten()
    x = np.repeat(np.expand_dims(coefs, -2), len(sigma0), axis=-2)
    x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
    return x


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 purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
    """Remove ell values from expansion for which all (ell, m) coefficients are zero."""
    def delete_zero_entries(l, l_arr, cfs):
        l_idx = np.where(l_arr == l)[0][0]
        m_indices = m_indices_at_l(l_arr, l_idx)
        if np.all(cfs[..., m_indices] == 0):
            return np.delete(l_arr, l_idx), np.delete(cfs, m_indices, axis=-1)
        return l_arr, cfs
    for l in l_array:
        l_array, coefs = delete_zero_entries(l, l_array, coefs)
    return l_array, coefs


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)