5 Commits 9cc0694af0 ... 1672ad5541

Author SHA1 Message Date
  gnidovec 1672ad5541 Added Janus symmetry support, improved plots 9 months ago
  gnidovec 3502b543e0 Updated plots 1 year ago
  gnidovec 9cb152786e Path plot style now part of PairRotationalPath 1 year ago
  gnidovec 58c7cb09f8 Corrected spherical bessel definitions 1 year ago
  gnidovec 4a3b7730df 2024-01-16 1 year ago

+ 444 - 0
analysis/dip_path_plot.py

@@ -0,0 +1,444 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from charged_shells.rotational_path import PairRotationalPath, PathEnergyPlot
+from charged_shells import expansion, interactions
+from charged_shells.parameters import ModelParams
+from pathlib import Path
+import json
+from plot_settings import *
+
+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)
+pi_to_three_halves_pi = np.linspace(np.pi, 3 * np.pi / 2, 100, endpoint=True)
+
+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)
+
+DipolePath3 = PairRotationalPath()
+DipolePath3.set_default_x_axis(zero_to_pi_half)
+DipolePath3.add_euler(beta2=np.pi/2, beta1=zero_to_pi_half, start_name="EP", end_name="EE")
+DipolePath3.add_euler(beta2=pi_half_to_pi, beta1=pi_half_to_pi, end_name="PP")
+DipolePath3.add_euler(beta2=pi_half_to_pi[::-1], beta1=np.pi, end_name="EP")
+DipolePath3.add_euler(beta2=pi_half_to_pi, beta1=pi_half_to_pi[::-1], end_name="EP")
+DipolePath3.add_euler(beta1=np.pi/2, beta2=pi_half_to_pi[::-1], alpha2=np.pi/2, end_name="tEE")
+DipolePath3.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half[::-1], end_name="EE")
+DipolePath3.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half[::-1], end_name="EP")
+
+
+def sections_plot(kappaR: float = 3, abar: float = 0.5, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+
+    path_plot.plot_sections(save_as=save_as)
+
+
+def kappaR_dependence(kappaR: Array, abar: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=0)
+
+    path_plot.plot(labels=[rf'$\kappa R$={kR}' for kR in kappaR],
+                   # norm_euler_angles={'beta2': np.pi},
+                   save_as=save_as)
+
+
+def abar_dependence(abar: Array, kappaR: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+
+    path_plot.plot(labels=[rf'$\bar a$={a}' for a in abar],
+                   # norm_euler_angles={'beta2': np.pi},
+                   save_as=save_as)
+
+
+def sigma0_dependence(sigma0: Array, kappaR: float, abar: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=sigma0)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+
+    path_plot.plot(labels=[rf'$\eta={s0 / sigma_tilde}$' for s0 in sigma0],
+                   # norm_euler_angles={'beta2': np.pi},
+                   save_as=save_as)
+
+
+def model_comparison(config_data: dict, save_as=None, save_data=False):
+    kappaR = 3
+    params = ModelParams(R=150, kappaR=kappaR)
+    a_bar = 0.5
+    sigma_tilde = 0.001
+
+    ex1 = expansion.MappedExpansionDipole(a_bar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign()
+
+    # matching other models to the mapped CSp model based on equal patch size in potential
+    # ex_gauss = quadrupole_model_mappings.ic_to_gauss(sigma_tilde, a_bar, params, l_max=30, sigma0=0)
+    # ex_gauss2 = ex_gauss.clone()
+    # ex_cap = quadrupole_model_mappings.ic_to_cap(sigma_tilde, a_bar, params, l_max=30, sigma0=0)
+    # ex_cap2 = ex_cap.clone()
+
+    # path plots for all models
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params)
+    energy_inv = path_plot_inv.evaluate_path()
+
+    # path_plot_gauss = PathEnergyPlot(ex_gauss, ex_gauss2, DipolePath3, dist=2., params=params)
+    # energy_gauss = path_plot_gauss.evaluate_path()
+    #
+    # path_plot_cap = PathEnergyPlot(ex_cap, ex_cap2, DipolePath3, dist=2., params=params)
+    # energy_cap = path_plot_cap.evaluate_path()
+
+    # peak_energy_sanity_check
+    # ex1new = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    # ex2new = ex1new.clone()
+    # pp_energy = interactions.charged_shell_energy(ex1new, ex2new, params)
+    # print(f'PP energy: {pp_energy}')
+
+    # Emanuele data
+    em_data = np.load(Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_A").joinpath("pathway.npz"))['arr_0']
+    # em_data = np.load(Path(config_data["emanuele_data"]).joinpath("FIG_7").joinpath("pathway.npz"))['arr_0']
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE")
+    #                 .joinpath("FIX_A").joinpath("ECC_0.25"))
+    # em_data = np.load(em_data_path.joinpath(f"EMME_{kappaR}.").joinpath("pathway.npz"))['arr_0']
+    em_data, em_data_inv = em_data[:int(len(em_data) / 2)], em_data[int(len(em_data) / 2):]
+
+    # if save_data:
+    #     np.savez(Path(config_data["figure_data"]).joinpath(f"fig_5_janus_kR{kappaR}.npz"),
+    #              ICi=em_data,
+    #              CSp=np.stack((x_axis, np.squeeze(energy))).T,
+    #              CSp_gauss=np.stack((x_axis, np.squeeze(energy_gauss))).T,
+    #              CSp_cap=np.stack((x_axis, np.squeeze(energy_cap))).T)
+
+    fig, ax = plt.subplots(figsize=0.5 * np.array([8.25, 4.125]))
+    ax.plot(em_data[:, 0], em_data[:, 1], label='ICi', c='tab:blue')
+    ax.plot(em_data_inv[:, 0], em_data_inv[:, 1], ls='--', c='tab:blue')
+    ax.plot(x_axis, np.squeeze(energy), label='CSp', c='tab:orange')
+    ax.plot(x_axis, np.squeeze(energy_inv), ls='--', c='tab:orange')
+    # ax.plot(x_axis, np.squeeze(energy_gauss), label='CSp - Gauss')
+    # ax.plot(x_axis, np.squeeze(energy_cap), label='CSp - cap')
+    # ax.plot(x_axis, em_data[:, 1] / np.squeeze(energy), label='CSp')
+    path_plot.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_kappaR_dependence(config_data: dict, kappaR: list[int], abar: float, sigma_tilde=0.001, save_as=None):
+
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_C")
+                    .joinpath("FIX_A").joinpath(f"ECC_{np.round(abar/2, 4)}"))
+
+    ic_data = []
+    ic_data_inv = []
+    for kR in kappaR:
+        em_data = np.load(em_data_path.joinpath(f"EMME_{kR}.").joinpath("pathway.npz"))['arr_0']
+        em_data, em_data_inv = em_data[:int(len(em_data) / 2)], em_data[int(len(em_data) / 2):]
+        ic_data.append(em_data)
+        ic_data_inv.append(em_data_inv)
+
+    params = ModelParams(R=150, kappaR=np.asarray(kappaR))
+
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=0)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=0)
+    energy_inv = path_plot_inv.evaluate_path()
+
+    labels = [rf'$\kappa R = {kR}$' for kR in [1, 3, 10]]
+
+    fig, ax = plt.subplots()
+    for d, d_inv, en, en_inv, label, c in zip(ic_data, ic_data_inv, energy.T, energy_inv.T, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(d_inv[:, 0], d_inv[:, 1], c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+        ax.plot(x_axis, en_inv, ls='--', c=c)
+    DipolePath3.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+    
+    
+def combined_abar_dependence(config_data: dict, kappaR: int, abar: list[float], sigma_tilde=0.001, save_as=None):
+    
+    em_data_path = Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_B").joinpath("FIX_M")
+
+    ic_data = []
+    ic_data_inv = []
+    for ab in abar:
+        em_data = np.load(em_data_path.joinpath(f"ECC_{np.round(ab/2, 4)}").joinpath(f"EMME_{kappaR}.").joinpath("pathway.npz"))['arr_0']
+        em_data, em_data_inv = em_data[:int(len(em_data) / 2)], em_data[int(len(em_data) / 2):]
+        ic_data.append(em_data)
+        ic_data_inv.append(em_data_inv)
+
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionDipole(np.asarray(abar), params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_inv = path_plot_inv.evaluate_path()
+
+    labels = [rf'$\bar a={a}$' for a in [0.3, 0.4, 0.5]]
+
+    fig, ax = plt.subplots()
+    for d, d_inv, en, en_inv, label, c in zip(ic_data, ic_data_inv, energy.T, energy_inv.T, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(d_inv[:, 0], d_inv[:, 1], c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+        ax.plot(x_axis, en_inv, ls='--', c=c)
+    DipolePath3.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_sigma0_dependence(config_data: dict, kappaR=3., abar=0.5, sigma0=(-0.0002, 0.00, 0.0002), sigma_tilde=0.001, save_as=None):
+
+    em_data_path = Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_D").joinpath("CHANGE_ZC")
+    undercharged = np.load(em_data_path.joinpath("ZC_-56").joinpath("pathway.npz"))['arr_0']
+    overcharged = np.load(em_data_path.joinpath("ZC_56").joinpath("pathway.npz"))['arr_0']
+
+    neutral_path = Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_B").joinpath("FIX_M")
+    neutral = np.load(neutral_path.joinpath(f"ECC_{np.round(abar/2, 4)}").joinpath(f"EMME_{int(kappaR)}.").joinpath("pathway.npz"))['arr_0']
+
+    undercharged, undercharged_inv = undercharged[:int(len(undercharged) / 2)], undercharged[int(len(undercharged) / 2):]
+    overcharged, overcharged_inv = overcharged[:int(len(overcharged) / 2)], overcharged[int(len(overcharged) / 2):]
+    neutral, neutral_inv = neutral[:int(len(neutral) / 2)], neutral[int(len(neutral) / 2):]
+
+    ic_data = [undercharged, neutral, overcharged]
+    ic_data_inv = [undercharged_inv, neutral_inv, overcharged_inv]
+
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=np.asarray(sigma0))
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign(exclude_00=True)
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_inv = path_plot_inv.evaluate_path()
+
+    labels = [rf'$\eta={s0/sigma_tilde}$' for s0 in sigma0]
+
+    fig, ax = plt.subplots()
+    for d, d_inv, en, en_inv, label, c in zip(ic_data, ic_data_inv, energy.T, energy_inv.T, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(d_inv[:, 0], d_inv[:, 1], c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+        ax.plot(x_axis, en_inv, ls='--', c=c)
+    DipolePath3.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_all(config_data: dict, save_as=None):
+
+    sigma_tilde = 0.001
+    kappaR_list = [1, 3, 10]
+    abar_list = [0.5, 0.4, 0.3]
+    sigma0_list = [-0.0002, 0.00, 0.0002]
+    kappaR = 3
+    abar = 0.5
+
+    em_data_kappaR = (Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_C")
+                      .joinpath("FIX_A").joinpath(f"ECC_{np.round(abar / 2, 4)}"))
+    
+    ic_data_kappaR = []
+    ic_data_kappaR_inv = []
+    for kR in kappaR_list:
+        em_data = np.load(em_data_kappaR.joinpath(f"EMME_{kR}.").joinpath("pathway.npz"))['arr_0']
+        em_data, em_data_inv = em_data[:int(len(em_data) / 2)], em_data[int(len(em_data) / 2):]
+        ic_data_kappaR.append(em_data)
+        ic_data_kappaR_inv.append(em_data_inv)
+
+    params = ModelParams(R=150, kappaR=np.asarray(kappaR_list))
+
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=0)
+    energy_kappaR = path_plot.evaluate_path()
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=0)
+    energy_kappaR_inv = path_plot_inv.evaluate_path()
+    x_axis_kappaR = path_plot.rot_path.stack_x_axes()
+    labels_kappaR = [rf'$\kappa R={kR}$' for kR in [1, 3, 10]]
+
+    em_data_abar = Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_B").joinpath("FIX_M")
+
+    ic_data_abar = []
+    ic_data_abar_inv = []
+    for ab in abar_list:
+        em_data = np.load(
+            em_data_abar.joinpath(f"ECC_{np.round(ab / 2, 4)}").joinpath(f"EMME_{kappaR}.").joinpath("pathway.npz"))[
+            'arr_0']
+        em_data, em_data_inv = em_data[:int(len(em_data) / 2)], em_data[int(len(em_data) / 2):]
+        ic_data_abar.append(em_data)
+        ic_data_abar_inv.append(em_data_inv)
+
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionDipole(np.asarray(abar_list), params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_abar = path_plot.evaluate_path()
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_abar_inv = path_plot_inv.evaluate_path()
+    x_axis_abar = path_plot.rot_path.stack_x_axes()
+    labels_abar = [rf'$\bar a={a}$' for a in abar_list]
+
+    em_data_charge = Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_D").joinpath("CHANGE_ZC")
+    undercharged = np.load(em_data_charge.joinpath("ZC_-56").joinpath("pathway.npz"))['arr_0']
+    overcharged = np.load(em_data_charge.joinpath("ZC_56").joinpath("pathway.npz"))['arr_0']
+    neutral_path = Path(config_data["emanuele_data"]).joinpath("FIG_5_JANUS").joinpath("FIG_5_JANUS_B").joinpath("FIX_M")
+    neutral = np.load(
+        neutral_path.joinpath(f"ECC_{np.round(abar / 2, 4)}").joinpath(f"EMME_{int(kappaR)}.").joinpath("pathway.npz"))[
+        'arr_0']
+    undercharged, undercharged_inv = undercharged[:int(len(undercharged) / 2)], undercharged[
+                                                                                int(len(undercharged) / 2):]
+    overcharged, overcharged_inv = overcharged[:int(len(overcharged) / 2)], overcharged[int(len(overcharged) / 2):]
+    neutral, neutral_inv = neutral[:int(len(neutral) / 2)], neutral[int(len(neutral) / 2):]
+
+    ic_data_sigma0 = [overcharged, neutral, undercharged]
+    ic_data_sigma0_inv = [overcharged_inv, neutral_inv, undercharged_inv]
+
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex1 = expansion.MappedExpansionDipole(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=np.asarray(sigma0_list))
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign()
+
+    path_plot = PathEnergyPlot(ex1, ex2, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_sigma0 = path_plot.evaluate_path()
+    path_plot_inv = PathEnergyPlot(ex1, ex3, DipolePath3, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_sigma0_inv = path_plot_inv.evaluate_path()
+    x_axis_sigma0 = path_plot.rot_path.stack_x_axes()
+    labels_sigma0 = [rf'$\eta={s0/sigma_tilde}$' for s0 in sigma0_list]
+
+    fig, axs = plt.subplots(3, 1, figsize=(6, 7.8))
+    for d, d_inv, en, en_inv, label, c in zip(ic_data_kappaR, ic_data_kappaR_inv, energy_kappaR.T, energy_kappaR_inv.T, labels_kappaR, COLOR_LIST):
+        axs[0].set_title('Screening', fontsize=15)
+        axs[0].plot(d[:, 0], d[:, 1], label=label, c=c)
+        axs[0].plot(x_axis_kappaR, en, ls='--', c=c)
+        axs[0].plot(d_inv[:, 0], d_inv[:, 1], c=c)
+        axs[0].plot(x_axis_kappaR, en_inv, ls='--', c=c)
+        DipolePath3.plot_style(fig, axs[0], size=None)
+        axs[0].get_legend().set_bbox_to_anchor((0.65, 1.07))
+    for d, d_inv, en, en_inv, label, c in zip(ic_data_abar, ic_data_abar_inv, energy_abar.T, energy_abar_inv.T, labels_abar, COLOR_LIST):
+        axs[1].set_title('Asymmetry', fontsize=15)
+        axs[1].plot(d[:, 0], d[:, 1], label=label, c=c)
+        axs[1].plot(x_axis_abar, en, ls='--', c=c)
+        axs[1].plot(d_inv[:, 0], d_inv[:, 1], c=c)
+        axs[1].plot(x_axis_abar, en_inv, ls='--', c=c)
+        DipolePath3.plot_style(fig, axs[1], size=None)
+        axs[1].get_legend().set_bbox_to_anchor((0.65, 1.05))
+    for d, d_inv, en, en_inv, label, c in zip(ic_data_sigma0, ic_data_sigma0_inv, energy_sigma0.T, energy_sigma0_inv.T, labels_sigma0, COLOR_LIST):
+        axs[2].set_title('Net charge', fontsize=15)
+        axs[2].plot(d[:, 0], d[:, 1], label=label, c=c)
+        axs[2].plot(x_axis_sigma0, en, ls='--', c=c)
+        axs[2].plot(d_inv[:, 0], d_inv[:, 1], c=c)
+        axs[2].plot(x_axis_sigma0, en_inv, ls='--', c=c)
+        DipolePath3.plot_style(fig, axs[2], size=None)
+        axs[2].get_legend().set_bbox_to_anchor((0.65, 1.04))
+    for ax in axs:
+        ax.yaxis.set_label_coords(-0.09, 0.5)
+    # axs[-1].set_xlabel('rotational path', fontsize=15)
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+if __name__ == '__main__':
+
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    # sections_plot(save_as=Path("/home/andraz/ChargedShells/Figures/dipole_test_path.png"))
+
+    # kappaR_dependence(np.array([3, 5]), 0.5,
+    #                   # save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_kappaR_dep.png")
+    #                   )
+
+    # abar_dependence(np.array([0.3, 0.4, 0.5]), 3,
+    #                 save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_abar_dep.png")
+    #                 )
+
+    # sigma0_dependence(np.array([-0.0002, 0.00, 0.0002]), 3, 0.5,
+    #                   save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_charge_dep_abar05_kappaR3.png")
+    #                   )
+
+    # model_comparison(config_data,
+    #                  # save_as=Path(config_data["figures"]).joinpath("Emanuele_data").joinpath('IC_CS_janus_path.pdf')
+    #                  )
+
+    # combined_kappaR_dependence(config_data, kappaR=[1, 3, 10], abar=0.5,
+    #                      # save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('janus_kappaR_dep.png')
+    #                      )
+    #
+    # combined_abar_dependence(config_data, kappaR=3, abar=[0.3, 0.4, 0.5],
+    #                          # save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('janus_abar_dep.png')
+    #                          )
+    #
+    # combined_sigma0_dependence(config_data,
+    #                      # save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('janus_charge_dep.png')
+    #                      )
+    #
+    combined_all(config_data,
+                 save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('janus_combined_dep.png')
+                 )

+ 301 - 0
analysis/energy_gap.py

@@ -0,0 +1,301 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from charged_shells import expansion, interactions, mapping
+from charged_shells.parameters import ModelParams
+from functools import partial
+from pathlib import Path
+from matplotlib import cm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from matplotlib.colors import TwoSlopeNorm
+from matplotlib.ticker import FuncFormatter
+import json
+
+Expansion = expansion.Expansion
+
+
+def energy_gap(ex1: Expansion, params: ModelParams, dist=2., match_expansion_axis_to_params=None):
+    ex2 = ex1.clone()
+    ex3 = ex1.clone()
+
+    ex2.rotate_euler(alpha=0, beta=np.pi/2, gamma=0)  # to get EP config between ex1 and ex2
+
+    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist),
+                                                     match_expansion_axis_to_params)
+    energy_ep = energy_fn(ex1, ex2, params)
+    energy_pp = energy_fn(ex1, ex3, params)
+
+    return (energy_pp - energy_ep) / energy_pp
+
+
+def abar_kappaR_dependence(save_as=None):
+
+    kappaR = np.linspace(0.01, 25, 25)
+    a_bar = np.array([0.2, 0.4, 0.6, 0.8])
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar[:, None], kappaR[None, :], 0.001)
+    # ex = expansion.MappedExpansionQuad(a_bar, kappaR, 0.001)
+    gap = energy_gap(ex, params, match_expansion_axis_to_params=1)
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for g, lbl in zip(gap, [rf'$\bar a={a}$' for a in a_bar]):
+        ax.plot(kappaR, g, label=lbl)
+    ax.legend(fontsize=17)
+    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'$\frac{V_{pp}-V_{ep}}{V_{pp}}$', fontsize=20)
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+def abar_kappaR_dependence2(save_as=None):
+
+    kappaR = np.array([1, 3, 10, 30])
+    a_bar = np.linspace(0.2, 0.8, 30)
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar[:, None], kappaR[None, :], 0.001)
+    # ex = expansion.MappedExpansionQuad(a_bar, kappaR, 0.001)
+    gap = energy_gap(ex, params, match_expansion_axis_to_params=1)
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for g, lbl in zip(gap.T, [rf'$\kappa R={kR}$' for kR in kappaR]):
+        ax.plot(a_bar, g, label=lbl)
+    ax.legend(fontsize=17)
+    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'$\frac{V_{pp}-V_{ep}}{V_{pp}}$', fontsize=20)
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+def charge_kappaR_dependence(a_bar, min_charge, max_charge, save_as=None, cmap=cm.jet):
+
+    kappaR = np.linspace(0.01, 10, 50)
+    sigma_tilde = 0.001
+    params = ModelParams(R=150, kappaR=kappaR)
+    charge = np.linspace(min_charge, max_charge, 100)
+    ex = expansion.MappedExpansionQuad(a_bar, kappaR, sigma_tilde, sigma0=charge)
+    # ex = expansion.MappedExpansionQuad(a_bar, kappaR, 0.001)
+    gap = energy_gap(ex, params, match_expansion_axis_to_params=0)
+
+    colors = cmap(np.linspace(0, 1, len(charge)))
+    sm = cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=np.min(charge) / sigma_tilde,
+                                                         vmax=np.max(charge) / sigma_tilde))
+    sm.set_array([])
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for g, c in zip(gap.T, colors):
+        ax.plot(kappaR, g, c=c)
+    # ax.legend(fontsize=17)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel(r'$\kappa R$', fontsize=20)
+    ax.set_ylabel(r'$(V_{pp}-V_{ep})/V_{pp}$', fontsize=20)
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    cax.tick_params(labelsize=15)
+    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
+    cbar.set_label(r'$\eta$', rotation=0, labelpad=15, fontsize=20)
+
+    # plt.tight_layout()
+    plt.subplots_adjust(left=0.1, right=0.9, top=0.95, bottom=0.12)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+def charge_kappaR_dependence_heatmap(a_bar, min_charge, max_charge, save_as=None, cmap=cm.jet):
+
+    kappaR = np.linspace(0.01, 10, 50)
+    params = ModelParams(R=150, kappaR=kappaR)
+    charge = np.linspace(min_charge, max_charge, 100)
+    ex = expansion.MappedExpansionQuad(a_bar, kappaR, 0.001, sigma0=charge)
+    # ex = expansion.MappedExpansionQuad(a_bar, kappaR, 0.001)
+    gap = energy_gap(ex, params, match_expansion_axis_to_params=0)
+
+    norm = TwoSlopeNorm(vmin=np.min(gap), vcenter=0, vmax=np.max(gap))
+    sm = cm.ScalarMappable(cmap=cmap, norm=norm)
+    sm.set_array([])
+
+    def y_formatter(x, pos):
+        return f"{charge[int(x)-1]:.2f}"
+
+    def x_formatter(x, pos):
+        return f"{kappaR[int(x)-1]:.2f}"
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(1))
+    ax.imshow(gap.T, cmap=cmap, origin='lower',
+              # extent=[kappaR.min(), kappaR.max(), charge.min(), charge.max()]
+              )
+    # ax.legend(fontsize=17)
+    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'$\tilde \sigma_0$', fontsize=15)
+
+    plt.gca().xaxis.set_major_formatter(FuncFormatter(x_formatter))
+    plt.gca().yaxis.set_major_formatter(FuncFormatter(y_formatter))
+
+    # ax.set_xticks(kappaR)
+    # ax.set_yticks(charge)
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
+    cbar.set_label(r'$\frac{V_{pp}-V_{ep}}{V_{pp}}$', rotation=90, labelpad=20, fontsize=12)
+
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+def IC_gap_plot(config_data: dict, save_as=None):
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_10"))
+    em_data = np.load(em_data_path.joinpath("relative_gap.npz"))
+    for k in list(em_data.keys()):
+        data = em_data[k]
+        print(k, data.shape)
+        for i in range(3):
+            print(np.unique(data[:, i]))
+        print('\n')
+
+
+def IC_gap_kappaR(config_data: dict, save_as=None):
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_10"))
+    em_data = np.load(em_data_path.joinpath("relative_gap.npz"))
+    data = em_data['fixA']
+
+    print(data)
+
+    sort = np.argsort(data[:, 2])
+    xdata = data[:, 2][sort]
+    ydata = data[:, 3][sort]
+
+    plt.plot(xdata, ydata)
+    plt.xlabel('kappaR')
+    plt.ylabel('gap')
+    plt.show()
+
+
+def IC_gap_abar(config_data: dict, save_as=None):
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_10"))
+    em_data = np.load(em_data_path.joinpath("relative_gap.npz"))
+    data = em_data['fixM']
+
+    print(data)
+
+    sort = np.argsort(data[:, 1])
+    xdata = data[:, 1][sort]
+    ydata = data[:, 3][sort]
+
+    plt.plot(xdata, ydata)
+    plt.xlabel('abar')
+    plt.ylabel('gap')
+    plt.show()
+
+
+def IC_gap_charge_at_abar(a_bar, config_data: dict, save_as=None, cmap=cm.coolwarm, which_change='changezp',
+                          eta_min: float = None, eta_max: float = None):
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_10"))
+    em_data = np.load(em_data_path.joinpath("relative_gap_ZC.npz"))
+    data = em_data[which_change]
+
+    sigma_tilde = 0.001
+
+    relevant_indices = data[:, 1] == a_bar
+    if not np.any(relevant_indices):
+        raise ValueError(f'No results for given a_bar = {a_bar}. Possible values: {np.unique(data[:, 1])}')
+
+    data = data[relevant_indices]
+
+    charge, inverse, counts = np.unique(data[:, 0], return_counts=True, return_inverse=True)
+    # print(f'All charge: {charge}')
+
+    eta = charge / sigma_tilde
+
+    if eta_min is None:
+        eta_min = np.min(eta)
+    if eta_max is None:
+        eta_max = np.max(eta)
+
+    def map_eta_to_unit(x):
+        return (x - eta_min) / (eta_max - eta_min)
+
+    # print(eta[0], eta[1])
+    # print(map_eta_to_unit(eta[0]), map_eta_to_unit(eta[-1]))
+
+    colors_linspace = np.linspace(map_eta_to_unit(eta[0]), map_eta_to_unit(eta[-1]), len(charge))
+    colors_linspace[colors_linspace > 1] = 1
+    colors_linspace[colors_linspace < 0] = 0
+
+    colors = cmap(colors_linspace)
+    sm = cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=eta_min, vmax=eta_max))
+    sm.set_array([])
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for i, c in enumerate(colors):
+        idx, = np.nonzero(inverse == i)
+        kR = data[idx, 2]
+        gap = data[idx, 3]
+        sort = np.argsort(kR)
+        ax.plot(kR[sort], gap[sort], c=c)
+
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel(r'$\kappa R$', fontsize=20)
+    ax.set_ylabel(r'$(V_{pp}-V_{ep})/V_{pp}$', fontsize=20)
+    ax.set_xlim(-0.25, 10.25)
+    ax.set_ylim(-4, 5)
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    cax.tick_params(labelsize=15)
+    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
+    cbar.set_label(r'$\eta$', rotation=0, labelpad=15, fontsize=20)
+
+    # plt.tight_layout()
+    plt.subplots_adjust(left=0.1, right=0.9, top=0.95, bottom=0.12)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+def test_gap(a_bar, kappaR, charge):
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar, kappaR, 0.001, sigma0=charge)
+    gap = energy_gap(ex, params, match_expansion_axis_to_params=None)
+    print(gap)
+
+
+def main():
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    # test_gap(0.3, 10, charge=-0.003)
+
+    # abar_kappaR_dependence(Path("/home/andraz/ChargedShells/Figures/full_amplitude_kappaR_dep.png"))
+    # abar_kappaR_dependence2(Path("/home/andraz/ChargedShells/Figures/full_amplitude_abar_dep.png"))
+
+    # charge_kappaR_dependence(a_bar=0.8, min_charge=-0.002, max_charge=0.002,
+    #                          save_as=Path("/home/andraz/ChargedShells/Figures/full_amplitude_charge_abar08.png"),
+    #                          cmap=cm.coolwarm)
+
+    # charge_kappaR_dependence_heatmap(a_bar=0.5, min_charge=-0.003, max_charge=0.003,
+    #                          save_as=Path("/home/andraz/ChargedShells/Figures/full_amplitude_heatmap_abar05.png"),
+    #                          cmap=cm.bwr)
+
+    # IC_gap_plot(config_data)
+
+    # IC_gap_kappaR(config_data)
+    # IC_gap_abar(config_data)
+
+    IC_gap_charge_at_abar(0.3, config_data, which_change='changezc', eta_min=-2, eta_max=2,
+                          save_as=Path("/home/andraz/ChargedShells/Figures/Emanuele_data/IC_full_amplitude_charge_abar03.png")
+                          )
+
+
+if __name__ == '__main__':
+
+    main()

