Custom Simulation and Advanced Tricks

This notebook will demonstrate how one can use the modular structure of fuse to change parts of the simulation by exchanging or adding new plugins. To get the most out of this notebook, you should start with an empty out_dir when setting up the simulation context. Additionaly the notebook will show how a map can be replaced with a dummy map to disable certain parts of the simulation.

Imports & Simulation Context

import fuse
import matplotlib.pyplot as plt

st = fuse.context.xenonnt_fuse_full_chain_simulation(
    output_folder="./fuse_data",
    simulation_config="sr0_dev",
)

st.set_config(
    {
        "path": "/project2/lgrandi/xenonnt/simulations/testing",
        "file_name": "pmt_neutrons_100.root",
        "entry_stop": 10,
        "debug": True,
    }
)

run_number = "00000"

Running the default simulation

First we will run the default microphysics simulation. As we set the debug config to True fuse will give use some more information during the simulation. For example, each plugin will print its version number and the used seed for random number generation.

st.make(run_number, "microphysics_summary")

Using a different plugin

First lets see how we can exchange a fuse plugin with a plugin that is included in fuse but not used by default. One example is the BBFYields plugin that can be used instead of NestYields. To add a plugin to our simulation context we can use the st.register function provided by strax.

st.register(fuse.plugins.micro_physics.yields.BBFYields)

Now that the new plugin is registered we can run the simulation again. Can you spot the difference in the debug output?

st.make(run_number, "microphysics_summary")

fuse should now tell you:

DEBUG:fuse.micro_physics.yields:Running BBFYields version 0.1.0 in debug mode

This way you can easily see that the new plugin is used. You do not need to worry that you could mix up data produced with different plugins as strax recognizes the changed simulation context and will rerun the needed simulation steps.

Building a new plugin

Now that we know how to exchange a plugin in the context we can build a new plugin and use it in the simulation. For our example we will replace the ElectronDrift plugin with a new plugin that does the simulation a “creative” way. You can find the plugin with some inline comments below.

import strax
import straxen
import numpy as np
import logging

logging.basicConfig(handlers=[logging.StreamHandler()])
log = logging.getLogger("fuse.detector_physics.electron_drift")

from fuse.common import FUSE_PLUGIN_TIMEOUT


