<jupyter_start><jupyter_text>This notebook demonstrates estimation of the marginal posteriors using MomentNetworks  Load modules<jupyter_code>!pip install emcee getdist
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import numpy as np
from scipy.stats import norm
from scipy.stats import multivariate_normal
from scipy import fftpack

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, LeakyReLU, Input
from tensorflow.keras.layers import concatenate, BatchNormalization
from tensorflow.keras.models import Model
from tensorflow.keras import optimizers
import os, sys
import random
import emcee as mc
import getdist
from getdist import plots, MCSamples
import scipy.interpolate as interpolate
import time
%matplotlib inline<jupyter_output><empty_output><jupyter_text>Generate high dim training data<jupyter_code>size=20
n_training = 100000
cov_signal_fourier = np.identity(size)*np.logspace(size,size*0.9,size)/np.sum(np.logspace(size,size*0.9,size))
cov_noise = np.identity(size)*0.1
cov_noise_inv = np.linalg.pinv(cov_noise)
training_z = np.random.multivariate_normal(np.zeros(size),cov_signal_fourier,n_training) + \
                  1j*np.random.multivariate_normal(np.zeros(size), cov_signal_fourier,n_training)
training_z = np.fft.fft(training_z).real

training_data = training_z + np.random.multivariate_normal(np.zeros(size),cov_noise,n_training)


training_z_obs = training_z[0]
training_data_obs = training_data[0]
training_z = training_z[1:]
training_data = training_data[1:]
cov_signal = np.identity(size)*0.

for i in range(n_training-1):
  cov_signal += np.outer(training_z[i],training_z[i]) 

cov_signal = cov_signal/np.float(n_training-1.)

cov_signal_inv = np.linalg.pinv(cov_signal)
plt.imshow(cov_signal_inv)
plt.plot(training_z_obs)
plt.scatter(np.arange(size),training_data_obs)<jupyter_output><empty_output><jupyter_text>Moment network variance estimation<jupyter_code>class simple_leaky:
    """
    A simple MLP with LeakyReLU activation
    """
    
    def __init__(self, input_size, output_size, learning_rate=None):
        """
        Initialisation
        :param map_size: size of square image (there are map_size**2 pixels)
        :param learning_rate: learning rate for the optimizer
        """
        self.input_size = input_size
        self.learning_rate = learning_rate
        self.output_size = output_size
        
    def model(self):
        print(self.input_size)
        
        input_data = (Input(shape=(self.input_size,)))

        x1 = Dense(100, input_dim=self.input_size, kernel_initializer='normal')(input_data)
        x1 = LeakyReLU(alpha=0.1)(x1)
        x2 = Dense(150, kernel_initializer='normal')(x1)
        x2 = LeakyReLU(alpha=0.1)(x2)
        x3 = Dense(100, kernel_initializer='normal')(x2)
        x3 = LeakyReLU(alpha=0.1)(x3)
        x4 = Dense(self.output_size, kernel_initializer='normal')(x3)
        

        dense_model = Model(input_data, x4)
        dense_model.summary()

        if self.learning_rate is None:
            dense_model.compile(optimizer='adam', loss='mse')
        else:
            dense_model.compile(optimizer=optimizers.Adam(lr=self.learning_rate), loss='mse')

        return dense_model
model_instance = simple_leaky(size,size, learning_rate=1e-4)
regression = model_instance.model()
history = regression.fit(training_data[:80000],
                               training_z[:80000],
                               epochs =20,batch_size =50, shuffle=True,
                               validation_data = (training_data[80000:],
                               training_z[80000:]))
_ = plt.figure(figsize = (14,6))
_ = plt.subplot(121), plt.plot(history.history['val_loss'], marker = 'o', label = 'validation')
_ = plt.plot(history.history['loss'], marker = 'x', label = 'training'), plt.legend()

_ = plt.subplot(122), plt.xscale('log')
_ = plt.ylim(0.1*np.min(history.history['val_loss']), 4*np.max(history.history['loss']))
_ = plt.plot(history.history['val_loss'], marker = 'o', label = 'validation')
_ = plt.plot(history.history['loss'], marker = 'x', label = 'training'), plt.legend(), plt.ylim(3e-4,5e-4)
predicted_mean = regression.predict(np.atleast_2d(training_data))
training_var_unknown_mean = (training_z-regression.predict(np.atleast_2d(training_data)))**2.
model_instance = simple_leaky(size,size, learning_rate=1e-4)
regression_var_unknown_mean = model_instance.model()
history = regression_var_unknown_mean.fit(training_data[:80000],
                               training_var_unknown_mean[:80000],
                               epochs =20,batch_size =50, shuffle=True,
                               validation_data = (training_data[80000:],
                               training_var_unknown_mean[80000:]))
predicted_var_obs = (regression_var_unknown_mean.predict(np.atleast_2d(training_data_obs[:]))[0])
print(predicted_var_obs)
(predicted_var_obs)<jupyter_output><empty_output><jupyter_text>Sample true posterior<jupyter_code>def log_prior(theta, data,cov_signal_inv):
    return -0.5*np.inner(theta,np.inner(cov_signal_inv,theta))

def log_likelihood(theta, data,cov_noise_inv):
    return -0.5*np.inner(theta-data,np.inner(cov_noise_inv,theta-data))

def log_posterior_delfi(theta, data,cov_signal_inv):
    return log_prior(theta, data,cov_signal_inv) + DelfiEnsemble.log_likelihood_stacked(theta,data=data) 

def log_posterior_likelihood(theta, data,cov_signal_inv,cov_noise_inv):
    return log_prior(theta, data,cov_signal_inv)  + log_likelihood(theta, data,cov_noise_inv)

def initial_parameters(theta, relative_sigma):
    """
    This is not randomise the initial position of the
    :param theta: list/array of parameter values
    :param relative_sigma: controls variance of random draws
    :return: the theta array but with random shifts
    """
    theta = np.array(theta)
    return np.random.normal(theta, np.abs(theta * relative_sigma))
n_dim=size
n_burn=100
n_steps = 100000
n_walkers = 2* n_dim



sampler_likelihood = mc.EnsembleSampler(n_walkers, n_dim, log_posterior_likelihood,
                             args=(training_data_obs,cov_signal_inv,cov_noise_inv,))
theta_expected_mean = [0.1]*size

theta0 = np.array([list(initial_parameters(theta_expected_mean, 1)) for i in range(n_walkers)])
t = time.time()
_ =sampler_likelihood.run_mcmc(theta0, n_steps) 

print(time.time() - t)

flat_samples_likelihood = sampler_likelihood.get_chain(discard=n_burn, flat=True)
print(flat_samples_likelihood.shape)
np.var(flat_samples_likelihood,axis=0)
names = ["x%s"%i for i in range(size)]
labels =  ["x_%s"%i for i in range(size)]
samples2 = MCSamples(samples=flat_samples_likelihood,names = names, labels = labels, label='lik',
               settings={'smooth_scale_2D':1})

g = plots.get_subplot_plotter()
g.triangle_plot([samples2], filled=True)<jupyter_output>Removed no burn in<jupyter_text>Compare MomentNetwork variance with known MCMC  $\Delta = \frac{1}{N} \sum_i^N \Big[ \frac{\mathrm{Var}_{MC} - \mathrm{Var}_{pred}}{\mathrm{Var}_{MC}} \Big]_{\theta_i}$<jupyter_code>_ = plt.scatter(np.var(flat_samples_likelihood,axis=0), predicted_var_obs)
delta_stat = np.mean(1.- predicted_var_obs/np.var(flat_samples_likelihood,axis=0))
print(delta_stat)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11743723418850323, 'fraction_numerical': 0.02099053368088901, 'mean_word_length': 4.590490797546012, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '183751', 'license': 'MIT', 'n_tokens_mistral': 2597, 'n_tokens_neox': 2480, 'n_words': 418, 'path': 'MomentNetwork_demo/dimensional_challenges_demo/delfi_moment_network.ipynb', 'repo_name': 'Pablo-Lemos/MomentNetworks', 'size': '7289'}