notebooks/nz.py

# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.16.2
#   kernelspec:
#     display_name: temps
#     language: python
#     name: temps
# ---

# # FIGURE 5 IN THE PAPER

# ## n(z) distributions

# %load_ext autoreload
# %autoreload 2

import pandas as pd
import numpy as np
from astropy.io import fits
from astropy.table import Table
import torch
from pathlib import Path

#matplotlib settings
from matplotlib import rcParams
import matplotlib.pyplot as plt
rcParams["mathtext.fontset"] = "stix"
rcParams["font.family"] = "STIXGeneral"

from temps.archive import Archive 
from temps.utils import nmad
from temps.temps_arch import EncoderPhotometry, MeasureZ
from temps.temps import TempsModule
from temps.plots import plot_nz

eval_methods=False

# ### LOAD DATA

#define here the directory containing the photometric catalogues
parent_dir = Path('/data/astro/scratch/lcabayol/insight/data/Euclid_EXT_MER_PHZ_DC2_v1.5')
modules_dir = Path('../data/models/')

# +
filename_valid='euclid_cosmos_DC2_S1_v2.1_valid_matched.fits'

hdu_list = fits.open(parent_dir / filename_valid)
cat = Table(hdu_list[1].data).to_pandas()
cat = cat[cat['FLAG_PHOT']==0]
cat = cat[cat['mu_class_L07']==1]
cat = cat[(cat['z_spec_S15'] > 0) | (cat['photo_z_L15'] > 0)]
cat = cat[cat['MAG_VIS']<25]

# -

ztarget = [cat['z_spec_S15'].values[ii] if cat['z_spec_S15'].values[ii]> 0 else cat['photo_z_L15'].values[ii] for ii in range(len(cat))]
specz_or_photo = [0 if cat['z_spec_S15'].values[ii]> 0 else 1 for ii in range(len(cat))]
ID = cat['ID']
VISmag = cat['MAG_VIS']
zsflag = cat['reliable_S15']

photoz_archive = Archive(path = parent_dir,only_zspec=False)
f, ferr = photoz_archive._extract_fluxes(catalogue= cat)
col, colerr = photoz_archive._to_colors(f, ferr)

# ### LOAD TRAINED MODELS AND EVALUATE PDFs AND REDSHIFT

if eval_methods:
    dfs = {}

    for il, lab in enumerate(['z','L15','DA']):
    
        nn_features = EncoderPhotometry()
        nn_features.load_state_dict(torch.load(modules_dir / f'modelF_{lab}.pt',map_location=torch.device('cpu')))
        nn_z = MeasureZ(num_gauss=6)
        nn_z.load_state_dict(torch.load(modules_dir / f'modelZ_{lab}.pt',map_location=torch.device('cpu')))
        
        temps_module = TempsModule(nn_features, nn_z)

        z, pz, odds = temps_module.get_pz(input_data=torch.Tensor(col), 
                                    return_pz=True)
        # Create a DataFrame with the desired columns
        df = pd.DataFrame(np.c_[ID, VISmag,z, odds, ztarget,zsflag, specz_or_photo], 
                          columns=['ID','VISmag','z','odds', 'ztarget','zsflag','S15_L15_flag'])

        # Calculate additional columns or operations if needed
        df['zwerr'] = (df.z - df.ztarget) / (1 + df.ztarget)

        # Drop any rows with NaN values
        df = df.dropna()

        # Assign the DataFrame to a key in the dictionary
        dfs[lab] = df


# ### LOAD CATALOGUES IF AVAILABLE

if not eval_methods:

    df_zs = pd.read_csv(parent_dir / 'predictions_specztraining.csv', header=0) 
    df_zsL15 = pd.read_csv(parent_dir / 'predictions_speczL15training.csv', header=0) 
    df_DA = pd.read_csv(parent_dir / 'predictions_speczDAtraining.csv', header=0) 


    dfs = {}
    dfs['z'] = df_zs
    dfs['L15'] = df_zsL15
    dfs['DA'] = df_DA

# +
import matplotlib.pyplot as plt
from matplotlib import gridspec

# Create figure and grid specification
fig = plt.figure(figsize=(8, 10))
gs = gridspec.GridSpec(5, 1, height_ratios=[0.1, 1, 1,1,1])

# Upper panel (very thin) with shaded areas
ax1 = plt.subplot(gs[0])
ax1.set_yticks([])

ax1.set_ylabel('Bins', fontsize=10)

