DOC: add docs for array backends

docs/array_backends.rst

@@ -0,0 +1,10 @@
+.. include:: ../examples/array_backends/README.rst
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Backend walkthroughs
+
+   array_backends/numpy
+   array_backends/cupy
+   array_backends/torch
+   array_backends/jax

docs/array_backends/cupy.rst

@@ -0,0 +1 @@
+.. include:: ../../examples/array_backends/cupy.rst

docs/array_backends/jax.rst

@@ -0,0 +1 @@
+.. include:: ../../examples/array_backends/jax.rst

docs/array_backends/numpy.rst

@@ -0,0 +1 @@
+.. include:: ../../examples/array_backends/numpy.rst

docs/array_backends/torch.rst

@@ -0,0 +1 @@
+.. include:: ../../examples/array_backends/torch.rst

docs/index.rst

@@ -36,6 +36,15 @@ and observation-conditioned proposals with a Kalman-filter reference.
 
    _generated/examples/particle_filtering/README
 
+Use the array backend examples to run differential evolution with NumPy, CuPy, or PyTorch,
+and to compile an SMC iteration with JAX's functional API.
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Array backends
+
+   array_backends
+
 The basic, adaptive mutation, flow preconditioning, persistent sampling, and
 waste-free SMC, and particle-filtering tutorials include plots alongside their
 explanations. Their

examples/array_backends/README.rst

