3.4.13. A simple python module#

This lesson will give you an idea about the structure and use of basf2 modules. Most of the modules in our software are implemented in C++ and are made available for analysis by modularAnalysis. This package consists of python wrapper functions around the C++ modules in order to use them in the python steering file. You have already learned about this in The basics.

C++ is very strong and fast, but usually much more complicated to read and write than Python. For this reason the basf2 framework provides the possibility to write modules also in Python. This can be very helpful if you want to investigate or test something.

To put your hands on this, simply copy the code into a python file and run it with basf2 my_python_module.py. It is nothing more than a steering file with your own class.

Minimal example#

Let’s begin with the following minimal example for a new python module. It is the “Hello World” of basf2 modules. The magic happens in the class MinModule(basf2.Module). In this basic example, the class only consists of one member function event that is called once for each event. We use the logging function basf2.B2INFO() to print our message. To execute the model, we need to create a path and generate dummy events. Then, our module is added to the path and we run the path.

#!/usr/bin/env python3

import basf2 as b2


class MinModule(b2.Module):
    """Very minimal class to print Hello World in each event"""

    def event(self):
        """Event function, called once for each event"""
        b2.B2INFO("Hello World!")


# create a path
main = b2.Path()

# set to generate 10 dummy events
main.add_module("EventInfoSetter", evtNumList=[10])

# and add our module
main.add_module(MinModule())

# run the path
b2.process(main)

You can see that implementing a minimal python module just takes 5 lines of code (3 without documentation) so it’s very nice for fast and quick prototyping.

Note

Python modules have to be implemented or imported in the steering file
Python modules are usually much slower then C++ modules but for many small tasks this does not make a significant difference.
These hacky modules will not appear in module list (basf2 -m)
Python modules can only be used in analysis code or private scripts. Only C++ modules can be added to the official reconstruction code that is run for HLT or for calibration.

Detailed usage#

Let’s extend the minimal example class above to show all methods which can be implemented by a Python module. As you have seen above all the member functions are optional.

#!/usr/bin/env python3

import basf2 as b2


class MinModule(b2.Module):
    """A minimal example of a basf2 module in python."""

    def __init__(self):
        """Constructor"""
        # call constructor of base class, required if you implement __init__
        # yourself!
        super().__init__()
        # and do whatever else is necessary like declaring member variables

    def initialize(self):
        """Called once in the beginning just before starting processing"""
        b2.B2INFO("initialize()")

    def beginRun(self):
        """Called every time a run changes before the actual events in that run
        are processed
        """
        b2.B2INFO("beginRun()")

    def event(self):
        """Called once for each event"""
        b2.B2INFO("event()")

    def endRun(self):
        """Called every time a run changes after the actual events in that run
        were processed
        """
        b2.B2INFO("endRun()")

    def terminate(self):
        """Called once after all the processing is complete"""
        b2.B2INFO("terminate()")


# create a path
main = b2.Path()

# set to generate 10 dummy events
main.add_module("EventInfoSetter", evtNumList=[10])

# and add our module
main.add_module(MinModule())

# run the path
b2.process(main)

Accessing Datastore Objects#

Datastore objects can be accessed via the PyStoreArray class and the PyStoreObj classes. Let’s create a small module which will print the event number and information on MCParticles, namely the PDG code. To have tracks available, we will use the ParticleGun module, which generates very simple events.

Exercise

Write a Python module that prints the number of charged particles per event and the total charge. Note: per default, the ParticleGun generates only one track per event, but you can adjust this.

Hint

You can find information on the Particle class in doxygen. The ParticleGun has the option nTracks.

Solution

#!/usr/bin/env python3

import basf2 as b2
from ROOT import Belle2


class AccessingDataStoreModule(b2.Module):
    """An example of a basf2 module in python which accesses things in the datastore."""

    def initialize(self):
        """Create a member to access event info and the MCParticles
        StoreArray
        """
        #: an example object from the datastore (the metadata collection for the event)
        self.eventinfo = Belle2.PyStoreObj("EventMetaData")
        #: an example array from the datastore (the list of MC particles)
        self.particles = Belle2.PyStoreArray("MCParticles")

    def event(self):
        """Print the number of charged particles and the total charge"""
        n_charged = 0
        total_charge = 0
        for particle in self.particles:
            charge = particle.getCharge()
            if charge:
                n_charged += 1
            total_charge += charge

        b2.B2INFO(
            f"Number of charged particles = {n_charged}, "
            f"total charge of event = {total_charge}"
        )


# create a path
main = b2.Path()

# generate events
main.add_module("EventInfoSetter", evtNumList=[10])

# generate events with 3 tracks (not all of them are charged tracks)
main.add_module("ParticleGun", nTracks=3)

# and add our module
main.add_module(AccessingDataStoreModule())

# run the path
b2.process(main)

Note

For PyStoreObj you can access most members of the underlying class directly, like eventinfo.getEvent() above. However if you want to get the object directly or want to access a member which also exists by the same name in PyStoreObj you can use the obj() member to get a reference to the underlying object itself: eventinfo.obj().getEvent()

More advanced examples#

framework/examples/cdcplotmodule.py - A full example that uses matplotlib to plot CDCSimHits
framework/examples/interactive_python.py drops into an interactive (i)python shell inside the event() function, allowing exploration of available objects and data structures
reconstruction/examples/plot_LL_diff.py - Gets PID log-likelihoods, uses relations to get corresponding MC truth and fills ROOT histograms accordingly

Stuck? We can help!

If you get stuck or have any questions to the online book material, the #starterkit-workshop channel in our chat is full of nice people who will provide fast help.

Refer to Collaborative Tools. for other places to get help if you have specific or detailed questions about your own analysis.

Improving things!

If you know how to do it, we recommend you to report bugs and other requests with GitLab. Make sure to use the documentation-training label of the basf2 project.

If you just want to give very quick feedback, use the last box “Quick feedback”.

Please make sure to be as precise as possible to make it easier for us to fix things! So for example:

typos (where?)
missing bits of information (what?)
bugs (what did you do? what goes wrong?)
too hard exercises (which one?)
etc.

If you are familiar with git and want to create your first merge request for the software, take a look at How to contribute. We’d be happy to have you on the team!

Quick feedback!

Author of this lesson

Pascal Schmolz

basf2 development documentation

A simple python module

Contents