<jupyter_start><jupyter_text>Bayesian NN (MLP) for binary classification(**The Bayesian Neural Network Model**) We will implement the following Bayesian model for the data:\begin{align}\mathbf{W} &\sim \mathcal{N}(0, 5^2 \mathbf{I}_{D\times D})\\z^{(n)} &= g_{\mathbf{W}}(\mathbf{X}^{(n)})\\Y^{(n)} &\sim Ber(\text{sigm}(z^{(n)}) ))\end{align}where $g_{\mathbf{W}}$ is a neural network with parameters $\mathbf{W}$ represented as a vector in $\mathbb{R}^{D}$ with $D$ being the total number of parameters (including biases).Implement the log of the joint distribution in `autograd`'s version of `numpy`, i.e. implement $\log \left[p(\mathbf{W})\prod_{n=1}^N p(Y^{(n)} |\mathbf{X}^{(n)} , \mathbf{W}) \right]$.***Hint:*** you'll need to write out the log of the various Gaussian pdf's and implement their formulae using `autograd`'s numpy functions.<jupyter_code>### We first build a MLP NN with numpy

class Feedforward:
    def __init__(self, architecture, random=None, weights=None):
        self.params = {'H': architecture['width'],
                       'L': architecture['hidden_layers'],
                       'D_in': architecture['input_dim'],
                       'D_out': architecture['output_dim'],
                       'activation_type': architecture['activation_fn_type'],
                       'activation_params': architecture['activation_fn_params']}

        self.D = (  (architecture['input_dim'] * architecture['width'] + architecture['width'])
                  + (architecture['output_dim'] * architecture['width'] + architecture['output_dim'])
                  + (architecture['hidden_layers'] - 1) * (architecture['width']**2 + architecture['width'])
                 )

        if random is not None:
            self.random = random
        else:
            self.random = np.random.RandomState(0)

        self.h = architecture['activation_fn']

        if weights is None:
            self.weights = self.random.normal(0, 1, size=(1, self.D))
        else:
            self.weights = weights

        self.objective_trace = np.empty((1, 1))
        self.weight_trace = np.empty((1, self.D))


    def forward(self, weights, x):
        ''' Forward pass given weights and input '''
        H = self.params['H']
        D_in = self.params['D_in']
        D_out = self.params['D_out']

        x = x.T.reshape((1,D_in,-1))
        assert weights.shape[1] == self.D

        if len(x.shape) == 2:
            assert x.shape[0] == D_in
            x = x.reshape((1, D_in, -1))
        else:
            assert x.shape[1] == D_in

        weights = weights.T

        #input to first hidden layer
        W = weights[:H * D_in].T.reshape((-1, H, D_in))
        b = weights[H * D_in:H * D_in + H].T.reshape((-1, H, 1))
        input = self.h(np.matmul(W, x) + b)
        index = H * D_in + H

        assert input.shape[1] == H

        #additional hidden layers
        for _ in range(self.params['L'] - 1):
            before = index
            W = weights[index:index + H * H].T.reshape((-1, H, H))
            index += H * H
            b = weights[index:index + H].T.reshape((-1, H, 1))
            index += H
            output = np.matmul(W, input) + b
            input = self.h(output)

            assert input.shape[1] == H

        #output layer
        W = weights[index:index + H * D_out].T.reshape((-1, D_out, H))
        b = weights[index + H * D_out:].T.reshape((-1, D_out, 1))
        output = np.matmul(W, input) + b
        assert output.shape[1] == self.params['D_out']
        output = output.reshape(-1) # flatten it into a vector, to match the output format of logistic regression??
        return 1 / (1 + np.exp(-output)) # add a sigmoid activation to the output
    
    def forward_VI(self, weights, x):
        ''' Forward pass given weights and input '''
        H = self.params['H']
        D_in = self.params['D_in']
        D_out = self.params['D_out']

        x = x.T.reshape((1,D_in,-1))
        assert weights.shape[1] == self.D

        if len(x.shape) == 2:
            assert x.shape[0] == D_in
            x = x.reshape((1, D_in, -1))
        else:
            assert x.shape[1] == D_in

        weights = weights.T

        #input to first hidden layer
        W = weights[:H * D_in].T.reshape((-1, H, D_in))
        b = weights[H * D_in:H * D_in + H].T.reshape((-1, H, 1))
        input = self.h(np.matmul(W, x) + b)
        index = H * D_in + H

        assert input.shape[1] == H

        #additional hidden layers
        for _ in range(self.params['L'] - 1):
            before = index
            W = weights[index:index + H * H].T.reshape((-1, H, H))
            index += H * H
            b = weights[index:index + H].T.reshape((-1, H, 1))
            index += H
            output = np.matmul(W, input) + b
            input = self.h(output)

            assert input.shape[1] == H

        #output layer
        W = weights[index:index + H * D_out].T.reshape((-1, D_out, H))
        b = weights[index + H * D_out:].T.reshape((-1, D_out, 1))
        output = np.matmul(W, input) + b
        assert output.shape[1] == self.params['D_out']
