Source code for datapipe_testbench.benchmarks.dl2

#!/usr/bin/env python3

"""Benchmarks using DL2 information."""

from hist import axis

from ..resolution import ResolutionBenchmark
from .classification import ClassifierBenchmark
from .energy_regression import EnergyRecoBenchmark

__all__ = [
    "ClassifierBenchmark",
    "EnergyRecoBenchmark",
    "EnergyResolutionBenchmark",
    "TelescopeEnergyResolutionBenchmark",
]




class EnergyResolutionBenchmark(ResolutionBenchmark):
    """
    Benchmark reconsructed energy resolution and bias per telescope.
    """

    def __init__(
        self,
        reco_energy_column: str = "RandomForestRegressor_energy",
        chunk_size: int = 100_000,
        max_chunks=None,
    ):
        """
        Construct a energy resolution metric.

        Parameters
        ----------
        reco_energy_column: str
            name of reconstructed energy column in the input.
        chunk_size: int
            how many events to load per chunk (smaller means lower memory
            footprint, but slower)
        max_chunks: int | None
            If specified, stop after this many chunks
        """
        super().__init__(
            input_data_level="dl2",
            reco_column=reco_energy_column,
            true_axis=axis.Regular(
                bins=40,
                start=0.003,
                stop=5000.0,
                name="true_energy",
                transform=axis.transform.log,
                metadata=dict(unit="TeV"),
            ),
            per_tel_type=False,
            chunk_size=chunk_size,
            max_chunks=max_chunks,
            bias_plot_range=(-5, 5),
        )





class TelescopeEnergyResolutionBenchmark(ResolutionBenchmark):
    """
    Benchmark final reconsructed energy resolution and bias.
    """

    def __init__(
        self,
        reco_energy_column: str = "RandomForestRegressor_tel_energy",
        chunk_size: int = 100_000,
        max_chunks=None,
    ):
        """
        Construct a energy resolution metric.

        Parameters
        ----------
        reco_energy_column: str
            name of reconstructed energy column in the input.
        chunk_size: int
            how many events to load per chunk (smaller means lower memory
            footprint, but slower)
        max_chunks: int | None
            If specified, stop after this many chunks
        """
        super().__init__(
            input_data_level="dl2",
            reco_column=reco_energy_column,
            true_axis=axis.Regular(
                bins=40,
                start=0.003,
                stop=5000.0,
                name="true_energy",
                transform=axis.transform.log,
                metadata=dict(unit="TeV"),
            ),
            per_tel_type=True,
            chunk_size=chunk_size,
            max_chunks=max_chunks,
            bias_plot_range=(-5, 5),
        )