cubie.batchsolving.arrays.BatchInputArrays

Batch Input Arrays Module.

This module provides classes for managing input arrays in batch integration operations, including containers for storing arrays and managers for handling memory allocation and data transfer between host and device.

Classes

`InputArrayContainer`([stride_order, ...])	Container for batch input arrays.
`InputArrays`(precision, chunks, chunk_axis, ...)	Manage batch integration input arrays between host and device.

class cubie.batchsolving.arrays.BatchInputArrays.InputArrayContainer(stride_order: tuple[str, ...] | None = None, initial_values: ndarray[Any, dtype[_ScalarType_co]] | FakeCUDAArray | None = None, parameters: ndarray[Any, dtype[_ScalarType_co]] | FakeCUDAArray | None = None, forcing_vectors: ndarray[Any, dtype[_ScalarType_co]] | FakeCUDAArray | None = None, memory_type: str = 'device')[source]

Bases: ArrayContainer

Container for batch input arrays.

This container holds the input arrays needed for batch integration, including initial values, parameters, and forcing vectors.

Parameters:

initial_values (ArrayTypes, optional) – Initial state values for the integration.
parameters (ArrayTypes, optional) – Parameter values for the integration.
forcing_vectors (ArrayTypes, optional) – Forcing function vectors for the integration.
stride_order (tuple[str, ...], default=("run", "variable")) – Order of array dimensions.
_memory_type (str, default="device") – Type of memory allocation.
_unchunkable (tuple[str, ...], default=('forcing_vectors',)) – Array names that cannot be chunked.

Notes

This class uses attrs for automatic initialization and validation. The _unchunkable attribute specifies which arrays should not be divided into chunks during memory management.

classmethod device_factory()[source]

Create a new device memory container.

Returns:: A new container configured for device memory.
Return type:: InputArrayContainer

forcing_vectors: ndarray[Any, dtype[_ScalarType_co]] | FakeCUDAArray | None

classmethod host_factory()[source]

Create a new host memory container.

Returns:: A new container configured for host memory.
Return type:: InputArrayContainer

initial_values: ndarray[Any, dtype[_ScalarType_co]] | FakeCUDAArray | None

parameters: ndarray[Any, dtype[_ScalarType_co]] | FakeCUDAArray | None

class cubie.batchsolving.arrays.BatchInputArrays.InputArrays(precision: type = <class 'numpy.float32'>, chunks: int = 0, chunk_axis: str = 'run', stream_group: str = 'default', memory_proportion: float | None = None, memory_manager: ~cubie.memory.mem_manager.MemoryManager = MemoryManager(totalmem=8589934592, registry={}, stream_groups=StreamGroups(groups={}, streams={}), _mode='passive', _allocator=<class 'cubie.cudasim_utils.FakeNumbaCUDAMemoryManager'>, _auto_pool=[], _manual_pool=[], _stride_order=('time', 'run', 'variable'), _queued_allocations={}), sizes: ~cubie.outputhandling.output_sizes.BatchInputSizes | None = NOTHING, host: ~cubie.batchsolving.arrays.BatchInputArrays.InputArrayContainer = NOTHING)[source]

Bases: BaseArrayManager

Manage batch integration input arrays between host and device.

This class manages the allocation, transfer, and synchronization of input arrays needed for batch integration operations. It handles initial values, parameters, and forcing vectors.

Parameters:

_sizes (BatchInputSizes, optional) – Size specifications for the input arrays.
host (InputArrayContainer) – Container for host-side arrays.
device (InputArrayContainer) – Container for device-side arrays.

Notes

This class is initialized with a BatchInputSizes instance (which is drawn from a solver instance using the from_solver factory method), which sets the allowable array heights from the ODE system’s data. Once initialized, the object can be updated with arguments (initial_values, parameters, forcing_vectors). Each call to the update method will:

Check if the input array has changed in shape or content since the last update
Queue allocation requests with the MemoryManager
Attach allocated arrays once received from MemoryManager

device: InputArrayContainer

property device_forcing_vectors

Get device forcing vectors array.

Returns:: The forcing vectors array from the device container.
Return type:: ArrayTypes

property device_initial_values

Get device initial values array.

Returns:: The initial values array from the device container.
Return type:: ArrayTypes

property device_parameters

Get device parameters array.

Returns:: The parameters array from the device container.
Return type:: ArrayTypes

finalise(host_indices)[source]

Copy out final states if they’re required.

Parameters:: host_indices (slice or array-like) – Indices for the chunk being finalized.

Notes

This method copies data from device back to host for the specified chunk indices.

property forcing_vectors

Get host forcing vectors array.

Returns:: The forcing vectors array from the host container.
Return type:: ArrayTypes

classmethod from_solver(solver_instance: BatchSolverKernel) → InputArrays[source]

Create an InputArrays instance from a solver.

Creates an empty instance from a solver instance, importing the heights of the parameters, initial values, and driver arrays from the ODE system for checking inputs against. Does not allocate host or device arrays.

Parameters:: solver_instance (BatchSolverKernel) – The solver instance to extract configuration from.
Returns:: A new InputArrays instance configured for the solver.
Return type:: InputArrays

host: InputArrayContainer

property initial_values

Get host initial values array.

Returns:: The initial values array from the host container.
Return type:: ArrayTypes

initialise(host_indices)[source]

Copy chunk of data to device.

Parameters:: host_indices (slice or array-like) – Indices for the chunk being initialized.

Notes

This method copies the appropriate chunk of data from host to device arrays before kernel execution.

property parameters

Get host parameters array.

Returns:: The parameters array from the host container.
Return type:: ArrayTypes

update(solver_instance, initial_values: ndarray[Any, dtype[_ScalarType_co]], parameters: ndarray[Any, dtype[_ScalarType_co]], forcing_vectors: ndarray[Any, dtype[_ScalarType_co]]) → None[source]

Set the initial values, parameters, and forcing vectors.

Queues allocation requests with the MemoryManager for batch processing.

Parameters:

solver_instance (BatchSolverKernel) – The solver instance providing configuration and sizing information.
initial_values (NDArray) – Initial state values for each integration run.
parameters (NDArray) – Parameter values for each integration run.
forcing_vectors (NDArray) – Forcing vectors for each integration run.

update_from_solver(solver_instance: BatchSolverKernel)[source]

Update size and precision from solver instance.

Parameters:: solver_instance (BatchSolverKernel) – The solver instance to update from.