#         output = output.reshape(-1) # flatten it into a vector, to match the output format of logistic regression??
        return 1 / (1 + np.exp(-output)) # add a sigmoid activation to the output
    
    
    
    
    

    def predict(self,x):
        y = self.forward(self.weights,x)
        return (y > 0.5).astype(np.int_)

    def predict_proba(self, x):
        return np.array([1-self.forward(self.weights,x), self.forward(self.weights,x)]).T


    def make_objective(self, x_train, y_train, reg_param):### modify it to do binary classification

#         def objective(W, t):
#             squared_error = np.linalg.norm(y_train - self.forward(W, x_train), axis=1)**2
#             if reg_param is None:
#                 sum_error = np.sum(squared_error)
#                 return sum_error
#             else:
#                 mean_error = np.mean(squared_error) + reg_param * np.linalg.norm(W)
#                 return mean_error
        
        def objective(W, t):
            binary_cross_entropy = - (1/len(y_train.reshape(-1)))*(np.dot(y_train, np.log(self.forward(W,x_train))) + np.dot(1 - y_train, np.log(1-self.forward(W,x_train))))
            # print(binary_cross_entropy.shape)
            
            return binary_cross_entropy


        return objective, grad(objective)

    def fit(self, x_train, y_train, params, reg_param=None):
        x_train = x_train.reshape((self.params['D_in'],-1))
        y_train = y_train.reshape((self.params['D_out'],-1))

        assert x_train.shape[0] == self.params['D_in']
        assert y_train.shape[0] == self.params['D_out']

        ### make objective function for training
        self.objective, self.gradient = self.make_objective(x_train, y_train, reg_param)

        ### set up optimization
        step_size = 0.01
        max_iteration = 5000
        check_point = 1000
        # weights_init = self.weights.reshape((1, -1))
        weights_init = self.weights
        mass = None
        optimizer = 'adam'
        random_restarts = 5

        if 'step_size' in params.keys():
            step_size = params['step_size']
        if 'max_iteration' in params.keys():
            max_iteration = params['max_iteration']
        if 'check_point' in params.keys():
            self.check_point = params['check_point']
        if 'init' in params.keys():
            weights_init = params['init']
        if 'call_back' in params.keys():
            call_back = params['call_back']
        if 'mass' in params.keys():
            mass = params['mass']
        if 'optimizer' in params.keys():
            optimizer = params['optimizer']
        if 'random_restarts' in params.keys():
            random_restarts = params['random_restarts']

        def call_back(weights, iteration, g):
            ''' Actions per optimization step '''
            objective = self.objective(weights, iteration)
            self.objective_trace = np.vstack((self.objective_trace, objective))
            self.weight_trace = np.vstack((self.weight_trace, weights))
            if iteration % check_point == 0:
                print("Iteration {} lower bound {}; gradient mag: {}".format(iteration, objective, np.linalg.norm(self.gradient(weights, iteration))))

        ### train with random restarts
        optimal_obj = 1e16
        optimal_weights = self.weights

        for i in range(random_restarts):
            if optimizer == 'adam':
