
import pandas as pd
import numpy as np
import matplotlib
%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import rc
import scipy
# from scipy import interpolate
import sys
import os 
sys.path.append(os.environ['HOME'] + '/ns_plot/scripts')
from plot_ascii import combine_files, plot_0D_data
from plot_hdf5 import *
from helper import interpol, find_correct_path, default_rc_params, find_tan_using_ellipse
from process_sims import move_files
from configparser import ConfigParser, ExtendedInterpolation
import json
sys.path.append(os.environ['HOME'] + "/ns_plot/scripts/inverse_distance_weighting")

import glob
import re

# this notebook is constructed using kuibit for data management and analysis
import kuibit
from kuibit.simdir import SimDir


default_rc_params()


# For debugging purpose
# os.environ["SIM_LOG_PATH"] = 
# os.environ["RUN_NAME"] = 
# os.environ["NUM_DIR_RUNNING"] =


sim_log = ConfigParser()
sim_log.optionxform = str
config = ConfigParser()
config.optionxform = str
sim_log.read(os.environ["SIM_LOG_PATH"])
# sim_log.read("/glade/u/home/yufengl/ns_plot/ucar/sim_log.txt")
# this parameter will be set by the master run scripts (parent process)
run_name = os.environ["RUN_NAME"]
# print(os.environ["SIM_LOG_PATH"], os.environ["RUN_NAME"])
root_dir = os.path.realpath(os.path.join("/glade/scratch/yufengl", "simulations", run_name))
work_dir = os.path.join(sim_log["system"]["WORKDIR"], run_name)
# num_dir = int(sim_log[run_name]['num_runs'])  # -1 case
num_dir = int(os.environ["NUM_DIR_RUNNING"])
perturb_type = sim_log[run_name]["perturb"]
sim = sim_log[run_name]['parfile']
machine = sim_log["system"]["machine"]
rns_model = sim_log[run_name]['rns_model']
        
config.read(f"/glade/u/home/yufengl/ns_plot/params/{rns_model}_params.txt")
Req = Rx = float(config['RNS_param']["Rx"])
Ry = float(config['RNS_param']["Ry"])
RHO0 = float(config['RNS_param']["RHO0"])
PRESS_C = float(config['RNS_param']["PRESS_C"])
OMEGA_rot = float(config['RNS_param']["OMEGA_rot"])
M_ADM = float(config['RNS_param']["M_ADM"])
M_0 = float(config['RNS_param']["M_0"])
J_cocal = float(config['RNS_param']["J_cocal"])
t_P = 2 * numpy.pi / OMEGA_rot
t_dyn = 1/ np.sqrt(RHO0)
num_grid = float(sim_log[run_name]['num_grid'])
# t_P = t_dyn
dx = dy = dz = Req / num_grid * (2**4)
# r_gw = json.loads(config['grid']['r_gw'])

output_num = num_dir
# if plot_perturb:
#   sim = sim + f'_{perturb_type}'


from pathlib import Path

plot_output_dir = Path("../", run_name, f"output-{output_num:04}")

if not os.path.exists(plot_output_dir):
    os.makedirs(plot_output_dir)
    os.mkdir(Path(plot_output_dir, 'scalar'))
    os.mkdir(Path(plot_output_dir, '1D'))
    os.mkdir(Path(plot_output_dir, '2D'))
    print("created all directories for", plot_output_dir)
print(plot_output_dir)

print(f"Scratch path: {root_dir}, No of sims work on {num_dir}")
print(f"Work path: {work_dir}")

../C01s17_nonperturb_0412/output-0010
Scratch path: /glade/scratch/yufengl/simulations/C01s17_nonperturb_0412, No of sims work on 10
Work path: /glade/work/yufengl/simdata/C01s17_nonperturb_0412


path_template = find_correct_path(work_dir, 0, sim)

found path /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-{num_dir:04}/Baikal_ILGRMHD_C01s17/{filename}


simdir = SimDir(work_dir) #, pickle_file=os.path.join(work_dir, "simdir.pickle"))


timeseries = simdir.ts
gf = simdir.gf


# simdir.save(os.path.join(work_dir, "simdir.pickle"))


ham_constraint = timeseries.norm2["HGF"].to_numpy() * dx**3
mom_constraint = np.linalg.norm([timeseries.norm2["MU0GF"].to_numpy()*dx**3, timeseries.norm2["MU1GF"].to_numpy()*dx**3,
                                timeseries.norm2["MU2GF"].to_numpy()*dx**3], axis=0)


# Hamiltonian constraint
print("plotting Hamiltonian constraint")

plt.figure(figsize=(8, 8))
plt.semilogy(timeseries.norm2["HGF"].t / t_P, ham_constraint, label=sim)
plt.xlabel(r'$t / t_P$')
plt.ylabel(r'log $|H|$')
plt.ylim([1E-6, 1e-3])
plt.legend()
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_ml_ham_norm2.pdf'), bbox_inches="tight")

