```python
import pyprob
from pyprob import Model
from pyprob import InferenceEngine
from pyprob.distributions import Normal
from pyprob.dis import ModelDIS

import torch
import numpy as np
import math
import matplotlib.pyplot as plt
%matplotlib inline
fig = plt.figure();
```


    <Figure size 432x288 with 0 Axes>


# Defining the Model


```python
class GaussianUnknownMean(Model):
    def __init__(self):
        super().__init__(name='Gaussian with unknown mean') # give the model a name
        self.prior_mean = 1
        self.prior_std = math.sqrt(5)
        self.likelihood_std = math.sqrt(2)

    def forward(self): # Needed to specifcy how the generative model is run forward
        # sample the (latent) mean variable to be inferred:
        mu = pyprob.sample(Normal(self.prior_mean, self.prior_std)) # NOTE: sample -> denotes latent variables

        # define the likelihood
        likelihood = Normal(mu, self.likelihood_std)

        # Lets add two observed variables
        # -> the 'name' argument is used later to assignment values:
        #pyprob.observe(likelihood, name='obs0') # NOTE: observe -> denotes observable variables

        # return the latent quantity of interest
        return mu
    
class GaussianDIS(GaussianUnknownMean, ModelDIS):
    def __init__(self):
        super().__init__()
        self.obs = 8
        #self.bool_func = bool_func
    def dummy_bernoulli(self, x):
        return True
    def forward(self):
        x = super().forward()
        likelihood = Normal(x, self.likelihood_std) 
        sample_obs = pyprob.sample(likelihood)
        euclidean_dist = np.abs(sample_obs - self.obs)
        delta_val = self.dummy_bernoulli(x)
        bool_func_dist = pyprob.distributions.Bernoulli(delta_val)
        pyprob.observe(bool_func_dist, name = "bool_func")
        dummy = pyprob.distributions.Normal(euclidean_dist,0.1)
        pyprob.observe(dummy, euclidean_dist, name = "distance") # Any other way to do this?
        return x

model = GaussianUnknownMean()
simulator = GaussianDIS()
```

# Finding the correct posterior analytically
Since all distributions are gaussians in this model, we can analytically compute the posterior and we can compare the true posterior to the inferenced one.

Assuming that the prior and likelihood are $p(x) = \mathcal{N}(\mu_0, \sigma_0)$ and $p(y|x) = \mathcal{N}(x, \sigma)$ respectively and, $y_1, y_2, \ldots y_n$ are the observed values, the posterior would be $p(x|y) = \mathcal{N}(\mu_p, \sigma_p)$ where,
$$
\begin{align}
\sigma_{p}^{2} & = \frac{1}{\frac{n}{\sigma^2} + \frac{1}{\sigma_{0}^{2}}} \\
\mu_p & = \sigma_{p}^{2} \left( \frac{\mu_0}{\sigma_{0}^{2}} + \frac{n\overline{y}}{\sigma^2} \right)
\end{align}
$$
The following class implements computing this posterior distribution. We also implement some helper functions and variables for plotting the correct posterior and prior.


```python
def plot_function(min_val, max_val, func, *args, **kwargs):
        x = np.linspace(min_val,max_val,int((max_val-min_val)*50))
        plt.plot(x, np.vectorize(func)(x), *args, **kwargs)

def get_dist_pdf(dist):
    return lambda x: math.exp(dist.log_prob(x))
        
class CorrectDistributions:
    def __init__(self, model):
        self.prior_mean = model.prior_mean
        self.prior_std = model.prior_std
        self.likelihood_std = model.likelihood_std
        self.prior_dist = Normal(self.prior_mean, self.prior_std)
        
    @property
    def observed_list(self):
        return self.__observed_list

    @observed_list.setter
    def observed_list(self, new_observed_list):
        self.__observed_list = new_observed_list
        self.construct_correct_posterior()
    
    def construct_correct_posterior(self):
        n = len(self.observed_list)
        posterior_var = 1/(n/self.likelihood_std**2 + 1/self.prior_std**2)
        posterior_mu = posterior_var * (self.prior_mean/self.prior_std**2 + n*np.mean(self.observed_list)/self.likelihood_std**2)
        self.posterior_dist = Normal(posterior_mu, math.sqrt(posterior_var))

    def prior_pdf(self, model, x):
        p = Normal(model.prior_mean,model.prior_stdd)
        return math.exp(p.log_prob(x))

    def plot_posterior(self, min_val, max_val):
        if not hasattr(self, 'posterior_dist'):
            raise AttributeError('observed values are not set yet, and posterior is not defined.')
        plot_function(min_val, max_val, get_dist_pdf(self.posterior_dist), label='correct posterior', color='orange')


    def plot_prior(self, min_val, max_val):
        plot_function(min_val, max_val, get_dist_pdf(self.prior_dist), label='prior', color='green')

```


