Feature/MCMC (#9)

* Add MCMC sampling for I3Module

* MCMC sampling: restructure and perform sampling in trafo space per default

* Add verbose keyword to reconstruction and MCMC sampling trays

* ic3-module: do not fill in missing seed values by default. Add option to I3Module settings

* MCMC sampling: write out quantiles and median

egenerator/ic3/reconstruction.py

@@ -1,5 +1,6 @@
 import os
 import tensorflow as tf
+import numpy as np
 
 from icecube import dataclasses, icetray
 from icecube.icetray.i3logging import log_info, log_warn
@@ -81,6 +82,11 @@ def __init__(self, context):
                           'the posterior for the sampling, otherwise the '
                           'samples will simply be set to the best fit point.',
                           False)
+        self.AddParameter('add_mcmc_samples',
+                          'Add samples from a Markov-Chain-Monte-Carlo. '
+                          'Settings for MCMC are defined in key '
+                          '`mcmc_settings`.',
+                          False)
         self.AddParameter('label_key',
                           'Only relevant if labels are being loaded. '
                           'The key from which to load labels.',
@@ -134,6 +140,22 @@ def __init__(self, context):
                           'parameters will default to '
                           '`minimize_parameter_default_value`.',
                           True)
+        self.AddParameter('missing_seed_value_dict',
+                          'If a certain model parameter does not exist in the '
+                          'frame key set via `seed_keys`, a default value may '
+                          'be specified with the `missing_seed_value_dict` '
+                          'dictionary. Entries have the format: '
+                          r'{parameter_name}: default_value.',
+                          {})
+        self.AddParameter('missing_seed_value',
+                          'If a model parameter key is not given in the '
+                          'provided seed_key, it can be replaced with a '
+                          'default value defined in `missing_seed_value`. '
+                          'Note that this value will be set for any and all '
+                          'parameters that are not found. Keep in mind that '
+                          'this will also work if the wrong frame key is '
+                          'provided by accident to `seed_keys`.',
+                          None)
         self.AddParameter('reco_optimizer_interface',
                           'The reconstruction interface to use. Options are: '
                           '"scipy" or "tfp" (tensorflow_probability).',
@@ -162,6 +184,22 @@ def __init__(self, context):
                               'add_per_dom_calculation': True,
                               'normalize_by_total_charge': False,
                           })