# Momentum constraint
print("plotting Momentum constraint")
plt.figure(figsize=(8, 8))
plt.semilogy(timeseries.norm2["MU0GF"].t / t_P, mom_constraint,  label=sim)
plt.xlabel(r'$t / t_P$')
plt.ylabel(r'log $|M|$')
plt.ylim([1E-9, 1E-3])
plt.legend()
# plt.tight_layout()
# plt.show()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_ml_mom_norm2.pdf'), bbox_inches="tight")

plotting Hamiltonian constraint
plotting Momentum constraint


# # COM movement
# print("plotting center of mass movement")
# if not plot_perturb:
#     center_loc = combine_files('hydro_analysis-hydro_analysis_rho_max_loc.maximum.asc', root_dir, num_dir)
# else:
#     center_loc = combine_files('hydro_analysis-hydro_analysis_rho_max_loc.maximum.asc', root_dir_perturb, num_dir_perturb)


plt.figure(figsize=(10,10))
plt.plot(timeseries.max["rho_b"].t / t_P, timeseries.max["rho_b"].to_numpy()/RHO0)
# plt.plot(rho_max_low[:, 1] * np.sqrt(RHO0), rho_max_low[:, -1]/RHO0)
# plt.plot(rho_un_press[:, 1] * np.sqrt(RHO0), rho_un_press[:, -1]/RHO0)
# plt.scatter(rho_max_without[:, 1]/(2*np.pi/0.03), rho_max_without[:, -1]/RHO0)
# plt.xlim([0, 0.02]) # this is happening since B = 0.5
# plt.ylim([0.99, 1.01])
plt.xlabel(r't/$t_{P}$')
plt.ylabel(r'$\rho_c / \rho_{c,t=0}$')
plt.savefig(Path(plot_output_dir, 'scalar', sim + '_rho_c.pdf'))


# max pressure vs. time
print("plotting maximum pressure vs. time")
plt.figure(figsize=(8, 8))
plt.plot(timeseries.max["P"].t / t_P, timeseries.max["P"].to_numpy() / PRESS_C, label=sim)
# plt.ylim([0, 1.1])
plt.xlabel(r'time $t/t_P$')
plt.ylabel(r'$p_c / p_{c,t=0}$')
plt.legend()
# plt.ylim([0.5, 1.5])
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_press_c.pdf'), bbox_inches="tight")

plotting maximum pressure vs. time


cons_sum = combine_files("illinoisgrmhd-grmhd_conservatives.sum.asc", path_template, num_dir)

File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0008/Baikal_ILGRMHD_C01s17/illinoisgrmhd-grmhd_conservatives.sum.asc
File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0009/Baikal_ILGRMHD_C01s17/illinoisgrmhd-grmhd_conservatives.sum.asc


rho_sum = cons_sum[:, [0, 1, 2]]
print("plotting rest mass over initial rest mass vs. time")
plt.figure(figsize=(8, 8))
plt.plot(rho_sum[:, 1] / t_P, rho_sum[:, 2] 
             * dx * dy * dz / M_0, label=sim)
# plt.ylim([0, 1.1])
plt.xlabel(r"$t / t_P$")
plt.ylabel(r"$M_0 / M_{0, cocal}$")
plt.legend()
plt.ylim([-1.5, 1.5])
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_rest_mass.pdf'), bbox_inches = "tight")

plotting rest mass over initial rest mass vs. time


print("plotting ADM mass over initial ADM mass vs. time")
plt.figure(figsize=(8, 8))
# adm_mass = combine_files('admmass-admmass_masses.maximum.asc', path_template, num_dir)
timeseries.max["ADMMass_VolumeMass[0]"].t
plt.plot(timeseries.max["ADMMass_VolumeMass[0]"].t / t_P, 
         timeseries.max["ADMMass_VolumeMass[0]"].to_numpy() / M_ADM, label=sim)
plt.xlabel(r"$t / t_P$")
plt.ylabel(r"$M_{ADM} / M_{ADM, cocal}$")
plt.legend()
plt.ylim([0.8, 1.2])
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_adm_mass.pdf'), bbox_inches="tight")

plotting ADM mass over initial ADM mass vs. time


print("plotting stellar modes")
# this should be from sim_log to compare perturbed sim with other sims
plot_perturb = False
if not plot_perturb:
    r_modes = combine_files('stellarmodes-cm.maximum.asc', path_template, num_dir)
else:
    r_modes = combine_files('stellarmodes-cm.maximum.asc', path_template_perturb, num_dir_perturb)
modes = np.zeros_like(r_modes)
# imodes = np.genfromtxt(path_template + 'stellarmodes-cm.inorm2.asc')
for i in range(2, 13, 2):
    modes[:, i] = np.sqrt(r_modes[:, i]**2 + r_modes[:, (i+1)]**2)
modes[:, 1] = r_modes[:, 1]

