Added Janus symmetry support, improved plots

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

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

analysis/patch_shape_comparison.py

@@ -1,11 +1,11 @@
 from charged_shells import expansion, potentials, patch_size
 import numpy as np
 from charged_shells.parameters import ModelParams
-import matplotlib.pyplot as plt
 from pathlib import Path
 from matplotlib.lines import Line2D
 import json
 import quadrupole_model_mappings
+from plot_settings import *
 
 
 def ic_cs_comparison():
@@ -20,7 +20,7 @@ def ic_cs_comparison():
     dist = 1
 
     ex_cp = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30,
-                                             sigma0=sigma0)
+                                          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,
@@ -122,7 +122,7 @@ def abar_comparison():
         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()
+    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)
@@ -131,10 +131,10 @@ def abar_comparison():
     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.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.savefig(Path("/home/andraz/ChargedShells/Figures/potential_amplitude_comparison.pdf"), dpi=300)
     plt.show()
 
 
@@ -153,7 +153,7 @@ def charge_comparsion():
                                              l_max=30, sigma0=charges)
     pot = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)
 
-    fig, ax = plt.subplots()
+    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)
@@ -162,7 +162,7 @@ def charge_comparsion():
     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.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)
@@ -173,20 +173,29 @@ def charge_comparsion():
 
 
 def ic_cs_comparison2(save_data=False):
-    target_patch_size = 0.92
+    # 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, 1001)
+    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)
+    # 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 = []
@@ -210,38 +219,41 @@ def ic_cs_comparison2(save_data=False):
                  CSp_eta_0=np.stack((theta, potential_cs[1])).T,
                  CSp_eta_pos=np.stack((theta, potential_cs[2])).T)
 
-    fig, ax = plt.subplots()
+    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'))
+    # 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=3)
+        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[::50], 1000 * p_ic.T[::50], marker='o', facecolors='none', edgecolors='k', s=75)
-        ax.scatter(theta[::50], 1000 * p_cs.T[::50], marker='x', c='k', s=75)
+        # 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=20)
-    ax.set_ylabel(r'$\phi$ [mV]', fontsize=20)
+    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$']
+    custom_labels = ['$0$', r'$\pi/4$', r'$\pi/2$', r'$3\pi/4$', r'$\pi$']
     plt.axhline(y=0, color='black', linestyle=':')
     plt.axvline(x=target_patch_size, color='black', linestyle=':')
-    plt.xticks(custom_ticks, custom_labels, fontsize=13)
+    plt.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=12)
+    # 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/potential_ic_cs_comparison_new.pdf"), dpi=300)
+    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)
+    # models_comparison(save_data=False)
     # ic_cs_comparison()
-    # ic_cs_comparison2()
+    ic_cs_comparison2()
 
     # abar_comparison()
     # charge_comparsion()

analysis/patch_size_dependence.py

@@ -1,20 +1,24 @@
-from matplotlib import pyplot as plt
 from pathlib import Path
 import numpy as np
 from charged_shells import expansion, patch_size
 from charged_shells.parameters import ModelParams
 import time
 import json
+from plot_settings import *
 
 
 def plot_abar_dependence(*, save=False, save_data=False):
     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)
     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)
 
+    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']
 
     if save_data:
@@ -26,16 +30,20 @@ def plot_abar_dependence(*, save=False, save_data=False):
         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()
+    fig, ax = plt.subplots(figsize=(3, 2.5))
     for patch, kR, marker in zip(ps.T, kappaR, markers):
-        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', marker=marker, markerfacecolor='none', markevery=5, linewidth=2.5, markersize=8)
+        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=20)
-    ax.set_ylabel(r'$\theta_0$', fontsize=20)
-    plt.legend(fontsize=18)
+    ax.set_xlabel(r'$\bar a$', fontsize=15)
+    ax.set_ylabel(r'$\theta_0$', fontsize=15)
+    plt.legend(handlelength=0.6, fontsize=13, frameon=False, loc='upper right', bbox_to_anchor=(0.56, 0.58))
     plt.tight_layout()
+    plt.axhline(y=mark_ps, color='black', linestyle=':')
+    plt.axvline(x=0.5, color='black', linestyle=':')
     if save:
-        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_new.pdf"), dpi=300)
+        plt.savefig(Path("/home/andraz/ChargedShells/Figures/final_figures/patch_size_potential.png"), dpi=300)
     plt.show()
 
 

analysis/path_plot.py

