2023-9-22

expansion.py

@@ -18,8 +18,11 @@ class Expansion:
     coeffs: Array
 
     def __post_init__(self):
+        """Validation of the given expansion."""
         if len(self.coeffs) != np.sum(2 * self.l_array + 1):
-            raise InvalidExpansion('Number of expansion coefficients does not match the provided l_list.')
+            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.')
 
     @cached_property
     def max_l(self) -> int:
@@ -34,13 +37,14 @@ class Expansion:
                 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):
+    def repeat_over_m(self, arr: Array, axis=0):
+        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)
+        return np.repeat(arr, 2 * self.l_array + 1, axis=axis)
 
 
-def sph_sym_expansion(l_array: Array, coeffs: Array):
+def sph_sym_expansion(l_array: Array, coeffs: Array) -> Array:
+    """Create full expansion array for spherically 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
@@ -49,18 +53,20 @@ def sph_sym_expansion(l_array: Array, coeffs: Array):
 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]))
+        l_array = np.array([0, 2, 4])
+        coeffs = sph_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):
-        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))
+        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
-        self.coeffs = sph_sym_expansion(self.l_array, coeffs)
+        coeffs = sph_sym_expansion(l_array, coeffs)
+        super().__init__(l_array, coeffs)
 
 
 if __name__ == '__main__':

functions.py

@@ -17,6 +17,10 @@ def sph_harm(l, m, theta, phi, **kwargs):
     return sps.sph_harm(m, l, phi, theta, **kwargs)
 
 
+def interaction_coeff_C(l, p, x):
+    return x * sps.iv(l + 1 / 2, x) * sps.iv(p + 1 / 2, x)
+
+
 def coefficient_Cpm(l, x) -> Array:
     return x * sps.kv(l + 1 / 2, x) * sps.iv(l + 1 / 2, x)
 

interactions.py

@@ -0,0 +1,91 @@
+import expansion
+import functions as fn
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+Array = np.ndarray
+Expansion = expansion.Expansion
+
+
+def prefactor(R: float, kappaR: float, c0: float):
+    return R * kappaR * 1e4 / (12.04 * c0)
+
+
+def c0(R: float, kappaR: float):
+    return 10 * kappaR ** 2 / (0.329 ** 2 * R ** 2)
+
+
+def expansions_to_common_l(ex1: Expansion, ex2: expansion) -> (Expansion, Expansion):
+    common_l_array = np.union1d(ex1.l_array, ex2.l_array)
+    missing_l1 = np.setdiff1d(common_l_array, ex1.l_array, assume_unique=True)
+    missing_l2 = np.setdiff1d(common_l_array, ex2.l_array, assume_unique=True)
+
+    fill_1 = np.zeros(np.sum(2 * missing_l1 + 1))
+    fill_2 = np.zeros(np.sum(2 * missing_l2 + 1))
+
+    full_l_array1, _ = ex1.lm_arrays
+    full_l_array2, _ = ex2.lm_arrays
+
+    # we search for where to place missing coeffs with the help of a boolean array and argmax function
+    bool1 = (full_l_array1[:, None] - missing_l1[None, :]) > 0
+    bool2 = (full_l_array2[:, None] - missing_l2[None, :]) > 0
+
+    # we set last element to True so that argmax returns last idx if all missing l > max_l
+    bool1[-1, :] = True
+    bool2[-1, :] = True
+
+    indices1 = np.argmax(bool1, axis=0)
+    indices2 = np.argmax(bool2, axis=0)
+
+    new_coeffs1 = np.insert(ex1.coeffs, np.repeat(indices1, 2 * missing_l1 + 1), fill_1)
+    new_coeffs2 = np.insert(ex2.coeffs, np.repeat(indices2, 2 * missing_l2 + 1), fill_2)
+
+    assert len(new_coeffs1) == len(new_coeffs2)
+
+    return Expansion(common_l_array, new_coeffs1), Expansion(common_l_array, new_coeffs2)
+
+
+def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, kappaR: float, R: float):
+
+    ex1, ex2 = expansions_to_common_l(ex1, ex2)
+
+    full_l_array, full_m_array = ex1.lm_arrays
+
+    coefficient_C = fn.interaction_coeff_C(ex1.l_array[:, None], ex2.l_array[None, :], kappaR)
+    full_coefficient_C = ex1.repeat_over_m(ex2.repeat_over_m(coefficient_C, axis=1), axis=0)
+
+    indices, _ = np.nonzero(full_m_array[:, None] == full_m_array[None, :])
+
+    flat_l = full_l_array[indices]
+    flat_m = full_m_array[indices]
+
+    flat_C = full_coefficient_C[indices1, indices2]
+
+    flat_sigma1 = ex1.coeffs[indices1]
+    flat_sigma2 = ex2.coeffs[indices2]
+
+    # charge_factor = -1 ** (flat_l + flat_m) * np.real(flat_sigma1 * np.conj(flat_sigma2) + (-1) ** (flat_l + flat_p) * )
+
+    return
+
+
+
+if __name__ == '__main__':
+
+    kappaR = 3
+    R = 150
+    ex1 = expansion.MappedExpansion(1, kappaR, 0.001, max_l=10)
+    ex2 = expansion.MappedExpansion(1, kappaR, 0.001, max_l=5)
+
+    dist = 2.
+
+    ex1, ex2 = expansions_to_common_l(ex1, ex2)
+    print(ex1.coeffs)
+    print(ex2.coeffs)
+
+    # energy = charged_shell_energy(ex1, ex2, dist, kappaR, R)
+    # print(potential)
+
+    # plt.plot(energy)
+    # plt.show()

potentials.py

@@ -38,8 +38,8 @@ def charged_shell_potential(theta: Array | float,
                  / fn.sph_bessel_k(ex.l_array, kappaR))
 
     return np.real(R / kappaR * np.sum(ex.repeat_over_m(l_factors)[None, :] * ex.coeffs
-                               * fn.sph_harm(l_array[None, :], m_array[None, :], theta[:, None], phi[:, None]),
-                               axis=1))
+                                       * fn.sph_harm(l_array[None, :], m_array[None, :], theta[:, None], phi[:, None]),
+                                       axis=1))
 
 
 if __name__ == '__main__':