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


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."""
        all_m_list = []
        for l in self.l_array:
            for i in range(2 * l + 1):
                all_m_list.append(-l + i)
        return np.repeat(self.l_array, 2 * self.l_array + 1), np.array(all_m_list)

    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):
        R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
        self.rotate(R_euler, rotate_existing=rotate_existing)

    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 MappedExpansion(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)
        # coeffs = np.full((1000, len(coeffs)), coeffs)
        super().__init__(l_array, coeffs)


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 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('rmn, qm -> rqn',
                                                   wm, coefs_reshaped[:, all_m_indices])
    return new_coefs.reshape(rotations.ndarray.shape[:-1] + coefs.shape)


if __name__ == '__main__':

    ex = MappedExpansion(0.44, 3, 1, 10)
    print(ex.coefs)

    # 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(rot_angles, rot_angles, rot_angles)
    t1 = time.time()
    print(ex.coefs.shape)
    print(t1 - t0)