Combine independent “analyses” into a joint distribution¶

The basic components in JMCTF are “analyses”. These are objects describing the joint PDF of a single “analysis” or “experiment”, for example an experiment searching for a signal in a series of Poisson “bins” or signal regions, or simply a normal random variable representing an experimental measurement. But these “analysis” objects do more than just describe the PDF of the experiment; they are responsible for constructing the TensorFlow probability model for the experiment (basically the PDF), for providing good starting guesses for distribution parameters based on sample realizations, for scaling parameters such that their MLEs have approximately unit variance (to help out the TensorFlow minimisation routines), for providing Asimov samples, and for keeping track of the sample structure. Extending JMCTF with new analysis objects is somewhat involved so it is covered separately.

JMCTF is designed for jointly analysing many independent experiments simultaneously. So to begin, let us create two simple experiments. First, a simple normal random variable, characterised by a mean and standard deviation. This can be created from the NormalAnalysis class:

>>> from JMCTF import NormalAnalysis, BinnedAnalysis, JointDistribution
>>> norm = NormalAnalysis("Test normal", 5, 2)

where the arguments are simply a name for the experiment, followed by an “observed” value, and the standard deviation of the normal random variable. The mean is treated as a free parameter so it is not specified. The “observed” value will only be used for final p-value computations, after all simulations are done, so it does not affect the simulations themselves.

Next, an experiment consisting of several independent Poisson random variables, whose means are characterised in terms of a “signal” parameter (either fixed, or to be fitted) and a second “background” parameter constrained by a normally distributed control measurement:

>>> bins = [("SR1", 10, 9, 2),
...         ("SR2", 50, 55, 4)]
>>> binned = BinnedAnalysis("Test binned", bins)

where the bin data consists of a name, an “observed” value, an expected background parameter, and an uncertainty parameter for the background. A full description is given in the BinnedAnalysis class documentation.

Not much can be done with these Analysis objects on their own; they act mainly as “blueprints” for experiments, to be used by other classes. The simplest such class is JointDistribution, which can do such things as jointly fit the parameters of many Analysis classes, and sample from the underlying distribution functions [1].

Creating a joint PDF object with JointDistribution is as simple as providing a list of component Analysis objects to the constructor:

>>> joint = JointDistribution([norm,binned])

Fitting the joint PDF to some data requires that the data be provided in an appropriately structured dictionary, so that samples can be assigned to the correct Analysis. The correct structure to use can be revealed by the get_sample_structure method:

>>> joint.get_sample_structure()
{'Test normal::x': 1, 'Test normal::x_theta': 1, 'Test binned::n': 2, 'Test binned::x': 2}

Here the dictionary keys are in the format analysis_name::random_variable_name and the values give the dimension of those random variables. For example we defined two bins in our BinnedAnalysis, so the dimension of the bin count random variable Test binned::n, and control variable Test binned::x, is 2.

Knowing this information, we can manually construct a sample for the full joint PDF and fit the joint PDF to it:

>>> my_sample = {'Test normal::x': 4.3, \
...              'Test normal::x_theta': 0, \
...              'Test binned::n': [9,53], \
...              'Test binned::x': [0,0]}

>>> q, joint_fitted, par_dicts = joint.fit_all(my_sample)
>>> print(par_dicts["all"])
{'Test normal':
 {'mu': <tf.Variable 'mu:0' shape=() dtype=float32, numpy=4.3>,
  'theta': <tf.Variable 'theta:0' shape=() dtype=float32, numpy=0.0>,
  'sigma_t': <tf.Tensor: id=57, shape=(), dtype=float32, numpy=0.0>},
'Test binned':
 {'s': <tf.Variable 's:0' shape=(2,) dtype=float32, numpy=array([ 0.        , -0.23735635], dtype=float32)>,
  'theta': <tf.Variable 'theta:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>}}

The output is not so pretty because the parameters are TensorFlow objects. We can convert them to numpy with the common.to_numpy() function for better viewing:

>>> from JMCTF.common import to_numpy
>>> print(to_numpy(par_dicts["all"]))
{'Test normal':
  {'mu': 4.3, 'theta': 0.0, 'sigma_t': 0.0},
 'Test binned':
  {'s': array([ 0.        , -0.23735635], dtype=float32),
   'theta': array([0., 0.], dtype=float32)}}

The fit_all method fits all the free parameters in the full joint distribution to the supplied samples, and returns q (negative 2 times the log probability density for the samples under the fitted model(s)), a new JointDistribution object fitted to the samples (the original remains unchanged), and three dictionaries of parameter values (“all” parameters, just the “fitted” parameters”, and just the “fixed” parameters) packed into a dictionary.

Here we have only fit the PDF to one sample, however the power of JMCTF really lies in fitting the PDF to lots of samples quickly. Of course manually creating lots of samples is tedious and not very useful; a more standard workflow is to actually sample the samples from the PDF under some fixed “null hypothesis” parameter values. We cover this in Sample from the joint distribution.

Footnotes

(Note, the crappy formatting of these footnotes is fixed in later versions of the sphinx_rtd_theme, and should go away as soon as readthedocs adopt a new release of it)

[1]	As you might guess from the name, the `JointDistribution` class inherits from `tensorflow_probability.JointDistributionNamed`. So all the standard log probability and sampling methods etc. from tensorflow_probability are available.