+ 25 - 0
analysis/expanison_coef_export.py

@@ -0,0 +1,25 @@
+from charged_shells import expansion, parameters
+from pathlib import Path
+import json
+import numpy as np
+
+
+def main():
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    abar = np.array([0.2, 0.3, 0.4, 0.5, 0.5, 0.7, 0.8])
+    params = parameters.ModelParams(kappaR=3, R=150)
+    ex = expansion.MappedExpansionDipole(abar, params.kappaR, 0.001, 30)
+    print(ex.coefs.shape)
+    l_arr, m_arr = ex.lm_arrays
+    lm = np.stack((l_arr, m_arr)).T
+    print(lm.shape)
+
+    np.savez(Path(config_data["expansion_data"]).joinpath("janus.npz"), abar=abar, lm=lm, coefs=ex.coefs)
+
+
+
+if __name__ == '__main__':
+
+    main()

+ 29 - 8
analysis/expansion_plot.py

@@ -1,8 +1,10 @@
-from charged_shells import expansion
+from charged_shells import expansion, parameters
 import numpy as np
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.pyplot as plt
 from pathlib import Path
 from pathlib import Path
 import plotly.graph_objects as go
 import plotly.graph_objects as go
+import json
+import quadrupole_model_mappings
 
 
 Expansion = expansion.Expansion
 Expansion = expansion.Expansion
 
 
@@ -30,7 +32,7 @@ def plot_theta_profile_multiple(ex_list: list[Expansion], label_list, phi: float
     plt.show()
     plt.show()
 
 
 
 
-def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100):
+def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100, save_as: Path = None):
     theta = np.linspace(0, np.pi, num_theta)
     theta = np.linspace(0, np.pi, num_theta)
     phi = np.linspace(0, 2 * np.pi, num_phi)
     phi = np.linspace(0, 2 * np.pi, num_phi)
     theta, phi = np.meshgrid(theta, phi)
     theta, phi = np.meshgrid(theta, phi)
@@ -42,21 +44,40 @@ def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100):
     z = np.cos(theta)
     z = np.cos(theta)
 
 
     # Create a heatmap on the sphere
     # Create a heatmap on the sphere
-    fig = go.Figure(data=go.Surface(x=x, y=y, z=z, surfacecolor=r, colorscale='Jet'))
+    fig = go.Figure(data=go.Surface(x=x, y=y, z=z, surfacecolor=r,
+                                    colorscale='RdBu', reversescale=True))
     fig.update_layout(scene=dict(aspectmode='data'))
     fig.update_layout(scene=dict(aspectmode='data'))
-    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
+    fig.update_layout(scene=dict(xaxis_title='', yaxis_title='', zaxis_title=''))
+
+    # Remove axes planes, background, ticks, and labels
+    fig.update_layout(scene=dict(xaxis=dict(showbackground=False, gridcolor='white', showticklabels=False, ticks=''),
+                                 yaxis=dict(showbackground=False, gridcolor='white', showticklabels=False, ticks=''),
+                                 zaxis=dict(showbackground=False, gridcolor='white', showticklabels=False, ticks='')))
+
+    # Adjust the width and height for higher resolution
+    fig.update_layout(width=1200, height=1200)
+
+    # Save as PNG with higher resolution
+    if save_as is not None:
+        fig.write_image(save_as, scale=3)  # Adjust the scale as needed
 
 
     fig.show()
     fig.show()
 
 
 
 
 def main():
 def main():
-    # ex = expansion.MappedExpansionQuad(0.44, 3, 1, 10)
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    params = parameters.ModelParams(kappaR=3, R=150)
+    # ex = expansion.MappedExpansionQuad(0.328, params.kappaR, 0.001, 30)
     # ex = expansion.Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
     # 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)
+    # ex = expansion.GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=2.676, sigma1=0.00044, l_max=30)
+    # ex = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.894, 0.00132, 50)
+    # ex = quadrupole_model_mappings.ic_to_gauss(0.001, 0.328, params, l_max=30)
+    ex = quadrupole_model_mappings.ic_to_cap(0.001, 0.328, params, l_max=50)
     # print(np.real(ex.coefs))
     # print(np.real(ex.coefs))
     # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
     # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
-    plot_charge_3d(ex)
+    plot_charge_3d(ex, save_as=Path(config_data["figures"]).joinpath("model_3D_cap.png"))
 
 
     # new_coeffs = expanison.expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
     # new_coeffs = expanison.expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
     # print(new_coeffs.shape)
     # print(new_coeffs.shape)

+ 218 - 37
analysis/patch_shape_comparison.py

@@ -1,63 +1,91 @@
-from charged_shells import expansion, functions as fn, potentials, patch_size
+from charged_shells import expansion, potentials, patch_size
 import numpy as np
 import numpy as np
 from charged_shells.parameters import ModelParams
 from charged_shells.parameters import ModelParams
-import matplotlib.pyplot as plt
-import scipy.special as sps
 from pathlib import Path
 from pathlib import Path
+from matplotlib.lines import Line2D
+import json
+import quadrupole_model_mappings
+from plot_settings import *
 
 
 
 
-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 ic_cs_comparison():
+    kappaR = 15
+    params = ModelParams(R=150, kappaR=kappaR)
+    sigma_m = 0.001
+    sigma0 = 0.0001
+    a_bar = 0.3
 
 
+    theta = np.linspace(0, np.pi, 1001)
+    phi = 0.
+    dist = 1
 
 
-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))))
+    ex_cp = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30,
+                                          sigma0=sigma0)
+    potential_cp = potentials.charged_shell_potential(theta, phi, dist, ex_cp, params)
 
 
+    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m + sigma0,
+                                                                (sigma_m, sigma_m), params, 30)
 
 
-if __name__ == '__main__':
+    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 * potential_cp.T, label='CSp - mapped', linewidth=2)
+    # ax.plot(1000 * potential_ic.T - 1000 * potential_cp.T)
+    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.xticks(custom_ticks, custom_labels, fontsize=13)
+    plt.legend(fontsize=12)
+    plt.tight_layout()
+    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison_overcharged05.pdf"), dpi=300)
+    plt.show()
 
 
+
+def models_comparison(save_data=False):
     target_patch_size = 0.92
     target_patch_size = 0.92
-    params = ModelParams(R=150, kappaR=3)
+    kappaR = 3
+    params = ModelParams(R=150, kappaR=kappaR)
     sigma_m = 0.001
     sigma_m = 0.001
