# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [default]
#     language: python
#     name: python2
# ---

# %load_ext autoreload
# %autoreload 2

from pearce.emulator import OriginalRecipe, ExtraCrispy
from pearce.mocks import cat_dict
import numpy as np
from os import path

import tensorflow as tf

import matplotlib
#matplotlib.use('Agg')
from matplotlib import pyplot as plt
# %matplotlib inline
import seaborn as sns
sns.set()

# The files include:
#
# cosmology_camb.dat : the input training cosmology, only 5 parameters: Om, Ob, sigma_8, h, n_s.    w and N_eff are not used here because the analytic method is only for LCDM.
#
# HOD_design_m_4_n_400_tinker.dat : 400 HOD designs for the training set.
#
# EH_test_200COSMO_tinker.dat : the 200 test cosmologies from Tinker.
#
# EH_test_200COSMO_tinker.dat : the 1000 test HODs, just use the first 200.
#
# Cosmo_err.dat : the fractional error of wp estimated from the test boxes.
#
# wp_clustering_emu: folder contains the wp data for training, the columns are rp, wp
#
# test_200COSMO_tinker_wp_clustering_emu: folder contains the wp data for test, same format as training set.
#
# example.py: is an example script for my GP modeling. you should fill out the missing places. My comment on line 31 may not be right: because 0-49 for line 1 and 50-99 for line 2 etc will result repeated HOD sampling for different cosmologies, 400/50=8<40, so a better choice might be just randomly choose 50 HOD for each cosmology. (edited)

dir = '/home/sean/Downloads/Zhongxu_data/for_Sean/'
cosmo_data_fname = 'EH_test_200COSMO_tinker.dat'
hod_data_fname = 'GP_test_HOD_1000.dat'

from os.path import join

cosmo_colnames = ['Om', 'Ob', 'sigma_8', 'h', 'n_s']
cosmo_data = np.loadtxt(join(dir, cosmo_data_fname), delimiter=' ')

hod_colnames = ['M1', 'alpha', 'Mmin', 'sigma_logM']
hod_data = np.loadtxt(join(dir, hod_data_fname), delimiter = ' ')

training_file = '/home/sean/PearceRedMagicXiCosmoFixedNd.hdf5'
#test_file = '/home/sean/PearceRedMagicXiCosmoTest.hdf5'

# + active=""
# training_file = '/home/sean/PearceRedMagicXiCosmo.hdf5'
# test_file = '/home/sean/PearceRedMagicXiCosmoTest.hdf5'
# -

em_method = 'nn'
split_method = 'random'

a = 1.0
z = 1.0/a - 1.0

# + active=""
# emu.scale_bin_centers
# -

fixed_params = {'z':z}#, 'r':17.0389993 }

# + active=""
# n_leaves, n_overlap = 50, 1
# emu = ExtraCrispy(training_file, n_leaves, n_overlap, split_method, method = em_method, fixed_params=fixed_params,
#                  custom_mean_function = None, downsample_factor = 1.0)
# -

emu = OriginalRecipe(training_file, method = em_method, fixed_params=fixed_params,
                    hyperparams = {'hidden_layer_sizes': (10),
                                 'activation': 'relu', 'verbose': True, 
                                    'tol': 1e-8, 'learning_rate_init':1e-4,\
                                   'max_iter':10, 'alpha':0, 'early_stopping':False, 'validation_fraction':0.3})

# + active=""
# #convert zhongxu's data to my format
# emu.get_param_names()

# + active=""
# my_cosmo_data = np.zeros((cosmo_data.shape[0], 7))
# my_hod_data = np.zeros((200, 4))
#
# my_cosmo_data[:,0] = cosmo_data[:,1]*(cosmo_data[:,3])**2
# my_cosmo_data[:,1] = cosmo_data[:,0]*(cosmo_data[:,3])**2 - my_cosmo_data[:,0]
# my_cosmo_data[:,2] = -1.0
# my_cosmo_data[:,3] = cosmo_data[:,4]
# #my_cosmo_data[:,4]
# my_cosmo_data[:, 5] = cosmo_data[:,3]*100
# my_cosmo_data[:, 6] = 3.046

