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()

analysis/patch_shape_comparison.py

@@ -0,0 +1,84 @@
+from charged_shells import expansion, functions as fn, potentials, patch_size
+import numpy as np
+from charged_shells.parameters import ModelParams
+import matplotlib.pyplot as plt
+import scipy.special as sps
+from pathlib import Path
+
+
+def point_to_gauss_map(sigma_m, a_bar, lbd, params: ModelParams):
+    return (sigma_m * fn.coefficient_Cim(2, params.kappaR) / fn.coefficient_Cpm(2, params.kappaR)
+            * np.sinh(lbd) / (lbd * fn.sph_bessel_i(2, lbd)) * a_bar ** 2)
+
+
+def point_to_cap_map(sigma_m, a_bar, theta0, params: ModelParams):
+    return (sigma_m * 10 * fn.coefficient_Cim(2, params.kappaR) / fn.coefficient_Cpm(2, params.kappaR) *
+            a_bar ** 2 / (sps.eval_legendre(1, np.cos(theta0)) - sps.eval_legendre(3, np.cos(theta0))))
+
+
+if __name__ == '__main__':
+
+    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_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)
+
+    def fn3(x):
+        return expansion.SphericalCap(theta0_k=x, sigma1=0.001, l_max=50, omega_k=np.array([[0, 0], [np.pi, 0]]))
+
+    a_bar = patch_size.inverse_potential_patch_size(target_patch_size, fn1, 0.5, params)
+    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_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)
+
+    cap_sigma = point_to_cap_map(sigma_m, a_bar, theta0, params)
+    ex_cap = expansion.SphericalCap(theta0_k=theta0, sigma1=cap_sigma, omega_k=np.array([[0, 0], [np.pi, 0]]), l_max=30)
+
+    theta = np.linspace(0, np.pi, 1001)
+    phi = 0.
+    dist = 1
+
+    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m, (sigma_m, sigma_m), params, 30)
+    potential1 = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)
+    potential2 = potentials.charged_shell_potential(theta, phi, dist, ex_gauss, params)
+    potential3 = potentials.charged_shell_potential(theta, phi, dist, ex_cap, params)
+    # print(potential.shape)
+    # print(potential)
+
+    # expansion.plot_theta_profile_multiple([ex_point, ex_gauss, ex_cap], ['IC', 'Gauss', 'cap'], num=1000)
+
+    fig, ax = plt.subplots()
+    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.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)
+    custom_ticks = [0, np.pi / 4, np.pi / 2, 3 * np.pi / 4, np.pi]
+    custom_labels = ['0', r'$\pi/4$', r'$\pi/2$', r'$3\pi/4$', r'$\pi$']
+    plt.axhline(y=0, color='black', linestyle=':')
+    plt.axvline(x=target_patch_size, color='black', linestyle=':')
+    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.pdf"), dpi=300)
+    plt.show()
+
+    # path_comparison = rotational_path.PathExpansionComparison([ex_point, ex_gauss, ex_cap],
+    #                                                           path_plot.QuadPath,
+    #                                                           dist=2,
+    #                                                           params=params)
+    # 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)

analysis/path_plot.py

@@ -0,0 +1,114 @@
+import numpy as np
+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)
+
+QuadPath = PairRotationalPath()
+QuadPath.set_default_x_axis(zero_to_pi_half)
+QuadPath.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half)
+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, alpha1=zero_to_pi_half[::-1])
+
+DipolePath = PairRotationalPath()
+DipolePath.set_default_x_axis(zero_to_pi_half)
+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(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)
+
+
+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(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]), 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)
+    # 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'))

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

charged_shells/expansion.py

