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