The core modules are now a package

analysis/expansion_plot.py

@@ -0,0 +1,70 @@
+from charged_shells import expansion
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+import plotly.graph_objects as go
+
+Expansion = expansion.Expansion
+
+
+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.T)
+    plt.show()
+
+
+def plot_theta_profile_multiple(ex_list: list[Expansion], label_list, phi: float = 0, num: int = 100,
+                                theta_start: float = 0, theta_end: float = np.pi):
+    theta_vals = np.linspace(theta_start, theta_end, num)
+
+    fig, ax = plt.subplots()
+    for ex, label in zip(ex_list, label_list):
+        ax.plot(theta_vals, ex.charge_value(theta_vals, phi).T, label=label)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\theta$', fontsize=13)
+    ax.set_ylabel(r'$\sigma$', fontsize=13)
+    plt.legend(fontsize=12)
+    plt.tight_layout()
+    plt.savefig(Path("/home/andraz/ChargedShells/Figures/charge_shape_comparison.png"), dpi=600)
+    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 main():
+    # ex = expansion.MappedExpansionQuad(0.44, 3, 1, 10)
+    # ex = expansion.Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
+    # ex = expansion.GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
+    ex = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.1, 30)
+    # print(np.real(ex.coefs))
+    # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
+    plot_charge_3d(ex)
+
+    # new_coeffs = expanison.expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
+    # print(new_coeffs.shape)
+    #
+    # newnew_coeffs = expansion.expansion_rotation(Quaternion(np.arange(16).reshape(4, 4)).normalized, new_coeffs, ex.l_array)
+    # print(newnew_coeffs.shape)
+
+
+if __name__ == '__main__':
+
+    main()

patch_shape_comparison.py → analysis/patch_shape_comparison.py

@@ -1,14 +1,9 @@
-import expansion
-import patch_size
-import potentials
+from charged_shells import expansion, functions as fn, potentials, patch_size
 import numpy as np
-from parameters import ModelParams
-import functions as fn
+from charged_shells.parameters import ModelParams
 import matplotlib.pyplot as plt
 import scipy.special as sps
 from pathlib import Path
-import rotational_path
-import path_plot
 
 
 def point_to_gauss_map(sigma_m, a_bar, lbd, params: ModelParams):
@@ -23,12 +18,12 @@ def point_to_cap_map(sigma_m, a_bar, theta0, params: ModelParams):
 
 if __name__ == '__main__':
 
-    target_patch_size = 0.9
-    params = ModelParams(R=150, kappaR=1)
+    target_patch_size = 0.92
+    params = ModelParams(R=150, kappaR=3)
     sigma_m = 0.001
 
     def fn1(x):
-        return expansion.MappedExpansionQuad(a_bar=x, kappaR=params.kappaR, sigma_m=sigma_m, l_max=30)
+        return expansion.MappedExpansionQuad(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)
 
     def fn2(x):
         return expansion.GaussianCharges(lambda_k=x, omega_k=np.array([[0, 0], [np.pi, 0]]), sigma1=0.001, l_max=30)
@@ -40,7 +35,7 @@ if __name__ == '__main__':
     lbd = patch_size.inverse_potential_patch_size(target_patch_size, fn2, 5, params)
     theta0 = patch_size.inverse_potential_patch_size(target_patch_size, fn3, 0.5, params)
 
-    ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_m=sigma_m, l_max=30)
+    ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)
 
     gauss_sigma = point_to_gauss_map(sigma_m, a_bar, lbd, params)
     ex_gauss = expansion.GaussianCharges(lambda_k=lbd, omega_k=np.array([[0, 0], [np.pi, 0]]), sigma1=gauss_sigma, l_max=30)
@@ -62,11 +57,11 @@ if __name__ == '__main__':
     # expansion.plot_theta_profile_multiple([ex_point, ex_gauss, ex_cap], ['IC', 'Gauss', 'cap'], num=1000)
 
     fig, ax = plt.subplots()
-    ax.plot(theta, 1000 * potential1.T, label='CS', linewidth=2)
+    ax.scatter(theta[::50], 1000 * potential_ic.T[::50], marker='o', label='ICi', facecolors='none', edgecolors='tab:red')
+    ax.plot(theta, 1000 * potential1.T, label='CSp - mapped', linewidth=2)
     # ax.plot(theta, potential_ic.T, label='IC', ls=':', linewidth=2, marker='o', markevery=50, mfc='none')
-    ax.scatter(theta[::50], 1000 * potential_ic.T[::50], marker='o', label='IC', facecolors='none', edgecolors='tab:red')
-    ax.plot(theta, 1000 * potential2.T, label='Gauss', linewidth=2, ls='--')
-    ax.plot(theta, 1000 * potential3.T, label='cap', linewidth=2, ls='--')
+    ax.plot(theta, 1000 * potential2.T, label='CSp - Gauss', linewidth=2, ls='--')
+    ax.plot(theta, 1000 * potential3.T, label='CSp - caps', linewidth=2, ls='--')
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
     ax.set_xlabel(r'$\theta$', fontsize=15)
     ax.set_ylabel(r'$\phi$ [mV]', fontsize=15)
@@ -77,7 +72,7 @@ if __name__ == '__main__':
     plt.xticks(custom_ticks, custom_labels, fontsize=13)
     plt.legend(fontsize=12)
     plt.tight_layout()
-    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison.png"), dpi=600)
+    plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison.pdf"), dpi=300)
     plt.show()
 
     # path_comparison = rotational_path.PathExpansionComparison([ex_point, ex_gauss, ex_cap],
@@ -87,6 +82,3 @@ if __name__ == '__main__':
     # path_comparison.plot(['IC', 'Gauss', 'cap'],
     #                      save_as=Path("/home/andraz/ChargedShells/Figures/energy_shape_comparison_kR1.png"))
 
-
-
-

analysis/patch_size_dependence.py

