charged-shells/expansion.py

import numpy as np
from dataclasses import dataclass
from functools import cached_property
import functions as fn

Array = np.ndarray


class InvalidExpansion(Exception):
    pass


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

    l_array: Array
    coeffs: Array

    def __post_init__(self):
        if len(self.coeffs) != np.sum(2 * self.l_array + 1):
            raise InvalidExpansion('Number of expansion coefficients does not match the provided l_list.')

    @cached_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):
        if not len(arr) == 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)


def sph_sym_expansion(l_array: Array, coeffs: Array):
    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


class Expansion24(Expansion):

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


class MappedExpansion(Expansion):

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


if __name__ == '__main__':

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