2023-9-24

expansion.py

@@ -24,7 +24,7 @@ class Expansion:
         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
+    @property
     def max_l(self) -> int:
         return max(self.l_array)
 
@@ -43,8 +43,8 @@ class Expansion:
         return np.repeat(arr, 2 * self.l_array + 1, axis=axis)
 
 
-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."""
+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
@@ -54,7 +54,7 @@ class Expansion24(Expansion):
 
     def __init__(self, sigma2: float, sigma4: float, sigma0: float = 0.):
         l_array = np.array([0, 2, 4])
-        coeffs = sph_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
+        coeffs = rot_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
         super().__init__(l_array, coeffs)
 
 
@@ -65,7 +65,7 @@ class MappedExpansion(Expansion):
         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 = sph_sym_expansion(l_array, coeffs)
+        coeffs = rot_sym_expansion(l_array, coeffs)
         super().__init__(l_array, coeffs)
 
 

interactions.py

@@ -1,19 +1,23 @@
 import expansion
+import parameters
 import functions as fn
 import numpy as np
+from py3nj import wigner3j
+import time
 import matplotlib.pyplot as plt
 
 
 Array = np.ndarray
 Expansion = expansion.Expansion
+ModelParams = parameters.ModelParams
 
 
-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 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):
@@ -31,61 +35,78 @@ def expansions_to_common_l(ex1: Expansion, ex2: expansion) -> (Expansion, Expans
     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)
+    # if all elements of bool are false, i.e. missing_l > all existing_l, additional l values come to the end
+    indices1 = np.where(np.any(bool1, axis=0), np.argmax(bool1, axis=0), full_l_array1.shape[0])
+    indices2 = np.where(np.any(bool2, axis=0), np.argmax(bool2, axis=0), full_l_array2.shape[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):
+def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: ModelParams):
 
     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)
+    coefficient_C = fn.interaction_coeff_C(ex1.l_array[:, None], ex2.l_array[None, :], params.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, :])
+    charge_factor = np.real(ex1.coeffs[:, None] * np.conj(ex2.coeffs[None, :]) +
+                            (-1) ** (full_l_array[:, None] + full_l_array[None, :])
+                            * ex1.coeffs[None, :] * ex1.coeffs[:, None])
 
-    flat_l = full_l_array[indices]
-    flat_m = full_m_array[indices]
+    indices_l, indices_p = np.nonzero(full_m_array[:, None] == full_m_array[None, :])
+    flat_l = full_l_array[indices_l]
+    flat_p = full_l_array[indices_p]
+    flat_m = full_m_array[indices_l]  # the same as full_m_array[indices_p]
 
-    flat_C = full_coefficient_C[indices1, indices2]
+    all_s_array = np.arange(2 * ex1.max_l + 1)
+    bessels = fn.sph_bessel_k(all_s_array, params.kappa * dist)
 
-    flat_sigma1 = ex1.coeffs[indices1]
-    flat_sigma2 = ex2.coeffs[indices2]
+    s_bool1 = np.abs(flat_l[:, None] - all_s_array[None, :]) <= flat_p[:, None]
+    s_bool2 = flat_p[:, None] <= (flat_l[:, None] + all_s_array[None, :])
+    indices_lpm, indices_s = np.nonzero(s_bool1 * s_bool2)
 
-    # charge_factor = -1 ** (flat_l + flat_m) * np.real(flat_sigma1 * np.conj(flat_sigma2) + (-1) ** (flat_l + flat_p) * )
+    l_vals = flat_l[indices_lpm]
+    p_vals = flat_p[indices_lpm]
+    m_vals = flat_m[indices_lpm]
+    C_vals = full_coefficient_C[indices_l, indices_p][indices_lpm]
+    charge_vals = charge_factor[indices_l, indices_p][indices_lpm]
+    s_vals = all_s_array[indices_s]
+    bessel_vals = bessels[indices_s]
 
-    return
+    lps_terms = (2 * s_vals + 1) * np.sqrt((2 * l_vals + 1) * (2 * p_vals + 1))
 
+    wigner1 = wigner3j(l_vals, s_vals, p_vals,
+                       0, 0, 0)
+    wigner2 = wigner3j(l_vals, s_vals, p_vals,
+                       -m_vals, 0, m_vals)
+
+    return (params.R ** 2 / (params.kappa * params.epsilon * params.epsilon0)
+            * np.sum(C_vals * (-1) ** (l_vals + m_vals) * charge_vals * lps_terms * bessel_vals * wigner1 * wigner2))
 
 
 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)
+    params = ModelParams(1, 3, 1, 1)
+    ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=5)
+    ex2 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=5)
 
     dist = 2.
 
-    ex1, ex2 = expansions_to_common_l(ex1, ex2)
-    print(ex1.coeffs)
-    print(ex2.coeffs)
+    # ex1, ex2 = expansions_to_common_l(ex1, ex2)
+    # print(ex1.coeffs)
+    # print(ex2.coeffs)
+
+    t0 = time.perf_counter()
+    energy = charged_shell_energy(ex1, ex2, dist, params)
+    t1 = time.perf_counter()
 
-    # energy = charged_shell_energy(ex1, ex2, dist, kappaR, R)
-    # print(potential)
+    print('energy: ', energy)
+    print('time: ', t1 - t0)
 
     # plt.plot(energy)
     # plt.show()

parameters.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelParams:
+    R: float
+    kappa: float
+    epsilon: float
+    epsilon0: float
+
+    @property
+    def kappaR(self):
+        return self.kappa * self.R

potentials.py

@@ -1,18 +1,20 @@
 import expansion
 import functions as fn
 import numpy as np
+import parameters
 import matplotlib.pyplot as plt
 
 
 Array = np.ndarray
+ModelParams = parameters.ModelParams
+Expansion = expansion.Expansion
 
 
 def charged_shell_potential(theta: Array | float,
                             phi: Array | float,
                             dist: float,
-                            ex: expansion.Expansion,
-                            kappaR: float,
-                            R: float) -> Array:
+                            ex: Expansion,
+                            params: ModelParams) -> Array:
     """
     Electrostatic potential around a charged shell with patches given by expansion over spherical harmonics.
 
