GPU Acceleration

The GPU module provides hardware-accelerated implementations of key tracking algorithms using CuPy (NVIDIA CUDA) or MLX (Apple Silicon). These implementations offer significant speedups (5-15x) for batch processing of multiple tracks.

The module automatically selects the best available backend:

  • On Apple Silicon (M1/M2/M3): Uses MLX if installed

  • On systems with NVIDIA GPUs: Uses CuPy if installed

  • Falls back to CPU (numpy) if no GPU backend is available

Installation

For NVIDIA CUDA GPUs:

pip install nrl-tracker[gpu]
# or directly:
pip install cupy-cuda12x

For Apple Silicon (M1/M2/M3):

pip install nrl-tracker[gpu-apple]
# or directly:
pip install mlx

Quick Start

Check GPU availability and backend:

from pytcl.gpu import is_gpu_available, get_backend, is_apple_silicon

if is_gpu_available():
    print(f"GPU available, using {get_backend()} backend")

if is_apple_silicon():
    print("Running on Apple Silicon")

Transfer arrays between CPU and GPU:

from pytcl.gpu import to_gpu, to_cpu
import numpy as np

# CPU array
x = np.random.randn(100, 4)

# Transfer to GPU (uses best available backend)
x_gpu = to_gpu(x)

# Transfer back to CPU
x_cpu = to_cpu(x_gpu)

Platform Detection

pytcl.gpu.utils.is_apple_silicon()[source]

Check if running on Apple Silicon (ARM64 Mac).

Returns:

True if running on Apple Silicon (M1, M2, M3, etc.).

Return type:

bool

Examples

>>> from pytcl.gpu.utils import is_apple_silicon
>>> isinstance(is_apple_silicon(), bool)
True
pytcl.gpu.utils.is_mlx_available()[source]

Check if MLX acceleration is available (Apple Silicon).

Returns True if: - Running on Apple Silicon (ARM64 Mac) - MLX is installed

Returns:

True if MLX acceleration is available.

Return type:

bool

Examples

>>> from pytcl.gpu.utils import is_mlx_available
>>> isinstance(is_mlx_available(), bool)
True
pytcl.gpu.utils.is_cupy_available()[source]

Check if CuPy (CUDA) acceleration is available.

Returns True if: - CuPy is installed - A CUDA-capable GPU is detected - CUDA runtime is functional

Returns:

True if CuPy acceleration is available.

Return type:

bool

Examples

>>> from pytcl.gpu.utils import is_cupy_available
>>> isinstance(is_cupy_available(), bool)
True
pytcl.gpu.utils.get_backend()[source]

Report which GPU backend this platform would use (detection only).

Priority here is MLX on Apple Silicon, then CuPy, then the string "numpy" for “neither”. Note the COMPUTE selector, pytcl.gpu._backend.get_compute_backend, tries CuPy first – on the rare machine with both stacks installed, this function reports mlx while batch algorithms run on CuPy. Nothing computes on the numpy result; the compute path raises DependencyError instead.

Returns:

One of “mlx”, “cupy”, or “numpy”.

Return type:

str

Examples

>>> from pytcl.gpu.utils import get_backend
>>> get_backend() in ("mlx", "cupy", "numpy")
True
pytcl.gpu.utils.is_gpu_available()[source]

Check if GPU acceleration is available.

Returns True if either: - MLX is available (Apple Silicon) - CuPy is available with a CUDA GPU

Returns:

True if GPU acceleration is available.

Return type:

bool

Examples

>>> from pytcl.gpu.utils import is_gpu_available
>>> isinstance(is_gpu_available(), bool)
True

Notes

The result is cached after the first call for performance. Use get_backend() to determine which backend is being used.

Array Operations

pytcl.gpu.utils.to_gpu(arr, dtype=None, backend=None)[source]

Transfer an array to GPU memory.

Automatically selects the best available backend (MLX on Apple Silicon, CuPy on NVIDIA GPUs) unless a specific backend is requested.

Parameters:
  • arr (array_like) – Input array (typically numpy).

  • dtype (dtype, optional) –

    Data type for the GPU array. If None, uses the input dtype.

    Honoured on CuPy only. MLX does not support float64, so on Apple Silicon every float array lands as float32 whatever is requested – _numpy_dtype_to_mlx maps float64 to float32. Results computed through the MLX backend therefore carry single precision regardless of this argument.

  • backend (str, optional) – Specific backend to use (“mlx”, “cupy”). If None, auto-selects.

Returns:

Array in GPU memory (cupy.ndarray or mlx.array).

Return type:

GPUArray

Raises:

Examples

>>> import numpy as np
>>> from pytcl.gpu.utils import to_gpu, is_gpu_available
>>> x = np.array([1.0, 2.0, 3.0])
>>> x_gpu = to_gpu(x)
>>> np.allclose(to_cpu(x_gpu), x)
True

Notes

If the input is already a GPU array, it is returned as-is (or converted to the requested dtype).

pytcl.gpu.utils.to_cpu(arr)[source]

Transfer an array from GPU to CPU memory.

Parameters:

arr (array_like, cupy.ndarray, or mlx.array) – Input array (numpy, cupy, or mlx).

Returns:

Array in CPU memory.

Return type:

numpy.ndarray

Examples

>>> import numpy as np
>>> from pytcl.gpu.utils import to_gpu, to_cpu, is_gpu_available
>>> x = np.array([1.0, 2.0, 3.0])
>>> if is_gpu_available():
...     x_gpu = to_gpu(x)
...     x_cpu = to_cpu(x_gpu)
...     np.allclose(x, x_cpu)
True

Notes

If the input is already a numpy array, it is returned as-is.

pytcl.gpu.utils.get_array_module(arr)[source]

Get the array module (numpy, cupy, or mlx.core) for the given array.

This function enables writing code that works with numpy, cupy, and mlx arrays by returning the appropriate module.

Parameters:

arr (array_like) – Input array (numpy, cupy, or mlx).