@@ -0,0 +1,302 @@
+from __future__ import annotations
+import numpy as np
+from dataclasses import dataclass
+from functools import cached_property
+import charged_shells.functions as fn
+import quaternionic
+import spherical
+import copy
+from scipy.special import eval_legendre
+
+
+Array = np.ndarray
+Quaternion = quaternionic.array
+
+
+class InvalidExpansion(Exception):
+    pass
+
+
+@dataclass
+class Expansion:
+    """Generic class for storing surface charge expansion coefficients."""
+
+    l_array: Array
+    coefs: Array
+    _starting_coefs: Array = None  # initialized with the __post_init__ method
+    _rotations: Quaternion = Quaternion([1., 0., 0., 0.])
+
+    def __post_init__(self):
+        if self.coefs.shape[-1] != np.sum(2 * self.l_array + 1):
+            raise InvalidExpansion('Number of expansion coefficients does not match the provided l_array.')
+        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.')
+        self.coefs = self.coefs.astype(np.complex128)
+        self._starting_coefs = np.copy(self.coefs)
+
+    def __getitem__(self, item):
+        return Expansion(self.l_array, self.coefs[item])
+
+    @property
+    def max_l(self) -> int:
+        return max(self.l_array)
+
+    @property
+    def shape(self):
+        return self.coefs.shape[:-1]
+
+    def flatten(self) -> Expansion:
+        new_expansion = self.clone()  # np.ndarray.flatten() also copies the array
+        new_expansion.coefs = new_expansion.coefs.reshape(-1, new_expansion.coefs.shape[-1])
+        new_expansion._rotations = new_expansion._rotations.reshape(-1, 4)
+        return new_expansion
+
+    def reshape(self, shape: tuple):
+        self.coefs = self.coefs.reshape(shape + (self.coefs.shape[-1],))
+        self._rotations = self._rotations.reshape(shape + (4,))
+
+    @cached_property
+    def lm_arrays(self) -> (Array, Array):
+        """Return l and m arrays containing all (l, m) pairs."""
+        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):
+            raise ValueError('Array length should be equal to the number of l in the expansion.')
+        return np.repeat(arr, 2 * self.l_array + 1, axis=axis)
+
+    def rotate(self, rotations: Quaternion, rotate_existing=False):
+        # TODO: rotations are currently saved wrong if we start form existing coefficients not the og ones
+        self._rotations = rotations
+        coefs = self.coefs if rotate_existing else self._starting_coefs
+        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.squeeze(np.real(np.sum(self.coefs[..., None] * fn.sph_harm(full_l_array[:, None],
+                                                                             full_m_array[:, None],
+                                                                             theta[None, :], phi[None, :]), axis=-2)))
+
+    def clone(self) -> Expansion:
+        return copy.deepcopy(self)
+
+
+class Expansion24(Expansion):
+
+    def __init__(self, sigma2: float, sigma4: float, sigma0: float = 0.):
+        l_array = np.array([0, 2, 4])
+        coefs = rot_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
+        super().__init__(l_array, coefs)
+
+
+class MappedExpansionQuad(Expansion):
+    """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, :])
+
+        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 | 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])
+        if equal_charges:
+            if lambda_k.ndim > 1:
+                raise ValueError(f'If equal_charges=True, lambda_k should be a 1D array, got shape {lambda_k.shape}')
+            lambda_k = np.full((omega_k.shape[0], lambda_k.shape[0]), lambda_k).T
+        if lambda_k.shape[-1] != omega_k.shape[0]:
+            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_lm_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, :, None], lambda_k[:, None, :])
+                    * 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 = 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 | 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)
+        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)
+        coefs_all = expansion_total_charge(coefs_all, sigma0)
+        super().__init__(l_array, np.squeeze(coefs_all))
+
+
+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):
+            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, coefs: Array) -> Array:
+    """Create full expansion array for rotationally symmetric distributions with only m=0 terms different form 0."""
+    full_coefs = np.zeros(coefs.shape[:-1] + (np.sum(2 * l_array + 1),))
+    full_coefs[..., np.cumsum(2 * l_array + 1) - l_array - 1] = coefs
+    return full_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)
+        if np.all(cfs[..., m_indices] == 0):
+            return np.delete(l_arr, l_idx), np.delete(cfs, m_indices, axis=-1)
+        return l_arr, cfs
+    for l in l_array:
+        l_array, coefs = delete_zero_entries(l, l_array, coefs)
+    return l_array, coefs
+
+
+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)
+
+
+def expansion_rotation(rotations: Quaternion, coefs: Array, l_array: Array):
+    """
+    General function for rotations of expansion coefficients using WignerD matrices. Combines all rotations
+    with each expansion given in coefs array.
+    :param rotations: Quaternion array, last dimension is 4
+    :param coefs: array of expansion coefficients
+    :param l_array: array of all ell values of the expansion
+    :return rotated coefficients, output shape is rotations.shape[:-1] + coefs.shape
+    """
+    rot_arrays = rotations.ndarray.reshape((-1, 4))
+    coefs_reshaped = coefs.reshape((-1, coefs.shape[-1]))
+    wigner_matrices = spherical.Wigner(np.max(l_array)).D(rot_arrays)
+    new_coefs = np.zeros((rot_arrays.shape[0],) + coefs_reshaped.shape, dtype=np.complex128)
+    for i, l in enumerate(l_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 = m_indices_at_l(l_array, i)
+        wm = wigner_matrices[:, Dlmn_slice].reshape((-1, 2*l+1, 2*l+1))
+        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)

functions.py → charged_shells/functions.py

@@ -2,9 +2,6 @@ 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)
 
@@ -17,14 +14,17 @@ 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 interaction_coef_C(l, p, x):
+    return x * sps.iv(l + 1 / 2, x) * sps.iv(p + 1 / 2, x)
 
 
-def coeff_C_diff(l, x):
-    return 1 / (x * sps.iv(l + 1 / 2, x) * sps.kv(l + 3 / 2, x))
+def coefficient_Cpm(l, x):
+    return x * sps.kv(l + 1 / 2, x) * sps.iv(l + 1 / 2, x)
+
 
+def coefficient_Cim(l, x):
+    return sps.kv(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, :]))
 
+def coef_C_diff(l, x):
+    return 1 / (x * sps.iv(l + 1 / 2, x) * sps.kv(l + 3 / 2, x))

charged_shells/interactions.py

@@ -0,0 +1,78 @@
+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 charged_shells.units_and_constants as uc
+
+
+Array = np.ndarray
+Expansion = expansion.Expansion
+ModelParams = parameters.ModelParams
+EnergyUnit = Literal['kT', 'eV', 'J']
+
+
+def energy_units(units: EnergyUnit, params: ModelParams) -> float:
+    match units:
+        case 'eV':
+            return 1 / (uc.CONSTANTS.e0 * uc.UNITS.voltage)
+        case 'kT':
+            return 1 / (params.temperature * uc.CONSTANTS.Boltzmann)
+        case 'J':
+            return uc.UNITS.energy
+
+
+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
+
+    full_l_array, full_m_array = ex1.lm_arrays
+
+    # determine indices of relevant elements in the sum
+    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]
+
+    charge_factor = np.real(ex1.coefs[..., indices_l] * np.conj(ex2.coefs[..., indices_p]) +
+                            (-1) ** (flat_l + flat_p) * ex1.coefs[..., indices_p] * np.conj(ex2.coefs[..., indices_l]))
+
+    all_s_array = np.arange(2 * ex1.max_l + 1)
+    bessels = fn.sph_bessel_k(all_s_array, params.kappa * dist)
+
+    # additional selection that excludes terms where Wigner 3j symbols are 0 by definition
+    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)
+
+    l_vals = flat_l[indices_lpm]
+    p_vals = flat_p[indices_lpm]
+    m_vals = flat_m[indices_lpm]
+    s_vals = all_s_array[indices_s]
+    bessel_vals = bessels[indices_s]
+
+    # While all other arrays are 1D, this one can have extra leading axes corresponding to leading dimensions
+    # of expansion coefficients. The last dimension is the same as other arrays.
+    charge_vals = charge_factor[..., indices_lpm]
+
+    lps_terms = (2 * s_vals + 1) * np.sqrt((2 * l_vals + 1) * (2 * p_vals + 1))
+
+    # the same combination of l, s, and p is repeated many times
+    _, unique_indices1, inverse1 = np.unique(np.stack((l_vals, s_vals, p_vals)),
+                                             axis=1, return_inverse=True, return_index=True)
+    wigner1 = wigner3j(2 * l_vals[unique_indices1], 2 * s_vals[unique_indices1], 2 * p_vals[unique_indices1],
+                       0, 0, 0)[inverse1]
+
+    # all the combinations (l, s, p, m) are unique
+    wigner2 = wigner3j(2 * l_vals, 2 * s_vals, 2 * p_vals,
+                       -2 * m_vals, 0, 2 * m_vals)
+
+    constants = params.R ** 2 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0) * energy_units(units, params)
+
+    C_vals = fn.interaction_coef_C(l_vals, p_vals, params.kappaR)
+    lspm_vals = C_vals * (-1) ** (l_vals + m_vals) * lps_terms * bessel_vals * wigner1 * wigner2
+    broadcasted_lspm_vals = np.broadcast_to(lspm_vals, charge_vals.shape)
+
+    return 0.5 * constants * np.sum(broadcasted_lspm_vals * charge_vals, axis=-1)

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

charged_shells/parameters.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+from dataclasses import dataclass
+import numpy as np
+import charged_shells.units_and_constants as uc
+
+
+Array = np.ndarray
+
+
+@dataclass(kw_only=True)
+class ModelParams:
+    R: float | Array
+    kappa: float | Array = None
+    kappaR: float | Array = None
+    c0: float | Array = None
+    epsilon: float = 80
+    temperature: float = 293
+
+    def __post_init__(self):
+        self.kappa, self.kappaR, self.c0 = screening_calculator(self.R, self.temperature, self.epsilon,
+                                                                self.kappa, self.kappaR, self.c0)
+
+
+def bjerrum(temp: float, epsilon: float) -> float:
+    return uc.CONSTANTS.e0 ** 2 / (4 * np.pi * epsilon * uc.CONSTANTS.epsilon0 * uc.CONSTANTS.Boltzmann * temp)
+
+
+def kappa_from_concentration(c0: float, temp: float, epsilon: float) -> float:
+    return np.sqrt(8 * np.pi * bjerrum(temp, epsilon) * c0)
+
+
+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))
+
+
+def screening_calculator(radius: float, temp: float, epsilon: float,
+                         kappa: float = None, kappaR: float = None, c0: float = None) -> (float, float, float):
+    if kappa is not None:
+        return kappa, kappa * radius, concentration_from_kappa(kappa, temp, epsilon)
+    elif kappaR is not None:
+        return kappaR / radius, kappaR, concentration_from_kappa(kappaR / radius, temp, epsilon)
+    elif c0 is not None:
+        kappa = kappa_from_concentration(c0, temp, epsilon)
+        return kappa, kappa * radius, c0
+    raise ValueError('One of the arguments kappa, kappaR or c0 should be different from None.')

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

charged_shells/potentials.py

@@ -0,0 +1,70 @@
+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
+from scipy.special import eval_legendre, kv
+
+
+Array = np.ndarray
+ModelParams = parameters.ModelParams
+Expansion = expansion.Expansion
+
+
+def charged_shell_potential(theta: Array | float,
+                            phi: Array | float,
+                            dist: float,
+                            ex: Expansion,
+                            params: ModelParams,
+                            print_idx: int = None) -> 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 params: ModelParams object specifying parameter values for the model
+    :param print_idx: if not None, print a single term for debugging purposes
+    """
+    theta, phi = np.broadcast_arrays(theta, phi)
+    angles_shape = theta.shape
+    theta = theta.reshape(-1)  # ensures that arrays are 1D
+    phi = phi.reshape(-1)
+
+    if not theta.shape == phi.shape:
+        raise ValueError('theta and phi arrays should have the same shape.')
+    l_array, m_array = ex.lm_arrays
+
+    dist = dist * params.R
+    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))
+    l_factors = ex.repeat_over_m(l_factors)
+
+    factors = l_factors[:, None] * ex.coefs[..., None] * fn.sph_harm(l_array[:, None], m_array[:, None], theta[None, :], phi[None, :])
+    if print_idx is not None:
+        print(l_array[print_idx], m_array[print_idx], np.real(factors[print_idx]))
+
+    pot = (1 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0) * np.real(np.sum(factors, axis=-2)))
+
+    return pot.reshape(ex.shape + angles_shape)
+
+
+def inverse_patchy_particle_potential(theta: Array | float,
+                                      dist: float,
+                                      abar: float,
+                                      Zc: float,
+                                      Zp: tuple[float, float],
+                                      params: ModelParams,
+                                      lmax: int = 20):
+    dist = dist * params.R
+
+    out0 = ((Zc + Zp[0] + Zp[1]) * (np.exp(params.kappaR) / (1. + params.kappaR)) *
+            (np.exp(-params.kappa * dist) / dist) / (params.epsilon * uc.CONSTANTS.epsilon0)) * params.R ** 2
+
+    for l in range(2, lmax, 2):
+        out1 = (Zp[0] * np.power(abar, l) * eval_legendre(l, np.cos(theta))
+                + Zp[1] * np.power(abar, l) * eval_legendre(l, np.cos(np.pi - theta)))
+        out0 += (((2 * l + 1.) * (kv(l + 0.5, params.kappa * dist) / kv(l + 1.5, params.kappaR))) * out1
+                 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0))
+
+    return out0

charged_shells/rotational_path.py

@@ -0,0 +1,186 @@
+import numpy as np
+from dataclasses import dataclass, field
+import quaternionic
+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
+Array = np.ndarray
+Expansion = expansion.Expansion
+
+
+@dataclass
+class PairRotationalPath:
+
+    rotations1: list[Quaternion] = field(default_factory=list)
+    rotations2: list[Quaternion] = field(default_factory=list)
+    x_axis: list[Array] = field(default_factory=list)
+    overlapping_last: bool = True
+    _default_x_axis: Array = None
+
+    def add(self, rotation1: Quaternion, rotation2: Quaternion, x_axis: Array = None):
+        rotation1, rotation2 = np.broadcast_arrays(rotation1, rotation2)
+        self.rotations1.append(Quaternion(rotation1))
+        self.rotations2.append(Quaternion(rotation2))
+        if x_axis is None:
+            x_axis = np.arange(len(rotation1)) if self._default_x_axis is None else self._default_x_axis
+        self.add_x_axis(x_axis)
+
+    def add_x_axis(self, x_axis: Array):
+        try:
+            last_x_val = self.x_axis[-1][-1]
+        except IndexError:
+            last_x_val = 0
+        if self.overlapping_last:
+            self.x_axis.append(x_axis + last_x_val)
+        else:
+            raise NotImplementedError('Currently only overlapping end points for x-axes are supported.')
+
+    def set_default_x_axis(self, default_x_axis: Array):
+        self._default_x_axis = default_x_axis
+
+    def add_euler(self, *, alpha1: Array = 0, beta1: Array = 0, gamma1: Array = 0,
+                  alpha2: Array = 0, beta2: Array = 0, gamma2: Array = 0,
+                  x_axis: Array = None):
+        R1_euler = quaternionic.array.from_euler_angles(alpha1, beta1, gamma1)
+        R2_euler = quaternionic.array.from_euler_angles(alpha2, beta2, gamma2)
+        self.add(Quaternion(R1_euler), Quaternion(R2_euler), x_axis)
+
+    def stack_rotations(self) -> (Quaternion, Quaternion):
+        return Quaternion(np.vstack(self.rotations1)), Quaternion(np.vstack(self.rotations2))
+
+    def stack_x_axes(self) -> Array:
+        return np.hstack(self.x_axis)
+
+
+@dataclass
+class PathEnergyPlot:
+    """Path comparison for a pair of charge distributions, possibly at different parameter values."""
+
+    expansion1: Expansion
+    expansion2: Expansion
+    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):
+            self.dist = np.array([self.dist])
+
+    def evaluate_energy(self):
+        energy = []
+        for dist in self.dist:
+            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):
+        rotations1, rotations2 = self.rot_path.stack_rotations()
+        self.expansion1.rotate(rotations1)
+        self.expansion2.rotate(rotations2)
+        return self.evaluate_energy()
+
+    def section_energies(self):
+        energy_list = []
+        for rot1, rot2 in zip(self.rot_path.rotations1, self.rot_path.rotations2):
+            self.expansion1.rotate(rot1)
+            self.expansion2.rotate(rot2)
+            energy_list.append(self.evaluate_energy())
+        return energy_list
+
+    def normalization(self, norm_euler_angles: dict):
+        if norm_euler_angles is None:
+            return np.array([1.])
+        self.expansion1.rotate_euler(alpha=norm_euler_angles.get('alpha1', 0),
+                                     beta=norm_euler_angles.get('beta1', 0),
+                                     gamma=norm_euler_angles.get('gamma1', 0))
+        self.expansion2.rotate_euler(alpha=norm_euler_angles.get('alpha2', 0),
+                                     beta=norm_euler_angles.get('beta2', 0),
+                                     gamma=norm_euler_angles.get('gamma2', 0))
+        return np.abs(self.evaluate_energy())
+
+    def evaluate_path(self, norm_euler_angles: dict = None):
+        energy = self.path_energy()
+        normalization = self.normalization(norm_euler_angles)
+        energy = energy / normalization[None, ...]
+        energy = energy.reshape(energy.shape[0], -1)
+        return np.squeeze(energy)
+
+    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)
+        x_axis = self.rot_path.stack_x_axes()
+
+        fig, ax = plt.subplots()
+        ax.axhline(y=0, c='k', linestyle=':')
+        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):
+        energy_list = self.section_energies()
+        normalization = self.normalization(norm_euler_angles)
+
+        fig, ax = plt.subplots()
+        ax.axhline(y=0, c='k', linestyle=':')
+        for x_axis, energy in zip(self.rot_path.x_axis, energy_list):
+            energy, norm = np.broadcast_arrays(energy, normalization)
+            ax.plot(x_axis, energy / norm)
+        # 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)
+        if save_as is not None:
+            plt.savefig(save_as, dpi=600)
+        plt.show()
+
+
+@dataclass
+class PathExpansionComparison:
+    """Path comparison for different charge distribution models."""
+
+    ex_list: list[Expansion]
+    rot_path: PairRotationalPath
+    dist: float
+    params: ModelParams
+    path_list: list[PathEnergyPlot] = field(default_factory=list)
+
+    def __post_init__(self):
+        for ex in self.ex_list:
+            self.path_list.append(PathEnergyPlot(ex, ex.clone(), self.rot_path, self.dist, self.params))
+
+    def evaluate_path(self, norm_euler_angles: dict = None):
+        path_vals = []
+        for path in self.path_list:
+            path_vals.append(path.evaluate_path(norm_euler_angles))
+        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)
+        x_axis = self.rot_path.stack_x_axes()
+
+        fig, ax = plt.subplots()
+        ax.axhline(y=0, c='k', linestyle=':')
+        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)
+        ax.set_ylabel('U', fontsize=13)
+        plt.tight_layout()
+        if save_as is not None:
+            plt.savefig(save_as, dpi=600)
+        plt.show()
+

