charged-shells/expansion.py

from __future__ import annotations
import numpy as np
from dataclasses import dataclass
from functools import cached_property
import functions as fn
import quaternionic
import spherical
import time
import copy
import matplotlib.pyplot as plt


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)

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

    @cached_property
    def lm_arrays(self) -> (Array, Array):
        """Return l and m arrays containing all (l, m) pairs."""
        return full_fm_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):
        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.sum(self.coefs[None, :] * fn.sph_harm(full_l_array[None, :], full_m_array[None, :],
                                                        theta[:, None], phi[:, None]), axis=1)

    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])
        coeffs = rot_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
        super().__init__(l_array, coeffs)


class MappedExpansionQuad(Expansion):

    def __init__(self, a_bar: float, kappaR: float, sigma_m: float, max_l: int = 20, sigma0: float = 0):
        l_array = np.array([l for l in range(max_l + 1) if l % 2 == 0])
        coeffs = (2 * sigma_m * fn.coeff_C_diff(l_array, kappaR)
                  * np.sqrt(4 * np.pi * (2 * l_array + 1)) * np.power(a_bar, l_array))
        coeffs[0] = sigma0
        coeffs = rot_sym_expansion(l_array, coeffs)
        super().__init__(l_array, coeffs)


class SmearedCharges(Expansion):

    def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int, sigma0: float = 0):
        omega_k = omega_k.reshape(-1, 2)
        if not isinstance(lambda_k, Array):
            lambda_k = np.full((len(omega_k),), lambda_k)
        if lambda_k.shape[-1] != omega_k.shape[0]:
            raise ValueError("Omega and lambda arrays should have the same length.")
        l_array = np.arange(l_max + 1)
        full_l_array, full_m_array = full_fm_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], lambda_k[None, :])
                    * np.conj(fn.sph_harm(full_l_array[:, None], full_m_array[:, None],
                                          theta_k[None, :], phi_k[None, :])))
        coefs = 4 * np.pi * sigma1 * np.sum(summands, axis=1)
        coefs[0] = sigma0
        super().__init__(l_array, coefs)


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


def coeffs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expansion:
    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 coeffs 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_coeffs = np.insert(expansion.coefs, np.repeat(indices, 2 * missing_l + 1), fill, axis=-1)
    return Expansion(target_l_array, new_coeffs)


def plot_theta_profile(ex: Expansion, phi=0, num=100):
    theta_vals = np.linspace(0, np.pi, num)
    charge = ex.charge_value(theta_vals, phi)
    plt.plot(theta_vals, charge)
    plt.show()


def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
    common_l_array = np.union1d(ex1.l_array, ex2.l_array)
    return coeffs_fill_missing_l(ex1, common_l_array),  coeffs_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 = np.arange(np.sum(2 * l_array[:i] + 1), np.sum(2 * l_array[:i + 1] + 1))
        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)


if __name__ == '__main__':

    # ex = MappedExpansionQuad(0.44, 3, 1, 10)
    # ex = Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
    ex = SmearedCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
    # print(ex.coefs)
    plot_theta_profile(ex)

    # new_coeffs = expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
    # print(new_coeffs.shape)
    #
    # newnew_coeffs = expansion_rotation(Quaternion(np.arange(16).reshape(4, 4)).normalized, new_coeffs, ex.l_array)
    # print(newnew_coeffs.shape)

    # rot_angles = np.linspace(0, np.pi, 1000)
    # t0 = time.time()
    # ex.rotate_euler(0, np.pi / 2, -1)
    # t1 = time.time()
    # print(ex.coefs)
    # print(t1 - t0)