Returns:

numpy, cupy, or mlx.core module, depending on the input array type.

Return type:

module

Examples

>>> import numpy as np
>>> from pytcl.gpu.utils import get_array_module
>>> x = np.array([1, 2, 3])
>>> xp = get_array_module(x)
>>> xp is np
True
>>> # With a CuPy array, on a machine that has CuPy
>>> import cupy as cp
>>> x_gpu = cp.array([1, 2, 3])
>>> get_array_module(x_gpu) is cp
True
>>> # With an MLX array, on a machine that has MLX
>>> import mlx.core as mx
>>> x_mlx = mx.array([1, 2, 3])
>>> xp = get_array_module(x_mlx)
>>> xp.__name__
'mlx.core'
pytcl.gpu.utils.ensure_gpu_array(arr, dtype=<class 'numpy.float64'>, backend=None)[source]

Ensure an array is on the GPU with the specified dtype.

Parameters:
  • arr (array_like) – Input array.

  • dtype (dtype) – Desired data type.

  • backend (str, optional) – Specific backend to use (“mlx”, “cupy”). If None, auto-selects.

Returns:

Array on GPU (cupy.ndarray or mlx.array). dtype is honoured on CuPy; on MLX the array is float32 regardless, so the declared float64 default is unreachable there.

Return type:

GPUArray

Examples

>>> import numpy as np
>>> from pytcl.gpu.utils import ensure_gpu_array, is_gpu_available
>>> x = np.array([1, 2, 3])
>>> x_gpu = ensure_gpu_array(x, dtype=np.float32)
>>> np.asarray(to_cpu(x_gpu)).dtype == np.float32
True

Memory Management

pytcl.gpu.utils.sync_gpu()[source]

Synchronize GPU operations.

This blocks until all pending GPU operations are complete. Useful for accurate timing measurements.

Examples

>>> from pytcl.gpu.utils import sync_gpu
>>> sync_gpu()  # returns once queued work has completed
pytcl.gpu.utils.get_gpu_memory_info()[source]

Get GPU memory usage information.

Returns:

Dictionary with keys: - ‘backend’: Backend in use (“mlx”, “cupy”, or “numpy”) - ‘free’: Free memory in bytes (if available) - ‘total’: Total memory in bytes (if available) - ‘used’: Used memory in bytes (if available)

Return type:

dict

Examples

>>> from pytcl.gpu.utils import get_gpu_memory_info
>>> info = get_gpu_memory_info()
>>> info["backend"] in ("mlx", "cupy", "numpy")
True
pytcl.gpu.utils.clear_gpu_memory()[source]

Clear GPU memory pools.

This frees cached memory blocks held by the GPU backend. Call this when you need to free GPU memory for other operations.

Examples

>>> from pytcl.gpu.utils import clear_gpu_memory, is_gpu_available
>>> if is_gpu_available():
...     # ... perform GPU operations ...
...     clear_gpu_memory()  # Free cached memory

Batch Kalman Filter

GPU-accelerated batch Kalman filter operations for processing multiple tracks in parallel. These functions provide 5-10x speedup compared to sequential CPU processing.

pytcl.gpu.kalman.batch_kf_predict(x, P, F, Q, B=None, u=None)[source]

Batch Kalman filter prediction for multiple tracks.

Performs the prediction step for N tracks in parallel on GPU:

x_pred[i] = F @ x[i] + B @ u[i] (if B, u provided) P_pred[i] = F @ P[i] @ F’ + Q

Parameters:
  • x (array_like) – Current state estimates, shape (n_tracks, state_dim).

  • P (array_like) – Current covariances, shape (n_tracks, state_dim, state_dim).

  • F (array_like) – State transition matrix, shape (state_dim, state_dim). Can also be (n_tracks, state_dim, state_dim) for track-specific matrices.

  • Q (array_like) – Process noise covariance, shape (state_dim, state_dim). Can also be (n_tracks, state_dim, state_dim) for track-specific noise.

  • B (array_like, optional) – Control input matrix, shape (state_dim, control_dim).

  • u (array_like, optional) – Control inputs, shape (n_tracks, control_dim).

Returns:

result – Named tuple with predicted states and covariances.

Return type:

BatchKalmanPrediction

Examples

>>> import numpy as np
>>> from pytcl.gpu.kalman import batch_kf_predict
>>> n_tracks = 100
>>> x = np.random.randn(n_tracks, 4)
>>> P = np.tile(np.eye(4) * 0.1, (n_tracks, 1, 1))
>>> F = np.array([[1, 1, 0, 0], [0, 1, 0, 0],
...               [0, 0, 1, 1], [0, 0, 0, 1]])
>>> Q = np.eye(4) * 0.01
>>> pred = batch_kf_predict(x, P, F, Q)
>>> pred.x.shape
(100, 4)
pytcl.gpu.kalman.batch_kf_update(x, P, z, H, R)[source]

Batch Kalman filter update for multiple tracks.

Performs the update step for N tracks in parallel on GPU:

y[i] = z[i] - H @ x[i] (innovation) S[i] = H @ P[i] @ H’ + R (innovation covariance) K[i] = P[i] @ H’ @ S[i]^{-1} (Kalman gain) x_upd[i] = x[i] + K[i] @ y[i] (updated state) P_upd[i] = (I - K[i] @ H) @ P[i] (updated covariance)

Parameters:
  • x (array_like) – Predicted state estimates, shape (n_tracks, state_dim).

  • P (array_like) – Predicted covariances, shape (n_tracks, state_dim, state_dim).

  • z (array_like) – Measurements, shape (n_tracks, meas_dim).

  • H (array_like) – Measurement matrix, shape (meas_dim, state_dim). Can also be (n_tracks, meas_dim, state_dim).

  • R (array_like) – Measurement noise covariance, shape (meas_dim, meas_dim). Can also be (n_tracks, meas_dim, meas_dim).

Returns:

result – Named tuple with updated states, covariances, and statistics.

Return type:

