2023-10-30

expansion.py

@@ -9,6 +9,8 @@ import time
 import copy
 import matplotlib.pyplot as plt
 from typing import Callable, TypeVar
+from scipy.special import eval_legendre
+import plotly.graph_objects as go
 
 
 Array = np.ndarray
@@ -114,7 +116,7 @@ class MappedExpansionQuad(Expansion):
 
         coefs = (2 * sigma_m * fn.coef_C_diff(l_array_expanded, kappaR)
                   * np.sqrt(4 * np.pi * (2 * l_array_expanded + 1)) * np.power(a_bar, l_array_expanded))
-        coefs[..., 0] = sigma0
+        coefs[..., 0] = sigma0 / np.sqrt(4 * np.pi)
         coefs = rot_sym_expansion(l_array, coefs)
         super().__init__(l_array, coefs)
 
@@ -122,7 +124,7 @@ class MappedExpansionQuad(Expansion):
 class GaussianCharges(Expansion):
 
     def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int,
-                 sigma0: float = 0, equal_charges=True):
+                 sigma0: float = 0, equal_charges: bool = True):
         omega_k = omega_k.reshape(-1, 2)
         if not isinstance(lambda_k, Array):
             lambda_k = np.array([lambda_k])
@@ -142,11 +144,45 @@ class GaussianCharges(Expansion):
                     * np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
                                           theta_k[None, None, :], phi_k[None, None, :])))
         coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
-        coefs[..., 0] = sigma0
+        coefs[..., 0] = sigma0 / np.sqrt(4 * np.pi)
         l_array, coefs = purge_unneeded_l(l_array, coefs)
         super().__init__(l_array, coefs)
 
 
+class SphericalCap(Expansion):
+
+    def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float = 0,
+                 equal_sizes: bool = True):
+        omega_k = omega_k.reshape(-1, 2)
+        if not isinstance(theta0_k, Array):
+            theta0_k = np.array(theta0_k)
+        if equal_sizes:
+            if theta0_k.ndim == 0:
+                theta0_k = np.full(omega_k.shape[0], theta0_k)
+            elif theta0_k.ndim == 1:
+                theta0_k = np.full((omega_k.shape[0], theta0_k.shape[0]), theta0_k)
+            else:
+                raise ValueError(f'If equal_charges=True, theta0_k should be a 1D array, got shape {theta0_k.shape}')
+        if theta0_k.shape[0] != omega_k.shape[0]:
+            raise ValueError("Number of charges (length of omega_k) should match the last dimension of theta0_k array.")
+        rotations = Quaternion(np.stack((np.cos(omega_k[..., 0] / 2),
+                                         np.sin(omega_k[..., 1]) * np.sin(omega_k[..., 0] / 2),
+                                         np.cos(omega_k[..., 1]) * np.sin(omega_k[..., 0] / 2),
+                                         np.zeros_like(omega_k[..., 0]))).T)
+        l_array = np.arange(l_max + 1)
+        coefs_l0 = -sigma1 * (np.sqrt(np.pi / (2 * l_array[None, :] + 1)) *
+                              (eval_legendre(l_array[None, :] + 1, np.cos(theta0_k[..., None]))
+                               - eval_legendre(l_array[None, :] - 1, np.cos(theta0_k[..., None]))))
+        coefs = rot_sym_expansion(l_array, coefs_l0)
+        coefs_all_single_caps = expansion_rotation(rotations, coefs, l_array)
+        # Rotating is implemented in such a way that it rotates every patch to every position,
+        # hence the redundancy of out of diagonal elements.
+        coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
+        if sigma0 is not None:
+            coefs_all[..., 0] = sigma0 / np.sqrt(4 * np.pi)
+        super().__init__(l_array, np.squeeze(coefs_all))
+
+
 def map_over_expansion(f: Callable[[Expansion, T], V]) -> Callable[[Expansion, T], V]:
     """Map a function f over the leading axes of an expansion. Uses for loops, so it is kinda slow."""
 