# Define the ranges for shaded areas
#z_ranges = [[0.15, 0.35], [0.35, 0.55], [0.55, 0.85], [0.85, 1.05], [1.05, 1.35], 
#                 [1.35, 1.55],# [1.55, 1.85], [1.85, 2], [2, 2.5], [2.5, 3], [3, 4]]
            
z_ranges = [[0.15, 0.5], [0.5, 1], [1, 1.5], [1.5,2]]#, [2, 3], [3,4]]#, 
                 #[1.35, 1.55],# [1.55, 1.85], [1.85, 2], [2, 2.5], [2.5, 3], [3, 4]]
            
colors = ['deepskyblue', 'forestgreen', 'coral', 'grey', 'pink', 'goldenrod',
          'cyan', 'seagreen', 'salmon', 'steelblue', 'orange']

# Plot shaded areas
x_values = [0, 1, 2]  # Example x values, adjust as needed
for i, (start, end) in enumerate(z_ranges):
    ax1.fill_betweenx(x_values, start, end, color=colors[i], alpha=0.5)

# Middle panel (equally thick)
ax2 = plt.subplot(gs[1])
for i, (start, end) in enumerate(z_ranges):
    dfplot_z = dfs['z'][(dfs['z']['ztarget'] > start) & (dfs['z']['ztarget'] < end)]
    ax2.hist(dfplot_z.ztarget, bins=50, color=colors[i], histtype='step', linestyle='-', density=True, range=(0, 4))

# Bottom panel (equally thick)
ax3 = plt.subplot(gs[2])
for i, (start, end) in enumerate(z_ranges):
    dfplot_z = dfs['z'][(dfs['z']['z'] > start) & (dfs['z']['z'] < end)]
    ax3.hist(dfplot_z.ztarget, bins=50, color=colors[i], histtype='step', linestyle='-', density=True, range=(0, 4))
    
# Bottom panel (equally thick)
ax4 = plt.subplot(gs[3])
for i, (start, end) in enumerate(z_ranges):
    dfplot_z = dfs['L15'][(dfs['L15']['z'] > start) & (dfs['L15']['z'] < end)]
    print(len(dfplot_z))
    ax4.hist(dfplot_z.ztarget, bins=50, color=colors[i], histtype='step', linestyle='-', density=True, range=(0, 4))
    
ax5 = plt.subplot(gs[4])
for i, (start, end) in enumerate(z_ranges):
    dfplot_z = dfs['DA'][(dfs['DA']['z'] > start) & (dfs['DA']['z'] < end)]
    ax5.hist(dfplot_z.ztarget, bins=50, color=colors[i], histtype='step', linestyle='-', density=True, range=(0, 4))

plt.tight_layout()
plt.show()

# -

def plot_nz(df_list, 
            zcuts = [0.1, 0.5, 1, 1.5, 2, 3, 4],
            save=False):
    # Plot properties
    plt.rcParams['font.family'] = 'serif'
    plt.rcParams['font.size'] = 16

    cmap = plt.get_cmap('Dark2')  # Choose a colormap for coloring lines
    
    # Create subplots
    fig, axs = plt.subplots(3, 1, figsize=(20, 8), sharex=True)

    for i, df in enumerate(df_list):
        dfplot = df_list[i].copy()  # Assuming df_list contains dataframes
        ax = axs[i]  # Selecting the appropriate subplot
        
        for iz in range(len(zcuts)-1):
            dfplot_z = dfplot[(dfplot['ztarget'] > zcuts[iz]) & (dfplot['ztarget'] < zcuts[iz + 1])]
            color = cmap(iz)  # Get a different color for each redshift
            
            zt_mean = np.median(dfplot_z.ztarget.values)
            zp_mean = np.median(dfplot_z.z.values)

            
            # Plot histogram on the selected subplot
            ax.hist(dfplot_z.z, bins=50, color=color, histtype='step', linestyle='-', density=True, range=(0, 4))
            ax.axvline(zt_mean, color=color, linestyle='-', lw=2)
            ax.axvline(zp_mean, color=color, linestyle='--', lw=2)
            
        ax.set_ylabel(f'Frequency', fontsize=14)
        ax.grid(False)
        ax.set_xlim(0, 3.5)
    
    axs[-1].set_xlabel(f'$z$', fontsize=18)
    
    if save:
        plt.savefig(f'nz_hist.pdf', dpi=300, bbox_inches='tight')
    
    plt.show()

plot_nz(df_list)