BatchKalmanUpdate

Examples

>>> import numpy as np
>>> from pytcl.gpu.kalman import batch_kf_update
>>> n_tracks = 100
>>> x = np.random.randn(n_tracks, 4)
>>> P = np.tile(np.eye(4) * 0.1, (n_tracks, 1, 1))
>>> z = np.random.randn(n_tracks, 2)  # position measurements
>>> H = np.array([[1, 0, 0, 0], [0, 0, 1, 0]])
>>> R = np.eye(2) * 0.5
>>> upd = batch_kf_update(x, P, z, H, R)
>>> upd.x.shape
(100, 4)
class pytcl.gpu.kalman.CuPyKalmanFilter(state_dim, meas_dim, F=None, H=None, Q=None, R=None)[source]

GPU-accelerated Linear Kalman Filter for batch processing.

This class provides a stateful interface for processing multiple tracks in parallel on the GPU. It maintains the filter matrices and provides methods for prediction and update.

Parameters:
  • state_dim (int) – Dimension of the state vector.

  • meas_dim (int) – Dimension of the measurement vector.

  • F (array_like, optional) – State transition matrix. If None, uses identity.

  • H (array_like, optional) – Measurement matrix. If None, measures first meas_dim states.

  • Q (array_like, optional) – Process noise covariance. If None, uses 0.01 * I.

  • R (array_like, optional) – Measurement noise covariance. If None, uses 1.0 * I.

Examples

>>> import numpy as np
>>> from pytcl.gpu.kalman import CuPyKalmanFilter
>>>
>>> # Create filter for 2D constant velocity model
>>> kf = CuPyKalmanFilter(
...     state_dim=4,  # [x, vx, y, vy]
...     meas_dim=2,   # [x, y]
...     F=np.array([[1, 1, 0, 0], [0, 1, 0, 0],
...                 [0, 0, 1, 1], [0, 0, 0, 1]]),
...     H=np.array([[1, 0, 0, 0], [0, 0, 1, 0]]),
...     Q=np.eye(4) * 0.1,
...     R=np.eye(2) * 1.0,
... )
>>>
>>> # Process batch of tracks
>>> n_tracks = 1000
>>> x = np.random.randn(n_tracks, 4)
>>> P = np.tile(np.eye(4), (n_tracks, 1, 1))
>>> z = np.random.randn(n_tracks, 2)
>>>
>>> # Predict and update
>>> x_pred, P_pred = kf.predict(x, P)
>>> result = kf.update(x_pred, P_pred, z)
__init__(state_dim, meas_dim, F=None, H=None, Q=None, R=None)[source]
predict(x, P, B=None, u=None)[source]

Perform batch prediction.

Parameters:
  • x (array_like) – State estimates, shape (n_tracks, state_dim).

  • P (array_like) – Covariances, shape (n_tracks, state_dim, state_dim).

  • B (array_like, optional) – Control input matrix.

  • u (array_like, optional) – Control inputs.

Returns:

  • x_pred (ndarray) – Predicted states.

  • P_pred (ndarray) – Predicted covariances.

Return type:

Tuple[ndarray[tuple[Any, …], dtype[floating]], ndarray[tuple[Any, …], dtype[floating]]]

update(x, P, z)[source]

Perform batch update.

Parameters:
  • x (array_like) – Predicted state estimates.

  • P (array_like) – Predicted covariances.

  • z (array_like) – Measurements.

Returns:

result – Update results including states, covariances, and statistics.

Return type:

BatchKalmanUpdate

predict_update(x, P, z, B=None, u=None)[source]

Combined batch prediction and update.

Parameters:
  • x (array_like) – Current state estimates.

  • P (array_like) – Current covariances.

  • z (array_like) – Measurements.

  • B (array_like, optional) – Control input matrix.

  • u (array_like, optional) – Control inputs.

Returns:

result – Update results.

Return type:

BatchKalmanUpdate

Batch Extended Kalman Filter

GPU-accelerated Extended Kalman Filter for nonlinear dynamics.

pytcl.gpu.ekf.batch_ekf_predict(x, P, f, F_jacobian, Q)[source]

Batch EKF prediction for multiple tracks.

Parameters:
  • x (array_like) – Current state estimates, shape (n_tracks, state_dim).

  • P (array_like) – Current covariances, shape (n_tracks, state_dim, state_dim).

  • f (callable) – Batched dynamics. Takes the whole (n_tracks, state_dim) device array and returns (n_tracks, state_dim). Called once.

  • F_jacobian (callable or None) – Batched Jacobian df/dx. Takes (n_tracks, state_dim) and returns (n_tracks, state_dim, state_dim). If None, computed numerically with 2 * state_dim evaluations of f over the whole batch.

  • Q (array_like) – Process noise covariance, shape (state_dim, state_dim) or (n_tracks, state_dim, state_dim).

Returns:

result – Predicted states and covariances, as device arrays of the active backend. Use pytcl.gpu.utils.to_cpu() to bring them back.

Return type:

BatchEKFPrediction

Raises:

DependencyError – If neither CuPy nor MLX is installed.

Examples

>>> import numpy as np
>>> from pytcl.gpu.ekf import batch_ekf_predict
>>> from pytcl.gpu.utils import get_array_module
>>> # Coordinated turn, evaluated for the whole batch at once
>>> def f_turn(x):
...     xp = get_array_module(x)
...     w = 0.01
...     return xp.stack([x[:, 0] + xp.cos(w) * x[:, 2],
...                      x[:, 1] + xp.sin(w) * x[:, 3],
...                      x[:, 2], x[:, 3]], axis=1)
>>> def F_jacobian(x):
...     xp = get_array_module(x)
...     w = 0.01
...     J = xp.array([[1.0, 0.0, xp.cos(w).item(), 0.0],
...                   [0.0, 1.0, xp.sin(w).item(), 0.0],
...                   [0.0, 0.0, 1.0, 0.0],
...                   [0.0, 0.0, 0.0, 1.0]])
...     return xp.broadcast_to(J, (x.shape[0], 4, 4))
>>> n_tracks = 30
>>> x = np.random.randn(n_tracks, 4) * 0.1
>>> P = np.tile(np.eye(4) * 0.01, (n_tracks, 1, 1))
>>> Q = np.eye(4) * 0.001
>>> result = batch_ekf_predict(x, P, f_turn, F_jacobian, Q)
>>> result.x.shape
(30, 4)