@@ -154,7 +190,7 @@ def map_over_expansion(f: Callable[[Expansion, T], V]) -> Callable[[Expansion, T
         og_shape = ex.shape
         flat_ex = ex.flatten()
         results = []
-        for i in range(np.prod(og_shape)):
+        for i in range(int(np.prod(og_shape))):
             results.append(f(flat_ex[i], *args, **kwargs))
         try:
             return np.array(results).reshape(og_shape + results[0].shape)
@@ -206,13 +242,32 @@ def purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
     return l_array, coefs
 
 
-def plot_theta_profile(ex: Expansion, phi=0, num=100):
-    theta_vals = np.linspace(0, np.pi, num)
+def plot_theta_profile(ex: Expansion, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
+    theta_vals = np.linspace(theta_start, theta_end, num)
     charge = ex.charge_value(theta_vals, phi)
-    plt.plot(theta_vals, charge)
+    plt.plot(theta_vals, charge.T)
     plt.show()
 
 
+def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100):
+    theta = np.linspace(0, np.pi, num_theta)
+    phi = np.linspace(0, 2 * np.pi, num_phi)
+    theta, phi = np.meshgrid(theta, phi)
+    r = ex.charge_value(theta.flatten(), phi.flatten()).reshape(theta.shape)
+
+    # Convert spherical coordinates to Cartesian coordinates
+    x = np.sin(theta) * np.cos(phi)
+    y = np.sin(theta) * np.sin(phi)
+    z = np.cos(theta)
+
+    # Create a heatmap on the sphere
+    fig = go.Figure(data=go.Surface(x=x, y=y, z=z, surfacecolor=r, colorscale='Jet'))
+    fig.update_layout(scene=dict(aspectmode='data'))
+    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
+
+    fig.show()
+
+
 def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
     common_l_array = np.union1d(ex1.l_array, ex2.l_array)
     return coefs_fill_missing_l(ex1, common_l_array),  coefs_fill_missing_l(ex2, common_l_array)
@@ -244,9 +299,11 @@ if __name__ == '__main__':
 
     # ex = MappedExpansionQuad(0.44, 3, 1, 10)
     # ex = Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
-    ex = GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
+    # ex = GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
+    ex = SphericalCap(np.array([[0, 0], [np.pi / 2, 0]]), 0.5, 0.1, 30)
     # print(ex.coefs)
-    plot_theta_profile(ex)
+    # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
+    plot_charge_3d(ex)
 
     # new_coeffs = expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
     # print(new_coeffs.shape)

patch_size.py

@@ -4,6 +4,7 @@ import expansion
 from matplotlib import pyplot as plt
 import parameters
 import potentials
+from pathlib import Path
 
 Expansion = expansion.Expansion
 Array = np.ndarray
@@ -46,7 +47,7 @@ def potential_patch_size(ex: Expansion, params: ModelParams,
 
     if meap is not None:
         return np.array(results).swapaxes(0, meap)  # return the params-matched axis to where it belongs
-    return np.array(results)
+    return np.squeeze(np.array(results))
 
 
 if __name__ == '__main__':
@@ -54,12 +55,18 @@ if __name__ == '__main__':
     a_bar = np.linspace(0.2, 0.8, 100)
     kappaR = np.array([0.26, 1, 3, 10, 26])
     params = ModelParams(R=150, kappaR=kappaR)
-    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_m=0.001, l_max=20, kappaR=kappaR[None, :])
+    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_m=0.001, l_max=30, kappaR=kappaR[None, :])
 
     # patch_size = charge_patch_size(ex)
-    patch_size_pot = potential_patch_size(ex, params, match_expansion_axis_to_params=1)
-
-    plt.plot(a_bar, patch_size_pot * 180 / np.pi, label=kappaR)
-    plt.legend()
+    patch_size = potential_patch_size(ex, params, match_expansion_axis_to_params=1)
+
+    fig, ax = plt.subplots()
+    for patch, kR in zip(patch_size.T, kappaR):
+        ax.plot(a_bar, patch * 180 / np.pi, label=rf'$\kappa$ = {kR}')
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\bar a$', fontsize=13)
+    ax.set_ylabel(r'$\theta$', fontsize=13)
+    plt.legend(fontsize=12)
+    plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential.png"), dpi=600)
     plt.show()
 

path_plot.py

@@ -2,7 +2,9 @@ import numpy as np
 from rotational_path import PairRotationalPath, PathEnergyPlot
 import expansion
 from parameters import ModelParams
+import patch_size
 from pathlib import Path
+import matplotlib.pyplot as plt
 
 zero_to_pi_half = np.linspace(0, np.pi/2, 100, endpoint=True)
 pi_half_to_pi = np.linspace(np.pi/2, np.pi, 100, endpoint=True)
@@ -18,41 +20,62 @@ QuadPath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half[::-1])
 
 DipolePath = PairRotationalPath()
 DipolePath.set_default_x_axis(zero_to_pi_half)