charged_shells/units_and_constants.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class BaseUnits:
+    distance: float
+    charge: float
+    voltage: float
+    concentrationM: float
+
+    @property
+    def charge_density(self):
+        return self.charge / (self.distance ** 2)
+
+    @property
+    def energy(self):
+        return self.charge * self.voltage
+
+    @property
+    def concentration(self):
+        return self.concentrationM * 6.02214076 * 1e23 / (10 * self.distance) ** 3
+
+
+@dataclass
+class Constants:
+    Boltzmann: float
+    epsilon0: float
+    e0: float
+
+
+base_units = {'distance': 1e-9,
+              'charge': 1.602176634 * 1e-19,
+              'voltage': 1,
+              'concentrationM': 1e-3}
+
+
+UNITS = BaseUnits(**base_units)
+
+# values of different constants in provided base units
+CONSTANTS = Constants(Boltzmann=1.380649 * 1e-23 / UNITS.energy,
+                      epsilon0=8.8541878128 * 1e-12 * UNITS.distance * UNITS.voltage / UNITS.charge,
+                      e0=1.602176634 * 1e-19 / UNITS.charge)

expansion.py

@@ -1,84 +0,0 @@
-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)
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-

potentials.py

@@ -1,59 +0,0 @@
-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()

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

tests/wigner_test.py

@@ -0,0 +1,11 @@
+import numpy as np
+from py3nj import wigner3j
+
+
+if __name__ == '__main__':
+
+    w3j = wigner3j(0, 38, 38,
+                   0, 38, -38)
+
+    print(w3j)
+