Notes

Both the dynamics and the covariance propagation stay on the device. The callback is invoked once for the batch rather than once per track.

pytcl.gpu.ekf.batch_ekf_update(x, P, z, h, H_jacobian, R)[source]

Batch EKF update for multiple tracks.

Parameters:
  • x (array_like) – Predicted state estimates, shape (n_tracks, state_dim).

  • P (array_like) – Predicted covariances, shape (n_tracks, state_dim, state_dim).

  • z (array_like) – Measurements, shape (n_tracks, meas_dim).

  • h (callable) – Batched measurement function. Takes (n_tracks, state_dim) and returns (n_tracks, meas_dim). Called once.

  • H_jacobian (callable or None) – Batched Jacobian dh/dx. Takes (n_tracks, state_dim) and returns (n_tracks, meas_dim, state_dim). If None, computed numerically.

  • R (array_like) – Measurement noise covariance.

Returns:

result – Update results including states, covariances, and statistics, as device arrays of the active backend.

Return type:

BatchEKFUpdate

Raises:

DependencyError – If neither CuPy nor MLX is installed.

Examples

>>> import numpy as np
>>> from pytcl.gpu.ekf import batch_ekf_update
>>> from pytcl.gpu.utils import get_array_module
>>> # Polar measurement from a Cartesian state, batched
>>> def h_polar(x):
...     xp = get_array_module(x)
...     r = xp.sqrt(x[:, 0] ** 2 + x[:, 1] ** 2)
...     theta = xp.arctan2(x[:, 1], x[:, 0])
...     return xp.stack([r, theta], axis=1)
>>> def H_jacobian(x):
...     xp = get_array_module(x)
...     r = xp.sqrt(x[:, 0] ** 2 + x[:, 1] ** 2)
...     row0 = xp.stack([x[:, 0] / r, x[:, 1] / r], axis=1)
...     row1 = xp.stack([-x[:, 1] / r**2, x[:, 0] / r**2], axis=1)
...     return xp.stack([row0, row1], axis=1)
>>> n_tracks = 20
>>> x = np.random.randn(n_tracks, 2)
>>> P = np.tile(np.eye(2), (n_tracks, 1, 1))
>>> z = np.random.randn(n_tracks, 2) * [100, 0.1]  # r, theta
>>> R = np.diag([10.0, 0.01])
>>> result = batch_ekf_update(x, P, z, h_polar, H_jacobian, R)
>>> result.x.shape
(20, 2)
class pytcl.gpu.ekf.CuPyExtendedKalmanFilter(state_dim, meas_dim, f, h, F_jacobian=None, H_jacobian=None, Q=None, R=None)[source]

GPU-accelerated Extended Kalman Filter for batch processing.

Despite the historical name, this class runs on whichever GPU backend is available: CuPy on CUDA devices, MLX on Apple Silicon.

Parameters:
  • state_dim (int) – Dimension of state vector.

  • meas_dim (int) – Dimension of measurement vector.

  • f (callable) – Nonlinear dynamics function f(x) -> x_next.

  • h (callable) – Nonlinear measurement function h(x) -> z.

  • F_jacobian (callable, optional) – Jacobian of dynamics. If None, computed numerically.

  • H_jacobian (callable, optional) – Jacobian of measurement. If None, computed numerically.

  • Q (array_like, optional) – Process noise covariance.

  • R (array_like, optional) – Measurement noise covariance.

Raises:

DependencyError – If neither CuPy nor MLX is installed.

Examples

>>> import numpy as np
>>> from pytcl.gpu.ekf import CuPyExtendedKalmanFilter
>>>
>>> # Nonlinear dynamics. Callbacks receive the WHOLE (n_tracks, dim)
>>> # batch as one device array, not a single state -- see the module
>>> # docstring's contract. Indexing x[0] here would take the first
>>> # track rather than the first state component, and raise on predict.
>>> def f(x):
...     return np.stack([x[:, 0] + x[:, 1], x[:, 1] * 0.99], axis=-1)
>>>
>>> def h(x):
...     return np.sqrt(x[:, 0] ** 2 + x[:, 1] ** 2)[:, None]
>>>
>>> ekf = CuPyExtendedKalmanFilter(
...     state_dim=2, meas_dim=1,
...     f=f, h=h,
...     Q=np.eye(2) * 0.01,
...     R=np.array([[0.1]]),
... )
__init__(state_dim, meas_dim, f, h, F_jacobian=None, H_jacobian=None, Q=None, R=None)[source]
predict(x, P)[source]

Perform batch EKF prediction.

update(x, P, z)[source]

Perform batch EKF update.

predict_update(x, P, z)[source]

Combined prediction and update.

Batch Unscented Kalman Filter

GPU-accelerated Unscented Kalman Filter for highly nonlinear systems.

pytcl.gpu.ukf.batch_ukf_predict(x, P, f, Q, alpha=0.001, beta=2.0, kappa=0.0)[source]

Batch UKF prediction for multiple tracks.

Parameters:
  • x (array_like) – Current state estimates, shape (n_tracks, state_dim).

  • P (array_like) – Current covariances, shape (n_tracks, state_dim, state_dim).

  • f (callable) – Batched dynamics. Takes (N, state_dim) and returns (N, state_dim), where N is n_tracks * (2 * state_dim + 1) sigma points. Called once.

  • Q (array_like) – Process noise covariance.

  • alpha (float) – Sigma point parameters.

  • beta (float) – Sigma point parameters.

  • kappa (float) – Sigma point parameters.