@@ -0,0 +1,199 @@
+Array backends
+==============
+
+Using array backends
+--------------------
+
+Sequince can sample with NumPy, CuPy, PyTorch, and JAX arrays. The particle
+arrays select the numerical backend, and the model callbacks must evaluate
+their log densities using compatible array operations. NumPy runs on CPU;
+CuPy requires CUDA; PyTorch and JAX can run on CPU or supported accelerators.
+Backend support also depends on the mutation method and other components you
+choose. The standard differential-evolution (DE) sampler used here supports
+all four backends.
+
+Eager and functional execution
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Sequince provides eager and functional ways to execute a sampler.
+
+In the eager API, call ``smc.run(initial_state, rng=rng, metadata=metadata)``.
+The sampler advances a mutable random generator as it applies weighting,
+resampling, and mutation steps. The NumPy, CuPy, and PyTorch examples use this
+approach.
+
+In the functional API, call ``smc.run_functional`` with an explicit random
+state. It returns both the sampling result and the updated random state.
+The JAX example uses ``compile_functional_iteration`` to compile one SMC
+iteration and passes that iteration to ``run_functional``. The outer sampling
+loop remains under host control. Callbacks in the compiled iteration must be
+JAX-traceable; avoid converting their arrays to NumPy or Python scalars.
+
+Random number generation with orng
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Sequince uses ``orng`` to provide a common random-number interface across
+array libraries. Its ``RandomGenerator`` exposes NumPy-style methods such as
+``uniform``, ``normal``, and ``choice``, while drawing with the selected
+backend. It is installed as a dependency of Sequince.
+
+The CuPy and PyTorch examples construct this generator explicitly::
+
+   from orng import RandomGenerator
+
+   rng = RandomGenerator(backend="cupy", seed=42)
+   particles = rng.uniform(-6.0, 6.0, size=(1000, 2))
+
+Here ``particles`` is a CuPy array on the GPU. With ``backend="torch"``,
+draws are PyTorch tensors; pass ``device="cpu"`` or ``device="cuda"`` to
+choose their device. Pass the same generator to ``smc.run(..., rng=rng)``
+so initialization, resampling, and DE proposals advance one random stream.
+The generator is mutable: drawing samples changes its internal state.
+
+The NumPy example uses ``np.random.default_rng(seed)`` directly. Sequince
+wraps this existing generator in ``orng`` internally, retaining its state
+after the initial prior draws. You can also use
+``RandomGenerator(backend="numpy", seed=seed)`` explicitly.
+
+The JAX functional example supplies a JAX random key rather than a mutable
+generator. Sequince's steps use ``orng`` internally to obtain backend random
+draws and return the updated key through the functional API. Keep using the
+returned key, rather than reusing an earlier key, for subsequent draws.
+Equal integer seeds across libraries do not imply identical random streams.
+
+Writing portable array code
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use ``array_api_compat.array_namespace(particles)`` to obtain a compatible
+namespace, conventionally named ``xp``, from an input array. Call operations
+such as ``xp.all`` and ``xp.where`` through that namespace to keep results
+backend-native. Constructors such as ``xp.full_like`` preserve an existing
+array's dtype and device. This lets you share model code without depending
+on Sequince-specific array helpers.
+
+``array-api-extra`` (imported as ``array_api_extra``) complements the standard
+namespace with portable helpers for operations outside the Array API standard.
+For example, ``array_api_extra.atleast_nd(particles, ndim=2, xp=xp)`` prepends
+axes when needed to give an array at least two dimensions. Its ``at`` helper
+also provides indexed updates compatible with immutable JAX arrays. Use these
+helpers when the standard namespace lacks an operation you need, passing the
+namespace from ``array_namespace`` through ``xp``.
+
+For models that need data-dependent branches or loops, ``acai`` complements
+``array_api_compat`` with ``cond``, ``while_loop``, ``fori_loop``, and ``scan``.
+For example, a Python ``if`` on an array value cannot be traced by JAX;
+``acai.cond`` dispatches to ``jax.lax.cond`` for JAX and uses eager Python
+control flow for NumPy, CuPy, and PyTorch. Both JAX branches must be traceable
+and return compatible structures, shapes, and dtypes. PyTorch execution is
+eager here; this does not guarantee compatibility with ``torch.compile``.
+Sequince uses ``acai`` internally for its functional control flow.
+
+``array_api_compat`` supplies compatible array namespaces, ``array_api_extra``
+adds portable array helpers, ``orng`` supplies random draws, and ``acai``
+supplies control flow. These libraries let model code work across backends
+without depending on Sequince-specific array helpers.
+
+The Himmelblau examples
+-----------------------
+
+The following examples run the same sampling problem on each backend so you
+can compare how arrays, random generators, and execution styles fit together.
+Start with the NumPy script for the complete eager workflow, then compare
+CuPy and PyTorch for device-specific execution and JAX for the functional API.
+
+We use the two-dimensional Himmelblau function, an optimization benchmark
+with four separated minima. Taking its negative as a log likelihood creates
+a multimodal sampling problem: the sampler must explore several distinct
+regions of parameter space. This synthetic target illustrates both backend
+choices and the challenge of retaining multiple modes during sampling.
+
+The prior is uniform on ``[-6, 6]`` in each coordinate. The unnormalized
+likelihood is ``exp(-H(x, y))``, where
+
+.. math::
+
+   H(x, y) = (x^2 + y - 11)^2 + (x + y^2 - 7)^2.
+
+The posterior has four separated peaks, making it useful for exploring mode
+coverage. By default, all examples use adaptive tempering, affine
+preconditioning, and ten standard differential-evolution (DE) sweeps per
+stage, selected with
+``mutation="differential-evolution"``. Each sweep updates two complementary
+particle groups. Each script defines its prior and likelihood using its own
+backend's array operations, preserving the particles' device.
+The JAX script also offers BlackJAX NUTS mutation for comparison.
+
+Each script is self-contained, defining its target callbacks, sampler setup,
+and result plotting. The backend pages highlight the parts to compare; the
+scripts contain the complete runnable examples.
+
+Running the examples
+~~~~~~~~~~~~~~~~~~~~
+
+Each backend has its own walkthrough and environment instructions:
+
+* :doc:`NumPy </array_backends/numpy>`: the complete eager workflow on CPU.
+* :doc:`CuPy </array_backends/cupy>`: eager sampling on CUDA.
+* :doc:`PyTorch </array_backends/torch>`: eager sampling on CPU or CUDA.
+* :doc:`JAX </array_backends/jax>`: functional execution with DE or BlackJAX NUTS.
+
+Edit ``n_particles``, ``seed``, and ``outdir`` at the top of each script.
+Defaults are 1,000 particles, seed 42, and ``outdir/<backend>`` relative to
+``examples/array_backends``.
+For JAX, the default output directory is ``outdir/jax/<mutation>``
+(``differential-evolution`` or ``blackjax``) so the DE and BlackJAX plots can be compared without overwriting one another.
+Each prints estimated log
+evidence, weighted posterior means and covariance, then saves ``posterior.png``
+and ``history.png``. These are manual examples;
+they are not part of the automated documentation output manifest and do not
+require a GPU to build their tutorial pages. Equal seeds across backends do not
+produce identical random streams or particle populations.
+
+Interpreting the results
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+The four minima of the Himmelblau function lie near ``(3, 2)``,
+``(-2.805, 3.131)``, ``(-3.779, -3.283)``, and ``(3.584, -1.848)``.
+Inspect the corner plot for all four modes; means and covariance alone do not
+describe this posterior well. Equal objective values at the minima do not imply
+equal integrated posterior mass around each peak.
+
+The history plot shows accumulated log evidence and the tempering parameter
+``beta`` approaching one. The evidence refers to the stated unnormalized
+likelihood and normalized box prior; no analytic evidence is supplied.
+Small populations can lose modes during resampling, and ten DE sweeps do not
+guarantee mixing. Increase ``n_particles`` and compare multiple seeds when
+checking mode coverage.
+
+Writing backend-agnostic model code
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The backend-specific scripts can also share their target callbacks. For the
+Himmelblau problem, obtain the namespace from the particles in the prior and
+keep the likelihood's arithmetic unchanged:
+
+.. code-block:: python
+
+   from math import log
+
+   from array_api_compat import array_namespace
+
+   def log_prior(particles, context):
+       xp = array_namespace(particles)
+       inside = xp.all((particles >= -6.0) & (particles <= 6.0), axis=-1)
+       density = xp.full_like(
+           particles[..., 0], -particles.shape[-1] * log(12.0)
+       )
+       return xp.where(inside, density, xp.full_like(density, -float("inf")))
+
+   def log_likelihood(particles, context):
+       x, y = particles[..., 0], particles[..., 1]
+       return -((x**2 + y - 11.0) ** 2 + (x + y**2 - 7.0) ** 2)
+
+``array_api_compat`` provides compatible namespaces for NumPy, CuPy, and
+PyTorch, and recognizes JAX's namespace. Operations through ``xp`` return
+backend-native arrays. Using ``full_like`` also preserves the input dtype
+and device. The Himmelblau likelihood already needs only indexing and
+arithmetic, so its implementation can be shared without a namespace lookup.
+These callbacks can replace the corresponding definitions in any of the
+four scripts; initialization and the choice of execution API stay explicit.

