diff --git a/docs/array_backends.rst b/docs/array_backends.rst new file mode 100644 index 0000000..26d507c --- /dev/null +++ b/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 diff --git a/docs/array_backends/cupy.rst b/docs/array_backends/cupy.rst new file mode 100644 index 0000000..67be6af --- /dev/null +++ b/docs/array_backends/cupy.rst @@ -0,0 +1 @@ +.. include:: ../../examples/array_backends/cupy.rst diff --git a/docs/array_backends/jax.rst b/docs/array_backends/jax.rst new file mode 100644 index 0000000..ccfe90b --- /dev/null +++ b/docs/array_backends/jax.rst @@ -0,0 +1 @@ +.. include:: ../../examples/array_backends/jax.rst diff --git a/docs/array_backends/numpy.rst b/docs/array_backends/numpy.rst new file mode 100644 index 0000000..b3515e2 --- /dev/null +++ b/docs/array_backends/numpy.rst @@ -0,0 +1 @@ +.. include:: ../../examples/array_backends/numpy.rst diff --git a/docs/array_backends/torch.rst b/docs/array_backends/torch.rst new file mode 100644 index 0000000..30d9dbb --- /dev/null +++ b/docs/array_backends/torch.rst @@ -0,0 +1 @@ +.. include:: ../../examples/array_backends/torch.rst diff --git a/docs/index.rst b/docs/index.rst index f80dd6a..54bc729 100644 --- a/docs/index.rst +++ b/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 diff --git a/examples/array_backends/README.rst b/examples/array_backends/README.rst new file mode 100644 index 0000000..337622a --- /dev/null +++ b/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 `: the complete eager workflow on CPU. +* :doc:`CuPy `: eager sampling on CUDA. +* :doc:`PyTorch `: eager sampling on CPU or CUDA. +* :doc:`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/`` relative to +``examples/array_backends``. +For JAX, the default output directory is ``outdir/jax/`` +(``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. diff --git a/examples/array_backends/cupy.rst b/examples/array_backends/cupy.rst new file mode 100644 index 0000000..2dfecc5 --- /dev/null +++ b/examples/array_backends/cupy.rst @@ -0,0 +1,42 @@ +CuPy backend +============ + +This example samples the Himmelblau target described in the +:doc:`array backends overview `. + +.. 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 ` 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 `. +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 ` for shared settings and guidance on +interpreting the plots. diff --git a/examples/array_backends/jax.rst b/examples/array_backends/jax.rst new file mode 100644 index 0000000..38cdec3 --- /dev/null +++ b/examples/array_backends/jax.rst @@ -0,0 +1,119 @@ +JAX backend +=========== + +This example samples the Himmelblau target described in the +:doc:`array backends overview `. + +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 +`_ 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 ` for shared settings and guidance on +interpreting the plots. diff --git a/examples/array_backends/numpy.rst b/examples/array_backends/numpy.rst new file mode 100644 index 0000000..96143a3 --- /dev/null +++ b/examples/array_backends/numpy.rst @@ -0,0 +1,61 @@ +NumPy backend +============= + +This example samples the Himmelblau target described in the +:doc:`array backends overview `. + +Start with ``run_numpy.py``. Its callbacks use NumPy operations on batches +of particles. The likelihood returns one value per particle: + +.. code-block:: python + + def log_likelihood(particles, context): + x, y = particles[:, 0], particles[:, 1] + return -((x**2 + y - 11.0) ** 2 + (x + y**2 - 7.0) ** 2) + +Initialize the population by drawing from the box prior. ``ParticleState`` +supplies equal initial weights, appropriate for these prior draws: + +.. code-block:: python + + rng = np.random.default_rng(seed) + particles = rng.uniform( + lower_bound, upper_bound, size=(n_particles, n_dimensions) + ) + initial_state = ParticleState(particles=particles) + +Configure tempering and DE, then run the sampler with the same generator: + +.. code-block:: python + + smc, _, metadata = make_standard_bayesian_smc( + log_prior=log_prior, + log_likelihood=log_likelihood, + target_ess_ratio=0.7, + always_resample=True, + mutation="differential-evolution", + mutation_n_steps=10, + particle_transform="affine", + ) + result = smc.run(initial_state, rng=rng, metadata=metadata) + +The prior callback includes the box normalization and returns negative infinity +outside the bounds. To adapt this example, replace the callbacks, dimensions, +bounds, and initial draws together. Each callback accepts a particle batch and +returns one log density per particle. Keep the prior normalized when estimating +evidence. + +Running the example +------------------- + +Run from ``examples/array_backends`` with the repository's ``default`` +Pixi environment, or another environment containing Sequince and its plotting +extras:: + + python run_numpy.py + +The default output directory is ``outdir/numpy``. + +Edit ``n_particles``, ``seed``, and ``outdir`` at the top of the script. See the +:doc:`overview ` for shared settings and guidance on +interpreting the plots. diff --git a/examples/array_backends/run_cupy.py b/examples/array_backends/run_cupy.py new file mode 100644 index 0000000..741e359 --- /dev/null +++ b/examples/array_backends/run_cupy.py @@ -0,0 +1,107 @@ +"""Sample a Himmelblau target with CuPy. See README.rst.""" + +from math import log +from pathlib import Path + +import cupy as cp +import matplotlib.pyplot as plt +from orng import RandomGenerator +from sequince import ParticleState, make_standard_bayesian_smc +from sequince.array_api import to_numpy +from sequince.log import configure_logging +from sequince.steps import StepContext + +n_particles = 1000 +seed = 42 +outdir = Path("outdir/cupy") + +n_dimensions = 2 +lower_bound, upper_bound = -6.0, 6.0 + + +def log_prior(particles: cp.ndarray, context: StepContext) -> cp.ndarray: + """Evaluate a normalized uniform prior. + + Parameters + ---------- + particles : array, shape (n_particles, n_dimensions) + CuPy particle coordinates. + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_density : array, shape (n_particles,) + Normalized box log density, or negative infinity outside the box. + """ + inside = cp.all( + (particles >= lower_bound) & (particles <= upper_bound), axis=-1 + ) + density = cp.full_like( + particles[:, 0], -n_dimensions * log(upper_bound - lower_bound) + ) + return cp.where(inside, density, cp.full_like(density, -float("inf"))) + + +def log_likelihood(particles: cp.ndarray, context: StepContext) -> cp.ndarray: + """Evaluate the negative Himmelblau function. + + Parameters + ---------- + particles : array, shape (n_particles, 2) + CuPy coordinates (x, y). + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_likelihood : array, shape (n_particles,) + Negative Himmelblau objective, defining an unnormalized likelihood. + """ + x, y = particles[:, 0], particles[:, 1] + return -((x**2 + y - 11.0) ** 2 + (x + y**2 - 7.0) ** 2) + + +def main() -> None: + """Initialize CuPy particles and run DE mutations through tempering.""" + configure_logging() + # orng supplies NumPy-style RNG methods that return CuPy arrays on CUDA. + # Pass this same mutable generator to SMC to continue the random stream. + rng = RandomGenerator(backend="cupy", seed=seed) + particles = rng.uniform( + lower_bound, upper_bound, size=(n_particles, n_dimensions) + ) + # Prior draws start with equal weights, supplied by ParticleState. + initial_state = ParticleState(particles=particles) + smc, _, metadata = make_standard_bayesian_smc( + log_prior=log_prior, + log_likelihood=log_likelihood, + target_ess_ratio=0.7, + always_resample=True, + mutation="differential-evolution", + mutation_n_steps=10, + particle_transform="affine", + ) + result = smc.run(initial_state, rng=rng, metadata=metadata) + + outdir.mkdir(parents=True, exist_ok=True) + print(f"Estimated log evidence: {result.final_state.log_evidence:.4f}") + print(f"Posterior mean: {to_numpy(result.weighted_mean())}") + print(f"Posterior covariance:\n{to_numpy(result.weighted_covariance())}") + fig = result.plot_corner() + fig.savefig(outdir / "posterior.png") + plt.close(fig) + fig, axes = plt.subplots(2, 1, figsize=(6, 6)) + result.plot_log_evidence(ax=axes[0]) + # Read the tempering parameter from each weighting step output. + betas = [float(to_numpy(output[0]["beta"])) for output in result.outputs] + axes[1].plot(betas) + axes[1].set_xlabel("Iteration") + axes[1].set_ylabel("Tempering parameter beta") + fig.tight_layout() + fig.savefig(outdir / "history.png") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/examples/array_backends/run_jax.py b/examples/array_backends/run_jax.py new file mode 100644 index 0000000..9c9484a --- /dev/null +++ b/examples/array_backends/run_jax.py @@ -0,0 +1,152 @@ +"""Run functional JAX SMC with DE and BlackJAX NUTS. See README.rst.""" + +from math import log +from pathlib import Path + +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +from sequince import ParticleState, make_standard_bayesian_smc +from sequince.array_api import to_numpy +from sequince.jax import compile_functional_iteration +from sequince.log import configure_logging +from sequince.steps import StepContext + +n_particles = 1000 +seed = 42 +outdir = Path("outdir/jax") + +# Change to "cuda" with GPU-enabled JAX to run on the GPU. +platform = "cpu" + +n_dimensions = 2 +lower_bound, upper_bound = -6.0, 6.0 + + +def log_prior(particles: jax.Array, context: StepContext) -> jax.Array: + """Evaluate a normalized uniform prior. + + Parameters + ---------- + particles : array, shape (..., n_dimensions) + A particle vector or a batch of particle coordinates. + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_density : array, shape (...,) + Normalized box log density, or negative infinity outside the box. + """ + inside = jnp.all( + (particles >= lower_bound) & (particles <= upper_bound), axis=-1 + ) + density = jnp.full_like( + particles[..., 0], -n_dimensions * log(upper_bound - lower_bound) + ) + return jnp.where(inside, density, jnp.full_like(density, -float("inf"))) + + +def log_likelihood(particles: jax.Array, context: StepContext) -> jax.Array: + """Evaluate the negative Himmelblau function. + + Parameters + ---------- + particles : array, shape (..., 2) + A coordinate vector (x, y) or a batch of such vectors. + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_likelihood : array, shape (...,) + Negative Himmelblau objective, defining an unnormalized likelihood. + """ + x, y = particles[..., 0], particles[..., 1] + return -((x**2 + y - 11.0) ** 2 + (x + y**2 - 7.0) ** 2) + + +def main() -> None: + """Compile an SMC iteration and thread an explicit JAX random key.""" + jax.config.update("jax_platforms", platform) + configure_logging() + for mutation in ("differential-evolution", "blackjax"): + output_dir = outdir / mutation + # Functional steps use orng internally, threading a JAX key explicitly. + # Split the key so initialization and SMC use separate random draws. + 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, + ) + ) + mutation_kwargs = {} + if mutation == "blackjax": + # Fixed NUTS settings for this example; no automatic warm-up is run. + mutation_kwargs = { + "algorithm": "nuts", + "algorithm_kwargs": { + "step_size": 0.02, + "inverse_mass_matrix": jnp.ones(n_dimensions), + }, + } + smc, _, metadata = make_standard_bayesian_smc( + log_prior=log_prior, + log_likelihood=log_likelihood, + target_ess_ratio=0.7, + always_resample=True, + mutation=mutation, + mutation_n_steps=10, + mutation_kwargs=mutation_kwargs, + particle_transform="affine", + # Keep Python logging outside the compiled iteration. + logging_frequency=0, + ) + # Compile one iteration without consuming a sampling transition. + 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) + # Retain the returned key for subsequent random operations. + print(f"Final random key: {jax.random.key_data(final_random_state)}") + print(f"JAX device: {result.final_state.particles.device}") + print(f"Mutation: {mutation}") + likelihood_evaluations = int( + result.evaluation_counts.log_likelihood_evaluations + ) + print(f"Total likelihood evaluations: {likelihood_evaluations:,}") + + output_dir.mkdir(parents=True, exist_ok=True) + print(f"Estimated log evidence: {result.final_state.log_evidence:.4f}") + print(f"Posterior mean: {to_numpy(result.weighted_mean())}") + print( + f"Posterior covariance:\n{to_numpy(result.weighted_covariance())}" + ) + fig = result.plot_corner() + fig.savefig(output_dir / "posterior.png") + plt.close(fig) + fig, axes = plt.subplots(2, 1, figsize=(6, 6)) + result.plot_log_evidence(ax=axes[0]) + # Functional runs retain beta in step outputs rather than state metadata. + betas = [ + float(to_numpy(output[0]["beta"])) for output in result.outputs + ] + axes[1].plot(betas) + axes[1].set_xlabel("Iteration") + axes[1].set_ylabel("Tempering parameter beta") + fig.tight_layout() + fig.savefig(output_dir / "history.png") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/examples/array_backends/run_numpy.py b/examples/array_backends/run_numpy.py new file mode 100644 index 0000000..b543393 --- /dev/null +++ b/examples/array_backends/run_numpy.py @@ -0,0 +1,105 @@ +"""Sample a Himmelblau target with NumPy. See README.rst.""" + +from math import log +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +from sequince import ParticleState, make_standard_bayesian_smc +from sequince.array_api import to_numpy +from sequince.log import configure_logging +from sequince.steps import StepContext + +n_particles = 1000 +seed = 42 +outdir = Path("outdir/numpy") + +n_dimensions = 2 +lower_bound, upper_bound = -6.0, 6.0 + + +def log_prior(particles: np.ndarray, context: StepContext) -> np.ndarray: + """Evaluate a normalized uniform prior. + + Parameters + ---------- + particles : array, shape (n_particles, n_dimensions) + NumPy particle coordinates. + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_density : array, shape (n_particles,) + Normalized box log density, or negative infinity outside the box. + """ + inside = np.all( + (particles >= lower_bound) & (particles <= upper_bound), axis=-1 + ) + density = np.full_like( + particles[:, 0], -n_dimensions * log(upper_bound - lower_bound) + ) + return np.where(inside, density, np.full_like(density, -float("inf"))) + + +def log_likelihood(particles: np.ndarray, context: StepContext) -> np.ndarray: + """Evaluate the negative Himmelblau function. + + Parameters + ---------- + particles : array, shape (n_particles, 2) + NumPy coordinates (x, y). + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_likelihood : array, shape (n_particles,) + Negative Himmelblau objective, defining an unnormalized likelihood. + """ + x, y = particles[:, 0], particles[:, 1] + return -((x**2 + y - 11.0) ** 2 + (x + y**2 - 7.0) ** 2) + + +def main() -> None: + """Initialize NumPy particles and run DE mutations through tempering.""" + configure_logging() + # Sequince wraps this NumPy generator in orng, preserving its RNG state. + rng = np.random.default_rng(seed) + particles = rng.uniform( + lower_bound, upper_bound, size=(n_particles, n_dimensions) + ) + # Prior draws start with equal weights, supplied by ParticleState. + initial_state = ParticleState(particles=particles) + smc, _, metadata = make_standard_bayesian_smc( + log_prior=log_prior, + log_likelihood=log_likelihood, + target_ess_ratio=0.7, + always_resample=True, + mutation="differential-evolution", + mutation_n_steps=10, + particle_transform="affine", + ) + result = smc.run(initial_state, rng=rng, metadata=metadata) + + outdir.mkdir(parents=True, exist_ok=True) + print(f"Estimated log evidence: {result.final_state.log_evidence:.4f}") + print(f"Posterior mean: {to_numpy(result.weighted_mean())}") + print(f"Posterior covariance:\n{to_numpy(result.weighted_covariance())}") + fig = result.plot_corner() + fig.savefig(outdir / "posterior.png") + plt.close(fig) + fig, axes = plt.subplots(2, 1, figsize=(6, 6)) + result.plot_log_evidence(ax=axes[0]) + # Read the tempering parameter from each weighting step output. + betas = [float(to_numpy(output[0]["beta"])) for output in result.outputs] + axes[1].plot(betas) + axes[1].set_xlabel("Iteration") + axes[1].set_ylabel("Tempering parameter beta") + fig.tight_layout() + fig.savefig(outdir / "history.png") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/examples/array_backends/run_torch.py b/examples/array_backends/run_torch.py new file mode 100644 index 0000000..82f28cb --- /dev/null +++ b/examples/array_backends/run_torch.py @@ -0,0 +1,112 @@ +"""Sample a Himmelblau target with PyTorch. See README.rst.""" + +from math import log +from pathlib import Path + +import matplotlib.pyplot as plt +import torch +from orng import RandomGenerator +from sequince import ParticleState, make_standard_bayesian_smc +from sequince.array_api import to_numpy +from sequince.log import configure_logging +from sequince.steps import StepContext + +n_particles = 1000 +seed = 42 +outdir = Path("outdir/torch") +device = "cpu" # Change to "cuda" for a CUDA-enabled PyTorch installation. + +n_dimensions = 2 +lower_bound, upper_bound = -6.0, 6.0 + + +def log_prior(particles: torch.Tensor, context: StepContext) -> torch.Tensor: + """Evaluate a normalized uniform prior. + + Parameters + ---------- + particles : array, shape (n_particles, n_dimensions) + PyTorch particle coordinates. + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_density : array, shape (n_particles,) + Normalized box log density, or negative infinity outside the box. + """ + inside = torch.all( + (particles >= lower_bound) & (particles <= upper_bound), axis=-1 + ) + density = torch.full_like( + particles[:, 0], -n_dimensions * log(upper_bound - lower_bound) + ) + return torch.where( + inside, density, torch.full_like(density, -float("inf")) + ) + + +def log_likelihood( + particles: torch.Tensor, context: StepContext +) -> torch.Tensor: + """Evaluate the negative Himmelblau function. + + Parameters + ---------- + particles : array, shape (n_particles, 2) + PyTorch coordinates (x, y). + context : StepContext + Runtime context; unused for this fixed target. + + Returns + ------- + log_likelihood : array, shape (n_particles,) + Negative Himmelblau objective, defining an unnormalized likelihood. + """ + x, y = particles[:, 0], particles[:, 1] + return -((x**2 + y - 11.0) ** 2 + (x + y**2 - 7.0) ** 2) + + +def main() -> None: + """Initialize PyTorch particles and run DE mutations through tempering.""" + configure_logging() + # orng returns PyTorch tensors on the selected device for every draw. + # Pass this same mutable generator to SMC to continue the random stream. + rng = RandomGenerator(backend="torch", seed=seed, device=device) + particles = rng.uniform( + lower_bound, upper_bound, size=(n_particles, n_dimensions) + ) + # Prior draws start with equal weights, supplied by ParticleState. + initial_state = ParticleState(particles=particles) + smc, _, metadata = make_standard_bayesian_smc( + log_prior=log_prior, + log_likelihood=log_likelihood, + target_ess_ratio=0.7, + always_resample=True, + mutation="differential-evolution", + mutation_n_steps=10, + particle_transform="affine", + ) + result = smc.run(initial_state, rng=rng, metadata=metadata) + + outdir.mkdir(parents=True, exist_ok=True) + print(f"Estimated log evidence: {result.final_state.log_evidence:.4f}") + print(f"Posterior mean: {to_numpy(result.weighted_mean())}") + print(f"Posterior covariance:\n{to_numpy(result.weighted_covariance())}") + fig = result.plot_corner() + fig.savefig(outdir / "posterior.png") + plt.close(fig) + fig, axes = plt.subplots(2, 1, figsize=(6, 6)) + result.plot_log_evidence(ax=axes[0]) + # Read the tempering parameter from each weighting step output. + betas = [float(to_numpy(output[0]["beta"])) for output in result.outputs] + axes[1].plot(betas) + axes[1].set_xlabel("Iteration") + axes[1].set_ylabel("Tempering parameter beta") + fig.tight_layout() + fig.savefig(outdir / "history.png") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/examples/array_backends/torch.rst b/examples/array_backends/torch.rst new file mode 100644 index 0000000..3360ba3 --- /dev/null +++ b/examples/array_backends/torch.rst @@ -0,0 +1,40 @@ +PyTorch backend +=============== + +This example samples the Himmelblau target described in the +:doc:`array backends overview `. + +In ``run_torch.py``, the callbacks use ``torch`` operations. The random +generator selects the device for both initialization and later random draws: + +.. code-block:: python + + rng = RandomGenerator( + backend="torch", seed=seed, device=device + ) + particles = rng.uniform( + lower_bound, upper_bound, size=(n_particles, n_dimensions) + ) + initial_state = ParticleState(particles=particles) + +Set ``device = "cpu"`` or, with a CUDA-enabled installation, +``device = "cuda"`` at the top of the script. +Array constructors such as ``torch.full_like`` in the prior callback preserve +the particles' device. This example uses eager execution. + +Running the example +------------------- + +Run from ``examples/array_backends`` with the repository's ``flow`` +Pixi environment, which provides PyTorch on CPU:: + + python run_torch.py + +With a CUDA-enabled PyTorch installation, change the ``device`` setting +to ``"cuda"`` before running the same command. + +The default output directory is ``outdir/torch``. + +Edit ``n_particles``, ``seed``, and ``outdir`` at the top of the script. See the +:doc:`overview ` for shared settings and guidance on +interpreting the plots. diff --git a/examples/gpu_sampling/cupy_examples/run_minipcn_cupy.py b/examples/gpu_sampling/cupy_examples/run_minipcn_cupy.py deleted file mode 100644 index 881f830..0000000 --- a/examples/gpu_sampling/cupy_examples/run_minipcn_cupy.py +++ /dev/null @@ -1,87 +0,0 @@ -import cupy as cp -import matplotlib.pyplot as plt -from orng import ArrayRNG -from sequince import ParticleState -from sequince.logging_utils import configure_logging -from sequince.prebuilt_samplers import make_standard_bayesian_smc - -# Number of dimensions in the multivariate normal model -dims = 4 -# Mean and standard deviation of the target normal distribution -mu = cp.array(0.0) -sigma = cp.array(1.0) -# Prior bounds -bounds = (cp.array(-10.0), cp.array(10.0)) -# True log evidence for the model with uniform prior over the bounds -true_log_evidence = -dims * cp.log(cp.array(20.0)) - - -def log_likelihood(particles: cp.ndarray, context) -> cp.ndarray: - """Log-likelihood for a multivariate standard normal observation.""" - return -0.5 * cp.sum(cp.square(particles - mu) / sigma**2, axis=-1) - ( - dims / 2 - ) * cp.log(2 * cp.pi * sigma**2) - - -def log_prior(particles: cp.ndarray, context) -> cp.ndarray: - """Uniform prior over a bounded region.""" - lower_bound = context.shared["lower_bound"] - upper_bound = context.shared["upper_bound"] - valid = (particles >= lower_bound) & (particles <= upper_bound) - # sum over dimensions; assign 0 if within bounds else -inf - return cp.where(valid, 0, -cp.inf).sum(axis=-1) - - -def main() -> None: - # Create a CuPy RNG - rng = ArrayRNG(backend="cupy", seed=42) - - # Enable logging - configure_logging() - - # Initialize particles from the prior - lower_bound, upper_bound = bounds - n_particles = 1000 - initial_particles = rng.uniform( - low=lower_bound, high=upper_bound, size=(n_particles, dims) - ) - initial_log_weights = cp.full( - (n_particles,), -cp.log(cp.array(n_particles)) - ) - - initial_state = ParticleState( - particles=initial_particles, log_weights=initial_log_weights - ) - - # Use the prebuilt SMC with MiniPCN mutation - # MiniPCN is the only sampler that supports the array API backends - smc, posterior, shared_context = make_standard_bayesian_smc( - log_prior=log_prior, - log_likelihood=log_likelihood, - target_ess_ratio=0.7, - max_temperature=1.0, - min_delta=1e-8, - always_resample=True, - mutation="minipcn", - mutation_n_steps=10, - particle_transform="affine", - auto_normalize=True, - ) - - # Add prior bounds to the shared context - shared_context["lower_bound"] = lower_bound - shared_context["upper_bound"] = upper_bound - - result = smc.run(initial_state, rng=rng, shared_context=shared_context) - - fig, axs = plt.subplots(2, 1, figsize=(4, 12)) - result.plot_log_evidence(ax=axs[0]) - result.plot_metadata_history("temperature", ax=axs[1]) - fig.savefig("state.png") - - fig = result.plot_corner() - fig.savefig("posterior.png") - - -if __name__ == "__main__": - main()