@@ -0,0 +1,138 @@
+from matplotlib import pyplot as plt
+from pathlib import Path
+import numpy as np
+from charged_shells import expansion, patch_size
+from charged_shells.parameters import ModelParams
+
+
+def plot_abar_dependence(*, save=False):
+    a_bar = np.linspace(0.2, 0.8, 101)
+    kappaR = np.array([0.1, 1, 3, 10, 30])
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])
+
+    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=1)
+
+    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    for patch, kR, marker in zip(ps.T, kappaR, markers):
+        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', marker=marker, markerfacecolor='none', markevery=5)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\bar a$', fontsize=15)
+    ax.set_ylabel(r'$\theta_0$', fontsize=15)
+    plt.legend(fontsize=14)
+    plt.tight_layout()
+    if save:
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential.pdf"), dpi=300)
+    plt.show()
+
+
+def plot_abar_charge_dependence(*, save=False):
+    a_bar = np.linspace(0.2, 0.8, 101)
+    kappaR = np.array([0.1, 1, 3, 10, 30])
+    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])
+
+    ps = patch_size.charge_patch_size(ex)
+
+    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    for patch, kR, marker in zip(ps.T, kappaR, markers):
+        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', marker=marker, markerfacecolor='none', markevery=5)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\bar a$', fontsize=15)
+    ax.set_ylabel(r'$\theta_0$', fontsize=15)
+    plt.legend(fontsize=14)
+    plt.tight_layout()
+    if save:
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_charge.pdf"), dpi=300)
+    plt.show()
+
+
+def plot_kappaR_charge_dependence(*, normalized=False, save=False):
+    a_bar = np.array([0.2, 0.4, 0.6, 0.8])
+    kappaR = np.linspace(0.01, 30, 100)
+    ex = expansion.MappedExpansionQuad(a_bar=a_bar[None, :], sigma_tilde=0.001, l_max=30, kappaR=kappaR[:, None])
+
+    ps = patch_size.charge_patch_size(ex)
+    if normalized:
+        ps = ps / ps[0][None, :]
+
+    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    for patch, abar, marker in zip(ps.T, a_bar, markers):
+        ax.plot(kappaR, patch, label=rf'$\bar a$ = {abar}', marker=marker, markerfacecolor='none', markevery=5)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\kappa R$', fontsize=15)
+    ax.set_ylabel(r'$\theta_0$', fontsize=15)
+    plt.legend(fontsize=14)
+    plt.tight_layout()
+    if save:
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_charge_nonmonotonicity.pdf"), dpi=300)
+    plt.show()
+
+
+def plot_sigma0_dependence(*, save=False):
+
+    # kappaR = np.array([0.1, 1, 3, 10, 30])
+    kappaR = np.linspace(0.1, 15, 101)
+    # sigma0 = np.linspace(-0.001, 0.0015, 6)
+    sigma0 = np.array([-0.001, 0, 0.001])
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar=0.8, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)
+
+    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0)
+
+    markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    for patch, sgm, marker in zip(ps.T, sigma0, markers):
+        ax.plot(kappaR, patch, label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=5)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\kappa R$', fontsize=15)
+    ax.set_ylabel(r'$\theta_0$', fontsize=15)
+    plt.legend(fontsize=14)
+    plt.tight_layout()
+    if save:
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_charge_abar08.pdf"), dpi=300)
+    plt.show()
+
+
+def plot_sigma0_dependence_relative(*, save=False):
+
+    # kappaR = np.array([0.1, 1, 3, 10, 30])
+    kappaR = np.linspace(0.1, 15, 101)
+    # sigma0 = np.linspace(-0.001, 0.0015, 6)
+    sigma0 = np.array([-0.001, 0, 0.001])
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar=0.3, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)
+    ex_neutral = expansion.MappedExpansionQuad(a_bar=0.3, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=0)
+
+    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0)
+    ps_neutral = patch_size.potential_patch_size(ex_neutral, params, match_expansion_axis_to_params=0)
+
+    markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    for patch, sgm, marker in zip(ps.T, sigma0, markers):
+        ax.plot(kappaR, patch / ps_neutral, label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=5)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+    ax.set_xlabel(r'$\kappa R$', fontsize=15)
+    ax.set_ylabel(r'$\theta_0$', fontsize=15)
+    plt.legend(fontsize=14)
+    plt.tight_layout()
+    if save:
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_charge_abar03_relative.pdf"), dpi=300)
+    plt.show()
+
+
+if __name__ == '__main__':
+
+    # plot_abar_dependence(save=False)
+    # plot_sigma0_dependence(save=True)
+    # plot_sigma0_dependence_relative(save=True)
+
+    # plot_abar_charge_dependence(save=True)
+    plot_kappaR_charge_dependence(normalized=True, save=True)

path_plot.py → analysis/path_plot.py

@@ -1,10 +1,10 @@
 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
+from charged_shells.rotational_path import PairRotationalPath, PathEnergyPlot
+from charged_shells import expansion, patch_size
+from charged_shells.parameters import ModelParams
+
+Array = np.ndarray
+
 
 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)
@@ -43,16 +43,49 @@ DipolePath2.add_euler(beta2=zero_to_pi_half, beta1=pi_half_to_pi[::-1], alpha1=n
 DipolePath2.add_euler(beta2=pi_half_to_pi, beta1=zero_to_pi_half[::-1], alpha1=np.pi)
 
 
+def kappaR_dependence(kappaR: Array, abar: float, sigma_tilde=0.001, save=False):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, 0.001, l_max=30)
+    ex2 = ex1.clone()
+
+    # 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, DipolePath2, dist=2., params=params, match_expansion_axis_to_params=None)
+
+    path_plot.plot(  # labels=[rf'$\theta$={180 * patch / np.pi:.1f}' for patch in np.array([0.5])],
+        # 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'))
+
+
+
 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.MappedExpansionQuad(0.4, params.kappaR, 0.001, l_max=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, 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.5]), 0.001, 20)
+    ex1 = expansion.GaussianCharges(omega_k=np.array([0, 0]), lambda_k=np.array([1, 5, 10]), sigma1=0.001, l_max=20)
+    # ex1 = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), np.array([0.5]), 0.001, 20)
     ex2 = ex1.clone()
 
     # charge_profile = ex1.charge_value(theta=np.linspace(0, np.pi, 100), phi=0)
@@ -66,7 +99,7 @@ if __name__ == '__main__':
     #                                      )
     # print(ps)
 
-    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params)
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath2, dist=2., params=params, match_expansion_axis_to_params=None)
 
     path_plot.plot(# labels=[rf'$\theta$={180 * patch / np.pi:.1f}' for patch in np.array([0.5])],
                    # norm_euler_angles={'beta2': np.pi},

analysis/peak_heigth.py

@@ -0,0 +1,62 @@
+import matplotlib.pyplot as plt
+from charged_shells import expansion, interactions, mapping
+from charged_shells.parameters import ModelParams
+import numpy as np
+from typing import Literal
+from pathlib import Path
+from functools import partial
+
+Array = np.ndarray
+Expansion = expansion.Expansion
+RotConfig = Literal['ep', 'pp']
+
+
+def peak_energy_plot(kappaR: Array,
+                     a_bar: Array,
+                     which: Literal['ep', 'pp'],
+                     R: float = 150,
+                     dist: float = 2.,
+                     l_max=20,
+                     save_as: Path = None):
+
+    params = ModelParams(R=R, kappaR=kappaR)
+
+    # energy = []
+    # for params in params.unravel():
+    #     ex1 = expansion.MappedExpansionQuad(a_bar, params.kappaR, 0.001, l_max=l_max)
+    #     ex2 = ex1.clone()
+    #     if which == 'ep':
+    #         ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
+    #     energy.append(interactions.charged_shell_energy(ex1, ex2, dist, params))
+    #
+    # energy = np.array(energy)
+
+    ex1 = expansion.MappedExpansionQuad(a_bar[:, None], params.kappaR[None, :], 0.001, l_max=l_max)
+    ex2 = ex1.clone()
+    if which == 'ep':
+        ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
+
+    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist), 1)
+    energy = energy_fn(ex1, ex2, params)
+
+    fig, ax = plt.subplots()
+    for en, ab in zip(energy, a_bar):
+        ax.plot(kappaR, en / en[0], label=rf'$\bar a = {ab:.1f}$')
+    ax.legend(fontsize=12)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel(r'$\kappa R$', fontsize=15)
+    ax.set_ylabel(rf'$\bar V$', fontsize=15)
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+if __name__ == '__main__':
+
+    kappaR = np.arange(1, 10, 0.1)
+    a_bar = np.arange(0.2, 0.8, 0.2)
+
+    peak_energy_plot(kappaR, a_bar, which='ep',
+                     # save_as=Path('/home/andraz/ChargedShells/Figures/nonmonotonicity_check_ep.pdf')
+                     )

analysis/sEE_minimum.py

@@ -0,0 +1,39 @@
+import numpy as np
+from charged_shells import expansion, interactions, mapping
+from charged_shells.parameters import ModelParams
+from functools import partial
+
+Expansion = expansion.Expansion
+
+
+def sEE_minimum(ex: Expansion, params: ModelParams, accuracy=1e-3, dist=2., match_expansion_axis_to_params=None):
+    ex2 = ex.clone()
+    zero_to_pi_half = np.linspace(0, np.pi / 2, int(1 / accuracy), endpoint=True)
+
+    ex.rotate_euler(alpha=0, beta=zero_to_pi_half, gamma=0)
+    ex2.rotate_euler(alpha=0, beta=zero_to_pi_half, gamma=0)
+
+    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist),
+                                                     match_expansion_axis_to_params)
+    energy = energy_fn(ex, ex2, params)
+
+    min_idx = np.argmin(energy, axis=0)
+    min_energy = np.min(energy, axis=0)
+
+    return min_energy, zero_to_pi_half[min_idx]
+
+
+def main():
+
+    kappaR = np.array([1, 3, 10])
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(0.4, kappaR, 0.001)
+    min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=1)
+    print(min_energy, min_angle + np.pi/2)
+
+
+
+if __name__ == '__main__':
+
+    main()
+

expansion.py → charged_shells/expansion.py

@@ -2,22 +2,15 @@ from __future__ import annotations
 import numpy as np
 from dataclasses import dataclass
 from functools import cached_property
-import functions as fn
+import charged_shells.functions as fn
 import quaternionic
 import spherical
-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
-from pathlib import Path
 
 
 Array = np.ndarray
 Quaternion = quaternionic.array
-T = TypeVar('T')
-V = TypeVar('V')
 
 
 class InvalidExpansion(Exception):
@@ -65,7 +58,7 @@ class Expansion:
     @cached_property
     def lm_arrays(self) -> (Array, Array):
         """Return l and m arrays containing all (l, m) pairs."""
-        return full_fm_arrays(self.l_array)
+        return full_lm_arrays(self.l_array)
 
     def repeat_over_m(self, arr: Array, axis=0) -> Array:
         if not arr.shape[axis] == len(self.l_array):
@@ -108,24 +101,49 @@ class Expansion24(Expansion):
 
 
 class MappedExpansionQuad(Expansion):
-
-    def __init__(self, a_bar: Array | float, kappaR: Array | float, sigma_m: float, l_max: int = 20, sigma0: float = 0):
+    """Expansion that matches the outside potential of a quadrupolar impermeable particle with point charges inside."""
+
+    def __init__(self,
+                 a_bar: Array | float,
+                 kappaR: Array | float,
+                 sigma_tilde: float,
+                 l_max: int = 20,
+                 sigma0: float | Array = 0):
+        """
+        :param a_bar: distance between the center and off center charges
+        :param kappaR: screening parameter
+        :param sigma_tilde: magnitude of off-center charges / 4pi R^2
+        :param l_max: maximal ell value for the expansion
+        :param sigma0: total (mean) charge density
+        """
         a_bar, kappaR = np.broadcast_arrays(a_bar, kappaR)
 
         l_array = np.array([l for l in range(l_max + 1) if l % 2 == 0])