@@ -20,9 +22,7 @@ def charged_shell_potential(theta: Array | float,
     :param phi: array of polar angles
     :param dist: distance between the particles in units of radius R
     :param ex: Expansion object detailing patch distribution
-    :param kappaR: screening value
-    :param R: particle radius in nm
-    :return: potential values in units of 1 / epsilon * epsilon0
+    :param params: ModelParams object specifying parameter values for the model
     """
     if isinstance(theta, float):
         theta = np.full_like(phi, theta)
@@ -34,25 +34,24 @@ def charged_shell_potential(theta: Array | float,
         raise ValueError('theta and phi arrays should have the same shape.')
     l_array, m_array = ex.lm_arrays
 
-    l_factors = (fn.coefficient_Cpm(ex.l_array, kappaR) * fn.sph_bessel_k(ex.l_array, kappaR * dist)
-                 / fn.sph_bessel_k(ex.l_array, kappaR))
+    l_factors = (fn.coefficient_Cpm(ex.l_array, params.kappaR) * fn.sph_bessel_k(ex.l_array, params.kappa * dist)
+                 / fn.sph_bessel_k(ex.l_array, params.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))
+    return (1 / (params.kappa * params.epsilon * params.epsilon0)
+            * np.real(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)))
 
 
 if __name__ == '__main__':
 
-    kappaR = 3
-    R = 150
-    ex = expansion.MappedExpansion(1, kappaR, 0.001, max_l=20)
+    params = ModelParams(1, 3, 1, 1)
+    ex = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=20)
 
     theta = np.linspace(0, np.pi, 1000)
     phi = 0.
     dist = 1.
 
-    potential = charged_shell_potential(theta, phi, dist, ex, kappaR, R)
+    potential = charged_shell_potential(theta, phi, dist, ex, params)
     # print(potential)
 
     plt.plot(potential)