<jupyter_start><jupyter_text>Apply CNN Classifier to DESI Spectra and visualize results with gradCAMMini-SV2 tiles from February-March 2020:- https://desi.lbl.gov/trac/wiki/TargetSelectionWG/miniSV2See also the DESI tile picker with (limited) SV0 tiles from March 2020:- https://desi.lbl.gov/svn/data/tiles/trunk/- https://desi.lbl.gov/svn/data/tiles/trunk/SV0.html<jupyter_code>import sys
sys.path.append('/global/homes/p/palmese/desi/timedomain/desitrip/py/')

from desispec.io import read_spectra, write_spectra
from desispec.spectra import Spectra
from desitarget.cmx.cmx_targetmask import cmx_mask

from desitrip.preproc import rebin_flux, rescale_flux

from astropy.io import fits
from astropy.table import Table, vstack, hstack

from glob import glob
from datetime import date

import os
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

from tensorflow import keras
mpl.rc('font', size=14)
# Set up BGS target bit selection.
cmx_bgs_bits = '|'.join([_ for _ in cmx_mask.names() if 'BGS' in _])<jupyter_output><empty_output><jupyter_text>Select a Date & Tile from the Tile Picker<jupyter_code># Choose tile 66000 from Mar 14, 2020
tile_id = 66003
obsdate = 20200315

# Access redux folder.
redux='/global/project/projectdirs/desi/spectro/redux/daily/tiles'
prefix_in='/'.join([redux, '{:05d}/{}'.format(tile_id, obsdate)])
if not os.path.isdir(prefix_in):
    print('{} does not exist.'.format(prefix_in))
# List zbest and coadd files.
# Data are stored by petal ID.
zbfiles = sorted(glob('{}/zbest*.fits'.format(prefix_in)))
cafiles = sorted(glob('{}/coadd*.fits'.format(prefix_in)))
# zbest files from redrock.
zbfiles
# Co-added spectra from multiple exposures + the three spectrograph arms.
cafiles<jupyter_output><empty_output><jupyter_text>Load the Keras ModelLoad a model trained on real or simulated data using the native Keras output format. In the future this could be updated to just load the Keras weights.<jupyter_code>tfmodel = '/global/homes/l/lehsani/timedomain/desitrip/docs/nb/models_9label_first/6_b65_e200_9label/b65_e200_9label_model'
#tfmodel = '/global/homes/s/sybenzvi/desi/timedomain/desitrip/docs/nb/6label_cnn_restframe'
if os.path.exists(tfmodel):
    classifier = keras.models.load_model(tfmodel)
else:
    classifier = None
    print('Sorry, could not find {}'.format(tfmodel))
if classifier is not None:
    classifier.summary()<jupyter_output>Model: "SNnet"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
Input_Spec (InputLayer)      [(None, 150, 1)]          0         
_________________________________________________________________
conv1d_20 (Conv1D)           (None, 150, 8)            48        
_________________________________________________________________
batch_normalization_20 (Batc (None, 150, 8)            32        
_________________________________________________________________
activation_20 (Activation)   (None, 150, 8)            0         
_________________________________________________________________
max_pooling1d_20 (MaxPooling (None, 75, 8)             0         
_________________________________________________________________
conv1d_21 (Conv1D)           (None, 75, 16)            656       
_____________________________________________________________[...]<jupyter_text>Loop Through Spectra and Classify<jupyter_code># Loop through zbest and coadd files for each petal.
# Extract the fibermaps, ZBEST tables, and spectra.
# Keep only BGS targets passing basic event selection.
allzbest = None
allfmap = None
allwave = None
allflux = None
allivar = None
allmask = None
allres  = None
for cafile, zbfile in zip(cafiles, zbfiles):
    # Access data per petal.
    zbest = Table.read(zbfile, 'ZBEST')
    fibermap = Table.read(zbfile, 'FIBERMAP')
    pspectra = read_spectra(cafile)

    # Apply standard event selection.
    isTGT = fibermap['OBJTYPE'] == 'TGT'
    isGAL = zbest['SPECTYPE'] == 'GALAXY'
    isBGS = fibermap['CMX_TARGET'] & cmx_mask.mask(cmx_bgs_bits) != 0
    select = isTGT & isGAL & isBGS

    # Accumulate spectrum data.
    if allzbest is None:
        allzbest = zbest[select]
        allfmap = fibermap[select]
        allwave = pspectra.wave['brz']
        allflux = pspectra.flux['brz'][select]
        allivar = pspectra.ivar['brz'][select]
        allmask = pspectra.mask['brz'][select]
        allres  = pspectra.resolution_data['brz'][select]
    else:
        allzbest = vstack([allzbest, zbest[select]])
        allfmap = vstack([allfmap, fibermap[select]])
        allflux = np.vstack([allflux, pspectra.flux['brz'][select]])
        allivar = np.vstack([allivar, pspectra.ivar['brz'][select]])
        allmask = np.vstack([allmask, pspectra.mask['brz'][select]])
        allres  = np.vstack([allres, pspectra.resolution_data['brz'][select]])
# Apply the DESITRIP preprocessing to selected spectra.
rewave, reflux, reivar = rebin_flux(allwave, allflux, allivar, allzbest['Z'],
                                    minwave=2500., maxwave=9500., nbins=150,
                                    log=True, clip=True)
rsflux = rescale_flux(reflux)
# Run the classifier on the spectra.
# The output layer uses softmax activation to produce an array of label probabilities.
# The classification is based on argmax(pred).
pred = classifier.predict(rsflux)
pred.shape
ymax = np.max(pred, axis=1)
fig, ax = plt.subplots(1,1, figsize=(6,4), tight_layout=True)
ax.hist(ymax, bins=np.linspace(0,1,51))
ax.set(xlabel='$\max{(y_\mathrm{pred})}$',
       ylabel='count',
       title='Tile {}, {}'.format(tile_id, obsdate));<jupyter_output><empty_output><jupyter_text>Selection on Classifier OutputTo be conservative we can select only spectra where the classifier is very confident in its output, e.g., ymax > 0.99. See the [CNN training notebook](https://github.com/desihub/timedomain/blob/master/desitrip/docs/nb/cnn_multilabel-restframe.ipynb) for the motivation behind this cut.<jupyter_code>idx = np.argwhere(ymax > 0.99)
labels = np.argmax(pred, axis=1)
idx.shape
label_names = ['Galaxy',
          'SN Ia',
          'SN Ib',
          'SN Ib/c',
          'SN Ic',
          'SN IIn',
          'SN IIL/P',
          'SN IIP',
          'KN']
# Randomly select 16 of the spectra and plot them.
#selection = sorted(np.random.choice(idx.flatten(), size=idx.shape[0], replace=False))
selection=idx<jupyter_output><empty_output><jupyter_text>GradCAM action happens hereAdapting from https://keras.io/examples/vision/grad_cam/<jupyter_code>import tensorflow as tf
last_conv_layer_name = "conv1d_23"
classifier_layer_names = [
"batch_normalization_23",
"activation_23",
"max_pooling1d_23",
"flatten_5",
"dense_5",
"dropout_5",
"Output_Classes"
]
def make_gradcam_heatmap(
    img_array, model, last_conv_layer_name, classifier_layer_names
):
    # First, we create a model that maps the input image to the activations
    # of the last conv layer
    last_conv_layer = model.get_layer(last_conv_layer_name)
    last_conv_layer_model = keras.Model(model.inputs, last_conv_layer.output)

    # Second, we create a model that maps the activations of the last conv
    # layer to the final class predictions
    classifier_input = keras.Input(shape=last_conv_layer.output.shape[1:])
    x = classifier_input
    for layer_name in classifier_layer_names:
        #print(layer_name,x.shape)
        x = model.get_layer(layer_name)(x)
    classifier_model = keras.Model(classifier_input, x)

    # Then, we compute the gradient of the top predicted class for our input image
    # with respect to the activations of the last conv layer
    with tf.GradientTape() as tape:
        # Compute activations of the last conv layer and make the tape watch it
        last_conv_layer_output = last_conv_layer_model(img_array)
        tape.watch(last_conv_layer_output)
        # Compute class predictions
        preds = classifier_model(last_conv_layer_output)
        top_pred_index = tf.argmax(preds[0])
        top_class_channel = preds[:, top_pred_index]

    # This is the gradient of the top predicted class with regard to
    # the output feature map of the last conv layer
    grads = tape.gradient(top_class_channel, last_conv_layer_output)
    # This is a vector where each entry is the mean intensity of the gradient
    # over a specific feature map channel
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1))
    #print(grads.shape,pooled_grads.shape)

    # We multiply each channel in the feature map array
    # by "how important this channel is" with regard to the top predicted class
    last_conv_layer_output = last_conv_layer_output.numpy()[0]
    pooled_grads = pooled_grads.numpy()
    for i in range(pooled_grads.shape[-1]):
        last_conv_layer_output[:, i] *= pooled_grads[i]

    # The channel-wise mean of the resulting feature map
    # is our heatmap of class activation
    heatmap = np.mean(last_conv_layer_output, axis=-1)

    #We apply ReLU here and select only elements>0
    # For visualization purpose, we will also normalize the heatmap between 0 & 1
    heatmap = np.maximum(heatmap, 0) / np.max(heatmap)
    return heatmap<jupyter_output><empty_output><jupyter_text>Apply GradCAM to one spectrum only<jupyter_code># Prepare image
preprocess_input = keras.applications.xception.preprocess_input
decode_predictions = keras.applications.xception.decode_predictions
#myarr=rsflux[:1,:] 
myarr=rsflux[selection[1],:] 

# Print what the top predicted class is
preds = classifier.predict(myarr)
#print("Predicted:", preds)

# Generate class activation heatmap
heatmap = make_gradcam_heatmap(
    myarr, classifier, last_conv_layer_name, classifier_layer_names
)


color='blue'
rewave_nbin_inblock=rewave.shape[0]/float(heatmap.shape[0])
first_bin=0
for i in range(1,heatmap.shape[0]+1):
    alpha=np.min([1,heatmap[i-1]+0.2])
    last_bin=int(i*rewave_nbin_inblock)
    plt.plot(rewave[first_bin:last_bin+1], myarr[0,first_bin:last_bin+1],c=color,alpha=alpha)
    first_bin=last_bin<jupyter_output><empty_output><jupyter_text>Apply GradCAM to all spectra classified as transients<jupyter_code>preprocess_input = keras.applications.xception.preprocess_input
decode_predictions = keras.applications.xception.decode_predictions

fig, axes = plt.subplots(4,4, figsize=(15,10), sharex=True, sharey=True,
                         gridspec_kw={'wspace':0, 'hspace':0})

for j, ax in zip(selection, axes.flatten()):
    myarr=rsflux[j,:] 

    # Print what the top predicted class is
    preds = classifier.predict(myarr)
    #print("Predicted:", preds)

    # Generate class activation heatmap
    heatmap = make_gradcam_heatmap(
        myarr, classifier, last_conv_layer_name, classifier_layer_names
    )

    color='blue'
    rewave_nbin_inblock=rewave.shape[0]/float(heatmap.shape[0])
    first_bin=0
    for i in range(1,heatmap.shape[0]+1):
        alpha=np.min([1,heatmap[i-1]+0.2])
        last_bin=int(i*rewave_nbin_inblock)
        if (i==1):
            ax.plot(rewave[first_bin:last_bin+1], myarr[0,first_bin:last_bin+1],c=color,alpha=alpha,\
                    label=label_names[labels[j[0]]]+'\nz={:.2f}'.format(allzbest[j[0]]['Z']))
        else:
            ax.plot(rewave[first_bin:last_bin+1], myarr[0,first_bin:last_bin+1],c=color,alpha=alpha)
        first_bin=last_bin
    ax.legend(fontsize=10)<jupyter_output><empty_output><jupyter_text>Plot spectra of objects classified as transientsPlot observed spectra<jupyter_code>testwave, testflux, testivar = rebin_flux(allwave, allflux, allivar,
                                    minwave=2500., maxwave=9500., nbins=150,
                                    log=True, clip=True)
fig, axes = plt.subplots(4,4, figsize=(15,10), sharex=True, sharey=True,
                         gridspec_kw={'wspace':0, 'hspace':0})

for j, ax in zip(selection, axes.flatten()):
    ax.plot(testwave, testflux[j[0]], alpha=0.7, label='label: '+label_names[labels[j[0]]]+'\nz={:.2f}'.format(allzbest[j[0]]['Z']))
    ax.set(xlim=(3500,9900),ylim=(-0.1,4))
    ax.fill_between([5600,6000],[-0.1,-0.1],[4,4],alpha=0.1,color='blue')
    ax.fill_between([7400,7800],[-0.1,-0.1],[4,4],alpha=0.1,color='blue')
    ax.legend(fontsize=10)

# for k in [0,1,2]:
#     axes[k,0].set(ylabel=r'flux [erg s$^{-1}$ cm$^{-1}$ $\AA^{-1}$]')
#     axes[2,k].set(xlabel=r'$\lambda_\mathrm{obs}$ [$\AA$]',
#                   xlim=(3500,9900))
fig.tight_layout();<jupyter_output><empty_output>