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_releases=deque([]), _idle_pinned=[], _pinned_generation=0, _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.

_after_allocation() → None

Back the host outputs once the run is known to transfer.

_ensure_host_arrays() → None

Give every output a host buffer the transfer can land in.

A buffer is kept when it has the sized shape and dtype and a backing the run partition accepts: pinned only unchunked, pageable only when the policy would not pin it, memmap always. A new pinned buffer is zeroed only for an output no transfer writes.

_host_fits(label: str, slot: ManagedArray) → bool

Whether the slot’s host array has the sized shape and dtype.

_request_shape(label: str) → tuple

Device shape of label from the sizes, never a host array.

_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, transfer_outputs: bool = True) → None

Size the device outputs for a run and queue their allocation.

Parameters:

• solver_instance – The solver instance providing configuration and sizing information.

• transfer_outputs – Whether the run copies its outputs to host buffers.

update_from_solver(solver_instance: BatchSolverKernel) → bool | None

Refresh sizes and precision from the solver.

Parameters:

solver_instance – The solver instance to update from.

Returns:

True when the sizes changed and every device output is queued for reallocation, None when they are unchanged.

Return type:

bool 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.