# |Cm|
plt.figure(figsize=(8, 8))
plt.semilogy(modes[:, 1] / t_P, modes[:, 4] / modes[:, 2], '-.', label='m=1')
plt.semilogy(modes[:, 1] / t_P, modes[:, 6] / modes[:, 2], ':', label='m=2')
plt.semilogy(modes[:, 1] / t_P, modes[:, 8] / modes[:, 2], '--', label='m=3')
plt.semilogy(modes[:, 1] / t_P, modes[:, 10] / modes[:, 2], '-', label='m=4')
# plt.semilogy(modes[:, 1], modes[:, 12] / modes[:,2], label='m=5')
plt.legend()
plt.xlabel(r"$t / t_P$")
plt.ylabel(r"$|C_m| / C_0$")
plt.ylim([1E-5, 1])
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_stellar_modes_magnitude.pdf'), bbox_inches="tight")

plotting stellar modes
File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0008/Baikal_ILGRMHD_C01s17/stellarmodes-cm.maximum.asc
File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0009/Baikal_ILGRMHD_C01s17/stellarmodes-cm.maximum.asc


# angular momentum
print("plotting angular momentum")
plt.figure(figsize=(8, 8))
angularmom = timeseries.max["Jphi"]
plt.plot(angularmom.t / t_P, angularmom.to_numpy() / J_cocal, label='Jphi int')
plt.legend()
plt.xlabel(r"$t / t_P$")
plt.ylabel(r'$J / J_{cocal}$')
# plt.ylim([0.8, 1.2])
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, 'scalar', sim +f'_{perturb_type}_output_{output_num}'  + '_angular_mom.pdf'), bbox_inches="tight")

plotting angular momentum


print('plotting the gw signal')
# radi = [1.99, 3.09, 4.42, 5.53, 6.63]
l = 2
m = 2
# m = 0, l
plt.figure(figsize=(8, 8))
# r_gw = [float(re.findall("\d\.\d+", name_path)[-1]) for name_path in glob.glob(path_template.format(num_dir=0, filename="mp_Psi4_l2_m2_r*.asc"))]
# r_gw.sort()
# for r in r_gw:
#     if not plot_perturb:
#         psi4 = combine_files("mp_Psi4_l"+str(l) +
#                             "_m"+str(m)+f"_r{r:.2f}.asc", path_template, num_dir)
#     else:
#         psi4 = combine_files("mp_Psi4_l"+str(l) +
#                         "_m"+str(m)+f"_r{r:.2f}.asc", path_template_perturb, num_dir_perturb)
#     plt.plot((psi4[:, 0] - r) / t_P, psi4[:, 1] * r, label=f'r={r}')
    
plt.plot((simdir.gws.outermost[(l, m)].t - simdir.gws.r_outer) / t_P, 
         simdir.gws.outermost[(l, m)].y * simdir.gws.r_outer, label=f'r={simdir.gws.r_outer}')
plt.xlabel(r'$(t-r)/t_P$')
plt.ylabel(r'$r \Psi_4 (l=2, m=2)$')
plt.legend()
# plt.tight_layout()
plt.savefig(Path(plot_output_dir, sim +f'_{perturb_type}_output_{output_num}'  + '_gw.pdf', bbox_inches="tight"))

plotting the gw signal

/glade/work/yufengl/conda-envs/py39/lib/python3.9/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part
  return np.asarray(x, float)


Iij = combine_files(
    "quadmoment-iij.maximum.asc", path_template, num_dir)
com = combine_files(
     "quadmoment-com.maximum.asc", path_template, num_dir)

File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0008/Baikal_ILGRMHD_C01s17/quadmoment-iij.maximum.asc
File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0009/Baikal_ILGRMHD_C01s17/quadmoment-iij.maximum.asc
File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0008/Baikal_ILGRMHD_C01s17/quadmoment-com.maximum.asc
File path does not exist /glade/work/yufengl/simdata/C01s17_nonperturb_0412/output-0009/Baikal_ILGRMHD_C01s17/quadmoment-com.maximum.asc


plt.figure(figsize=(10, 10))
t_dyn = 1 / np.sqrt(RHO0)
plt.plot(Iij[:, 1] / t_dyn, Iij[:, 2], label="Ixx")
plt.plot(Iij[:, 1]/ t_dyn, Iij[:, 3], label="Ixy")
plt.plot(Iij[:, 1]/ t_dyn, Iij[:, 4], label="Ixz")
plt.plot(Iij[:, 1]/ t_dyn, Iij[:, 5], label="Iyy")
plt.plot(Iij[:, 1]/ t_dyn, Iij[:, 6], label="Iyz")
plt.plot(Iij[:, 1]/ t_dyn, Iij[:, 7], label="Izz")
# plt.plot(Iij[:, 1], Iij[:, 2], label="Ixx")
plt.xlabel(r"$t_{dyn}$")
plt.ylabel(r"$I_{ij}$")
plt.legend()

<matplotlib.legend.Legend at 0x7f67b9234fa0>


times = cactus_hdf5_get_timeinfo(path_template.format(num_dir=num_dir, filename='illinoisgrmhd-grmhd_primitives_allbutbi.x.h5'))
iterations_init = sorted(times.keys())[0]
iterations_final = sorted(times.keys())[-1]


rho_x_1d = gf.x["rho_b"]
rho_y_1d = gf.y["rho_b"]
# Need to use 2D data to ensure that both x and y axis 
# data exists 
times = gf.xy["rho_b"].times
iterations = gf.xy["rho_b"].iterations

