Adding Summary Metrics

CuBIE’s summary metrics system is extensible via the @register_metric decorator. This page walks through implementing a custom metric.

File Location

Built-in metrics live in src/cubie/outputhandling/summarymetrics/. Each metric is a separate module.

The @register_metric Decorator

Register a new metric by subclassing SummaryMetric and decorating with @register_metric:

from cubie.outputhandling.summarymetrics.metrics import (
    SummaryMetric,
    MetricFuncCache,
    register_metric,
    summary_metrics,
)

@register_metric(summary_metrics)
class MyCustomMetric(SummaryMetric):
    def __init__(self, precision):
        super().__init__(
            buffer_size=2,
            output_size=1,
            name="my_custom",
            precision=precision,
            unit_modification="[custom]",
            sample_summaries_every=None,
        )

    def build(self) -> MetricFuncCache:
        # Return a MetricFuncCache with update and save callables
        ...

__init__ Parameters

buffer_size

Number of scratch elements per variable needed during accumulation. Can be an int or a callable that receives the number of summarised variables.

output_size

Number of output elements per variable written to the result. Can also be a callable.

name

String identifier used in output_types lists.

unit_modification

Label suffix appended to variable names in the summary legend (e.g. "[mean]", "[max]").

sample_summaries_every

If not None, the metric requests sub-step sampling at this interval.

The build() Method

build() must return a MetricFuncCache containing two Numba CUDA device functions:

update(buffer, value, t, dt)

Called at every summary sample point during integration. Accumulates data into buffer.

save(buffer, output, n_samples)

Called at the end of the solve. Finalises the accumulated data and writes to output.

Both functions operate on per-variable slices and must be compatible with Numba’s CUDA compilation.

Example: Implementing a “Count” Metric

A metric that counts how many samples were accumulated:

from numba import cuda

@register_metric(summary_metrics)
class CountMetric(SummaryMetric):
    def __init__(self, precision):
        super().__init__(
            buffer_size=1,
            output_size=1,
            name="count",
            precision=precision,
            unit_modification="[count]",
        )

    def build(self) -> MetricFuncCache:
        @cuda.jit(device=True)
        def update(buffer, value, t, dt):
            buffer[0] += 1.0

        @cuda.jit(device=True)
        def save(buffer, output, n_samples):
            output[0] = buffer[0]

        return MetricFuncCache(update=update, save=save)

After adding the file, the metric is available as output_types=["count"].