2023-10-02

expansion.py

@@ -7,6 +7,7 @@ import quaternionic
 import spherical
 import time
 import copy
+import matplotlib.pyplot as plt
 
 
 Array = np.ndarray
@@ -41,11 +42,7 @@ class Expansion:
     @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)
+        return full_fm_arrays(self.l_array)
 
     def repeat_over_m(self, arr: Array, axis=0) -> Array:
         if not arr.shape[axis] == len(self.l_array):
@@ -58,9 +55,21 @@ class Expansion:
         self.coefs = expansion_rotation(rotations, coefs, self.l_array)
 
     def rotate_euler(self, alpha: Array, beta: Array, gamma: Array, rotate_existing=False):
+        # TODO: additional care required on the convention used to transform euler angles to quaternions
+        # TODO: might be off for a minus sign for each? angle !!
         R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
         self.rotate(R_euler, rotate_existing=rotate_existing)
 
+    def charge_value(self, theta: Array | float, phi: Array | float):
+        if not isinstance(theta, Array):
+            theta = np.array([theta])
+        if not isinstance(phi, Array):
+            phi = np.array([phi])
+        theta, phi = np.broadcast_arrays(theta, phi)
+        full_l_array, full_m_array = self.lm_arrays
+        return np.sum(self.coefs[None, :] * fn.sph_harm(full_l_array[None, :], full_m_array[None, :],
+                                                        theta[:, None], phi[:, None]), axis=1)
+
     def clone(self) -> Expansion:
         return copy.deepcopy(self)
 
@@ -73,7 +82,7 @@ class Expansion24(Expansion):
         super().__init__(l_array, coeffs)
 
 
-class MappedExpansion(Expansion):
+class MappedExpansionQuad(Expansion):
 
     def __init__(self, a_bar: float, kappaR: float, sigma_m: float, max_l: int = 20, sigma0: float = 0):
         l_array = np.array([l for l in range(max_l + 1) if l % 2 == 0])
@@ -81,10 +90,38 @@ class MappedExpansion(Expansion):
                   * np.sqrt(4 * np.pi * (2 * l_array + 1)) * np.power(a_bar, l_array))
         coeffs[0] = sigma0
         coeffs = rot_sym_expansion(l_array, coeffs)
-        # coeffs = np.full((1000, len(coeffs)), coeffs)
         super().__init__(l_array, coeffs)
 
 
+class SmearedCharges(Expansion):
+
+    def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int, sigma0: float = 0):
+        omega_k = omega_k.reshape(-1, 2)
+        if not isinstance(lambda_k, Array):
+            lambda_k = np.full((len(omega_k),), lambda_k)
+        if lambda_k.shape[-1] != omega_k.shape[0]:
+            raise ValueError("Omega and lambda arrays should have the same length.")
+        l_array = np.arange(l_max + 1)
+        full_l_array, full_m_array = full_fm_arrays(l_array)
+        theta_k = omega_k[:, 0]
+        phi_k = omega_k[:, 1]
+        summands = (lambda_k[None, :] / np.sinh(lambda_k[None, :])
+                    * fn.sph_bessel_i(full_l_array[:, None], lambda_k[None, :])
+                    * np.conj(fn.sph_harm(full_l_array[:, None], full_m_array[:, None],
+                                          theta_k[None, :], phi_k[None, :])))
+        coefs = 4 * np.pi * sigma1 * np.sum(summands, axis=1)
+        coefs[0] = sigma0
+        super().__init__(l_array, coefs)
+
+
+def full_fm_arrays(l_array: Array) -> (Array, Array):
+    all_m_list = []
+    for l in l_array:
+        for i in range(2 * l + 1):
+            all_m_list.append(-l + i)
+    return np.repeat(l_array, 2 * l_array + 1), np.array(all_m_list)
+
+
 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))
@@ -103,6 +140,13 @@ def coeffs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expans
     return Expansion(target_l_array, new_coeffs)
 
 
