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 * from dataclasses import dataclass Array = np.ndarray Expansion = expansion.Expansion RotConfig = Literal['ep', 'pp', 'sep'] @dataclass class Peak: rot_config: str y_label: str ex1: Expansion ex2: Expansion ici_data_column: int log_y: bool kappaR_axis_in_expansion: int = None class PeakEP(Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.ici_data_column = 4 self.y_label = r'$|V_\mathrm{EP}|$' if log_y else r'$V_\mathrm{EP}$' self.ex1 = ex.clone() self.ex2 = ex.clone() self.ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0) self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion class PeakPP(Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.ici_data_column = 3 self.y_label = r'$V_\mathrm{PP}$' self.ex1 = ex.clone() self.ex2 = ex.clone() self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion class PeakSEP(Peak): def __init__(self, ex: Expansion, log_y: bool = False, kappaR_axis_in_expansion: int = None): self.ici_data_column = 5 self.y_label = r'$|V_\mathrm{sEP}|$' if log_y else r'$V_\mathrm{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=3 * np.pi / 4, gamma=0) self.log_y = log_y self.kappaR_axis_in_expansion = kappaR_axis_in_expansion def get_charge_energy_dicts(peak: Peak, params: ModelParams, emanuele_data: Array, sigma0: Array, abar_cs: list, sigma_tilde=0.001): abar_ic, inverse, counts = np.unique(emanuele_data[:, 1], return_counts=True, return_inverse=True) energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2), peak.kappaR_axis_in_expansion) energy = energy_fn(peak.ex1, peak.ex2, params) data_dict_ic = {} data_dict_cs = {} k = 0 for i in range(len(abar_ic)): ab = np.around(abar_ic[i], 5) if ab in np.around(abar_cs, 5): idx, = np.nonzero(inverse == i) charge = emanuele_data[idx, 0] en = emanuele_data[idx, peak.ici_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, ici_data: Array, kappaR: Array, abar_cs: list, max_kappaR: float = 50, only_loaded: bool = False): abar_ic, inverse = np.unique(ici_data[:, 1], return_inverse=True) energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2), peak.kappaR_axis_in_expansion) if not only_loaded: energy = energy_fn(peak.ex1, peak.ex2, params) else: energy = np.full((len(abar_cs), len(kappaR)), 1.) 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 = ici_data[idx, 2][ici_data[idx, 2] <= max_kappaR] en = ici_data[idx, peak.ici_data_column][ici_data[idx, 1] <= max_kappaR] sort = np.argsort(kR) if peak.log_y: data_dict_ic[ab] = np.stack((kR[sort], np.abs(en)[sort])).T data_dict_cs[ab] = np.stack((kappaR, np.abs(energy[k]))).T else: data_dict_ic[ab] = np.stack((kR[sort], en[sort])).T data_dict_cs[ab] = np.stack((kappaR, energy[k])).T k += 1 return data_dict_ic, data_dict_cs def get_abar_energy_dicts(peak: Peak, params: ModelParams, ici_data: Array, a_bar: Array, kR_cs: list, max_abar: float = 0.8): kR_ic, inverse = np.unique(ici_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 = ici_data[idx, 1][ici_data[idx, 1] <= max_abar] en = ici_data[idx, peak.ici_data_column][ici_data[idx, 1] <= max_abar] sort = np.argsort(abar) if peak.log_y: data_dict_ic[kr] = np.stack((abar[sort], np.abs(en)[sort])).T data_dict_cs[kr] = np.stack((a_bar, np.abs(energy[k]))).T else: data_dict_ic[kr] = np.stack((abar[sort], en[sort])).T data_dict_cs[kr] = np.stack((a_bar, energy[k])).T k += 1 return data_dict_ic, data_dict_cs def save_data(data_dict_ic, data_dict_cs, a_bar: list): if save_data: ic_save = {str(key): val for key, val in data_dict_ic.items()} cs_save = {str(key): val for key, val in data_dict_cs.items()} ic_save['abar'] = a_bar cs_save['abar'] = a_bar np.savez(FIGURES_DATA_PATH.joinpath("fig_11_IC.npz"), **ic_save) np.savez(FIGURES_DATA_PATH.joinpath("fig_11_CS.npz"), **cs_save) def IC_peak_energy_plot(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 = (ICI_DATA_PATH.joinpath("FIG_11")) em_data_path = (ICI_DATA_PATH.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 = charge_distributions.create_mapped_quad_expansion(np.array(a_bar)[:, None], params.kappaR[None, :], 0.00099, 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) corrected_data = {} for key, data in data_dict_ic.items(): corrected_data[key] = data[:, 1] * higher_multipole_ic_correction(data[:, 0], float(key), higher_l=4) corrected_data2 = {} for key, data in data_dict_ic.items(): corrected_data2[key] = data[:, 1] * higher_multipole_ic_correction(data[:, 0], float(key), higher_l=6) colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color']) fig = plt.figure(figsize=(2, 1.8)) gs = gridspec.GridSpec(1, 1, figure=fig) gs.update(left=0.25, right=0.95, top=0.97, bottom=0.21) ax = fig.add_subplot(gs[0, 0]) i = 0 for (key, data1), data2, cd, cd2 in zip(data_dict_ic.items(), data_dict_cs.values(), corrected_data.values(), corrected_data2.values()): i += 1 if i == 1: 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.plot(data1[:, 0], cd, ls='-', c='r', label=r'$l=4$') ax.plot(data1[:, 0], cd2, ls='-', c='g', label=r'$l=6$') ax.legend(fontsize=9, ncol=1, frameon=False, handlelength=0.7, loc='upper right', bbox_to_anchor=(0.7, 0.52), # bbox_to_anchor=(0.42, 0.9) ) ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11) ax.set_xlabel(r'$\kappa R$', fontsize=11) ax.set_ylabel(peak.y_label, fontsize=11) 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_abar_plot( 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 = (ICI_DATA_PATH.joinpath("FIG_11")) em_data_path = (ICI_DATA_PATH.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 = charge_distributions.create_mapped_quad_expansion(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(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 = (ICI_DATA_PATH.joinpath("FIG_11")) em_data_path = (ICI_DATA_PATH.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 = charge_distributions.create_mapped_quad_expansion(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(a_bar: list, R: float = 150, save_as: Path = None, min_kappaR: float = 0.01, max_kappaR: float = 50., ): # em_data_path = (ICI_DATA_PATH.joinpath("FIG_11")) em_data_path = (ICI_DATA_PATH.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 = charge_distributions.create_mapped_quad_expansion(np.sort(np.array(a_bar))[:, None], # sorting necessary as it is also in energy_dicts() params.kappaR[None, :], 0.00099, 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, only_loaded=False) data_ic.append(dict_ic) data_cs.append(dict_cs) legend_coords = [(0.55, 0.45), (0.55, 0.43), (0.55, 0.63)] # fig, axs = plt.subplots(3, 1, figsize=(2, 5.4)) fig = plt.figure(figsize=(2, 5.4)) gs = gridspec.GridSpec(3, 1, figure=fig) gs.update(left=0.25, right=0.95, top=0.99, bottom=0.07, hspace=0.3) axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[2, 0])] 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:.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) ax.axvline(x=10, c='k', ls=':', linewidth=1) ax.axvline(x=3, c='k', ls=':', linewidth=1) ax.axvline(x=1, c='k', ls=':', linewidth=1) ax.legend(fontsize=9, 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=11) # if ax == axs[-1]: ax.set_xlabel(r'$\kappa R$', 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.25, 0.5) ax.xaxis.set_label_coords(0.5, -0.12) # plt.subplots_adjust(left=0.1) # plt.tight_layout() # plt.subplots_adjust(left=0.25, right=0.97, top=0.85, bottom=0.15) if save_as is not None: plt.savefig(save_as, dpi=300) plt.show() 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_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.000297, 0.000297, 300) sigma_tilde = 0.00099 ex = charge_distributions.create_mapped_quad_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_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)) fig = plt.figure(figsize=(2, 5.4)) gs = gridspec.GridSpec(3, 1, figure=fig) gs.update(left=0.25, right=0.95, top=0.99, bottom=0.07, hspace=0.3) axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[2, 0])] 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) ax.legend(fontsize=9, 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=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) # plt.subplots_adjust(left=0.3) # plt.tight_layout() if save_as is not None: plt.savefig(save_as, dpi=300) plt.show() def higher_multipole_ic_correction(kappaR, abar, higher_l: int = 4): return (functions.sph_bessel_k(3, kappaR) * functions.sph_bessel_i(2, kappaR * abar) * abar ** (higher_l - 2) / (functions.sph_bessel_k(higher_l + 1, kappaR) * functions.sph_bessel_i(higher_l, kappaR * abar))) def main(): a_bar = [0.2, 0.5, 0.8] # IC_peak_energy_plot(a_bar=a_bar, which='pp', min_kappaR=0.01, max_kappaR=50, # save_as=FIGURES_PATH.joinpath('ICi_data').joinpath('peak_pp_kappaR_higher_correction_abar02.png'), # save_data=False, # quad_correction=False, # log_y=False # ) # IC_peak_energy_kappaR_combined_plot(a_bar, # save_as=FIGURES_PATH.joinpath('final_figures').joinpath('peak_combined_kappaR.png') # ) a_bar = [0.1, 0.2, 0.3] # IC_peak_energy_charge_plot(a_bar=a_bar, which='sep', # # save_as=FIGURES_PATH.joinpath('ICi_data').joinpath('peak_sep_charge.png'), # ) IC_peak_energy_charge_combined_plot(a_bar, save_as=FIGURES_PATH.joinpath('final_figures').joinpath('peak_combined_charge.png') ) kappaR = [0.01, 3.02407, 30] # IC_peak_energy_abar_plot(kappaR=kappaR, which='sep', min_abar=0.2, max_abar=0.8, # save_as=FIGURES_PATH.joinpath('ICi_data').joinpath('peak_sep_abar.png'), # save_data=False, # quad_correction=False, # log_y=False # ) if __name__ == '__main__': main()