# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# %matplotlib inline

import cosmo_metric_utils as cmu
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

import matplotlib.colors as colors
import matplotlib.cm as cmx
from mpl_toolkits.axes_grid1 import make_axes_locatable

import os
import glob

# +
#dist_loc_base = '/media/RESSPECT/data/PLAsTiCC/for_metrics/wfd/distances/omprior_0.01_flat/emille_samples/' #mu_photoIa_plasticc*'
#file_extension = 'WFD'

dist_loc_base = '/media/RESSPECT/data/PLAsTiCC/for_metrics/ddf/distances/omprior_0.01_flat/emille_samples/' #mu_photoIa_plasticc*'
file_extension = 'DDF'


# -

from collections import OrderedDict
remap_dict = OrderedDict({
                          'perfect3000': 'Perfect', 
                          'fiducial3000': 'Fiducial', 
                          'random3000fail2998': 'Random',
                          'random3000': 'Random',
                          'all_objs_survived_SALT2_DDF' : 'All SALT',
                          'all_objs_survived_SALT2_WFD': 'All SALT',
                          '50SNIa50SNII': 'SN-II 50',
                          '68SNIa32SNII': 'SN-II 32',
                          '72SNIa28SNII': 'SN-II 28',
                          '75SNIa25SNII': 'SN-II 25', 
                          '90SNIa10SNII': 'SN-II 10',
                          '95SNIa5SNII': 'SN-II 5',
                          '98SNIa2SNII': 'SN-II 2',
                          '99SNIa1SNII': 'SN-II 1',
                          '50SNIa50SNIbc': 'SN-Ibc 50',
                          '68SNIa32SNIbc': 'SN-Ibc 32',
                          '75SNIa25SNIbc': 'SN-Ibc 25',
                          '83SNIa17SNIbc': 'SN-Ibc 17',
                          '90SNIa10SNIbc': 'SN-Ibc 10',
                          '95SNIa5SNIbc': 'SN-Ibc 5',
                          '98SNIa2SNIbc': 'SN-Ibc 2',
                          '99SNIa1SNIbc': 'SN-Ibc 1',
                          '50SNIa50SNIax': 'SN-Iax 50',
                          '68SNIa32SNIax': 'SN-Iax 32',
                          '75SNIa25SNIax': 'SN-Iax 25',
                          '86SNIa14SNIax': 'SN-Iax 14',
                          '90SNIa10SNIax': 'SN-Iax 10',
                          '94SNIa6SNIax': 'SN-Iax 6',
                          '95SNIa5SNIax': 'SN-Iax 5',
                          '97SNIa3SNIax': 'SN-Iax 3',
                          '98SNIa2SNIax': 'SN-Iax 2',
                          '99SNIa1SNIax': 'SN-Iax 1',
                          '71SNIa29SNIa-91bg': 'SN-Ia-91bg 29',
                          '75SNIa25SNIa-91bg': 'SN-Ia-91bg 25',
                          '90SNIa10SNIa-91bg': 'SN-Ia-91bg 10',
                          '95SNIa5SNIa-91bg': 'SN-Ia-91bg 5',
                          '98SNIa2SNIa-91bg': 'SN-Ia-91bg 2',
                          '99SNIa1SNIa-91bg': 'SN-Ia-91bg 1',
                          '99.8SNIa0.2SNIa-91bg': 'SN-Ia-91bg 0.2',
                          '57SNIa43AGN': 'AGN 43',
                          '75SNIa25AGN': 'AGN 25',
                          '90SNIa10AGN': 'AGN 10',
                          '94SNIa6AGN': 'AGN 6',
                          '95SNIa5AGN': 'AGN 5',
                          '98SNIa2AGN': 'AGN 2',
                          '99SNIa1AGN': 'AGN 1',
                          '99.9SNIa0.1AGN': 'AGN 0.1',
                          '83SNIa17SLSN-I': 'SNLS-I 17',
                          '90SNIa10SLSN-I': 'SNLS-I 10',
                          '95SNIa5SLSN-I': 'SNLS-I 5',
                          '98SNIa2SLSN-I': 'SNLS-I 2',
                          '99SNIa1SLSN-I': 'SNLS-I 1',
                          '99.9SNIa0.1SLSN': 'SNLS-I 0.1',
                          '95SNIa5TDE': 'TDE 5',
                          '98SNIa2TDE': 'TDE 2',
                          '99SNIa1TDE': 'TDE 1',
                          '99.6SNIa0.4TDE': 'TDE 0.4',
                          '99.1SNIa0.9CART': 'CART 0.9',
                          '99.7SNIa0.3CART': 'CART 0.3'
              })

all_shapes = {'SNIa-91bg': 'o',
              'SNIax': 's',
              'SNII': 'd',
              'SNIbc': 'X',
              'SLSN-I': 'v',
              'AGN': '^',
              'TDE': '<',
              'KN': '>',
              'CART': 'v'}