+    sigma0 = 0  # taking this total charge parameter nonzero causes cap model to fail in finding the appropriate theta0
 
 
-    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]]))
+    def fn(x):
+        return expansion.MappedExpansionQuad(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30, sigma0=sigma0)
 
 
-    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)
+    a_bar = patch_size.inverse_potential_patch_size(target_patch_size, fn, 0.5, params)
+    ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30, sigma0=sigma0)
 
 
-    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)
+    ex_gauss = quadrupole_model_mappings.ic_to_gauss(sigma_m, a_bar, params, l_max=30, sigma0=sigma0)
+    ex_cap = quadrupole_model_mappings.ic_to_cap(sigma_m, a_bar, params, l_max=50, sigma0=sigma0)
 
 
     theta = np.linspace(0, np.pi, 1001)
     theta = np.linspace(0, np.pi, 1001)
     phi = 0.
     phi = 0.
     dist = 1
     dist = 1
 
 
-    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m, (sigma_m, sigma_m), params, 30)
+    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m + sigma0,
+                                                                (sigma_m, sigma_m), params, 30)
     potential1 = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)
     potential1 = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)
     potential2 = potentials.charged_shell_potential(theta, phi, dist, ex_gauss, params)
     potential2 = potentials.charged_shell_potential(theta, phi, dist, ex_gauss, params)
     potential3 = potentials.charged_shell_potential(theta, phi, dist, ex_cap, params)
     potential3 = potentials.charged_shell_potential(theta, phi, dist, ex_cap, params)
     # print(potential.shape)
     # print(potential.shape)
     # print(potential)
     # print(potential)
 
 
+    if save_data:
+        with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+            config_data = json.load(config_file)
+        np.savez(Path(config_data["figure_data"]).joinpath("fig_6_shape_models.npz"),
+                 ICi=np.stack((theta, potential_ic)).T,
+                 CSp_mapped=np.stack((theta, potential1)).T,
+                 CSp_gauss=np.stack((theta, potential2)).T,
+                 CSp_caps=np.stack((theta, potential3)).T)
+
     # expansion.plot_theta_profile_multiple([ex_point, ex_gauss, ex_cap], ['IC', 'Gauss', 'cap'], num=1000)
     # expansion.plot_theta_profile_multiple([ex_point, ex_gauss, ex_cap], ['IC', 'Gauss', 'cap'], num=1000)
 
 
     fig, ax = plt.subplots()
     fig, ax = plt.subplots()
-    ax.scatter(theta[::50], 1000 * potential_ic.T[::50], marker='o', label='ICi', facecolors='none', edgecolors='tab:red')
+    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, 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, 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 * potential2.T, label='CSp - Gauss', linewidth=2, ls='--')
@@ -75,10 +103,163 @@ if __name__ == '__main__':
     plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison.pdf"), dpi=300)
     plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison.pdf"), dpi=300)
     plt.show()
     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"))
+
+def abar_comparison():
+    target_patch_sizes = [0.8, 0.85, 0.92]
+    params = ModelParams(R=150, kappaR=3)
+    sigma_m = 0.001
+
+    theta = np.linspace(0, np.pi, 1001)
+    phi = 0.
+    dist = 1
+
+    def fn1(x):
+        return expansion.MappedExpansionQuad(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)
+
+    pots = []
+    for ps in target_patch_sizes:
+        a_bar = patch_size.inverse_potential_patch_size(ps, fn1, 0.5, params)
+        ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)
+        pots.append(potentials.charged_shell_potential(theta, phi, dist, ex_point, params))
+
+    fig, ax = plt.subplots(figsize=(4.125, 3))
+    for pot, ps in zip(pots, target_patch_sizes):
+        ax.plot(theta, 1000 * pot, label=fr'$\theta_0={180 / np.pi * ps:.1f}^o$', linewidth=2)
+    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.xticks(custom_ticks, custom_labels, fontsize=15)
+    plt.legend(fontsize=12)
+    plt.tight_layout()
+    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_amplitude_comparison.pdf"), dpi=300)
+    plt.show()
+
+
+def charge_comparsion():
+
+    charges = np.array([-0.001, 0, 0.001, 0.002])
+    a_bar = 0.6
+    params = ModelParams(R=150, kappaR=10)
+    sigma_m = 0.001
+
+    theta = np.linspace(0, np.pi, 1001)
+    phi = 0.
+    dist = 1
+
+    ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m,
+                                             l_max=30, sigma0=charges)
+    pot = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)
+
+    fig, ax = plt.subplots(figsize=(4.125, 3))
+    for p, lbl in zip(pot, [fr'$\sigma_0={c}$' for c in charges]):
+        ax.plot(theta, 1000 * p, linewidth=2, label=lbl)
+    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.xticks(custom_ticks, custom_labels, fontsize=15)
+    plt.legend(fontsize=12)
+    plt.tight_layout()
+    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_charge_comparison.pdf"), dpi=300)
+    plt.show()
+
+
+# TODO: comparison of patch shape at different kappaR and a neutral particle, with fixed patch size
+
+
+def ic_cs_comparison2(save_data=False):
+    # target_patch_size = 0.92
+    kappaR = 3
+    params = ModelParams(R=150, kappaR=kappaR)
+    sigma_m = 0.001
+    sigma0_array = np.array([-0.0002, 0, 0.0002])
+
+    theta = np.linspace(0, np.pi, 1201)
+    phi = 0.
+    dist = 1
+
+    def fn1(x):
+        return expansion.MappedExpansionQuad(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30, sigma0=0)
+    # a_bar is determined only for a neutral particle (patch size only well-defined in this case)
+    # a_bar = patch_size.inverse_potential_patch_size(target_patch_size, fn1, 0.5, params)
+    # print(a_bar)
+
+    a_bar = 0.5
+    target_patch_size = patch_size.potential_patch_size(expansion.MappedExpansionQuad(a_bar=a_bar,
+                                                                                      kappaR=params.kappaR,
+                                                                                      sigma_tilde=sigma_m,
+                                                                                      l_max=30,
+                                                                                      sigma0=0),
+                                                        params)
+
+    potential_ic = []
+    potential_cs = []
+    for sigma0 in sigma0_array:
+
+        ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR,
+                                                 sigma_tilde=sigma_m, l_max=30, sigma0=sigma0)
+
+        potential_ic.append(potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m + sigma0,
+                                                                    (sigma_m, sigma_m), params, 30))
+        potential_cs.append(potentials.charged_shell_potential(theta, phi, dist, ex_point, params))
+
+    if save_data:
+        with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+            config_data = json.load(config_file)
+        np.savez(Path(config_data["figure_data"]).joinpath("fig_6a.npz"),
+                 ICi_eta_neg=np.stack((theta, potential_ic[0])).T,
+                 ICi_eta_0=np.stack((theta, potential_ic[1])).T,
+                 ICi_eta_pos=np.stack((theta, potential_ic[2])).T,
+                 CSp_eta_neg=np.stack((theta, potential_cs[0])).T,
+                 CSp_eta_0=np.stack((theta, potential_cs[1])).T,
+                 CSp_eta_pos=np.stack((theta, potential_cs[2])).T)
+
+    fig, ax = plt.subplots(figsize=(3, 2.5))
+    lines = []
+    # lines.append(plt.scatter([], [], marker='x', c='k', label='CSp'))
+    # lines.append(plt.scatter([], [], marker='o', facecolor='none', edgecolor='k', label='ICi'))
+    for p_ic, p_cs, charge in zip(potential_ic, potential_cs, sigma0_array):
+        l, = ax.plot(theta, 1000 * p_cs.T, label=rf'$\eta = {charge/sigma_m}$', linewidth=2, zorder=0)
+        l2, = ax.plot(theta, 1000 * p_cs.T, linewidth=2, zorder=0, c='k', ls='--', dashes=(5, 5))
+        lines.append(l)
+        # ax.scatter(theta[::100], 1000 * p_ic.T[::100], marker='o', facecolors='none', edgecolors='k', s=50)
+        # ax.scatter(theta[::100], 1000 * p_cs.T[::100], marker='x', c='k', s=50)
+    # ax.plot(theta, potential_ic.T, label='IC', ls=':', linewidth=2, marker='o', markevery=50, mfc='none')
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    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=15)
+
+    # lines.append(Line2D([0], [0], color='k', label='CSp'))
+    # lines.append(plt.scatter([], [], marker='o', facecolor='none', edgecolor='red', label='ICi'))
+    plt.legend(handles=lines, fontsize=14, handlelength=0.8, frameon=False,
+               bbox_to_anchor=(0.57, 1.0), loc='upper center'
+               )
+
+    plt.tight_layout()
+    plt.savefig(Path("/home/andraz/ChargedShells/Figures/final_figures/potential_ic_cs_comparison.png"), dpi=300)
+    plt.show()
+
+
+def main():
+    # models_comparison(save_data=False)
+    # ic_cs_comparison()
+    ic_cs_comparison2()
+
+    # abar_comparison()
+    # charge_comparsion()
+
+
+if __name__ == '__main__':
+
+    main()
 
 

+ 55 - 23
analysis/patch_size_dependence.py

@@ -1,30 +1,49 @@
-from matplotlib import pyplot as plt
 from pathlib import Path
 from pathlib import Path
 import numpy as np
 import numpy as np
 from charged_shells import expansion, patch_size
 from charged_shells import expansion, patch_size
 from charged_shells.parameters import ModelParams
 from charged_shells.parameters import ModelParams
+import time
+import json
+from plot_settings import *
 
 
 
 
-def plot_abar_dependence(*, save=False):
+def plot_abar_dependence(*, save=False, save_data=False):
     a_bar = np.linspace(0.2, 0.8, 101)
     a_bar = np.linspace(0.2, 0.8, 101)
-    kappaR = np.array([0.1, 1, 3, 10, 30])
+    kappaR = np.array([1, 3, 10])
     params = ModelParams(R=150, kappaR=kappaR)
     params = ModelParams(R=150, kappaR=kappaR)
     ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])
     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)
     ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=1)
 
 
+    marked_point_ex = expansion.MappedExpansionQuad(a_bar=0.5, sigma_tilde=0.001, l_max=30, kappaR=3)
+    mark_ps = patch_size.potential_patch_size(marked_point_ex, ModelParams(R=150, kappaR=3),
+                                              match_expansion_axis_to_params=None)
+
     markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']
     markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']
 
 
-    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    if save_data:
+        data_dict = {}
+        for key, patch in zip(["kR=01", "kR=1", "kR=3", "kR=10", "kR=30"], ps.T):
+            data_dict[key] = np.stack((a_bar, patch)).T
+        with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+            config_data = json.load(config_file)
+        np.savez(Path(config_data["figure_data"]).joinpath("fig_6.npz"), **data_dict)
+
+    # fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
+    fig, ax = plt.subplots(figsize=(3, 2.5))
     for patch, kR, marker in zip(ps.T, kappaR, markers):
     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.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', linewidth=2.5,
+                # marker=marker, markerfacecolor='none', markersize=8, markevery=10
+                )
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
     ax.set_xlabel(r'$\bar a$', fontsize=15)
     ax.set_xlabel(r'$\bar a$', fontsize=15)
     ax.set_ylabel(r'$\theta_0$', fontsize=15)
     ax.set_ylabel(r'$\theta_0$', fontsize=15)
-    plt.legend(fontsize=14)
+    plt.legend(handlelength=0.6, fontsize=13, frameon=False, loc='upper right', bbox_to_anchor=(0.56, 0.58))
     plt.tight_layout()
     plt.tight_layout()
+    plt.axhline(y=mark_ps, color='black', linestyle=':')
+    plt.axvline(x=0.5, color='black', linestyle=':')
     if save:
     if save:
-        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential.pdf"), dpi=300)
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/final_figures/patch_size_potential.png"), dpi=300)
     plt.show()
     plt.show()
 
 
 
 
@@ -77,47 +96,53 @@ def plot_kappaR_charge_dependence(*, normalized=False, save=False):
 def plot_sigma0_dependence(*, save=False):
 def plot_sigma0_dependence(*, save=False):
 
 
     # kappaR = np.array([0.1, 1, 3, 10, 30])
     # kappaR = np.array([0.1, 1, 3, 10, 30])
-    kappaR = np.linspace(0.1, 15, 101)
+    kappaR = np.linspace(0.1, 15, 201)
     # sigma0 = np.linspace(-0.001, 0.0015, 6)
     # sigma0 = np.linspace(-0.001, 0.0015, 6)
-    sigma0 = np.array([-0.001, 0, 0.001])
+    sigma0 = np.array([-0.0002, 0, 0.0002])
     params = ModelParams(R=150, kappaR=kappaR)
     params = ModelParams(R=150, kappaR=kappaR)
-    ex = expansion.MappedExpansionQuad(a_bar=0.8, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)
+    ex = expansion.MappedExpansionQuad(a_bar=0.3, 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)
+    t0 = time.perf_counter()
+    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0, skip_nozero_cases=True)
+    t1 = time.perf_counter()
+    print(f'Time: {t1 - t0}')
 
 
     markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']
     markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']
 
 
     fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
     fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
     for patch, sgm, marker in zip(ps.T, sigma0, markers):
     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)
+        nonzero_ps = np.nonzero(patch)
+        ax.plot(kappaR[nonzero_ps], patch[nonzero_ps], label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=10)
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
     ax.set_xlabel(r'$\kappa R$', fontsize=15)
     ax.set_xlabel(r'$\kappa R$', fontsize=15)
     ax.set_ylabel(r'$\theta_0$', fontsize=15)
     ax.set_ylabel(r'$\theta_0$', fontsize=15)
-    plt.legend(fontsize=14)
+    plt.legend(fontsize=14, loc='upper right')
     plt.tight_layout()
     plt.tight_layout()
     if save:
     if save:
-        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_charge_abar08.pdf"), dpi=300)
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_charge_abar03.pdf"), dpi=300)
     plt.show()
     plt.show()
 
 
 
 
 def plot_sigma0_dependence_relative(*, save=False):
 def plot_sigma0_dependence_relative(*, save=False):
 
 
     # kappaR = np.array([0.1, 1, 3, 10, 30])
     # kappaR = np.array([0.1, 1, 3, 10, 30])
-    kappaR = np.linspace(0.1, 15, 101)
+    kappaR = np.linspace(0.1, 15, 201)
     # sigma0 = np.linspace(-0.001, 0.0015, 6)
     # sigma0 = np.linspace(-0.001, 0.0015, 6)
-    sigma0 = np.array([-0.001, 0, 0.001])
+    sigma0 = np.array([-0.0002, 0, 0.0002])
     params = ModelParams(R=150, kappaR=kappaR)
     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 = expansion.MappedExpansionQuad(a_bar=0.8, 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)
     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 = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0, skip_nozero_cases=True)
     ps_neutral = patch_size.potential_patch_size(ex_neutral, 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', '+']
     markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']
 
 
     fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
     fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
     for patch, sgm, marker in zip(ps.T, sigma0, markers):
     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)
+        nonzero_ps = np.nonzero(patch)
+        ax.plot(kappaR[nonzero_ps], patch[nonzero_ps] / ps_neutral[nonzero_ps],
+                label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=10)
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
     ax.set_xlabel(r'$\kappa R$', fontsize=15)
     ax.set_xlabel(r'$\kappa R$', fontsize=15)
     ax.set_ylabel(r'$\theta_0$', fontsize=15)
     ax.set_ylabel(r'$\theta_0$', fontsize=15)
@@ -128,11 +153,18 @@ def plot_sigma0_dependence_relative(*, save=False):
     plt.show()
     plt.show()
 
 
 
 
-if __name__ == '__main__':
+def main():
+
+    plot_abar_dependence(save=True)
 
 
-    # plot_abar_dependence(save=False)
     # plot_sigma0_dependence(save=True)
     # plot_sigma0_dependence(save=True)
+
     # plot_sigma0_dependence_relative(save=True)
     # plot_sigma0_dependence_relative(save=True)
 
 
     # plot_abar_charge_dependence(save=True)
     # plot_abar_charge_dependence(save=True)
-    plot_kappaR_charge_dependence(normalized=True, save=True)
+    # plot_kappaR_charge_dependence(normalized=True, save=True)
+
+
+if __name__ == '__main__':
+
+    main()

+ 0 - 114
analysis/path_plot.py

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

+ 460 - 36
analysis/peak_heigth.py

@@ -1,62 +1,486 @@
-import matplotlib.pyplot as plt
-from charged_shells import expansion, interactions, mapping
+from charged_shells import expansion, interactions, mapping, functions
 from charged_shells.parameters import ModelParams
 from charged_shells.parameters import ModelParams
 import numpy as np
 import numpy as np
 from typing import Literal
 from typing import Literal
 from pathlib import Path
 from pathlib import Path
 from functools import partial
 from functools import partial
+import json
+from plot_settings import *
+from dataclasses import dataclass
 
 
 Array = np.ndarray
 Array = np.ndarray
 Expansion = expansion.Expansion
 Expansion = expansion.Expansion
-RotConfig = Literal['ep', 'pp']
+RotConfig = Literal['ep', 'pp', 'sep']
 
 
 
 
-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):
+@dataclass
+class Peak:
+    rot_config: str
+    y_label: str
+    ex1: Expansion
+    ex2: Expansion
+    emanuele_data_column: int
+    log_y: bool
+    kappaR_axis_in_expansion: int = None
 
 
+
+class PeakEP(Peak):
+
+    def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None):
+        self.emanuele_data_column = 4
+        self.y_label = r'$|U_{EP}|$' if log_y else r'$U_{EP}$'
+        self.ex1 = ex.clone()
+        self.ex2 = ex.clone()
+        self.ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
+        self.log_y = log_y
+        self.kappaR_axis_in_expansion = kappaR_axis_in_expansion
+
+
+class PeakPP(Peak):
+
+    def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None):
+        self.emanuele_data_column = 3
+        self.y_label = r'$U_{PP}$'
+        self.ex1 = ex.clone()
+        self.ex2 = ex.clone()
+        self.log_y = log_y
+        self.kappaR_axis_in_expansion = kappaR_axis_in_expansion
+
+
+class PeakSEP(Peak):
+
+    def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None):
+        self.emanuele_data_column = 5
+        self.y_label = r'$|U_{sEP}|$' if log_y else r'$U_{sEP}$'
+        self.ex1 = ex.clone()
+        self.ex2 = ex.clone()
+        self.ex1.rotate_euler(alpha=0, beta=np.pi / 4, gamma=0)
+        self.ex2.rotate_euler(alpha=0, beta=np.pi / 4, gamma=0)
+        self.log_y = log_y
+        self.kappaR_axis_in_expansion = kappaR_axis_in_expansion
+
+
+def get_charge_energy_dicts(peak: Peak, params: ModelParams, emanuele_data: Array, sigma0: Array, abar_cs: list, sigma_tilde=0.001):
+    abar_ic, inverse, counts = np.unique(emanuele_data[:, 1], return_counts=True, return_inverse=True)
+
+    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2),
+                                                     peak.kappaR_axis_in_expansion)
+    energy = energy_fn(peak.ex1, peak.ex2, params)
+
+    data_dict_ic = {}
+    data_dict_cs = {}
+    k = 0
+    for i in range(len(abar_ic)):
+        ab = np.around(abar_ic[i], 5)
+        if ab in np.around(abar_cs, 5):
+            idx, = np.nonzero(inverse == i)
+            charge = emanuele_data[idx, 0]
+            en = emanuele_data[idx, peak.emanuele_data_column]
+            sort = np.argsort(charge)
+            if peak.log_y:
+                data_dict_ic[ab] = np.stack((charge[sort] / 280 + 2, np.abs(en)[sort])).T
+                data_dict_cs[ab] = np.stack((sigma0 / sigma_tilde, np.abs(energy[k]))).T
+            else:
+                data_dict_ic[ab] = np.stack((charge[sort] / 280 + 2, en[sort])).T
+                data_dict_cs[ab] = np.stack((sigma0 / sigma_tilde, energy[k])).T
+            k += 1
+
+    return data_dict_ic, data_dict_cs
+
+
+def get_kappaR_energy_dicts(peak: Peak, params: ModelParams, emanuele_data: Array, kappaR: Array, abar_cs: list, max_kappaR: float = 50):
+    abar_ic, inverse = np.unique(emanuele_data[:, 1], return_inverse=True)
+
+    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2),
+                                                     peak.kappaR_axis_in_expansion)
+    energy = energy_fn(peak.ex1, peak.ex2, params)
+
+    data_dict_ic = {}
+    data_dict_cs = {}
+    k = 0
+    for i in range(len(abar_ic)):
+        ab = np.around(abar_ic[i], 5)
+        if ab in np.around(abar_cs, 5):
+            idx, = np.nonzero(inverse == i)
+            kR = emanuele_data[idx, 2][emanuele_data[idx, 2] <= max_kappaR]
+            en = emanuele_data[idx, peak.emanuele_data_column][emanuele_data[idx, 1] <= max_kappaR]
+            sort = np.argsort(kR)
+            if peak.log_y:
+                data_dict_ic[ab] = np.stack((kR[sort], np.abs(en)[sort])).T
+                data_dict_cs[ab] = np.stack((kappaR, np.abs(energy[k]))).T
+            else:
+                data_dict_ic[ab] = np.stack((kR[sort], en[sort])).T
+                data_dict_cs[ab] = np.stack((kappaR, energy[k])).T
+            k += 1
+
+    return data_dict_ic, data_dict_cs
+
+
+def get_abar_energy_dicts(peak: Peak, params: ModelParams, emanuele_data: Array, a_bar: Array, kR_cs: list, max_abar: float = 0.8):
+
+    kR_ic, inverse = np.unique(emanuele_data[:, 2], return_inverse=True)
+
+    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2),
+                                                     peak.kappaR_axis_in_expansion)
+    energy = energy_fn(peak.ex1, peak.ex2, params)
+
+    data_dict_ic = {}
+    data_dict_cs = {}
+    k = 0
+    for i in range(len(kR_ic)):
+        kr = np.around(kR_ic[i], 5)
+        if kr in np.around(kR_cs, 5):
+            idx, = np.nonzero(inverse == i)
+            abar = emanuele_data[idx, 1][emanuele_data[idx, 1] <= max_abar]
+            en = emanuele_data[idx, peak.emanuele_data_column][emanuele_data[idx, 1] <= max_abar]
+            sort = np.argsort(abar)
+            if peak.log_y:
+                data_dict_ic[kr] = np.stack((abar[sort], np.abs(en)[sort])).T
+                data_dict_cs[kr] = np.stack((a_bar, np.abs(energy[k]))).T
+            else:
+                data_dict_ic[kr] = np.stack((abar[sort], en[sort])).T
+                data_dict_cs[kr] = np.stack((a_bar, energy[k])).T
+            k += 1
+
+    return data_dict_ic, data_dict_cs
+
+
+def save_data(data_dict_ic, data_dict_cs, a_bar: list):
+    if save_data:
+        ic_save = {str(key): val for key, val in data_dict_ic.items()}
+        cs_save = {str(key): val for key, val in data_dict_cs.items()}
+        ic_save['abar'] = a_bar
+        cs_save['abar'] = a_bar
+        with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+            config_data = json.load(config_file)
+        np.savez(Path(config_data["figure_data"]).joinpath("fig_11_IC.npz"), **ic_save)
+        np.savez(Path(config_data["figure_data"]).joinpath("fig_11_CS.npz"), **cs_save)
+
+
+def IC_peak_energy_plot(config_data: dict,
+                        a_bar: list,
+                        which: RotConfig,
+                        min_kappaR: float = 0.01,
+                        max_kappaR: float = 50.,
+                        R: float = 150,
+                        save_as: Path = None,
+                        save_data: bool = False,
+                        quad_correction: bool = False,
+                        log_y: bool = True):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_PANELS_BDF"))
+    em_data = np.load(em_data_path.joinpath("newpair_energy.npz"))
+    data = em_data['fixA']
+
+    num = 100 if which == 'sep' else 30
+    kappaR = np.geomspace(min_kappaR, max_kappaR, num)
     params = ModelParams(R=R, kappaR=kappaR)
     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()
+    ex = expansion.MappedExpansionQuad(np.array(a_bar)[:, None], params.kappaR[None, :], 0.001, l_max=20)
+
     if which == 'ep':
     if which == 'ep':
-        ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
+        peak = PeakEP(ex, log_y, kappaR_axis_in_expansion=1)
+    elif which == 'pp':
+        peak = PeakPP(ex, log_y, kappaR_axis_in_expansion=1)
+    elif which == 'sep':
+        peak = PeakSEP(ex, log_y, kappaR_axis_in_expansion=1)
+    else:
+        raise ValueError
+
+    data_dict_ic, data_dict_cs = get_kappaR_energy_dicts(peak, params, data, kappaR, a_bar, max_kappaR)
 
 
-    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist), 1)
-    energy = energy_fn(ex1, ex2, params)
+    colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
 
 
-    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)
+    fig, ax = plt.subplots(figsize=(4.25, 3))
+    for (key, data1), data2 in zip(data_dict_ic.items(), data_dict_cs.values()):
+        current_color = next(colors)
+        ax.plot(data1[:, 0], data1[:, 1], label=rf'$\bar a = {key:.2f}$', c=current_color)
+        ax.plot(data2[:, 0], data2[:, 1], ls='--', c=current_color)
+    ax.legend(fontsize=14, ncol=1, frameon=False, handlelength=0.7, loc='upper right',
+              bbox_to_anchor=(0.42, 0.42),
+              # bbox_to_anchor=(0.42, 0.9)
+              )
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
     ax.set_xlabel(r'$\kappa R$', fontsize=15)
     ax.set_xlabel(r'$\kappa R$', fontsize=15)
-    ax.set_ylabel(rf'$\bar V$', fontsize=15)
+
+    ax.set_ylabel(peak.y_label, fontsize=15)
+    if log_y:
+        ax.set_yscale('log')
+    ax.set_xscale('log')
     plt.tight_layout()
     plt.tight_layout()
     if save_as is not None:
     if save_as is not None:
         plt.savefig(save_as, dpi=600)
         plt.savefig(save_as, dpi=600)
     plt.show()
     plt.show()
 
 
 
 
-if __name__ == '__main__':
+def IC_peak_energy_abar_plot(config_data: dict,
+                        kappaR: list,
+                        which: RotConfig,
+                        min_abar: float = 0.01,
+                        max_abar: float = 50.,
+                        R: float = 150,
+                        save_as: Path = None,
+                        save_data: bool = False,
+                        quad_correction: bool = False,
+                        log_y: bool = True):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_PANELS_BDF"))
+    em_data = np.load(em_data_path.joinpath("newpair_energy.npz"))
+    data = em_data['fixA']
+
+    num = 100 if which == 'sep' else 30
+    a_bar = np.linspace(min_abar, max_abar, num)
+    params = ModelParams(R=R, kappaR=np.array(kappaR))
+
+    ex = expansion.MappedExpansionQuad(np.array(a_bar)[None, :], params.kappaR[:, None], 0.001, l_max=20)
+
+    if which == 'ep':
+        peak = PeakEP(ex, log_y, kappaR_axis_in_expansion=0)
+    elif which == 'pp':
+        peak = PeakPP(ex, log_y, kappaR_axis_in_expansion=0)
+    elif which == 'sep':
+        peak = PeakSEP(ex, log_y, kappaR_axis_in_expansion=0)
+    else:
+        raise ValueError
+
+    data_dict_ic, data_dict_cs = get_abar_energy_dicts(peak, params, data, a_bar, kappaR, max_abar)
+
+    colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
+
+    fig, ax = plt.subplots(figsize=(4.25, 3))
+    for (key, data1), data2 in zip(data_dict_ic.items(), data_dict_cs.values()):
+        current_color = next(colors)
+        ax.plot(data1[:, 0], data1[:, 1], label=rf'$\kappa R = {key:.2f}$', c=current_color)
+        ax.plot(data2[:, 0], data2[:, 1], ls='--', c=current_color)
+    ax.legend(fontsize=14, ncol=1, frameon=False, handlelength=0.7,
+              # loc='lower right',
+              # bbox_to_anchor=(0.42, 0.42),
+              # bbox_to_anchor=(0.42, 0.9)
+              )
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel(r'$\bar a$', fontsize=15)
+
+    ax.set_ylabel(peak.y_label, fontsize=15)
+    if log_y:
+        ax.set_yscale('log')
+    # ax.set_xscale('log')
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
 
 
-    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')
-                     )
+def IC_peak_energy_charge_plot(config_data: dict,
+                        a_bar: list,
+                        which: RotConfig,
+                        max_kappaR: float = 30.,
+                        R: float = 150,
+                        save_as: Path = None,
+                        save_data: bool = False,
+                        quad_correction: bool = False,
+                        log_y: bool = False):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_PANELS_BDF"))
+    em_data = np.load(em_data_path.joinpath("newpair_energy.npz"))
+    data = em_data['changeZc']
+
+    params = ModelParams(R=R, kappaR=3)
+    sigma0 = np.linspace(-0.0003, 0.0003, 300)
+    sigma_tilde = 0.001
+
+    ex = expansion.MappedExpansionQuad(np.array(a_bar), kappaR=3, sigma0=sigma0, l_max=20, sigma_tilde=sigma_tilde)
+
+    if which == 'ep':
+        peak = PeakEP(ex, log_y)
+    elif which == 'pp':
+        peak = PeakPP(ex, log_y)
+    elif which == 'sep':
+        peak = PeakSEP(ex, log_y)
+    else:
+        raise ValueError
+
+    data_dict_ic, data_dict_cs = get_charge_energy_dicts(peak, params, data, sigma0, a_bar, sigma_tilde)
+
+    colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
+
+    fig, ax = plt.subplots(figsize=(4.25, 3))
+    for (key, data1), data2 in zip(data_dict_ic.items(), data_dict_cs.values()):
+        current_color = next(colors)
+        ax.plot(data1[:, 0], data1[:, 1], label=rf'$\bar a = {key:.2f}$', c=current_color)
+        ax.plot(data2[:, 0], data2[:, 1], ls='--', c=current_color)
+    ax.legend(fontsize=14, ncol=1, frameon=False, handlelength=0.7, loc='upper right',
+              # bbox_to_anchor=(0.42, 0.95),
+              # bbox_to_anchor=(0.7, 1),
+              bbox_to_anchor=(0.42, 1),
+              )
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel(r'$\eta$', fontsize=15)
+
+    ax.set_ylabel(peak.y_label, fontsize=15)
+    if log_y:
+        ax.set_yscale('log')
+    # ax.set_xscale('log')
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def IC_peak_energy_kappaR_combined_plot(config_data: dict,
+                        a_bar: list,
+                        R: float = 150,
+                        save_as: Path = None,
+                        min_kappaR: float = 0.01,
+                        max_kappaR: float = 50.,
+                        ):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_PANELS_BDF"))
+    em_data = np.load(em_data_path.joinpath("newpair_energy.npz"))
+    data = em_data['fixA']
+
+    num = 50
+    kappaR = np.geomspace(min_kappaR, max_kappaR, num)
+    params = ModelParams(R=R, kappaR=kappaR)
+
+    ex = expansion.MappedExpansionQuad(np.sort(np.array(a_bar))[:, None],  # sorting necessary as it is also in energy_dicts()
+                                       params.kappaR[None, :], 0.001, l_max=20)
+
+    peak_ep = PeakEP(ex, log_y=True, kappaR_axis_in_expansion=1)
+    peak_pp = PeakPP(ex, log_y=True, kappaR_axis_in_expansion=1)
+    peak_sep = PeakSEP(ex, log_y=False, kappaR_axis_in_expansion=1)
+    peaks = [peak_ep, peak_pp, peak_sep]
+
+    data_ic = []
+    data_cs = []
+    for peak in peaks:
+        dict_ic, dict_cs = get_kappaR_energy_dicts(peak, params, data, kappaR, a_bar, max_kappaR)
+        data_ic.append(dict_ic)
+        data_cs.append(dict_cs)
+
+    legend_coords = [(0.58, 0.45), (0.58, 0.43), (0.58, 0.9)]
+
+    fig, axs = plt.subplots(3, 1, figsize=(3, 7.8))
+    for ax, data_dict_ic, data_dict_cs, peak, lc in zip(axs, data_ic, data_cs, peaks, legend_coords):
+        colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
+        for ab in a_bar:
+            key = np.around(ab, 5)
+            current_color = next(colors)
+            ax.plot(data_dict_ic[key][:, 0], data_dict_ic[key][:, 1], label=rf'$\bar a = {key:.2f}$', c=current_color)
+            ax.plot(data_dict_cs[key][:, 0], data_dict_cs[key][:, 1], ls='--', c=current_color)
+        ax.axvline(x=10, c='k', ls=':')
+        ax.axvline(x=3, c='k', ls=':')
+        ax.axvline(x=1, c='k', ls=':')
+        ax.legend(fontsize=13, ncol=1, frameon=False, handlelength=0.7, loc='upper right', bbox_to_anchor=lc)
+        ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+        # if ax == axs[-1]:
+        ax.set_xlabel(r'$\kappa R$', fontsize=15)
+
+        ax.set_ylabel(peak.y_label, fontsize=15)
+        if peak.log_y:
+            ax.set_yscale('log')
+        ax.set_xscale('log')
+        ax.yaxis.set_label_coords(-0.2, 0.5)
+    plt.subplots_adjust(left=0.3)
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def IC_peak_energy_charge_combined_plot(config_data: dict,
+                        a_bar: list,
+                        R: float = 150,
+                        save_as: Path = None,
+                        log_y: bool = False):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_PANELS_BDF"))
+    em_data = np.load(em_data_path.joinpath("newpair_energy.npz"))
+    data = em_data['changeZc']
+
+    params = ModelParams(R=R, kappaR=3)
+    sigma0 = np.linspace(-0.0003, 0.0003, 300)
+    sigma_tilde = 0.001
+
+    ex = expansion.MappedExpansionQuad(np.sort(np.array(a_bar)),  # sorting necessary as it is also in energy_dicts()
+                                       kappaR=3, sigma0=sigma0, l_max=20, sigma_tilde=sigma_tilde)
+
+    peak_ep = PeakEP(ex, log_y)
+    peak_pp = PeakPP(ex, log_y)
+    peak_sep = PeakSEP(ex, log_y)
+    peaks = [peak_ep, peak_pp, peak_sep]
+
+    data_ic = []
+    data_cs = []
+    for peak in peaks:
+        dict_ic, dict_cs = get_charge_energy_dicts(peak, params, data, sigma0, a_bar, sigma_tilde)
+        data_ic.append(dict_ic)
+        data_cs.append(dict_cs)
+
+    fig, axs = plt.subplots(3, 1, figsize=(3, 7.8))
+    for ax, data_dict_ic, data_dict_cs, peak in zip(axs, data_ic, data_cs, peaks):
+        colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
+        for ab in a_bar:
+            key = np.around(ab, 5)
+            current_color = next(colors)
+            ax.plot(data_dict_ic[key][:, 0], data_dict_ic[key][:, 1], label=rf'$\bar a = {key:.2f}$', c=current_color)
+            ax.plot(data_dict_cs[key][:, 0], data_dict_cs[key][:, 1], ls='--', c=current_color)
+        ax.legend(fontsize=13, ncol=1, frameon=False, handlelength=0.7, loc='upper right',
+                  bbox_to_anchor=(0.77, 1.03),
+                  )
+        ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+        # if ax == axs[-1]:
+        ax.set_xlabel(r'$\eta$', fontsize=15)
+
+        ax.set_ylabel(peak.y_label, fontsize=15)
+        if peak.log_y:
+            ax.set_yscale('log')
+        # ax.set_xscale('log')
+        ax.yaxis.set_label_coords(-0.2, 0.5)
+    plt.subplots_adjust(left=0.3)
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def main():
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    # a_bar = [0.2, 0.5, 0.8]
+    # IC_peak_energy_plot(config_data, a_bar=a_bar, which='ep', max_kappaR=50,
+    #                     # save_as=Path('/home/andraz/ChargedShells/Figures/Emanuele_data/peak_pp_kappaR.png'),
+    #                     save_data=False,
+    #                     quad_correction=False,
+    #                     log_y=True
+    #                     )
+    # IC_peak_energy_kappaR_combined_plot(config_data, a_bar,
+    #                                     save_as=Path(
+    #                                         '/home/andraz/ChargedShells/Figures/final_figures/peak_combined_kappaR.png')
+    #                                     )
+
+    # a_bar = [0.1, 0.2, 0.3]
+    # IC_peak_energy_charge_plot(config_data, a_bar=a_bar, which='sep',
+    #                            # save_as=Path('/home/andraz/ChargedShells/Figures/Emanuele_data/peak_sep_charge.png'),
+    #                            )
+    # IC_peak_energy_charge_combined_plot(config_data, a_bar,
+    #                                     save_as=Path('/home/andraz/ChargedShells/Figures/final_figures/peak_combined_charge.png')
+    #                                     )
+
+    kappaR = [0.01, 3.02407, 30]
+    IC_peak_energy_abar_plot(config_data, kappaR=kappaR, which='sep', min_abar=0.2, max_abar=0.8,
+                        save_as=Path('/home/andraz/ChargedShells/Figures/Emanuele_data/peak_sep_abar.png'),
+                        save_data=False,
+                        quad_correction=False,
+                        log_y=False
+                        )
+
+
+if __name__ == '__main__':
+    
+    main()

+ 16 - 0
analysis/plot_settings.py

@@ -0,0 +1,16 @@
+from matplotlib import cm
+import matplotlib.pyplot as plt
+from itertools import cycle
+import matplotlib.legend as mlegend
+from matplotlib.legend_handler import HandlerBase
+
+# all the common imports and settings for plotting
+
+plt.rcParams['text.usetex'] = True
+plt.rcParams["font.family"] = "Times New Roman"
+
+COLOR_LIST = [cm.cividis(0.15), cm.autumn(0.5), cm.summer(0),
+              cm.gist_heat(0.5), cm.cool(0.65), cm.gray(0.5), cm.winter(0.2)]
+
+plt.rcParams['axes.prop_cycle'] = plt.cycler(color=COLOR_LIST)
+COLORS = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])

+ 540 - 0
analysis/quad_path_plot.py

@@ -0,0 +1,540 @@
+import numpy as np
+from matplotlib.lines import Line2D
+
+from charged_shells.rotational_path import PairRotationalPath, PathEnergyPlot
+from charged_shells import expansion
+from charged_shells.parameters import ModelParams
+from pathlib import Path
+import json
+import quadrupole_model_mappings
+from plot_settings import *
+
+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, start_name="EP", end_name="EE")
+QuadPath.add_euler(beta1=zero_to_pi_half[::-1], beta2=zero_to_pi_half[::-1], end_name="PP")
+QuadPath.add_euler(beta1=zero_to_pi_half, end_name="EP")
+QuadPath.add_euler(beta1=zero_to_pi_half[::-1], beta2=zero_to_pi_half, end_name="EP")
+QuadPath.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half, alpha2=np.pi/2, end_name="tEE")
+QuadPath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half[::-1], end_name="EE")
+
+
+def model_comparison(config_data: dict, save_as=None, save_data=False):
+    kappaR = 3
+    params = ModelParams(R=150, kappaR=kappaR)
+    a_bar = 0.5
+    sigma_tilde = 0.001
+
+    ex1 = expansion.MappedExpansionQuad(a_bar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    # matching other models to the mapped CSp model based on equal patch size in potential
+    # ex_gauss = quadrupole_model_mappings.ic_to_gauss(sigma_tilde, a_bar, params, l_max=30, sigma0=0)
+    # ex_gauss2 = ex_gauss.clone()
+    # ex_cap = quadrupole_model_mappings.ic_to_cap(sigma_tilde, a_bar, params, l_max=30, sigma0=0)
+    # ex_cap2 = ex_cap.clone()
+
+    # path plots for all models
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    # path_plot_gauss = PathEnergyPlot(ex_gauss, ex_gauss2, QuadPath, dist=2., params=params)
+    # energy_gauss = path_plot_gauss.evaluate_path()
+    #
+    # path_plot_cap = PathEnergyPlot(ex_cap, ex_cap2, QuadPath, dist=2., params=params)
+    # energy_cap = path_plot_cap.evaluate_path()
+
+    # peak_energy_sanity_check
+    # ex1new = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    # ex2new = ex1new.clone()
+    # pp_energy = interactions.charged_shell_energy(ex1new, ex2new, params)
+    # print(f'PP energy: {pp_energy}')
+
+    # Emanuele data
+    em_data = np.load(Path(config_data["emanuele_data"]).joinpath("FIG_3C").joinpath("pathway.npz"))['arr_0']
+    # em_data = np.load(Path(config_data["emanuele_data"]).joinpath("FIG_7").joinpath("pathway.npz"))['arr_0']
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE")
+    #                 .joinpath("FIX_A").joinpath("ECC_0.25"))
+    # em_data = np.load(em_data_path.joinpath(f"EMME_{kappaR}.").joinpath("pathway.npz"))['arr_0']
+
+    if save_data:
+        np.savez(Path(config_data["figure_data"]).joinpath(f"fig_7_kR{kappaR}.npz"),
+                 ICi=em_data,
+                 CSp=np.stack((x_axis, np.squeeze(energy))).T,
+                 # CSp_gauss=np.stack((x_axis, np.squeeze(energy_gauss))).T,
+                 # CSp_cap=np.stack((x_axis, np.squeeze(energy_cap))).T
+                 )
+
+    fig, ax = plt.subplots(figsize=(8.25, 3))
+    ax.plot(em_data[:, 0], em_data[:, 1], label='ICi', c=COLOR_LIST[1])
+    ax.plot(x_axis, np.squeeze(energy), label='CSp', c=COLOR_LIST[1], ls='--')
+    # ax.plot(x_axis, np.squeeze(energy_gauss), label='CSp - Gauss')
+    # ax.plot(x_axis, np.squeeze(energy_cap), label='CSp - cap')
+    # ax.plot(x_axis, em_data[:, 1] / np.squeeze(energy), label='CSp')
+    path_plot.plot_style(fig, ax, size=(8.25, 3.5))
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def kappaR_dependence(kappaR: Array, abar: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=0)
+
+    path_plot.plot(labels=[rf'$\kappa R$={kR}' for kR in kappaR],
+                   # norm_euler_angles={'beta2': np.pi},
+                   save_as=save_as)
+
+
+def abar_dependence(abar: Array, kappaR: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)
+
+    path_plot.plot(labels=[rf'$\bar a$={a}' for a in abar],
+                   # norm_euler_angles={'beta2': np.pi},
+                   save_as=save_as)
+
+
+def sigma0_dependence(sigma0: Array, kappaR: float, abar: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=sigma0)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)
+
+    path_plot.plot(labels=[rf'$\eta={s0 / sigma_tilde}$' for s0 in sigma0],
+                   # norm_euler_angles={'beta2': np.pi},
+                   save_as=save_as)
+
+
+def distance_dependence(dist: Array, kappaR: float, abar: float, sigma_tilde=0.001, save_as=None):
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    plots = []
+    for d in dist:
+        path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=d, params=params)
+        x = d * kappaR
+        plots.append(path_plot.evaluate_path() * np.exp(x) * x)
+
+    x_axis = path_plot.rot_path.stack_x_axes()
+    labels = [rf'$\rho/R ={d}$' for d in dist]
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for pl, lbl in zip(plots, labels):
+        ax.plot(x_axis, pl, label=lbl)
+    QuadPath.plot_style(fig, ax)
+    ax.set_ylabel(r'$U \kappa\rho e^{\kappa\rho}$', fontsize=15)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def IC_kappaR_dependence(config_data: dict, save_as=None):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE")
+    #                 .joinpath("FIX_A").joinpath("ECC_0.25"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("FIXEDCHARGE")
+                    .joinpath("FIX_A").joinpath("ECC_0.25"))
+    kR1 = np.load(em_data_path.joinpath("EMME_1.").joinpath("pathway.npz"))['arr_0']
+    kR3 = np.load(em_data_path.joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
+    kR10 = np.load(em_data_path.joinpath("EMME_10.").joinpath("pathway.npz"))['arr_0']
+
+    labels = [rf'$\kappa R$={kR}' for kR in [1, 3, 10]]
+
+    fig, ax = plt.subplots()
+    ax.plot(kR1[:, 0], kR1[:, 1], label=labels[0])
+    ax.plot(kR3[:, 0], kR3[:, 1], label=labels[1])
+    ax.plot(kR10[:, 0], kR10[:, 1], label=labels[2])
+    QuadPath.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def IC_abar_dependence(config_data: dict, save_as=None):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE").joinpath("FIX_M"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("FIXEDCHARGE").joinpath("FIX_M"))
+    a03 = np.load(em_data_path.joinpath("ECC_0.15").joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
+    a04 = np.load(em_data_path.joinpath("ECC_0.2").joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
+    a05 = np.load(em_data_path.joinpath("ECC_0.25").joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
+
+    labels =[rf'$\bar a$={a}' for a in [0.3, 0.4, 0.5]]
+
+    fig, ax = plt.subplots()
+    ax.plot(a03[:, 0], a03[:, 1], label=labels[0])
+    ax.plot(a04[:, 0], a04[:, 1], label=labels[1])
+    ax.plot(a05[:, 0], a05[:, 1], label=labels[2])
+    QuadPath.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def IC_sigma0_dependence(config_data: dict, save_as=None):
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("CHARGE_ZC"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("CHARGE_ZC"))
+    undercharged = np.load(em_data_path.joinpath("ZC_-277.27").joinpath("pathway.npz"))['arr_0']
+    neutral = np.load(em_data_path.joinpath("ZC_-560").joinpath("pathway.npz"))['arr_0']
+    overchargerd = np.load(em_data_path.joinpath("ZC_-842.74").joinpath("pathway.npz"))['arr_0']
+
+    labels = [rf'$\eta={eta}$' for eta in [-0.1, 0, 0.1]]
+
+    fig, ax = plt.subplots()
+    ax.plot(overchargerd[:, 0], overchargerd[:, 1], label=labels[0])
+    ax.plot(neutral[:, 0], neutral[:, 1], label=labels[1])
+    ax.plot(undercharged[:, 0], undercharged[:, 1], label=labels[2])
+    QuadPath.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_distance_dependence(config_data: dict, dist: Array = 2 * np.array([1., 1.15, 1.3, 1.45]),
+                                 kappaR: float = 3,
+                                 abar: float = 0.5,
+                                 sigma_tilde=0.001,
+                                 save_as=None):
+
+    # em_data_path = Path(config_data["emanuele_data"]).joinpath("FIG_12")
+    em_data_path = Path(config_data["emanuele_data"]).joinpath("FIG_3D_LONG_DIST")
+    em_data = np.load(em_data_path.joinpath("pathway_fig12A.npz"))
+
+    em_data_d2 = np.load(Path(config_data["emanuele_data"]).joinpath("FIG_3C").joinpath("pathway.npz"))['arr_0']
+
+    ic_data = [em_data_d2]
+    for key, d in em_data.items():
+        ic_data.append(d)
+
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    plots = []
+    for d in dist:
+        path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=d, params=params)
+        plots.append(path_plot.evaluate_path())
+
+    x_axis = path_plot.rot_path.stack_x_axes()
+    labels = [rf'$\rho/R ={d}$' for d in dist]
+
+    # additional legend
+    line1 = Line2D([0], [0], color='black', linewidth=1, label='ICi')
+    line2 = Line2D([0], [0], color='black', linestyle='--', linewidth=1, label='CSp')
+
+    fig, ax = plt.subplots()
+    for i, (d, en, label, c) in enumerate(zip(ic_data, plots, labels, COLORS)):
+        if i < 3:
+            ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+            ax.plot(x_axis, en, ls='--', c=c)
+    QuadPath.plot_style(fig, ax)
+    main_legend = ax.get_legend()
+    extra_legend = ax.legend(handles=[line1, line2], loc='upper left', fontsize=15, frameon=False)
+    ax.add_artist(main_legend)
+    ax.add_artist(extra_legend)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_rescaled_distance_dependence(config_data: dict,
+                                          dist: Array = 2 * np.array([1, 1.5, 2, 3, 5, 10]),
+                                          kappaR: float = 3,
+                                          abar: float = 0.5,
+                                          sigma_tilde=0.001,
+                                          save_as=None):
+
+    # em_data_path = Path(config_data["emanuele_data"]).joinpath("FIG_12")
+    em_data_path = Path(config_data["emanuele_data"]).joinpath("FIG_3D_LONG_DIST")
+    em_data = np.load(em_data_path.joinpath("pathway_fig12B.npz"))
+    ic_data = []
+    for key, d in em_data.items():
+        # if key == 'kr10':
+        #     continue
+        print(key, np.max(d[:, 1]))
+        ic_data.append(d)
+
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    plots = []
+    for d in dist:
+        path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=d, params=params)
+        x = d * kappaR
+        plots.append(path_plot.evaluate_path() * np.exp(x) * x)
+        # plots.append(path_plot.evaluate_path())
+
+    x_axis = path_plot.rot_path.stack_x_axes()
+    labels = [rf'$\rho/R ={d}$' for d in dist]
+
+    fig, ax = plt.subplots()
+    for d, en, label, c in zip(ic_data, plots, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+    QuadPath.plot_style(fig, ax)
+    ax.set_ylabel(r'$U \kappa\rho e^{\kappa\rho}$', fontsize=15)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_kappaR_dependence(config_data: dict, kappaR: list[int], abar: float, sigma_tilde=0.001, save_as=None):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE")
+    #                 .joinpath("FIX_A").joinpath(f"ECC_{np.round(abar/2, 4)}"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("FIXEDCHARGE")
+                    .joinpath("FIX_A").joinpath(f"ECC_{np.round(abar/2, 4)}"))
+
+    ic_data = []
+    for kR in kappaR:
+        ic_data.append(np.load(em_data_path.joinpath(f"EMME_{kR}.").joinpath("pathway.npz"))['arr_0'])
+
+    params = ModelParams(R=150, kappaR=np.asarray(kappaR))
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=0)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    labels = [rf'$\kappa R={kR}$' for kR in [1, 3, 10]]
+
+    fig, ax = plt.subplots()
+    for d, en, label, c in zip(ic_data, energy.T, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+    QuadPath.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_abar_dependence(config_data: dict, kappaR: int, abar: list[float], sigma_tilde=0.001, save_as=None):
+
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE").joinpath("FIX_M"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("FIXEDCHARGE").joinpath("FIX_M"))
+
+    ic_data = []
+    for ab in abar:
+        ic_data.append(np.load(em_data_path.joinpath(f"ECC_{np.round(ab/2, 4)}").
+                               joinpath(f"EMME_{kappaR}.").joinpath("pathway.npz"))['arr_0'])
+
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(np.asarray(abar), params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    labels = [rf'$\bar a={a}$' for a in [0.3, 0.4, 0.5]]
+
+    fig, ax = plt.subplots()
+    for d, en, label, c in zip(ic_data, energy.T, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+    QuadPath.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_sigma0_dependence(config_data: dict, kappaR=3., abar=0.5, sigma0=(0.0002, 0.00, -0.0002), sigma_tilde=0.001, save_as=None):
+    # em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("CHARGE_ZC"))
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("CHARGE_ZC"))
+    undercharged = np.load(em_data_path.joinpath("ZC_-503").joinpath("pathway.npz"))['arr_0']
+    neutral = np.load(em_data_path.joinpath("ZC_-560").joinpath("pathway.npz"))['arr_0']
+    overchargerd = np.load(em_data_path.joinpath("ZC_-617").joinpath("pathway.npz"))['arr_0']
+    ic_data = [undercharged, neutral, overchargerd]
+
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=np.asarray(sigma0))
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy = path_plot.evaluate_path()
+    x_axis = path_plot.rot_path.stack_x_axes()
+
+    labels = [rf'$\eta={s0/sigma_tilde}$' for s0 in sigma0]
+
+    fig, ax = plt.subplots()
+    for d, en, label, c in zip(ic_data, energy.T, labels, COLORS):
+        ax.plot(d[:, 0], d[:, 1], label=label, c=c)
+        ax.plot(x_axis, en, ls='--', c=c)
+    QuadPath.plot_style(fig, ax)
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def combined_all(config_data: dict, save_as=None):
+
+    sigma_tilde = 0.001
+    kappaR_list = [1, 3, 10]
+    abar_list = [0.5, 0.4, 0.3]
+    sigma0_list = [0.0002, 0.00, -0.0002]
+    kappaR = 3
+    abar = 0.5
+
+
+    em_data_kappaR = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("FIXEDCHARGE")
+                    .joinpath("FIX_A").joinpath(f"ECC_{np.round(abar/2, 4)}"))
+
+    ic_data_kappaR = []
+    for kR in kappaR_list:
+        ic_data_kappaR.append(np.load(em_data_kappaR.joinpath(f"EMME_{kR}.").joinpath("pathway.npz"))['arr_0'])
+
+    params = ModelParams(R=150, kappaR=np.asarray(kappaR_list))
+
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=0)
+    energy_kappaR = path_plot.evaluate_path()
+    x_axis_kappaR = path_plot.rot_path.stack_x_axes()
+    labels_kappaR = [rf'$\kappa R={kR}$' for kR in [1, 3, 10]]
+
+
+    em_data_abar = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("FIXEDCHARGE").joinpath("FIX_M"))
+
+    ic_data_abar = []
+    for ab in abar_list:
+        ic_data_abar.append(np.load(em_data_abar.joinpath(f"ECC_{np.round(ab/2, 4)}").
+                               joinpath(f"EMME_{kappaR}.").joinpath("pathway.npz"))['arr_0'])
+
+    params = ModelParams(R=150, kappaR=kappaR)
+
+    ex1 = expansion.MappedExpansionQuad(np.asarray(abar_list), params.kappaR, sigma_tilde, l_max=30)
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_abar = path_plot.evaluate_path()
+    x_axis_abar = path_plot.rot_path.stack_x_axes()
+    labels_abar = [rf'$\bar a={a}$' for a in abar_list]
+
+
+    em_data_charge = (Path(config_data["emanuele_data"]).joinpath("FIG_4_Panels_ACE").joinpath("CHARGE_ZC"))
+    undercharged = np.load(em_data_charge.joinpath("ZC_-503").joinpath("pathway.npz"))['arr_0']
+    neutral = np.load(em_data_charge.joinpath("ZC_-560").joinpath("pathway.npz"))['arr_0']
+    overchargerd = np.load(em_data_charge.joinpath("ZC_-617").joinpath("pathway.npz"))['arr_0']
+    ic_data_sigma0 = [undercharged, neutral, overchargerd]
+
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=np.asarray(sigma0_list))
+    ex2 = ex1.clone()
+
+    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)
+    energy_sigma0 = path_plot.evaluate_path()
+    x_axis_sigma0 = path_plot.rot_path.stack_x_axes()
+
+    labels_sigma0 = [rf'$\eta={s0/sigma_tilde}$' for s0 in sigma0_list]
+
+    fig, axs = plt.subplots(3, 1, figsize=(6, 7.8))
+    for d, en, label, c in zip(ic_data_kappaR, energy_kappaR.T, labels_kappaR, COLOR_LIST):
+        axs[0].set_title('Screening', fontsize=15)
+        axs[0].plot(d[:, 0], d[:, 1], label=label, c=c)
+        axs[0].plot(x_axis_kappaR, en, ls='--', c=c)
+    QuadPath.plot_style(fig, axs[0], size=None)
+    for d, en, label, c in zip(ic_data_abar, energy_abar.T, labels_abar, COLOR_LIST):
+        axs[1].set_title('Asymmetry', fontsize=15)
+        axs[1].plot(d[:, 0], d[:, 1], label=label, c=c)
+        axs[1].plot(x_axis_abar, en, ls='--', c=c)
+    QuadPath.plot_style(fig, axs[1], size=None)
+    for d, en, label, c in zip(ic_data_sigma0, energy_sigma0.T, labels_sigma0, COLOR_LIST):
+        axs[2].set_title('Net charge', fontsize=15)
+        axs[2].plot(d[:, 0], d[:, 1], label=label, c=c)
+        axs[2].plot(x_axis_sigma0, en, ls='--', c=c)
+    for ax in axs:
+        ax.yaxis.set_label_coords(-0.09, 0.5)
+    # axs[-1].set_xlabel('rotational path', fontsize=15)
+    QuadPath.plot_style(fig, axs[2], size=None)
+    for ax in axs:
+        ax.get_legend().set_bbox_to_anchor((0.65, 1))
+    if save_as is not None:
+        plt.savefig(save_as, dpi=300)
+    plt.show()
+
+
+def main():
+
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    # model_comparison(config_data, save_data=False,
+    #     save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('quad_og_comparison.png')
+    # )
+
+    # kappaR_dependence(np.array([1, 3, 10]), 0.5,
+    #                   # save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_kappaR_dep.png")
+    #                   )
+    #
+    # abar_dependence(np.array([0.3, 0.4, 0.5]), 3,
+    #                 save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_abar_dep.png")
+    #                 )
+
+    # sigma0_dependence(np.array([-0.0002, 0.00, 0.0002]), 3, 0.5,
+    #                   save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_charge_dep_abar05_kappaR3.png")
+    #                   )
+
+    # distance_dependence(dist=np.array([2, 3, 4, 6, 10, 20]), kappaR=3, abar=0.5,
+    #                     # save_as=Path(config_data["figures"]).joinpath('quadrupole_distance_dep.png')
+    #                     )
+
+    # IC_kappaR_dependence(config_data,
+    #                      save_as=Path(config_data["figures"]).joinpath("Emanuele_data").joinpath('IC_quadrupole_kappaR_dep.png')
+    #                      )
+    #
+    # IC_abar_dependence(config_data, save_as=Path(config_data["figures"]).joinpath("Emanuele_data").
+    #                    joinpath('IC_quadrupole_abar_dep.png'))
+    #
+    # IC_sigma0_dependence(config_data, save_as=Path(config_data["figures"]).joinpath("Emanuele_data").
+    #                      joinpath('IC_quadrupole_charge_dep_abar05_kappaR3.png'))
+
+    # combined_kappaR_dependence(config_data, kappaR=[1, 3, 10], abar=0.5,
+    #                      save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('quad_kappaR_dep.png')
+    #                      )
+
+    # combined_sigma0_dependence(config_data,
+    #                      save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('quad_charge_dep.png')
+    #                      )
+
+    # combined_abar_dependence(config_data, kappaR=3, abar=[0.3, 0.4, 0.5],
+    #                          save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('quad_abar_dep.png')
+    #                          )
+
+    # combined_rescaled_distance_dependence(config_data)
+
+    combined_distance_dependence(config_data,
+                                 save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('quad_dist_dep.png')
+                                 )
+
+    # combined_all(config_data,
+    #              save_as=Path(config_data["figures"]).joinpath("final_figures").joinpath('quad_combined_dep.png')
+    #              )
+
+
+if __name__ == '__main__':
+
+    main()
+

