Store simulation/training results about brian2modelfitting HOT 4 CLOSED

import pickle

posterior = ...

with open("/path/to/posterior.pkl", "wb") as handle:
    pickle.dump(posterior, handle)

posterior.net.state_dict()

import torch
from sbi import utils as utils
from sbi import analysis as analysis
from sbi.inference import SNPE, prepare_for_sbi, simulate_for_sbi


def simulator(params):
    return params + torch.randn(params.shape) * 0.1


# data
prior = utils.BoxUniform(low=-2*torch.ones(3), high=2*torch.ones(3))
observation = torch.zeros(3)

# learning the density estimator and building the posterior
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNPE(prior)
theta, x = simulate_for_sbi(simulator, proposal=prior, num_simulations=500)
density_estimator = inference.append_simulations(theta, x).train()
og_posterior = inference.build_posterior(density_estimator)

# sampling from the posterior
samples = og_posterior.sample((10000,), x=observation)
log_probability = og_posterior.log_prob(samples, x=observation)
_ = analysis.pairplot(samples)

# save the density estimator state dictionary
torch.save(og_posterior.net.state_dict(), 'psd.pth')

# build new "empty" posterior
new_posterior = inference.build_posterior()

# load the state dictionary
new_posterior.net.load_state_dict(torch.load('psd.pth'))

# sampling from the new posterior
samples = new_posterior.sample((10000,), x=observation)
log_probability = new_posterior.log_prob(samples, x=observation)
_ = analysis.pairplot(samples)