# Mapping the percent contaminated to the colormap.
## size corresponds to remap_dict
color_nums = np.array([1, 1, 1, 1, 1, 1,                    # Special
                       50, 32, 28, 25, 10, 5, 2, 1,         # II
                       50, 32, 25, 17, 10, 5, 2, 1,         # Ibc
                       50, 32, 25, 14, 10, 6, 5, 3, 2, 1,   # Iax
                       29, 25, 10, 5, 2, 1, 1,              # 91bg
                       43, 25, 10, 6, 5, 2, 1, 1,           # AGN
                       17, 10, 5, 2, 1, 1,                  # SNLS
                       5, 2, 1, 1,                          # TDE
                       1, 1,                                # CART
                      ])  

# Color map
rainbow = cm = plt.get_cmap('plasma_r')
cNorm  = colors.LogNorm(vmin=1, vmax=52) #colors.Normalize(vmin=0, vmax=50)
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=rainbow)
color_map = scalarMap.to_rgba(np.arange(1, 52))

# +
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(5, 10), sharey=True)
tick_lbls = []

ax1.axvline(-1, color='c', ls='--')
df = pd.read_csv(dist_loc_base + 'stan_input_salt2mu_lowz_withbias_perfect3000.csv')
sig_perf = cmu.fisher_results(df['z'].values, df['muerr'].values)[0]

ax1.axvspan(-1 - sig_perf[1], 
            -1 + sig_perf[1], alpha=0.1, color='grey')
ax2.axvline(0, color='k')


file_base = dist_loc_base + 'stan_input_salt2mu_lowz_withbias_'
i = 0
i_list = []
for j, (a, c) in enumerate(zip(remap_dict, color_nums)):
    class_ = str.split(remap_dict[a])[0]

    if '91bg' in class_:
        class_ = 'SNIa-91bg'
    else:
        class_ = class_.replace('-', '')
        
    #mfc='none'
    if 'DDF' in file_extension:
        if 'fiducial' in a:
            mfc = 'tab:blue'
        elif 'random' in a:
            mfc = 'tab:red'
        elif 'perfect' in a:
            mfc = 'k'
        else:
            mfc = color_map[c]
    if 'WFD' in file_extension:
        mfc = "none"
    try:
        file = glob.glob(file_base + a + '.csv')
        df = pd.read_csv(str(file[0]))
        sig = cmu.fisher_results(df['z'].values, df['muerr'].values)[0]
        if 'perfect' in a:
            ax1.plot(-1, -i, ms=10, color='k', marker='*', mfc=mfc)
            ax1.plot([-1 - sig[1], -1 + sig[1]], [-i, -i], "|-", ms=10, color='k', mfc=mfc)
        elif 'random' in a:
            ax1.plot(-1, -i, ms=10, color='tab:red', marker='*', mfc=mfc)
            ax1.plot([-1 - sig[1], -1 + sig[1]], [-i, -i], "|-", ms=10, color='tab:red')
            ax2.plot((sig[1]-sig_perf[1])/sig_perf[1], -i, 'o', color='tab:red', marker='*', ms=10, mfc=mfc)
        elif 'fiducial' in a:
            ax1.plot(-1, -i, ms=10, color='tab:blue', marker='*', mfc=mfc)
            ax1.plot([-1 - sig[1], -1 + sig[1]], [-i, -i], "|-", ms=10, color='tab:blue')
            ax2.plot((sig[1]-sig_perf[1])/sig_perf[1], -i, color='tab:blue', marker='*', ms=10, mfc=mfc)
        else:
            ax1.plot(-1, -i, marker=all_shapes[class_], ms=10, color=color_map[c], mfc=mfc)
            ax1.plot([-1 - sig[1], -1 + sig[1]], [-i, -i], "|-", ms=10, color=color_map[c])
            ax2.plot((sig[1]-sig_perf[1])/sig_perf[1], -i, color=color_map[c], marker=all_shapes[class_], ms=10, mfc=mfc)
        tick_lbls.append(remap_dict[a])
        i_list.append(-i)
        i +=2
        if 'random' in a or '99SNIa1' in a:
            if 'DDF' in file_extension:
                if 'AGN' in a or '91bg' in a or 'CART' in a:
                    continue
            else:
                i_list.append(i)
                i += 1.6
                tick_lbls.append('')
        elif 'SNIa0' in a and 'CART' not in a:
            i_list.append(-i)
            i += 1.6
            tick_lbls.append('')
    except:
        print("Missing: ", a)

tick_locs = tick_locs = i_list[::-1]#np.arange(-len(tick_lbls)+1, 1)
ax1.set_yticks(tick_locs)
ax1.set_yticklabels(tick_lbls[::-1], fontsize=13)

ax1.set_ylim(i_list[-1]-0.7, i_list[0]+0.7)
#(-len(tick_lbls)+0.5, 0.5)
#ax1.set_xscale('log')
ax1.set_xlabel('Fisher Matrix', fontsize=13)
ax2.set_xlabel('Fractional difference', fontsize=13)

yticks = ax1.yaxis.get_major_ticks()
yticks2 = ax2.yaxis.get_major_ticks()

#ticks = [-4, -11, -15, -21]
#for t in ticks:
#    yticks[t].set_visible(False)
#    yticks2[t].set_visible(False)

#plt.savefig('fisher_matrix_' + file_extension + '_20210603.pdf', bbox_inches='tight')
# -



StarcoderdataJupyterScriptsDedupFiltered