<jupyter_start><jupyter_text>Run section 1 and 2 to prepare data and load model.Section 3 does the loglik estimation and the confidence intervals. There may be OOM issues when dealing with the long test files. If so, then the loglik will need to be computed in batches (not shown here).<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
import sklearn
import time
from pickle import load
import tensorflow as tf
from tensorflow import keras
if not tf.config.list_physical_devices('GPU'):
    print("No GPU was detected. LSTMs and CNNs can be very slow without a GPU.")

tf.random.set_seed(42)
K = keras.backend

from sklearn.preprocessing import StandardScaler,PowerTransformer
import math
from scipy.stats import multivariate_normal
from scipy.stats import norm
import os
import pandas as pd
import pickle
from helper import *
K.set_floatx("float32")<jupyter_output>No GPU was detected. LSTMs and CNNs can be very slow without a GPU.<jupyter_text>Run the below if using a GPU.<jupyter_code>#using laptop gpu
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], enable=True)
k = 8
J = 32
save_time_step = 0.005
h=1 
c=10
b=10<jupyter_output><empty_output><jupyter_text>1. Prepare data<jupyter_code>history_length = 100
test_seq_length = 1000
#load train, validation and test datasets 
train_dataset = np.load("../../data/truth_run/training_dataset.npy")
valid_dataset = np.load("../../data/truth_run/val_dataset.npy")
test_dataset   = np.load("../../data/truth_run/climate_eval_dataset.npy")
test_dataset_23 = np.load("../../data/truth_run/climate_change_exp/full_test_set.npy")
x_train = train_dataset[:,:,0]
x_valid = valid_dataset[:,:,0]
x_test = test_dataset[:,:,0]
x_test_23 = test_dataset_23[:,:,0]
##### Functions to work out the exact U for each x #########

def _rhs_X_dt(X, F,U,dt=0.005):
    """Compute the right hand side of the X-ODE."""

    dXdt = (-np.roll(X, 1,axis=1) * (np.roll(X, 2,axis=1) - np.roll(X, -1,axis=1)) -
                X + F - U)

    return dt * dXdt 


def U(Xt,Xt_1,F,dt=0.005):
    k1_X = _rhs_X_dt(Xt, F,U=0)
    k2_X = _rhs_X_dt(Xt + k1_X / 2, F,U=0)
    Xt_1_pred = k2_X + Xt 
    #print(Xt_1_pred)
    Ut = (Xt_1_pred - Xt_1 )/dt

    return Ut
u_t = U(x_train[:-1,:],x_train[1:,:],20)    
u_t_valid = U(x_valid[:-1,:],x_valid[1:,:],20)  
u_t_test = U(x_test[:-1,:],x_test[1:,:],20)
u_t_test_23 = U(x_test_23[:-1,:],x_test_23[1:,:],23)

training_dataset = np.stack([x_train[:-1],u_t],axis=2)
valid_dataset = np.stack([x_valid[:-1],u_t_valid],axis=2)
test_dataset =  np.stack([x_test[:-1],u_t_test],axis=2)
test_dataset_23 =  np.stack([x_test_23[:-1],u_t_test_23],axis=2)
def prepare_datasets_for_RNN(dataset,history_length):
    max_index = (dataset.shape[0]-1)//history_length
    dataset = dataset[:(max_index*history_length +1),:,:] 
    dataset_shape = dataset.shape[0]
    last_elements = dataset[-1,:,:]
    remaining_dataset = dataset[:-1,:,:]
    reshaped = remaining_dataset.reshape(-1,history_length,k,2)
    add_on = reshaped[1:,:1,:,:]
    last_elements = last_elements.reshape(1,1,k,2)
    add_on_combined = np.concatenate((add_on,last_elements),axis=0)
    concat = np.concatenate((reshaped,add_on_combined),axis=1)
    concat = concat.transpose((2,0,1,3)).reshape((-1,history_length+1,2),order="F")
    return concat.astype("float32")
def prep_for_loglik_estimation(dataset,history_length,exp_its,conf_its):
    dataset = dataset.astype("float32")
    max_index = (dataset.shape[0]-1)//history_length
    dataset = dataset[:(max_index*history_length +1),:,:] 
    dataset_shape = dataset.shape[0]
    last_elements = dataset[-1,:,:]
    remaining_dataset = dataset[:-1,:,:]
    reshaped = remaining_dataset.reshape(-1,history_length,k,2)
    add_on = reshaped[1:,:1,:,:]
    last_elements = last_elements.reshape(1,1,k,2)
    add_on_combined = np.concatenate((add_on,last_elements),axis=0)
    concat = np.concatenate((reshaped,add_on_combined),axis=1)
    new = np.repeat(concat.transpose((0,2,1,3)),exp_its,axis=0).reshape(-1,history_length+1,2)
    new = np.tile(new,(conf_its,1,1))
    return new.astype("float32")

train_nn_features = prepare_datasets_for_RNN(training_dataset,history_length)
valid_nn_features = prepare_datasets_for_RNN(valid_dataset,test_seq_length)
test_for_loglik = prep_for_loglik_estimation(test_dataset[:],6000,1,1)
test_for_loglik_23 = prep_for_loglik_estimation(test_dataset_23[:],6000,1,1)
test_for_loglik_23.shape
x_mean = np.mean(train_nn_features[:,:,0])
x_std = np.std(train_nn_features[:,:,0])
u_mean = np.mean(train_nn_features[:,:,1])
u_std = np.std(train_nn_features[:,:,1])
#scaling
train_nn_features[:,:,0] = (train_nn_features[:,:,0] - x_mean)/x_std
train_nn_features[:,:,1] = (train_nn_features[:,:,1] - u_mean)/u_std
#scaling
valid_nn_features[:,:,0] = (valid_nn_features[:,:,0] - x_mean)/x_std
valid_nn_features[:,:,1] = (valid_nn_features[:,:,1] - u_mean)/u_std
#scaling
test_for_loglik[:,:,0] = (test_for_loglik[:,:,0] - x_mean)/x_std
test_for_loglik[:,:,1] = (test_for_loglik[:,:,1] - u_mean)/u_std
#scaling
test_for_loglik_23[:,:,0] = (test_for_loglik_23[:,:,0] - x_mean)/x_std
test_for_loglik_23[:,:,1] = (test_for_loglik_23[:,:,1] - u_mean)/u_std
test_for_loglik_23_tf = tf.convert_to_tensor(test_for_loglik_23)
test_for_loglik_tf = tf.convert_to_tensor(test_for_loglik)<jupyter_output><empty_output><jupyter_text>2. Load model<jupyter_code>class Sampling(keras.layers.Layer):
    def call(self, inputs):
        mean, log_var = inputs
        return K.random_normal(tf.shape(log_var)) * K.exp(log_var / 2) + mean

z_shape = 34
hidden_state_size_bi = 32
encoder_hidden_state_size=32

generator = keras.models.load_model("gan_generator_final.h5")
h_encoder = keras.models.load_model("h_encoder_final.h5",custom_objects={
           "Sampling":Sampling})
bi_rnn = keras.models.load_model("bi_rnn_final.h5")
h_encoder_first = keras.models.load_model("h_encoder_first_final.h5",custom_objects={
           "Sampling":Sampling})

