Complete tempered likelihood implementation

[simeoncarstens]

Currently, Chainsail supports only a single family of tempered distributions, namely the Boltzmann distribution / canonical ensemble-like $p(x|\beta) \propto p(x)^\beta$. To make Chainsail more appealing for Bayesian statistics use cases, we want to support a similar tempered distribution family, but that only applies a temperature to the likelihood of a posterior distribution. This looks as follows:

$$ p(x|D, \beta) \propto p(D|x)^\beta p(x) $$

I started implementing this a while ago: we already have a class LikelihoodTemperedPosterior, defined in https://github.com/tweag/resaas/blob/main/lib/common/chainsail/common/tempering/tempered_distributions.py#L108 that implements this family of distributions.

To make this fully work, the following things have to be taken care of (roughly in order):

• extend the job specification (in https://github.com/tweag/resaas/blob/main/lib/common/chainsail/common/spec.py) to accomodate the new tempered distribution family, named something useful such as "likelihood_tempered". This involves
  • extending the Enum (https://github.com/tweag/resaas/blob/main/lib/common/chainsail/common/spec.py#L217),
  • specifying an initial schedule parameter schema for the new tempered distribution in https://github.com/tweag/resaas/blob/main/lib/common/chainsail/common/spec.py#L229; likely the existing BoltzmannInitialScheduleParametersSchema can be reused,
  • adding the new parameter schema (if not reusing a previous one) to the list of allowed types for the initial_schedule_parameters field in the JobSpec (https://github.com/tweag/resaas/blob/main/lib/common/chainsail/common/spec.py#L301),
• make the controller able to work with the new tempered distribution:
  • add an argument tempered_distribution_family to the constructors in https://github.com/tweag/resaas/blob/main/app/controller/chainsail/controller/__init__.py#L126 / https://github.com/tweag/resaas/blob/main/app/controller/chainsail/controller/__init__.py#L385 and make a field that holds its value,
  • in the run scripts in https://github.com/tweag/resaas/blob/main/lib/runners/rexfw/chainsail/runners/rexfw, pass the tempered distribution family value from the JobSpec to the controller (e.g., in https://github.com/tweag/resaas/blob/main/app/controller/chainsail/controller/run_local.py#L83),
  • pass that value to _config_template_from_params in https://github.com/tweag/resaas/blob/main/app/controller/chainsail/controller/__init__.py#L279 and add it to the job config dictionary in that function,
  • allow instantiation of optimization-related objects (https://github.com/tweag/resaas/blob/main/app/controller/chainsail/controller/__init__.py#L73 ) for the new tempered distribution,
• in the MPI rexfw runner, replace the hardcoding of the BoltzmannTemperedDistribution (https://github.com/tweag/resaas/blob/main/lib/runners/rexfw/chainsail/runners/rexfw/mpi.py#L210) by a choice of the correct tempered distribution based on the config file entry.

At this point, you should be able to run the controller locally (https://github.com/tweag/resaas/blob/main/app/controller/README.md#L30), specifying your new tempered distribution family in a job.json file.

A final step is then changing the front-end JavaScript code (https://github.com/tweag/resaas/blob/main/app/client/pages/job.js) to add a dropdown field that allows the user to choose between the existing Boltzmann tempered distribution and the likelihood-tempered distribution in the job submission form.

[simeoncarstens]

All done via #377 and #383.