@@ -1,267 +0,0 @@
-import numpy as np
-from charged_shells.rotational_path import PairRotationalPath, PathEnergyPlot
-from charged_shells import expansion
-from charged_shells.parameters import ModelParams
-import matplotlib.pyplot as plt
-from pathlib import Path
-import json
-import quadrupole_model_mappings
-
-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")
-
-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 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_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=plt.figaspect(0.5))
-    ax.plot(em_data[:, 0], em_data[:, 1], label='ICi')
-    ax.plot(x_axis, np.squeeze(energy), label='CSp')
-    # 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=600)
-    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=600)
-    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"))
-    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=600)
-    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"))
-    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=600)
-    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"))
-    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=600)
-    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("Emanuele_data").joinpath('IC_CS_quadrupole_path.pdf')
-    # )
-
-    # 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.001, 0.00, 0.001]), 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('quadrupole_kappaR_dep.png')
-    #                      )
-    #
-    # IC_abar_dependence(config_data, save_as=Path(config_data["figures"]).joinpath("Emanuele_data").
-    #                    joinpath('quadrupole_abar_dep.png'))
-    #
-    IC_sigma0_dependence(config_data, save_as=Path(config_data["figures"]).joinpath("Emanuele_data").
-                         joinpath('quadrupole_charge_dep_abar05_kappaR3.png'))
-
-
-if __name__ == '__main__':
-
-    main()
-

analysis/peak_heigth.py

@@ -1,4 +1,3 @@
-import matplotlib.pyplot as plt
 from charged_shells import expansion, interactions, mapping, functions
 from charged_shells.parameters import ModelParams
 import numpy as np
@@ -6,95 +5,147 @@ from typing import Literal
 from pathlib import Path
 from functools import partial
 import json
-from itertools import cycle
+from plot_settings import *
+from dataclasses import dataclass
 
 Array = np.ndarray
 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
 
-    params = ModelParams(R=R, kappaR=kappaR)
 
-    ex1 = expansion.MappedExpansionQuad(a_bar[:, None], params.kappaR[None, :], 0.001, l_max=l_max)
-    ex2 = ex1.clone()
-    if which == 'ep':
-        ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
+class PeakEP(Peak):
 
-    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist), 1)
-    energy = energy_fn(ex1, ex2, params)
+    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
 
-    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.plot(kappaR, en, label=rf'$\bar a = {ab:.1f}$')
-    ax.legend(fontsize=12)
-    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
-    ax.set_xlabel(r'$\kappa R$', fontsize=15)
-    ax.set_ylabel(rf'$\bar V$', fontsize=15)
-    plt.tight_layout()
-    if save_as is not None:
-        plt.savefig(save_as, dpi=600)
-    plt.show()
 
+class PeakPP(Peak):
 
-def IC_peak_energy_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):
+    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
 
-    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
-    em_data = np.load(em_data_path.joinpath("pair_energy_fig11.npz"))
-    data = em_data['fixA']
 
-    if which == 'ep':
-        column_idx = 4
-    elif which == 'pp':
-        column_idx = 3
-    else:
-        raise ValueError
+class PeakSEP(Peak):
 
-    abar, inverse, counts = np.unique(data[:, 1], return_counts=True, return_inverse=True)
-    all_kappaR = np.unique(data[:, 2])
+    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
 
-    kappaR = np.geomspace(np.min(all_kappaR), max_kappaR, 30)
-    params = ModelParams(R=R, kappaR=kappaR)
 
-    ex1 = expansion.MappedExpansionQuad(np.array(a_bar)[:, None], params.kappaR[None, :], 0.001, l_max=20)
-    ex2 = ex1.clone()
-    if which == 'ep':
-        ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
+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), 1)
-    energy = energy_fn(ex1, ex2, params)
+    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)):
-        ab = np.around(abar[i], 5)
-        if ab in np.around(a_bar, 5):
+    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 = data[idx, 2][data[idx, 2] <= max_kappaR]
-            en = data[idx, column_idx][data[idx, 2] <= max_kappaR]
-            if quad_correction:
-                en *= extra_correction(ab, kR)
+            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)
-            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
+            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()}
@@ -105,20 +156,57 @@ def IC_peak_energy_plot(config_data: dict,
         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)
+
+    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=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)
+
     colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
 
-    fig, ax = plt.subplots()
+    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=2)
+    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.set_xlabel(r'$\kappa R$', fontsize=15)
 
-    y_label = r'$U_{pp}$' if which == 'pp' else r'$|U_{ep}|$'
-    ax.set_ylabel(y_label, fontsize=15)
-    ax.set_yscale('log')
+    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:
@@ -126,39 +214,272 @@ def IC_peak_energy_plot(config_data: dict,
     plt.show()
 
 
-def extra_correction(a_bar, kappaR):
-    y = a_bar * kappaR
-    # return (kappaR ** 2 * a_bar ** 4 * functions.sph_bessel_k(1, kappaR) / functions.sph_bessel_k(3, kappaR) *
-    #         np.sinh(y) / (-3 * y * np.cosh(y) + (3 + y ** 2) * np.sinh(y)))
-    return (a_bar ** 2 * functions.sph_bessel_k(1, kappaR) / functions.sph_bessel_k(3, kappaR)
-            * functions.sph_bessel_i(0, y) / functions.sph_bessel_i(2, y))
+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()
+
+
+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)
 