+        self.AddParameter('mcmc_settings',
+                          'Only relevant if `add_mcmc_samples` is set '
+                          ' to "True". Defines settings for MCMC sampling.',
+                          {
+                            'mcmc_num_chains': 10,
+                            'mcmc_method': 'HamiltonianMonteCarlo',
+                            'mcmc_num_results': 1000,
+                            'mcmc_num_burnin_steps': 100,
+                            'mcmc_num_steps_between_results': 0,
+                            'mcmc_num_parallel_iterations': 1,
+                          })
+        self.AddParameter('mcmc_quantiles',
+                          'Only relevant if `add_mcmc_samples` is set '
+                          ' to "True". Defines quantiles of MCMC samples to '
+                          'add to result frame key of reconstruction.',
+                          [0.5, 0.68, 0.9])
 
     def Configure(self):
         """Configures Module and loads model from file.
@@ -176,6 +214,7 @@ def Configure(self):
         self.add_circular_err = self.GetParameter('add_circular_err')
         self.add_covariances = self.GetParameter('add_covariances')
         self.add_goodness_of_fit = self.GetParameter('add_goodness_of_fit')
+        self.add_mcmc_samples = self.GetParameter('add_mcmc_samples')
         self.label_key = self.GetParameter('label_key')
         self.snowstorm_key = self.GetParameter('snowstorm_key')
         self.num_threads = self.GetParameter('num_threads')
@@ -195,6 +234,12 @@ def Configure(self):
         self.tf_optimizer_settings = self.GetParameter('tf_optimizer_settings')
         self.goodness_of_fit_settings = \
             self.GetParameter('goodness_of_fit_settings')
+        self.mcmc_settings = self.GetParameter('mcmc_settings')
+        self.mcmc_quantiles = self.GetParameter('mcmc_quantiles')
+
+        self.missing_seed_value_dict = \
+            self.GetParameter('missing_seed_value_dict')
+        self.missing_seed_value = self.GetParameter('missing_seed_value')
 
         if 'reconstruct_samples' not in self.goodness_of_fit_settings:
             self.goodness_of_fit_settings['reconstruct_samples'] = True
@@ -236,6 +281,8 @@ def Configure(self):
             additional_loss_modules=additional_loss_modules,
             misc_setting_updates={
                 'seed_names': self.seed_keys,
+                'missing_value_dict': self.missing_seed_value_dict,
+                'missing_value': self.missing_seed_value,
             },
             label_setting_updates={
                 'label_key': self.label_key,
@@ -280,6 +327,12 @@ def Configure(self):
                               range(self.manager.models[0].num_parameters)]
         for name, value in self.minimize_parameter_dict.items():
             fit_parameter_list[self.manager.models[0].get_index(name)] = value
+        self.fit_parameter_list = fit_parameter_list
+
+        self.fitted_parameters = []
+        for i, n in enumerate(self.manager.models[0].parameter_names):
+            if self.fit_parameter_list[i]:
+                self.fitted_parameters.append(n)
 
         # parameter input signature
         parameter_tensor_name = 'x_parameters'
@@ -360,6 +413,19 @@ def Configure(self):
                 scipy_optimizer_settings=self.scipy_optimizer_settings,
             )
 
+        # add MCMC
+        if self.add_mcmc_samples:
+            self.reco_tray.add_module(
+                'MarkovChainMonteCarlo',
+                name='MarkovChainMonteCarlo',
+                fit_parameter_list=fit_parameter_list,
+                seed_tensor_name='reco',
+                reco_key='reco',
+                minimize_in_trafo_space=self.minimize_in_trafo_space,
+                parameter_tensor_name=parameter_tensor_name,
+                **self.mcmc_settings
+            )
+
     def Physics(self, frame):
         """Apply Event-Generator model to physics frames.
 
@@ -484,6 +550,41 @@ def Physics(self, frame):
             result_dict['runtime_circular_err'] = float(
                 results['CircularizedAngularUncertainty']['runtime'])
 
+        # write MCMC samples to frame
+        if self.add_mcmc_samples:
+            mcmc_res = results['MarkovChainMonteCarlo']
+            num_accepted = len(mcmc_res['log_prob_values'])
+            result_dict['MCMC_acceptance_ratio'] = mcmc_res['acceptance_ratio']
+
+            # fill in median and quantiles
+            for i, n in enumerate(self.fitted_parameters):
+
+                if num_accepted > 0:
+                    values = mcmc_res['samples'][:, i]
+                else:
+                    # create some dummy values so that we fill in NaNs
+                    values = np.ones(10) * float('nan')
+
+                result_dict['MCMC_{}_median'.format(n)] = float(
+                    np.median(values))
+                for q in self.mcmc_quantiles:
+                    result_dict['MCMC_{}_q{:3.3f}_lower'.format(n, q)] = float(
+                        np.quantile(values, 0.5 - 0.5*q))
+                    result_dict['MCMC_{}_q{:3.3f}_upper'.format(n, q)] = float(
+                        np.quantile(values, 0.5 + 0.5*q))
+
+            if num_accepted > 0:
+                # create vectors for output quantities
+                vectors = {}
+                for i, n in enumerate(self.fitted_parameters):
+                    vectors[n] = dataclasses.I3VectorFloat(
+                        mcmc_res['samples'][:, i])
+                vectors['log_prob_values'] = dataclasses.I3VectorFloat(
+                    mcmc_res['log_prob_values'])
+
+                for n, vector in vectors.items():
+                    frame[self.output_key + '_MCMC_' + n] = vector
+
         # save to frame
         frame[self.output_key] = result_dict
         if self.i3_mapping: