OutputArrays

class cubie.batchsolving.arrays.BatchOutputArrays.OutputArrays(precision: type = <class 'numpy.float32'>, chunks: int = 0, stream_group: str = 'default', memory_proportion: float | None = None, memory_manager: MemoryManager = MemoryManager(totalmem=8589934592, registry={}, stream_groups=StreamGroups(groups={}, streams={}), _mode='passive', _auto_pool=[], _manual_pool=[], _queued_allocations={}, _group_chunk_parameters={}, _usage_clock=0, _pending_teardowns=[], pinned_max_bytes=8589934592, allocation_granule_bytes=33554432, _pinned_lock=<unlocked _thread.lock object>, _pinned_live_bytes=0, _pinned_retained_bytes=0, _pinned_releases=deque([]), _device_probe_error=None), memory_owner: object | None = None, num_runs: int = 1, sizes: BatchOutputSizes = NOTHING, host: OutputArrayContainer = NOTHING)

Bases: BaseArrayManager

Manage batch integration output arrays between host and device.

This class manages the allocation, transfer, and synchronization of output arrays generated during batch integration operations. It handles state trajectories, observables, summary statistics, and per-run status codes.

Parameters:

• _sizes (cubie.outputhandling.output_sizes.BatchOutputSizes) – Size specifications for the output arrays.

• host (cubie.batchsolving.arrays.BatchOutputArrays.OutputArrayContainer) – Container for host-side arrays.

• device (cubie.batchsolving.arrays.BatchOutputArrays.OutputArrayContainer) – Container for device-side arrays.

Notes

This class is initialized with a BatchOutputSizes instance (which is drawn from a solver instance using the from_solver factory method), which sets the allowable 3D array sizes from the ODE system’s data and run settings. Once initialized, the object can be updated with a solver instance to update the expected sizes, check the cache, and allocate if required.

_stage_array(array_name, device_array, host_array, stream) → None

Stage one device output through pooled pinned buffers.

The host target may be a strided view (a chunk slice or a memmap), so completed blocks are written through strided assignment; flattening such a view would silently copy it and discard the writeback. Blocks come from staging_blocks(), and the buffer is trimmed to the host block’s shape because the device array can carry extra run-axis padding on the final chunk. Each block is handed to the writeback watcher with its own event: the watcher copies it to the host target and releases the buffer as soon as its transfer lands, so this method never blocks on the stream and the drain of one chunk overlaps the next chunk’s kernel.

_teardown_cleanups()

Return writeback cleanup calls.

device: OutputArrayContainer

property device_iteration_counters: NDArray | FakeCUDAArray | None

Device iteration counters output array.

property device_observable_summaries: NDArray | FakeCUDAArray | None

Device observable summary output array.

property device_observables: NDArray | FakeCUDAArray | None

Device observables output array.

property device_state: NDArray | FakeCUDAArray | None

Device state output array.

property device_state_summaries: NDArray | FakeCUDAArray | None

Device state summary output array.

property device_status_codes: NDArray | FakeCUDAArray | None

Device status code output array.

finalise(chunk_index: int, stream=None) → None

Queue device-to-host transfers for a chunk.

Parameters:

chunk_index – Indices for the chunk being finalized.

Notes

Host slices are made contiguous before transfer to ensure compatible strides with device arrays. For chunked mode, data is transferred to pooled pinned buffers and submitted to the watcher thread for async writeback. For non-chunked mode, the writeback call is made immediately (but will happen asynchronously).

classmethod from_solver(solver_instance: BatchSolverKernel) → OutputArrays

Create an OutputArrays instance from a solver.

Does not allocate arrays, just sets up size specifications.

Parameters:

solver_instance – The solver instance to extract configuration from.

Returns:

A new OutputArrays instance configured for the solver.

Return type:

OutputArrays

host: OutputArrayContainer

initialise(chunk_index: int, stream=None) → None

Initialize device arrays before kernel execution.

Parameters:

chunk_index – Indices for the chunk being initialized.

Notes

No initialization to zeros is needed unless chunk calculations in time leave a dangling sample at the end, which is possible but not expected.

property iteration_counters: NDArray | FakeCUDAArray | None

Host iteration counters output array.

property observable_summaries: NDArray | FakeCUDAArray | None

Host observable summary output array.

property observables: NDArray | FakeCUDAArray | None

Host observables output array.

reset() → None

Drain transfers and clear all arrays and staging buffers.

property state: NDArray | FakeCUDAArray | None

Host state output array.

property state_summaries: NDArray | FakeCUDAArray | None

Host state summary output array.

property status_codes: NDArray | FakeCUDAArray | None

Host status code output array.

update(solver_instance: BatchSolverKernel) → None

Update output arrays from solver instance.

Parameters:

solver_instance – The solver instance providing configuration and sizing information.

update_from_solver(solver_instance: BatchSolverKernel) → Dict[str, NDArray[floating]]

Update sizes and precision from solver, returning new host arrays.

Only creates new pinned arrays when existing arrays do not match the expected shape and dtype. This avoids expensive pinned memory allocation on repeated solver runs with identical configurations.

Parameters:

solver_instance – The solver instance to update from.

Returns:

Host arrays with updated shapes for update_host_arrays, or None when sizes are unchanged and the current host arrays already match.

Return type:

dict[str, numpy.ndarray] or None

wait_pending(timeout: float | None = None) → None

Wait for all pending async writebacks to complete.

Parameters:

timeout – Maximum seconds to wait. None waits indefinitely.

Notes

Only applies to chunked mode with watcher-based writebacks.