<jupyter_start><jupyter_text>Defining the modelFirst, 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.<jupyter_code>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
        pyprob.observe(likelihood, name='obs1')

        # return the latent quantity of interest
        return mu
    
model = GaussianUnknownMean()<jupyter_output><empty_output><jupyter_text>Finding the correct posterior analyticallySince 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.<jupyter_code>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')

correct_dists = CorrectDistributions(model)<jupyter_output><empty_output><jupyter_text>Prior distributionWe 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.<jupyter_code>prior = model.prior_distribution(num_traces=1000)<jupyter_output>Time spent  | Time remain.| Progress             | Trace     | Traces/sec
0d:00:00:00 | 0d:00:00:00 | #################### | 1000/1000 | 1,293.24<jupyter_text>We can plot a historgram of these samples that are held by the `Empirical` distribution.<jupyter_code>prior.plot_histogram(show=False, alpha=0.75, label='emprical prior')
correct_dists.plot_prior(min(prior.values_numpy()),max(prior.values_numpy()))
plt.legend();<jupyter_output><empty_output><jupyter_text>Posterior inference with importance samplingFor a given set of observations, we can get samples from the posterior distribution.<jupyter_code>correct_dists.observed_list = [8, 9] # 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
                                         )<jupyter_output>Time spent  | Time remain.| Progress             | Trace     | Traces/sec
0d:00:00:02 | 0d:00:00:00 | #################### | 5000/5000 | 1,840.04<jupyter_text>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 with the prior distribution.<jupyter_code>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();<jupyter_output><empty_output><jupyter_text>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).<jupyter_code>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();<jupyter_output><empty_output><jupyter_text>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:<jupyter_code>print(posterior.sample())
print(posterior.mean)
print(posterior.stddev)
print(posterior.expectation(lambda x: torch.sin(x)))<jupyter_output>tensor(0.4846)<jupyter_text>Inference compilationInference 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.<jupyter_code>model.learn_inference_network(num_traces=20000,
                              observe_embeddings={'obs0' : {'dim' : 32},
                                                  'obs1': {'dim' : 32}},
                              inference_network=pyprob.InferenceNetwork.LSTM)<jupyter_output>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:33 | 1    | 20,032    | +2.40e+00 | +1.10e+00 | [32m+1.26e+00[0m | 0d:00:00:05 | 1,003.0<jupyter_text>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.<jupyter_code># 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
                                         )
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();<jupyter_output><empty_output><jupyter_text>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.<jupyter_code>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();<jupyter_output><empty_output><jupyter_text>Inference compilation performs amortized inferece 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.<jupyter_code>correct_dists.observed_list = [12, 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]}
                                         )
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();
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();<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09317844784972813, 'fraction_numerical': 0.02051408798813643, 'mean_word_length': 3.6598848368522074, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 48, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '129726', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 3580, 'n_tokens_neox': 3266, 'n_words': 949, 'path': 'examples/gaussian_unknown_mean.ipynb', 'repo_name': 'ammunk/pyprob', 'size': '12138'}