Returns:

result – Predicted states and covariances.

Return type:

BatchUKFPrediction

Warns:

RuntimeWarning – If the active backend computes in float32 (MLX) and alpha is below 1e-2, where the O(1/alpha**2) Merwe weights destroy every significant digit. See Notes.

Notes

On the float32 MLX backend the maximum relative error against the float64 CPU UKF, on a linear problem where the UKF reduces to the Kalman filter, is 5.8e+01 at alpha=1e-3, 4.6e-03 at 1e-2, 1.9e-05 at 0.1 and 1.9e-06 at 1.0 (the O(1/alpha**2) Merwe weights amplify float32 rounding). Use alpha >= 0.1 on MLX. In float64 the same errors are 1.9e-10, 1.1e-12, 2.1e-14 and 6.6e-16. See the module Notes.

Examples

>>> import numpy as np
>>> from pytcl.gpu.ukf import batch_ukf_predict
>>> # Nonlinear dynamics example
>>> from pytcl.gpu.utils import get_array_module
>>> def f_dynamics(x):
...     xp = get_array_module(x)
...     return xp.stack([x[:, 0] + 0.1 * x[:, 1], x[:, 1] * 0.99], axis=1)
>>> n_tracks = 50
>>> x = np.random.randn(n_tracks, 2)
>>> P = np.tile(np.eye(2) * 0.01, (n_tracks, 1, 1))
>>> Q = np.eye(2) * 0.001
>>> result = batch_ukf_predict(x, P, f_dynamics, Q)
>>> result.x.shape
(50, 2)
pytcl.gpu.ukf.batch_ukf_update(x, P, z, h, R, alpha=0.001, beta=2.0, kappa=0.0)[source]

Batch UKF update for multiple tracks.

Parameters:
  • x (array_like) – Predicted state estimates, shape (n_tracks, state_dim).

  • P (array_like) – Predicted covariances, shape (n_tracks, state_dim, state_dim).

  • z (array_like) – Measurements, shape (n_tracks, meas_dim).

  • h (callable) – Batched measurement function. Takes (N, state_dim) and returns (N, meas_dim), where N is n_tracks * (2 * state_dim + 1) sigma points. Called once.

  • R (array_like) – Measurement noise covariance.

  • alpha (float) – Sigma point parameters.

  • beta (float) – Sigma point parameters.

  • kappa (float) – Sigma point parameters.

Returns:

result – Update results.

Return type:

BatchUKFUpdate

Warns:

RuntimeWarning – If the active backend computes in float32 (MLX) and alpha is below 1e-2, where the O(1/alpha**2) Merwe weights destroy every significant digit. See Notes.

Notes

On the float32 MLX backend the maximum relative error against the float64 CPU UKF, on a linear problem where the UKF reduces to the Kalman filter, is 5.8e+01 at alpha=1e-3, 4.6e-03 at 1e-2, 1.9e-05 at 0.1 and 1.9e-06 at 1.0 (the O(1/alpha**2) Merwe weights amplify float32 rounding). Use alpha >= 0.1 on MLX. In float64 the same errors are 1.9e-10, 1.1e-12, 2.1e-14 and 6.6e-16. See the module Notes.

Examples

>>> import numpy as np
>>> from pytcl.gpu.ukf import batch_ukf_update
>>> # Nonlinear measurement example
>>> from pytcl.gpu.utils import get_array_module
>>> def h_measurement(x):  # Range-only
...     xp = get_array_module(x)
...     return xp.sqrt(x[:, 0] ** 2 + x[:, 1] ** 2)[:, None]
>>> n_tracks = 40
>>> x = np.random.randn(n_tracks, 2)
>>> P = np.tile(np.eye(2), (n_tracks, 1, 1))
>>> z = np.random.randn(n_tracks, 1) * 10 + 100
>>> R = np.array([[1.0]])
>>> result = batch_ukf_update(x, P, z, h_measurement, R)
>>> result.x.shape
(40, 2)
class pytcl.gpu.ukf.CuPyUnscentedKalmanFilter(state_dim, meas_dim, f, h, Q=None, R=None, alpha=0.001, beta=2.0, kappa=0.0)[source]

GPU-accelerated Unscented Kalman Filter for batch processing.

Runs on whichever compute backend is available (CuPy on CUDA, MLX on Apple Silicon); the name is retained for backwards compatibility.

Parameters:
  • state_dim (int) – Dimension of state vector.

  • meas_dim (int) – Dimension of measurement vector.

  • f (callable) – Nonlinear dynamics function.

  • h (callable) – Nonlinear measurement function.

  • Q (array_like, optional) – Process noise covariance.

  • R (array_like, optional) – Measurement noise covariance.

  • alpha (float) – Spread of sigma points (default 1e-3). On the float32 MLX backend use 0.1 or larger; see the module Notes.

  • beta (float) – Prior knowledge parameter (default 2.0).

  • kappa (float) – Secondary scaling (default 0.0).

Examples

>>> import numpy as np
>>> from pytcl.gpu.ukf import CuPyUnscentedKalmanFilter
>>>
>>> from pytcl.gpu.utils import get_array_module
>>> def f(x):
...     xp = get_array_module(x)
...     return xp.stack([x[:, 0] + x[:, 1], x[:, 1]], axis=1)
>>>
>>> def h(x):
...     xp = get_array_module(x)
...     return xp.sqrt(x[:, 0] ** 2 + x[:, 1] ** 2)[:, None]
>>>
>>> ukf = CuPyUnscentedKalmanFilter(
...     state_dim=2, meas_dim=1,
...     f=f, h=h,
... )
__init__(state_dim, meas_dim, f, h, Q=None, R=None, alpha=0.001, beta=2.0, kappa=0.0)[source]
predict(x, P)[source]

Perform batch UKF prediction.

update(x, P, z)[source]

Perform batch UKF update.

predict_update(x, P, z)[source]

Combined prediction and update.

GPU Particle Filter

GPU-accelerated particle filtering with efficient resampling algorithms.

pytcl.gpu.particle_filter.gpu_resample_systematic(weights, seed=None)[source]

GPU-accelerated systematic resampling.

Systematic resampling uses a single random number to select particles, resulting in low variance and O(N) complexity.

Parameters:
  • weights (array_like) – Normalized particle weights, shape (n_particles,).

  • seed (int, optional) – Seed for the single uniform draw. If None, the backend’s global random state is used.

Returns:

indices – Resampled particle indices, shape (n_particles,).

Return type:

ndarray

Examples

>>> import numpy as np
>>> from pytcl.gpu.particle_filter import gpu_resample_systematic
>>> weights = np.array([0.1, 0.3, 0.4, 0.2])
>>> indices = gpu_resample_systematic(weights)
>>> # Particles 1 and 2 will be selected more often

Notes

Every particle is selected either floor(n * w_i) or ceil(n * w_i) times, i.e. |count_i - n * w_i| < 1.

pytcl.gpu.particle_filter.gpu_resample_multinomial(weights, seed=None)[source]

GPU-accelerated multinomial resampling.

Multinomial resampling samples particles independently according to their weights.

Parameters:
  • weights (array_like) – Normalized particle weights, shape (n_particles,).

  • seed (int, optional) – Seed for the uniform draws. If None, the backend’s global random state is used.

Returns:

indices – Resampled particle indices, shape (n_particles,).

Return type:

ndarray

Examples

>>> import numpy as np
>>> from pytcl.gpu.particle_filter import gpu_resample_multinomial
>>> from pytcl.gpu.utils import to_cpu
>>> weights = np.array([0.1, 0.4, 0.5])
>>> indices = np.asarray(to_cpu(gpu_resample_multinomial(weights)))
>>> indices.shape
(3,)
>>> bool(np.all(indices < 3))
True

Notes

Multinomial resampling has higher variance than systematic resampling but is simpler and can be more efficient on GPU for certain sizes.

pytcl.gpu.particle_filter.gpu_resample_stratified(weights, seed=None)[source]

GPU-accelerated stratified resampling.

Stratified resampling divides the CDF into N equal strata and samples one particle from each stratum.

Parameters:
  • weights (array_like) – Normalized particle weights, shape (n_particles,).

  • seed (int, optional) – Seed for the uniform draws. If None, the backend’s global random state is used.

Returns:

indices – Resampled particle indices, shape (n_particles,).

Return type:

ndarray

pytcl.gpu.particle_filter.gpu_effective_sample_size(weights)[source]

Compute effective sample size on GPU.

ESS = 1 / sum(w_i^2)

Parameters:

weights (array_like) – Normalized particle weights.

Returns:

ess – Effective sample size.

Return type:

float

Examples

>>> import numpy as np
>>> from pytcl.gpu.particle_filter import gpu_effective_sample_size
>>> weights = np.array([0.1, 0.2, 0.3, 0.4])
>>> ess = gpu_effective_sample_size(weights)
>>> ess > 0
True
>>> ess <= len(weights)
True
pytcl.gpu.particle_filter.gpu_normalize_weights(log_weights)[source]

Normalize log weights to proper weights on GPU.

Uses log-sum-exp trick for numerical stability.

Parameters:

log_weights (array_like) – Unnormalized log weights, shape (n_particles,).

Returns:

  • weights (ndarray) – Normalized weights, shape (n_particles,).

  • log_likelihood (float) – Log of the normalization constant.

Return type:

Tuple[ndarray[tuple[Any, …], dtype[floating[Any]]], float]

Examples

>>> import numpy as np
>>> from pytcl.gpu.particle_filter import gpu_normalize_weights
>>> log_w = np.array([-1.0, -0.5, -2.0])
>>> from pytcl.gpu.utils import to_cpu
>>> weights, log_likelihood = gpu_normalize_weights(log_w)
>>> bool(np.allclose(np.asarray(to_cpu(weights)).sum(), 1.0))
True
>>> bool(np.isclose(float(np.asarray(to_cpu(log_likelihood))),
...                 np.log(np.exp(log_w).sum())))
True
class pytcl.gpu.particle_filter.CuPyParticleFilter(n_particles, state_dim, resample_method='systematic', resample_threshold=0.5)[source]

GPU-accelerated Bootstrap Particle Filter.

This class implements the Sequential Importance Resampling (SIR) particle filter with GPU acceleration.

Parameters:
  • n_particles (int) – Number of particles.

  • state_dim (int) – Dimension of state vector.

  • resample_method (str) – Resampling method: ‘systematic’, ‘multinomial’, or ‘stratified’.

  • resample_threshold (float) – ESS threshold for resampling (as fraction of n_particles).

Variables:
  • particles (GPUArray) – Current particle states, shape (n_particles, state_dim).

  • weights (GPUArray) – Current particle weights, shape (n_particles,).

Examples

>>> import numpy as np
>>> from pytcl.gpu.particle_filter import CuPyParticleFilter
>>>
>>> pf = CuPyParticleFilter(n_particles=1000, state_dim=4)
>>> pf.initialize(np.zeros(4), np.eye(4))
>>> dynamics_fn = lambda particles: particles * 0.99
>>> likelihood_fn = lambda particles, z: np.exp(-0.5 * (particles[:, 0] - z) ** 2)
>>> for measurement in (0.1, 0.2, 0.3):
...     pf.predict(dynamics_fn)
...     _ = pf.update(measurement, likelihood_fn)
>>> pf.get_estimate().shape
(4,)
__init__(n_particles, state_dim, resample_method='systematic', resample_threshold=0.5)[source]
initialize(mean, cov)[source]

Initialize particles from Gaussian distribution.

Parameters:
  • mean (array_like) – Mean state, shape (state_dim,).

  • cov (array_like) – Covariance matrix, shape (state_dim, state_dim).

initialize_uniform(low, high)[source]

Initialize particles from uniform distribution.