# + active=""
# from classy import Class
# cosmo = Class()
#
# for i, row in enumerate(cosmo_data):
#     Om, Ob, sigma_8, h, n_s = row
#     params = {
#             'output': 'mPk',
#             'sigma8': sigma_8,
#             'n_s': n_s,
#             'h': h,
#             'non linear': 'halofit',
#             'omega_b': Ob*h*h,
#             'omega_cdm': (Om-Ob)*h*h,
#             'z_pk': 0.0}
#
#
#     cosmo.set(params)
#     cosmo.compute()
#     #print cosmo.pm
#     val = cosmo.get_current_derived_parameters(['ln10^{10}A_s'])['ln10^{10}A_s']
#     my_cosmo_data[i,4] = val

# + active=""
# my_hod_data[:,0] = hod_data[:200,3]
# my_hod_data[:,1] = np.log10(hod_data[:200,2])
# my_hod_data[:,2] = np.log10(hod_data[:200,0])
# my_hod_data[:,3] = hod_data[:200,1]

# +
clustering_dir = 'test_200COSMO_tinker_wp_clustering_emu/'
from glob import glob

clustering_files = sorted(glob(join(dir, clustering_dir) + '*') )
# -

nbins = 9
zx = np.zeros((len(clustering_files)*nbins, 12))
zy = np.zeros((len(clustering_files)*nbins,))

# +
for i, cf in enumerate(clustering_files):
    if i%1000==0:
        print i
    data = np.loadtxt(cf, delimiter = ' ')
    rs = np.log10(data[:,0])
    wp = np.log10(data[:,1])
    fbase = cf.split('/')[-1]
    split_fbase = fbase.split('_')
    cosmo, hod = int(split_fbase[1]), int(split_fbase[3])
    
    zx[i*nbins:(i+1)*nbins, :7] = my_cosmo_data[cosmo]
    zx[i*nbins:(i+1)*nbins, 7:-1] = my_hod_data[hod]
    zx[i*nbins:(i+1)*nbins, -1] = rs
    zy[i*nbins:(i+1)*nbins] = wp
    
    
# -

np.savetxt('zx.npy', zx)
np.savetxt('zy.npy', zy)

zx = np.loadtxt('zx.npy')
zy = np.loadtxt('zy.npy')

# +
idxs = np.random.choice(emu.x.shape[0], size = int(emu.x.shape[0]*1.0), replace=False)

x_train, y_train,yerr_train = emu.x[idxs, :],emu.y[idxs],emu.yerr[idxs]

y_train = y_train*(emu._y_std + 1e-5) + emu._y_mean
yerr_train = yerr_train*(emu._y_std+1e-5)
# -

idxs

len(emu.get_param_names())

unique_cosmos = np.unique(x_train[:, :7], axis =0)#*(emu._x_std[:7]+1e-5) + emu._x_mean[:7]

unique_cosmos.shape

left_out_cosmo = unique_cosmos[0]
is_loc = np.all(x_train[:,:7] == left_out_cosmo, axis = 1)
x_test = x_train[is_loc]
x_train = x_train[~is_loc]
y_test = y_train[is_loc]
y_train = y_train[~is_loc]
yerr_test = yerr_train[is_loc]
yerr_train = yerr_train[~is_loc]


# + active=""
# x_test, y_test, ycov_test, _ = emu.get_data(test_file, fixed_params, None, False)
# x_test = (x_test - emu._x_mean)/(emu._x_std+1e-5)
#
# #split_ycov = np.dsplit(ycov_test, ycov_test.shape[-1])
#         #fullcov = block_diag(*[yc[:,:,0] for yc in split_ycov])
# #yerr_test = np.sqrt(np.hstack(np.diag(syc[:,:,0]) for syc in split_ycov))

# + active=""
# from sklearn.model_selection import train_test_split
# x_train, x_test, y_train, y_test, yerr_train, _ = train_test_split(x_train, y_train,yerr_train, test_size = 0.3, shuffle = True)

