<jupyter_start><jupyter_text>Load trained models<jupyter_code>energy_pipeline = comp.load_trained_model('linearregression_energy_{}'.format(config))
pipeline_str = 'xgboost_comp_{}_{}-groups'.format(config, num_groups)
comp_pipeline = comp.load_trained_model(pipeline_str)
# Load DataFrame with saved prior distributions
df_file = os.path.join(unfolding_dir,
                       'unfolding-df_{}-groups.hdf'.format(num_groups))
df_priors = pd.read_hdf(df_file)
df_priors.head()
# Load fitted efficiencies and calculate effective areas
eff_path = os.path.join(comp.paths.comp_data_dir, config, 'efficiencies',
                        'efficiency_fit_num_groups_{}.hdf'.format(num_groups))
df_eff = pd.read_hdf(eff_path)
df_eff.head()
# Get simulation thrown areas for each energy bin
thrown_radii = comp.simfunctions.get_sim_thrown_radius(energybins.log_energy_midpoints)
thrown_area = np.pi * thrown_radii.max()**2
thrown_area
efficiencies, efficiencies_err = comp.get_detector_efficiencies(config=config,
                                                                num_groups=num_groups,
                                                                sigmoid='flat',
                                                                pyunfold_format=True)<jupyter_output><empty_output><jupyter_text>Load and process data<jupyter_code>print('Loading data into memory...')
df_data = comp.load_data(config=config,
                         processed=True,
                         columns=feature_list,
                         n_jobs=15,
                         verbose=True)
X_data = df_data[feature_list].values
print('Running composition classifications...')
df_data['pred_comp_target'] = comp_pipeline.predict(X_data)
print('Running energy regressions...')
df_data['reco_log_energy'] = energy_pipeline.predict(X_data)
counts_observed = {}
counts_observed_err = {}

for comp_idx, composition in enumerate(comp_list):
    print(composition)
    # Filter out events that don't pass composition & energy mask
    pred_comp_mask = df_data['pred_comp_target'] == comp_idx
    mask = pred_comp_mask
    
    energies = df_data.loc[mask, 'reco_log_energy'].values
    comp_counts, _ = np.histogram(energies, bins=energybins.log_energy_bins)
    counts_observed[composition] = comp_counts
    counts_observed_err[composition] = np.sqrt(comp_counts)

counts_observed_err['total'] = np.sqrt(np.sum(counts_observed_err[composition]**2 for composition in comp_list))
# Calculate total counts
counts_observed['total'] = np.sum(counts_observed[composition] for composition in comp_list)

# Format observed counts, detection efficiencies, and priors for PyUnfold use
counts_pyunfold = np.empty(num_groups * len(energybins.energy_midpoints))
counts_err_pyunfold = np.empty(num_groups * len(energybins.energy_midpoints))
for idx, composition in enumerate(comp_list):
    counts_pyunfold[idx::num_groups] = counts_observed[composition]
    counts_err_pyunfold[idx::num_groups] = counts_observed_err[composition]
counts_observed
# # Effective area
# eff_area, eff_area_err = {}, {}
# for composition in comp_list + ['total']:
#     eff_area[composition] = df_eff['eff_median_{}'.format(composition)].values * thrown_area
#     eff_area_err[composition] = df_eff['eff_err_low_{}'.format(composition)].values * thrown_area

# Solid angle
theta_max = 40 if config == 'IC79.2010' else 65
solid_angle = np.pi*np.sin(np.deg2rad(theta_max))**2

# Livetime
livetime, livetime_err = comp.get_detector_livetime(config=config)
def counts_to_flux(counts, counts_err=None, composition=None):
    return comp.get_flux(counts, counts_err,
                         energybins=energybins.energy_bins,
                         eff_area=thrown_area, 
                         livetime=livetime,
                         livetime_err=livetime_err,
                         solid_angle=solid_angle,
                         scalingindex=2.7)
