In [ ]:
%matplotlib inline
from __future__ import print_function, division

import os, sys, glob
import collections 
import numpy as np
import matplotlib.pyplot as plt

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

import ipywidgets as widgets
from ipywidgets import RadioButtons, fixed
from ipywidgets import interact, interactive, interact_manual

from IPython.display import display
import traitlets

## 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 import data_dir_default
from clustered_SNe.analysis.constants import m_proton, pc, yr, M_solar, \
                                   metallicity_solar
from clustered_SNe.analysis.parse import Overview, RunSummary, \
                                         extract_masses_momenta_raw

from clustered_SNe.analysis.visualize_helpers import plot_momentum

from clustered_SNe.analysis.database_helpers import session, \
                                                    Simulation, \
                                                    Simulation_Inputs, \
                                                    Simulation_Status, \
                                                    extract_masses_momenta

In [ ]:
from IPython.core.display import display, HTML


if widgets.version_info[0]==5:
    display(HTML("""
    <style>
    .widget-hbox .widget-label { min-width: 225px; }
    .widget-text { width: 600px; }
    .widget-button { width: 300px; }
    </style>"""))

In [ ]:
data_dir = data_dir_default

In [ ]:
metallicities_given_density = {}

for simulation in session.query(Simulation):
    if not simulation.background_density in metallicities_given_density:
        metallicities_given_density[simulation.background_density] = set()
    metallicities_given_density[simulation.background_density].add(simulation.metallicity)

densities = np.array(list(metallicities_given_density.keys()))
densities.sort()
log_densities = np.log10(densities / m_proton)
density_label_map = collections.OrderedDict() # so labels will be in order
for density in densities:
    log_density = np.log10(density / m_proton)
    density_label_map[str(round(log_density, 4))] = density

slice = {}
def extract_masses_momenta_wrapper(density):
    metallicities = np.array(list(metallicities_given_density[density]))
    metallicities.sort()
    metallicity_label_map = collections.OrderedDict()
    for metallicity in metallicities:
        log_metallicity = np.log10(metallicity / metallicity_solar)
        metallicity_label_map[str(round(log_metallicity, 4))] = metallicity
    extract_masses_momenta_wrapper.previous_widget.widget.close()
    v = interact_manual(extract_masses_momenta,
                        density = fixed(density),
                        metallicity = widgets.Dropdown(
                            options = metallicity_label_map,
                            description = 'Log10 (Z / Z_solar):',
                            value=list(metallicity_label_map.values())
                                    [np.abs([float(k)
                                             for k in metallicity_label_map]) \
                                     .argmin()
                                    ]
                            ),
                       )
    extract_masses_momenta_wrapper.previous_widget = v
    def update_slice(widget):
        slice.clear()
        slice["masses"] = v.widget.result[0]
        slice["momenta"] = v.widget.result[1]
        slice["ids"] = v.widget.result[2]
    v.widget.children[-1].on_click(update_slice)
    return v
tmp = widgets.Box()
tmp.widget = widgets.Box()
extract_masses_momenta_wrapper.previous_widget = tmp

selected_run_widget = widgets.Text()
selected_run_widget.visible = False

In [ ]:
w = interact(extract_masses_momenta_wrapper,
             density= widgets.Dropdown(options=density_label_map,
                             description='Log10 (density / m_p):',
                             value=list(density_label_map.values())
                                     [np.abs([float(k) 
                                              for k in density_label_map]
                                            ) \
                                      .argmin()
                                     ]
                                     )
             )

In [ ]:
from bokeh import mpl
from bokeh.plotting import output_file, show, figure
from bokeh.models import HoverTool, TapTool, CustomJS
from bokeh.io import output_notebook
output_notebook()

from bokeh.models import ColumnDataSource
N_SNe = np.array([session.query(Simulation).get(id).num_SNe for id in slice["ids"]])
masses=slice["masses"] / M_solar
momenta=slice["momenta"] / (N_SNe * 100 * M_solar) / (1000 * 100) 
momenta[~np.isfinite(momenta)] = 1 # log scale can't handle 0
statuses = [session.query(Simulation_Status).get(id).status for id in slice["ids"]]
colors = []
for status in statuses:
    if status=="Complete":
        colors.append("green")
    elif (status=="Ready") or (status=="Running"):
        colors.append("orange")
    else:
        colors.append("red")
source = ColumnDataSource(data=dict(masses=masses,
                                    momenta=momenta,
                                    ids=slice["ids"], 
                                    colors=colors,
                                    statuses=statuses,
                                    N_SNe=N_SNe))

p = figure(title="momentum scaling", 
           x_axis_type="log", 
           y_axis_type="log",
           x_range= [1, N_SNe.max()*2],
           y_range= [10**3, momenta.max()*2],
           tools="reset,hover,tap, wheel_zoom")
p.xaxis.axis_label = "N_SNe"
p.yaxis.axis_label = "Momentum / (N_SNe * 100 * M_sol) [km / s]"
p.scatter("N_SNe","momenta", 
          source=source,
          size=10,
          legend="data", 
          color="colors")


hover = p.select(dict(type=HoverTool))
hover.tooltips = {
    "Mass": "@masses",
    "Momentum / N_SNe / M_sol": "@momenta",
    "ids": "@ids",
    "status": "@statuses",
    "N_SNe": "@N_SNe",
}

source.callback = CustomJS(code="""
        var inds = cb_obj.get('selected')['1d'].indices;
        var d = cb_obj.get('data');
        for (i=0; i<inds.length; i++) {
            console.log("dataset id: " + d['ids'][inds[i]])
            // this is an unsecure hack, but it's safe enough if you are running your own code
            msg = "selected_run_widget.value = '" + d['ids'][inds[i]] + "'"
            console.log(msg);
            IPython.notebook.kernel.execute(msg);
            
            // copy id to clipboard
            var id_textarea = document.createElement("textarea")
            id_textarea.value = d['ids'][inds[i]]
            document.body.appendChild(id_textarea)
            id_textarea.select()
            document.execCommand('copy')
            document.body.removeChild(id_textarea)

    }
    """)


show(p)

In [ ]:
run_summary = RunSummary()
def show_momentum_evolution(data_dir, id):
    run_summary.replace_with(RunSummary(data_dir=data_dir, id=id))
    show_momentum_evolution.previous_widget.close()
    u = interactive(plot_momentum,
             run_summary=fixed(run_summary),
             x_axis = RadioButtons(options=["time", "checkpoints"]),
             clear_previous=fixed(True),
             distplot=fixed(True), 
             y_axis_scaling=fixed("mass"),
             label=fixed(""),
             plot_times_in_Myr=fixed(True),
             with_momentum_axvline=fixed(True))
    display(u)
    show_momentum_evolution.previous_widget = u

show_momentum_evolution.previous_widget = widgets.Box()


w = widgets.interact_manual(show_momentum_evolution,
                     data_dir = fixed(data_dir),
                     id=widgets.Text())
l = traitlets.link((selected_run_widget, "value"), (w.widget.children[0], "value"))

In [ ]:
def set_status(id, status="Error"):
    assert(status in Simulation_Status.possible_statuses)
    session.query(Simulation_Status).get(id).status = status
    print("Id '"+ id + "' set to: ", status)
w_e = widgets.interact_manual(set_status, 
                              id = widgets.Text(), 
                              status = RadioButtons(options=list(Simulation_Status.possible_statuses)))
traitlets.link((selected_run_widget, "value"), (w_e.widget.children[0], "value"))

In [ ]:


In [ ]: