<jupyter_start><jupyter_text>In this notebook we fix the Bayes estimator, as it's not working. As a baseline, we take the sk-learn algo. Those are the values we need to hit.<jupyter_code># General imports
import numpy as np
import torch

# DeepMoD stuff
from multitaskpinn import DeepMoD
from multitaskpinn.model.func_approx import Siren, NN
from multitaskpinn.model.library import Library1D
from multitaskpinn.model.constraint import LeastSquares
from multitaskpinn.model.sparse_estimators import Threshold 

from phimal_utilities.data import Dataset
from phimal_utilities.data.burgers import BurgersDelta
from sklearn.linear_model import BayesianRidge

if torch.cuda.is_available():
    device ='cuda'
else:
    device = 'cpu'

# Settings for reproducibility
np.random.seed(42)
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


%load_ext autoreload
%autoreload 2
# 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)
X, y = dataset.create_dataset(x_grid.reshape(-1, 1), t_grid.reshape(-1, 1), n_samples=1000, noise=0.2, random=True, normalize=False)
X, y = X.to(device), y.to(device)
        

network = NN(2, [30, 30, 30, 30, 30], 1)
library = Library1D(poly_order=2, diff_order=3) # Library function
estimator = Threshold(0.1) # Sparse estimator 
constraint = LeastSquares() # How to constrain
model = DeepMoD(network, library, estimator, constraint).to(device) # Putting it all in the model
# Getting data out
prediction, time_derivs, thetas = model(X)
t = time_derivs[0].cpu().detach().numpy()
theta = thetas[0].cpu().detach().numpy()
# Getting sklearn results
sk_reg = BayesianRidge(fit_intercept=False, compute_score=True, alpha_1=0, alpha_2=0, lambda_1=0, lambda_2=0)
sk_reg.fit(theta, t)
# Precision, noise level
print(sk_reg.alpha_, 1 / sk_reg.alpha_)
# Precision of prior, std of prior
print(sk_reg.lambda_, 1 / sk_reg.lambda_)
# Found coeffs
print(sk_reg.coef_)
# What would the reg cost be; probably similar to the noise level
np.mean((t - theta @ sk_reg.coef_[:, None])**2) # beware; sk learn coeff output is 1d and will give wrong results!
# And the score aka the log evidence
sk_reg.scores_[-1]<jupyter_output><empty_output><jupyter_text>Which sounds reasonable but it isn't at all; now let's implement our own and see if we can get it to work.<jupyter_code># Getting data out
prediction, time_derivs, thetas = model(X)
t = time_derivs[0]
Theta = thetas[0]

# Let's use the found alpha and beta to check if the code is good; later we see if we can find it by optimizing
alpha = torch.tensor(sk_reg.lambda_) # precision of weights; follow sk learn init
beta = torch.tensor(sk_reg.alpha_) # precision of noise

M = Theta.shape[1]
N = Theta.shape[0]
# Posterior std and mean
A = torch.eye(M).to(Theta.device) * alpha + beta * Theta.T @ Theta  
mn = beta * torch.inverse(A) @ Theta.T @ t
# Difference between posterior means; seems within numerical acc.
np.abs(mn.detach().cpu().numpy() - sk_reg.coef_[:, None])
# Difference between posterior std; seems within numerical acc.
np.mean(np.abs(sk_reg.sigma_ - Sn.detach().cpu().numpy()))
1/2 * (M * torch.log(alpha) 
 + N * torch.log(beta)
 - beta * (t - Theta @ mn).T @ (t - Theta @ mn) - alpha * mn.T @ mn 
 - torch.trace(torch.log(A))
 - N * np.log(2*np.pi))<jupyter_output><empty_output><jupyter_text>Which is close enough given numerical accuracy<jupyter_code>def neg_marginal_LL(Theta, t, alpha, beta):
    M = Theta.shape[1]
    N = Theta.shape[0]
    
    
    # Posterior std and mean
    A = torch.eye(M).to(Theta.device) * alpha + beta * Theta.T @ Theta  
    mn = beta * torch.inverse(A) @ Theta.T @ t
    
    loss = -1/2 * (M * torch.log(alpha) 
         + N * torch.log(beta)
         - beta * (t - Theta @ mn).T @ (t - Theta @ mn) - alpha * mn.T @ mn 
         - torch.trace(torch.log(A))
         - N * np.log(2*np.pi))
    
    return loss
neg_marginal_LL(Theta, t, torch.tensor(sk_reg.lambda_), torch.tensor(sk_reg.alpha_))
#Now let's try to optimize using gradient descent
prediction, time_derivs, thetas = model(X)
t = time_derivs[0].detach().cpu()
Theta = thetas[0].detach().cpu()

alpha = torch.nn.Parameter(torch.log(torch.tensor(1.0))) # precision of weights; follow sk learn init
beta = torch.nn.Parameter(torch.log(torch.tensor(1 / torch.var(t)))) # precision of noise

optimizer = torch.optim.Adam([{'params': alpha}, {'params': beta}])

print(alpha, beta)
for iteration in torch.arange(100000):
    loss = neg_marginal_LL(Theta, t, torch.exp(alpha), torch.exp(beta))
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if iteration % 1000 == 0:
        print(loss)
torch.exp(alpha)
torch.exp(beta)
print(sk_reg.lambda_, sk_reg.alpha_)<jupyter_output>0.014547715291435611 624313.8692322593
{'fraction_non_alphanumeric': 0.13026584867075663, 'fraction_numerical': 0.02556237218813906, 'mean_word_length': 4.247854077253219, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 12, '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': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '392892', 'license': 'MIT', 'n_tokens_mistral': 1869, 'n_tokens_neox': 1778, 'n_words': 480, 'path': 'experiments/KS/Bayes/Fixing_Bayes.ipynb', 'repo_name': 'GJBoth/MultiTaskPINN', 'size': '4890'}