In [ ]:
%matplotlib inline

from __future__ import print_function, division

import os, sys
import glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import ipywidgets as widgets
from ipywidgets import interact, interactive, interact_manual
from ipywidgets import FloatSlider, IntSlider, Dropdown
from ipywidgets import Checkbox, RadioButtons, fixed 

from IPython.display import display, Math, clear_output

from astropy.convolution import convolve, Gaussian1DKernel

import seaborn as sns
sns.set(context="poster", font_scale=1.5)

## import from local files

## Boilerplate path hack to give access to full clustered_SNe package
import sys, os
if __package__ is None:
    if os.pardir not in sys.path[0]:
        file_dir = os.getcwd()
        sys.path.insert(0, os.path.join(file_dir, 
                                        os.pardir, 
                                        os.pardir))

from clustered_SNe.analysis.constants import m_proton, pc, yr, M_solar, \
                                   gamma, E_0, metallicity_solar
    
from clustered_SNe.analysis.sedov.dimensionalize_sedov import dimensionalized_sedov
from clustered_SNe.analysis.sedov.closed_form_sedov import SedovSolution
from clustered_SNe.analysis.parse import RunSummary, Overview, RunSummary, cols


general_string_format = ".2e"

from clustered_SNe.analysis.visualize_helpers import plotter, single_run, \
                                                     conduction_comparisons, \
                                                     parameter_study_wrapper, \
                                                     SNe_distplot, \
                                                     plot_zones, \
                                                     plot_shock_location, \
                                                     plot_energy, \
                                                     plot_momentum, \
                                                     plot_luminosity



        
# Holds the most recent run, if more analysis is desired
run_summary = RunSummary()

In [ ]:
# these ids probably have to be changed... Probably better to write a script to deal with ids?

ids_1e2_M_solar = ["2177f", "0590d", "639be"]
ids_1e3_M_solar = ["7247a", "ff5b5", "de6b1"]
ids_1e4_M_solar = ["30faf", "cccbd", "c390f"]
ids_1e5_M_solar = ["2bf32", "46aed", "fa9c8"]

mass_loss_label_map = {"uniform"   : "constant",
                       "disappear" : "disappear",
                       "none"      : r"$\mathtt{AGORA}$"}

data_dir = "../saved_runs/mass_loss_comparisons/"

def mass_loss_comparisons_last_checkpoint(ids, data_dir="../src"):
    
    if not os.path.exists(data_dir):
        raise FileNotFoundError("No directory found named: "+ data_dir)
        
    plt.figure()
    mass_loss_labels = []
    
    for id in ids:
    
        run_summary = RunSummary(data_dir=data_dir, id=id)
        sedov_solution = SedovSolution(E_0,
                                       run_summary.overview.background_density, 
                                       run_summary.overview.metallicity)

        #### PASS TO PLOTTER ####
        num_checkpoints = len(run_summary.filenames)

#         log_R_max = round(np.log10(run_summary.df["Radius"].max()), 2)
#         log_R_min = max(log_R_max-4, 
#                         round(np.log10(run_summary.df["Radius"].min()), 2)+1)


        R_min = run_summary.df["Radius"].min()
        R_min = 0
        R_max = run_summary.df["Radius"].max()
        
        plt.title("{0:.0e}".format(run_summary.overview.cluster_mass / M_solar) + r" $M_\odot$ cluster" )

        plotter(run_summary,
                sedov_solution,
                x_axis_variable  = "Radius",
                y_axis_variable  = "Density",
                with_Sedov       = False,
                highlight_timestep_limiting_cell = False,
                xlim             = (R_min, R_max), 
                checkpoint_index = num_checkpoints-1,
                show_at_end = False)
        mass_loss_labels.append(mass_loss_label_map[run_summary.overview.mass_loss])


    plt.legend(mass_loss_labels, loc="best", title="mass loss")
    plt.setp(plt.gca().get_legend().get_title(),
             fontsize=plt.rcParams["legend.fontsize"])

    plot_filename = "plots/last_checkpoint_{0:.0e}_Msun".format(run_summary.overview.cluster_mass / M_solar)
    plt.savefig(plot_filename + ".eps")
    plt.savefig(plot_filename + ".pdf")
    plt.savefig(plot_filename + ".png")
    clear_output()
#     plt.close()

def mass_loss_comparisons_before_second_blast(ids, data_dir="../src"):
    
    if not os.path.exists(data_dir):
        raise FileNotFoundError("No directory found named: "+ data_dir)
        
    plt.figure()
    mass_loss_labels = []
    
    for id in ids:
    
        run_summary = RunSummary(data_dir=data_dir, id=id)
        sedov_solution = SedovSolution(E_0,
                                       run_summary.overview.background_density, 
                                       run_summary.overview.metallicity)

        #### PASS TO PLOTTER ####
        num_checkpoints = len(run_summary.filenames)
        plot_checkpoint = np.argmin(run_summary.times < run_summary.overview.SNe_times[1]) - 1