-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=pi_half_to_pi)
-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)
+DipolePath.add_euler(beta2=pi_half_to_pi[::-1])
+DipolePath.add_euler(beta2=zero_to_pi_half[::-1])
+DipolePath.add_euler(beta2=zero_to_pi_half, beta1=zero_to_pi_half)
+DipolePath.add_euler(beta2=np.pi/2, beta1=np.pi/2, alpha2=zero_to_pi_half)
+DipolePath.add_euler(beta2=np.pi/2, alpha2=np.pi/2, beta1=pi_half_to_pi)
+DipolePath.add_euler(beta2=np.pi/2, beta1=pi_half_to_pi[::-1], alpha1=np.pi)
+DipolePath.add_euler(beta2=zero_to_pi_half[::-1], beta1=pi_half_to_pi, alpha1=np.pi)
+DipolePath.add_euler(beta2=zero_to_pi_half, beta1=pi_half_to_pi[::-1], alpha1=np.pi)
+DipolePath.add_euler(beta2=pi_half_to_pi, beta1=zero_to_pi_half[::-1], alpha1=np.pi)
 
 DipolePath2 = PairRotationalPath()
 DipolePath2.set_default_x_axis(zero_to_pi_half)
-DipolePath2.add_euler(beta1=pi_half_to_pi[::-1])
-DipolePath2.add_euler(beta1=zero_to_pi_half[::-1])
-DipolePath2.add_euler(beta1=zero_to_pi_half, beta2=zero_to_pi_half)
-DipolePath2.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half)
-DipolePath2.add_euler(beta1=np.pi/2, alpha1=np.pi/2, beta2=pi_half_to_pi)
-DipolePath2.add_euler(beta1=zero_to_pi_half[::-1], beta2=pi_half_to_pi[::-1])
-DipolePath2.add_euler(beta1=zero_to_pi_half[::-1], beta2=np.pi)
-DipolePath2.add_euler(beta1=zero_to_pi_half, beta2=pi_half_to_pi[::-1], alpha2=np.pi)
-DipolePath2.add_euler(beta1=pi_half_to_pi, beta2=zero_to_pi_half[::-1], alpha2=np.pi)
+DipolePath2.add_euler(beta2=pi_half_to_pi[::-1])
+DipolePath2.add_euler(beta2=zero_to_pi_half[::-1])
+DipolePath2.add_euler(beta2=zero_to_pi_half, beta1=zero_to_pi_half)
+DipolePath2.add_euler(beta2=np.pi/2, beta1=np.pi/2, alpha2=zero_to_pi_half)
+DipolePath2.add_euler(beta2=np.pi/2, alpha2=np.pi/2, beta1=pi_half_to_pi)
+DipolePath2.add_euler(beta2=zero_to_pi_half[::-1], beta1=pi_half_to_pi[::-1])
+DipolePath2.add_euler(beta2=zero_to_pi_half[::-1], beta1=np.pi)
+DipolePath2.add_euler(beta2=zero_to_pi_half, beta1=pi_half_to_pi[::-1], alpha1=np.pi)
+DipolePath2.add_euler(beta2=pi_half_to_pi, beta1=zero_to_pi_half[::-1], alpha1=np.pi)
 
 
 if __name__ == '__main__':
 