comp_list + ['total']
UnfoldedFlux = namedtuple('UnfoldedFlux', ['flux', 'flux_sys_err', 'flux_stat_err'])
def get_unfolded_fluxes(df_file, return_sys_err=False, return_stat_err=False):
    print('df_file = {}'.format(df_file))
    # Load simulation and train composition classifier
    df_sim_train, df_sim_test = comp.load_sim(df_file=df_file,
                                              config=config,
                                              energy_reco=False,
                                              log_energy_min=None,
                                              log_energy_max=None,
                                              test_size=0.5,
                                              n_jobs=10,
                                              verbose=True)
    
    print('Running energy reconstruction...')
    for df in [df_sim_train, df_sim_test]:
        X = df_sim_train[feature_list].values
        # Energy reconstruction
        df['reco_log_energy'] = energy_pipeline.predict(df[feature_list].values)
        df['reco_energy'] = 10**df['reco_log_energy']
        
    print('Running composition classifications...')
    pred_target = comp_pipeline.predict(df_sim_test[feature_list].values)
    
    print('Making response matrix...')
    log_reco_energy_sim_test = df_sim_test['reco_log_energy']
    log_true_energy_sim_test = df_sim_test['MC_log_energy']
    true_target = df_sim_test['comp_target_{}'.format(num_groups)].values

    response, response_err = comp.response_matrix(true_energy=log_true_energy_sim_test,
                                                  reco_energy=log_reco_energy_sim_test,
                                                  true_target=true_target,
                                                  pred_target=pred_target,
                                                  efficiencies=efficiencies,
                                                  efficiencies_err=efficiencies_err,
                                                  energy_bins=energybins.log_energy_bins)
    
#     fig, ax = plt.subplots()
#     im = ax.imshow(response, origin='lower', cmap='viridis')
#     ax.plot([0, response.shape[0] - 1], [0, response.shape[1] - 1], marker='None', ls=':', color='white')
#     comp.plotting.colorbar(im, label='Normalized response matrix')

#     ax.set_xlabel('True bins')
#     ax.set_ylabel('Reconstructed bins')

#     plt.show()
    
    # Run unfolding for each of the priors
    prior_name = 'H4a'
    prior = None if prior_name == 'uniform' else df_priors['{}_prior'.format(prior_name)]
    # priors = 'Jeffreys' if prior_name == 'Jeffreys' else df['{}_prior'.format(prior_name)]
    unfolded_result = pyunfold.iterative_unfold(data=counts_pyunfold,
                                                data_err=counts_err_pyunfold,
                                                response=response,
                                                response_err=response_err,
                                                efficiencies=efficiencies,
                                                efficiencies_err=efficiencies_err,
                                                ts='ks',
                                                ts_stopping=0.005,
                                                prior=prior,
                                                return_iterations=True)
    
    counts, counts_sys_err, counts_stat_err = comp.unfolded_counts_dist(unfolded_result,
                                                                        num_groups=num_groups)
    unfolded_fluxes = {}
    unfolded_sys_err = {}
    unfolded_stat_err = {}
    for composition in comp_list + ['total']:

        flux, flux_sys_err = counts_to_flux(counts=counts[composition],
                                            counts_err=counts_sys_err[composition],
                                            composition=composition)
        flux, flux_stat_err = counts_to_flux(counts=counts[composition],
                                             counts_err=counts_stat_err[composition],
                                             composition=composition)
        
        unfolded_fluxes[composition] = flux
        unfolded_sys_err[composition] = flux_sys_err
        unfolded_stat_err[composition] = flux_stat_err
    
    output = [unfolded_fluxes]
    if return_sys_err:
        output.append(unfolded_sys_err)
    if return_stat_err:
        output.append(unfolded_stat_err)
    
    if len(output) == 1:
        output = output[0]

    return output<jupyter_output><empty_output><jupyter_text>Set up mapping between dataset name and file locations<jupyter_code>sim_files = {'nominal': '/data/user/jbourbeau/composition/IC86.2012/sim/processed_hdf/nominal/*.hdf',
             'snow_lambda_up': '/data/user/jbourbeau/composition/IC86.2012/sim/processed_hdf/snow_lambda_2.45/*.hdf',
             'snow_lambda_down': '/data/user/jbourbeau/composition/IC86.2012/sim/processed_hdf/snow_lambda_2.05/*.hdf',
             'vem_cal_down': '/data/user/jbourbeau/composition/IC86.2012/sim/sim_dataframe_vem_cal_down.hdf',
             'vem_cal_up': '/data/user/jbourbeau/composition/IC86.2012/sim/sim_dataframe_vem_cal_up.hdf',
             'dom_eff_up': '/data/user/jbourbeau/composition/IC86.2012/sim/processed_hdf/dom_eff_1.08504/*.hdf',
             'dom_eff_down': '/data/user/jbourbeau/composition/IC86.2012/sim/processed_hdf/dom_eff_0.86625/*.hdf',
            }

unfolded_results = {}
for sys_name, f in sim_files.items():
    unfolded_results[sys_name] = get_unfolded_fluxes(f)
nominal_files = '/data/user/jbourbeau/composition/IC86.2012/sim/processed_hdf/nominal/*.hdf'
nominal_flux, nominal_flux_sys_err = get_unfolded_fluxes(nominal_files, return_sys_err=True)
unfolded_results['unfolding_sys_up'] = {}
unfolded_results['unfolding_sys_down'] = {}
for composition in comp_list + ['total']:
    unfolded_results['unfolding_sys_up'][composition] = nominal_flux[composition] + nominal_flux_sys_err[composition]
    unfolded_results['unfolding_sys_down'][composition] = nominal_flux[composition] - nominal_flux_sys_err[composition]
unfolded_results['unfolding_sys_up']
sys_names = ['snow_lambda', 'vem_cal', 'dom_eff', 'unfolding_sys']
sys_names
sys_labels = {'snow_lambda': 'IceTop snow correction',
              'vem_cal': 'IceTop VEM calibration',
              'dom_eff': 'In-ice uncertainties',
              'nominal': 'Nominal',
              'unfolding_sys': 'Unfolding systematic'}
sys_markers = {'snow_lambda': '*',
               'vem_cal': 'X',
               'dom_eff': '.',
               'nominal': 'o',
               'unfolding_sys': '^'}
fig, ax = plt.subplots()

sys_color_dict = {'light': sns.color_palette('Blues', len(sys_names) + 1).as_hex(),
                  'heavy': sns.color_palette('Oranges', len(sys_names) + 1).as_hex(),
                  'total': sns.color_palette('Greens', len(sys_names) + 1).as_hex()}

for name, flux in unfolded_results.items():
#     marker = plot_marker(name)
    
    if 'snow_lambda' in name:
        marker = '^'
    elif 'dom_eff' in name:
        marker = 'X'
    elif 'vem_cal' in name:
        marker = 'v'
    else:
        marker = 'o'
    
    
    for idx, composition in enumerate(comp_list + ['total']):
        if composition == 'total':
            label_terms = name.split('_')
            stop_idx = -1 if len(label_terms) > 1 else None
            label = '_'.join(label_terms[:stop_idx])
            label = sys_labels[label]
        else:
            label = ''
        
        if 'down' in name:
            label = ''
            

        ax.plot(energybins.log_energy_midpoints, flux[composition],
                color=color_dict[composition],
                alpha=0.7, 
                ls=':' if 'nominal' in name else 'None',
                lw=0.8,
                marker=marker,
                label=label)

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=0.5, which='both')

# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
# ax.legend(loc='upper center',
#           bbox_to_anchor=(0.5, 1.35),
#           ncol=3)
ax.legend(title='Detector systematic uncertainties')

outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'detector-systematics-flux-comparison.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
nominal_flux = unfolded_results['nominal']
for sys_type in sys_names:
    
    sys_results = {key:value for key, value in unfolded_results.items() if sys_type in key}
    fig, ax = plt.subplots()
    
    for sys_name, flux in sys_results.items():
        if 'up' in sys_name:
            marker = '^'
        else:
            marker = 'v'    
    
        for idx, composition in enumerate(comp_list + ['total']):
            ax.plot(energybins.log_energy_midpoints, flux[composition],
                    color=color_dict[composition],
                    alpha=0.7, 
                    ls='None',
                    marker=marker,
                   )
    for idx, composition in enumerate(comp_list + ['total']):
        # Plot nominal flux
        ax.plot(energybins.log_energy_midpoints, nominal_flux[composition],
                color=color_dict[composition],
                alpha=0.7, 
                ls=':',
                lw=0.8,
                marker='o',
                label=composition
               )

    ax.set_yscale('log', nonposy='clip')
    ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
    ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
    ax.set_xlim(6.4, 7.8)
    ax.set_ylim(1e3, 1e5)
    ax.grid(lw=0.5, which='both')

    # ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    # ax.legend(loc='upper center',
    #           bbox_to_anchor=(0.5, 1.35),
    #           ncol=3)
    title = sys_labels[sys_type]
    ax.legend(loc='lower left', title=title)

    outfile = os.path.join(comp.paths.figures_dir,
                           'systematics',
                           '{}-flux-comparison.png'.format(sys_type))
    comp.check_output_dir(outfile)
    plt.savefig(outfile)

    plt.show()
energybins.log_energy_midpoints
bin_width = energybins.log_energy_bin_width[0]
num_sys = len(sys_names)
step = bin_width / num_sys
step
nominal_flux = unfolded_results['nominal']

sys_color_dict = {'light': sns.color_palette('Blues', len(sys_names) + 1).as_hex(),
                  'heavy': sns.color_palette('Oranges', len(sys_names) + 2).as_hex()[:-1],
                  'total': sns.color_palette('Greens', len(sys_names) + 1).as_hex()}


fig, ax = plt.subplots()
for sys_idx, sys_type in enumerate(sys_names):
    
    sys_results = {key:value for key, value in unfolded_results.items() if sys_type in key}
    
    for sys_name, flux in sys_results.items():
        if 'up' in sys_name:
            marker = '^'
        else:
            marker = 'v'    
    
        for idx, composition in enumerate(comp_list + ['total']):
            sys_fluxes = [value[composition] for value in sys_results.values()]
            sys_flux_max = np.max(sys_fluxes, axis=0)
            sys_flux_min = np.min(sys_fluxes, axis=0)
            x = energybins.log_energy_bins[:-1] + sys_idx * step
            
            if composition == 'total' and 'up' in sys_name:
                label = sys_labels[sys_type]
            else:
                label = ''
            
            for idx in range(len(x)):
                ax.fill_between([x[idx], x[idx]+step],
                                sys_flux_max[idx],
                                sys_flux_min[idx],
                                color=sys_color_dict[composition][sys_idx],
                                alpha=0.7,
                                label=label
                                )
#                 # Plot markers for different systematics
#                 ax.plot(x[idx] + step / 2, nominal_flux[composition][idx],
#                         color=color_dict[composition],
#                         alpha=1, 
#                         ls='None',
#                         marker=sys_markers[sys_type],
#                         label=label
#                        )
                if label:
                    label = ''

# Plot nominal flux
for idx, composition in enumerate(comp_list + ['total']):
    if composition == 'total':
        label = 'Nominal'
    else:
        label = ''
    comp.plot_steps(energybins.log_energy_bins, nominal_flux[composition],
                    color=color_dict[composition],
                    alpha=1, 
                    ls='-',
                    lw=0.9,
                    ax=ax,
                    label=label
                   )

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')

major_ticks = energybins.log_energy_bins
ax.set_xticks(major_ticks, minor=True)

ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)

ax.grid(lw=0.5,
        which='minor')

# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
# ax.legend(loc='upper center',
#           bbox_to_anchor=(0.5, 1.35),
#           ncol=3)
# title = sys_labels[sys_type]
ax.legend(loc='lower left',
#           title='Systematic uncertainties',
         )
outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'systematics-flux-comparison-step.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
nominal_flux = unfolded_results['nominal']

sys_color_dict = {'light': sns.color_palette('Blues', len(sys_names) + 1).as_hex(),
                  'heavy': sns.color_palette('Oranges', len(sys_names) + 2).as_hex()[:-1],
                  'total': sns.color_palette('Greens', len(sys_names) + 1).as_hex()}


fig, ax = plt.subplots()
for sys_idx, sys_type in enumerate(sys_names):
    
    sys_results = {key:value for key, value in unfolded_results.items() if sys_type in key}
    
    for sys_name, flux in sys_results.items():
        if 'up' in sys_name:
            marker = '^'
        else:
            marker = 'v'    
    
        for idx, composition in enumerate(comp_list + ['total']):
            sys_fluxes = [value[composition] for value in sys_results.values()]
            sys_flux_max = np.max(sys_fluxes, axis=0)
            sys_flux_min = np.min(sys_fluxes, axis=0)
            
            sys_abs_deviation = np.max([np.abs(nominal_flux[composition] - sys_flux_max),
                                        np.abs(nominal_flux[composition] - sys_flux_min)],
                                       axis=0)
            
            x = energybins.log_energy_bins[:-1] + sys_idx * step
            
            if composition == 'total' and 'up' in sys_name:
                label = sys_labels[sys_type]
            else:
                label = ''
            
            for idx in range(len(x)):
                ax.fill_between([x[idx], x[idx]+step],
                                nominal_flux[composition][idx] + sys_abs_deviation[idx],
                                nominal_flux[composition][idx] - sys_abs_deviation[idx],
                                color=sys_color_dict[composition][sys_idx],
                                alpha=0.7,
                                label=label
                                )
                if label:
                    label = ''

# Plot nominal flux
for idx, composition in enumerate(comp_list + ['total']):
    if composition == 'total':
        label = 'Nominal'
    else:
        label = ''
    comp.plot_steps(energybins.log_energy_bins, nominal_flux[composition],
                    color=color_dict[composition],
                    alpha=1, 
                    ls='-',
                    lw=0.9,
                    ax=ax,
                    label=label
                   )

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')

major_ticks = energybins.log_energy_bins
ax.set_xticks(major_ticks, minor=True)

ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)

ax.grid(lw=0.5,
        which='minor')

# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
# ax.legend(loc='upper center',
#           bbox_to_anchor=(0.5, 1.35),
#           ncol=3)
# title = sys_labels[sys_type]
ax.legend(loc='lower left',
#           title='Systematic uncertainties',
         )
outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'systematics-flux-comparison-step-max-deviation.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
nominal_flux = unfolded_results['nominal']