+def plot_theta_profile(ex: Expansion, phi=0, num=100):
+    theta_vals = np.linspace(0, np.pi, num)
+    charge = ex.charge_value(theta_vals, phi)
+    plt.plot(theta_vals, charge)
+    plt.show()
+
+
 def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
     common_l_array = np.union1d(ex1.l_array, ex2.l_array)
     return coeffs_fill_missing_l(ex1, common_l_array),  coeffs_fill_missing_l(ex2, common_l_array)
@@ -125,15 +169,18 @@ def expansion_rotation(rotations: Quaternion, coefs: Array, l_array: Array):
         Dlmn_slice = np.arange(l * (2 * l - 1) * (2 * l + 1) / 3, (l + 1) * (2 * l + 1) * (2 * l + 3) / 3).astype(int)
         all_m_indices = np.arange(np.sum(2 * l_array[:i] + 1), np.sum(2 * l_array[:i + 1] + 1))
         wm = wigner_matrices[:, Dlmn_slice].reshape((-1, 2*l+1, 2*l+1))
-        new_coefs[..., all_m_indices] = np.einsum('rmn, qm -> rqn',
+        new_coefs[..., all_m_indices] = np.einsum('rnm, qm -> rqn',
                                                    wm, coefs_reshaped[:, all_m_indices])
     return new_coefs.reshape(rotations.ndarray.shape[:-1] + coefs.shape)
 
 
 if __name__ == '__main__':
 
-    ex = MappedExpansion(0.44, 3, 1, 10)
-    print(ex.coefs)
+    # ex = MappedExpansionQuad(0.44, 3, 1, 10)
+    # ex = Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
+    ex = SmearedCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
+    # print(ex.coefs)
+    plot_theta_profile(ex)
 
     # new_coeffs = expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
     # print(new_coeffs.shape)
@@ -141,12 +188,12 @@ if __name__ == '__main__':
     # newnew_coeffs = expansion_rotation(Quaternion(np.arange(16).reshape(4, 4)).normalized, new_coeffs, ex.l_array)
     # print(newnew_coeffs.shape)
 
-    rot_angles = np.linspace(0, np.pi, 1000)
-    t0 = time.time()
-    ex.rotate_euler(rot_angles, rot_angles, rot_angles)
-    t1 = time.time()
-    print(ex.coefs.shape)
-    print(t1 - t0)
+    # rot_angles = np.linspace(0, np.pi, 1000)
+    # t0 = time.time()
+    # ex.rotate_euler(0, np.pi / 2, -1)
+    # t1 = time.time()
+    # print(ex.coefs)
+    # print(t1 - t0)
 
 
 

interactions.py

@@ -82,7 +82,7 @@ def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: Mo
 if __name__ == '__main__':
 
     params = ModelParams(R=150, kappaR=3)
-    ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=20)
+    ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, max_l=20)
     ex2 = ex1.clone()
 
     dist = 2.

parameters.py

@@ -26,6 +26,7 @@ def kappa_from_concentration(c0: float, temp: float, epsilon: float) -> float:
 
 
 def concentration_from_kappa(kappa: float, temp: float, epsilon: float):
+    # TODO: what is with the units of c0 calculated here?
     return kappa ** 2 / (8 * np.pi * bjerrum(temp, epsilon))
 
 

potentials.py

@@ -47,11 +47,11 @@ def charged_shell_potential(theta: Array | float,
 if __name__ == '__main__':
 
     params = ModelParams(R=150, kappaR=3)
-    ex = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=10)
+    ex = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, max_l=10)
 
     theta = np.linspace(0, np.pi, 1000)
     phi = 0.
-    dist = 2
+    dist = 1
 
     potential = charged_shell_potential(theta, phi, dist, ex, params)
     # print(potential)

rotational_path.py

@@ -18,21 +18,22 @@ class PairRotationalPath:
     rotations2: list[Quaternion] = field(default_factory=list)
 
     def add(self, rotation1: Quaternion, rotation2: Quaternion):
