cubie.outputhandling.save_summaries

CUDA device function factory for saving summary metrics.

This module implements a recursive function chaining approach for CUDA device functions that save summary metrics. It compiles only the requested summary functions to optimize memory usage and performance.

Notes

This implementation is based on the “chain” approach by sklam from https://github.com/numba/numba/issues/3405. The approach allows iterating through jitted functions without passing them as an iterable, which isn’t supported by Numba.

The process consists of: 1. A recursive chain_metrics function that builds a chain of summary functions 2. A save_summary_factory that applies the chained functions to each variable

Functions

`chain_metrics`(metric_functions, ...[, ...])	Recursively chain summary metric functions for CUDA execution.
`save_summary_factory`(buffer_sizes, ...)	Factory function for creating CUDA device functions to save summary metrics.

cubie.outputhandling.save_summaries.chain_metrics(metric_functions: ~typing.Sequence, buffer_offsets, buffer_sizes, output_offsets, output_sizes, function_params, inner_chain=<numba.cuda.simulator.kernel.FakeCUDAKernel object>)[source]

Recursively chain summary metric functions for CUDA execution.

This function builds a recursive chain of summary metric functions, where each function in the sequence is wrapped with the previous functions to create a single callable that executes all metrics.

Parameters:

metric_functions (Sequence) – Sequence of CUDA device functions for summary metrics.
buffer_offsets (Sequence) – Buffer memory offsets for each metric function.
buffer_sizes (Sequence) – Buffer sizes required by each metric function.
output_offsets (Sequence) – Output array offsets for each metric function.
output_sizes (Sequence) – Output sizes for each metric function.
function_params (Sequence) – Parameters required by each metric function.
inner_chain (callable, default do_nothing) – Previously chained function to execute before current function.

Returns:

CUDA device function that executes all chained metrics.

Return type:

callable

Notes

The function uses recursion to build a chain where each level executes the inner chain first, then the current metric function. This ensures all requested metrics are computed in the correct order.

Factory function for creating CUDA device functions to save summary metrics.

This factory generates a CUDA device function that applies chained summary metric calculations to all requested state and observable variables.

Parameters:

buffer_sizes (SummariesBufferSizes) – Object containing buffer size information for summary calculations.
summarised_state_indices (Sequence[int] or ArrayLike) – Indices of state variables to include in summary calculations.
summarised_observable_indices (Sequence[int] or ArrayLike) – Indices of observable variables to include in summary calculations.
summaries_list (Sequence[str]) – List of summary metric names to compute.

Returns:

Compiled CUDA device function for saving summary metrics.

Return type:

callable

Notes

The generated function iterates through all specified state and observable variables, applying the chained summary metrics to each variable’s buffer and saving results to the appropriate output arrays.