-    # kappaR = np.arange(0.5, 50, 0.1)
-    # a_bar = np.arange(0.2, 0.8, 0.2)
-    # peak_energy_plot(kappaR, a_bar, which='pp',
-    #                  # save_as=Path('/home/andraz/ChargedShells/Figures/nonmonotonicity_check_ep.pdf')
-    #                  )
+    # 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.2, 0.32, 0.52, 0.8]
-    # a_bar = list(np.arange(0.2, 0.81, 0.04))
-    IC_peak_energy_plot(config_data, a_bar=a_bar, which='pp', max_kappaR=50,
-                        save_as=Path('/home/andraz/ChargedShells/Figures/Emanuele_data/both_nonmonotonicity_check_pp.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
+                        quad_correction=False,
+                        log_y=False
                         )
 
-    # kr = np.linspace(0.1, 100, 1000)
-    # for ab in [0.2, 0.32, 0.52, 0.8]:
-    #     plt.plot(kr, extra_correction(ab, kr))
-    # plt.xscale('log')
-    # plt.yscale('log')
-    # plt.show()
-
 
 if __name__ == '__main__':
     

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'])

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

charged_shells/expansion.py

@@ -65,17 +65,26 @@ class Expansion:
         return np.repeat(arr, 2 * self.l_array + 1, axis=axis)
 
     def rotate(self, rotations: Quaternion, rotate_existing=False):
-        # TODO: rotations are currently saved wrong if we start form existing coefficients not the og ones
+        if rotate_existing:
+            raise NotImplementedError("Rotation of possibly already rotated coefficients is not yet supported.")
         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: might be off for a minus sign for each? angle !!
         R_euler = quaternionic.array.from_euler_angles(alpha, beta, gamma)
         self.rotate(R_euler, rotate_existing=rotate_existing)
 
+    def 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):
         if not isinstance(theta, Array):
             theta = np.array([theta])
@@ -103,14 +112,17 @@ class Expansion24(Expansion):
         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,
                  a_bar: Array | float,
                  kappaR: Array | float,
                  sigma_tilde: float,
-                 l_max: int = 20,
+                 l_array: Array,
                  sigma0: float | Array = 0):
         """
         :param a_bar: distance between the center and off center charges
@@ -125,18 +137,41 @@ class MappedExpansionQuad(Expansion):
             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)
 
-        l_array = np.array([l for l in range(l_max + 1) if l % 2 == 0])
         a_bar_bc2, kappaR_bc2, l_array_expanded = np.broadcast_arrays(a_bar_bc[..., None],
-                                                                     kappaR_bc[..., None],
-                                                                     l_array[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))
+        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])
+        super().__init__(a_bar, kappaR, sigma_tilde, l_array, sigma0)
+
+
+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):
     """Expansion for a collection of smeared charges on the sphere."""
 
@@ -170,7 +205,7 @@ class GaussianCharges(Expansion):
                     * np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
                                           theta_k[None, None, :], phi_k[None, None, :])))
         coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
-        coefs = expansion_total_charge(coefs, sigma0)
+        coefs = expansion_total_charge(coefs, sigma0, l_array)
         l_array, coefs = purge_unneeded_l(l_array, coefs)
         super().__init__(l_array, coefs)
 
@@ -214,7 +249,7 @@ class SphericalCap(Expansion):
         # Rotating is implemented in such a way that it rotates every patch to every position,
         # hence the redundancy of out of diagonal elements.
         coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
-        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))
 
 
@@ -238,8 +273,11 @@ def rot_sym_expansion(l_array: Array, coefs: Array) -> Array:
     return full_coefs
 
 
-def expansion_total_charge(coefs: Array, sigma0: float | Array | None):
+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."""
+    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:
         return coefs
     if not isinstance(sigma0, Array):

charged_shells/rotational_path.py

@@ -65,14 +65,20 @@ class PairRotationalPath:
     def stack_x_axes(self) -> Array:
         return np.hstack(self.x_axis)
 
-    def plot_style(self, fig, ax, energy_units: interactions.EnergyUnit = 'kT'):
+    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=':')
-        ax.legend(fontsize=18)
+        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=20)
-        ax.set_xticks(list(self.ticks.keys()), list(self.ticks.values()), fontsize=18)
-        fig.set_size_inches(plt.figaspect(0.5))
+        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()
 
 
@@ -95,8 +101,9 @@ class PathEnergyPlot:
         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'):
-        self.rot_path.plot_style(fig, ax, energy_units)
+    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):
         energy = []
@@ -143,7 +150,9 @@ class PathEnergyPlot:
         energy = self.evaluate_path(norm_euler_angles=norm_euler_angles)
         x_axis = self.rot_path.stack_x_axes()
 
-        fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
+        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)
         self.plot_style(fig, ax, energy_units=self.units)
         if save_as is not None:
@@ -158,7 +167,7 @@ class PathEnergyPlot:
         for x_axis, energy in zip(self.rot_path.x_axis, energy_list):
             energy, norm = np.broadcast_arrays(energy, normalization)
             ax.plot(x_axis, energy / norm)
-        self.plot_style(fig, ax, energy_units=self.units)
+        self.plot_style(fig, ax, energy_units=self.units, legend=False)
         if save_as is not None:
             plt.savefig(save_as, dpi=600)
         plt.show()

config.json

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

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