<jupyter_start><jupyter_text>In this notebook we test our implementation of bayesian regression. We build a gradient-descent version and compare it to the sk-learn result to see if it works.<jupyter_code># General imports
import numpy as np
import torch

from deepymod.data import Dataset
from deepymod.data.burgers import BurgersDelta
from sklearn.linear_model import BayesianRidge<jupyter_output><empty_output><jupyter_text>Making data<jupyter_code># Making dataset
v = 0.1
A = 1.0

x = np.linspace(-3, 4, 100)
t = np.linspace(0.5, 5.0, 50)
x_grid, t_grid = np.meshgrid(x, t, indexing='ij')
dataset = Dataset(BurgersDelta, v=v, A=A)

y = dataset.time_deriv(x_grid.reshape(-1, 1), t_grid.reshape(-1, 1)) # observations
X = dataset.library(x_grid.reshape(-1, 1), t_grid.reshape(-1, 1), poly_order=2, deriv_order=3) # covariates

print(y.shape, X.shape)
y += np.std(y) * 0.1 * np.random.randn(*y.shape)<jupyter_output><empty_output><jupyter_text>Baseline<jupyter_code>regressor = BayesianRidge(fit_intercept=False, compute_score=True, alpha_1=0, alpha_2=0, lambda_1=0, lambda_2=0)
regressor.fit(X, y.squeeze())
baseline_coeffs = regressor.coef_[:, None]
print(baseline_coeffs)
baseline_noise_precision = regressor.alpha_
print(baseline_noise_precision)
baseline_prior_precision = regressor.lambda_
print(baseline_prior_precision)
regressor.scores_<jupyter_output><empty_output><jupyter_text>Own implementation<jupyter_code>X = torch.tensor(X)
y = torch.tensor(y)
alpha_ = torch.tensor(baseline_noise_precision)
lambda_ = torch.tensor(baseline_prior_precision)

N = X.shape[0]
M = X.shape[1]
Sn = torch.inverse(lambda_ * torch.eye(M) + alpha_ * X.T @ X)
mn = alpha_ * Sn @ X.T @ y
print(mn)<jupyter_output>tensor([[-1.8028e-15],
        [-5.9827e-14],
        [ 1.0000e-01],
        [ 2.4078e-15],
        [ 1.1935e-14],
        [-1.0000e+00],
        [-1.2396e-13],
        [ 3.3003e-15],
        [-7.8382e-14],
        [-4.2205e-13],
        [ 1.0732e-13],
        [-7.6553e-15]], dtype=torch.float64)<jupyter_text>This matches the SK learn values, so thats correct. Now to calculate the neg LL<jupyter_code>%%time
mu_post = X @ mn
sigma_post = alpha_ * torch.eye(N) + X @ Sn @ X.T 
L = torch.inverse(sigma_post)
log_p = 1/2 * (-torch.trace(torch.log(L)) - (y - mu_post).T @ L @ (y - mu_post) - N * np.log(2*np.pi))
print(log_p)<jupyter_output>tensor([[158446.8950]], dtype=torch.float64)<jupyter_text>Seems close enough to the sklearn implementation, although we can make many things much more efficient (e.g. use woodbury inversion etc). Now let's try and do gradient descent:<jupyter_code>alpha_ = torch.nn.Parameter(1/torch.var(y))
lambda_ = torch.nn.Parameter(torch.ones(1))

X = torch.tensor(X)
y = torch.tensor(y)

N = X.shape[0]
M = X.shape[1]
optimizer = torch.optim.Adam([alpha_, lambda_], lr=1.0)
max_epochs = 1e4
for epoch in torch.arange(max_epochs):
    Sn = torch.inverse(lambda_ * torch.eye(M) + alpha_ * X.T @ X)
    mn = alpha_ * Sn @ X.T @ y
    
    mu_post = X @ mn
    sigma_post = alpha_ * torch.eye(N) + X @ Sn @ X.T 
    L = torch.inverse(sigma_post)

    log_p = 1/2 * (-torch.trace(torch.log(L)) - (y - mu_post).T @ L @ (y - mu_post) - N * np.log(2*np.pi))
    loss = -log_p
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 1e3 == 0:
        print(log_p)<jupyter_output>tensor([[4565.8581]], dtype=torch.float64, grad_fn=<MulBackward0>)<jupyter_text>Own implementation, efficient:<jupyter_code>alpha_ = torch.tensor(baseline_noise_precision)
lambda_ = torch.tensor(baseline_prior_precision)

N = X.shape[0]
M = X.shape[1]