# + active=""
# pnames = emu.get_param_names()
# for i in xrange(x_train.shape[1]):
#     for j in xrange(i+1,x_train.shape[1]):
#         plt.scatter(x_train[:,i], x_train[:,j])
#         plt.scatter(x_test[:,i], x_test[:,j])
#         plt.title('%s vs %s'%(pnames[i], pnames[j]))
#         plt.show();

# + active=""
# plt.plot(x_np[:emu.n_bins, -1:], y_np[:emu.n_bins])
# -

def n_layer_fc(x, hidden_sizes, training=False, l = 1e-8):
    initializer = tf.variance_scaling_initializer(scale=2.0)
    regularizer = tf.contrib.layers.l2_regularizer(l)
    fc_output = tf.layers.dense(x, hidden_sizes[0], activation=tf.nn.relu,
                                 kernel_initializer = initializer, kernel_regularizer = regularizer)
                                 #kernel_regularizer = tf.nn.l2_loss)
    #fc2_output = tf.layers.dense(fc1_output, hidden_sizes[1], activation=tf.nn.relu,
    #                             kernel_initializer = initializer, kernel_regularizer = regularizer)
    for size in hidden_sizes[1:]:
        fc_output = tf.layers.dense(fc_output, size, activation=tf.nn.relu, kernel_initializer=initializer,
                                 kernel_regularizer = regularizer)
    pred = tf.layers.dense(fc_output, 1, kernel_initializer=initializer, 
                              kernel_regularizer = regularizer)[:,0]#,
    return pred


def novel_fc(x, hidden_sizes, training=False, l = (1e-6, 1e-6, 1e-6), p = (0.5, 0.5, 0.5),\
             n_cosmo_params = 7, n_hod_params = 4):
    
    cosmo_sizes, hod_sizes, cap_sizes = hidden_sizes
    
    if type(l) is float:
        cosmo_l, hod_l, cap_l = l, l, l
    else:
        cosmo_l, hod_l, cap_l = l
        
    if type(p) is float:
        cosmo_p, hod_p, cap_p = p,p,p
    else:
        cosmo_p, hod_p, cap_p = p
    
    initializer = tf.variance_scaling_initializer(scale=2.0)
    
    #onlly for duplicating r
    n_params = n_cosmo_params+n_hod_params
    cosmo_x = tf.slice(x, [0,0], [-1, n_cosmo_params])
    cosmo_x = tf.concat(values=[cosmo_x, tf.slice(x, [0, n_params-1], [-1, -1]) ], axis = 1)
    #print tf.shape(cosmo_x)
    #print tf.shape(tf.slice(x, [0, n_params-1], [-1, -1]))
    hod_x = tf.slice(x, [0, n_cosmo_params], [-1, -1])
    
    cosmo_regularizer = tf.contrib.layers.l2_regularizer(cosmo_l)
    cosmo_out = cosmo_x
    
    for size in cosmo_sizes:
        fc_output = tf.layers.dense(cosmo_out, size,
                                 kernel_initializer = initializer,\
                                    kernel_regularizer = cosmo_regularizer)
        bd_out = tf.layers.dropout(fc_output, cosmo_p, training = training)
        bn_out = tf.layers.batch_normalization(bd_out, axis = -1, training=training)
        cosmo_out = tf.nn.relu(bn_out)#tf.nn.leaky_relu(bn_out, alpha=0.01)
    
    hod_regularizer = tf.contrib.layers.l1_regularizer(hod_l)
    hod_out = hod_x
    
    for size in hod_sizes:
        fc_output = tf.layers.dense(hod_out, size,
                                 kernel_initializer = initializer,\
                                    kernel_regularizer = hod_regularizer)
        bd_out = tf.layers.dropout(fc_output, hod_p, training = training)
        bn_out = tf.layers.batch_normalization(bd_out, axis = -1, training=training)
        hod_out = tf.nn.relu(bn_out)#tf.nn.leaky_relu(bn_out, alpha=0.01)
    
    cap_out=tf.concat(values=[cosmo_out, hod_out], axis = 1)
    
    return cap_out


def pretrain_cap(cap_input, hidden_sizes, training=False, l = (1e-6, 1e-6, 1e-6), p = (0.5, 0.5, 0.5)):
    initializer = tf.variance_scaling_initializer(scale=2.0)

    cosmo_sizes, hod_sizes, cap_sizes = hidden_sizes
    
    if type(l) is float:
        cosmo_l, hod_l, cap_l = l, l, l
    else:
        cosmo_l, hod_l, cap_l = l
        
    if type(p) is float:
        cosmo_p, hod_p, cap_p = p,p,p
    else:
        cosmo_p, hod_p, cap_p = p
    
    cap_out=cap_input
    cap_regularizer = tf.contrib.layers.l2_regularizer(cap_l)

    for size in cap_sizes:
        fc_output = tf.layers.dense(cap_out, size,
                                 kernel_initializer = initializer,\
                                    kernel_regularizer = cap_regularizer)
        bd_out = tf.layers.dropout(fc_output, cap_p, training = training)
        bn_out = tf.layers.batch_normalization(bd_out, axis = -1, training=training)
        cap_out = tf.nn.relu(bn_out)#tf.nn.leaky_relu(bn_out, alpha=0.01)
    
    pred = tf.layers.dense(cap_out, 1, kernel_initializer=initializer, 
                              kernel_regularizer = cap_regularizer)[:,0]#,
    return pred


def final_cap(cap_input, hidden_sizes, training=False, l = (1e-6, 1e-6, 1e-6), p = (0.5, 0.5, 0.5)):
    initializer = tf.variance_scaling_initializer(scale=2.0)

    cosmo_sizes, hod_sizes, cap_sizes = hidden_sizes
    
    if type(l) is float:
        cosmo_l, hod_l, cap_l = l, l, l
    else:
        cosmo_l, hod_l, cap_l = l
        
    if type(p) is float:
        cosmo_p, hod_p, cap_p = p,p,p
    else:
        cosmo_p, hod_p, cap_p = p
    
    cap_out=cap_input
    cap_regularizer = tf.contrib.layers.l2_regularizer(cap_l)

    for size in cap_sizes:
        fc_output = tf.layers.dense(cap_out, size,
                                 kernel_initializer = initializer,\
                                    kernel_regularizer = cap_regularizer)
        bd_out = tf.layers.dropout(fc_output, cap_p, training = training)
        bn_out = tf.layers.batch_normalization(bd_out, axis = -1, training=training)
        cap_out = tf.nn.relu(bn_out)#tf.nn.leaky_relu(bn_out, alpha=0.01)
    
    pred = tf.layers.dense(cap_out, 1, kernel_initializer=initializer, 
                              kernel_regularizer = cap_regularizer)[:,0]#,
    return pred


def optimizer_init_fn(learning_rate = 1e-7):
    return tf.train.AdamOptimizer(learning_rate)#, beta1=0.9, beta2=0.999, epsilon=1e-6)


from sklearn.metrics import r2_score, mean_squared_error


def check_accuracy(sess, val_data,batch_size, x, weights, preds, is_training=None):
    val_x, val_y = val_data
    perc_acc, scores = [],[]
    for idx in xrange(0, val_x.shape[0], batch_size):
        feed_dict = {x: val_x[idx:idx+batch_size],
                     is_training: 0}
        y_pred = sess.run(preds, feed_dict=feed_dict)
        #print y_pred.shape, val_y[idx:idx+batch_size].shape
        score = r2_score(val_y[idx:idx+batch_size], y_pred)
        scores.append(score)
        
        perc_acc = np.mean(emu._y_std*np.abs(val_y[idx:idx+batch_size]-y_pred)/np.abs(emu._y_std*val_y[idx:idx+batch_size] + emu._y_mean) )

        
    print 'Val score: %.6f, %.2f %% Loss'%(np.mean(np.array(scores)), 100*np.mean(np.array(perc_acc)))

device = '/device:GPU:0'
#device = '/cpu:0'
def train(model_init_fn, optimizer_init_fn,num_params, pretrain_data, train_data, val_data, hidden_sizes,\
               num_pretrain_epochs = 500, num_epochs=1000, batch_size = 200, l = 1e-6, p = 0.5, print_every=10):
    tf.reset_default_graph()
    pretrain = True
    with tf.device(device):
        # Construct the computational graph we will use to train the model. We
        # use the model_init_fn to construct the model, declare placeholders for
        # the data and labels
        x = tf.placeholder(tf.float32, [None,num_params])
        y = tf.placeholder(tf.float32, [None])
        weights = tf.placeholder(tf.float32, [None])
        
        is_training = tf.placeholder(tf.bool, name='is_training')
        
        cap_input = model_init_fn(x, hidden_sizes, is_training, l = l, p=p)
        
        if pretrain:
            preds = pretrain_cap(cap_input, hidden_sizes, is_training, l=l, p=p)
        else:
            preds = final_cap(cap_input, hidden_sizes, is_training, l=l, p=p)

        # Compute the loss like we did in Part II
        #loss = tf.reduce_mean(loss)
        
    with tf.device('/cpu:0'):
        loss = tf.losses.mean_squared_error(labels=y,\
                   predictions=preds, weights = weights)#weights?
        #loss = tf.losses.absolute_difference(labels=y, predictions=preds, weights = tf.abs(1.0/y))#weights?
        pass
    with tf.device(device):
        optimizer = optimizer_init_fn()
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            train_op = optimizer.minimize(loss)
        
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        #t = 0
        pretrain_x, pretrain_y = pretrain_data
        rand_idxs = range(pretrain_x.shape[0])
        for epoch in range(num_pretrain_epochs):
            #print('Starting epoch %d' % epoch)
            np.random.shuffle(rand_idxs)
            losses = []
            for idx in xrange(0, pretrain_x.shape[0], batch_size):
                feed_dict = {x: pretrain_x[rand_idxs[idx:idx+batch_size]],\
                             y: pretrain_y[rand_idxs[idx:idx+batch_size]],\
                             weights: np.ones_like(pretrain_y[rand_idxs[idx:idx+batch_size]]),\
                             is_training:1}
                loss_np, _ = sess.run([loss, train_op], feed_dict=feed_dict)
                losses.append(loss_np)
                
            if epoch % print_every == 0:
                loss_avg = np.mean(np.array(losses))
                print('Epoch %d, loss = %e' % (epoch, loss_avg))
                check_accuracy(sess, val_data, batch_size, x, weights, preds, is_training=is_training)
        
        pretrain = False 
        train_x, train_y, train_yerr = train_data
        rand_idxs = range(train_x.shape[0])
        for epoch in range(num_epochs):
            #print('Starting epoch %d' % epoch)
            
            
            np.random.shuffle(rand_idxs)
            losses = []
            for idx in xrange(0, train_x.shape[0], batch_size):
                yerrbatch = train_yerr[rand_idxs[idx:idx+batch_size]]
                _bs = yerrbatch.shape[0]
                feed_dict = {x: train_x[rand_idxs[idx:idx+batch_size]],\
                             y: train_y[rand_idxs[idx:idx+batch_size]] + np.random.randn(_bs)*yerrbatch,\
                             weights: 1/yerrbatch,\
                             is_training:1}
                loss_np, _ = sess.run([loss, train_op,], feed_dict=feed_dict)
                losses.append(loss_np)
                
            if epoch % print_every == 0:
                loss_avg = np.mean(np.array(losses))
                print('Epoch %d, loss = %e' % (epoch, loss_avg))
                check_accuracy(sess, val_data, batch_size, x, weights, preds, is_training=is_training)
            #t += 1

train(novel_fc, optimizer_init_fn, x_train.shape[1],\
           (zx, zy), (x_train, y_train, yerr_train), (x_test, y_test),\
           [(100,100), (200,100,200), (500,100)], num_pretrain_epochs = 500,  num_epochs= int(1e3),\
           batch_size = 200, l = (1e-6, 1e-8, 1e-8), p = (0.333, 0.1, 0.1),\
           print_every = 100)

np.abs(emu.goodness_of_fit(training_file, statistic = 'log_frac')).mean()

np.abs(emu.goodness_of_fit(training_file, statistic = 'frac')).mean()

fit_idxs = np.argsort(gof.mean(axis = 1))

emu.goodness_of_fit(training_file).mean()#, statistic = 'log_frac')).mean()

model = emu._emulator

ypred = model.predict(emu.x)

plt.hist( np.log10( (emu._y_std+1e-5)*np.abs(ypred-emu.y)/np.abs((emu._y_std+1e-5)*emu.y+emu._y_mean) ))

( (emu._y_std+1e-5)*np.abs(ypred-emu.y)/np.abs((emu._y_std+1e-5)*emu.y+emu._y_mean) ).mean()

emu._y_mean, emu._y_std

for idx in fit_idxs[:10]:
    print gof[idx].mean()
    print (ypred[idx*emu.n_bins:(idx+1)*emu.n_bins]-emu.y[idx*emu.n_bins:(idx+1)*emu.n_bins])/emu.y[idx*emu.n_bins:(idx+1)*emu.n_bins]
    plt.plot(emu.scale_bin_centers, ypred[idx*emu.n_bins:(idx+1)*emu.n_bins], label = 'Emu')
    plt.plot(emu.scale_bin_centers, emu.y[idx*emu.n_bins:(idx+1)*emu.n_bins], label = 'True')
    plt.legend(loc='best')
    plt.xscale('log')
    plt.show()

print dict(zip(emu.get_param_names(), emu.x[8*emu.n_bins, :]*emu._x_std+emu._x_mean))

emu.get_param_names()

# + active=""
# #print emu.x.shape
# #print emu.downsample_x.shape
# if hasattr(emu, "_emulators"):
#     print emu._emulators[0]._x.shape
# else:
#     print emu._emulator._x.shape
# -

emu._ordered_params

# + active=""
# x, y, y_pred = emu.goodness_of_fit(training_file, statistic = 'log_frac')

# + active=""
# x, y, y_pred

# + active=""
# N = 25
# for _y, yp in zip(y[:N], y_pred[:N]):
#     #plt.plot(emu.scale_bin_centers ,  (_y - yp)/yp ,alpha = 0.3, color = 'b')
#     
#     plt.plot(emu.scale_bin_centers, _y, alpha = 0.3, color = 'b')
#     plt.plot(emu.scale_bin_centers, yp, alpha = 0.3, color = 'r')
#     
# plt.loglog();

# + active=""
# %%timeit
# #truth_file = '/u/ki/swmclau2/des/PearceRedMagicWpCosmoTest.hdf5'
# gof = emu.goodness_of_fit(training_file, N = 100, statistic = 'log_frac')
# -

gof = emu.goodness_of_fit(training_file, statistic = 'frac')
print gof.mean()

for row in gof:
    print row

gof = emu.goodness_of_fit(training_file, statistic = 'frac')
print gof.mean()

model = emu._emulator

model.score(emu.x, emu.y)

# +
ypred = model.predict(emu.x)

np.mean(np.abs(ypred-emu.y)/emu.y)

# +
plt.plot(emu.scale_bin_centers, np.abs(gof.mean(axis = 0)) )
plt.plot(emu.scale_bin_centers, np.ones_like(emu.scale_bin_centers)*0.01)
plt.plot(emu.scale_bin_centers, np.ones_like(emu.scale_bin_centers)*0.05)
plt.plot(emu.scale_bin_centers, np.ones_like(emu.scale_bin_centers)*0.1)


plt.loglog();
# -

plt.plot(emu.scale_bin_centers, np.abs(gof.T),alpha = 0.1, color = 'b')
plt.plot(emu.scale_bin_centers, np.ones_like(emu.scale_bin_centers)*0.01, lw = 2, color = 'k')
plt.loglog();

gof[:,i].shape

# +
for i in xrange(gof.shape[1]):
    plt.hist(np.log10(gof[:, i]), label = str(i), alpha = 0.2);
    
plt.legend(loc='best')
plt.show()
# -



StarcoderdataJupyterScriptsDedupFiltered