charged-shells/analysis/sEE_minimum.py

import matplotlib.pyplot as plt
import numpy as np
from charged_shells import expansion, interactions, mapping, charge_distributions
from charged_shells.parameters import ModelParams
from functools import partial
from config import *
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable
from labellines import labelLine, labelLines
from matplotlib.colors import TwoSlopeNorm
from matplotlib.ticker import FuncFormatter

Expansion = expansion.Expansion


def sEE_minimum(ex: Expansion, params: ModelParams, accuracy=1e-2, dist=2., match_expansion_axis_to_params=None,
                angle_start: float = 0, angle_stop: float = np.pi / 2):
    ex2 = ex.clone()
    angle_range = np.linspace(angle_start, angle_stop, int((angle_stop - angle_start) / accuracy), endpoint=True)
    # angle_range = np.linspace(0.1, 0.7, int(1 / accuracy), endpoint=True)

    ex.rotate_euler(alpha=0, beta=angle_range, gamma=0)
    ex2.rotate_euler(alpha=0, beta=angle_range, gamma=0)

    if match_expansion_axis_to_params is not None:
        match_expansion_axis_to_params += 1

    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist),
                                                     match_expansion_axis_to_params)
    energy = energy_fn(ex, ex2, params)

    min_idx = np.argmin(energy, axis=0)
    min_energy = np.min(energy, axis=0)

    return min_energy, angle_range[min_idx]


def contours():

    kappaR = np.linspace(1, 10, 20)
    a_bar = np.linspace(0.2, 0.6, 20)
    params = ModelParams(R=150, kappaR=kappaR)
    ex = charge_distributions.create_mapped_quad_expansion(a_bar[:, None], kappaR[None, :], 0.001)
    min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=1, accuracy=0.01)

    kR_mesh, a_mesh = np.meshgrid(kappaR, a_bar)

    print(np.min(min_angle), np.max(min_angle))
    print(np.min(min_energy), np.max(min_energy))

    plt.contourf(kR_mesh, a_mesh, min_angle)
    # plt.imshow(min_angle)
    plt.show()

    plt.contourf(kR_mesh, a_mesh, min_energy, np.array([-9, -5, -2.5, -1, -0.5, -0.2, 0]))
    # plt.imshow(min_energy)
    plt.show()


def kappaR_dependence(kappaR, save_as=None, cmap=cm.jet):

    a_bar = np.linspace(0.12, 0.8, 15)
    params = ModelParams(R=150, kappaR=kappaR)

    ex = charge_distributions.create_mapped_quad_expansion(a_bar[None, :], kappaR[:, None], 0.001)
    min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=0, accuracy=0.001,
                                        angle_start=0.5, angle_stop=1.)

    kappaR_alt = np.array([0.01, 2, 5, 10, 50])
    params_alt = ModelParams(R=150, kappaR=kappaR_alt)
    a_bar_alt = np.linspace(np.min(a_bar) - 0.05, np.max(a_bar) + 0.05, 20)
    ex2 = charge_distributions.create_mapped_quad_expansion(a_bar_alt[:, None], kappaR_alt[None, :], 0.001)
    min_energy_alt, min_angle_alt = sEE_minimum(ex2, params_alt, match_expansion_axis_to_params=1, accuracy=0.001,
                                                angle_start=0.5, angle_stop=1.)

    colors = cmap(np.linspace(0, 1, len(a_bar)))
    sm = cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=np.min(a_bar), vmax=np.max(a_bar)))
    sm.set_array([])

    # Transformation to the position along our quadrupolar rotational path
    min_angle = np.pi - min_angle
    min_angle_alt = np.pi - min_angle_alt

    kappaR_labels = [rf'$\kappa R={kR:.1f}$' for kR in kappaR_alt]

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
    for me, ma, lbl in zip(min_energy_alt.T, min_angle_alt.T, kappaR_labels):
        ax.plot(ma, me, ls=':', c='k', label=lbl)
    labelLines(ax.get_lines(), align=False, fontsize=15,  xvals=(2.35, 2.65))

    for me, ma, lbl, c in zip(min_energy.T, min_angle.T, [rf'$\bar a={a:.2f}$' for a in a_bar], colors):
        ax.plot(ma, me, label=lbl, c=c)
    # ax.plot(min_angle, min_energy)
    # ax.legend(fontsize=17)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
    ax.set_xlabel('pry angle', fontsize=20)
    ax.set_ylabel('U', fontsize=20)
    ax.set_xlim(2.2, 2.65)
    ax.set_ylim(-31, 1)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.05)
    cax.tick_params(labelsize=15)
    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
    cbar.set_label(r'$\bar a$', rotation=0, labelpad=15, fontsize=20)

    plt.tight_layout()
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def IC_kappaR_dependence(which_kappa_lines: list,
                         save_as=None, cmap=cm.jet, file_suffix=""):
    em_data_path = ICI_DATA_PATH.joinpath("FIG_SUPP_sEE")
    em_data = np.load(em_data_path.joinpath(f"fig9{file_suffix}.npz"))['arr_0']
    # print(em_data.shape)

    a_bar, indices, counts = np.unique(em_data[:, 0], return_counts=True, return_inverse=True)
    print(a_bar)

    if not np.all(counts == counts[0]):
        raise ValueError("Data not reshapable.")

    colors = cmap(np.linspace(0, 1, len(a_bar)))
    sm = cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=np.min(a_bar), vmax=np.max(a_bar)))
    sm.set_array([])

    min_energy = em_data[:, 3].reshape(-1, counts[0])
    min_angle = em_data[:, 2].reshape(-1, counts[0])

    kappa_line_energy = []
    kappa_line_angle = []
    kappaR_labels = []
    for kappa in which_kappa_lines:
        ind = em_data[:, 1] == kappa
        if np.sum(ind) == 0:
            print(f'No lines with kR={kappa} in data. Available values: {np.unique(em_data[:, 1])}')
            continue
        kappa_line_energy.append(em_data[:, 3][ind])
        kappa_line_angle.append(em_data[:, 2][ind])
        kappaR_labels.append(rf'$\kappa R={kappa:.1f}$')


    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
    for me, ma, lbl in zip(kappa_line_energy, kappa_line_angle, kappaR_labels):
        sort = np.argsort(ma)
        ax.plot(ma[sort], me[sort], ls=':', c='k', label=lbl)

    labelLines(ax.get_lines(), align=False, fontsize=15, xvals=(2.4, 2.65))

    for me, ma, lbl, c in zip(min_energy, min_angle, [rf'$\bar a={a:.2f}$' for a in a_bar], colors):
        ax.plot(ma, me, label=lbl, c=c)
    # ax.plot(min_angle, min_energy)
    # ax.legend(fontsize=17)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
    ax.set_xlabel('pry angle', fontsize=20)
    ax.set_ylabel('U', fontsize=20)
    ax.set_xlim(2.2, 2.65)
    ax.set_ylim(-41, 1)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.05)
    cax.tick_params(labelsize=15)
    cbar = fig.colorbar(sm, cax=cax, orientation='vertical')
    cbar.set_label(r'$\bar a$', rotation=0, labelpad=15, fontsize=20)

    plt.tight_layout()
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


# def charge_dependence(charge, save_as=None):
#
#     a_bar = np.array([0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])
#     params = ModelParams(R=150, kappaR=kappaR)
#
#     ex = expansion.MappedExpansionQuad(a_bar[None, :], kappaR[:, None], 0.001)
#     min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=0, accuracy=0.001,
#                                         angle_start=0.5, angle_stop=1.)
#
#     kappaR_alt = np.array([0.01, 2, 5, 10, 50])
#     params_alt = ModelParams(R=150, kappaR=kappaR_alt)
#     a_bar_alt = np.linspace(np.min(a_bar) - 0.05, np.max(a_bar) + 0.05, 20)
#     ex2 = expansion.MappedExpansionQuad(a_bar_alt[:, None], kappaR_alt[None, :], 0.001)
#     min_energy_alt, min_angle_alt = sEE_minimum(ex2, params_alt, match_expansion_axis_to_params=1, accuracy=0.001,
#                                                 angle_start=0.5, angle_stop=1.)
#
#     fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
#     for me, ma, lbl in zip(min_energy_alt.T, min_angle_alt.T, [rf'$\bar a={a:.2f}$' for a in a_bar]):
#         ax.plot(ma, me, ls=':', c='k')
#     for me, ma, lbl in zip(min_energy.T, min_angle.T, [rf'$\bar a={a:.2f}$' for a in a_bar]):
#         ax.plot(ma, me, label=lbl)
#     # ax.plot(min_angle, min_energy)
#     ax.legend(fontsize=17)
#     ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
#     ax.set_xlabel('angle', fontsize=15)
#     ax.set_ylabel('U', fontsize=20)
#     ax.set_xlim(0.56, 0.93)
#     ax.set_ylim(-20, 1)
#     plt.tight_layout()
#     if save_as is not None:
#         plt.savefig(save_as, dpi=600)
#     plt.show()


def main():

    # contours()

    # kappaR_dependence(kappaR=np.linspace(0.1, 30, 20),
    #                   save_as=Path(config_data["figures"]).joinpath('sEE_min_kappaR_abar.png')
    #                   )

    IC_kappaR_dependence(which_kappa_lines=[0.01, 3., 5., 10., 50.], file_suffix="_1",
                         # save_as=FIGURES_PATH.joinpath("ICi_data").joinpath('IC_sEE_min_kappaR_abar.png')
                         )


if __name__ == '__main__':

    main()