Sn = lambda_ * torch.eye(M) + alpha_ * X.T @ X
mn = alpha_ * torch.inverse(Sn) @ X.T @ y
E = alpha_ / 2 * torch.sum((y - X @ mn)**2) + lambda_ / 2 * mn.T @ mn
log_p = M / 2 * torch.log(lambda_) + N / 2 * torch.log(alpha_) - E - 1/2 * torch.trace(torch.log(Sn)) - N / 2 * np.log(2*np.pi)
alpha_ = torch.nn.Parameter(1/torch.var(y))
lambda_ = torch.nn.Parameter(torch.ones(1))

X = torch.tensor(X)
y = torch.tensor(y)

N = X.shape[0]
M = X.shape[1]
optimizer = torch.optim.Adam([alpha_, lambda_])
max_epochs = 1e5
for epoch in torch.arange(max_epochs):
    A = lambda_ * torch.eye(M) + alpha_ * X.T @ X
    mn = alpha_ * torch.inverse(A) @ X.T @ y
    
    E = alpha_ / 2 * torch.sum((y - X @ mn)**2) + lambda_ / 2 * mn.T @ mn
    log_p = M / 2 * torch.log(lambda_) + N / 2 * torch.log(alpha_) - E - 1/2 * torch.trace(torch.log(Sn)) - N / 2 * np.log(2*np.pi)
    loss = -log_p
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 1e4 == 0:
        print(log_p)
alpha_
lambda_<jupyter_output><empty_output><jupyter_text>Cool, it works. Now let's optimize for big and small values:<jupyter_code>X = torch.tensor(X)
y = torch.tensor(y)

N = X.shape[0]
M = X.shape[1]


a = torch.nn.Parameter(-torch.log(torch.var(y)))
l = torch.nn.Parameter(torch.zeros(M))
optimizer = torch.optim.Adam([a, l], lr=1.0)
max_epochs = 1e4
for epoch in torch.arange(max_epochs):
    lambda_ = torch.exp(l)
    alpha_ = torch.exp(a)
    
    A = lambda_ * torch.eye(M) + alpha_ * X.T @ X
    mn = alpha_ * torch.inverse(A) @ X.T @ y
    
    E = alpha_ * torch.sum((y - X @ mn)**2) + lambda_ * mn.T @ mn
    loss = E + torch.sum(torch.log(torch.diag(A))) - (M * l + N * a)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 1e3 == 0:
        print(loss)
mn
lambda_
alpha_
a<jupyter_output><empty_output><jupyter_text>Sparse bayesian learning<jupyter_code>X = torch.tensor(X)
y = torch.tensor(y)

N = X.shape[0]
M = X.shape[1]


a = torch.nn.Parameter(-torch.log(torch.var(y)))
l = torch.nn.Parameter(torch.zeros(M, dtype=torch.float64))
optimizer = torch.optim.Adam([a, l], lr=1.0)
max_epochs = 1e4
for epoch in torch.arange(max_epochs):
    lambda_ = torch.exp(l)
    alpha_ = torch.exp(a)
    
    A = torch.diag(lambda_) + alpha_ * X.T @ X
    mn = alpha_ * torch.inverse(A) @ X.T @ y
    
    E = alpha_ * torch.sum((y - X @ mn)**2) + mn.T @ torch.diag(lambda_) @ mn
    loss = E + torch.sum(torch.log(torch.diag(A))) - (torch.sum(l) + N * a)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 1e3 == 0:
        print(loss)
X = torch.tensor(X)
y = torch.tensor(y)

N = X.shape[0]
M = X.shape[1]


a = torch.nn.Parameter(-torch.log(torch.var(y)))
l = torch.nn.Parameter(torch.zeros(M, dtype=torch.float64))
threshold = 1e4
optimizer = torch.optim.Adam([a, l], lr=1.0)
max_epochs = 1e4
for epoch in torch.arange(max_epochs):
    lambda_ = torch.exp(l)
    alpha_ = torch.exp(a)
    
    A = torch.diag(lambda_) + alpha_ * X.T @ X
    mn = (lambda_ < threshold)[:, None] * (alpha_ * torch.inverse(A) @ X.T @ y)
    
    E = alpha_ * torch.sum((y - X @ mn)**2) + mn.T @ torch.diag(lambda_) @ mn
    loss = E + torch.sum(torch.log(torch.diag(A)[lambda_ < threshold])) - (torch.sum(l[lambda_ < threshold]) + N * a)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 1e3 == 0:
        print(loss)
torch.diag(A)
mn
l[lambda_ < threshold]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16274150331406007, 'fraction_numerical': 0.03215343393033423, 'mean_word_length': 3.41044776119403, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 35, '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': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '286953', 'license': 'MIT', 'n_tokens_mistral': 3055, 'n_tokens_neox': 2913, 'n_words': 588, 'path': 'notebooks/Bayes/.ipynb_checkpoints/Bayesian_regression-checkpoint.ipynb', 'repo_name': 'GJBoth/MultiTaskPINN', 'size': '7091'}