Writing a module in python

Example

Python is the easiest language to write a new module in. Here's an example, which assumes you already have a python code called my_calculation.py that does your calculation.

Let's assume is need one fixed parameter telling it what to do, called "mode", and it uses a single cosmological parameter input, "omega_m" and has one cosmological parameter output, "cluster_mass".

#!python
from cosmosis.datablock import names, option_section
import my_calculation

# We have a collection of commonly used pre-defined block section names.
# If none of the names here is relevant for your calculation you can use any
# string you want instead.
cosmo = names.cosmological_parameters

def setup(options):
    #This function is called once per processor per chain.
    #It is a chance to read any fixed options from the configuration file,
    #load any data, or do any calculations that are fixed once.

    #Use this syntax to get a single parameter from the ini file section
    #for this module.  There is no type checking here - you get whatever the user
    #put in.

    mode = options[option_section, "mode"]

    #The call above will crash if "mode" is not found in the ini file.
    #Sometimes you want a default if nothing is found:
    high_accuracy = options.get(option_section, "high_accuracy", default=False)

    #Now you have the input options you can do any useful preparation
    #you want.  Maybe load some data, or do a one-off calculation.
    loaded_data = my_calculation.prepare_something(mode)

    #Whatever you return here will be saved by the system and the function below
    #will get it back.  You could return 0 if you won't need anything.
    return loaded_data


def execute(block, config):
    #This function is called every time you have a new sample of cosmological and other parameters.
    #It is the main workhorse of the code. The block contains the parameters and results of any 
    #earlier modules, and the config is what we loaded earlier.

    # Just a simple rename for clarity.
    loaded_data = config

    #This loads a value from the section "cosmological_parameters" that we read above.
    omega_m = block[cosmo, "omega_m"]

    # Do the main calculation that is the purpose of this module.
    # It is good to make this execute function as simple as possible
    cluster_mass = my_calculation.compute_something(omega_m, loaded_data)

    # Now we have got a result we save it back to the block like this.
    block[cosmo, "cluster_mass"] = cluster_mass

    #We tell CosmoSIS that everything went fine by returning zero
    return 0

def cleanup(config):
    # Usually python modules do not need to do anything here.
    # We just leave it in out of pedantic completeness.
    pass

Reference for cosmosis DataBlock python methods

The "block" object you get passed in python has these methods that you can use to load or save data:

Simple access functions

The simplest way to get, put, or replace data is with the usual python access method:

#!python

block["section_name", "value_name"] = x
x = block["section_name", "value_name"]

The first will set the "value_name" in the section "section_name" of the block to the value of x. If something with that name already exists then it will be replaced.

The second will load the same things back from the block. An exception (cosmosis.BlockError) is raised if the named thing is not found.

GET functions

#!python

block.get("section_name", "value_name", default=x)
block.get_int("section_name", "value_name", default=x)
block.get_double("section_name", "value_name", default=x)
block.get_complex("section_name", "value_name", default=x)
block.get_bool("section_name", "value_name", default=x)
block.get_string("section_name", "value_name", default=x)
block.get_int_array_1d("section_name", "value_name", default=x)
block.get_double_array_1d("section_name", "value_name", default=x)

These functions give you a bit more control over what you load from the block. They also allow you to specify a default value to be used if the named thing is not found.

The specifically typed functions (all except the first) allow you to define what type of value you want returned - an error will be given if things are the wrong type; this can be useful for debugging.

The "array" functions return numpy (numerical python) arrays.
The "bool" function returns python True or False. The "string" function returns text (any length).

PUT and REPLACE functions

#!python

block.put("section_name", "value_name", value)
block.put_int("section_name", "value_name", value)
block.put_double("section_name", "value_name", value)
block.put_complex("section_name", "value_name", value)
block.put_bool("section_name", "value_name", value)
block.put_string("section_name", "value_name", value)
block.put_int_array_1d("section_name", "value_name", value)
block.put_double_array_1d("section_name", "value_name", value)

block.replace("section_name", "value_name", value)
block.replace_int("section_name", "value_name", value)
block.replace_double("section_name", "value_name", value)
block.replace_complex("section_name", "value_name", value)
block.replace_bool("section_name", "value_name", value)
block.replace_string("section_name", "value_name", value)
block.replace_int_array_1d("section_name", "value_name", value)
block.replace_double_array_1d("section_name", "value_name", value)

The put and replace functions work very similarly to each other. They save things to the block, for later usage.

The "put" functions raise an error (cosmosis.BlockError) if the thing they try to put already exists (to stop you accidentally overwriting something) and the "replace" functions raise an error if it does not already exist.

The typed functions will raise an error if you pass something in with the wrong type.

Grids

A common pattern is to want to tabulate a function of two variables, such as P(k, z). In that case we need to save or load three values, a 1D array k, another 1D array z, and a 2D array P. CosmoSIS has some helper functions to do this. They also ensure that if another module loads the grid in the opposite order (they want P(z, k) instead of P(k, z) then this is handled correctly. You can use these functions:

#!python

block.put_grid("section_name", "name_x", x, "name_y", y, "name_z", z)
block.replace_grid("section_name", "name_x", x, "name_y", y, "name_z", z)
x, y, z = block.get_grid("section_name", "name_x", "name_y", "name_z", z)

where z is 2D and x, y are 1D. These functions will check that the sizes of your arrays are as expected.

Debugging functions

These functions can help with debugging code as you are writing it. You wouldn't normally use them when running long chains but you might when fixing errors of validating things.

#!python
block.save_to_directory("./dirname")

This function saves the entire contents of the block to the named directory. Subdirectories will be the different sections, and inside those vectors will be saved as their own files, and a values.txt file will contain the scalar quantities.

The "test" sampler uses this at the end of its run to save the outputs.

#!python
block.report_failures()
block.print_log()

These functions give a log of what things have been read and written from this block. The report_failures function just tells you things that have gone wrong with this block, and the print_log function tells you every access and whether it worked or not.

If you set "debug=T" in the test sampler you will get a log like this at the end of the run.

Exploring the block

These functions let you check whether or not things are saved in the block already. In a typical situation you would not need them; the would be useful only if you were doing something special like deciding what to calculate based on what was available.

#!python

block.has_section("section_name")
block.has_value("section_name", "value_name")
block.sections()
block.keys()

The has_value and has_section functions return True or False according to whether the named thing(s) exist in the block. The sections() function returns a list of all the sections in the block, and the keys() function returns a list of (section_name, value_name) pairs.

Wiki

cosmosis / modules_python