+ 53 - 0
analysis/quadrupole_model_mappings.py

@@ -0,0 +1,53 @@
+from charged_shells import patch_size, expansion, parameters
+from charged_shells import functions as fn
+from scipy.special import eval_legendre
+import numpy as np
+
+
+ModelParams = parameters.ModelParams
+Array = np.ndarray
+
+
+def point_to_gauss_magnitude(sigma_tilde: Array, a_bar: Array, lbd: Array, kappaR: Array):
+    return (sigma_tilde * fn.coefficient_Cim(2, kappaR) / fn.coefficient_Cpm(2, kappaR)
+            * np.sinh(lbd) / (lbd * fn.sph_bessel_i(2, lbd)) * a_bar ** 2)
+
+
+def point_to_cap_magnitude(sigma_tilde: Array, a_bar: Array, theta0: Array, kappaR: Array):
+    return (sigma_tilde * 10 * fn.coefficient_Cim(2, kappaR) / fn.coefficient_Cpm(2, kappaR) *
+            a_bar ** 2 / (eval_legendre(1, np.cos(theta0)) - eval_legendre(3, np.cos(theta0))))
+
+
+def ic_to_gauss(sigma_tilde, a_bar, params: ModelParams, l_max: int = 30,
+                sigma0: float = 0) -> expansion.GaussianCharges:
+
+    ex_mapped = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR,
+                                              sigma_tilde=sigma_tilde, l_max=30, sigma0=sigma0)
+    target_patch_size = patch_size.potential_patch_size(ex_mapped, params)
+    sigma0_mapped = expansion.net_charge_map(sigma0, params.kappaR)
+
+    def fn_gauss(x):
+        return expansion.GaussianCharges(lambda_k=x, omega_k=np.array([[0, 0], [np.pi, 0]]),
+                                         sigma1=sigma_tilde, l_max=l_max, sigma0=sigma0_mapped)
+
+    lbd = patch_size.inverse_potential_patch_size(target_patch_size, fn_gauss, 5, params)
+    gauss_sigma = point_to_gauss_magnitude(sigma_tilde, a_bar, lbd, params.kappaR)
+    return expansion.GaussianCharges(lambda_k=lbd, omega_k=np.array([[0, 0], [np.pi, 0]]),
+                                     sigma1=gauss_sigma, l_max=l_max, sigma0=sigma0_mapped)
+
+
+def ic_to_cap(sigma_tilde, a_bar, params: ModelParams, l_max: int = 30, sigma0: float = 0) -> expansion.SphericalCap:
+
+    ex_mapped = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR,
+                                              sigma_tilde=sigma_tilde, l_max=30, sigma0=sigma0)
+    target_patch_size = patch_size.potential_patch_size(ex_mapped, params)
+    sigma0_mapped = expansion.net_charge_map(sigma0, params.kappaR)
+
+    def fn_cap(x):
+        return expansion.SphericalCap(theta0_k=x, sigma1=sigma_tilde, l_max=l_max,
+                                      omega_k=np.array([[0, 0], [np.pi, 0]]), sigma0=sigma0_mapped)
+
+    theta0 = patch_size.inverse_potential_patch_size(target_patch_size, fn_cap, 0.5, params)
+    cap_sigma = point_to_cap_magnitude(sigma_tilde, a_bar, theta0, params.kappaR)
+    return expansion.SphericalCap(theta0_k=theta0, sigma1=cap_sigma,
+                                  omega_k=np.array([[0, 0], [np.pi, 0]]), l_max=l_max, sigma0=0)

+ 199 - 10
analysis/sEE_minimum.py

@@ -1,17 +1,30 @@
+import matplotlib.pyplot as plt
 import numpy as np
 import numpy as np
 from charged_shells import expansion, interactions, mapping
 from charged_shells import expansion, interactions, mapping
 from charged_shells.parameters import ModelParams
 from charged_shells.parameters import ModelParams
 from functools import partial
 from functools import partial
+from pathlib import Path
+import json
+from matplotlib import cm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from labellines import labelLine, labelLines
+from matplotlib.colors import TwoSlopeNorm
+from matplotlib.ticker import FuncFormatter
 
 
 Expansion = expansion.Expansion
 Expansion = expansion.Expansion
 
 
 
 
-def sEE_minimum(ex: Expansion, params: ModelParams, accuracy=1e-3, dist=2., match_expansion_axis_to_params=None):
+def sEE_minimum(ex: Expansion, params: ModelParams, accuracy=1e-2, dist=2., match_expansion_axis_to_params=None,
+                angle_start: float = 0, angle_stop: float = np.pi / 2):
     ex2 = ex.clone()
     ex2 = ex.clone()
-    zero_to_pi_half = np.linspace(0, np.pi / 2, int(1 / accuracy), endpoint=True)
+    angle_range = np.linspace(angle_start, angle_stop, int((angle_stop - angle_start) / accuracy), endpoint=True)
+    # angle_range = np.linspace(0.1, 0.7, 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)
+    ex.rotate_euler(alpha=0, beta=angle_range, gamma=0)
+    ex2.rotate_euler(alpha=0, beta=angle_range, gamma=0)
+
+    if match_expansion_axis_to_params is not None:
+        match_expansion_axis_to_params += 1
 
 
     energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist),
     energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist),
                                                      match_expansion_axis_to_params)
                                                      match_expansion_axis_to_params)
@@ -20,17 +33,193 @@ def sEE_minimum(ex: Expansion, params: ModelParams, accuracy=1e-3, dist=2., matc
     min_idx = np.argmin(energy, axis=0)
     min_idx = np.argmin(energy, axis=0)
     min_energy = np.min(energy, axis=0)
     min_energy = np.min(energy, axis=0)
 
 
-    return min_energy, zero_to_pi_half[min_idx]
+    return min_energy, angle_range[min_idx]
 
 
 
 
-def main():
+def contours():
+
+    kappaR = np.linspace(1, 10, 20)
+    a_bar = np.linspace(0.2, 0.6, 20)
+    params = ModelParams(R=150, kappaR=kappaR)
+    ex = expansion.MappedExpansionQuad(a_bar[:, None], kappaR[None, :], 0.001)
+    min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=1, accuracy=0.01)
+
+    kR_mesh, a_mesh = np.meshgrid(kappaR, a_bar)
+
+    print(np.min(min_angle), np.max(min_angle))
+    print(np.min(min_energy), np.max(min_energy))
+
+    plt.contourf(kR_mesh, a_mesh, min_angle)
+    # plt.imshow(min_angle)
+    plt.show()
+
+    plt.contourf(kR_mesh, a_mesh, min_energy, np.array([-9, -5, -2.5, -1, -0.5, -0.2, 0]))
+    # plt.imshow(min_energy)
+    plt.show()
+
+
+def kappaR_dependence(kappaR, save_as=None, cmap=cm.jet):
 
 
-    kappaR = np.array([1, 3, 10])
+    a_bar = np.linspace(0.12, 0.8, 15)
     params = ModelParams(R=150, kappaR=kappaR)
     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)
 
 
+    ex = expansion.MappedExpansionQuad(a_bar[None, :], kappaR[:, None], 0.001)
+    min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=0, accuracy=0.001,
+                                        angle_start=0.5, angle_stop=1.)
+
+    kappaR_alt = np.array([0.01, 2, 5, 10, 50])
+    params_alt = ModelParams(R=150, kappaR=kappaR_alt)
+    a_bar_alt = np.linspace(np.min(a_bar) - 0.05, np.max(a_bar) + 0.05, 20)
+    ex2 = expansion.MappedExpansionQuad(a_bar_alt[:, None], kappaR_alt[None, :], 0.001)
+    min_energy_alt, min_angle_alt = sEE_minimum(ex2, params_alt, match_expansion_axis_to_params=1, accuracy=0.001,
+                                                angle_start=0.5, angle_stop=1.)
+
+    colors = cmap(np.linspace(0, 1, len(a_bar)))
+    sm = cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=np.min(a_bar), vmax=np.max(a_bar)))
+    sm.set_array([])
+
+    # Transformation to the position along our quadrupolar rotational path
+    min_angle = np.pi - min_angle
+    min_angle_alt = np.pi - min_angle_alt
+
+    kappaR_labels = [rf'$\kappa R={kR:.1f}$' for kR in kappaR_alt]
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for me, ma, lbl in zip(min_energy_alt.T, min_angle_alt.T, kappaR_labels):
+        ax.plot(ma, me, ls=':', c='k', label=lbl)
+    labelLines(ax.get_lines(), align=False, fontsize=15,  xvals=(2.35, 2.65))
+
+    for me, ma, lbl, c in zip(min_energy.T, min_angle.T, [rf'$\bar a={a:.2f}$' for a in a_bar], colors):
+        ax.plot(ma, me, label=lbl, c=c)
+    # ax.plot(min_angle, min_energy)
+    # ax.legend(fontsize=17)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel('pry angle', fontsize=20)
+    ax.set_ylabel('U', fontsize=20)
+    ax.set_xlim(2.2, 2.65)
+    ax.set_ylim(-31, 1)
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    cax.tick_params(labelsize=15)
+    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
+    cbar.set_label(r'$\bar a$', rotation=0, labelpad=15, fontsize=20)
+
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+def IC_kappaR_dependence(config_data: dict, which_kappa_lines: list,
+                         save_as=None, cmap=cm.jet, file_suffix=""):
+    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_9"))
+    em_data = np.load(em_data_path.joinpath(f"fig9{file_suffix}.npz"))['arr_0']
+    # print(em_data.shape)
+
+    a_bar, indices, counts = np.unique(em_data[:, 0], return_counts=True, return_inverse=True)
+    print(a_bar)
+
+    if not np.all(counts == counts[0]):
+        raise ValueError("Data not reshapable.")
+
+    colors = cmap(np.linspace(0, 1, len(a_bar)))
+    sm = cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=np.min(a_bar), vmax=np.max(a_bar)))
+    sm.set_array([])
+
+    min_energy = em_data[:, 3].reshape(-1, counts[0])
+    min_angle = em_data[:, 2].reshape(-1, counts[0])
+
+    kappa_line_energy = []
+    kappa_line_angle = []
+    kappaR_labels = []
+    for kappa in which_kappa_lines:
+        ind = em_data[:, 1] == kappa
+        if np.sum(ind) == 0:
+            print(f'No lines with kR={kappa} in data. Available values: {np.unique(em_data[:, 1])}')
+            continue
+        kappa_line_energy.append(em_data[:, 3][ind])
+        kappa_line_angle.append(em_data[:, 2][ind])
+        kappaR_labels.append(rf'$\kappa R={kappa:.1f}$')
+
+
+    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+    for me, ma, lbl in zip(kappa_line_energy, kappa_line_angle, kappaR_labels):
+        sort = np.argsort(ma)
+        ax.plot(ma[sort], me[sort], ls=':', c='k', label=lbl)
+
+    labelLines(ax.get_lines(), align=False, fontsize=15, xvals=(2.4, 2.65))
+
+    for me, ma, lbl, c in zip(min_energy, min_angle, [rf'$\bar a={a:.2f}$' for a in a_bar], colors):
+        ax.plot(ma, me, label=lbl, c=c)
+    # ax.plot(min_angle, min_energy)
+    # ax.legend(fontsize=17)
+    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+    ax.set_xlabel('pry angle', fontsize=20)
+    ax.set_ylabel('U', fontsize=20)
+    ax.set_xlim(2.2, 2.65)
+    ax.set_ylim(-41, 1)
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    cax.tick_params(labelsize=15)
+    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
+    cbar.set_label(r'$\bar a$', rotation=0, labelpad=15, fontsize=20)
+
+    plt.tight_layout()
+    if save_as is not None:
+        plt.savefig(save_as, dpi=600)
+    plt.show()
+
+
+# def charge_dependence(charge, save_as=None):
+#
+#     a_bar = np.array([0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])
+#     params = ModelParams(R=150, kappaR=kappaR)
+#
+#     ex = expansion.MappedExpansionQuad(a_bar[None, :], kappaR[:, None], 0.001)
+#     min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=0, accuracy=0.001,
+#                                         angle_start=0.5, angle_stop=1.)
+#
+#     kappaR_alt = np.array([0.01, 2, 5, 10, 50])
+#     params_alt = ModelParams(R=150, kappaR=kappaR_alt)
+#     a_bar_alt = np.linspace(np.min(a_bar) - 0.05, np.max(a_bar) + 0.05, 20)
+#     ex2 = expansion.MappedExpansionQuad(a_bar_alt[:, None], kappaR_alt[None, :], 0.001)
+#     min_energy_alt, min_angle_alt = sEE_minimum(ex2, params_alt, match_expansion_axis_to_params=1, accuracy=0.001,
+#                                                 angle_start=0.5, angle_stop=1.)
+#
+#     fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+#     for me, ma, lbl in zip(min_energy_alt.T, min_angle_alt.T, [rf'$\bar a={a:.2f}$' for a in a_bar]):
+#         ax.plot(ma, me, ls=':', c='k')
+#     for me, ma, lbl in zip(min_energy.T, min_angle.T, [rf'$\bar a={a:.2f}$' for a in a_bar]):
+#         ax.plot(ma, me, label=lbl)
+#     # ax.plot(min_angle, min_energy)
+#     ax.legend(fontsize=17)
+#     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
+#     ax.set_xlabel('angle', fontsize=15)
+#     ax.set_ylabel('U', fontsize=20)
+#     ax.set_xlim(0.56, 0.93)
+#     ax.set_ylim(-20, 1)
+#     plt.tight_layout()
+#     if save_as is not None:
+#         plt.savefig(save_as, dpi=600)
+#     plt.show()
+
+
+def main():
+
+    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
+        config_data = json.load(config_file)
+
+    # contours()
+
+    # kappaR_dependence(kappaR=np.linspace(0.1, 30, 20),
+    #                   save_as=Path(config_data["figures"]).joinpath('sEE_min_kappaR_abar.png')
+    #                   )
+
+    IC_kappaR_dependence(config_data, which_kappa_lines=[0.01, 3., 5., 10., 50.], file_suffix="_1",
+                         save_as=Path(config_data["figures"]).joinpath("Emanuele_data").joinpath('IC_sEE_min_kappaR_abar.png')
+                         )
 
 
 
 
 if __name__ == '__main__':
 if __name__ == '__main__':