examples/array_backends/cupy.rst

@@ -0,0 +1,42 @@
+CuPy backend
+============
+
+This example samples the Himmelblau target described in the
+:doc:`array backends overview </array_backends>`.
+
+.. note::
+
+   This example requires an NVIDIA GPU with a working CUDA driver and a
+   compatible CuPy installation. It cannot run on CPU; use the
+   :doc:`NumPy example </array_backends/numpy>` for CPU sampling.
+
+In ``run_cupy.py``, the prior uses ``cupy`` operations and ``orng`` draws
+the initial particles on CUDA:
+
+.. code-block:: python
+
+   rng = RandomGenerator(backend="cupy", seed=seed)
+   particles = rng.uniform(
+       lower_bound, upper_bound, size=(n_particles, n_dimensions)
+   )
+   initial_state = ParticleState(particles=particles)
+
+The sampler configuration and eager run call follow the
+:doc:`NumPy example </array_backends/numpy>`.
+Keep model calculations on the CuPy arrays; copying particles to the host
+inside a callback would introduce device transfers during sampling.
+
+Running the example
+-------------------
+
+Run from ``examples/array_backends`` in a Python environment containing
+Sequince, its plotting extras, and CuPy compatible with the installed CUDA
+runtime::
+
+   python run_cupy.py
+
+The default output directory is ``outdir/cupy``.
+
+Edit ``n_particles``, ``seed``, and ``outdir`` at the top of the script. See the
+:doc:`overview </array_backends>` for shared settings and guidance on
+interpreting the plots.

examples/array_backends/jax.rst