max_reflvl = 4


iteration_init = iterations[np.argmin(np.abs(iterations - iterations_init))]
iteration_final = iterations[np.argmin(np.abs(iterations - iterations_final))]


# principle_axes = {}
# for i, it in enumerate(iterations):
#     Qmat = np.array([[Iij[i, 2], Iij[i, 3]], 
#                     [Iij[i, 3], Iij[i, 5]]])
#     eigenvalues, eigenvectors = np.linalg.eig(Qmat)
#     idx = eigenvalues.argsort()[::-1]
#     #     print(eigenvalues[idx], eigenvectors[:, idx])
#     eigenvectors = eigenvectors[:, idx]
#     principle_axes[it] = eigenvectors


print("plotting density along x, y axis")
fig, ax = plt.subplots(1, 2, figsize=(20.0, 8.0))

# plot three iterations
for it, time in zip(iterations[-300::100], times[-300::100]):
    data_x = rho_x_1d[it].get_level(max_reflvl)
    data_y = rho_y_1d[it].get_level(max_reflvl)
    ax[0].plot(data_x.coordinates_meshgrid()[0] / Req, data_x.data / RHO0, label=fr't={time/t_P:.3g} $t_P$')
    ax[1].plot(data_y.coordinates_meshgrid()[0] / Req, data_y.data / RHO0, label=fr't={time/t_P:.3g} $t_P$')
    
ax[0].set_xlabel(r"$x /R_x$")
ax[0].set_ylabel(r"$\rho_x / \rho_{c,t=0}$")
ax[0].set_xlim([-1.2, 1.2])
ax[0].set_ylim([0, 1.2])
ax[0].legend(loc='upper right')
# plt.savefig(Path(plot_output_dir, sim+ f'_output_{num_dir}' +  '_rho_x' + '_plot.pdf', bbox_inches="tight")

ax[1].set_xlabel(r"$y /R_x$")
ax[1].set_ylabel(r"$\rho_y / \rho_{c,t=0}$")
ax[1].set_xlim([-1.2, 1.2])
ax[1].set_ylim([0, 1.2])
ax[1].legend(loc='upper right')
plt.savefig(Path(plot_output_dir, "1D", sim+ f'_output_{num_dir}' +  '_rho_x_y' + '_plot.pdf'), bbox_inches="tight")

plotting density along x, y axis


