andraz-gnidovec 1 год назад
Сommit
2e36dd8c74
3 измененных файлов с 173 добавлено и 0 удалено
  1. 84 0
      expansion.py
  2. 30 0
      functions.py
  3. 59 0
      potentials.py

+ 84 - 0
expansion.py

@@ -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)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

+ 30 - 0
functions.py

@@ -0,0 +1,30 @@
+import numpy as np
+import scipy.special as sps
+
+
+Array = np.ndarray
+
+
+def sph_bessel_i(n, x, **kwargs):
+    return sps.spherical_in(n, x, **kwargs)
+
+
+def sph_bessel_k(n, x, **kwargs):
+    return 2 / np.pi * sps.spherical_kn(n, x, **kwargs)
+
+
+def sph_harm(l, m, theta, phi, **kwargs):
+    return sps.sph_harm(m, l, phi, theta, **kwargs)
+
+
+def coefficient_Cpm(l, x) -> Array:
+    return x * sps.kv(l + 1 / 2, x) * sps.iv(l + 1 / 2, x)
+
+
+def coeff_C_diff(l, x):
+    return 1 / (x * sps.iv(l + 1 / 2, x) * sps.kv(l + 3 / 2, x))
+
+
+if __name__ == '__main__':
+    print(sph_bessel_k(np.array([1, 2, 3])[:, None], np.array([0.1, 0.2, 0.5, 0.6])[None, :]))
+

+ 59 - 0
potentials.py

@@ -0,0 +1,59 @@
+import expansion
+import functions as fn
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+Array = np.ndarray
+
+
+def charged_shell_potential(theta: Array | float,
+                            phi: Array | float,
+                            dist: float,
+                            ex: expansion.Expansion,
+                            kappaR: float,
+                            R: float) -> Array:
+    """
+    Electrostatic potential around a charged shell with patches given by expansion over spherical harmonics.
+
+    :param theta: array of azimuthal angles
+    :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
+    """
+    if isinstance(theta, float):
+        theta = np.full_like(phi, theta)
+
+    if isinstance(phi, float):
+        phi = np.full_like(theta, phi)
+
+    if not theta.shape == phi.shape:
+        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))
+
+    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))
+
+
+if __name__ == '__main__':
+
+    kappaR = 3
+    R = 150
+    ex = expansion.MappedExpansion(1, 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)
+    # print(potential)
+
+    plt.plot(potential)
+    plt.show()