@@ -0,0 +1,119 @@
+JAX backend
+===========
+
+This example samples the Himmelblau target described in the
+:doc:`array backends overview </array_backends>`.
+
+In ``run_jax.py``, the callbacks use ``jax.numpy``. The script runs DE and
+BlackJAX NUTS in turn, starting each run from the same seed. Both use
+the same functional execution path. Split the initial key so prior
+initialization and sampling use separate random draws:
+
+.. code-block:: python
+
+   particle_key, random_state = jax.random.split(jax.random.key(seed))
+   initial_state = ParticleState(
+       particles=jax.random.uniform(
+           particle_key,
+           shape=(n_particles, n_dimensions),
+           minval=lower_bound,
+           maxval=upper_bound,
+       )
+   )
+
+The DE configuration matches the eager examples. For BlackJAX, configure
+NUTS through the sampler's mutation options:
+
+.. code-block:: python
+
+   mutation = "blackjax"
+   mutation_kwargs = {
+       "algorithm": "nuts",
+       "algorithm_kwargs": {
+           "step_size": 0.02,
+           "inverse_mass_matrix": jnp.ones(n_dimensions),
+       },
+   }
+
+Pass these as ``mutation=mutation`` and ``mutation_kwargs=mutation_kwargs``
+to ``make_standard_bayesian_smc``, keeping ``mutation_n_steps=10``. For NUTS,
+this means ten transitions per particle per tempering stage; each transition
+can involve several leapfrog steps. NUTS uses gradients of the log target,
+whereas DE uses differences between complementary particles. These runs do
+not have equal computational costs.
+The JAX script prints ``Total likelihood evaluations`` for either method,
+using ``result.evaluation_counts.log_likelihood_evaluations`` to compare
+their recorded likelihood work.
+
+BlackJAX evaluates and differentiates one particle at a time. The JAX
+callbacks therefore use ``particles[..., 0]`` and ``particles[..., 1]`` so
+they accept both an individual vector and a particle batch.
+
+The NUTS step size and inverse mass matrix are fixed in the affine-transformed
+coordinates; this example does not run BlackJAX warm-up or tune them. The box
+prior has a hard boundary, and NUTS moves are local, so inspect mode coverage
+across seeds rather than assuming it will move between separated peaks.
+
+For either choice, set ``logging_frequency=0`` to keep Python logging outside
+the compiled iteration. ``compile_functional_iteration`` uses JAX to compile
+the weighting, resampling, and mutation steps of one SMC iteration together.
+Passing the resulting callable to ``run_functional`` reuses that compiled
+computation at each tempering stage, reducing Python overhead within the
+iteration. The outer loop stays in Python, where Sequince checks when to
+stop and collects results.
+
+The helper initializes the functional state and context before compiling,
+and checks that the iteration preserves their structure, shapes, and dtypes.
+Compilation does not consume a sampling transition. Compile the iteration,
+then execute it through the functional API:
+
+.. code-block:: python
+
+   iteration = compile_functional_iteration(
+       smc, initial_state, random_state, metadata=metadata
+   )
+   result, final_random_state = smc.run_functional(
+       initial_state,
+       random_state=random_state,
+       metadata=metadata,
+       iteration=iteration,
+   )
+   jax.block_until_ready(result.final_state.particles)
+
+JAX can dispatch computations asynchronously: Python may receive an array
+before the device has finished computing its values.
+``jax.block_until_ready(result.final_state.particles)`` waits for the final
+particle array to be ready before the script reports results and plots them.
+It synchronizes with the device without copying the particles to NumPy.
+If you time sampling, put the end of the timing interval after this call so
+the measurement includes device execution rather than only dispatch. See
+JAX's `asynchronous dispatch guide
+<https://docs.jax.dev/en/latest/async_dispatch.html>`_ for more detail.
+
+Retain ``final_random_state`` for subsequent random operations. The script
+reports the device used by its final particles.
+
+Running the example
+-------------------
+
+Run from ``examples/array_backends`` with the repository's ``jax`` Pixi
+environment, which includes JAX on CPU and BlackJAX::
+
+   python run_jax.py
+
+For another Python environment, install Sequince's ``blackjax`` extra to run
+both methods. To run only DE, remove ``"blackjax"`` from the loop in ``main``.
+
+The script selects CPU with ``platform = "cpu"`` at the top of the file.
+For a GPU run, use an environment with GPU-enabled JAX and change that line
+to ``platform = "cuda"``. The script applies this setting through
+``jax.config.update("jax_platforms", platform)``, overriding the
+repository environment's ``JAX_PLATFORMS=cpu`` selection.
+
+Outputs go to ``outdir/jax/differential-evolution`` and
+``outdir/jax/blackjax``, so the plots can be compared without overwriting one
+another.
+
+Edit ``n_particles``, ``seed``, and ``outdir`` at the top of the script. See the
+:doc:`overview </array_backends>` for shared settings and guidance on
+interpreting the plots.