#                 adam(self.gradient, weights_init, step_size=step_size, num_iters=max_iteration, callback=call_back)
                sgd(self.gradient, weights_init, step_size=step_size, num_iters=max_iteration, callback=call_back)

            local_opt = np.min(self.objective_trace[-100:])
            if local_opt < optimal_obj:
                opt_index = np.argmin(self.objective_trace[-100:])
                self.weights = self.weight_trace[-100:][opt_index].reshape((1, -1))
            weights_init = self.random.normal(0, 1, size=(1, self.D))

        self.objective_trace = self.objective_trace[1:]
        self.weight_trace = self.weight_trace[1:]
###relu activation
activation_fn_type = 'relu'
activation_fn = lambda x: np.maximum(np.zeros(x.shape), x)

###neural network model design choices
width = 5
hidden_layers = 1
input_dim = 2
output_dim = 1

architecture = {'width': width,
               'hidden_layers': hidden_layers,
               'input_dim': input_dim,
               'output_dim': output_dim,
               'activation_fn_type': 'relu',
               'activation_fn_params': 'rate=1',
               'activation_fn': activation_fn}

#instantiate a Feedforward neural network object
# nn = Feedforward(architecture, random=random) # initiate W randomly

# ###define design choices in gradient descent
# params = {'step_size':1e-3, 
#           'max_iteration':15000, 
#           'random_restarts':1}

# #fit my neural network to minimize BCE on the given data  ### we are not using this and will use sklearn for maximum likelihood inference
# nn.fit(x_train, y_train, params)
## todo: add a function to unpack weights for NN 
#instantiate a Feedforward neural network object, initiate W with the MLE solution learnt by sklearn

# copy the weights from the sklearn model
MLE_weights = []
for a,b in zip(model_nn.coefs_,model_nn.intercepts_):
    MLE_weights = MLE_weights + list(np.array(a).flatten(order= 'F')) + list(np.array(b).flatten())
MLE_weights = np.array(MLE_weights).reshape(1,-1)

#set random state to make the experiments replicable
rand_state = 0
random = np.random.RandomState(rand_state)
nn = Feedforward(architecture, random=random, weights = MLE_weights) 
#Visualize the decision boundary of the nn with the MLE solution

models = [nn]

fig, ax = plt.subplots(1, 1, figsize=(10, 5))
ax = plot_decision_boundary(x[:,:], y, models, ax, poly_degree=1,shaded =True)
ax.set_xlabel('x_1')
ax.set_ylabel('x_2')
ax.legend(loc='best')
plt.show()
# Implement the log of the joint distribution
def log_joint(x,y,weights,nn):
    
    N = x.shape[0]# number of samples
    
    output = nn.forward(weights, x)
    
    w_cov = 25*np.eye(nn.D)
    w_inv = np.linalg.inv(w_cov)
    
    log_prior = -0.5*nn.D*np.log(2*np.pi) - 0.5 * np.log(np.linalg.det(w_cov)) - 0.5*np.matmul(np.matmul(weights,w_inv),weights.T)
    
    log_condition = - (1/N)* (np.dot(y, np.log(1e-10 + output)) + np.dot(1 - y, np.log(1e-10 + 1-output)))
    return log_prior + log_condition
# HMC 

def leap_frog(position_init, momentum_init, step_size, leapfrog_steps, potential_energy, kinetic_energy):        # initialize position
        position = position_init
        # initialize momentum
        momentum = momentum_init

        # compute the gradient of potential energy
        potential_energy_grad = grad(potential_energy)

        # compute the gradient of kinetic energy
        kinetic_energy_grad = grad(kinetic_energy)

        # print(momentum)
        # print(position)
        # full leap frog steps
        for _ in range(leapfrog_steps - 1):
            # half step update of momentum
#             print(momentum)
#             print(position)
            momentum = momentum - step_size * 0.5 * potential_energy_grad(position)
#             print(potential_energy_grad(position))
#             print(momentum)
#             print(position)
            # full step update of position
            position = position + step_size * kinetic_energy_grad(momentum)
#             print(kinetic_energy_grad(momentum))
            # print(momentum)
#             print(position)
            # half step update of momentum
            momentum = momentum - step_size * 0.5 * potential_energy_grad(position)
            # print(position)
#             print(potential_energy_grad(position))
            assert not np.any(np.isnan(position))
            assert not np.any(np.isnan(momentum))

        # print(momentum)
        # print(position)
        # Reverse momentum
        momentum = -momentum
        return position, momentum

def hmc(position_current, momentum_current, step_size, leapfrog_steps, potential_energy, kinetic_energy, momentum_sampler):

        ### Define total energy
        total_energy = lambda q, p: potential_energy(q) + kinetic_energy(p)

        ### Sample momentum
        # Make sure that the distribution from which you sample the momentum matches 
        # with your kinetic energy function
        momentum_current = momentum_sampler(1) 

        # print(momentum_current)
        # print(position_current)
        ### Simulate Hamiltonian dynamics using Leap Frog
        position_proposal, momentum_proposal = leap_frog(position_current.copy(), momentum_current.copy(), step_size, leapfrog_steps, potential_energy, kinetic_energy)

        # compute total energy in current position and proposal position
        current_total_energy = total_energy(position_current, momentum_current)
        proposal_total_energy = total_energy(position_proposal, momentum_proposal)

        ### Metropolis Hastings Step
        # comute accept probability
        accept_prob = min(1,np.exp(current_total_energy - proposal_total_energy))

        # accept proposal with accept probability
        if np.random.rand() < accept_prob:
            accept = 1
            position_current = np.copy(position_proposal)
            momentum_current = np.copy(momentum_proposal)
        else:
            accept = 0
            
        return position_current, momentum_current, accept
### Setting up the energy functions for HMC
potential_energy = lambda q: -log_joint(x_train,y_train,q,nn)
kinetic_energy = lambda p: np.sum(0.5*np.power(p,2))
momentum_sampler = lambda n: np.random.normal(0,1,size = nn.D)

### Setting up hyperparameters of HMC
#initialize the position vector for HMC sampling with the weights of the trained neural network
position_init = nn.weights
position_current = position_init

momentum_current = momentum_sampler(1)
step_size = 0.005
leapfrog_steps = 20

### Running HMC to sample from a target distribution
total_trials = 10000 # 5000
burn_in = 0.1
thin = 2
samples = [position_current]
accepts = 0.

for t in range(total_trials):
    position_current, momentum_current, accept = hmc(position_current, momentum_current, step_size, leapfrog_steps, potential_energy, kinetic_energy, momentum_sampler)
    samples.append(position_current)
    accepts += accept
    if t % 1000 == 0 and t > 0:
        print('{} iterations, accept rate: {}'.format(t, accepts / (t + 1.)))

accept_rate = accepts / total_trials
samples = np.array(samples[int(total_trials*burn_in)::thin])
fig, ax = plt.subplots(21, 2, figsize=(10, 45))

for i in range(21):
    ax[i, 0].hist(samples[:, 0,i], bins=50, density=True, label=f"coef{i}")
    ax[i, 1].plot(samples[:, 0,i], label=f"coef{i}")

    ax[i, 0].legend()
    ax[i, 1].legend()
#Visualize the decision boundaries corresponding to the posterior samples from your Bayesian Logistic Regression model (see In-Class Exercise 10.13)
# posterior_samples_index = np.random.choice(np.arange(len(samples)), size=100)
posterior_samples_index = np.random.choice(np.arange(len(samples)), size=100)
models_nn = []