def sample_from_encoder(xu_seq,encoder,first_encoder,encoder_hidden_state_size,bi_rnn):
    
    length = xu_seq.shape[1]
    batch_shape = xu_seq.shape[0]
    
    h_sequence = tf.TensorArray(dtype=tf.float32,size=0,dynamic_size=True)
    h_mean_out = tf.TensorArray(dtype=tf.float32,size=0,dynamic_size=True)
    h_log_var_out = tf.TensorArray(dtype=tf.float32,size=0,dynamic_size=True)
    
    u_summary = bi_rnn(xu_seq[:,:-1,:])

    h_mean1,h_log_var1,h_prev = first_encoder.predict([xu_seq[:,0,0],u_summary[:,0,0]])
    h_sequence = h_sequence.write(0,h_prev)
    h_mean_out = h_mean_out.write(0,h_mean1)
    h_log_var_out = h_log_var_out.write(0,h_log_var1)

    
    hidden_state_1 = tf.zeros(shape=(batch_shape,encoder_hidden_state_size))
    hidden_state_2 = tf.zeros(shape=(batch_shape,encoder_hidden_state_size))    
    
    for n in tf.range(0,length-2):
        h_mean,h_log_var,h_sample,state,state2 = encoder.predict([h_prev,u_summary[:,n+1:n+2,:],xu_seq[:,n+1:n+2,:1],
                                                    hidden_state_1,hidden_state_2])
        
        h_sequence = h_sequence.write(n+1,h_sample)
        h_prev = h_sample
        h_mean_out = h_mean_out.write(n+1,h_mean)
        h_log_var_out = h_log_var_out.write(n+1,h_log_var) 
        hidden_state_1 = state  
        hidden_state_2 = state2     
        
    h_sequence = h_sequence.stack()        
    h_mean_out_enc = h_mean_out.stack()
    h_log_var_out = h_log_var_out.stack()
    h_sequence = tf.transpose(h_sequence[:,:,0,:],[1,0,2])            
    h_mean_out_enc = tf.transpose(h_mean_out_enc[:,:,0,:],[1,0,2])
    h_log_var_out = tf.transpose(h_log_var_out[:,:,0,:],[1,0,2])

    return h_sequence,h_mean_out_enc,h_log_var_out<jupyter_output><empty_output><jupyter_text>3. Loglik estimations<jupyter_code>def eval_loglik_gaussian_u_cond_h(xu_seq,h_encoding,sigma):
    x_array =  xu_seq[:,:-1,:1]
    x_array_reshape = tf.reshape(x_array,(-1,1))
    u_array_reshape = tf.reshape(xu_seq[:,:-1,1:2],(-1,1))
    h_encoding_reshape = tf.reshape(h_encoding,(-1,z_shape))
    mean_u = generator([x_array_reshape,h_encoding_reshape])
    term = -K.log((sigma**2) *2*math.pi) - tf.math.divide((u_array_reshape-mean_u),sigma)**2
    loglik = 0.5*term
    loglik = tf.reshape(loglik,(-1,xu_seq.shape[1]-1,1))
    return tf.reduce_sum(loglik,axis=1) #take sum over time
    
def eval_loglik_gaussian_h_encoder(h_encoding,h_mean,h_logvar):
    term1 = -(1/2)*(K.log(2*math.pi) + h_logvar)
    term2 = -((h_encoding-h_mean)**2)/(2*K.exp(h_logvar))
    loglik = term1+term2
    loglik = tf.reduce_sum(loglik,axis=[2],keepdims=True) #sum over the h dimensions 
    return tf.reduce_sum(loglik,axis=1) #sum over time 
def eval_loglik_gaussian_h_gan(h_encoding):
    #### term 1 ####
    ### h drawn from normal(0,1) ####
    term = tf.reduce_sum(0.5*(-K.log((1**2) *2*math.pi) - tf.math.divide((h_encoding[:,:1,:]),1)**2),axis=2) #sum over z dimensions

    #### term 2 #####
    #### loglik for the rest of the markovian seq ####
    array = h_encoding[:,1:,:]
    phi = 0.7486
    mean = h_encoding[:,:-1,:]*phi
    sigma = (1-phi**2)**0.5
    
    term2 = 0.5*(-K.log((sigma**2) *2*math.pi) - tf.math.divide((array-mean),sigma)**2)
    term2 = tf.reduce_sum(term2,axis=[2]) #sum over the h dimensions 
    
    loglik_array = tf.concat([term,term2 ],axis=1)
    loglik = tf.reduce_sum(loglik_array,axis=1,keepdims=True) #sum over t
                          
    return loglik