Parameters:
  • low (array_like) – Lower bounds, shape (state_dim,).

  • high (array_like) – Upper bounds, shape (state_dim,).

predict(dynamics_fn, *args, **kwargs)[source]

Propagate particles through dynamics.

Parameters:
  • dynamics_fn (callable) – Function that takes particles (N, state_dim) and returns propagated particles (N, state_dim).

  • *args (Any) – Additional arguments passed to dynamics_fn.

  • **kwargs (Any) – Additional arguments passed to dynamics_fn.

Notes

The dynamics function receives backend arrays (CuPy or MLX). It should return arrays of the same type.

update(measurement, likelihood_fn)[source]

Update weights based on measurement likelihood.

Parameters:
  • measurement (array_like) – Measurement vector.

  • likelihood_fn (callable) – Function that computes likelihood for each particle. Takes (particles, measurement) and returns likelihoods (n_particles,).

Returns:

log_likelihood – Log of the marginal likelihood (normalization constant).

Return type:

float

get_estimate()[source]

Compute weighted mean estimate.

Returns:

estimate – Weighted mean state, shape (state_dim,).

Return type:

ndarray

get_covariance()[source]

Compute weighted covariance estimate.

Returns:

cov – Weighted covariance, shape (state_dim, state_dim).

Return type:

ndarray

get_ess()[source]

Get current effective sample size.

get_state()[source]

Get current filter state.

Returns:

state – Named tuple with particles, weights, and ESS.

Return type:

ParticleFilterState

get_particles_cpu()[source]

Get particles on CPU.

get_weights_cpu()[source]

Get weights on CPU.

GPU Matrix Utilities

GPU-accelerated matrix operations commonly used in tracking algorithms.

pytcl.gpu.matrix_utils.gpu_cholesky(A, lower=True)[source]

GPU-accelerated Cholesky decomposition.

Computes L such that A = L @ L.T (lower=True) or A = U.T @ U (lower=False).

Parameters:
  • A (array_like) – Symmetric positive definite matrix, shape (n, n) or batch (k, n, n).

  • lower (bool) – If True, return lower triangular. If False, return upper triangular.

Returns:

L – Cholesky factor, same shape as A.

Return type:

ndarray

Raises:

numpy.linalg.LinAlgError – If matrix is not positive definite.

Notes

The factor is validated on the host before it is returned, because neither GPU backend raises reliably for non-positive-definite input. This forces a device synchronization.

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_cholesky
>>> from pytcl.gpu.utils import to_cpu
>>> A = np.array([[4, 2], [2, 3]])
>>> L = to_cpu(gpu_cholesky(A))
>>> np.allclose(L @ L.T, A)
True
pytcl.gpu.matrix_utils.gpu_cholesky_safe(A, lower=True, regularization=1e-10)[source]

GPU Cholesky decomposition with fallback for non-positive-definite matrices.

If standard Cholesky fails, adds regularization to diagonal and retries.

Parameters:
  • A (array_like) – Symmetric matrix, shape (n, n) or batch (k, n, n).

  • lower (bool) – Return lower (True) or upper (False) triangular factor.

  • regularization (float) – Amount to add to diagonal if matrix is not positive definite.

Returns:

  • L (ndarray) – Cholesky factor.

  • success (bool) – True if succeeded without regularization.

Return type:

Tuple[ndarray[tuple[Any, …], dtype[floating[Any]]], bool]

Notes

Diagonal regularization only repairs a matrix that is positive semi-definite but singular; it cannot repair an indefinite one. When the regularized retry also fails, the factor is instead computed for the nearest positive definite matrix, obtained by flooring the eigenvalues of A at regularization. This function therefore always returns a factor and never raises for a non-positive-definite input; success reports whether any repair was needed.

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_cholesky_safe
>>> A = np.array([[1, 2], [2, 1]])  # Not positive definite
>>> L, success = gpu_cholesky_safe(A)
>>> success
False
pytcl.gpu.matrix_utils.gpu_qr(A, mode='reduced')[source]

GPU-accelerated QR decomposition.

Computes A = Q @ R where Q is orthogonal and R is upper triangular.

Parameters:
  • A (array_like) – Matrix to decompose, shape (m, n) or batch (k, m, n).

  • mode (str) – ‘reduced’ (default) or ‘complete’.

Returns:

  • Q (ndarray) – Orthogonal matrix.

  • R (ndarray) – Upper triangular matrix.

Return type:

Tuple[ndarray[tuple[Any, …], dtype[floating[Any]]], ndarray[tuple[Any, …], dtype[floating[Any]]]]

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_qr
>>> from pytcl.gpu.utils import to_cpu
>>> A = np.random.randn(4, 3)
>>> Q, R = gpu_qr(A)
>>> np.allclose(to_cpu(Q) @ to_cpu(R), A)
True
pytcl.gpu.matrix_utils.gpu_solve(A, b)[source]

GPU-accelerated linear system solve.

Solves A @ x = b for x.

Parameters:
  • A (array_like) – Coefficient matrix, shape (n, n) or batch (k, n, n).

  • b (array_like) – Right-hand side, shape (n,) or (n, m) or batch (k, n).

Returns:

x – Solution vector/matrix.

Return type:

ndarray

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_solve
>>> from pytcl.gpu.utils import to_cpu
>>> A = np.array([[3, 1], [1, 2]])
>>> b = np.array([9, 8])
>>> x = to_cpu(gpu_solve(A, b))
>>> np.allclose(A @ x, b)
True
pytcl.gpu.matrix_utils.gpu_inv(A)[source]

GPU-accelerated matrix inversion.

Parameters:

A (array_like) – Matrix to invert, shape (n, n) or batch (k, n, n).

Returns:

A_inv – Inverse matrix.

Return type:

ndarray

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_inv
>>> from pytcl.gpu.utils import to_cpu
>>> A = np.array([[1, 2], [3, 4]])
>>> A_inv = to_cpu(gpu_inv(A))
>>> np.allclose(A @ A_inv, np.eye(2))
True
pytcl.gpu.matrix_utils.gpu_eigh(A)[source]

GPU-accelerated eigendecomposition for symmetric matrices.

Computes eigenvalues and eigenvectors of symmetric matrix A.

Parameters:

A (array_like) – Symmetric matrix, shape (n, n) or batch (k, n, n).

Returns:

  • eigenvalues (ndarray) – Eigenvalues in ascending order.

  • eigenvectors (ndarray) – Corresponding eigenvectors as columns.

Return type:

Tuple[ndarray[tuple[Any, …], dtype[floating[Any]]], ndarray[tuple[Any, …], dtype[floating[Any]]]]

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_eigh
>>> from pytcl.gpu.utils import to_cpu
>>> A = np.array([[2, 1], [1, 2]])
>>> eigvals, eigvecs = gpu_eigh(A)
>>> bool(np.allclose(np.asarray(to_cpu(eigvals)), [1.0, 3.0]))
True
pytcl.gpu.matrix_utils.gpu_matrix_sqrt(A)[source]

GPU-accelerated matrix square root for positive definite matrices.

Computes S such that S @ S = A using eigendecomposition.

Parameters:

A (array_like) – Symmetric positive definite matrix.

Returns:

S – Matrix square root.

Return type:

ndarray

Examples

>>> import numpy as np
>>> from pytcl.gpu.matrix_utils import gpu_matrix_sqrt
>>> from pytcl.gpu.utils import to_cpu
>>> A = np.array([[4, 0], [0, 9]])
>>> S = to_cpu(gpu_matrix_sqrt(A))
>>> np.allclose(S @ S, A)
True
class pytcl.gpu.matrix_utils.MemoryPool[source]

GPU memory pool manager for efficient memory allocation.

Wraps CuPy’s memory pool on NVIDIA GPUs and MLX’s allocator on Apple Silicon, adding monitoring and limit management. With no GPU backend installed every method is a no-op.

Examples

>>> from pytcl.gpu.matrix_utils import MemoryPool
>>> pool = MemoryPool()
>>> stats = pool.get_stats()
>>> sorted(stats)
['device_total', 'free', 'total', 'used']
>>>
>>> # Free cached memory
>>> pool.free_all()
__init__()[source]

Initialize memory pool manager.

get_stats()[source]

Get memory pool statistics.

Returns:

stats – Dictionary with ‘used’, ‘total’, ‘free’, and ‘device_total’ bytes.

Return type:

dict

Examples

>>> from pytcl.gpu.matrix_utils import get_memory_pool
>>> pool = get_memory_pool()
>>> stats = pool.get_stats()
>>> sorted(stats)
['device_total', 'free', 'total', 'used']
>>> stats['used'] >= 0
True
free_all()[source]

Free all cached memory blocks.

Clears the memory pool cache, which can help free up GPU memory when operations are complete.

Examples

>>> from pytcl.gpu.matrix_utils import get_memory_pool
>>> pool = get_memory_pool()
>>> # After allocations
>>> pool.free_all()  # Clear cached blocks
set_limit(limit=None)[source]

Set memory pool limit.

Parameters:

limit (int or None) – Maximum bytes to allocate. None restores the backend default (unlimited on CuPy).

Examples

>>> from pytcl.gpu.matrix_utils import get_memory_pool
>>> pool = get_memory_pool()
>>> # Limit to 2 GB
>>> pool.set_limit(2 * 1024**3)
>>> # Reset to the backend default
>>> pool.set_limit(None)
limit_memory(max_bytes)[source]

Context manager for temporary memory limit.

Parameters:

max_bytes (int) – Maximum bytes allowed during context.

Examples

>>> pool = MemoryPool()
>>> with pool.limit_memory(10**9):  # 1GB limit
...     # Operations here have limited memory
...     pass

Example: Batch Track Processing

Process multiple tracks in parallel using GPU acceleration:

import numpy as np
from pytcl.gpu import (
    is_gpu_available,
    to_gpu,
    to_cpu,
    batch_kf_predict,
    batch_kf_update,
)

if not is_gpu_available():
    raise RuntimeError("GPU not available")

# Simulate 1000 tracks with 4D state (x, vx, y, vy)
n_tracks = 1000
state_dim = 4
meas_dim = 2

# Initial states and covariances
x = np.random.randn(n_tracks, state_dim)
P = np.tile(np.eye(state_dim), (n_tracks, 1, 1))

# System matrices
dt = 0.1
F = np.array([
    [1, dt, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, dt],
    [0, 0, 0, 1]
])
Q = np.eye(state_dim) * 0.1
H = np.array([[1, 0, 0, 0], [0, 0, 1, 0]])
R = np.eye(meas_dim) * 0.5

# Transfer to GPU
x_gpu = to_gpu(x)
P_gpu = to_gpu(P)

# Batch predict (all 1000 tracks at once!)
pred_result = batch_kf_predict(x_gpu, P_gpu, F, Q)

# Generate measurements
z = np.random.randn(n_tracks, meas_dim)

# Batch update
upd_result = batch_kf_update(
    pred_result.x, pred_result.P, z, H, R
)

# Transfer results back to CPU
x_updated = to_cpu(upd_result.x)
P_updated = to_cpu(upd_result.P)

print(f"Processed {n_tracks} tracks in batch")

Performance Notes

The GPU implementations achieve significant speedups for:

  • Large batch sizes: Processing 100+ tracks simultaneously

  • Large particle counts: Particle filters with 1000+ particles

  • Matrix operations: Cholesky, QR, and eigendecompositions

For small batch sizes (< 10 tracks), CPU implementations may be faster due to GPU transfer overhead.

Backend Differences

CuPy (NVIDIA CUDA): - Full float64 (double precision) support - Explicit memory pool management - CUDA stream synchronization

MLX (Apple Silicon): - Optimized for float32 (single precision) - Automatic memory management - Lazy evaluation with explicit sync via mx.eval()