from matplotlib import gridspec from charged_shells import expansion, interactions, mapping, functions, charge_distributions from charged_shells.parameters import ModelParams import numpy as np from typing import Literal from functools import partial from config import * from plot_settings import * import peak_heigth Array = np.ndarray Expansion = expansion.Expansion class JanusPeakPP(peak_heigth.Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.emanuele_data_column = 3 self.y_label = r'$V_\mathrm{PP, 1}$' self.ex1 = ex.clone() self.ex2 = ex.clone() self.ex1.rotate_euler(beta=np.pi, alpha=0, gamma=0) self.ex2.rotate_euler(beta=np.pi, alpha=0, gamma=0) self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion class JanusPeakPPinv(peak_heigth.Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.emanuele_data_column = 4 self.y_label = r'$V_\mathrm{PP, 2}$' self.ex1 = ex.clone() self.ex2 = ex.clone().inverse_sign(exclude_00=True) self.ex1.rotate_euler(beta=np.pi, alpha=0, gamma=0) self.ex2.rotate_euler(beta=np.pi, alpha=0, gamma=0) self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion class JanusPeakEP(peak_heigth.Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.emanuele_data_column = 4 self.y_label = r'$V_{EP, 1}$' self.ex1 = ex.clone() self.ex2 = ex.clone() self.ex1.rotate_euler(beta=np.pi/2, alpha=0, gamma=0) self.ex2.rotate_euler(beta=np.pi/2, alpha=0, gamma=0) self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion class JanusPeakEPinv(peak_heigth.Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.emanuele_data_column = 4 self.y_label = r'$V_{EP, 2}$' self.ex1 = ex.clone() self.ex2 = ex.clone().inverse_sign(exclude_00=True) self.ex1.rotate_euler(beta=np.pi/2, alpha=0, gamma=0) self.ex2.rotate_euler(beta=np.pi/2, alpha=0, gamma=0) self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion def get_charge_energy_dicts(peak: peak_heigth.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) # abar_ic = [0.1, 0.2, 0.3] energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2), peak.kappaR_axis_in_expansion) energy = energy_fn(peak.ex1, peak.ex2, params) data_dict_ic = {} data_dict_cs = {} k = 0 for i in range(len(abar_ic)): ab = np.around(abar_ic[i], 5) if ab in np.around(abar_cs, 5): idx, = np.nonzero(inverse == i) charge = emanuele_data[idx, 0] en = emanuele_data[idx, peak.emanuele_data_column] sort = np.argsort(charge) if peak.log_y: data_dict_ic[ab] = np.stack((charge[sort], 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], 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 IC_peak_energy_charge_combined_plot(a_bar: list, R: float = 150, save_as: Path = None, log_y: bool = False): # em_data_path = (ICI_DATA_PATH.joinpath("FIG_11")) em_data_path = (ICI_DATA_PATH.joinpath("FIG_5_JANUS")).joinpath("FIG_5_JANUS_NEW_CHARGE") em_data = np.load(em_data_path.joinpath("pair_energy_PP_janus.npz")) data = em_data['changezc'] params = ModelParams(R=R, kappaR=3) sigma0 = np.linspace(-0.00099, 0.00099, 300) sigma_tilde = 0.00099 ex = charge_distributions.create_mapped_dipolar_expansion(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_pp = JanusPeakPP(ex, log_y) peak_pp_inv = JanusPeakPPinv(ex, log_y) peaks = [peak_pp, peak_pp_inv] # peak_ep = JanusPeakEP(ex, log_y) # peak_ep_inv = JanusPeakEPinv(ex, log_y) # peaks = [peak_ep, peak_ep_inv] 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)) fig = plt.figure(figsize=(4, 1.7)) gs = gridspec.GridSpec(1, 2, figure=fig) gs.update(left=0.125, right=0.975, top=0.99, bottom=0.21, wspace=0.35) axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1])] 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:.1f}$', c=current_color, linewidth=1.5) ax.plot(data_dict_cs[key][:, 0], data_dict_cs[key][:, 1], ls='--', c=current_color, linewidth=1.5, # label=rf'$\bar a = {key:.1f}$' ) ax.legend(fontsize=9, ncol=1, frameon=False, handlelength=0.7, loc='upper right', bbox_to_anchor=(0.75, 1.03), ) ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11) # if ax == axs[-1]: ax.set_xlabel(r'$\eta$', fontsize=11) ax.set_ylabel(peak.y_label, fontsize=11) if peak.log_y: ax.set_yscale('log') # ax.set_xscale('log') ax.yaxis.set_label_coords(-0.18, 0.5) ax.xaxis.set_label_coords(0.5, -0.12) axs[0].set_ylim(-23, 23) # 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(): a_bar = [0.2, 0.5, 0.8] # a_bar = [0.1, 0.2, 0.3] IC_peak_energy_charge_combined_plot(a_bar, save_as=FIGURES_PATH.joinpath('final_figures').joinpath('janus_peak_combined_charge.png') ) if __name__ == '__main__': main()