charged-shells/analysis/peak_heigth.py

import matplotlib.pyplot as plt
from charged_shells import expansion, interactions, mapping
from charged_shells.parameters import ModelParams
import numpy as np
from typing import Literal
from pathlib import Path
from functools import partial

Array = np.ndarray
Expansion = expansion.Expansion
RotConfig = Literal['ep', 'pp']


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

    params = ModelParams(R=R, kappaR=kappaR)

    # energy = []
    # for params in params.unravel():
    #     ex1 = expansion.MappedExpansionQuad(a_bar, params.kappaR, 0.001, l_max=l_max)
    #     ex2 = ex1.clone()
    #     if which == 'ep':
    #         ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
    #     energy.append(interactions.charged_shell_energy(ex1, ex2, dist, params))
    #
    # energy = np.array(energy)

    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)

    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist), 1)
    energy = energy_fn(ex1, ex2, params)

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


if __name__ == '__main__':

    kappaR = np.arange(1, 10, 0.1)
    a_bar = np.arange(0.2, 0.8, 0.2)

    peak_energy_plot(kappaR, a_bar, which='ep',
                     # save_as=Path('/home/andraz/ChargedShells/Figures/nonmonotonicity_check_ep.pdf')
                     )