<jupyter_start><jupyter_text>SphereNN Interpretability Analysis  Loading dependencies, data and model<jupyter_code>%load_ext autoreload
%autoreload 2
import os
import pickle
import numpy as np
import healpy as hp
import matplotlib.pyplot as plt
import matplotlib as mpl
import tensorflow as tf
import tensorflow.keras as keras
from deep_dss.helpers import *
import warnings

warnings.filterwarnings(action='once')

# Run on GPU.
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Defaults
plt.rcParams.update({
'lines.color':'black',
'font.family':'serif',
'font.weight':'normal',
'text.color':'black',
'text.usetex': True,
'axes.edgecolor':'black',
'axes.linewidth':1.0,
'axes.titlesize':'x-large',
'axes.labelsize':'x-large',
'axes.labelcolor':'black',
'xtick.labelsize':'x-large',
'xtick.minor.width':1.0,
'xtick.major.width':1.0,
'ytick.major.size':7,
'ytick.minor.size':4,
'ytick.major.pad':8,
'ytick.minor.pad':8,
'ytick.labelsize':'x-large',
'ytick.minor.width':1.0,
'ytick.major.width':1.0,
'legend.numpoints':1,
'legend.fontsize':'x-large',
'legend.shadow':False,
'legend.frameon':False})
config = "k"
channels = 1
noiseless_m = False
noiseless_kg = False
free_bias = False
gaussian = False

order = 2
nside = 1024

val_set = "TEST"

exp_name = "spherenn-v3-simple-noisy-fixed-kappa"
num_id = 2
checkpoint_path = "../checkpoints/spherenn/{0}/{0}-{1}".format(exp_name, num_id)
checkpoint_dir = "../checkpoints/spherenn/{0}".format(exp_name)
# !mkdir $checkpoint_dir
log_dir = "../log/{0}".format(exp_name)
# !mkdir $log_dir
fig_dir = "../figures/{0}".format(exp_name)
# !mkdir $fig_dir
def num_cosmologies(dataset):
    if dataset == "TRAINLITE":
        return 16
    if dataset == "TESTLITE":
        return 4
    if dataset == "TEST":
        return 21
    return 45

def generate_reshaped_data(dataset):
    print("Generating data for {0}!".format(dataset))
    num_cosmos = num_cosmologies(dataset)
    data = split_count_and_lensing_maps_by_dataset(dataset, config=config, noiseless_m=noiseless_m,
                                                  noiseless_kg=noiseless_kg, free_bias=free_bias, gaussian=gaussian)
    data["x"] = np.reshape(data["x"], (12*(order**2)*num_cosmos, (nside//order)**2, channels))
    data["y"] = np.reshape(data["y"], (12*(order**2)*num_cosmos, 1, 1))
    return data
def build_model():
    return keras.Sequential([
        keras.Input(shape=(262144, channels)),
        keras.layers.Conv1D(64, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(128, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(256, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(256, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(256, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(256, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(256, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(256, 4, strides=4, activation='relu'),
        keras.layers.Conv1D(1, 4, strides=4, activation='relu'),
    ])

def load_model():
    model = build_model()
    model.load_weights(checkpoint_path+"-Q4")
    return model
model = load_model()
model.summary()
data = generate_reshaped_data("TEST")<jupyter_output>Generating data for TEST!<jupyter_text>Gradient Calculation: Example<jupyter_code>ex_input_batch = data["x"][:32]
ex_input_batch.shape
ex_label = data["y"][:32]
ex_label.shape
ex_input_tensor = tf.convert_to_tensor(ex_input_batch, dtype=tf.float32)
ex_label_tensor = tf.convert_to_tensor(ex_label, dtype=tf.float32)
loss_function = tf.keras.losses.MeanAbsoluteError()
with tf.GradientTape() as tape:
    tape.watch(ex_input_tensor)
    output = model(ex_input_tensor)
    loss = loss_function(ex_label_tensor, output)
ex_grads = tape.gradient(loss, ex_input_tensor)
ex_grads
def zscale(inputs):
    return (inputs - np.mean(inputs))/np.std(inputs)
ex_zscaled_grads = zscale(ex_grads)
ex_zscaled_grads
def full_map(partial_map):
    fmap = np.zeros(hp.nside2npix(nside))
    fmap[:(nside//order)**2] = np.copy(partial_map)
    return fmap
hp.mollview(full_map(ex_input_batch[0,:,0]), nest=True, title="Example Input Map, $\sigma_8=1.144$")
hp.mollview(full_map(ex_zscaled_grads[0,:,0]), nest=True, 
            title=r"Example zscaled gradients, $\frac{\partial NN}{\partial x}$, $\sigma_8=1.144$")
ex_input_x_grads = np.multiply(ex_input_batch, ex_grads)
ex_zscaled_input_x_grads = zscale(ex_input_x_grads)
hp.mollview(full_map(ex_zscaled_input_x_grads[0,:,0]), nest=True,
           title=r"Example zscaled input x gradients, $\frac{\partial NN}{\partial x} \cdot x$, $\sigma_8=1.144$")<jupyter_output>/Users/adiraju/opt/anaconda3/envs/deep_dss/lib/python3.8/site-packages/healpy/projaxes.py:920: MatplotlibDeprecationWarning: You are modifying the state of a globally registered colormap. In future versions, you will not be able to modify a registered colormap in-place. To remove this warning, you can make a copy of the colormap first. cmap = copy.copy(mpl.cm.get_cmap("viridis"))
  newcm.set_over(newcm(1.0))
/Users/adiraju/opt/anaconda3/envs/deep_dss/lib/python3.8/site-packages/healpy/projaxes.py:202: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  aximg = self.imshow(<jupyter_text>Binned saliency scores: example<jupyter_code>ex_zscaled_input_map = zscale(ex_input_batch[0])
ex_input_map = ex_input_batch[0]
ex_zscaled_input_map
ex_zscaled_grads_map = ex_zscaled_grads[0,:,0].numpy()
ex_zscaled_grads_map = np.reshape(ex_zscaled_grads_map, ex_zscaled_input_map.shape)
ex_zscaled_grads_map
ex_zscaled_input_x_grads_map = ex_zscaled_input_x_grads[0,:,0]
ex_zscaled_input_x_grads_map = np.reshape(ex_zscaled_input_x_grads_map, ex_zscaled_input_map.shape)
ex_zscaled_input_x_grads_map
from astropy.visualization import hist
fig, ax = plt.subplots()
hist(ex_zscaled_input_map, bins="scott", ax=ax, histtype="stepfilled", alpha=0.8, density=True)
# ax.set_xlim([-2.5, 2.5])
ax.set_xlabel("Convergence, $\kappa$")
ax.set_ylabel("Probability Density")
ax.set_title("Pixel Histogram $\sigma_8=1.144$")
plt.tight_layout()
fig.show()
fig, ax = plt.subplots()
hist(ex_zscaled_grads_map, bins="scott", ax=ax, histtype="stepfilled", alpha=0.8, density=True)
# ax.set_xlim([-2.5, 2.5])
ax.set_xlabel(r"Gradient, $\frac{\partial NN}{\partial x}$")
ax.set_ylabel("Probability Density")
ax.set_title("Gradient Histogram $\sigma_8=1.144$")
plt.tight_layout()
fig.show()
fig, ax = plt.subplots()
hist(ex_zscaled_input_x_grads_map, bins="scott", ax=ax, histtype="stepfilled", alpha=0.8, density=True)
# ax.set_xlim([-2.5, 2.5])
ax.set_xlabel(r"Input x Gradient, $\frac{\partial NN}{\partial x} \cdot x$")
ax.set_ylabel("Probability Density")
ax.set_title("Input x Gradient Histogram $\sigma_8=1.144$")
plt.tight_layout()
from scipy import stats
stats.describe(ex_zscaled_input_map)
x = np.ones(262144)
x.shape[0]
def square_sum(x):
    return np.sum(np.multiply(x, x))/x.shape[0]
bin_ss_grads = stats.binned_statistic(np.reshape(ex_zscaled_input_map, (262144)),
                       np.reshape(ex_zscaled_grads_map, (262144)), statistic=square_sum, bins=40, range=(-20,20))
bin_mean_grads = stats.binned_statistic(np.reshape(ex_zscaled_input_map, (262144)),
                       np.reshape(ex_zscaled_grads_map, (262144)), statistic='mean', bins=40, range=(-20,20))
bin_ss_input_x_grads = stats.binned_statistic(np.reshape(ex_zscaled_input_map, (262144)),
                       np.reshape(ex_zscaled_input_x_grads_map, (262144)), 
                                           statistic=square_sum, bins=40, range=(-20,20))
bin_mean_input_x_grads = stats.binned_statistic(np.reshape(ex_zscaled_input_map, (262144)),
                       np.reshape(ex_zscaled_input_x_grads_map, (262144)), 
                                           statistic='mean', bins=40, range=(-20,20))
bin_ss_grads
bin_mean_grads
bin_ss_input_x_grads
bin_mean_input_x_grads
plt.plot(np.linspace(-20, 20, num=40), bin_ss_grads.statistic)
# plt.xticks(np.arange(-20, 20))
plt.xlabel("Binned r.m.s. convergence, $\kappa$")
plt.ylabel("Average squared scaled gradient")
plt.title("Average squared scaled gradient by binned r.m.s. $\kappa$")
plt.plot(np.linspace(-20, 20, num=40), bin_ss_input_x_grads.statistic)
# plt.xticks(np.arange(-20, 20))
plt.xlabel("Binned r.m.s. convergence, $\kappa$")
plt.ylabel("Average squared scaled input x gradient")
plt.title("Average squared scaled input x gradient by binned r.m.s. $\kappa$")
plt.plot(np.linspace(-20, 20, num=40), bin_mean_grads.statistic)
# plt.xticks(np.arange(-20, 20))
plt.xlabel("Binned r.m.s. convergence, $\kappa$")
plt.ylabel("Mean scaled gradient")
plt.title("Mean scaled gradient by binned r.m.s. $\kappa$")
plt.plot(np.linspace(-20, 20, num=40), bin_mean_input_x_grads.statistic)
# plt.xticks(np.arange(-20, 20))
plt.xlabel("Binned r.m.s. convergence, $\kappa$")
plt.ylabel("Mean scaled input x gradient")
plt.title("Mean scaled input x gradient by binned r.m.s. $\kappa$")
stats.describe(ex_input_map)
np.mean(ex_input_map)
np.std(ex_input_map)<jupyter_output><empty_output><jupyter_text>Two-point correlation saliency scores: example<jupyter_code>ex_input_map
ex_zscaled_input_map
ex_zscaled_grads_map
ex_zscaled_input_x_grads_map
import treecorr
pixel_dec = np.pi/2 - hp.pix2ang(nside, 
                                 np.arange(hp.nside2npix(nside)))[0]
pixel_ra = hp.pix2ang(nside, np.arange(hp.nside2npix(nside)))[1]
empty_map = np.zeros(hp.nside2npix(nside))
ex_map_ring = hp.reorder(full_map(ex_input_map[:,0]), n2r=True)
ex_map_catalogue = treecorr.Catalog(ra=pixel_ra, dec=pixel_dec, ra_units='rad', dec_units='rad',
                                      k=ex_map_ring)
sampled = np.zeros((nside//order)**2)
import random
hp.nside2resol(nside, arcmin=True)
hp.nside2resol(order, arcmin=True)
def new_sample():
    pixel_index = random.randint(0, (nside//order)**2 - 1)
    while sampled[pixel_index] == 1:
        pixel_index = random.randint(0, (nside//order)**2 - 1)
    sampled[pixel_index] = 1
    print("Sampling index ", pixel_index)
    
    empty_map[hp.nest2ring(nside, pixel_index)] = ex_input_map[pixel_index, 0]
    print("Generating catalogue")
    empty_map_catalogue = treecorr.Catalog(ra=pixel_ra, dec=pixel_dec, ra_units='rad', dec_units='rad',
                                      k=empty_map)
    
    print("Computing correlation")
    kk = treecorr.KKCorrelation(min_sep=5, max_sep=1750, nbins=120, sep_units='arcmin')
    kk.process(ex_map_catalogue, empty_map_catalogue)
    kk.write("2pcf.txt")
new_sample()
with open("dummy.txt", "r") as f:
    xi = float(f.readlines()[-1].split()[3])
hp.query_disc(1024, hp.pix2vec(1024, hp.nest2ring(1024, 135521)), 5*np.pi/(60*180)).shape
def two_pcf_saliency_estimate(index, radius, bin_width=3.5):
    pixel_set = hp.query_disc(nside, hp.pix2vec(nside, hp.nest2ring(nside, index)), radius*np.pi/(60*180)).shape
    small_set = hp.query_disc(nside, hp.pix2vec(nside, hp.nest2ring(nside, index)),
                              (radius-bin_width)*np.pi/(60*180)).shape
    annular_set = np.setdiff1d(pixel_set, small_set)
    nnxi = ex_grads.numpy()[0,index,0]
    fxi = np.sum(full_map(np.reshape(ex_input_map, (262144)))[annular_set])
    if fxi == 0:
        return np.inf
    return (nside//order)**2 * nnxi / fxi
two_pcf_saliency_estimate(135531, 10)
plt.plot(np.arange(100), np.geomspace(5, 1750, num=100), 'b-', np.arange(100), np.linspace(5, 1750, num=100), 'r-')
%timeit two_pcf_saliency_estimate(100, 1500)
# Linearly spaced bin values

num_lin_bins = 150
pixel_freq = 1000
lin_saliencies = np.zeros((num_lin_bins, (nside//order)**2))

period = 10000

for n, r in zip(np.arange(num_lin_bins), np.linspace(5, 1750, num=num_lin_bins)):
    for i in range(0, (nside//order)**2, pixel_freq):
        if i % period == 0: print("n: ", n, "r: ", r, "i: ", i, end="\r")
        lin_saliencies[n, i] = two_pcf_saliency_estimate(i, r)
valid_saliencies = []
for i in range(num_lin_bins):
    valid_saliencies.append([])
    for j in np.where(lin_saliencies[i] != np.inf)[0]:
        valid_saliencies[i].append(lin_saliencies[i,j])
    valid_saliencies[i] = np.array(valid_saliencies[i])
valid_saliencies
mean_valid_saliencies = [np.mean(valid_saliencies[i]) for i in range(num_lin_bins)]
plt.plot(np.linspace(5, 1750, num=num_lin_bins), mean_valid_saliencies)
plt.xlabel("Scale, r (arcmin)")
plt.ylabel("Saliency, S")
plt.title(r"Mean single-pixel saliency estimates for $\xi(r)$, $N=262$")
for n, r in zip(np.arange(num_lin_bins), np.linspace(5, 1750, num=num_lin_bins)):
    print(n, r)
stats.describe(valid_saliencies[49])
fig, ax = plt.subplots()
# hist(lin_saliencies[49,:], bins="scott", ax=ax, histtype="stepfilled", density=True)
ax.hist(valid_saliencies[49], bins=200)
# ax.set_xlim([-2.5, 2.5])
ax.set_xlabel("Saliency, S")
ax.set_ylabel("Probability Density")
ax.set_title(r"Single-Pixel Saliency Histogram $\sigma_8=1.144$, $n=1$, $r=578.86$")
from scipy.special import comb

average_saliencies = []

k = 262
n = 50
m = 10000
M = comb(k, n)
m = min(m, M)

for i in range(m):
    pixel_set = np.random.choice(valid_saliencies[49], (n), replace=False)
    average_saliencies.append(np.mean(pixel_set))
average_saliencies
fig, ax = plt.subplots()
# hist(lin_saliencies[49,:], bins="scott", ax=ax, histtype="stepfilled", density=True)
ax.hist(average_saliencies, bins=200)
# ax.set_xlim([-2.5, 2.5])
ax.set_xlabel("Saliency, S")
ax.set_ylabel("Probability Density")
ax.set_title(r"Multi-Pixel Saliency Histogram $\sigma_8=1.144$, $n=20$, $r=578.86$")
fig, ax = plt.subplots()
# hist(lin_saliencies[49,:], bins="scott", ax=ax, histtype="stepfilled", density=True)
ax.hist(average_saliencies, bins=200)
# ax.set_xlim([-2.5, 2.5])
ax.set_xlabel("Saliency, S")
ax.set_ylabel("Probability Density")
ax.set_title(r"Multi-Pixel Saliency Histogram $\sigma_8=1.144$, $n=10$, $r=578.86$")
def bin_score_averages(b, n):
    average_saliencies = []

    k = 262
    m = 10000
    M = comb(k, n)
    m = min(m, M)

    for i in range(m):
        pixel_set = np.random.choice(lin_saliencies[b,:], (n), replace=False)
        average_saliencies.append(np.mean(pixel_set))
        
    print(stats.describe(average_saliencies))
    
    fig, ax = plt.subplots()
    # hist(lin_saliencies[49,:], bins="scott", ax=ax, histtype="stepfilled", density=True)
    ax.hist(average_saliencies, bins=200)
    # ax.set_xlim([-2.5, 2.5])
    ax.set_xlabel("Saliency, S")
    ax.set_ylabel("Probability Density")
    ax.set_title(r"Multi-Pixel Saliency Histogram $\sigma_8=1.144$, $n={0}$, $r={1}$".format(n, 
                                                                                             np.linspace(5, 1750, num=num_lin_bins)[b]))
    fig.show()
bin_score_averages(20, 20)
bin_score_averages(109, 2)
stats.describe(lin_saliencies[109,:])
lin_saliencies[109,:]
test_saliencies = np.zeros((1260, 150, 262144))
data["x"]
def two_pcf_saliency_estimate(index, radius, bin_width=3.5):
    pixel_set = hp.query_disc(nside, hp.pix2vec(nside, hp.nest2ring(nside, index)), radius*np.pi/(60*180)).shape
    small_set = hp.query_disc(nside, hp.pix2vec(nside, hp.nest2ring(nside, index)),
                              (radius-bin_width)*np.pi/(60*180)).shape
    annular_set = np.setdiff1d(pixel_set, small_set)
    nnxi = ex_grads.numpy()[0,index,0]
    fxi = ex_input_map[index] * np.sum(full_map(np.reshape(ex_input_map, (262144)))[annular_set])
    if fxi == 0:
        return np.inf
    return (nside//order)**2 * nnxi / fxi
for i in range(21):
    print("On cosmology", i)
    input_batch = tf.convert_to_tensor(data["x"][i*48:(i+1)*48], dtype=tf.float32)
    label_batch = tf.convert_to_tensor(data["y"][i*48:(i+1)*48], dtype=tf.float32)
    
    print("Computing gradients")
    with tf.GradientTape() as tape:
        tape.watch(input_batch)
        output = model(input_batch)
        loss = loss_function(label_batch, output)
        
    grads = tape.gradient(loss, input_batch)
    
    for j in range(48):
        print("On map", j)
        for n, r in zip(np.arange(num_lin_bins), np.linspace(5, 1750, num=num_lin_bins)):
            for k in range(0, (nside//order)**2, pixel_freq):
                if k % period == 0: print("n: ", n, "r: ", r, "k: ", k, end="\r")
                pixel_set = hp.query_disc(nside, hp.pix2vec(nside, hp.nest2ring(nside, k)), r*np.pi/(60*180)).shape
                small_set = hp.query_disc(nside, hp.pix2vec(nside, hp.nest2ring(nside, k)),
                                          (r-3.5)*np.pi/(60*180)).shape
                annular_set = np.setdiff1d(pixel_set, small_set)
                nnxi = grads.numpy()[i*48+j,k,0]
                fxi = data["x"][i*48+j,k,0] * np.sum(full_map(np.reshape(data["x"][i*48+j], (262144)))[annular_set])
                if fxi == 0:
                    test_saliencies[i*48+j,n,k] = np.inf
                else:
                    test_saliencies[i*48+j,n,k] = (nside//order)**2 * nnxi / fxi<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1569913889051608, 'fraction_numerical': 0.04071231913771894, 'mean_word_length': 4.4042418486862935, 'pattern_counts': {'":': 3, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 15, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '215312', 'license': 'MIT', 'n_tokens_mistral': 7200, 'n_tokens_neox': 6597, 'n_words': 1110, 'path': 'notebooks/.ipynb_checkpoints/Interpretability Analysis v1.0-checkpoint.ipynb', 'repo_name': 'adiraju21s/deep_dss', 'size': '17071'}