for index in posterior_samples_index:
    sample = np.array(samples[index]).reshape(1,-1)
    model = Feedforward(architecture, weights = sample)
    models_nn.append(model)
# Estimate the epistemic uncertainty of each test data point

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
ax = plot_decision_boundary(x[:,:], y, models_nn, ax, poly_degree=1)

for i in range(len(test_points)):
    pt = test_points[i]
    if i == 0:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black', label='test data')
    else:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black')
    

for c,test_point in enumerate(test_points):
    test_pred_prob = np.zeros(100)
    test_point_ = test_point
    test_point = np.array(test_point).reshape(1,-1)
    
    for c,index in enumerate(posterior_samples_index):
        sample = samples[index]
        model = Feedforward(architecture, weights = sample)
        test_pred_prob[c] = model.predict_proba(test_point)
        
    if c == 0:    
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange', label = 'ttt')
    else:
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange')
        
    print(f' estimated epistemic uncertainty of test point {test_point} (the variance of the predict probability):{np.var(test_pred_prob)}')
    
ax.set_xlabel('x_1')
ax.set_ylabel('x_2')
ax.legend(loc='best')
plt.show()
def black_box_variational_inference(logprob, D, num_samples):
    
    """
    Implements http://arxiv.org/abs/1401.0118, and uses the
    local reparameterization trick from http://arxiv.org/abs/1506.02557
    code taken from:
    https://github.com/HIPS/autograd/blob/master/examples/black_box_svi.py
    """

    def unpack_params(params):
        # Variational dist is a diagonal Gaussian.
        mean, log_std = params[:D], params[D:]
        return np.array(mean), np.array(log_std)

    def negative_gaussian_entropy(log_std):
        return 0.5 * D * (1.0 + np.log(2*np.pi)) + np.sum(log_std)

    rs = np.random.RandomState(0)
    def variational_objective(params, t):
        """Provides a stochastic estimate of the variational lower bound."""
        mean, log_std = unpack_params(params)
        samples = rs.randn(num_samples, D) * np.exp(log_std) + mean
        lower_bound = negative_gaussian_entropy(log_std) + np.mean(logprob(samples, t))
        return -lower_bound

    gradient = grad(variational_objective)

    return variational_objective, gradient, unpack_params
# def log_joint(x,y,weights,nn):
    
#     N = x.shape[0]# number of samples
    
#     output = nn.forward(weights, x)
    
#     w_cov = 25*np.eye(nn.D)
#     w_inv = np.linalg.inv(w_cov)
    
#     log_prior = -0.5*nn.D*np.log(2*np.pi) - 0.5 * np.log(np.linalg.det(w_cov)) - 0.5*np.matmul(np.matmul(weights,w_inv),weights.T)
    
#     log_condition = - (1/N)* (np.dot(y, np.log(1e-10 + output)) + np.dot(1 - y, np.log(1e-10 + 1-output)))
#     return log_prior + log_condition
def variational_inference_BNN(Sigma_W, y_train, x_train, forward, S, max_iteration, step_size, verbose,init_W = None):
    '''implements wrapper for variational inference via bbb for bayesian NN'''
    D = Sigma_W.shape[0]
    Sigma_W_inv = np.linalg.inv(Sigma_W)
    Sigma_W_det = np.linalg.det(Sigma_W)
    variational_dim = D
    N = x_train.shape[0]
    
    #define the log prior on the model parameters
    def log_prior(W):
        constant_W = -0.5 * (D * np.log(2 * np.pi) + np.log(Sigma_W_det))
        exponential_W = -0.5 * np.diag(np.dot(np.dot(W, Sigma_W_inv), W.T))
        log_p_W = constant_W + exponential_W
#         print(log_p_W)
        return log_p_W

    #define the log likelihood
