charged-shells/analysis/patch_size_dependence.py

from matplotlib import pyplot as plt
from pathlib import Path
import numpy as np
from charged_shells import expansion, patch_size
from charged_shells.parameters import ModelParams
import time
import json


def plot_abar_dependence(*, save=False, save_data=False):
    a_bar = np.linspace(0.2, 0.8, 101)
    kappaR = np.array([0.1, 1, 3, 10, 30])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])

    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=1)

    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']

    if save_data:
        data_dict = {}
        for key, patch in zip(["kR=01", "kR=1", "kR=3", "kR=10", "kR=30"], ps.T):
            data_dict[key] = np.stack((a_bar, patch)).T
        with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
            config_data = json.load(config_file)
        np.savez(Path(config_data["figure_data"]).joinpath("fig_6.npz"), **data_dict)

    # fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    fig, ax = plt.subplots()
    for patch, kR, marker in zip(ps.T, kappaR, markers):
        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', marker=marker, markerfacecolor='none', markevery=5, linewidth=2.5, markersize=8)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
    ax.set_xlabel(r'$\bar a$', fontsize=20)
    ax.set_ylabel(r'$\theta_0$', fontsize=20)
    plt.legend(fontsize=18)
    plt.tight_layout()
    if save:
        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_new.pdf"), dpi=300)
    plt.show()


def plot_abar_charge_dependence(*, save=False):
    a_bar = np.linspace(0.2, 0.8, 101)
    kappaR = np.array([0.1, 1, 3, 10, 30])
    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])

    ps = patch_size.charge_patch_size(ex)

    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, kR, marker in zip(ps.T, kappaR, markers):
        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', marker=marker, markerfacecolor='none', markevery=5)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\bar a$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14)
    plt.tight_layout()
    if save:
        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_charge.pdf"), dpi=300)
    plt.show()


def plot_kappaR_charge_dependence(*, normalized=False, save=False):
    a_bar = np.array([0.2, 0.4, 0.6, 0.8])
    kappaR = np.linspace(0.01, 30, 100)
    ex = expansion.MappedExpansionQuad(a_bar=a_bar[None, :], sigma_tilde=0.001, l_max=30, kappaR=kappaR[:, None])

    ps = patch_size.charge_patch_size(ex)
    if normalized:
        ps = ps / ps[0][None, :]

    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, abar, marker in zip(ps.T, a_bar, markers):
        ax.plot(kappaR, patch, label=rf'$\bar a$ = {abar}', marker=marker, markerfacecolor='none', markevery=5)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14)
    plt.tight_layout()
    if save:
        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_charge_nonmonotonicity.pdf"), dpi=300)
    plt.show()


def plot_sigma0_dependence(*, save=False):

    # kappaR = np.array([0.1, 1, 3, 10, 30])
    kappaR = np.linspace(0.1, 15, 201)
    # sigma0 = np.linspace(-0.001, 0.0015, 6)
    sigma0 = np.array([-0.0002, 0, 0.0002])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = expansion.MappedExpansionQuad(a_bar=0.3, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)

    t0 = time.perf_counter()
    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0, skip_nozero_cases=True)
    t1 = time.perf_counter()
    print(f'Time: {t1 - t0}')

    markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, sgm, marker in zip(ps.T, sigma0, markers):
        nonzero_ps = np.nonzero(patch)
        ax.plot(kappaR[nonzero_ps], patch[nonzero_ps], label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=10)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14, loc='upper right')
    plt.tight_layout()
    if save:
        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_charge_abar03.pdf"), dpi=300)
    plt.show()


def plot_sigma0_dependence_relative(*, save=False):

    # kappaR = np.array([0.1, 1, 3, 10, 30])
    kappaR = np.linspace(0.1, 15, 201)
    # sigma0 = np.linspace(-0.001, 0.0015, 6)
    sigma0 = np.array([-0.0002, 0, 0.0002])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = expansion.MappedExpansionQuad(a_bar=0.8, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)
    ex_neutral = expansion.MappedExpansionQuad(a_bar=0.3, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=0)

    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0, skip_nozero_cases=True)
    ps_neutral = patch_size.potential_patch_size(ex_neutral, params, match_expansion_axis_to_params=0)

    markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, sgm, marker in zip(ps.T, sigma0, markers):
        nonzero_ps = np.nonzero(patch)
        ax.plot(kappaR[nonzero_ps], patch[nonzero_ps] / ps_neutral[nonzero_ps],
                label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=10)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14)
    plt.tight_layout()
    if save:
        plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential_charge_abar03_relative.pdf"), dpi=300)
    plt.show()


def main():

    plot_abar_dependence(save=True)

    # plot_sigma0_dependence(save=True)

    # plot_sigma0_dependence_relative(save=True)

    # plot_abar_charge_dependence(save=True)
    # plot_kappaR_charge_dependence(normalized=True, save=True)


if __name__ == '__main__':

    main()