+        rotation1, rotation2 = np.broadcast_arrays(rotation1, rotation2)
         self.rotations1.append(rotation1)
         self.rotations2.append(rotation2)
 
-    def add_euler(self, alpha1: Array = 0, beta1: Array = 0, gamma1: Array = 0,
+    def add_euler(self, *, alpha1: Array = 0, beta1: Array = 0, gamma1: Array = 0,
                   alpha2: Array = 0, beta2: Array = 0, gamma2: Array = 0):
         R1_euler = quaternionic.array.from_euler_angles(alpha1, beta1, gamma1)
         R2_euler = quaternionic.array.from_euler_angles(alpha2, beta2, gamma2)
-        R1_euler, R2_euler = np.broadcast_arrays(R1_euler, R2_euler)
         self.add(Quaternion(R1_euler), Quaternion(R2_euler))
 
     def get_rotations(self) -> (Quaternion, Quaternion):
         return Quaternion(np.vstack(self.rotations1)), Quaternion(np.vstack(self.rotations2))
 
 
-zero_to_pi_half = np.linspace(0, np.pi/2, 1000)
+zero_to_pi_half = np.linspace(0, np.pi/2, 100, endpoint=False)
+pi_half_to_pi = np.linspace(np.pi/2, np.pi, 100, endpoint=False)
 
 QuadPath = PairRotationalPath()
 QuadPath.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half)
@@ -40,21 +41,36 @@ QuadPath.add_euler(beta1=zero_to_pi_half[::-1], beta2=zero_to_pi_half[::-1])
 QuadPath.add_euler(beta1=zero_to_pi_half)
 QuadPath.add_euler(beta1=zero_to_pi_half[::-1], beta2=zero_to_pi_half)
 QuadPath.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half, alpha2=np.pi/2)
-QuadPath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha2=zero_to_pi_half[::-1])
+QuadPath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half[::-1])
+
+DipolePath = PairRotationalPath()
+DipolePath.add_euler(beta1=pi_half_to_pi[::-1])
+DipolePath.add_euler(beta1=zero_to_pi_half[::-1])
+DipolePath.add_euler(beta1=zero_to_pi_half, beta2=zero_to_pi_half)
+DipolePath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half)
+DipolePath.add_euler(beta1=np.pi/2, alpha1=np.pi/2, beta2=zero_to_pi_half)
+DipolePath.add_euler(beta1=np.pi/2, beta2=pi_half_to_pi[::-1], alpha2=np.pi)
+DipolePath.add_euler(beta1=zero_to_pi_half[::-1], beta2=pi_half_to_pi, alpha2=np.pi)
+DipolePath.add_euler(beta1=zero_to_pi_half, beta2=pi_half_to_pi[::-1], alpha2=np.pi)
+DipolePath.add_euler(beta1=pi_half_to_pi, beta2=zero_to_pi_half[::-1], alpha2=np.pi)
 
 
 if __name__ == '__main__':
 
     params = ModelParams(R=150, kappaR=3)
-    ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=15)
+    # ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, max_l=20)
     # ex1 = expansion.Expansion24(sigma2=0.001, sigma4=0)
+    # ex1 = expansion.Expansion(l_array=np.array([1]), coefs=expansion.rot_sym_expansion(np.array([1]), np.array([0.001])))
+    # ex1 = expansion.SmearedCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=5, sigma1=0.001, l_max=10)
+    ex1 = expansion.SmearedCharges(omega_k=np.array([np.pi, 0]), lambda_k=3, sigma1=0.001, l_max=10)
     ex2 = ex1.clone()
 
-    rotations1, rotations2 = QuadPath.get_rotations()
+    # rotations1, rotations2 = QuadPath.get_rotations()
+    rotations1, rotations2 = DipolePath.get_rotations()
     ex1.rotate(rotations1)
     ex2.rotate(rotations2)
 
-    dist = 3
+    dist = 2
 
     energy = interactions.charged_shell_energy(ex1, ex2, dist, params)