#     def log_lklhd(W):
#         S = W.shape[0]
#         constant = (-np.log(sigma_y) - 0.5 * np.log(2 * np.pi)) * N
#         exponential = -0.5 * sigma_y**-2 * np.sum((y_train.reshape((1, 1, N)) - forward(W, x_train))**2, axis=2).flatten()
#         return constant + exponential
    def log_lklhd(W):
#         print(forward(W, x_train).shape)
        output = forward(W, x_train)
        
#         print((y_train.reshape(-1,1) * np.log(output))[-3:])
        log_condition = np.sum(y_train * np.log(output + 1e-10) + (1- y_train) * (np.log(1-output+1e-10)),axis = 2)
#         print(log_condition.shape)


#         output = forward(W, x_train)
#         argument = (1e-10+output)**y_train + (1e-10 + 1-output)**(1-y_train)
#         log_condition = sp.special.logsumexp(argument, axis=2)
#         print(log_condition.shape)
        return log_condition
    #define the log joint density
    log_density = lambda w, t: log_lklhd(w) + log_prior(w)

    #build variational objective.
    objective, gradient, unpack_params = black_box_variational_inference(log_density, D, num_samples=S)

    def callback(params, t, g):
        if verbose:
            if  t % 500 == 0:
                print("Iteration {} lower bound {}; gradient mag: {}".format(t, -objective(params, t), np.linalg.norm(gradient(params, t))))

    print("Optimizing variational parameters...")
    #initialize variational parameters
    if init_W is not None:
        init_mean = init_W.flatten()
    else:
        init_mean = np.ones(D)
    init_log_std = -1 * np.ones(D)
    init_var_params = np.concatenate([init_mean, init_log_std])
    
    #perform gradient descent using adam (a type of gradient-based optimizer)
    variational_params = adam(gradient, init_var_params, step_size=step_size, num_iters=max_iteration, callback=callback)
    
    return variational_params 
###relu activation
activation_fn_type = 'relu'
activation_fn = lambda x: np.maximum(np.zeros(x.shape), x)

###neural network model design choices
width = 3
hidden_layers = 2
input_dim = 2
output_dim = 1

architecture = {'width': width,
               'hidden_layers': hidden_layers,
               'input_dim': input_dim,
               'output_dim': output_dim,
               'activation_fn_type': 'relu',
               'activation_fn_params': 'rate=1',
               'activation_fn': activation_fn}

nn = Feedforward(architecture, random=random) 

D = nn.D
weight_noise = 3**2
Sigma_W = weight_noise * np.eye(D)


BNN_paras = variational_inference_BNN(Sigma_W, y_train, x_train, nn.forward_VI, S = 100, max_iteration=100000, step_size=1e-3, verbose=True)
mean, log_std = BNN_paras[:D], BNN_paras[D:]
samples = np.random.randn(1000, D) * np.exp(log_std) + mean

posterior_samples_index = np.random.choice(np.arange(len(samples)), size=100)
models_nn = []

for index in posterior_samples_index:
    sample = np.array(samples[index]).reshape(1,-1)
    model = Feedforward(architecture, weights = sample)
    models_nn.append(model)
# Estimate the epistemic uncertainty of each test data point

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
ax = plot_decision_boundary(x[:,:], y, models_nn, ax, poly_degree=1)

for i in range(len(test_points)):
    pt = test_points[i]
    if i == 0:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black', label='test data')
    else:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black')
    

for c,test_point in enumerate(test_points):
    test_pred_prob = np.zeros(100)
    test_point_ = test_point
    test_point = np.array(test_point).reshape(1,-1)
    
    for c,index in enumerate(posterior_samples_index):
        sample = samples[index].reshape(1,-1)
        model = Feedforward(architecture, weights = sample)
        test_pred_prob[c] = model.predict_proba(test_point)[:,1]
        
    if c == 0:    
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange', label = 'ttt')
    else:
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange')
        
    print(f' estimated epistemic uncertainty of test point {test_point} (the variance of the predict probability):{np.var(test_pred_prob)}')
    
