18.1. Background Generator#

The package pxd.background_generator implements a proof-of-concept generator of PXD data for background overlay on the fly. The generator is a stand-alone PXDBackgroundGenerator module integrated into the current background overlay setup. Each instance of background overlay in detector simulation requires input data to combine with simulated signal data. Background overlay input data — a background sample of detector outputs — is high in volume. Computing sites must obtain background overlay input samples before data production. Assuming the correlation between the output of the PXD and other detectors in background samples is negligible, the volume of data distributed is reduced by using generated PXD output.

The module is designed to be incorporated into the detector simulation and used to provide partial background overlay input data in conjunction with data from background overlay files.

18.1.1. Usage#

Warning

The module is a proof of concept not intended for production.

Background overlay is enabled by providing one or more files with input samples. In the current setup, background overlay uses exclusively data from the background overlay files:

import basf2
from background import get_background_files
from simulation import add_simulation

path = basf2.create_path()
path.add_module("EventInfoSetter", evtNumList=[128])
path.add_module("EvtGenInput")
add_simulation(path, bkgfiles=get_background_files())
basf2.process(path)

The module is integrated to discard the PXD output from the sampled data and provide background overlay with new input generated on the fly. To use the module, create a new module instance. Note that a Python environment with PyTorch is required — it is recommended to use the environment included with the externals. Use the drop-in replacement function incorporating the module instance instead of the standard function to initialize the simulation:

import basf2
from background import get_background_files
from pxd.background_generator import PXDBackgroundGenerator, inject_simulation

module = PXDBackgroundGenerator(model="resnet")
add_simulation = inject_simulation(module)

path = basf2.create_path()
path.add_module("EventInfoSetter", evtNumList=[128])
path.add_module("EvtGenInput")
add_simulation(path, bkgfiles=get_background_files())
basf2.process(path)

See also

Example steering file provided in $BELLE2_RELEASE_DIR/pxd/examples/background_generator.py.

The module implements different generator models. A module instance uses one of the implemented generator models.

18.1.2. Models#

The generator is based on generative machine learning models pre-trained on samples of PXD output from background overlay files. The implemented generator models convnet and resnet are neural networks of different architectures inspired by WGAN-GP. To initialize a generator model, a checkpoint file with pre-trained model parameters must be provided. If no checkpoint is specified, the module transfers and uses the default checkpoint from the conditions database.

The PXD consists of forty modules, referred to as sensors. Each sensor embeds a \(250 \times 768\) pixel matrix. Sensors are approximated to have pixels with matching dimensions and equally spaced. This approximation is necessary to allow the same representation of all sensor readout by grayscale images — heat maps. Accordingly, a PXD readout and output corresponds to forty grayscale images. Training instances of the default checkpoints for both models are based on the same training dataset of \(3.6 \times 10^6\) images. The images were extracted from nine background overlay files with release version 01-00-00 with simulated background samples. In the future, it is hoped that more checkpoints with model parameters trained on different background datasets will be provided to generate run-dependent samples.

Generally, a PXD output is produced from forty images specific for each sensor. The generator models currently implemented generate one image at a time and do not allow specifying the sensor for which the image is to be generated. For this reason, the module is limited to produce PXD output from forty arbitrary images, some belonging to the same sensor. The limitation is expected to be resolved with the implementation of new class-conditional models that are currently being developed.

18.1.3. Development#

Model Implementation

To implement a new generator model, choose a unique model name and add the name into the collection pxd.background_generator.models.MODELS. Create a Python module with the same name in the package pxd.background_generator.models. The module must contain the definition of a model class and a generation function.

The model class defines the neural network model structure — it must be named Model and inherit from the torch.nn.Module base class. The generation function defines how to use a model instance to generate a PXD output sample — it must be named generate, act on a model instance, and return a torch.Tensor with shape (40, 250, 768) and torch.uint8 values.

18.1.4. Reference#

class pxd.background_generator.PXDBackgroundGenerator(model: str = 'convnet', checkpoint: Union[None, str, pathlib.Path] = None, seed: Union[None, int] = None, nintra: int = 1, ninter: int = 1, globaltag: str = 'PXDBackgroundGenerator')[source]#

Generates PXD background samples for background overlay on the fly.

Parameters

  • model (str) – Name of the generator model to use - either “convnet” or “resnet”, defaults to “convnet” (optional)

  • checkpoint (str) – Path to the checkpoint file with weights for the selected model, defaults to None - use the default checkpoint from the conditions database (optional)

  • seed (int) – Integer number in the interval \([-2^{63}, 2^{63} - 1]\) used internally as the initial seed, defaults to None - derive a deterministic seed from the value returned by basf2.get_random_seed() (optional)

  • nintra (int) – Number of intra-op threads to be utilized for the generation, defaults to 1 (optional)

  • ninter (int) – Number of inter-op threads to be utilized for the generation, defaults to 1 (optional)

  • globaltag (str) – Global tag of the conditions database providing the default checkpoints stored as payloads, defaults to “PXDBackgroundGenerator” (optional)

pxd.background_generator.inject_simulation(module: Union[None, pybasf2.Module] = None) → Callable[source]#

Incorporate a module instance into add_simulation() after BGOverlayInput.

Parameters

module (basf2.Module, optional) – Module instance to be incorporated, defaults to None - return unmodified function

Returns

Drop-in replacement function for add_simulation()