-        a_bar, kappaR, l_array_expanded = np.broadcast_arrays(a_bar[..., None], kappaR[..., None], l_array[None, :])
+        a_bar, kappaR, l_array_expanded = np.broadcast_arrays(a_bar[..., None],
+                                                              kappaR[..., None],
+                                                              l_array[None, :])
 
-        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 / np.sqrt(4 * np.pi)
+        coefs = (2 * sigma_tilde * 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 = np.squeeze(rot_sym_expansion(l_array, coefs))
+        coefs = expansion_total_charge(coefs, sigma0)
         super().__init__(l_array, coefs)
 
 
 class GaussianCharges(Expansion):
+    """Expansion for a collection of smeared charges on the sphere."""
 
     def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int,
-                 sigma0: float = 0, equal_charges: bool = True):
+                 sigma0: float | Array = 0, equal_charges: bool = True):
+        """
+        :param omega_k: array of positions (theta, phi) of all charges
+        :param lambda_k: smear parameter for each charge or smear for different cases (if equal_charges = True)
+        :param sigma1: scaling
+        :param l_max: maximal ell value for the expansion
+        :param sigma0: total (mean) charge density
+        :param equal_charges: if this is False, length of lambda_k should be N. If True, theta0_k array will be treated
+            as different expansion cases
+        """
         omega_k = omega_k.reshape(-1, 2)
         if not isinstance(lambda_k, Array):
             lambda_k = np.array([lambda_k])
@@ -137,7 +155,7 @@ class GaussianCharges(Expansion):
             raise ValueError("Number of charges (length of omega_k) should match the last dimension of lambda_k array.")
         lambda_k = lambda_k.reshape(-1, omega_k.shape[0])
         l_array = np.arange(l_max + 1)
-        full_l_array, full_m_array = full_fm_arrays(l_array)
+        full_l_array, full_m_array = full_lm_arrays(l_array)
         theta_k = omega_k[:, 0]
         phi_k = omega_k[:, 1]
         summands = (lambda_k[:, None, :] / np.sinh(lambda_k[:, None, :])
@@ -145,15 +163,25 @@ 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 / np.sqrt(4 * np.pi)
+        coefs = expansion_total_charge(coefs, sigma0)
         l_array, coefs = purge_unneeded_l(l_array, coefs)
         super().__init__(l_array, coefs)
 
 
 class SphericalCap(Expansion):
+    """Expansion for a collection of spherical caps."""
 
-    def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float = 0,
+    def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float | Array = 0,
                  equal_sizes: bool = True):
+        """
+        :param omega_k: array of positions (theta, phi) of all spherical caps
+        :param theta0_k: sizes of each spherical caps or cap sizes for different cases (if equal_sizes = True)
+        :param sigma1: charge magnitude for the single cap, currently assumes that this is equal for all caps
+        :param l_max: maximal ell value for the expansion
+        :param sigma0: total (mean) charge density
+        :param equal_sizes: if this is False, length of theta0_k should be N. If True, theta0_k array will be treated as
+            different expansion cases
+        """
         omega_k = omega_k.reshape(-1, 2)
         if not isinstance(theta0_k, Array):
             theta0_k = np.array(theta0_k)
@@ -179,29 +207,12 @@ class SphericalCap(Expansion):
         # 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)
