```python
import pyprob
from pyprob import Model
from pyprob.distributions import Normal

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


    <Figure size 432x288 with 0 Axes>


# Defining the model
First, we define the model as a probabilistic program inheriting from `pyprob.Model`. Models inherit from `torch.nn.Module` and can be potentially trained with gradient-based optimization (not covered in this example).

The `forward` function can have any number and type of arguments as needed.


```python
%matplotlib inline
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 specify 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
        pyprob.observe(likelihood, name='obs1')

        # return the latent quantity of interest
        return mu
    
model = GaussianUnknownMean()
```

# 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):
    '''
    Apply a function repetitively to tuples within a vector
    '''
    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)
        # As per the analytical method we construct the posterior given by the formula
        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')

correct_dists = CorrectDistributions(model)
```

# Prior distribution
We inspect the prior distribution to see if it behaves in the way we intended. First we construct an `Empirical` distribution with forward samples from the model.

Note: Extra arguments passed to `prior_distribution` will be forwarded to model's `forward` function.


```python
prior = model.prior_distribution(num_traces=1000)
```

    Time spent  | Time remain.| Progress             | Trace     | Traces/sec
    0d:00:00:00 | 0d:00:00:00 | #################### | 1000/1000 | 1,621.42       


We can plot a histogram of these samples that are held by the `Empirical` distribution.


```python
# Compare our assumed prior against the actual prior from the model - sanity check
prior.plot_histogram(show=False, alpha=0.75, label='empirical prior')
correct_dists.plot_prior(min(prior.values_numpy()),max(prior.values_numpy()))
plt.legend();
```

# Posterior inference with importance sampling
For a given set of observations, we can get samples from the posterior distribution.


```python
correct_dists.observed_list = [4, 5] # Observations
# sample from posterior (5000 samples)
posterior = model.posterior_distribution(
                                         num_traces=5000, # the number of samples estimating the posterior
                                         inference_engine=pyprob.InferenceEngine.IMPORTANCE_SAMPLING, # specify which inference engine to use
                                         observe={'obs0': correct_dists.observed_list[0],
                                                  'obs1': correct_dists.observed_list[1]} # assign values to the observed values
                                         )
```

    Time spent  | Time remain.| Progress             | Trace     | Traces/sec
    0d:00:00:02 | 0d:00:00:00 | #################### | 5000/5000 | 1,807.42       


Regular importance sampling uses proposals from the prior distribution. We can see this by plotting the histogram of the posterior distribution without using the importance weights. As expected, this is the same as the prior distribution.


```python
posterior_unweighted = posterior.unweighted()
posterior_unweighted.plot_histogram(show=False, alpha=0.75, label='empirical proposal')
correct_dists.plot_prior(min(posterior_unweighted.values_numpy()),
                         max(posterior_unweighted.values_numpy()))
correct_dists.plot_posterior(min(posterior_unweighted.values_numpy()),
                             max(posterior_unweighted.values_numpy()))
plt.legend();
```

When we do use the weights, we end up with the correct posterior distribution. The following shows the sampled posterior with the correct posterior (orange curve).


```python
posterior.plot_histogram(show=False, alpha=0.75, bins=50, label='inferred posterior')
correct_dists.plot_posterior(min(posterior.values_numpy()),
                             max(posterior_unweighted.values_numpy()))
plt.legend();
```

In practice, it is advised to use methods of the `Empirical` posterior distribution instead of dealing with the weights directly, which ensures that the weights are used in the correct way.

For instance, we can get samples from the posterior, compute its mean and standard deviation, and evaluate expectations of a function under the distribution:


```python
print(posterior.sample())
```

    tensor(4.1385)



```python
print(posterior.mean)
```

    tensor(3.9196)



```python
print(posterior.stddev)
```

    tensor(0.9020)



```python
print(posterior.expectation(lambda x: torch.sin(x)))
```

    tensor(-0.4711)