print("plotting rotation curve across x axis")
# Omega = (0, - v_z/r, v_y/r) on axis
fig, ax = plt.subplots(1, 2, figsize=(25., 8.))
for it, time in zip(iterations[-300::100], times[-300::100]):
    # Hydrobase
    lapse_on_xaxis = np.vstack([gf.x["alp"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                gf.x["alp"][it].get_level(max_reflvl).data]).T
    shift_y_on_xaxis = np.vstack([gf.x["betay"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                  gf.x["betay"][it].get_level(max_reflvl).data]).T
    vel_y_on_xaxis = np.vstack([gf.x["vel[1]"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                gf.x["vel[1]"][it].get_level(max_reflvl).data]).T
    rest_vel_x = lapse_on_xaxis[:, 1] * vel_y_on_xaxis[:, 1] - shift_y_on_xaxis[:, 1]
    omega_x = rest_vel_x / vel_y_on_xaxis[:, 0]
    ax[0].plot(vel_y_on_xaxis[:, 0] / Req, omega_x / OMEGA_rot, label=fr't={time/t_P:.3g} $t_P$')
    # ILGRMHD
    vel_y_on_xaxis = np.vstack([gf.x["vy"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                gf.x["vy"][it].get_level(max_reflvl).data]).T
    rest_vel_x = vel_y_on_xaxis[:, 1]
    omega_x = rest_vel_x / vel_y_on_xaxis[:, 0]
    ax[1].plot(vel_y_on_xaxis[:, 0] / Req, omega_x / OMEGA_rot, label=fr't={time/t_P:.3g} $t_P$')
ax[0].set_xlim([-1.2, 1.2])
ax[0].set_ylim([0, 1.5])
ax[0].set_xlabel(r'$x / R_x$')
ax[0].set_ylabel(r'$\Omega_{z,Hydrobase} / \Omega_{cocal}$')
ax[0].legend(loc='lower right')
ax[1].set_xlim([-1.2, 1.2])
ax[1].set_ylim([0, 1.5])
ax[1].set_xlabel(r'$x / R_x$')
ax[1].set_ylabel(r'$\Omega_{z,IGM} / \Omega_{cocal}$')
ax[1].legend(loc='lower right')
plt.savefig(Path(plot_output_dir, "1D", sim+ f'_output_{num_dir}' +  '_omega_z_on_xaxis' + '_plot.pdf'), bbox_inches='tight')

plotting rotation curve across x axis


# plot velocity on z axis
# ILGRMHD
plt.figure(figsize=(10, 8.))
for it, time in zip(iterations[-300::100], times[-300::100]):
    vel_z_on_xaxis = np.vstack([gf.x["vz"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                gf.x["vz"][it].get_level(max_reflvl).data]).T
    rest_vel_z = vel_z_on_xaxis[:, 1]
    omega_y = - rest_vel_z / vel_z_on_xaxis[:, 0]
    plt.plot(vel_z_on_xaxis[:, 0] / Req, omega_y / OMEGA_rot, label=fr't={time/t_P:.3g} $t_P$')
plt.xlim([-1.2, 1.2])
plt.ylim([-0.2, 0.2])
plt.xlabel(r'$x / R_x$')
plt.ylabel(r'$\Omega_{y,IGM} / \Omega_{cocal}$')
plt.legend(loc='upper right')
plt.savefig(Path(plot_output_dir, "1D", sim+ f'_output_{num_dir}' +  '_omega_y_on_xaxis' + '_plot.pdf'), bbox_inches='tight')


print("plotting lapse across x")
fig, ax = plt.subplots(1, 1, figsize=(18, 8.))
for it, time in zip(iterations[-300::100], times[-300::100]):
    lapse_on_xaxis = np.vstack([gf.x["alp"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                gf.x["alp"][it].get_level(max_reflvl).data]).T
    ax.plot(lapse_on_xaxis[:, 0] / Req, lapse_on_xaxis[:, 1], label=fr't={time/t_P:.3g} $t_P$')
    ax.set_xlabel(r'$x / R_x$')
    ax.set_ylabel(r'$\alpha$')
    ax.set_xlim([-1.2, 1.2])
    ax.legend(loc='upper right')
    # lapse_on_yaxis = cactus_hdf5_get_1d_data(
    #     f'{root_dir}/output-000{num_dir}/admbase-lapse.y.h5', 'ADMBASE::alp', it)
    # ax[1].plot(lapse_on_yaxis[:, 0] / Req, lapse_on_yaxis[:, 1])
    # ax[1].xlabel(r'$y / R_x$')
    # ax[1].ylabel(r'\alpha')
    # ax[1].set_xlim([-1.2 * Req, 1.2 * Req])
    # ax[1].legend(loc='upper right')
# plt.show()
plt.savefig(Path(plot_output_dir, "1D", sim+ f'_output_{num_dir}' +  '_lapse_x' + '_line_plot.pdf'), bbox_inches='tight')
# plt.close()

plotting lapse across x


# xy_2d_path = path_template.format(num_dir=num_dir, filename='illinoisgrmhd-grmhd_primitives_allbutbi.xy.h5')
# rho_in_xy_plane_init = cactus_hdf5_get_2d_data(xy_2d_path, 'ILLINOISGRMHD::rho_b', iteration_init)
# rho_in_xy_plane_final = cactus_hdf5_get_2d_data(xy_2d_path, 'ILLINOISGRMHD::rho_b', iteration_end)


# plt.figure(figsize=(16, 8))
# lin = np.linspace(-Req * 2, Req * 2, 100)
# plt.subplot(1, 2, 1)
# plt.scatter(rho_in_xy_plane_init[:, 0], rho_in_xy_plane_init[:, 1], c=np.log(rho_in_xy_plane_init[:, 2]/RHO0))
# plt.plot(principle_axes[iteration_init][1, 0] * lin, principle_axes[iteration_init][1, 1] * lin, label='x axis')
# plt.plot(principle_axes[iteration_init][0, 0] * lin, principle_axes[iteration_init][0, 1] * lin, label='y axis')
# plt.legend()
# plt.xlim([-2 * Req, 2 * Req])
# plt.ylim([-2 * Req, 2 * Req])

# plt.subplot(1, 2, 2)
# plt.scatter(rho_in_xy_plane_final[:, 0], rho_in_xy_plane_final[:, 1], c=np.log(rho_in_xy_plane_final[:, 2]/RHO0))
# plt.plot(principle_axes[iteration_end][1, 0] * lin, principle_axes[iteration_end][1, 1] * lin, label='x axis')
# plt.plot(principle_axes[iteration_end][0, 0] * lin, principle_axes[iteration_end][0, 1] * lin, label='y axis')
# plt.legend()
# plt.xlim([-2 * Req, 2 * Req])
# plt.ylim([-2 * Req, 2 * Req])


vmin, vmax = 0, 1


print("plotting density xy")
# plt.figure(figsize=(24, 10))
# plt.subplot(1, 2, 1)
fig, ax = plt.subplots(1, 2, figsize=(24, 10))
rho_xy_2d = gf.xy["rho_b"][iteration_init].get_level(max_reflvl)
im = ax[0].pcolormesh(*rho_xy_2d.coordinates_meshgrid()[::-1], 
                      (rho_xy_2d.data/RHO0), cmap='jet', vmin=vmin, vmax=vmax)
# ax[0].legend(loc='lower left')
ax[0].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[0].set_xlim(xmin,xmax)
ax[0].set_ylim(ymin,ymax)
ax[0].set_xlabel("x [M]")
ax[0].set_ylabel("y [M]")
ax[0].set_title(f'{rho_xy_2d.time /t_P:.3f} $t_P$', 
               fontsize=20)
fig.colorbar(im, ax=ax[0])
x0, y0 = rho_xy_2d.coordinates_meshgrid()[::-1]
rho0 = rho_xy_2d.data

rho_xy_2d = gf.xy["rho_b"][iteration_final].get_level(max_reflvl)
im = ax[1].pcolormesh(*rho_xy_2d.coordinates_meshgrid()[::-1], 
                      (rho_xy_2d.data/RHO0), cmap='jet', vmin=vmin, vmax=vmax)
# ax[0].legend(loc='lower left')
ax[1].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[1].set_xlim(xmin,xmax)
ax[1].set_ylim(ymin,ymax)
ax[1].set_xlabel("x [M]")
ax[1].set_ylabel("y [M]")
ax[1].set_title(fr'{rho_xy_2d.time/t_P:.3f} $t_P$',
                fontsize=20)
fig.colorbar(im, ax=ax[1])
xf, yf = rho_xy_2d.coordinates_meshgrid()[::-1]
rhof = rho_xy_2d.data

plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_rho_xy' + '_pseudocolor_plot.pdf'), bbox_inches="tight")
plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_rho_xy' + '_pseudocolor_plot.png'), bbox_inches="tight")

plotting density xy


fig = plt.figure(figsize=(10.0,8.0))
ax1 = fig.add_subplot(1,1,1)
ax1.grid(color='black', linestyle=':', linewidth=0.5)
levels = [0.4, 0.5, 0.7, 0.9, 0.95]
levcol = [ 'green', 'cyan', 'brown', 'orange', 'magenta' ]
contour1 = ax1.contour(x0/Req, y0/Req, rho0/RHO0, levels=levels, colors=levcol, linestyles='dotted')
contour2 = ax1.contour(xf/Req, yf/Req, rhof/RHO0, levels=levels, colors=levcol, linestyles='-')
# ax1.clabel(contour1, levels, inline=True, inline_spacing=0, fontsize=14)
# ax1.clabel(contour2, levels, inline=True, inline_spacing=1, fontsize=14)
ax1.set_xlim([-1, 1])
ax1.set_ylim([-1, 1])
ax1.set_title(r"$\rho_{xy}$ of initial time (dotted) vs. final time (solid)")
# ax1.legend()
plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_vel_xy' + '_contour.png'), bbox_inches="tight")


print("plotting density xz")
# plt.figure(figsize=(24, 10))
# plt.subplot(1, 2, 1)
fig, ax = plt.subplots(1, 2, figsize=(24, 10))
rho_xz_2d = gf.xz["rho_b"][iteration_init].get_level(max_reflvl)
im = ax[0].pcolormesh(*rho_xz_2d.coordinates_meshgrid()[::-1], 
                      (rho_xz_2d.data/RHO0), cmap='jet', vmin=vmin, vmax=vmax)
# ax[0].legend(loc='lower left')
ax[0].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[0].set_xlim(xmin,xmax)
ax[0].set_ylim(ymin,ymax)
ax[0].set_xlabel("x [M]")
ax[0].set_ylabel("z [M]")
ax[0].set_title(fr'{gf.xz["rho_b"].time_at_iteration(iteration_init)/t_P:.4f} $t_P$', 
               fontsize=20)
fig.colorbar(im, ax=ax[0])
x0, z0 = rho_xz_2d.coordinates_meshgrid()[::-1]
rho0 = rho_xz_2d.data

rho_xz_2d = gf.xz["rho_b"][iteration_final].get_level(max_reflvl)
im = ax[1].pcolormesh(*rho_xz_2d.coordinates_meshgrid()[::-1], 
                      (rho_xz_2d.data/RHO0), cmap='jet', vmin=vmin, vmax=vmax)
# ax[0].legend(loc='lower left')
ax[1].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[1].set_xlim(xmin,xmax)
ax[1].set_ylim(ymin,ymax)
ax[1].set_xlabel("x [M]")
ax[1].set_ylabel("z [M]")
ax[1].set_title(fr'{gf.xz["rho_b"].time_at_iteration(iteration_final)/t_P:.4f} $t_P$',
                fontsize=20)
fig.colorbar(im, ax=ax[1])
xf, zf = rho_xz_2d.coordinates_meshgrid()[::-1]
rhof = rho_xz_2d.data

plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_rho_xz' + '_pseudocolor_plot.pdf'), bbox_inches="tight")
plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_rho_xz' + '_pseudocolor_plot.png'), bbox_inches="tight")

plotting density xz


# xz plane 

fig = plt.figure(figsize=(10.0,8.0))
ax1 = fig.add_subplot(1,1,1)
ax1.grid(color='black', linestyle=':', linewidth=0.5)
levels = [0.4, 0.5, 0.7, 0.9, 0.95]
levcol = [ 'green', 'cyan', 'brown', 'orange', 'magenta' ]
contour1 = ax1.contour(x0/Req, z0/Req, rho0/RHO0, levels=levels, colors=levcol, linestyles='dotted')
contour2 = ax1.contour(xf/Req, zf/Req, rhof/RHO0, levels=levels, colors=levcol, linestyles='-')
ax1.set_xlim([-1, 1])
ax1.set_ylim([-1, 1])
ax1.set_title(r"$\rho_{xz}$ of initial time (dotted) vs. final time (solid)")
plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_vel_xz' + '_contour.png'), bbox_inches="tight")


# get quiver plot using unstructured interpolation of all data (slow)
print("plotting velocity xy")
fig, ax = plt.subplots(1, 2, figsize=(20.0, 8.0))
rho_xy_2d = gf.xy["rho_b"][iteration_init].get_level(max_reflvl)
velx_in_xy_plane_init = gf.xy["vx"][iteration_init].get_level(max_reflvl)
vely_in_xy_plane_init = gf.xy["vy"][iteration_init].get_level(max_reflvl)
im = ax[0].pcolormesh(*rho_xy_2d.coordinates_meshgrid()[::-1], 
                      (rho_xy_2d.data/RHO0), cmap='jet', 
                      label=fr'{velx_in_xy_plane_init.time /t_P:.3f} $t_P$', vmin=vmin, vmax=vmax)
y, x = velx_in_xy_plane_init.coordinates_meshgrid()
# create a true false array as a mask to sample grid
mask = np.zeros_like(x, dtype=bool)
mask[::10, ::10] = True
ax[0].quiver(x[mask], y[mask], velx_in_xy_plane_init.data[mask], vely_in_xy_plane_init.data[mask], 
             label="t=%.3g" % (velx_in_xy_plane_init.time / t_P), 
             angles='xy',scale=6,width=0.005,color='black')
ax[0].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[0].set_xlim(xmin,xmax)
ax[0].set_ylim(ymin,ymax)
ax[0].set_xlabel("x [M]")
ax[0].set_ylabel("y [M]")
fig.colorbar(im, ax=ax[0])
ax[0].set_title("3-velocity in xy plane")
ax[0].legend(loc='upper right')


rho_xy_2d = gf.xy["rho_b"][iteration_final].get_level(max_reflvl)
velx_in_xy_plane_final = gf.xy["vx"][iteration_final].get_level(max_reflvl)
vely_in_xy_plane_final = gf.xy["vy"][iteration_final].get_level(max_reflvl)
im = ax[1].pcolormesh(*rho_xy_2d.coordinates_meshgrid()[::-1], 
                      (rho_xy_2d.data/RHO0), cmap='jet', 
                      label=fr'{rho_xy_2d.time /t_P:.3f} $t_P$', vmin=vmin, vmax=vmax )
y, x = velx_in_xy_plane_final.coordinates_meshgrid()
ax[1].quiver(x[mask], y[mask], velx_in_xy_plane_final.data[mask], vely_in_xy_plane_final.data[mask], 
             label="t=%.3g" % (velx_in_xy_plane_final.time / t_P), 
             angles='xy',scale=6,width=0.005,color='black')
ax[1].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[1].set_xlim(xmin,xmax)
ax[1].set_ylim(ymin,ymax)
ax[1].set_xlabel("x [M]")
ax[1].set_ylabel("y [M]")
ax[1].set_title("3-velocity in xy plane")
ax[1].legend(loc='upper right')
fig.colorbar(im, ax=ax[1])

plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_vel_xy' + '_quiver_plot.pdf'), bbox_inches="tight")
plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_vel_xy' + '_quiver_plot.png'), bbox_inches="tight")

plotting velocity xy


# get quiver plot using unstructured interpolation of all data (slow)
print("plotting velocity xz")
fig, ax = plt.subplots(1, 2, figsize=(20.0, 8.0))
rho_xz_2d = gf.xz["rho_b"][iteration_init].get_level(max_reflvl)
velx_in_xz_plane_init = gf.xz["vx"][iteration_init].get_level(max_reflvl)
velz_in_xz_plane_init = gf.xz["vy"][iteration_init].get_level(max_reflvl)
im = ax[0].pcolormesh(*rho_xz_2d.coordinates_meshgrid()[::-1], 
                      (rho_xz_2d.data/RHO0), cmap='jet', 
                      label=fr'{velx_in_xz_plane_init.time /t_P:.3f} $t_P$', vmin=vmin, vmax=vmax )
y, x = velx_in_xz_plane_init.coordinates_meshgrid()
# create a true false array as a mask to sample grid
mask = np.zeros_like(x, dtype=bool)
mask[::10, ::10] = True
ax[0].quiver(x[mask], y[mask], velx_in_xz_plane_init.data[mask], velz_in_xz_plane_init.data[mask], 
             label="t=%.3g" % (velx_in_xz_plane_init.time / t_P), 
             angles='xy',scale=6,width=0.005,color='black')
ax[0].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[0].set_xlim(xmin,xmax)
ax[0].set_ylim(ymin,ymax)
ax[0].set_xlabel("x [M]")
ax[0].set_ylabel("z [M]")
fig.colorbar(im, ax=ax[0])
ax[0].set_title("3-velocity in xz plane")
ax[0].legend(loc='upper right')


rho_xz_2d = gf.xz["rho_b"][iteration_final].get_level(max_reflvl)
velx_in_xz_plane_final = gf.xz["vx"][iteration_final].get_level(max_reflvl)
velz_in_xz_plane_final = gf.xz["vy"][iteration_final].get_level(max_reflvl)
im = ax[1].pcolormesh(*rho_xz_2d.coordinates_meshgrid()[::-1], 
                      (rho_xz_2d.data/RHO0), cmap='jet', 
                      label=fr'{rho_xz_2d.time /t_P:.3f} $t_P$', vmin=vmin, vmax=vmax )
y, x = velx_in_xz_plane_final.coordinates_meshgrid()
ax[1].quiver(x[mask], y[mask], velx_in_xz_plane_final.data[mask], velz_in_xz_plane_final.data[mask], 
             label="t=%.3g" % (velx_in_xz_plane_final.time / t_P), 
             angles='xy',scale=6,width=0.005,color='black')
ax[1].set_aspect('equal')
xmin = ymin = -1.5 * Req
xmax = ymax = 1.5 * Req
ax[1].set_xlim(xmin,xmax)
ax[1].set_ylim(ymin,ymax)
ax[1].set_xlabel("x [M]")
ax[1].set_ylabel("z [M]")
ax[1].set_title("3-velocity in xz plane")
ax[1].legend(loc='upper right')
fig.colorbar(im, ax=ax[1])

plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_vel_xz' + '_quiver_plot.pdf'), bbox_inches="tight")
plt.savefig(Path(plot_output_dir, "2D", sim+ f'_output_{num_dir}' +  '_vel_xz' + '_quiver_plot.png'), bbox_inches="tight")

plotting velocity xz


# plot velocity on z axis
# ILGRMHD
plt.figure(figsize=(10, 8.))
for it, time in zip(iterations[-300::100], times[-300::100]):
    vel_z_on_xaxis = np.vstack([gf.x["vz"][it].get_level(max_reflvl).coordinates_meshgrid()[0], 
                                gf.x["vz"][it].get_level(max_reflvl).data]).T
    rest_vel_z = vel_z_on_xaxis[:, 1]
    # omega_y = - rest_vel_z / vel_z_on_xaxis[:, 0]
    plt.plot(vel_z_on_xaxis[:, 0] / Req, rest_vel_z, label=fr't={time/t_P:.3g} $t_P$')
plt.xlim([-1.2, 1.2])
# plt.ylim([-0.2, 0.2])
plt.xlabel(r'$x / R_x$')
plt.ylabel(r'$v_z$ on x-axis')
plt.legend(loc='upper right')
plt.show()
# plt.savefig(Path(plot_output_dir, "1D", sim+ f'_output_{num_dir}' +  '_omega_y_on_xaxis' + '_plot.pdf'), bbox_inches='tight')


# sample_velxi = idw.tree(velx_in_xz_plane_init[:, :2], velx_in_xz_plane_init[:, 2])
# sample_velzi = idw.tree(velz_in_xz_plane_init[:, :2], velz_in_xz_plane_init[:, 2])

# sample_velxf = idw.tree(velx_in_xz_plane_final[:, :2], velx_in_xz_plane_final[:, 2])
# sample_velzf = idw.tree(velz_in_xz_plane_final[:, :2], velz_in_xz_plane_final[:, 2])

# fig = plt.figure(figsize=(10.0, 10.0))
# ax1 = fig.add_subplot(1,1,1)
# ax1.grid(color='black', linestyle=':', linewidth=0.5)
# levels = [0.4, 0.5, 0.7, 0.9, 0.95]
# levcol = [ 'green', 'cyan', 'brown', 'orange', 'magenta' ]
# contour1 = ax1.contour(x20i/Req, z20i/Req, rho20i/RHO0, levels=levels, colors=levcol, linestyles='dotted')
# contour2 = ax1.contour(x2fi/Req, z2fi/Req, rho2fi/RHO0, levels=levels, colors=levcol, linestyles='-')

# for cont in contour1.allsegs:
#     coord = cont[0]
#     velxi = sample_velxi(coord)
#     velzi = sample_velzi(coord)
#     ax1.quiver(coord[:, 0], coord[:, 1], velxi, velzi, angles='xy',scale=3,width=0.005,color='red')

# for cont in contour2.allsegs:
#     coord = cont[0]
#     velxf = sample_velxf(coord)
#     velzf = sample_velzf(coord)
#     ax1.quiver(coord[:, 0], coord[:, 1], velxf, velzf, angles='xy',scale=3,width=0.005,color='blue')

# ax1.set_xlim([-1, 1])
# ax1.set_ylim([-1, 1])
# ax1.set_aspect("equal")


simdir.save(os.path.join(work_dir, "simdir.pickle"))

Scalar diagnostics¶

Central density¶

Rest mass and ADM mass¶

Stellar modes¶

Angular momentum $J_\phi$¶

GW $\psi_4$¶

QuadMoments¶

1D output¶

density change along x-axis¶

$\Omega_x$ Rotation curve on x-axis¶

Lapse on x axis $\alpha_x$¶

2D output¶

Density in xy plane $\rho_{xy}$¶

Density in xz plane $\rho_{xz}$¶

Velocity on xy plane $v_{xy}$¶

Velocity on xz plane $v_{xz}$¶