Writing Cores and Likelihoods#

[4]:

import py21cmmc as p21c
from py21cmmc import mcmc

[2]:

p21c.__version__

[2]:

'0.1.0'

How do I write a Core Module?#

One of main aims of 21CMMC is to be extensible, so that the basic capabilities provided can be easily applied to new situations. For example, one use-case which may require a custom Core module (the concept of Core and Likelihood modules is introduced in the MCMC Introduction tutorial) is where the “observation” which is being analyzed is not a theoretical simulation box, but a set of interferometric visibilities at predefined baselines.

Since our philosophy is that a Core module should “construct” a model such that it approximates the pre-reduced “observations”, the conversion of the 21cmFAST cube into a set of visibilities should be defined in a new Core.

In principle, the Core module does not need to be subclassed from any particular class, so long as it implements a minimal API. However, we recommend always subclassing from py21cmmc.mcmc.CoreBase, which sets default methods for several of these API components (some of which should almost never have to be changed). For a vanilla Core module, the following methods/attributes are highly recommended to be over-ridden/defined:

  • __init__: this is the only place that user-input can be obtained. This method should almost certainly just save input parameters to the instance with the same name, modulo setting default parameters and simple logic.

  • build_model_data(ctx): this is the heart of the Core. It is what produces the actual model quantities, which are saved back to the ctx object. The current MCMC parameters are available via ctx.getParams(), and the model quantities should be saved back to the ctx via ctx.add("key", value). These model quantities should be deterministic, as they convey a model, not a mock. They may however include quantities of interest for determining probability distributions on the model.

The following methods have defaults, but may be required to be over-ridden for particular applications:

  • setup: this method, if defined, is run once only (if embedded in a Chain, the Chain checks if it has been previously run and disallows running again) before the beginning of an MCMC. It receives no parameters. In this sense, it is almost equivalent to tacking the operations onto the end of __init__. The difference is that when setup is run, it is generally assumed that the core has been embedded into a Chain, so that access to other cores (and, if loaded earlier in the core sequence, their respective instance attributes) is available.

  • convert_model_to_mock(ctx): this method takes the model output from build_model_dataand computes a mock simulation from it (i.e. it adds the requisite randomness). If not over-ridden, it is a no-op, which implies that data is considered to be deterministic (as far as this core goes). The method is not invoked in a standard MCMC run, but can be used to generate mock data for consistency tests.

  • __eq__(self, other): this method should determine if this core instance is identical to another core instance. It is used for checking whether the an MCMC chain can be continued from file (i.e. if it is consistent with what has already been run). It is defined in CoreBase by checking each of the input parameters to __init__, as saved in each instance either as its own name, or its name prefaced by “_”. This list of instance attributes can be supplemented using the extra_definining_attributes class attribute, and can be filtered with the ignore_attributes class attribute. It is probably rare that the __eq__ method should be required to be overwritten; usually these class attributes should be all that is required.

We can see the various methods/attributes made available by default:

[6]:

cls = mcmc.CoreBase()
help(cls)

Help on CoreBase in module py21cmmc.mcmc.core object:

class CoreBase(ModuleBase)
 |  CoreBase(store=None)
 |
 |  Method resolution order:
 |      CoreBase
 |      ModuleBase
 |      builtins.object
 |
 |  Methods defined here:
 |
 |  __call__(self, ctx)
 |      Call the class. By default, it will just build model data, with no stochasticity.
 |
 |  __init__(self, store=None)
 |      Initialize self.  See help(type(self)) for accurate signature.
 |
 |  build_model_data(self, ctx)
 |      Passed a standard context object, should construct model data and place it in the context.
 |
 |      The data generated by this method should ideally be *deterministic*, so that input parameters (which are
 |      inherently contained in the `ctx` object) map uniquely to output data. The addition of stochasticity in order
 |      to produce mock data is done in the :meth:`~convert_model_to_mock` method. All data necessary to full evaulate
 |      probabilities of mock data from the model data should be determined in this method (including model
 |      uncertainties, if applicable).
 |
 |      Parameters
 |      ----------
 |      ctx : dict-like
 |          The context, from which parameters are accessed.
 |
 |      Returns
 |      -------
 |      dct : dict
 |          A dictionary of data which was simulated.
 |
 |  convert_model_to_mock(self, ctx)
 |      Given a context object containing data from :meth:`~build_model_data`, generate random mock data, which should
 |      represent an exact forward-model of the process under investigation.
 |
 |      Parameters
 |      ----------
 |      ctx : dict-like
 |          The context, from which parameters and other simulated model data can be accessed.
 |
 |  prepare_storage(self, ctx, storage)
 |      Add variables to special dict which cosmoHammer will automatically store with the chain.
 |
 |  simulate_mock(self, ctx)
 |      Generate all mock data and add it to the context.
 |
 |  ----------------------------------------------------------------------
 |  Methods inherited from ModuleBase:
 |
 |  __eq__(self, other)
 |      Return self==value.
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors inherited from ModuleBase:
 |
 |  __dict__
 |      dictionary for instance variables (if defined)
 |
 |  __weakref__
 |      list of weak references to the object (if defined)
 |
 |  chain
 |      A reference to the LikelihoodComputationChain of which this Core is a part.
 |
 |  parameter_names
 |
 |  ----------------------------------------------------------------------
 |  Data and other attributes inherited from ModuleBase:
 |
 |  __hash__ = None
 |
 |  extra_defining_attributes = []
 |
 |  ignore_attributes = []