ax.set_xlabel('x_1')
ax.set_ylabel('x_2')
ax.legend(loc='best')
plt.show()<jupyter_output>estimated epistemic uncertainty of test point [[-5 -5]] (the variance of the predict probability):0.0329768556295834
 estimated epistemic uncertainty of test point [[-5.   3.5]] (the variance of the predict probability):0.058923708015756644
 estimated epistemic uncertainty of test point [[5 5]] (the variance of the predict probability):3.316374938389455e-22
 estimated epistemic uncertainty of test point [[ 5.  -3.5]] (the variance of the predict probability):0.1143617841685994
 estimated epistemic uncertainty of test point [[-1  1]] (the variance of the predict probability):0.08346166511153788
 estimated epistemic uncertainty of test point [[-1.  -1.5]] (the variance of the predict probability):0.0005413064003127623
 estimated epistemic uncertainty of test point [[1.  1.5]] (the variance of the predict probability):2.1717879016920705e-08<jupyter_text>relu activationactivation_fn_type = 'relu'activation_fn = lambda x: np.maximum(np.zeros(x.shape), x)neural network model design choiceswidth = 5hidden_layers = 3input_dim = 2output_dim = 1architecture = {'width': width,               'hidden_layers': hidden_layers,               'input_dim': input_dim,               'output_dim': output_dim,               'activation_fn_type': 'relu',               'activation_fn_params': 'rate=1',               'activation_fn': activation_fn}<jupyter_code># try BBVI on two NN architectures
experiments = [(5,3),(15,1),]
# results = []
results_w1 = []
for ex in experiments:
    print(ex)
    width, hidden_layers = ex
    
    architecture = {'width': width,
               'hidden_layers': hidden_layers,
               'input_dim': input_dim,
               'output_dim': output_dim,
               'activation_fn_type': 'relu',
               'activation_fn_params': 'rate=1',
               'activation_fn': activation_fn}
    
    nn = Feedforward(architecture, random=random) 

    D = nn.D
    sigma_y = 0.5**2
    weight_noise = 1**2
    Sigma_W = weight_noise * np.eye(D)


    BNN_paras = variational_inference_BNN(Sigma_W, y_train, x_train, nn.forward_VI, S = 100, max_iteration= 100000, step_size=1e-3, verbose=True)
    
    results_w1.append(BNN_paras)
results_w1[0].shape
width = 5
hidden_layers = 3
input_dim = 2
output_dim = 1

architecture = {'width': width,
               'hidden_layers': hidden_layers,
               'input_dim': input_dim,
               'output_dim': output_dim,
               'activation_fn_type': 'relu',
               'activation_fn_params': 'rate=1',
               'activation_fn': activation_fn}
nn = Feedforward(architecture, random=random)
BNN_paras = results_w1[0]
D = nn.D
mean, log_std = BNN_paras[:D], BNN_paras[D:]
samples = np.random.randn(1000, D) * np.exp(log_std) + mean

posterior_samples_index = np.random.choice(np.arange(len(samples)), size=100)
models_nn = []

for index in posterior_samples_index:
    sample = np.array(samples[index]).reshape(1,-1)
    model = Feedforward(architecture, weights = sample)
    models_nn.append(model)

# Estimate the epistemic uncertainty of each test data point

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
ax = plot_decision_boundary(x[:,:], y, models_nn, ax, poly_degree=1)

for i in range(len(test_points)):
    pt = test_points[i]
    if i == 0:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black', label='test data')
    else:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black')
    

