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
from typing import Callable, TypeVar
from scipy.special import eval_legendre
import plotly.graph_objects as go
from pathlib import Path


Array = np.ndarray
Quaternion = quaternionic.array
T = TypeVar('T')
V = TypeVar('V')


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_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):
        # 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):

    def __init__(self, a_bar: Array | float, kappaR: Array | float, sigma_m: float, l_max: int = 20, sigma0: float = 0):
        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_m * 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[..., 0] = sigma0 / np.sqrt(4 * np.pi)
        coefs = np.squeeze(rot_sym_expansion(l_array, coefs))
        super().__init__(l_array, coefs)


class GaussianCharges(Expansion):

    def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int,
                 sigma0: float = 0, equal_charges: bool = True):
        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_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, :, 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[..., 0] = sigma0 / np.sqrt(4 * np.pi)
        l_array, coefs = purge_unneeded_l(l_array, coefs)
        super().__init__(l_array, coefs)


class SphericalCap(Expansion):

    def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float = 0,
                 equal_sizes: bool = True):
        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)
        if sigma0 is not None:
            coefs_all[..., 0] = sigma0 / np.sqrt(4 * np.pi)
        super().__init__(l_array, np.squeeze(coefs_all))


def map_over_expansion(f: Callable[[Expansion, T], V]) -> Callable[[Expansion, T], V]:
    """Map a function f over the leading axes of an expansion. Uses for loops, so it is kinda slow."""

    def mapped_f(ex: Expansion, *args, **kwargs):
        og_shape = ex.shape
        flat_ex = ex.flatten()
        results = []
        for i in range(int(np.prod(og_shape))):
            results.append(f(flat_ex[i], *args, **kwargs))
        try:
            return np.array(results).reshape(og_shape + results[0].shape)
        except AttributeError:
            return np.array(results).reshape(og_shape)

    return mapped_f


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, 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 coefs_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 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 m_indices_at_l(l_arr: Array, l_idx: int):
    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):
    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 plot_theta_profile(ex: Expansion, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
    theta_vals = np.linspace(theta_start, theta_end, num)
    charge = ex.charge_value(theta_vals, phi)
    plt.plot(theta_vals, charge.T)
    plt.show()


def plot_theta_profile_multiple(ex_list: list[Expansion], label_list, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
    theta_vals = np.linspace(theta_start, theta_end, num)

    fig, ax = plt.subplots()
    for ex, label in zip(ex_list, label_list):
        ax.plot(theta_vals, ex.charge_value(theta_vals, phi).T, label=label)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\theta$', fontsize=13)
    ax.set_ylabel(r'$\sigma$', fontsize=13)
    plt.legend(fontsize=12)
    plt.tight_layout()
    plt.savefig(Path("/home/andraz/ChargedShells/Figures/charge_shape_comparison.png"), dpi=600)
    plt.show()


def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100):
    theta = np.linspace(0, np.pi, num_theta)
    phi = np.linspace(0, 2 * np.pi, num_phi)
    theta, phi = np.meshgrid(theta, phi)
    r = ex.charge_value(theta.flatten(), phi.flatten()).reshape(theta.shape)

    # Convert spherical coordinates to Cartesian coordinates
    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)

    # Create a heatmap on the sphere
    fig = go.Figure(data=go.Surface(x=x, y=y, z=z, surfacecolor=r, colorscale='Jet'))
    fig.update_layout(scene=dict(aspectmode='data'))
    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))

    fig.show()


def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
    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)


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 = GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
    ex = SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.1, 70)
    # print(np.real(ex.coefs))
    # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
    plot_charge_3d(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)