charged-shells/analysis/peak_heigth_janus.py

from matplotlib import gridspec
from charged_shells import expansion, interactions, mapping, functions
from charged_shells.parameters import ModelParams
import numpy as np
from typing import Literal
from pathlib import Path
from functools import partial
import json
from plot_settings import *
from dataclasses import dataclass
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(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_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 = expansion.MappedExpansionDipole(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():
    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
        config_data = json.load(config_file)

    a_bar = [0.2, 0.5, 0.8]
    # IC_peak_energy_kappaR_combined_plot(config_data, a_bar,
    #                                     save_as=Path(
    #                                         '/home/andraz/ChargedShells/Figures/final_figures/peak_combined_kappaR.png')
    #                                     )

    # a_bar = [0.1, 0.2, 0.3]
    # IC_peak_energy_charge_combined_plot(config_data, a_bar,
    #                                     save_as=Path('/home/andraz/ChargedShells/Figures/final_figures/janus_peak_combined_charge.png')
    #                                     )

if __name__ == '__main__':
    main()