```python
correct_dists = CorrectDistributions(model)
correct_dists.observed_list=[8]
```

# Training

For a sanity check, start with one training round.


```python
simulator.train(iterations = 1,
    importance_sample_size=5000, 
    proposal_mixture_components=5,
    observe_embeddings={'bool_func': {'dim': 1, 'depth': 1}} # Dummy value as we currently have to observe something
)
```

    Creating new inference network...
    Observable bool_func: reshape not specified, using shape torch.Size([]).
    Observable bool_func: using embedding dim torch.Size([1]).
    Observable bool_func: observe embedding not specified, using the default FEEDFORWARD.
    Observable bool_func: using embedding depth 1.
    Observe embedding dimension: 1
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    New layers, address: 16__forward__mu__Normal__1, distribution: Normal
    New layers, address: 28__forward__sample_obs__Normal__1, distribution: Normal
    Total addresses: 2, parameters: 308
    0d:00:00:00 | 1    | 300       | +4.03e+00 | [1m[32m+3.95e+00[0m | [1m[31m+4.16e+00[0m | [1m[32m0d:00:00:00[0m | +1.00e-03 | 1,006.0 
    Stop condition reached. num_traces: 500
    0d:00:00:00 | 1    | 500       | +4.03e+00 | +3.95e+00 | [31m+4.10e+00[0m | 0d:00:00:00 | +1.00e-03 | 1,163.2                               
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:20 | 0d:00:00:00 | #################### | 5000/5000 | 4807.00 | 240.37       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:00 | 0d:00:00:00 | #################### | 500/500 | 1,032.81       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:00 | 1    | 900       | +4.03e+00 | +3.95e+00 | [32m+4.99e+00[0m | 0d:00:00:00 | +1.00e-03 | 770.5 
    Stop condition reached. num_traces: 500
    0d:00:00:01 | 1    | 1,000     | +4.03e+00 | +3.95e+00 | [31m+5.14e+00[0m | 0d:00:00:00 | +1.00e-03 | 852.5 
    Training iterations 1  epsilon 1.92  ESS 500.0



```python
prior = simulator.prior_results(num_traces = 2000)
posterior_check = simulator.posterior(num_traces = 2000,
inference_engine = InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK,
observe={"bool_func": 1})
posterior_check = simulator.update_DIS_posterior_weights(posterior_check)
results_check = np.array([x.result for x in posterior_check.values])
```

    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:03 | 0d:00:00:00 | #################### | 2000/2000 | 2000.00 | 643.56       
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:08 | 0d:00:00:00 | #################### | 2000/2000 | 1920.54 | 230.02       


Now do 50 more iterations.


```python
simulator.train(iterations = 50,
    importance_sample_size=5000, 
    proposal_mixture_components=5,
    observe_embeddings={'bool_func': {'dim': 1, 'depth': 1}} # Dummy value as we currently have to observe something
)
```

    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:18 | 0d:00:00:00 | #################### | 5000/5000 | 4801.09 | 271.60       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:00 | 0d:00:00:00 | #################### | 500/500 | 1,204.59       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:01 | 1    | 1,300     | +4.03e+00 | +3.95e+00 | [32m+4.96e+00[0m | 0d:00:00:00 | +1.00e-03 | 439.9   
    Stop condition reached. num_traces: 500
    0d:00:00:01 | 1    | 1,500     | +4.03e+00 | +3.95e+00 | [32m+5.01e+00[0m | 0d:00:00:00 | +1.00e-03 | 989.9   
    Training iterations 1  epsilon 1.30  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:20 | 0d:00:00:00 | #################### | 5000/5000 | 4817.46 | 249.91       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:00 | 0d:00:00:00 | #################### | 500/500 | 829.38       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:02 | 1    | 1,800     | +4.03e+00 | +3.95e+00 | [32m+4.99e+00[0m | 0d:00:00:00 | +1.00e-03 | 693.1 
    Stop condition reached. num_traces: 500
    0d:00:00:02 | 1    | 2,000     | +4.03e+00 | +3.95e+00 | [31m+5.02e+00[0m | 0d:00:00:00 | +1.00e-03 | 351.1 
    Training iterations 2  epsilon 1.05  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:57 | 0d:00:00:00 | #################### | 5000/5000 | 4819.09 | 87.53        
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:02 | 0d:00:00:00 | #################### | 500/500 | 234.92       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:04 | 1    | 2,400     | +4.03e+00 | +3.95e+00 | [31m+4.99e+00[0m | 0d:00:00:02 | +1.00e-03 | 171.6 
    Stop condition reached. num_traces: 500
    0d:00:00:05 | 1    | 2,500     | +4.03e+00 | +3.95e+00 | [31m+5.05e+00[0m | 0d:00:00:02 | +1.00e-03 | 173.9 
    Training iterations 3  epsilon 0.90  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:01:19 | 0d:00:00:00 | #################### | 5000/5000 | 4805.33 | 62.69       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:02 | 0d:00:00:00 | #################### | 500/500 | 209.79       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:06 | 1    | 2,900     | +4.03e+00 | +3.95e+00 | [31m+4.99e+00[0m | 0d:00:00:01 | +1.00e-03 | 351.6 
    Stop condition reached. num_traces: 500
    0d:00:00:06 | 1    | 3,000     | +4.03e+00 | +3.95e+00 | [32m+4.83e+00[0m | 0d:00:00:01 | +1.00e-03 | 481.1 
    Training iterations 4  epsilon 0.79  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:55 | 0d:00:00:00 | #################### | 5000/5000 | 4784.43 | 90.89       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:01 | 0d:00:00:00 | #################### | 500/500 | 499.84       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:07 | 1    | 3,300     | +4.03e+00 | +3.95e+00 | [32m+4.96e+00[0m | 0d:00:00:00 | +1.00e-03 | 325.8 
    Stop condition reached. num_traces: 500
    0d:00:00:08 | 1    | 3,500     | +4.03e+00 | +3.95e+00 | [32m+4.80e+00[0m | 0d:00:00:01 | +1.00e-03 | 178.2 
    Training iterations 5  epsilon 0.71  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:01:11 | 0d:00:00:00 | #################### | 5000/5000 | 4758.18 | 69.80       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:01 | 0d:00:00:00 | #################### | 500/500 | 263.98       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:09 | 1    | 3,900     | +4.03e+00 | +3.95e+00 | [32m+4.77e+00[0m | 0d:00:00:01 | +1.00e-03 | 238.9 
    Stop condition reached. num_traces: 500
    0d:00:00:10 | 1    | 4,000     | +4.03e+00 | +3.95e+00 | [31m+4.90e+00[0m | 0d:00:00:01 | +1.00e-03 | 250.5 
    Training iterations 6  epsilon 0.65  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:01:08 | 0d:00:00:00 | #################### | 5000/5000 | 4746.82 | 73.42       
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:01 | 0d:00:00:00 | #################### | 500/500 | 415.50       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:11 | 1    | 4,400     | +4.03e+00 | +3.95e+00 | [31m+4.87e+00[0m | 0d:00:00:01 | +1.00e-03 | 353.2 
    Stop condition reached. num_traces: 500
    0d:00:00:12 | 1    | 4,500     | +4.03e+00 | +3.95e+00 | [32m+4.82e+00[0m | 0d:00:00:01 | +1.00e-03 | 358.4 
    Training iterations 7  epsilon 0.60  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:59 | 0d:00:00:00 | #################### | 5000/5000 | 4705.87 | 84.16        
    OfflineDataset at: .
    Num. traces      : 500
    Sorted on disk   : False
    No pre-computed hashes found, generating: ./pyprob_hashes
    Hashing offline dataset for sorting
    Time spent  | Time remain.| Progress             | Traces  | Traces/sec
    0d:00:00:01 | 0d:00:00:00 | #################### | 500/500 | 351.63       
    Sorting offline dataset
    Sorting done
    Num. trace types : 1
    Trace hash	Count
    2.94810057	500
    
    Continuing to train existing inference network...
    Total number of parameters: 308
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Learn.rate| Traces/sec
    0d:00:00:13 | 1    | 4,900     | +4.03e+00 | +3.95e+00 | [32m+4.70e+00[0m | 0d:00:00:01 | +1.00e-03 | 348.4 
    Stop condition reached. num_traces: 500
    0d:00:00:13 | 1    | 5,000     | +4.03e+00 | +3.95e+00 | [31m+4.86e+00[0m | 0d:00:00:01 | +1.00e-03 | 378.2 
    Training iterations 8  epsilon 0.55  ESS 500.0
    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:36 | 0d:00:00:36 | ##########---------- | 2521/5000 | 2344.68 | 68.41       


```python
posterior = simulator.posterior(num_traces = 2000,
inference_engine = InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK,
observe={"bool_func": 1})
posterior = simulator.update_DIS_posterior_weights(posterior)
results = np.array([x.result for x in posterior.values])
```

    Time spent  | Time remain.| Progress             | Trace     | ESS    | Traces/sec
    0d:00:00:06 | 0d:00:00:00 | #################### | 2000/2000 |  99.34 | 296.71       


# Plots

First do checks for single training round. This is to make sure that the algorithm isn't somehow 'cheating'


```python
posterior_check = simulator.update_DIS_posterior_weights(posterior_check)
results_check = np.array([x.result for x in posterior_check.values])
```


```python
plt.hist(results_check, bins=50, weights=np.array(posterior_check.weights), label='posterior',density=True)
plt.hist(results_check, bins=50, label='proposal',density = True)
correct_dists.plot_prior(min(prior.values_numpy()),max(prior.values_numpy()))
correct_dists.plot_posterior(min(results), max(results))
plt.legend()
```

Now check full results


```python
plt.hist(results, bins=50, weights=np.array(posterior.weights), label='posterior',density=True)
plt.hist(results, bins=50, label='proposal',density = True)
correct_dists.plot_prior(min(prior.values_numpy()),max(prior.values_numpy()))
correct_dists.plot_posterior(min(results), max(results))
plt.legend()
```

# Roughs


```python
importance_weights = np.exp(np.array([x.log_importance_weight for x in posterior.values]))
importance_weights = importance_weights/sum(importance_weights)
weights = posterior.weights
probs = posterior._categorical.probs
print(importance_weights[:10])
print(np.array(weights[:10].detach()))
print(np.array(probs[:10].detach()))
```

    [4.43934743e-04 1.12809762e-03 6.99781531e-05 3.08431125e-04
     2.07236835e-04 1.65965433e-04 2.79483327e-04 1.51536002e-03
     4.46125723e-04 3.05139416e-04]
    [5.16142642e-11 2.43493937e-16 3.47589046e-11 2.37514305e-08
     7.04764763e-05 8.20254264e-05 1.54813143e-05 2.09902897e-24
     4.43681561e-06 7.53436269e-09]
    [5.16142642e-11 2.43493937e-16 3.47589046e-11 2.37514305e-08
     7.04764763e-05 8.20254264e-05 1.54813143e-05 2.09902897e-24
     4.43681561e-06 7.53436269e-09]



```python
print(importance_weights.sum())
print(weights.sum())
```

    1.0000000000000004
    tensor(0.0158, dtype=torch.float64)



```python
plt.hist(results, bins=20, weights=importance_weights, label='posterior')
plt.hist(results, bins=20, label='proposal',density = True)
correct_dists.plot_prior(min(prior.values_numpy()),max(prior.values_numpy()))
correct_dists.plot_posterior(min(results), max(results))
plt.legend()
```

Why do importance weights bring the posterior back to the prior?



a01ec8da5cda1a41c2d1da384038e1550eb82218
ipynb
Jupyter Notebook
examples/Gaussian_DIS.ipynb
SRagy/pyprob
c8fc67d5f4ad78d565a86c2f86abe10b73549f18
BSD-2-Clause
examples/Gaussian_DIS.ipynb
SRagy/pyprob
c8fc67d5f4ad78d565a86c2f86abe10b73549f18
BSD-2-Clause
examples/Gaussian_DIS.ipynb
SRagy/pyprob
c8fc67d5f4ad78d565a86c2f86abe10b73549f18
BSD-2-Clause
Qwen/Qwen-72B
1. YES
2. YES