+        coefs_all = expansion_total_charge(coefs_all, sigma0)
         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."""
-
-    def mapped_f(ex: Expansion, *args, **kwargs):
-        og_shape = ex.shape
-        flat_ex = ex.flatten()
-        results = []
-        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)
-        except AttributeError:
-            return np.array(results).reshape(og_shape)
-
-    return mapped_f
-
-
-def full_fm_arrays(l_array: Array) -> (Array, Array):
+def full_lm_arrays(l_array: Array) -> (Array, Array):
+    """From an array of l_values get arrays of ell and m that give you all pairs (ell, m)."""
     all_m_list = []
     for l in l_array:
         for i in range(2 * l + 1):
@@ -216,22 +227,30 @@ def rot_sym_expansion(l_array: Array, coefs: Array) -> Array:
     return full_coefs
 
 
-def coefs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expansion:
-    missing_l = np.setdiff1d(target_l_array, expansion.l_array, assume_unique=True)
-    fill = np.zeros(np.sum(2 * missing_l + 1))
-    full_l_array1, _ = expansion.lm_arrays
-    # we search for where to place missing coefs with the help of a boolean array and argmax function
-    comparison_bool = (full_l_array1[:, None] - missing_l[None, :]) > 0
-    indices = np.where(np.any(comparison_bool, axis=0), np.argmax(comparison_bool, axis=0), full_l_array1.shape[0])
-    new_coefs = np.insert(expansion.coefs, np.repeat(indices, 2 * missing_l + 1), fill, axis=-1)
-    return Expansion(target_l_array, new_coefs)
+def expansion_total_charge(coefs: Array, sigma0: float | Array):
+    """Adds a new axis to the expansion coefficients that modifies expansion based on given net charge density."""
+    if sigma0 is None:
+        return coefs
+    if not isinstance(sigma0, Array):
+        x = copy.deepcopy(coefs)
+        x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
+        return x
+    sigma0 = sigma0.flatten()
+    x = np.repeat(np.expand_dims(coefs, -2), len(sigma0), axis=-2)
+    x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
+    return x
 
 
 def m_indices_at_l(l_arr: Array, l_idx: int):
+    """
+    For a given l_array and index l_idx for some ell in this array, get indices of all (ell, m) coefficients
+    in coefficients array.
+    """
     return np.arange(np.sum(2 * l_arr[:l_idx] + 1), np.sum(2 * l_arr[:l_idx+1] + 1))
 
 
 def purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
+    """Remove ell values from expansion for which all (ell, m) coefficients are zero."""
     def delete_zero_entries(l, l_arr, cfs):
         l_idx = np.where(l_arr == l)[0][0]
         m_indices = m_indices_at_l(l_arr, l_idx)
@@ -243,48 +262,20 @@ def purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
     return l_array, coefs
 
 
-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.T)
-    plt.show()
-
-
-def plot_theta_profile_multiple(ex_list: list[Expansion], label_list, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
-    theta_vals = np.linspace(theta_start, theta_end, num)
-
-    fig, ax = plt.subplots()
-    for ex, label in zip(ex_list, label_list):
-        ax.plot(theta_vals, ex.charge_value(theta_vals, phi).T, label=label)
-    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
-    ax.set_xlabel(r'$\theta$', fontsize=13)
-    ax.set_ylabel(r'$\sigma$', fontsize=13)
-    plt.legend(fontsize=12)
-    plt.tight_layout()
-    plt.savefig(Path("/home/andraz/ChargedShells/Figures/charge_shape_comparison.png"), dpi=600)
-    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 coefs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expansion:
+    """Explicitly add missing expansion coefficients so that expansion includes all ell values from the target array."""
+    missing_l = np.setdiff1d(target_l_array, expansion.l_array, assume_unique=True)
+    fill = np.zeros(np.sum(2 * missing_l + 1))
+    full_l_array1, _ = expansion.lm_arrays
+    # we search for where to place missing coefs with the help of a boolean array and argmax function
+    comparison_bool = (full_l_array1[:, None] - missing_l[None, :]) > 0
+    indices = np.where(np.any(comparison_bool, axis=0), np.argmax(comparison_bool, axis=0), full_l_array1.shape[0])
+    new_coefs = np.insert(expansion.coefs, np.repeat(indices, 2 * missing_l + 1), fill, axis=-1)
+    return Expansion(target_l_array, new_coefs)
 
 
 def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
+    """Explicitly add zero expansion coefficients so that both expansions include coefficients for the same ell."""
     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)
 
@@ -309,42 +300,3 @@ def expansion_rotation(rotations: Quaternion, coefs: Array, l_array: Array):
         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 = 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 = SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.1, 70)
-    # print(np.real(ex.coefs))
-    # 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)
-    #
-    # 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(0, np.pi / 2, -1)
-    # t1 = time.time()
-    # print(ex.coefs)
-    # print(t1 - t0)
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-

interactions.py → charged_shells/interactions.py

@@ -1,11 +1,10 @@
-import expansion
-import parameters
-import functions as fn
+from charged_shells import expansion, parameters
+import charged_shells.functions as fn
 from py3nj import wigner3j
 import numpy as np
 import time
 from typing import Literal
-import units_and_constants as uc
+import charged_shells.units_and_constants as uc
 
 
 Array = np.ndarray
@@ -24,7 +23,7 @@ def energy_units(units: EnergyUnit, params: ModelParams) -> float:
             return uc.UNITS.energy
 
 
-def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: ModelParams, units: EnergyUnit = 'kT'):
+def charged_shell_energy(ex1: Expansion, ex2: Expansion, params: ModelParams, dist: float = 2, units: EnergyUnit = 'kT'):
 
     ex1, ex2 = expansion.expansions_to_common_l(ex1, ex2)
     dist = dist * params.R
@@ -77,26 +76,3 @@ def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: Mo
     broadcasted_lspm_vals = np.broadcast_to(lspm_vals, charge_vals.shape)
 
     return 0.5 * constants * np.sum(broadcasted_lspm_vals * charge_vals, axis=-1)
-
-
-if __name__ == '__main__':
-
-    params = ModelParams(R=150, kappaR=3)
-    ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=20)
-    ex2 = ex1.clone()
-
-    dist = 2.
-
-    # 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()
-
-    print('energy: ', energy)
-    print('time: ', t1 - t0)
-
-    # plt.plot(energy)
-    # plt.show()

charged_shells/mapping.py

@@ -0,0 +1,102 @@
+from typing import Callable, TypeVar
+import numpy as np
+from charged_shells.expansion import Expansion
+from charged_shells.parameters import ModelParams
+
+T = TypeVar('T')
+V = TypeVar('V')
+Array = np.ndarray
+
+
+def map_over_expansion(f: Callable[[Expansion, T], V]) -> Callable[[Expansion, T], V]:
+    """Map a function f over all leading axes of an expansion. Uses for loops, so it is kinda slow."""
+
+    def mapped_f(ex: Expansion, *args, **kwargs):
+        og_shape = ex.shape
+        flat_ex = ex.flatten()
+        results = []
+        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)
+        except AttributeError:
+            return np.array(results).reshape(og_shape)
+
+    return mapped_f
+
+
+def unravel_params(params: ModelParams) -> list[ModelParams]:
+    if isinstance(params.R, Array) and isinstance(params.kappa, Array):
+        # if this is to be implemented, watch out for implementations of mapping expansions that depend
+        # on one of the parameters in ModelParams over other functions that also take the same ModelParameters
+        raise NotImplementedError("Currently only unravel over a single parameter is supported. ")
+    if isinstance(params.R, Array):
+        return [ModelParams(R=r, kappa=params.kappa) for r in params.R]
+    if isinstance(params.kappa, Array):
+        return [ModelParams(R=params.R, kappa=kappa) for kappa in params.kappa]
+    if not (isinstance(params.R, Array) or isinstance(params.kappa, Array)):
+        return [params]
+
+
+SingleExpansionFn = Callable[[Expansion, ModelParams], T]
+TwoExpansionsFn = Callable[[Expansion, Expansion, ModelParams], T]
+
+
+def parameter_map_single_expansion(f: SingleExpansionFn,
+                                   match_expansion_axis_to_params: int = None) -> SingleExpansionFn:
+
+    meap = match_expansion_axis_to_params  # just a shorter variable name
+
+    def mapped_f(ex: Expansion, params: ModelParams):
+        params_list = unravel_params(params)
+        if meap is not None:
+            expansion_list = [Expansion(ex.l_array, np.take(ex.coefs, i, axis=meap)) for i in range(ex.shape[meap])]
+        else:
+            expansion_list = [ex for _ in params_list]
+
+        if not len(expansion_list) == len(params_list):
+            raise ValueError(f'Axis of expansion that is supposed to match params does not have the same length, got '
+                             f'len(unraveled params) = {len(params_list)} and '
+                             f'expansion.shape[{meap}] = {len(expansion_list)}')
+
+        results = []
+        for exp, prms in zip(expansion_list, params_list):
+            results.append(f(exp, prms))
+
+        if meap is not None:
+            return np.array(results).swapaxes(0, meap)  # return the params-matched axis to where it belongs
+        return np.squeeze(np.array(results))
+
+    return mapped_f
+
+
+def parameter_map_two_expansions(f: TwoExpansionsFn,
+                                 match_expansion_axis_to_params: int = None) -> TwoExpansionsFn:
+
+    meap = match_expansion_axis_to_params  # just a shorter variable name
+
+    def mapped_f(ex1: Expansion, ex2: Expansion, params: ModelParams):
+        params_list = unravel_params(params)
+        if meap is not None:
+            expansion_list1 = [Expansion(ex1.l_array, np.take(ex1.coefs, i, axis=meap)) for i in range(ex1.shape[meap])]
+            expansion_list2 = [Expansion(ex2.l_array, np.take(ex2.coefs, i, axis=meap)) for i in range(ex2.shape[meap])]
+        else:
+            expansion_list1 = [ex1 for _ in params_list]
+            expansion_list2 = [ex2 for _ in params_list]
+
+        if not (len(expansion_list1) == len(params_list) and len(expansion_list2) == len(params_list)):
+            raise ValueError(f'Axis of at least one of the expansions that is supposed to match params '
+                             f'does not have the same length, got '
+                             f'len(unraveled params) = {len(params_list)} and '
+                             f'expansion1.shape[{meap}] = {len(expansion_list1)} and '
+                             f'expansion2.shape[{meap}] = {len(expansion_list2)}')
+
+        results = []
+        for exp1, exp2, prms in zip(expansion_list1, expansion_list2, params_list):
+            results.append(f(exp1, exp2, prms))
+
+        if meap is not None:
+            return np.array(results).swapaxes(0, meap)  # return the params-matched axis to where it belongs
+        return np.squeeze(np.array(results))
+
+    return mapped_f

parameters.py → charged_shells/parameters.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from dataclasses import dataclass
 import numpy as np
-import units_and_constants as uc
+import charged_shells.units_and_constants as uc
 
 
 Array = np.ndarray
@@ -20,15 +20,6 @@ class ModelParams:
         self.kappa, self.kappaR, self.c0 = screening_calculator(self.R, self.temperature, self.epsilon,
                                                                 self.kappa, self.kappaR, self.c0)
 
-    def unravel(self) -> list[ModelParams]:
-        params_list = []
-        all_r = np.array([self.R]) if not isinstance(self.R, Array) else self.R
-        all_kappa = np.array([self.kappa]) if not isinstance(self.kappa, Array) else self.kappa
-        for r in all_r:
-            for kappa in all_kappa:
-                params_list.append(ModelParams(R=r, kappa=kappa))
-        return params_list
-
 
 def bjerrum(temp: float, epsilon: float) -> float:
     return uc.CONSTANTS.e0 ** 2 / (4 * np.pi * epsilon * uc.CONSTANTS.epsilon0 * uc.CONSTANTS.Boltzmann * temp)

charged_shells/patch_size.py

@@ -0,0 +1,46 @@
+import numpy as np
+from scipy.optimize import bisect, root_scalar
+from charged_shells import expansion, mapping, parameters, potentials
+from typing import Callable
+
+Expansion = expansion.Expansion
+Array = np.ndarray
+ModelParams = parameters.ModelParams
+
+
+@mapping.map_over_expansion
+def charge_patch_size(ex: Expansion, phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2):
+    return bisect(lambda theta: ex.charge_value(theta, phi), theta0, theta1)
+
+
+def potential_patch_size(ex: Expansion, params: ModelParams,
+                         phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2,
+                         match_expansion_axis_to_params: int = None):
+
+    # this is more complicate to map over leading axes of the expansion as potential also depends on model parameters,
+    # with some, such as kappaR, also being the parameters of the expansion in the first place. When mapping,
+    # we must therefore provide the expansion axis that should match the collection of parameters in params.
+
+    @mapping.map_over_expansion
+    def potential_zero(exp: Expansion, prms: ModelParams):
+        return bisect(lambda theta: potentials.charged_shell_potential(theta, phi, 1, exp, prms), theta0, theta1)
+
+    return mapping.parameter_map_single_expansion(potential_zero, match_expansion_axis_to_params)(ex, params)
+
+
+def inverse_potential_patch_size(target_patch_size: float,
+                                 ex_generator: Callable[[float], Expansion],
+                                 x0: float,
+                                 params: ModelParams, **ps_kwargs):
+
+    def patch_size_dif(x):
+        ex = ex_generator(x)
+        return potential_patch_size(ex, params, **ps_kwargs) - target_patch_size
+
+    root_result = root_scalar(patch_size_dif, x0=x0)
+
+    if not root_result.converged:
+        raise ValueError('No convergence. Patches of desired size might not be achievable in the given model. '
+                         'Conversely, a common mistake might be target patch size input in degrees.')
+
+    return root_result.root

potentials.py → charged_shells/potentials.py

@@ -1,9 +1,7 @@
-import expansion
-import functions as fn
+from charged_shells import expansion, parameters
+from charged_shells import functions as fn
+from charged_shells import units_and_constants as uc
 import numpy as np
-import parameters
-import units_and_constants as uc
-import matplotlib.pyplot as plt
 from scipy.special import eval_legendre, kv
 
 
@@ -70,26 +68,3 @@ def inverse_patchy_particle_potential(theta: Array | float,
                  / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0))
 
     return out0
-
-
-if __name__ == '__main__':
-
-    params = ModelParams(R=150, kappaR=3)
-    ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=30)
-    # ex2 = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.003, l_max=70)
-
-    theta = np.linspace(0, np.pi, 1000)
-    phi = 0.
-    dist = 1
-
-    potential_ic = inverse_patchy_particle_potential(theta, dist, 0.44, -0.002, (0.001, 0.001), params)
-
-    potential1 = charged_shell_potential(theta, phi, dist, ex1, params)
-    # potential2 = charged_shell_potential(theta, phi, dist, ex2, params)
-    # print(potential.shape)
-    # print(potential)
-
-    plt.plot(theta, potential_ic)
-    plt.plot(theta, potential1.T)
-    # plt.plot(theta, potential2.T)
-    plt.show()

rotational_path.py → charged_shells/rotational_path.py

@@ -1,11 +1,11 @@
 import numpy as np
 from dataclasses import dataclass, field
 import quaternionic
-import expansion
-from parameters import ModelParams
-import interactions
+from charged_shells import expansion, interactions, mapping
+from charged_shells.parameters import ModelParams
 import matplotlib.pyplot as plt
 from pathlib import Path
+from functools import partial
 
 
 Quaternion = quaternionic.array
@@ -66,6 +66,8 @@ class PathEnergyPlot:
     rot_path: PairRotationalPath
     dist: float | Array
     params: ModelParams
+    match_expansion_axis_to_params: int | None = None
+    units: interactions.EnergyUnit = 'kT'
 
     def __post_init__(self):
         if not isinstance(self.dist, Array):
@@ -74,8 +76,10 @@ class PathEnergyPlot:
     def evaluate_energy(self):
         energy = []
         for dist in self.dist:
-            for params in self.params.unravel():
-                energy.append(interactions.charged_shell_energy(self.expansion1, self.expansion2, dist, params))
+            energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy,
+                                                                     dist=dist, units=self.units),
+                                                             self.match_expansion_axis_to_params)
+            energy.append(energy_fn(self.expansion1, self.expansion2, self.params))
         return np.squeeze(np.stack(energy, axis=-1))
 
     def path_energy(self):
@@ -162,7 +166,7 @@ class PathExpansionComparison:
         path_vals = []
         for path in self.path_list:
             path_vals.append(path.evaluate_path(norm_euler_angles))
-        return np.stack(path_vals)
+        return np.squeeze(np.stack(path_vals))
 
     def plot(self, labels: list[str] = None, norm_euler_angles: dict = None, save_as: Path = None):
         energy = self.evaluate_path(norm_euler_angles=norm_euler_angles)
@@ -170,7 +174,7 @@ class PathExpansionComparison:
 
         fig, ax = plt.subplots()
         ax.axhline(y=0, c='k', linestyle=':')
-        ax.plot(x_axis, np.squeeze(energy).T, label=labels)
+        ax.plot(x_axis, 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)

patch_size.py

@@ -1,96 +0,0 @@
-import numpy as np
-from scipy.optimize import bisect, root_scalar
-import expansion
-from matplotlib import pyplot as plt
-import parameters
-import potentials
-from pathlib import Path
-from typing import Callable
-
-Expansion = expansion.Expansion
-Array = np.ndarray
-ModelParams = parameters.ModelParams
-
-
-@expansion.map_over_expansion
-def charge_patch_size(ex: Expansion, phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2):
-    return bisect(lambda theta: ex.charge_value(theta, phi), theta0, theta1)
-
-
-def potential_patch_size(ex: Expansion, params: ModelParams,
-                         phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2,
-                         match_expansion_axis_to_params: int = None):
-
-    # this is more complicate to map over leading axes of the expansion as potential also depends on model parameters,
-    # with some, such as kappaR, also being the parameters of the expansion in the first place. When mapping,
-    # we must therefore provide the expansion axis that should match the collection of parameters in params.
-
-    meap = match_expansion_axis_to_params  # just a shorter variable name
-
-    @expansion.map_over_expansion
-    def potential_zero(exp: Expansion, prms: ModelParams):
-        return bisect(lambda theta: potentials.charged_shell_potential(theta, phi, 1, exp, prms), theta0, theta1)
-
-    params_list = params.unravel()
-    if meap is not None:
-        expansion_list = [Expansion(ex.l_array, np.take(ex.coefs, i, axis=meap)) for i in range(ex.shape[meap])]
-    else:
-        expansion_list = [ex for _ in params_list]
-
-    if not len(expansion_list) == len(params_list):
-        raise ValueError(f'Axis of expansion that is supposed to match params does not have the same length, got '
-                         f'len(params.unravel()) = {len(params_list)} and '
-                         f'expansion.shape[{meap}] = {len(expansion_list)}')
-
-    results = []
-    for exp, prms in zip(expansion_list, params_list):
-        results.append(potential_zero(exp, prms))
-
-    if meap is not None:
-        return np.array(results).swapaxes(0, meap)  # return the params-matched axis to where it belongs
-    return np.squeeze(np.array(results))
-
-
-def inverse_potential_patch_size(target_patch_size: float,
-                                 ex_generator: Callable[[float], Expansion],
-                                 x0: float,
-                                 params: ModelParams, **ps_kwargs):
-
-    def patch_size_dif(x):
-        ex = ex_generator(x)
-        return potential_patch_size(ex, params, **ps_kwargs) - target_patch_size
-
-    root_result = root_scalar(patch_size_dif, x0=x0)
-
-    if not root_result.converged:
-        raise ValueError('No convergence. Patches of desired size might not be achievable in the given model. '
-                         'Conversely, a common mistake might be target patch size input in degrees.')
-
-    return root_result.root
-
-
-if __name__ == '__main__':
-
-    a_bar = np.linspace(0.2, 0.8, 101)
-    kappaR = np.array([0.1, 1, 3, 10, 30])
-    params = ModelParams(R=150, kappaR=kappaR)
-    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_m=0.001, l_max=30, kappaR=kappaR[None, :])
-
-    # fn = lambda x: expansion.SphericalCap(theta0_k=x, sigma1=0.001, l_max=30, omega_k=np.array([[0, 0], [np.pi, 0]]))
-    # print(inverse_potential_patch_size(48.8 * np.pi / 180, fn, 0.4, params))
-
-    # patch_size = charge_patch_size(ex)
-    patch_size = potential_patch_size(ex, params, match_expansion_axis_to_params=1)
-
-    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']
-
-    fig, ax = plt.subplots()
-    for patch, kR, marker in zip(patch_size.T, kappaR, markers):
-        ax.plot(a_bar, patch, label=rf'$\kappa$ = {kR}', marker=marker, markerfacecolor='none', markevery=5)
-    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_0$', fontsize=13)
-    plt.legend(fontsize=12)
-    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential.png"), dpi=600)
-    plt.show()
-

peak_heigth.py

@@ -1,54 +0,0 @@
-import matplotlib.pyplot as plt
-import expansion
-import interactions
-from parameters import ModelParams
-import numpy as np
-from typing import Literal
-from pathlib import Path
-
-Array = np.ndarray
-Expansion = expansion.Expansion
-RotConfig = Literal['ep', 'pp']
-
-
-def peak_energy_plot(kappaR: Array,
-                     a_bar: Array,
-                     which: Literal['ep', 'pp'],
-                     R: float = 150,
-                     dist: float = 2.,
-                     l_max=20,
-                     save_as: Path = None):
-
-    params = ModelParams(R=R, kappaR=kappaR)
-
-    energy = []
-    for params in params.unravel():
-        ex1 = expansion.MappedExpansionQuad(a_bar, params.kappaR, 0.001, l_max=l_max)
-        ex2 = ex1.clone()
-        if which == 'ep':
-            ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
-        energy.append(interactions.charged_shell_energy(ex1, ex2, dist, params))
-
-    energy = np.array(energy)
-
-    fig, ax = plt.subplots()
-    for en, ab in zip(energy.T, a_bar):
-        ax.plot(kappaR, en / en[0], label=rf'$\bar a = {ab:.1f}$')
-    ax.legend(fontsize=12)
-    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
-    ax.set_xlabel(r'$\kappa R$', fontsize=13)
-    ax.set_ylabel(rf'$\bar V$', fontsize=13)
-    plt.tight_layout()
-    if save_as is not None:
-        plt.savefig(save_as, dpi=600)
-    plt.show()
-
-
-if __name__ == '__main__':
-
-    kappaR = np.arange(1, 10, 0.1)
-    a_bar = np.arange(0.2, 0.8, 0.2)
-
-    peak_energy_plot(kappaR, a_bar, which='pp',
-                     save_as=Path('/home/andraz/ChargedShells/Figures/nonmonotonicity_check_pp.pdf')
-                     )

potential_tests.py

@@ -1,38 +0,0 @@
-from interactions import charged_shell_energy
-import expansion
-from parameters import ModelParams
-import numpy as np
-import matplotlib.pyplot as plt
-
-
-def v22_distance_test():
-
-    params = ModelParams(R=10, kappaR=3.29)
-    ex0 = expansion.Expansion24(1, 0, 0)
-    ex1 = ex0.clone()
-
-    ex0.rotate_euler(0, np.array([0, 0, np.pi / 2]), 0)
-    ex1.rotate_euler(0, np.array([0, np.pi / 2, np.pi / 2]), 0)
-
-    dist = np.linspace(2, 3.2, 100)
-    energy_array = np.zeros((dist.shape[0], 3))
-    for i, d in enumerate(dist):
-        energy_array[i, ...] = charged_shell_energy(ex0, ex1, d, params)
-
-    print(charged_shell_energy(ex0, ex1, 2., params))
-
-    plt.plot(dist, energy_array)
-    plt.show()
-
-
-if __name__ == '__main__':
-
-    v22_distance_test()
-
-
-
-
-
-
-
-

tests/interactions_test.py

@@ -0,0 +1,53 @@
+from charged_shells import expansion, interactions
+from charged_shells.parameters import ModelParams
+import time
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def v22_distance_test():
+
+    params = ModelParams(R=10, kappaR=3.29)
+    ex0 = expansion.Expansion24(1, 0, 0)
+    ex1 = ex0.clone()
+
+    ex0.rotate_euler(0, np.array([0, 0, np.pi / 2]), 0)
+    ex1.rotate_euler(0, np.array([0, np.pi / 2, np.pi / 2]), 0)
+
+    dist = np.linspace(2, 3.2, 100)
+    energy_array = np.zeros((dist.shape[0], 3))
+    for i, d in enumerate(dist):
+        energy_array[i, ...] = interactions.charged_shell_energy(ex0, ex1, params, d)
+
+    print(interactions.charged_shell_energy(ex0, ex1, params, dist=2.))
+
+    plt.plot(dist, energy_array)
+    plt.show()
+
+
+def timing():
+    params = ModelParams(R=150, kappaR=3)
+    ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=20)
+    ex2 = ex1.clone()
+
+    dist = 2.
+
+    # ex1, ex2 = expansions_to_common_l(ex1, ex2)
+    # print(ex1.coeffs)
+    # print(ex2.coeffs)
+
+    t0 = time.perf_counter()
+    energy = interactions.charged_shell_energy(ex1, ex2, params, dist)
+    t1 = time.perf_counter()
+
+    print('energy: ', energy)
+    print('time: ', t1 - t0)
+
+    # plt.plot(energy)
+    # plt.show()
+
+
+if __name__ == '__main__':
+
+    v22_distance_test()
+    timing()

tests/potential_tests.py

@@ -0,0 +1,39 @@
+from charged_shells import expansion, potentials
+from charged_shells.parameters import ModelParams
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def main():
+    params = ModelParams(R=150, kappaR=3)
+    ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=30)
+    # ex2 = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.003, l_max=70)
+
+    theta = np.linspace(0, np.pi, 1000)
+    phi = 0.
+    dist = 1
+
+    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, 0.44, -0.002, (0.001, 0.001), params)
+
+    potential1 = potentials.charged_shell_potential(theta, phi, dist, ex1, params)
+    # potential2 = charged_shell_potential(theta, phi, dist, ex2, params)
+    # print(potential.shape)
+    # print(potential)
+
+    plt.plot(theta, potential_ic)
+    plt.plot(theta, potential1.T)
+    # plt.plot(theta, potential2.T)
+    plt.show()
+
+
+if __name__ == '__main__':
+
+    main()
+
+
+
+
+
+
+
+