for c,test_point in enumerate(test_points):
    test_pred_prob = np.zeros(100)
    test_point_ = test_point
    test_point = np.array(test_point).reshape(1,-1)
    
    for c,index in enumerate(posterior_samples_index):
        sample = samples[index].reshape(1,-1)
        model = Feedforward(architecture, weights = sample)
        test_pred_prob[c] = model.predict_proba(test_point)[:,1]
        
    if c == 0:    
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange', label = 'ttt')
    else:
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange')
        
    print(f' estimated epistemic uncertainty of test point {test_point} (the variance of the predict probability):{np.var(test_pred_prob)}')
    
ax.set_xlabel('x_1')
ax.set_ylabel('x_2')
ax.legend(loc='best')
plt.show()
width = 15
hidden_layers = 1
input_dim = 2
output_dim = 1

architecture = {'width': width,
               'hidden_layers': hidden_layers,
               'input_dim': input_dim,
               'output_dim': output_dim,
               'activation_fn_type': 'relu',
               'activation_fn_params': 'rate=1',
               'activation_fn': activation_fn}
nn = Feedforward(architecture, random=random)
BNN_paras = results_w1[1]
D = nn.D
mean, log_std = BNN_paras[:D], BNN_paras[D:]
samples = np.random.randn(1000, D) * np.exp(log_std) + mean

posterior_samples_index = np.random.choice(np.arange(len(samples)), size=100)
models_nn = []

for index in posterior_samples_index:
    sample = np.array(samples[index]).reshape(1,-1)
    model = Feedforward(architecture, weights = sample)
    models_nn.append(model)

# Estimate the epistemic uncertainty of each test data point

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
ax = plot_decision_boundary(x[:,:], y, models_nn, ax, poly_degree=1)

for i in range(len(test_points)):
    pt = test_points[i]
    if i == 0:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black', label='test data')
    else:
        ax.scatter(pt[0], pt[1], alpha=1., s=50, color='black')
    

for c,test_point in enumerate(test_points):
    test_pred_prob = np.zeros(100)
    test_point_ = test_point
    test_point = np.array(test_point).reshape(1,-1)
    
    for c,index in enumerate(posterior_samples_index):
        sample = samples[index].reshape(1,-1)
        model = Feedforward(architecture, weights = sample)
        test_pred_prob[c] = model.predict_proba(test_point)[:,1]
        
    if c == 0:    
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange', label = 'ttt')
    else:
        ax.scatter(test_point_[0], test_point_[1], alpha=.4, s=100000*np.var(test_pred_prob), color='orange')
        
    print(f' estimated epistemic uncertainty of test point {test_point} (the variance of the predict probability):{np.var(test_pred_prob)}')
    
ax.set_xlabel('x_1')
ax.set_ylabel('x_2')
ax.legend(loc='best')
plt.show()<jupyter_output>estimated epistemic uncertainty of test point [[-5 -5]] (the variance of the predict probability):4.97490031347715e-06
 estimated epistemic uncertainty of test point [[-5.   3.5]] (the variance of the predict probability):0.12270997642894013
 estimated epistemic uncertainty of test point [[5 5]] (the variance of the predict probability):5.61530663194464e-12
 estimated epistemic uncertainty of test point [[ 5.  -3.5]] (the variance of the predict probability):0.15051943622617775
 estimated epistemic uncertainty of test point [[-1  1]] (the variance of the predict probability):0.10612898340688391
 estimated epistemic uncertainty of test point [[-1.  -1.5]] (the variance of the predict probability):0.002234075754900715
 estimated epistemic uncertainty of test point [[1.  1.5]] (the variance of the predict probability):0.00011235224543082482
{'fraction_non_alphanumeric': 0.12886369394476818, 'fraction_numerical': 0.027584241195844947, 'mean_word_length': 2.990521925944648, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 1, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '440696', 'license': 'MIT', 'n_tokens_mistral': 10989, 'n_tokens_neox': 10212, 'n_words': 2424, 'path': 'BBVI_BNN/BNN.ipynb', 'repo_name': 'john-vastola/ML-from-scratch-seminar', 'size': '31576'}