#         log_R_max = round(np.log10(run_summary.df["Radius"].max()), 2)
#         log_R_min = max(log_R_max-4, 
#                         round(np.log10(run_summary.df["Radius"].min()), 2)+1)

        R_min = run_summary.df["Radius"].min()
        R_min = 0
        R_max = run_summary.df["Radius"].max()

        plt.title("{0:.0e}".format(run_summary.overview.cluster_mass / M_solar) + r" $M_\odot$ cluster" )

        plotter(run_summary,
                sedov_solution,
                x_axis_variable  = "Radius",
                y_axis_variable  = "Density",
                with_Sedov       = False,
                highlight_timestep_limiting_cell = False,
                xlim             = (R_min, R_max), 
                checkpoint_index = plot_checkpoint,
                show_at_end = False)
        mass_loss_labels.append(mass_loss_label_map[run_summary.overview.mass_loss])

    
    plt.legend(mass_loss_labels, loc="best", title="mass loss")
    plt.setp(plt.gca().get_legend().get_title(),
             fontsize=plt.rcParams["legend.fontsize"])

    plot_filename = "plots/before_second_blast_{0:.0e}_Msun".format(run_summary.overview.cluster_mass / M_solar)
    plt.savefig(plot_filename + ".eps")
    plt.savefig(plot_filename + ".pdf")
    plt.savefig(plot_filename + ".png")
    clear_output()
#     plt.close()

def mass_loss_comparisons_momentum(ids, data_dir="../src"):
    
    if not os.path.exists(data_dir):
        raise FileNotFoundError("No directory found named: "+ data_dir)
    
    plt.figure()
    mass_loss_labels = []

    for id in ids:
    
        run_summary = RunSummary(data_dir=data_dir, id=id)
        sedov_solution = SedovSolution(E_0,
                                       run_summary.overview.background_density, 
                                       run_summary.overview.metallicity)

        #### PASS TO PLOTTER ####
        num_checkpoints = len(run_summary.filenames)

        log_R_max = round(np.log10(run_summary.df["Radius"].max()), 2)
        log_R_min = max(log_R_max-4, 
                        round(np.log10(run_summary.df["Radius"].min()), 2)+1)


        plt.title("{0:.0e}".format(run_summary.overview.cluster_mass / M_solar) + r" $M_\odot$ cluster" )
    
        plot_momentum(run_summary, "time", clear_previous=False, distplot=False)
        mass_loss_labels.append(mass_loss_label_map[run_summary.overview.mass_loss])


    plt.legend(mass_loss_labels, loc="best", title="mass loss")
    plt.setp(plt.gca().get_legend().get_title(),
             fontsize=plt.rcParams["legend.fontsize"])
    
    SNe_distplot(run_summary, "time")

    plot_filename = "plots/momentum_{0:.0e}_Msun".format(run_summary.overview.cluster_mass / M_solar)
    plt.savefig(plot_filename + ".eps")
    plt.savefig(plot_filename + ".pdf")
    plt.savefig(plot_filename + ".png")
    clear_output()
#     plt.close()

    
    
def mass_loss_comparisons_R_shock(ids, data_dir="../src"):
    
    if not os.path.exists(data_dir):
        raise FileNotFoundError("No directory found named: "+ data_dir)
    
    plt.figure()
    mass_loss_labels = []
    
    for id in ids:
    
        run_summary = RunSummary(data_dir=data_dir, id=id)

        plt.title("{0:.0e}".format(run_summary.overview.cluster_mass / M_solar) + r" $M_\odot$ cluster" )
        shock_location_plot(run_summary, clear_previous=False, distplot=False)
        mass_loss_labels.append(mass_loss_label_map[run_summary.overview.mass_loss])

    plt.legend(mass_loss_labels, loc="best", title="mass loss")
    plt.setp(plt.gca().get_legend().get_title(),
             fontsize=plt.rcParams["legend.fontsize"])
    
    SNe_distplot(run_summary, "time")
    
    plot_filename = "plots/R_shock_{0:.0e}_Msun".format(run_summary.overview.cluster_mass / M_solar)
    plt.savefig(plot_filename + ".eps")
    plt.savefig(plot_filename + ".pdf")
    plt.savefig(plot_filename + ".png")
    clear_output()
    
    
    
plt.close()  
mass_loss_comparisons_before_second_blast(ids=ids_1e2_M_solar,
                                          data_dir=data_dir)

    
mass_loss_comparisons_last_checkpoint(ids_1e2_M_solar,
                                      data_dir=data_dir)
mass_loss_comparisons_last_checkpoint(ids_1e3_M_solar,
                                      data_dir=data_dir)
mass_loss_comparisons_last_checkpoint(ids_1e4_M_solar,
                                      data_dir=data_dir)
mass_loss_comparisons_last_checkpoint(ids_1e5_M_solar,
                                      data_dir=data_dir)

mass_loss_comparisons_momentum(ids_1e2_M_solar,
                               data_dir=data_dir)
mass_loss_comparisons_momentum(ids_1e3_M_solar,
                               data_dir=data_dir)
mass_loss_comparisons_momentum(ids_1e4_M_solar,
                               data_dir=data_dir)
mass_loss_comparisons_momentum(ids_1e5_M_solar,
                               data_dir=data_dir)

mass_loss_comparisons_R_shock(ids_1e2_M_solar,
                              data_dir=data_dir)
mass_loss_comparisons_R_shock(ids_1e3_M_solar,
                              data_dir=data_dir)
mass_loss_comparisons_R_shock(ids_1e4_M_solar,
                              data_dir=data_dir)
mass_loss_comparisons_R_shock(ids_1e5_M_solar,
                              data_dir=data_dir)

In [ ]: