<jupyter_start><jupyter_text>Formulating Problem<jupyter_code># Create a model
model = pints.toy.LogisticModel()

n_parameters = model.n_parameters()
n_outputs = model.n_outputs()
real_parameters = [0.15, 500]

values, times, noise = emutils.simulate(model,
                                        parameters = real_parameters)

# Create an object with links to the model and time series
problem = pints.SingleOutputProblem(model, times, values)

# Create a log-likelihood function (adds an extra parameter!)
real_log_likelihood = pints.KnownNoiseLogLikelihood(problem, noise)

print("Number of parameters: ", n_parameters)
print("Parameter values: ", real_parameters)
print("Example problem values:\n", problem.values()[:5])<jupyter_output>Number of parameters:  2
Parameter values:  [0.15, 500]
Example problem values:
 [14.36723662 -0.73010642 19.80292564 42.90095131  0.91011587]<jupyter_text>Creating and normalizing training data<jupyter_code># creating distribution on bounds
real_params_lower = [0.1, 400]
real_params_upper = [0.2, 600]

bounds = pints.RectangularBoundaries(lower = real_params_lower, upper = real_params_upper)
log_prior = pints.UniformLogPrior(bounds)
# make training sample
train_size = 5000
train_input = log_prior.sample(train_size)
train_target = np.apply_along_axis(real_log_likelihood, 1, train_input)
# normalize data between -0.5 and 0.5
def normalize(x):
    mean, std = np.mean(x, axis=0), np.std(x, axis=0)
    return (x - mean) / std

def denormalize(y, mean, std):
    return y * std + mean
    
train_input_normalized = normalize(train_input)
train_target_normalized = normalize(train_target)

# define functions for future use
train_input_mean = np.mean(train_input, axis=0)
train_input_std = np.std(train_input, axis=0)

train_target_mean = np.mean(train_target, axis=0)
train_target_std = np.std(train_target, axis=0)

def normalize_input(x):
    return (x - train_input_mean) / train_input_std

def denormalize_prediction(x_norm):
    return denormalize(x_norm, train_target_mean, train_target_std)
train_target_mean<jupyter_output><empty_output><jupyter_text>Bayesian NN<jupyter_code>tf.Session().run(tf.sqrt())
# network parameters
bnn_shape = [n_parameters, 5, 5, 1]
activations = tf.nn.relu

# variational inference parameters
n_iter=2000
n_samples=10
class BNNEmulator(emupints.Emulator):
    
    def __init__(self, log_likelihood, X, y, **kwargs):
        super(BNNEmulator, self).__init__(log_likelihood, X, y, **kwargs)
        
        self.graph = tf.Graph()
        self.sess = tf.Session(graph=self.graph)

        
    def __call__(self, x):
        # convert to np array
        if type(x) != np.ndarray:
            x = np.asarray(x)
        # reshape input to fit placeholder
        n_params = self.n_parameters()
        x = x.reshape((x.size // n_params, n_params))
        
        feed_dict = {self._x : x}
        
        ans = self.sess.run(self._nns, feed_dict=feed_dict)

        # for call return an average of predictions
        return np.mean(ans)


    def predict(self, x):
        # convert to np array
        if type(x) != np.ndarray:
            x = np.asarray(x)
        n_params = self.n_parameters()
        x = x.reshape((x.size // n_params, n_params))

        feed_dict = {self._x : x}
        
        # TODO: compare different ways of sampling posterior
        sample = self._nns.eval(feed_dict=feed_dict)
        return np.mean(sample), np.std(sample)
    

    def set_parameters(self, shape, activations=tf.tanh, posterior_n_samples = 10):
        # if only one activation function provided make list
        if not hasattr(activations, '__len__'):
            activations = [activations] * (len(shape) - 2)
            activations.append(tf.identity)

        self._activations = activations
            
        self.shape = shape
        
        self._weights = []
        self._biases = []
        self._q_weights = []
        self._q_biases = []
        
        self._qp_dict = {}
        
        with self.graph.as_default():
            with self.sess.as_default() as sess:
                # define variables
                for i in range(len(shape)-1):
                    # initialize weights for layer i with mean 0 and variance sqrt(2/shape[i])
                    W_i = Normal(loc=tf.zeros([shape[i], shape[i+1]]), 
                                 scale= tf.sqrt(2 / shape[i]),
                                 name='W_' + str(i),
                                )
                    b_i = Normal(loc=tf.zeros(shape[i+1]), 
                                 scale=tf.ones(shape[i+1]),
                                 name='b_' + str(i),
                                )
                    with tf.variable_scope("qW_" + str(i)):
                        loc = tf.get_variable("loc", [shape[i], shape[i+1]])
                        scale = tf.nn.softplus(tf.get_variable("scale", [shape[i], shape[i+1]]))
                        qW_i = Normal(loc=loc, scale=scale)

                    with tf.variable_scope("qb_" + str(i)):
                        loc = tf.get_variable("loc", [shape[i+1]])
                        scale = tf.nn.softplus(tf.get_variable("scale", [shape[i+1]]))
                        qb_i = Normal(loc=loc, scale=scale)

                    # store created variables in list
                    self._weights.append(W_i)
                    self._biases.append(b_i)
                    self._q_weights.append(qW_i)
                    self._q_biases.append(qb_i)
                    # relate variables to their approximations
                    self._qp_dict[W_i] = qW_i
                    self._qp_dict[b_i] = qb_i

                # make graph
                self._x = tf.placeholder(tf.float32, shape=(None, self.n_parameters()))

                self._z = Normal(
                    loc=self._make_nn(self._weights, self._biases, self._activations),
                    scale=0.1 * tf.ones(tf.shape(self._x)[0]),
                    name="z"
                )

            # sample weights and biases from approximated weights to make NNs
            nns = []
            for _ in range(posterior_n_samples):
                qweights_sample = [
                    qW.sample() for qW in self._q_weights
                ]
                qbiases_sample = [
                    qb.sample() for qb in self._q_biases
                ]
                nns.append(self._make_nn(qweights_sample,
                                         qbiases_sample, 
                                         self._activations))
            # eval _nns for prediction
            self._nns = tf.stack(nns)

    #np.sqrt(2 / layers_dims[l-1])


    # performs variational inference 
    def fit(self, n_iter=1000, n_samples=5, learning_rate=None):        
        # data for input and target
        data = {
            self._x : self._X,
            self._z : self._y.reshape(-1)
        }
        
        with self.graph.as_default():
            with self.sess.as_default() as sess:
                init = tf.global_variables_initializer()
                sess.run(init)
                # infer parameters through variational inference
                inference = ed.KLqp(self._qp_dict, data=data)
                inference.run(n_iter=n_iter, n_print=10, n_samples=n_samples, logdir='log')
                # for debugging
                self._inference = inference

            
    def _make_nn(self, weights, biases, activations):
        W_0, b_0, a_0 = weights[0], biases[0], activations[0]

        h = a_0(tf.matmul(self._x, W_0) + b_0)
        for W_i, b_i, a_i in zip(weights[1:], biases[1:], activations[1:]):
            h = a_i(tf.matmul(h, W_i) + b_i)
            
        return tf.reshape(h, [-1])
    
    def __str__():
        pass
# create emulator for this instance
emu_nn = BNNEmulator(real_log_likelihood, train_input_normalized, train_target_normalized)
# specify shape of NN
emu_nn.set_parameters(bnn_shape, activations=activations)
emu_nn.fit(n_iter=n_iter, n_samples=n_samples)
emu_nn_normalized = lambda x: denormalize_prediction(emu_nn(normalize_input(x)))
emu_nn_normalized([0.15, 500])<jupyter_output><empty_output><jupyter_text>Performance test<jupyter_code>def mean_error(y_pred, y_true):
    return np.mean(np.abs(y_true - y_pred))
# look at mean error
test_size = 200
test_input = log_prior.sample(test_size)

test_true = np.apply_along_axis(real_log_likelihood, 1, test_input)
test_pred = emu_nn_normalized(test_input)

print("Mean error:", mean_error(test_pred, test_true))
# generate data for surfaces
test_splits = 20 # number of splits along each axis
r_grid, k_grid, test_data = emutils.generate_grid(bounds.lower(), 
                                                  bounds.upper(), 
                                                  test_splits)
emulator_grid = emutils.predict_grid(emu_nn_normalized, test_data)
real_grid = emutils.predict_grid(real_log_likelihood, test_data)

plt.figure(figsize = (10, 5))
ax = emuplt.surface(r_grid, k_grid, emulator_grid, 
                    title = "True log_likelihood",
                    alpha = 0.9,
                    cmap="Blues",
                    x_label = "r (growth rate)",
                    y_label = "k (carrying capacity)"
                   )

ax.plot_surface(r_grid, k_grid, real_grid, cmap="Reds", alpha = .5)
emu_nn_normalized([0.140, 400])
%timeit real_log_likelihood([0.15, 500])
%timeit emu_nn([0.15, 500])<jupyter_output>22.2 µs ± 745 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
909 µs ± 9.95 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)<jupyter_text>Performing MCMC<jupyter_code># Create a posterior log-likelihood (log(likelihood * prior))
log_posterior = pints.LogPosterior(emu_nn, log_prior)
real_log_posterior = pints.LogPosterior(real_log_likelihood, log_prior)
xs = [
    np.array(real_parameters) * 0.9,
    np.array(real_parameters) * 1.05,
    np.array(real_parameters) * 1.15,
]
xs
chains = pints.mcmc_sample(log_posterior, 3, xs)
import pints.plot
pints.plot.trace(chains)
plt.show()
chains = pints.mcmc_sample(real_log_posterior, 3, xs)
pints.plot.trace(chains)
plt.show()
chain = chains[0]

emu_prediction = np.apply_along_axis(emu_nn, 1, chain).flatten()
real_prediction = np.apply_along_axis(real_log_likelihood, 1, chain).flatten()

diffs = np.abs(real_prediction - emu_prediction)
iters = np.linspace(0, 10000, len(chain))
plt.figure(figsize=(10, 5))
plt.title("Emulator and model absolute differences along one MCMC chain")
plt.xlabel("Number of iterations")
plt.ylabel("Likelihood")
plt.plot(iters, diffs, color = "Black")
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10908578408343513, 'fraction_numerical': 0.02067086347834257, 'mean_word_length': 2.336677115987461, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '212019', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3636, 'n_tokens_neox': 3363, 'n_words': 717, 'path': 'studies/legacy/BNN-emulator-development.ipynb', 'repo_name': 'MaksymPetyak/emupints', 'size': '10643'}