# A fuse plugin is a python class that inherits from strax.Plugin
# As naming convention we use CamelCase for the class name
class CosineElectronDrift(strax.Plugin):
    # Each plugin has a version number
    # If the version number changes, fuse will know that it need to re-simulate the data
    __version__ = "0.1.0"

    # You need to tell fuse and strax what the plugin needs as input
    # In this case we need the microphysics_summary
    # If you need more than one input, you can use a tuple
    depends_on = "microphysics_summary"

    # You need to tell fuse and strax what the plugin provides as output
    # In this case we provide drifted_electrons
    # You can later use st.make(run_number, "drifted_electrons") to run the simulation
    provides = "drifted_electrons"

    # You need to tell fuse and strax what the data looks like
    # Data of the same data_kind can be combined via "horizontal" concatenation and need
    # to have the same output length.
    data_kind = "interactions_in_roi"

    # You also need to tell strax what columns the data has
    # A column needs a name and a numpy data type.
    # For this example we will not deal with the other colums usually present in drifted_electrons
    # This can lead to problems later on in the simulation but is fine for this example
    dtype = [
        ("n_electron_interface", np.int64),
    ]
    dtype = dtype + strax.time_fields

    # We need to disable automatic rechunking for fuse plugins
    # As fuse is going from "leightweigt" data to "heavy" data,
    # automatic rechunking can lead to problems in later plugins
    rechunk_on_save = False

    # We need to specify when we want to save the data
    save_when = strax.SaveWhen.TARGET

    # strax uses a rather short timeout, lets increase it as
    # some of the fuse simulation steps can take a while
    input_timeout = FUSE_PLUGIN_TIMEOUT

    # We need to tell strax what config options the plugin needs
    # We will use the great URLConfigs that are a part of straxen
    debug = straxen.URLConfig(
        default=False,
        type=bool,
        track=False,
        help="Show debug informations",
    )

    deterministic_seed = straxen.URLConfig(
        default=True,
        type=bool,
        help="Set the random seed from lineage and run_id, or pull the seed from the OS.",
    )

    # For our example we will need the tpc_length
    tpc_length = straxen.URLConfig(
        type=(int, float),
        help="tpc_length",
    )

    # And a new variable we will call n_periods
    n_periods = straxen.URLConfig(
        type=(int, float),
        help="n_periods",
    )

    # If you want to prepare something before we start to run the compute method
    # you can put it into the setup method. The setup method is called once while the
    # compute method is called independently for each chunk
    def setup(self):
        # Lets convert the tpc_length and n_periods into a scaling factor
        self.scaling_factor = self.tpc_length / self.n_periods

        # All plugins can report problmes or debug information via the logging feature
        # You can set the log level via the debug config option.
        # WARNING messages are always shown whild DEBUG messages are only shown if debug is True
        if self.debug:
            log.setLevel("DEBUG")
            log.debug(f"Running ElectronDrift version {self.__version__} in debug mode")
        else:
            log.setLevel("WARNING")

        # Many plugins need to generate random numbers for simulation the corresponding physics process
        # In fuse we want to make sure that the simulation is reproducible.
        # Therefore we have the default setting of deterministic_seed = True
        # In this case the random seed is generated from the run_id and the lineage
        # The lineage includes all plugins and their verions that are connected to the input of the
        # current plugin as well as all tracked config options and the strax version.
        # The run_id is a user input. More on the deterministic seed can be found in
        # a dedicated notebook.
        # Please make sure that you use the random number generator self.rng when you need random numbers
        # later in the plugin.
        if self.deterministic_seed:
            hash_string = strax.deterministic_hash((self.run_id, self.lineage))
            seed = int(hash_string.encode().hex(), 16)
            self.rng = np.random.default_rng(seed=seed)
            log.debug(f"Generating random numbers from seed {seed}")
        else:
            self.rng = np.random.default_rng()
            log.debug(f"Generating random numbers with seed pulled from OS")

    # The compute method is the heart of the plugin. It is executed for each chunk of input data and
    # must produce data in the format specified in the self.dtype variable.
    def compute(self, interactions_in_roi):
        # For your new plugin you would put your new simulation code here. In this example
        # We will do some stupid calculation of the drift time.

        # Make sure you only apply the plugin to interactions that have some electrons
        # Adapt this to match your plugin-input
        mask = interactions_in_roi["electrons"] > 0

        # Make sure your plugin can handle empty inputs
        if len(interactions_in_roi[mask]) == 0:
            return np.zeros(0, self.dtype)

        # Build the output array with the correct length and dtype
        result = np.zeros(len(interactions_in_roi), dtype=self.dtype)
        # We do not want to change the timing of the interactions, so we just take them from the input
        result["time"] = interactions_in_roi["time"]
        result["endtime"] = interactions_in_roi["endtime"]

        # Lets do some stupid calculation of the drifted electrons using a cosine function
        n_electron_interface = (
            interactions_in_roi[mask]["electrons"]
            * np.cos(interactions_in_roi[mask]["z"] / self.scaling_factor * np.pi) ** 2
        )
        # Lets add some noise and make sure to use the random number generator defined in the setup method
        n_electron_interface = n_electron_interface * self.rng.normal(
            1, 0.1, len(n_electron_interface)
        )
        result["n_electron_interface"][mask] = n_electron_interface

        return result

Now that our new plugin is defined we can register it, adjust the config and then try to run it.

st = fuse.context.xenonnt_fuse_full_chain_simulation(
    output_folder="./fuse_data",
    simulation_config="sr0_dev",
)

# We will use a different file now with a little more statistics
st.set_config(
    {
        "path": (
            "/project2/lgrandi/xenonnt/simulations/lead-214/high_energy_er_fullchain_Pb214_simulations/TPC_Pb214_lower/G4outsTPCXenonProgeny1_2000_nochain_GS_42/"
        ),
        "file_name": "nT_TPC_Pb214_2000_nochain_50.root",
        "n_periods": 5,
    }
)

run_number = "00000"

st.register(CosineElectronDrift)

st.make(run_number, "drifted_electrons")

data = st.get_df(run_number, ["microphysics_summary", "drifted_electrons"])

Now that our new plugin produced some data we can try to visualize the simulation output. We will calculate the mean n_electron_interface for different z-slices and plot the result.

def average_value_in_interval(edge_start, edge_end, df, value):
    value_mean_in_interval = []
    for begin, end in zip(edge_start, edge_end):
        data_in_slice = df[(df.z >= begin) & (df.z < end)][value]

        if len(data_in_slice) > 0:
            value_mean_in_interval.append(np.nanmean(data_in_slice))

        else:
            value_mean_in_interval.append(0)

    value_mean_in_interval = np.array(value_mean_in_interval)

    return value_mean_in_interval


bin_edges = np.linspace(-150, -1, 150)
bin_edges_start = bin_edges[:-1]
bin_edges_end = bin_edges[1:]
bin_centers = (bin_edges_start + bin_edges_end) / 2

electrons_reaching_the_interface_in_z_slice = average_value_in_interval(
    bin_edges_start, bin_edges_end, data, "n_electron_interface"
)
electrons_at_interaction_site_in_z_slice = average_value_in_interval(
    bin_edges_start, bin_edges_end, data, "electrons"
)


plt.plot(
    bin_centers,
    electrons_reaching_the_interface_in_z_slice,
    color="purple",
    label="Electrons reaching the interface",
    lw=0,
    marker="o",
    markersize=3,
)

plt.plot(
    bin_centers,
    electrons_at_interaction_site_in_z_slice,
    color="orange",
    label="Electrons at interaction site",
    lw=0,
    marker="o",
    markersize=3,
)

plt.legend()
plt.xlabel("z [cm]")
plt.ylabel("Number of electrons")
plt.show()

Tipps & Tricks for plugin development

In this section I will collect some tipps and tricks that might be useful when developing new plugins.

1. Running the compute method in the notebook

strax has the functionality to access a plugin directly. This way you can test your plugins methods in the notebook without relying on st.make to run the plugin. First we can get the plugin completely initialized by calling st.get_single_plugin. Then we can call the compute method with the needed input arguments.

plugin = st.get_single_plugin(run_number, "drifted_electrons")

Make sure that your input data has the correct numpy format. Dataframes are not supported.

microphysics_summary = st.get_array(run_number, ["microphysics_summary"])

plugin_output = plugin.compute(microphysics_summary)
print(plugin_output[0:10])

Replacing a map with a dummy map

Sometimes you might want to disable a certain feature of the simulation without going into the code and changing a plugin. One example can be that you might want to turn of the effects of a certain map. In this example we will use the constant_dummy_map URLConfig protocol to disable the effect of the s2_correction map.

We can set up a new context. This time we will override the config option for the s2_correction map with the dummy map option. Instead of the real map values, the dummy map will return a constant value that can be configured by the user. In our case we set it to 1. You could also set it to other values if you want to test the effect of the map on the simulation.

st = fuse.context.xenonnt_fuse_full_chain_simulation(
    output_folder="./fuse_data",
    simulation_config="sr0_dev",
)

st.set_config(
    {
        "path": "/project2/lgrandi/xenonnt/simulations/testing",
        "file_name": "pmt_neutrons_100.root",
        "entry_stop": 10,
        "s2_correction_map": "constant_dummy_map://1",
    }
)

run_number = "00000"

To get a feeling how the map affects your simulation, lets take a look at the number of s2 photons.

st.make(run_number, "s2_photons")
s2_photons = st.get_df(run_number, "s2_photons")

Now that we have the data, lets change the map to a different value and rerun the simulation.

st.set_config(
    {
        "s2_correction_map": "constant_dummy_map://2",
    }
)

st.make(run_number, "s2_photons")
s2_photons_changed_constant = st.get_df(run_number, "s2_photons")

And now we can plot the data again. You should see that the distribution of s2 photons is different.

import numpy as np
import matplotlib.pyplot as plt

bins = np.linspace(0, 100, 100)

plt.hist(s2_photons["n_s2_photons"], bins=bins, alpha=0.5, label="Original")
plt.hist(
    s2_photons_changed_constant["n_s2_photons"], bins=bins, alpha=0.5, label="Changed constant"
)

plt.legend()
plt.xlabel("Number of S2 photons")
plt.ylabel("Counts")
plt.show()