sys_color_dict = {'light': sns.color_palette('Blues', len(sys_names) + 1).as_hex(),
                  'heavy': sns.color_palette('Oranges', len(sys_names) + 2).as_hex()[:-1],
                  'total': sns.color_palette('Greens', len(sys_names) + 1).as_hex()}


fig, ax = plt.subplots()
for sys_idx, sys_type in enumerate(sys_names):
    
    sys_results = {key:value for key, value in unfolded_results.items() if sys_type in key}
    
    for sys_name, flux in sys_results.items():
        if 'up' in sys_name:
            marker = '^'
        else:
            marker = 'v'    
    
        for idx, composition in enumerate(comp_list + ['total']):
            sys_fluxes = [value[composition] for value in sys_results.values()]
            sys_fluxes.append(nominal_flux[composition])
            sys_flux_max = np.max(sys_fluxes, axis=0)
            sys_flux_min = np.min(sys_fluxes, axis=0)
            
            sys_std_deviation = np.std(sys_fluxes, axis=0)
            
            x = energybins.log_energy_bins[:-1] + sys_idx * step
            
            if composition == 'total' and 'up' in sys_name:
                label = sys_labels[sys_type]
            else:
                label = ''
            
            for idx in range(len(x)):
                ax.fill_between([x[idx], x[idx]+step],
                                nominal_flux[composition][idx] + sys_std_deviation[idx],
                                nominal_flux[composition][idx] - sys_std_deviation[idx],
                                color=sys_color_dict[composition][sys_idx],
                                alpha=0.7,
                                label=label
                                )
                if label:
                    label = ''

# Plot nominal flux
for idx, composition in enumerate(comp_list + ['total']):
    if composition == 'total':
        label = 'Nominal'
    else:
        label = ''
    comp.plot_steps(energybins.log_energy_bins, nominal_flux[composition],
                    color=color_dict[composition],
                    alpha=1, 
                    ls='-',
                    lw=0.9,
                    ax=ax,
                    label=label
                   )

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')

major_ticks = energybins.log_energy_bins
ax.set_xticks(major_ticks, minor=True)

ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)

ax.grid(lw=0.5,
        which='minor')

# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
# ax.legend(loc='upper center',
#           bbox_to_anchor=(0.5, 1.35),
#           ncol=3)
# title = sys_labels[sys_type]
ax.legend(loc='lower left',
#           title='Systematic uncertainties',
         )
outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'systematics-flux-comparison-step-std.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
for name, flux in unfolded_results.items():
    
    fig, ax = plt.subplots()
    marker_dict = {'nominal': 'o', 'up': '^', 'down': 'v'}
    for idx, composition in enumerate(comp_list + ['total']):
        if composition == 'total':
            label = name.capitalize().replace('_', ' ')
        else:
            label = ''
        
        sys = (flux[composition] - nominal_flux[composition]) / nominal_flux[composition]
        comp.plot_steps(energybins.log_energy_bins,
                        sys,
                        color=color_dict[composition],
#                         ls='None',
#                         marker=marker,
                        label=composition)
    ax.axhline(0, ls='--', marker='None', color='k')
#     ax.set_yscale('log', nonposy='clip')
    ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
    ax.set_ylabel('Fractional Difference')
    ax.set_xlim(6.4, 7.8)
    ax.set_ylim(-1, 1)
    ax.grid(lw=0.5, which='both')

    ax.legend(title=label)

    # outfile = os.path.join(comp.paths.figures_dir,
    #                        'systematics',
    #                        'combined-flux-comparison.png')
    # comp.check_output_dir(outfile)
    # plt.savefig(outfile)

    plt.show()
s_up = set(i for i in b.data_vars if 'up' in i)
s_down = set(i for i in b.data_vars if 'down' in i)
s_up, s_down
up_err = np.sqrt(np.sum([(b[i] - b['nominal'])**2 for i in s_up], axis=1))
down_err = np.sqrt(np.sum([(b[i] - b['nominal'])**2 for i in s_down], axis=1))
up_err
np.asarray([up_err[0], down_err[0]])
fig, ax = plt.subplots()
colors = ['k', 'red']
names = ['VEM cal.', 'Snow $\mathrm{\lambda}$']
marker_dict = {'nominal': 'o', 'up': '^', 'down': 'v'}
for idx, composition in enumerate(comp_list + ['total']):
    if composition == 'total':
        label = name.capitalize().replace('_', ' ')
    else:
        label = ''

    comp.plot_steps(energybins.log_energy_bins,
                    b['nominal'][idx],
                    yerr=[down_err[idx], up_err[idx]],
                    color=color_dict[composition],
                    fillcolor='gray',
                    ls='-',
                    label=composition.capitalize(),
                    ax=ax)

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=0.5, which='both')
ax.legend()

outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'net-systematics-flux.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
s = set([i.replace('_up', '').replace('_down', '') for i in b.data_vars if i != 'nominal'])
s
len(s)
dw = 0.1 / len(s)
dw / 2
fig, ax = plt.subplots()
colors = ['k', 'red']
names = ['VEM cal.', 'Snow $\mathrm{\lambda}$']
marker_dict = {'nominal': 'o', 'up': '^', 'down': 'v'}
for name in s:
    up_flux = b[name+'_up']
    down_flux = b[name+'_down']
    for idx, composition in enumerate(comp_list + ['total']):
        
        if composition == 'total':
            label = name.capitalize().replace('_', ' ')
        else:
            label = ''
        for j, i in enumerate(energybins.log_energy_midpoints):
            x = [i - 0.01, i + 0.01]
            ax.fill_between(x,
                            [up_flux[idx].values[j]]*2,
                            [down_flux[idx].values[j]]*2,
                            color=color_dict[composition],
                            alpha=0.3)
#                             label=label)

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=0.5, which='both')
ax.legend()

# outfile = os.path.join(comp.paths.figures_dir,
#                        'systematics',
#                        'combined-flux-comparison.png')
# comp.check_output_dir(outfile)
# plt.savefig(outfile)

plt.show()
[up_flux[idx].values[i]]*2
fig, ax = plt.subplots()
colors = ['k', 'red']
names = ['VEM cal.', 'Snow $\mathrm{\lambda}$']
marker_dict = {'nominal': 'o', 'up': '^', 'down': 'v'}
for name in b.data_vars:
    flux = b[name]
    marker = plot_marker(name)
    for idx, composition in enumerate(comp_list + ['total']):
        if composition == 'total':
            label = name.capitalize().replace('_', ' ')
        else:
            label = ''

        ax.plot(energybins.log_energy_midpoints, flux[idx],
                color=color_dict[composition],
                ls='None',
                marker=marker,
                label=label)

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=0.5, which='both')
ax.legend()

# outfile = os.path.join(comp.paths.figures_dir,
#                        'systematics',
#                        'combined-flux-comparison.png')
# comp.check_output_dir(outfile)
# plt.savefig(outfile)

plt.show()
fig, ax = plt.subplots()
systematics_files = ['vem_calibration_systematic_flux.pkl', 'snow_lambda_systematic_flux.pkl']
colors = ['k', 'red']
names = ['VEM cal.', 'Snow $\mathrm{\lambda}$']
marker_dict = {'nominal': 'o', 'up': '^', 'down': 'v'}
for f, sys_color, sys_name in zip(systematics_files, colors, names):
    unfolded_results = joblib.load(f)
    linestyles = iter(['-', ':', '-.', '--'])
    markers = iter(['.', '^', 'X'])
    for shift_type, unfolded_systematic in unfolded_results.items():
        marker = marker_dict[shift_type]
        for composition in comp_list + ['total']:
            if composition == 'total':
                label = shift_type.capitalize() if shift_type == 'nominal' else '{} {}'.format(sys_name, shift_type.capitalize())
            else:
                label = ''
            
            if shift_type == 'nominal':
                comp.plot_steps(energybins.log_energy_bins, unfolded_systematic[composition].flux,
                                yerr=unfolded_systematic[composition].flux_sys_err,
                                color=color_dict[composition],
                                ls='-',
                                ax=ax)

                ax.errorbar(energybins.log_energy_midpoints, unfolded_systematic[composition].flux,
                            yerr=unfolded_systematic[composition].flux_stat_err,
                            color=color_dict[composition],
                            ls='None',
                            marker=marker,
                            capsize=20,
                            elinewidth=1,
                            label=label)
            else:
                ax.plot(energybins.log_energy_midpoints, unfolded_systematic[composition].flux,
                        color=sys_color,
                        ls='None',
                        marker=marker,
                        label=label)

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=0.5, which='both')
ax.legend()

outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'combined-flux-comparison.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
fig, ax = plt.subplots(figsize=(10, 6))
# unfolded_counts = {}
linestyles = iter(['-', ':', '-.', '--'])
markers = iter(['o', '^', 'X'])
vem_systematic = []
for shift_type in ['down', 'up']:
    shift_systematic = {}
    marker = next(markers)
    ls = next(linestyles)
    for composition in comp_list + ['total']:
        nominal = unfolded_counts['nominal'][composition]
        deviation = (unfolded_counts[shift_type][composition] - nominal) / nominal
        energy_mask = np.logical_and(energybins.log_energy_midpoints < 7.8,
                                     energybins.log_energy_midpoints > 6.4)
        ax.plot(energybins.log_energy_midpoints[energy_mask], deviation[energy_mask],
                color=color_dict[composition],
                marker=marker,
                ls=ls,
                label='{} 0.2 m {}'.format(shift_type, composition))
        shift_systematic[composition] = np.mean(np.abs(deviation[energy_mask]))
        
    vem_systematic.append(shift_systematic)

ax.axhline(0.0, marker='None', ls='-.', color='k')
ax.set(xlabel='$\mathrm{\log_{10}(E/GeV)}$',
       xlim=(6.4, 7.8),
       ylabel='\% Difference w.r.t. Nominal')
ax.grid()
# ax.legend()
ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))

outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'snow-lambda-counts-deviation.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
fig, ax = plt.subplots()
# fig, ax = plt.subplots(figsize=(10, 6))
for composition in comp_list + ['total']:
    nominal = unfolded_counts['nominal'][composition]
    
    max_counts = np.stack((unfolded_counts['down'][composition], unfolded_counts['up'][composition])).max(axis=0)
    min_counts = np.stack((unfolded_counts['down'][composition], unfolded_counts['up'][composition])).min(axis=0)
    
    deviation_upper = (max_counts - nominal) / nominal
    deviation_lower = (min_counts - nominal) / nominal
    energy_mask = np.logical_and(energybins.log_energy_midpoints < 7.8,
                                 energybins.log_energy_midpoints > 6.4)
    ax.fill_between(energybins.log_energy_midpoints[energy_mask],
                    deviation_upper[energy_mask],
                    deviation_lower[energy_mask],
                    color=color_dict[composition],
                    alpha=0.4,
                    label=composition.capitalize())

ax.axhline(0.0, marker='None', ls='-.', color='k')
ax.set(xlabel='$\mathrm{\log_{10}(E/GeV)}$',
       xlim=(6.4, 7.8),
       ylabel='\% Difference w.r.t. Nominal')
ax.grid()
ax.legend(loc='upper left')
# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))

outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'snow-lambda-counts-deviation-bands.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()
fig, axes = plt.subplots(ncols=1, nrows=3,
                         sharey=True, sharex=True,
                         gridspec_kw={'hspace':0.15})

for idx, (composition, ax) in enumerate(zip(comp_list + ['total'], axes)):
   
    nominal = unfolded_counts['nominal'][composition]
    
    max_counts = np.stack((unfolded_counts['down'][composition], unfolded_counts['up'][composition])).max(axis=0)
    min_counts = np.stack((unfolded_counts['down'][composition], unfolded_counts['up'][composition])).min(axis=0)
    
    deviation_upper = (max_counts - nominal) / nominal
    deviation_lower = (min_counts - nominal) / nominal
    energy_mask = np.logical_and(energybins.log_energy_midpoints < 7.8,
                                 energybins.log_energy_midpoints > 6.4)
    ax.fill_between(energybins.log_energy_midpoints[energy_mask],
                    deviation_upper[energy_mask],
                    deviation_lower[energy_mask],
                    color=color_dict[composition],
                    alpha=0.4,
                    label=composition.capitalize())

    ax.axhline(0.0, marker='None', ls='-.', color='k')
    ax.set_xlim(6.4, 7.8)
    if idx == 1:
        ax.set_ylabel('\% Difference w.r.t. Nominal')
    ax.grid()
    ax.legend(loc='lower left')
    
ax.set(xlabel='$\mathrm{\log_{10}(E/GeV)}$')
# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))

outfile = os.path.join(comp.paths.figures_dir,
                       'systematics',
                       'snow-lambda-counts-deviation-bands-2.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)
plt.tight_layout()
plt.show()
vem_systematic = pd.DataFrame.from_records(vem_systematic, index=['down', 'up'])
vem_systematic
# fig, ax = plt.subplots()
# im = ax.imshow(response, origin='lower', cmap='viridis')
# ax.plot([0, response.shape[0] - 1], [0, response.shape[1] - 1], marker='None', ls=':', color='white')
# comp.plotting.colorbar(im, label='Normalized response matrix')

# ax.set_xlabel('True bins')
# ax.set_ylabel('Reconstructed bins')

# plt.show()
# Run unfolding for each of the priors
names = ['uniform', 'H3a', 'H4a', 'simple_power_law']
# names = ['Jeffreys', 'H3a', 'H4a', 'Polygonato']
unfolding_results = {}
for prior_name in names:
    prior = None if prior_name == 'uniform' else df_priors['{}_prior'.format(prior_name)]
    # priors = 'Jeffreys' if prior_name == 'Jeffreys' else df['{}_prior'.format(prior_name)]
    df_unfolding_iter = pyunfold.iterative_unfold(data=counts_pyunfold,
                                                  data_err=counts_err_pyunfold,
                                                  response=response,
                                                  response_err=response_err,
                                                  efficiencies=efficiencies,
                                                  efficiencies_err=efficiencies_err,
                                                  ts='ks',
                                                  ts_stopping=0.005,
                                                  prior=prior,
                                                  return_iterations=True)
#                                                   callbacks=[pyunfold.callbacks.SplineRegularizer(degree=1, smooth=256)])
    unfolding_results[prior_name] = df_unfolding_iter
names
fig, ax = plt.subplots()
# fig, ax = plt.subplots(figsize=(10, 5))
linestyles = ['-', ':', '-.', '--']
markers = ['.', '*', '^', 'X']
for prior_name, ls, marker in zip(names, linestyles, markers):
#     if prior_name != 'H4a':
#         continue
    counts, counts_sys_err, counts_stat_err = comp.unfolded_counts_dist(unfolding_results[prior_name],
                                                                        num_groups=num_groups)
    for composition in comp_list + ['total']:
#         print(composition)
#         print(counts[composition])
        flux, flux_sys_err = counts_to_flux(counts=counts[composition],
                                            counts_err=counts_sys_err[composition],
                                            composition=composition)
        flux, flux_stat_err = counts_to_flux(counts=counts[composition],
                                             counts_err=counts_stat_err[composition],
                                             composition=composition)
        comp.plot_steps(energybins.log_energy_bins, flux,
                        yerr=flux_sys_err,
                        color=color_dict[composition],
                        ls=ls,
                        ax=ax)
#         label = '{} ({})'.format(composition, prior_name.replace('_', '-'))
        label = prior_name.replace('_', ' ') if composition == 'total' else ''
        ax.errorbar(energybins.log_energy_midpoints, flux,
                    yerr=flux_stat_err,
                    color=color_dict[composition],
                    ls='None',
                    marker=marker,
                    capsize=20,
                    elinewidth=1,
                    label=label)
ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=1, which='both')
ax.legend()
# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
# ax.legend(loc='upper center',
#           bbox_to_anchor=(0.5, 1.75),
#           ncol=2,
#           fancybox=True)

# outfile = os.path.join(comp.paths.figures_dir, 'skymaps', 'allsky.png')
# comp.check_output_dir(outfile)
# plt.savefig(outfile)

plt.show()
unfolding_results[prior_name]
fig, ax = plt.subplots()
# fig, ax = plt.subplots(figsize=(10, 5))

prior_name = 'H4a'
marker = '.'
ls = '-'
counts, counts_sys_err, counts_stat_err = comp.unfolded_counts_dist(unfolding_results[prior_name],
                                                                    num_groups=num_groups)
for composition in comp_list + ['total']:
    flux, flux_sys_err = counts_to_flux(counts=counts[composition],
                                        counts_err=counts_sys_err[composition],
                                        composition=composition)
    flux, flux_stat_err = counts_to_flux(counts=counts[composition],
                                         counts_err=counts_stat_err[composition],
                                         composition=composition)
    comp.plot_steps(energybins.log_energy_bins, flux,
                    yerr=flux_sys_err,
                    color=color_dict[composition],
                    fillalpha=0.4,
                    ls=ls,
                    label=composition,
                    ax=ax)
    ax.errorbar(energybins.log_energy_midpoints, flux,
                yerr=flux_stat_err,
                color=color_dict[composition],
                ls='None',
                marker=marker,
                capsize=20,
                elinewidth=1)

ax.set_yscale('log', nonposy='clip')
ax.set_xlabel('$\mathrm{\log_{10}(E/GeV)}$')
ax.set_ylabel('$\mathrm{ E^{2.7} \ J(E) \ [GeV^{1.7} m^{-2} sr^{-1} s^{-1}]}$')
ax.set_xlim(6.4, 7.8)
ax.set_ylim(1e3, 1e5)
ax.grid(lw=1, which='both')

ax.legend()
# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
# ax.legend(loc='upper center',
#           bbox_to_anchor=(0.5, 1.75),
#           ncol=2,
#           fancybox=True)

outfile = os.path.join(comp.paths.figures_dir, 'skymaps', 'allsky.png')
comp.check_output_dir(outfile)
plt.savefig(outfile)

plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13665028174079846, 'fraction_numerical': 0.013547536266634696, 'mean_word_length': 2.068648370245015, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 105, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '340426', 'license': 'MIT', 'n_tokens_mistral': 14062, 'n_tokens_neox': 13048, 'n_words': 2097, 'path': 'systematics/systematics.ipynb', 'repo_name': 'jrbourbeau/cr-composition', 'size': '41705'}