def loglik_exact_approx(testing_sequence,sigma,it):

    """
    
    it is number for approx
    
    """
    
    combined = tf.zeros(shape=(int(testing_sequence.shape[0]/k),1))
    
    print("expectation_outer_loop")
    
    for n in tf.range(it):
        
        print(n)
    
        h_sequence,h_mean_out_enc,z_log_var_out = sample_from_encoder(testing_sequence,h_encoder,h_encoder_first,
                                                    encoder_hidden_state_size,bi_rnn)

        print("generated")
        exponent = eval_loglik_gaussian_u_cond_h(testing_sequence,h_sequence,sigma) +  eval_loglik_gaussian_h_gan(h_sequence) \
          - eval_loglik_gaussian_h_encoder(h_sequence,h_mean_out_enc,z_log_var_out)

        reshape_exponent = tf.reshape(exponent,(-1,k))
        exponent_summed = tf.reduce_sum(reshape_exponent,axis=1)
        
        reshaped = tf.reshape(exponent_summed,(-1,1)) 
        
        combined = tf.concat([combined,reshaped],axis=1)
        
    result = combined[:,1:]
    
    most_positive_alpha = np.max(result,axis=1)
    remaining_factor = result[:,:] - most_positive_alpha.reshape(-1,1)[:,:]
    second_positive = np.sort(remaining_factor,axis=1)[:,-2]

    x = np.exp(second_positive.astype("float64"))
    term3 = x - (x**2)/2    
    
    loglik = -np.log(float(result.shape[1])) + most_positive_alpha +term3
    
    avg_loglik = np.sum(loglik) / ((testing_sequence.shape[0])*(testing_sequence.shape[1]-1))
    
    print(avg_loglik)
    
    return avg_loglik
def confidence_interval_arrays(num_its_conf,testing_sequence,sigma,it_mc):
    
    #testing_sequence must have testing_sequence.shape[0] be a multiple of K
    
    array = np.zeros(shape=(num_its_conf,1))
    
    for i in tf.range(num_its_conf):
        
        print("confidence iteration",i)
        
        loglik = loglik_exact_approx(testing_sequence,sigma,it_mc)
        array[i,0] = loglik
    return array

def rx_star(array):
    return np.random.choice(array,size=len(array))<jupyter_output><empty_output><jupyter_text>F = 20<jupyter_code>a = confidence_interval_arrays(50,test_for_loglik_tf,0.001,50)
# np.save("loglik_20.npy",a)
a = np.load("loglik_20.npy")
dt = 0.005

a_with_correction = a -np.log(dt*u_std)<jupyter_output><empty_output><jupyter_text>Note the correction is explained in the Appendix of the paper. We need to add a scaling factor to get loglik of x sequence. And just like for the RNN, we also include u_std here as we scale the variables when training the models (and just as how the factor dt comes from a change of variables for likelihood, so too does the u_std factor as we change from the variable U_scaled back to U_unscaled).<jupyter_code>a_sample = [np.mean(rx_star(a_with_correction.reshape(-1))) for _ in range(10000)]
np.mean(a_with_correction)
lo, hi = np.quantile(a_sample,[0.025,0.975])
print(lo)
print(hi)<jupyter_output>-37.69651479547297
-37.690796590341705<jupyter_text>F = 23<jupyter_code>a = confidence_interval_arrays(50,test_for_loglik_23_tf,0.001,50)
# np.save("loglik_23.npy",a)
a = np.load("loglik_23.npy")
dt = 0.005

a_with_correction = a -np.log(dt*u_std)
a_sample = [np.mean(rx_star(a_with_correction.reshape(-1))) for _ in range(10000)]
np.mean(a_with_correction)
lo, hi = np.quantile(a_sample,[0.025,0.975])
print(lo)
print(hi)<jupyter_output>-44.86401038957839
-44.85502655197918
{'fraction_non_alphanumeric': 0.1465503815605046, 'fraction_numerical': 0.038389658931630585, 'mean_word_length': 4.569384215091067, 'pattern_counts': {'":': 2, '<': 30, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '70626', 'license': 'MIT', 'n_tokens_mistral': 5291, 'n_tokens_neox': 5098, 'n_words': 746, 'path': 'saved_models/gan/importance_sampler_for_gan_results.ipynb', 'repo_name': 'raghul-parthipan/l96_rnn', 'size': '12842'}