# Inference compilation
Inference compilation is a technique where a deep neural network is used for parameterizing the proposal distribution in importance sampling (https://arxiv.org/abs/1610.09900). This neural network, which we call inference network, is automatically generated and trained with data sampled from the model.

We can learn an inference network for our model.


```python
model.learn_inference_network(num_traces=20000,
                              observe_embeddings={'obs0' : {'dim' : 32},
                                                  'obs1': {'dim' : 32}},
                              inference_network=pyprob.InferenceNetwork.LSTM)
```

    Creating new inference network...
    Observable obs0: observe embedding not specified, using the default FEEDFORWARD.
    Observable obs0: embedding depth not specified, using the default 2.
    Observable obs1: observe embedding not specified, using the default FEEDFORWARD.
    Observable obs1: embedding depth not specified, using the default 2.
    Observe embedding dimension: 64
    Train. time | Epoch| Trace     | Init. loss| Min. loss | Curr. loss| T.since min | Traces/sec
    New layers, address: 16__forward__mu__Normal__1, distribution: Normal
    Total addresses: 1, distribution types: 1, parameters: 1,643,583
    0d:00:00:31 | 1    | 20,032    | +2.21e+00 | +1.15e+00 | [32m+1.27e+00[0m | 0d:00:00:10 | 638.0                              


We now construct the posterior distribution using samples from inference compilation, using the trained inference network.

A much smaller number of samples are enough (200 vs. 5000) because the inference network provides good proposals based on the given observations. We can see that the proposal distribution given by the inference network is doing a job much better than the prior, by plotting the posterior samples without the importance weights, for a selection of observations.


```python
# sample from posterior (200 samples)
posterior = model.posterior_distribution(
    num_traces=200, # the number of samples estimating the posterior
    inference_engine=pyprob.InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK, # specify which inference engine to use
    observe={'obs0': correct_dists.observed_list[0],
             'obs1': correct_dists.observed_list[1]} # assign values to the observed values
)
```

    Time spent  | Time remain.| Progress             | Trace   | Traces/sec
    0d:00:00:00 | 0d:00:00:00 | #################### | 200/200 | 316.34       



```python
posterior_unweighted = posterior.unweighted()
posterior_unweighted.plot_histogram(
    show=False, bins=50, alpha=0.75, label='empirical proposal')

correct_dists.plot_posterior(min(posterior.values_numpy()),
                             max(posterior.values_numpy()))
plt.legend();
```

We can see that the proposal distribution given by the inference network is already a good estimate to the true posterior which makes the inferred posterior a much better estimate than the prior, even using much less number of samples.


```python
posterior.plot_histogram(show=False, bins=50, alpha=0.75, label='inferred posterior')
correct_dists.plot_posterior(min(posterior.values_numpy()),
                             max(posterior.values_numpy()))
plt.legend();
```

Inference compilation performs amortized inference which means, the same trained network provides proposal distributions for any observed values.

We can try performing inference using the same trained network with different observed values.


```python
correct_dists.observed_list = [8, 10] # New observations

posterior = model.posterior_distribution(
    num_traces=200,
    inference_engine=pyprob.InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK, # specify which inference engine to use
    observe={'obs0': correct_dists.observed_list[0],
          'obs1': correct_dists.observed_list[1]}
)
```

    Time spent  | Time remain.| Progress             | Trace   | Traces/sec
    0d:00:00:00 | 0d:00:00:00 | #################### | 200/200 | 362.70       



```python
posterior_unweighted = posterior.unweighted()
posterior_unweighted.plot_histogram(show=False, bins=50, alpha=0.75, label='empirical proposal')

correct_dists.plot_posterior(min(posterior.values_numpy()),
                       max(posterior.values_numpy()))
plt.legend();
```


```python
posterior.plot_histogram(show=False, bins=50, alpha=0.75, label='inferred posterior')
correct_dists.plot_posterior(min(posterior.values_numpy()),
                       max(posterior.values_numpy()))
plt.legend();

```


```python

```

6487fd6ebac99520782bd977d89ca3ff676915e4
ipynb
Jupyter Notebook
examples/gaussian_unknown_mean.ipynb
SwapneelM/pyprob
4d93441ea838c3491a49050ae05d218a34708e6d
BSD-2-Clause
examples/gaussian_unknown_mean.ipynb
SwapneelM/pyprob
4d93441ea838c3491a49050ae05d218a34708e6d
BSD-2-Clause
examples/gaussian_unknown_mean.ipynb
SwapneelM/pyprob
4d93441ea838c3491a49050ae05d218a34708e6d
BSD-2-Clause
Qwen/Qwen-72B
1. YES
2. YES