charged-shells/analysis/peak_heigth.py

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