|
@@ -0,0 +1,84 @@
|
|
|
+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)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|