+ 77 - 27
charged_shells/expansion.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 from __future__ import annotations
 import numpy as np
 import numpy as np
 from dataclasses import dataclass
 from dataclasses import dataclass
-from functools import cached_property
 import charged_shells.functions as fn
 import charged_shells.functions as fn
 import quaternionic
 import quaternionic
 import spherical
 import spherical
@@ -55,7 +54,7 @@ class Expansion:
         self.coefs = self.coefs.reshape(shape + (self.coefs.shape[-1],))
         self.coefs = self.coefs.reshape(shape + (self.coefs.shape[-1],))
         self._rotations = self._rotations.reshape(shape + (4,))
         self._rotations = self._rotations.reshape(shape + (4,))
 
 
-    @cached_property
+    @property
     def lm_arrays(self) -> (Array, Array):
     def lm_arrays(self) -> (Array, Array):
         """Return l and m arrays containing all (l, m) pairs."""
         """Return l and m arrays containing all (l, m) pairs."""
         return full_lm_arrays(self.l_array)
         return full_lm_arrays(self.l_array)
@@ -66,17 +65,26 @@ class Expansion:
         return np.repeat(arr, 2 * self.l_array + 1, axis=axis)
         return np.repeat(arr, 2 * self.l_array + 1, axis=axis)
 
 
     def rotate(self, rotations: Quaternion, rotate_existing=False):
     def rotate(self, rotations: Quaternion, rotate_existing=False):
-        # TODO: rotations are currently saved wrong if we start form existing coefficients not the og ones
+        if rotate_existing:
+            raise NotImplementedError("Rotation of possibly already rotated coefficients is not yet supported.")
         self._rotations = rotations
         self._rotations = rotations
-        coefs = self.coefs if rotate_existing else self._starting_coefs
-        self.coefs = expansion_rotation(rotations, coefs, self.l_array)
+        self.coefs = expansion_rotation(rotations, self._starting_coefs, self.l_array)
 
 
-    def rotate_euler(self, alpha: Array, beta: Array, gamma: Array, rotate_existing=False):
+    def rotate_euler(self, alpha: Array = 0, beta: Array = 0, gamma: Array = 0, rotate_existing=False):
         # TODO: additional care required on the convention used to transform euler angles to quaternions
         # 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 !!
         # TODO: might be off for a minus sign for each? angle !!
         R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
         R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
         self.rotate(R_euler, rotate_existing=rotate_existing)
         self.rotate(R_euler, rotate_existing=rotate_existing)
 
 
+    def inverse_sign(self, exclude_00: bool = False):
+        if self.l_array[0] == 0 and exclude_00:
+            self.coefs[..., 1:] = -self.coefs[..., 1:]
+            self._starting_coefs[..., 1:] = -self._starting_coefs[..., 1:]
+            return self
+        self.coefs = -self.coefs
+        self._starting_coefs = -self._starting_coefs
+        return self
+
     def charge_value(self, theta: Array | float, phi: Array | float):
     def charge_value(self, theta: Array | float, phi: Array | float):
         if not isinstance(theta, Array):
         if not isinstance(theta, Array):
             theta = np.array([theta])
             theta = np.array([theta])
@@ -94,20 +102,27 @@ class Expansion:
 
 
 class Expansion24(Expansion):
 class Expansion24(Expansion):
 
 
-    def __init__(self, sigma2: float, sigma4: float, sigma0: float = 0.):
+    def __init__(self, sigma2: float | Array, sigma4: float | Array, sigma0: float | Array = 0.):
         l_array = np.array([0, 2, 4])
         l_array = np.array([0, 2, 4])
-        coefs = rot_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
+        try:
+            broadcasted_arrs = np.broadcast_arrays(np.sqrt(4*np.pi) * sigma0, sigma2, sigma4)
+        except ValueError:
+            raise ValueError("Given sigma0, sigma2 and sigma4 arrays cannot be broadcast to a common shape.")
+        coefs = rot_sym_expansion(l_array, np.stack(broadcasted_arrs, axis=-1))
         super().__init__(l_array, coefs)
         super().__init__(l_array, coefs)
 
 
 
 
-class MappedExpansionQuad(Expansion):
-    """Expansion that matches the outside potential of a quadrupolar impermeable particle with point charges inside."""
+class MappedExpansionSym(Expansion):
+    """
+    Expansion that matches the outside potential of an impermeable particle with
+    a rotationally symmetric point charge distribution inside (specifically, either quadrupolar or dipolar).
+    """
 
 
     def __init__(self,
     def __init__(self,
                  a_bar: Array | float,
                  a_bar: Array | float,
                  kappaR: Array | float,
                  kappaR: Array | float,
                  sigma_tilde: float,
                  sigma_tilde: float,
-                 l_max: int = 20,
+                 l_array: Array,
                  sigma0: float | Array = 0):
                  sigma0: float | Array = 0):
         """
         """
         :param a_bar: distance between the center and off center charges
         :param a_bar: distance between the center and off center charges
@@ -116,18 +131,45 @@ class MappedExpansionQuad(Expansion):
         :param l_max: maximal ell value for the expansion
         :param l_max: maximal ell value for the expansion
         :param sigma0: total (mean) charge density
         :param sigma0: total (mean) charge density
         """
         """
-        a_bar, kappaR = np.broadcast_arrays(a_bar, kappaR)
+        if not isinstance(sigma0, Array):
+            sigma0 = np.array([sigma0])
+        if sigma0.ndim > 1:
+            raise ValueError(f'Sigma0 parameter cannot be an array of dimensions larger than 1, got dim={sigma0.ndim}')
+        a_bar_bc, kappaR_bc = np.broadcast_arrays(a_bar, kappaR)
+
+        a_bar_bc2, kappaR_bc2, l_array_expanded = np.broadcast_arrays(a_bar_bc[..., None],
+                                                                      kappaR_bc[..., None],
+                                                                      l_array[None, :])
+        coefs = (2 * sigma_tilde * fn.coef_C_diff(l_array_expanded, kappaR_bc2)
+                 * np.sqrt(4 * np.pi * (2 * l_array_expanded + 1)) * np.power(a_bar_bc2, l_array_expanded))
+        coefs = np.squeeze(rot_sym_expansion(l_array, coefs))
+        kappaR_bc3, sigma0_bc3 = np.broadcast_arrays(kappaR_bc[..., None], sigma0[None, :])
+        coefs = expansion_total_charge(coefs, net_charge_map(sigma0_bc3, kappaR_bc3), l_array)
+        super().__init__(l_array, np.squeeze(coefs))
+
 
 
+class MappedExpansionQuad(MappedExpansionSym):
+
+    def __init__(self,
+                 a_bar: Array | float,
+                 kappaR: Array | float,
+                 sigma_tilde: float,
+                 l_max: int = 20,
+                 sigma0: float | Array = 0):
         l_array = np.array([l for l in range(l_max + 1) if l % 2 == 0])
         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, :])
+        super().__init__(a_bar, kappaR, sigma_tilde, l_array, sigma0)
 
 
-        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 MappedExpansionDipole(MappedExpansionSym):
+
+    def __init__(self,
+                 a_bar: Array | float,
+                 kappaR: Array | float,
+                 sigma_tilde: float,
+                 l_max: int = 20,
+                 sigma0: float | Array = 0):
+        l_array = np.array([l for l in range(l_max + 1) if l % 2 == 1 or l == 0])
+        super().__init__(a_bar, kappaR, sigma_tilde, l_array, sigma0)
 
 
 
 
 class GaussianCharges(Expansion):
 class GaussianCharges(Expansion):
@@ -163,7 +205,7 @@ class GaussianCharges(Expansion):
                     * np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
                     * np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
                                           theta_k[None, None, :], phi_k[None, None, :])))
                                           theta_k[None, None, :], phi_k[None, None, :])))
         coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
         coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
-        coefs = expansion_total_charge(coefs, sigma0)
+        coefs = expansion_total_charge(coefs, sigma0, l_array)
         l_array, coefs = purge_unneeded_l(l_array, coefs)
         l_array, coefs = purge_unneeded_l(l_array, coefs)
         super().__init__(l_array, coefs)
         super().__init__(l_array, coefs)
 
 
@@ -207,10 +249,14 @@ class SphericalCap(Expansion):
         # Rotating is implemented in such a way that it rotates every patch to every position,
         # Rotating is implemented in such a way that it rotates every patch to every position,
         # hence the redundancy of out of diagonal elements.
         # hence the redundancy of out of diagonal elements.
         coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
         coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
-        coefs_all = expansion_total_charge(coefs_all, sigma0)
+        coefs_all = expansion_total_charge(coefs_all, sigma0, l_array)
         super().__init__(l_array, np.squeeze(coefs_all))
         super().__init__(l_array, np.squeeze(coefs_all))
 
 
 
 
+def net_charge_map(sigma0: float | Array, kappaR: float | Array):
+    return sigma0 * np.exp(kappaR) / np.sinh(kappaR) * kappaR / (1 + kappaR)
+
+
 def full_lm_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)."""
     """From an array of l_values get arrays of ell and m that give you all pairs (ell, m)."""
     all_m_list = []
     all_m_list = []
@@ -227,17 +273,21 @@ def rot_sym_expansion(l_array: Array, coefs: Array) -> Array:
     return full_coefs
     return full_coefs
 
 
 
 
-def expansion_total_charge(coefs: Array, sigma0: float | Array):
+def expansion_total_charge(coefs: Array, sigma0: float | Array | None, l_vals: Array):
     """Adds a new axis to the expansion coefficients that modifies expansion based on given net charge density."""
     """Adds a new axis to the expansion coefficients that modifies expansion based on given net charge density."""
+    if l_vals[0] != 0:
+        raise NotImplementedError("To modify total charge, the first expansion term should be l=0. "
+                                  "Adding l=0 term not yet implemented (TODO).")
     if sigma0 is None:
     if sigma0 is None:
         return coefs
         return coefs
     if not isinstance(sigma0, Array):
     if not isinstance(sigma0, Array):
         x = copy.deepcopy(coefs)
         x = copy.deepcopy(coefs)
-        x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
+        x[..., 0] = sigma0 * np.sqrt(4 * np.pi)
         return x
         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)
+
+    # insert new axis in the 2nd-to-last place and repeat the expansion data over this new axis
+    x = np.repeat(np.expand_dims(coefs, -2), sigma0.shape[-1], axis=-2)
+    x[..., 0] = sigma0 * np.sqrt(4 * np.pi)
     return x
     return x
 
 
 
 

+ 1 - 1
charged_shells/functions.py

@@ -15,7 +15,7 @@ def sph_harm(l, m, theta, phi, **kwargs):
 
 
 
 
 def interaction_coef_C(l, p, x):
 def interaction_coef_C(l, p, x):
-    return x * sps.iv(l + 1 / 2, x) * sps.iv(p + 1 / 2, x)
+    return np.pi / 2 * x * sps.iv(l + 1 / 2, x) * sps.iv(p + 1 / 2, x)
 
 
 
 
 def coefficient_Cpm(l, x):
 def coefficient_Cpm(l, x):

+ 12 - 3
charged_shells/interactions.py

@@ -2,7 +2,6 @@ from charged_shells import expansion, parameters
 import charged_shells.functions as fn
 import charged_shells.functions as fn
 from py3nj import wigner3j
 from py3nj import wigner3j
 import numpy as np
 import numpy as np
-import time
 from typing import Literal
 from typing import Literal
 import charged_shells.units_and_constants as uc
 import charged_shells.units_and_constants as uc
 
 
@@ -16,7 +15,8 @@ EnergyUnit = Literal['kT', 'eV', 'J']
 def energy_units(units: EnergyUnit, params: ModelParams) -> float:
 def energy_units(units: EnergyUnit, params: ModelParams) -> float:
     match units:
     match units:
         case 'eV':
         case 'eV':
-            return 1 / (uc.CONSTANTS.e0 * uc.UNITS.voltage)
+            # return 1 / (uc.CONSTANTS.e0 * uc.UNITS.voltage)
+            return 1.
         case 'kT':
         case 'kT':
             return 1 / (params.temperature * uc.CONSTANTS.Boltzmann)
             return 1 / (params.temperature * uc.CONSTANTS.Boltzmann)
         case 'J':
         case 'J':
@@ -36,6 +36,13 @@ def charged_shell_energy(ex1: Expansion, ex2: Expansion, params: ModelParams, di
     flat_p = full_l_array[indices_p]
     flat_p = full_l_array[indices_p]
     flat_m = full_m_array[indices_l]  # the same as full_m_array[indices_p]
     flat_m = full_m_array[indices_l]  # the same as full_m_array[indices_p]
 
 
+    relevant_pairs, = np.nonzero(flat_l >= flat_p)
+    flat_l = flat_l[relevant_pairs]
+    flat_p = flat_p[relevant_pairs]
+    flat_m = flat_m[relevant_pairs]
+    indices_l = indices_l[relevant_pairs]
+    indices_p = indices_p[relevant_pairs]
+
     charge_factor = np.real(ex1.coefs[..., indices_l] * np.conj(ex2.coefs[..., 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]))
                             (-1) ** (flat_l + flat_p) * ex1.coefs[..., indices_p] * np.conj(ex2.coefs[..., indices_l]))
 
 
@@ -75,4 +82,6 @@ def charged_shell_energy(ex1: Expansion, ex2: Expansion, params: ModelParams, di
     lspm_vals = C_vals * (-1) ** (l_vals + m_vals) * lps_terms * bessel_vals * wigner1 * wigner2
     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)
     broadcasted_lspm_vals = np.broadcast_to(lspm_vals, charge_vals.shape)
 
 
-    return 0.5 * constants * np.sum(broadcasted_lspm_vals * charge_vals, axis=-1)
+    rescale_at_equal_lp = np.where(l_vals == p_vals, 0.5, 1)
+
+    return constants * np.sum(rescale_at_equal_lp * broadcasted_lspm_vals * charge_vals, axis=-1)

+ 22 - 31
charged_shells/mapping.py

@@ -36,6 +36,17 @@ def unravel_params(params: ModelParams) -> list[ModelParams]:
         return [ModelParams(R=params.R, kappa=kappa) for kappa in params.kappa]
         return [ModelParams(R=params.R, kappa=kappa) for kappa in params.kappa]
     if not (isinstance(params.R, Array) or isinstance(params.kappa, Array)):
     if not (isinstance(params.R, Array) or isinstance(params.kappa, Array)):
         return [params]
         return [params]
+    raise NotImplementedError
+
+
+def unravel_expansion_over_axis(ex: Expansion, axis: int | None, param_list_len: int) -> list[Expansion]:
+    if axis is None:
+        return [ex for _ in range(param_list_len)]
+    axis_len = ex.shape[axis]
+    if axis_len != param_list_len:
+        raise ValueError(f'Parameter list has different length than the provided expansion axis, '
+                         f'got param_list_len={param_list_len} and axis_len={axis_len}.')
+    return [Expansion(ex.l_array, np.take(ex.coefs, i, axis=axis)) for i in range(axis_len)]
 
 
 
 
 SingleExpansionFn = Callable[[Expansion, ModelParams], T]
 SingleExpansionFn = Callable[[Expansion, ModelParams], T]
@@ -45,58 +56,38 @@ TwoExpansionsFn = Callable[[Expansion, Expansion, ModelParams], T]
 def parameter_map_single_expansion(f: SingleExpansionFn,
 def parameter_map_single_expansion(f: SingleExpansionFn,
                                    match_expansion_axis_to_params: int = None) -> 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):
     def mapped_f(ex: Expansion, params: ModelParams):
         params_list = unravel_params(params)
         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)}')
+        expansion_list = unravel_expansion_over_axis(ex, match_expansion_axis_to_params, len(params_list))
 
 
         results = []
         results = []
         for exp, prms in zip(expansion_list, params_list):
         for exp, prms in zip(expansion_list, params_list):
             results.append(f(exp, prms))
             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
+        if match_expansion_axis_to_params is not None:
+            # return the params-matched axis to where it belongs
+            return np.moveaxis(np.array(results), 0, match_expansion_axis_to_params)
         return np.squeeze(np.array(results))
         return np.squeeze(np.array(results))
 
 
     return mapped_f
     return mapped_f
 
 
 
 
 def parameter_map_two_expansions(f: TwoExpansionsFn,
 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
+                                 match_expansion_axis_to_params: int = None,
+                                 ) -> TwoExpansionsFn:
 
 
     def mapped_f(ex1: Expansion, ex2: Expansion, params: ModelParams):
     def mapped_f(ex1: Expansion, ex2: Expansion, params: ModelParams):
         params_list = unravel_params(params)
         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)}')
+        expansion_list1 = unravel_expansion_over_axis(ex1, match_expansion_axis_to_params, len(params_list))
+        expansion_list2 = unravel_expansion_over_axis(ex2, match_expansion_axis_to_params, len(params_list))
 
 
         results = []
         results = []
         for exp1, exp2, prms in zip(expansion_list1, expansion_list2, params_list):
         for exp1, exp2, prms in zip(expansion_list1, expansion_list2, params_list):
             results.append(f(exp1, exp2, prms))
             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
+        if match_expansion_axis_to_params is not None:
+            # return the params-matched axis to where it belongs
+            return np.moveaxis(np.array(results), 0, match_expansion_axis_to_params)
         return np.squeeze(np.array(results))
         return np.squeeze(np.array(results))
 
 
     return mapped_f
     return mapped_f

+ 9 - 2
charged_shells/patch_size.py

@@ -15,7 +15,8 @@ def charge_patch_size(ex: Expansion, phi: float = 0, theta0: Array | float = 0,
 
 
 def potential_patch_size(ex: Expansion, params: ModelParams,
 def potential_patch_size(ex: Expansion, params: ModelParams,
                          phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2,
                          phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2,
-                         match_expansion_axis_to_params: int = None):
+                         match_expansion_axis_to_params: int = None,
+                         skip_nozero_cases: bool = False):
 
 
     # this is more complicate to map over leading axes of the expansion as potential also depends on model parameters,
     # 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,
     # with some, such as kappaR, also being the parameters of the expansion in the first place. When mapping,
@@ -23,7 +24,13 @@ def potential_patch_size(ex: Expansion, params: ModelParams,
 
 
     @mapping.map_over_expansion
     @mapping.map_over_expansion
     def potential_zero(exp: Expansion, prms: ModelParams):
     def potential_zero(exp: Expansion, prms: ModelParams):
-        return bisect(lambda theta: potentials.charged_shell_potential(theta, phi, 1, exp, prms), theta0, theta1)
+        try:
+            return bisect(lambda theta: potentials.charged_shell_potential(theta, phi, 1, exp, prms), theta0, theta1)
+        except ValueError:
+            if skip_nozero_cases:
+                return 0.
+            else:
+                raise ValueError("Potential has no zero on the given interval.")
 
 
     return mapping.parameter_map_single_expansion(potential_zero, match_expansion_axis_to_params)(ex, params)
     return mapping.parameter_map_single_expansion(potential_zero, match_expansion_axis_to_params)(ex, params)
 
 

+ 43 - 18
charged_shells/rotational_path.py

@@ -19,18 +19,20 @@ class PairRotationalPath:
     rotations1: list[Quaternion] = field(default_factory=list)
     rotations1: list[Quaternion] = field(default_factory=list)
     rotations2: list[Quaternion] = field(default_factory=list)
     rotations2: list[Quaternion] = field(default_factory=list)
     x_axis: list[Array] = field(default_factory=list)
     x_axis: list[Array] = field(default_factory=list)
+    ticks: dict[float, str] = field(default_factory=dict)
     overlapping_last: bool = True
     overlapping_last: bool = True
     _default_x_axis: Array = None
     _default_x_axis: Array = None
 
 
-    def add(self, rotation1: Quaternion, rotation2: Quaternion, x_axis: Array = None):
+    def add(self, rotation1: Quaternion, rotation2: Quaternion, x_axis: Array = None,
+            start_name: str | float = None, end_name: str | float = None):
         rotation1, rotation2 = np.broadcast_arrays(rotation1, rotation2)
         rotation1, rotation2 = np.broadcast_arrays(rotation1, rotation2)
         self.rotations1.append(Quaternion(rotation1))
         self.rotations1.append(Quaternion(rotation1))
         self.rotations2.append(Quaternion(rotation2))
         self.rotations2.append(Quaternion(rotation2))
         if x_axis is None:
         if x_axis is None:
             x_axis = np.arange(len(rotation1)) if self._default_x_axis is None else self._default_x_axis
             x_axis = np.arange(len(rotation1)) if self._default_x_axis is None else self._default_x_axis
-        self.add_x_axis(x_axis)
+        self.add_x_axis(x_axis, start_name, end_name)
 
 
-    def add_x_axis(self, x_axis: Array):
+    def add_x_axis(self, x_axis: Array, start_name: str | float = None, end_name: str | float = None):
         try:
         try:
             last_x_val = self.x_axis[-1][-1]
             last_x_val = self.x_axis[-1][-1]
         except IndexError:
         except IndexError:
@@ -39,16 +41,23 @@ class PairRotationalPath:
             self.x_axis.append(x_axis + last_x_val)
             self.x_axis.append(x_axis + last_x_val)
         else:
         else:
             raise NotImplementedError('Currently only overlapping end points for x-axes are supported.')
             raise NotImplementedError('Currently only overlapping end points for x-axes are supported.')
+        # adding ticks to x_axis
+        s_n = x_axis[0] + last_x_val if start_name is None else start_name  # defaults to just numbers
+        e_n = x_axis[-1] + last_x_val if end_name is None else end_name
+        start_position = float(self.x_axis[-1][0])
+        end_position = float(self.x_axis[-1][-1])
+        if start_position not in self.ticks or start_name is not None:  # allows overwriting previous end with new start
+            self.ticks[start_position] = s_n
+        self.ticks[end_position] = e_n
 
 
     def set_default_x_axis(self, default_x_axis: Array):
     def set_default_x_axis(self, default_x_axis: Array):
         self._default_x_axis = default_x_axis
         self._default_x_axis = default_x_axis
 
 
     def add_euler(self, *, alpha1: Array = 0, beta1: Array = 0, gamma1: Array = 0,
     def add_euler(self, *, alpha1: Array = 0, beta1: Array = 0, gamma1: Array = 0,
-                  alpha2: Array = 0, beta2: Array = 0, gamma2: Array = 0,
-                  x_axis: Array = None):
+                  alpha2: Array = 0, beta2: Array = 0, gamma2: Array = 0, **add_kwargs):
         R1_euler = quaternionic.array.from_euler_angles(alpha1, beta1, gamma1)
         R1_euler = quaternionic.array.from_euler_angles(alpha1, beta1, gamma1)
         R2_euler = quaternionic.array.from_euler_angles(alpha2, beta2, gamma2)
         R2_euler = quaternionic.array.from_euler_angles(alpha2, beta2, gamma2)
-        self.add(Quaternion(R1_euler), Quaternion(R2_euler), x_axis)
+        self.add(Quaternion(R1_euler), Quaternion(R2_euler), **add_kwargs)
 
 
     def stack_rotations(self) -> (Quaternion, Quaternion):
     def stack_rotations(self) -> (Quaternion, Quaternion):
         return Quaternion(np.vstack(self.rotations1)), Quaternion(np.vstack(self.rotations2))
         return Quaternion(np.vstack(self.rotations1)), Quaternion(np.vstack(self.rotations2))
@@ -56,6 +65,22 @@ class PairRotationalPath:
     def stack_x_axes(self) -> Array:
     def stack_x_axes(self) -> Array:
         return np.hstack(self.x_axis)
         return np.hstack(self.x_axis)
 
 
+    def plot_style(self, fig, ax,
+                   energy_units: interactions.EnergyUnit = 'kT',
+                   legend: bool = True,
+                   size: tuple | None = (6, 2.5),
+                   legend_loc: str = 'upper left'):
+        ax.axhline(y=0, c='k', linestyle=':')
+        if legend:
+            ax.legend(fontsize=15, frameon=False, loc=legend_loc, bbox_to_anchor=(0.6, 1))
+        ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
+        # ax.set_xlabel('angle', fontsize=15)
+        ax.set_ylabel(f'$U [{energy_units}]$', fontsize=15)
+        ax.set_xticks(list(self.ticks.keys()), list(self.ticks.values()), fontsize=15)
+        if size is not None:
+            fig.set_size_inches(size)
+        plt.tight_layout()
+
 
 
 @dataclass
 @dataclass
 class PathEnergyPlot:
 class PathEnergyPlot:
@@ -72,6 +97,13 @@ class PathEnergyPlot:
     def __post_init__(self):
     def __post_init__(self):
         if not isinstance(self.dist, Array):
         if not isinstance(self.dist, Array):
             self.dist = np.array([self.dist])
             self.dist = np.array([self.dist])
+        # we add 1 to match_expansion_axis as rotations take the new leading axis
+        if self.match_expansion_axis_to_params is not None:
+            self.match_expansion_axis_to_params += 1
+
+    def plot_style(self, fig, ax, energy_units: interactions.EnergyUnit = 'kT', legend: bool = True,
+                   size: tuple = (8.25, 4.125)):
+        self.rot_path.plot_style(fig, ax, energy_units, legend=legend, size=size)
 
 
     def evaluate_energy(self):
     def evaluate_energy(self):
         energy = []
         energy = []
@@ -118,14 +150,11 @@ class PathEnergyPlot:
         energy = self.evaluate_path(norm_euler_angles=norm_euler_angles)
         energy = self.evaluate_path(norm_euler_angles=norm_euler_angles)
         x_axis = self.rot_path.stack_x_axes()
         x_axis = self.rot_path.stack_x_axes()
 
 
-        fig, ax = plt.subplots()
-        ax.axhline(y=0, c='k', linestyle=':')
+        print(plt.figaspect(0.5) * 8.25)
+
+        fig, ax = plt.subplots(figsize=plt.figaspect(0.5) * 8.25)
         ax.plot(x_axis, np.squeeze(energy), label=labels)
         ax.plot(x_axis, np.squeeze(energy), label=labels)
-        ax.legend(fontsize=12)
-        ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
-        ax.set_xlabel('angle', fontsize=13)
-        ax.set_ylabel('U', fontsize=13)
-        plt.tight_layout()
+        self.plot_style(fig, ax, energy_units=self.units)
         if save_as is not None:
         if save_as is not None:
             plt.savefig(save_as, dpi=600)
             plt.savefig(save_as, dpi=600)
         plt.show()
         plt.show()
@@ -135,14 +164,10 @@ class PathEnergyPlot:
         normalization = self.normalization(norm_euler_angles)
         normalization = self.normalization(norm_euler_angles)
 
 
         fig, ax = plt.subplots()
         fig, ax = plt.subplots()
-        ax.axhline(y=0, c='k', linestyle=':')
         for x_axis, energy in zip(self.rot_path.x_axis, energy_list):
         for x_axis, energy in zip(self.rot_path.x_axis, energy_list):
             energy, norm = np.broadcast_arrays(energy, normalization)
             energy, norm = np.broadcast_arrays(energy, normalization)
             ax.plot(x_axis, energy / norm)
             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)
+        self.plot_style(fig, ax, energy_units=self.units, legend=False)
         if save_as is not None:
         if save_as is not None:
             plt.savefig(save_as, dpi=600)
             plt.savefig(save_as, dpi=600)
         plt.show()
         plt.show()

+ 6 - 0
charged_shells/units.json

@@ -0,0 +1,6 @@
+{
+  "distance": 1e-09,
+  "charge": 1.6021766340000001e-19,
+  "voltage": 1,
+  "concentrationM": 0.001
+}

+ 4 - 4
charged_shells/units_and_constants.py

@@ -1,4 +1,6 @@
 from dataclasses import dataclass
 from dataclasses import dataclass
+from pathlib import Path
+import json
 
 
 
 
 @dataclass
 @dataclass
@@ -28,10 +30,8 @@ class Constants:
     e0: float
     e0: float
 
 
 
 
-base_units = {'distance': 1e-9,
-              'charge': 1.602176634 * 1e-19,
-              'voltage': 1,
-              'concentrationM': 1e-3}
+with open(Path(__file__).resolve().parent.joinpath('units.json')) as f:
+    base_units = json.load(f)
 
 
 
 
 UNITS = BaseUnits(**base_units)
 UNITS = BaseUnits(**base_units)

+ 6 - 0
config.json

@@ -0,0 +1,6 @@
+{
+  "emanuele_data":  "/home/andraz/ChargedShells/data_Emanuele/2024-06-24",
+  "figures": "/home/andraz/ChargedShells/Figures",
+  "figure_data": "/home/andraz/ChargedShells/Figures/Data",
+  "expansion_data": "/home/andraz/ChargedShells/ExpansionData/"
+}

+ 54 - 0
tests/expansion_mapping_test.py

@@ -0,0 +1,54 @@
+import unittest
+import numpy as np
+from charged_shells import expansion, interactions, parameters, units_and_constants, mapping, potentials
+from functools import partial
+
+
+EPSILON_0 = units_and_constants.CONSTANTS.epsilon0
+
+
+class IsotropicTest(unittest.TestCase):
+
+    def setUp(self):
+        self.charge = 560
+        self.kappaR = np.array([0.1, 1, 3, 5, 10, 20])
+        self.radius = 150
+        self.dist = 2 * self.radius
+        self.sigma0 = self.charge / (4 * np.pi * self.radius ** 2)
+        self.ex1 = expansion.MappedExpansionQuad(0, self.kappaR, 0, l_max=10, sigma0=self.sigma0)
+        self.ex2 = self.ex1.clone()
+        self.params = parameters.ModelParams(kappaR=self.kappaR, R=self.radius)
+
+    def test_potential(self):
+        theta = np.linspace(0, np.pi, 100)
+
+        def cs_potential_fn(ex, params):
+            return potentials.charged_shell_potential(theta, 0., self.dist / self.radius, ex, params)
+
+        mapped_cs_potential_fn = mapping.parameter_map_single_expansion(cs_potential_fn, 0)
+        cs_potential = mapped_cs_potential_fn(self.ex1, self.params)
+
+        ic_potential = []
+        for p in mapping.unravel_params(self.params):
+            ic_potential.append(potentials.inverse_patchy_particle_potential(theta, 2., 0, self.sigma0,
+                                                                             (0., 0.), p, lmax=10))
+        ic_potential = np.array(ic_potential)
+
+        np.testing.assert_almost_equal(cs_potential, ic_potential)
+
+    def test_interaction(self):
+
+        int_analytic = (self.charge ** 2 / (4 * np.pi * EPSILON_0 * self.params.epsilon) *
+                        (np.exp(self.params.kappaR) / (1 + self.params.kappaR)) ** 2 * np.exp(-self.params.kappa * self.dist) / self.dist)
+
+        energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy,
+                                                                 dist=self.dist/self.radius, units='eV'),
+                                                         0)
+        int_comp = energy_fn(self.ex1, self.ex2, self.params)
+
+        print(int_analytic)
+        print(int_comp)
+        print(int_comp / int_analytic)
+
+        np.testing.assert_almost_equal(int_comp, int_analytic)
+

+ 22 - 2
tests/interactions_test.py

@@ -25,6 +25,24 @@ def v22_distance_test():
     plt.show()
     plt.show()
 
 
 
 
+def quadrupole_variation_test():
+
+    params = ModelParams(R=10, kappaR=3.29)
+    sigma2 = np.array([0.45, 0.5, 0.55, 0.6, 0.65])
+    ex0 = expansion.Expansion24(sigma2, 0, sigma0=0.1)
+    ex1 = ex0.clone()
+
+    ex1.rotate_euler(0, np.pi / 2, 0)
+
+    dist = np.linspace(2, 3.2, 100)
+    energy_array = np.zeros((dist.shape[0], len(sigma2)))
+    for i, d in enumerate(dist):
+        energy_array[i, ...] = interactions.charged_shell_energy(ex0, ex1, params, d)
+
+    plt.plot(dist, energy_array)
+    plt.show()
+
+
 def timing():
 def timing():
     params = ModelParams(R=150, kappaR=3)
     params = ModelParams(R=150, kappaR=3)
     ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=20)
     ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=20)
@@ -49,5 +67,7 @@ def timing():
 
 
 if __name__ == '__main__':
 if __name__ == '__main__':
 
 
-    v22_distance_test()
-    timing()
+    # v22_distance_test()
+    # timing()
+
+    quadrupole_variation_test()

+ 21 - 0
tests/janus_interaction_test.py

@@ -0,0 +1,21 @@
+from charged_shells import expansion, interactions, parameters
+import numpy as np
+
+
+def charge_energy_test():
+
+    params = parameters.ModelParams(R=150, kappaR=3)
+    ex1 = expansion.MappedExpansionDipole(0.5, params.kappaR, 0.001, l_max=30, sigma0=0.0002)
+    ex2 = ex1.clone()
+    ex3 = ex1.clone().inverse_sign(exclude_00=True)
+
+    ex1.rotate_euler(beta=np.pi / 2)
+
+    e1 = interactions.charged_shell_energy(ex1, ex2, params, dist=2)
+    e2 = interactions.charged_shell_energy(ex1, ex3, params, dist=2)
+
+    print(e1, e2)
+
+
+if __name__ == '__main__':
+    charge_energy_test()