+    # kappaR = np.array([1, 3, 10])
     params = ModelParams(R=150, kappaR=3)
     # ex1 = expansion.MappedExpansionQuad(np.array([0.35, 0.44, 0.6]), 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.GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=np.array([5]), sigma1=0.001, l_max=10)
-    ex1 = expansion.GaussianCharges(omega_k=np.array([np.pi, 0]), lambda_k=np.array([1, 5, 10, 15]), sigma1=0.001, l_max=10)
+    # ex1 = expansion.GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=np.array([5, 10]), sigma1=0.001, l_max=10)
+    # ex1 = expansion.GaussianCharges(omega_k=np.array([0, 0]), lambda_k=np.array([1, 5, 10, 15]), sigma1=0.001, l_max=20)
+    ex1 = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), np.array([0.3, 0.5, 1]), 0.001, 30)
     ex2 = ex1.clone()
 
-    path_plot = PathEnergyPlot(ex1, ex2, DipolePath2, dist=2., params=params)
-    path_plot.plot(labels=[rf'$\lambda$={l}' for l in [1, 5, 10, 15]],
-                   # norm_euler_angles={'beta1': np.pi/2}
-                   )
+    # charge_profile = ex1.charge_value(theta=np.linspace(0, np.pi, 100), phi=0)
+    # plt.plot(charge_profile.T)
+    # plt.show()
+
+    expansion.plot_theta_profile(ex1, num=1000, theta_end=np.pi, phi=0)
+
+    ps = patch_size.potential_patch_size(ex1, params, theta1=np.pi / 2,
+                                         # match_expansion_axis_to_params=1
+                                         )
+    print(ps)
+
+    # path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params)
+
+    # path_plot.plot(labels=[rf'$\theta$={180 * patch / np.pi:.1f}' for patch in np.array([0.3, 0.5, 1.0])],
+    #                # norm_euler_angles={'beta2': np.pi},
+    #                # save_as=Path("/home/andraz/ChargedShells/Figures/dipole_path_kappaR3.png")
+    #                )
+
+    # path_plot.plot(labels=[rf'$\kappa R$={kR}' for kR in kappaR],
+    #                # norm_euler_angles={'beta2': np.pi},
+    #                # save_as=Path("/home/andraz/ChargedShells/Figures/dipole_path_lambda5_norm2.png")
+    #                )
+
     # path_plot.plot_sections(save_as=Path('/home/andraz/ChargedShells/Figures/dipole_path2.png'))

rotational_path.py

@@ -102,7 +102,7 @@ class PathEnergyPlot:
                                      gamma=norm_euler_angles.get('gamma2', 0))
         return np.abs(self.evaluate_energy())
 
-    def plot(self, labels: list[str] = None, norm_euler_angles: dict = None):
+    def plot(self, labels: list[str] = None, norm_euler_angles: dict = None, save_as: Path = None):
         energy = self.path_energy()
         normalization = self.normalization(norm_euler_angles)
         energy = energy / normalization[None, ...]
@@ -111,11 +111,14 @@ class PathEnergyPlot:
 
         fig, ax = plt.subplots()
         ax.axhline(y=0, c='k', linestyle=':')
-        ax.plot(x_axis, energy, label=labels)
+        ax.plot(x_axis, np.squeeze(energy), label=labels)
         ax.legend(fontsize=12)
         ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
         ax.set_xlabel('angle', fontsize=13)
         ax.set_ylabel('U', fontsize=13)
+        plt.tight_layout()
+        if save_as is not None:
+            plt.savefig(save_as, dpi=600)
         plt.show()
 
     def plot_sections(self, norm_euler_angles: dict = None, save_as: Path = None):