text: import os
  import shutil
  import math
  import torch
  import numpy as np
  import matplotlib.pyplot as plt
  
  import pyprob
  from pyprob import Model, InferenceEngine
  from pyprob.distributions import Uniform, Normal
  plt.switch_backend('agg')
  
  
  class GaussianWithUnknownMeanMarsaglia(Model):
      def __init__(self, prior_mean=1, prior_stddev=math.sqrt(5), likelihood_stddev=math.sqrt(2), replace=True, *args, **kwargs):
          self.prior_mean = prior_mean
          self.prior_stddev = prior_stddev
          self.likelihood_stddev = likelihood_stddev
          self.replace = replace
          super().__init__('Gaussian with unknown mean (Marsaglia)', *args, **kwargs)
  
      def marsaglia(self, mean, stddev):
          uniform = Uniform(-1, 1)
          s = 1
          i = 0
          while True:
              x = pyprob.sample(uniform, replace=self.replace)
              y = pyprob.sample(uniform, replace=self.replace)
              s = x*x + y*y
              i += 1
              if float(s) < 1:
                  pyprob.tag(x, name='x_accepted')
                  pyprob.tag(y, name='y_accepted')
                  pyprob.tag(s, name='s_accepted')
                  break
              else:
                  pyprob.tag(x, name='x_rejected')
                  pyprob.tag(y, name='y_rejected')
                  pyprob.tag(s, name='s_rejected')
          pyprob.tag(i, name='iterations')
          return mean + stddev * (x * torch.sqrt(-2 * torch.log(s) / s))
  
      def forward(self):
          mu = self.marsaglia(self.prior_mean, self.prior_stddev)
          likelihood = Normal(mu, self.likelihood_stddev)
          pyprob.tag(mu, name='mu')
          pyprob.observe(likelihood, name='obs0')
          pyprob.observe(likelihood, name='obs1')
          return mu
  
  
  def produce_results(results_dir):
      train_traces_max = 100000
      train_traces_resolution = 6
      infer_traces = 1000
      train_traces_step = int(train_traces_max / (train_traces_resolution - 1))
      observes = [{'obs0': 1, 'obs1': 1}, {'obs0': 3, 'obs1': 4}, {'obs0': 8, 'obs1': 9}]
  
      if os.path.exists(results_dir):
          shutil.rmtree(results_dir)
      pyprob.util.create_path(results_dir, directory=True)
  
      model_replace_true = GaussianWithUnknownMeanMarsaglia(replace=True)
      model_replace_false = GaussianWithUnknownMeanMarsaglia(replace=False)
  
      traces = []
      model_replace_true_loss = []
      model_replace_false_loss = []
  
      model_replace_true_is_posterior_ess = []
      model_replace_false_is_posterior_ess = []
      for observe in observes:
          model_replace_true_is_posterior = model_replace_true.posterior_results(infer_traces, observe=observe)
          model_replace_false_is_posterior = model_replace_false.posterior_results(infer_traces, observe=observe)
          model_replace_true_is_posterior_ess.append(float(model_replace_true_is_posterior.effective_sample_size) / infer_traces)
          model_replace_false_is_posterior_ess.append(float(model_replace_false_is_posterior.effective_sample_size) / infer_traces)
  
      model_replace_true_ic_posterior_ess = np.zeros([len(observes), train_traces_resolution])
      model_replace_false_ic_posterior_ess = np.zeros([len(observes), train_traces_resolution])
      model_replace_true.learn_inference_network(0, batch_size=1, observe_embeddings={'obs0': {}, 'obs1': {}}, inference_network=pyprob.InferenceNetwork.LSTM)
      model_replace_false.learn_inference_network(0, batch_size=1, observe_embeddings={'obs0': {}, 'obs1': {}}, inference_network=pyprob.InferenceNetwork.LSTM)
  
      for j in range(train_traces_resolution):
          train_traces = j * train_traces_step
          print('\ntrain_traces: {}/{}'.format(train_traces, train_traces_max))
          model_replace_true_loss.append(float(model_replace_true._inference_network._history_train_loss[-1]))
          model_replace_false_loss.append(float(model_replace_false._inference_network._history_train_loss[-1]))
          for i, observe in enumerate(observes):
              model_replace_true_ic_posterior = model_replace_true.posterior_results(infer_traces, observe=observe, inference_engine=InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK)
              model_replace_false_ic_posterior = model_replace_false.posterior_results(infer_traces, observe=observe, inference_engine=InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK)
              model_replace_true_ic_posterior_ess[i, j] = float(model_replace_true_ic_posterior.effective_sample_size) / infer_traces
              model_replace_false_ic_posterior_ess[i, j] = float(model_replace_false_ic_posterior.effective_sample_size) / infer_traces
  
          model_replace_true.learn_inference_network(train_traces_step)
          model_replace_false.learn_inference_network(train_traces_step)
          traces.append(train_traces)
  
      for i, observe in enumerate(observes):
          plot_file_name = os.path.join(results_dir, 'obs_{}_{}.pdf'.format(observe['obs0'], observe['obs1']))
          print('Saving result plot to: {}'.format(plot_file_name))
          fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 8))
          ax1.hlines(model_replace_true_is_posterior_ess[i], xmin=0, xmax=train_traces_max, label='IS, Replace=True')
          ax1.hlines(model_replace_false_is_posterior_ess[i], xmin=0, xmax=train_traces_max, label='IS, Replace=False')
          ax1.plot(traces, model_replace_true_ic_posterior_ess[i], label='IC, Replace=True')
          ax1.plot(traces, model_replace_false_ic_posterior_ess[i], label='IC, Replace=False')
          ax1.legend(loc='best')
          ax1.set_ylabel('Normalized ESS')
          ax2.plot(traces, model_replace_true_loss, label='Replace=True')
          ax2.plot(traces, model_replace_false_loss, label='Replace=False')
          ax2.set_ylabel('Loss')
          ax2.set_xlabel('Training traces')
          ax2.legend(loc='best')
          plt.tight_layout()
          fig.savefig(plot_file_name)
  
  
  if __name__ == '__main__':
      pyprob.set_random_seed(1)
  
      current_dir = os.path.dirname(os.path.abspath(__file__))
      print('Current dir: {}'.format(current_dir))
  
      results_dir = os.path.join(current_dir, 'gum_marsaglia')
      produce_results(results_dir=results_dir)
  
      print('Done')
  