A Visibility Core Example#

Back to our proposed example of creating a core which evaluates visibilities. There are two ways of doing this. The first is to subclass an original Core class from 21CMMC (in this case, the CoreLightConeModule will be required).

A minimal example would then be:

[ ]:

class MyVisibilityCore(core.CoreLightConeModule):
    def __init__(self, frequencies, baselines,
                 *args, **kwargs # always include *args, **kwargs and call super()
                ):
        # Call super to initialise standard lightcone module
        super().__init__(*args, **kwargs)

        # Add other user-dependent quantities
        self.frequencies = frequencies
        self.baselines = baselines

    def build_model_data(self, ctx):
        # Call the LightConeModule model builder, and add its model to the context.
        super().build_model_data(ctx)

        # Convert the lightcone into visibilities
        vis = convert_lightcone_to_vis(ctx.get("lightcone"))

        # Importantly, add useful quantities to the context, so they are available
        # to the likelihood
        ctx.add("visibilities", vis)

        # Also could clean up any quantities that we don't care about anymore.
        # But it might just be better to leave it there.
        ctx.remove("lightcone")

Another method hinges on the fact that multiple cores can be loaded for a given MCMC chain. In this method, the defined Core is entirely separate from the LightConeModule, but both must be explicitly loaded:

[ ]:

class MyVisibilityCore(mcmc.CoreBase):
    required_cores = [mcmc.CoreLightConeModule]

    def __init__(self, frequencies, baselines, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.frequencies = frequencies
        self.baselines = baselines

    def build_model_data(self, ctx):
        lightcone = ctx.get("lightcone") # Assumes that the LightConeModule is loaded before this one.

        # Convert the lightcone into visibilities
        vis = convert_lightcone_to_vis(lightcone)

        # Importantly, add useful quantities to the context, so they are available
        # to the likelihood
        ctx.add("visibilities", vis)

Which of these is preferred is difficult to determine. The first option is simpler, in that it uses well-worn methods for overwriting components of the original. It can also be generally assumed that if the visibilities are to be used, the original simulation will not be required. It also means that when it comes to actually running the MCMC, only the MyVisibilityCore will need to be loaded. On the other hand, this method requires passing parameters to MyVisibilityCore that actually are just passed through to CoreLightConeModule, which is perhaps a bit unclean.

TLDR; either will work.

We could supplement either of the above with an additional convert_model_to_mock method, which might add additional complex Gaussian noise to each visibility.

How do I write a Likelihood Module?#

If you have read the “Write My Own Core” section just above, then you know how important extensibility is to 21CMMC. This is especially true when it comes to the Likelihood modules (a brief intro to the idea behind cores and likelihoods is given in the MCMC intro).

Likelihoods should inherit from py21cmmc.mcmc.BaseLikelihood, or if you want a simple way to support reading in data and noise from files, py21cmmc.mcmc.BaseLikelihoodFile.

Along with the standard __init__ and setup methods which have very similar functions to their Core counterparts (see previous section), it is recommended to overwrite the following methods:

  • reduce_data(ctx): this takes the data produced from all cores, and reduces it to “final form”. This “final form” is essentially defined as the maximally reduced form that can be applied to either data or model separately, before calculating the likelihood from their combination. For instance, it may be to perform an FFT to obtain the power spectrum. It returns the reduced data as a dictionary of quantities.

  • computeLikelihood(model): this function takes the model dictionary output by reduce_data and computes a likelihood from it, using the self.data attribute of the instance. The reason for separating these two methods is simple: it allows applying the reduce_data method on both data and model, as well as simulating mock datasets.

  • store(model, storage): a method which is called on every MCMC iteration, and passed the model output from reduce_data. It should save this data to a storage dictionary, and it will subsequently be saved to the storage chain file.

If using the BaseLikelihoodFile class, an extra set of methods are available to be over-ridden which define how data and measurement noise should be read in from file. See the docstring or the FAQ on adding noise for details on these methods.

The simplest example of a likelihood would be something like the following:

[ ]:

class MySimpleLikelihood(mcmc.LikelihoodBaseFile):
    def reduce_data(self, ctx):
        k, power_spectrum = convert_to_power(ctx.get("lightcone"))
        return dict(power_spectrum=power_spectrum, k=k)

    def computeLikelihood(self, model):
        return np.sum((self.data["power_spectrum"] - model['power_spectrum'])**2 / (2 * self.noise(['variance'])))

Since we are using the LikelihoodBaseFile, the data is automatically read in from a user-input file. The default read method assumes that the file is a numpy format (.npz), and contains exactly the same quantities that are returned by reduce_data (which allows very easily for simulating mock data and outputting to file, and subsequently reading it back in).

The self.data attribute is defined during setup, and filled with the contents of the input data file, and likewise the self.noise attribute (it comes from a separate noisefile).