Persistent sampling

Persistent sampling#

This example demonstrates how to use persistent sampling in sequince to sample an n-dimensional Rosenbrock target.

Persistent sampling is a variation of SMC proposed by Karamanis and Seljak (2025) in Persistent sampling: enhancing the efficiency of sequential Monte Carlo. It retains and reweights populations from previous stages.

Run the example from the repository root:

python persistent_sampling/rosenbrock.py

For a larger run:

python persistent_sampling/rosenbrock.py \
    --dimensions 16 \
    --n-particles 1024 \
    --target-ess 2048 \
    --target-total-ess 8192 \
    --mutation-steps 32

The script prints diagnostics from the final weighted persistent state, including its archive size, effective sample size, weighted mean, best Rosenbrock potential, and log evidence. Unless --no-plot is supplied, it also saves a weighted corner plot as examples/persistent_sampling/rosenbrock_<D>d.png. Use --output PATH to choose a different location.

Run the following for all available options:

python persistent_sampling/rosenbrock.py --help

Preconditioning#

First read Flow preconditioning for an explicit flow-transform setup. This example also shows how to use the normalizing-flow-based preconditioning described in Karamanis et al. (2022), Accelerating astronomical and cosmological inference with Preconditioned Monte Carlo, combined with persistent sampling from Karamanis and Seljak (2025).

The --flow option enables learned normalizing-flow preconditioning:

python persistent_sampling/rosenbrock.py --dimensions 6 --flow

The flow option requires sequince-density and its flow dependencies.

Example plots#

Posterior marginals and joint distributions from the persistent Rosenbrock run, using the script defaults.

Posterior marginals and joint distributions from the persistent Rosenbrock run, using the script defaults.#

Run provenance.

Complete script#

Download the script.

  1"""Sample a configurable-dimensional Rosenbrock target persistently.
  2
  3For example, run a six-dimensional problem with::
  4
  5    python persistent_sampling/rosenbrock.py --dimensions 6
  6
  7The target is the bounded density ``exp(-Rosenbrock(x) / scale)`` under a
  8uniform prior. A corner plot of the final weighted persistent population is
  9saved unless ``--no-plot`` is supplied. Pass ``--flow`` to use optional learned
 10normalizing-flow preconditioning instead of the default bounded transform.
 11"""
 12
 13from __future__ import annotations
 14
 15import argparse
 16import json
 17from pathlib import Path
 18
 19import numpy as np
 20from sequince import (
 21    ParticleState,
 22    PersistentParticleState,
 23    make_persistent_bayesian_smc,
 24)
 25from sequince.log import configure_logging
 26from sequince.plot import plot_corner
 27from sequince.steps import StepContext
 28from sequince.transforms import BoundedToUnboundedTransform
 29
 30documentation_output: Path | None = None
 31
 32
 33def rosenbrock(particles: np.ndarray) -> np.ndarray:
 34    """Return the Rosenbrock potential for each particle."""
 35    values = np.asarray(particles, dtype=float)
 36    return np.sum(
 37        100.0 * (values[:, 1:] - values[:, :-1] ** 2.0) ** 2.0
 38        + (1.0 - values[:, :-1]) ** 2.0,
 39        axis=1,
 40    )
 41
 42
 43def make_likelihood_and_prior(
 44    dimensions: int,
 45    *,
 46    bound: float,
 47    scale: float,
 48) -> tuple[
 49    callable[[np.ndarray, StepContext], np.ndarray],
 50    callable[[np.ndarray, StepContext], np.ndarray],
 51]:
 52    """Construct a bounded Rosenbrock target."""
 53    log_prior_volume = dimensions * np.log(2.0 * bound)
 54
 55    def log_prior(particles: np.ndarray, context=None) -> np.ndarray:
 56        del context
 57        values = np.asarray(particles, dtype=float)
 58        inside = np.all(np.abs(values) <= bound, axis=1)
 59        return np.where(inside, -log_prior_volume, -np.inf)
 60
 61    def log_likelihood(particles: np.ndarray, context=None) -> np.ndarray:
 62        del context
 63        return -rosenbrock(particles) / scale
 64
 65    return (
 66        log_prior,
 67        log_likelihood,
 68    )
 69
 70
 71def build_parser() -> argparse.ArgumentParser:
 72    """Build the command-line parser."""
 73    parser = argparse.ArgumentParser(description=__doc__)
 74    parser.add_argument("--dimensions", type=int, default=4)
 75    parser.add_argument("--n-particles", type=int, default=256)
 76    parser.add_argument(
 77        "--target-ess",
 78        type=float,
 79        default=None,
 80        help="Absolute persistent ESS (default: twice n-particles).",
 81    )
 82    parser.add_argument(
 83        "--target-total-ess",
 84        type=float,
 85        default=None,
 86        help="Optional final persistent ESS accumulated at beta=1.",
 87    )
 88    parser.add_argument("--mutation-steps", type=int, default=8)
 89    parser.add_argument(
 90        "--flow",
 91        action="store_true",
 92        help="Use optional learned normalizing-flow preconditioning.",
 93    )
 94    parser.add_argument("--bound", type=float, default=5.0)
 95    parser.add_argument(
 96        "--rosenbrock-scale",
 97        type=float,
 98        default=1.0,
 99        help="Divide the Rosenbrock potential by this positive value.",
100    )
101    parser.add_argument("--seed", type=int, default=20260820)
102    parser.add_argument(
103        "--output",
104        type=Path,
105        default=None,
106        help="Plot path (default: persistent_sampling/rosenbrock_<D>d.png).",
107    )
108    parser.add_argument("--no-plot", action="store_true")
109    return parser
110
111
112def main() -> None:
113    """Run persistent SMC and report archive diagnostics."""
114    parser = build_parser()
115    args = parser.parse_args()
116    if args.dimensions < 2:
117        parser.error("--dimensions must be at least 2")
118    if args.n_particles < 1:
119        parser.error("--n-particles must be positive")
120    if args.bound <= 0.0:
121        parser.error("--bound must be positive")
122    if args.rosenbrock_scale <= 0.0:
123        parser.error("--rosenbrock-scale must be positive")
124    if args.mutation_steps < 1:
125        parser.error("--mutation-steps must be positive")
126    target_ess = (
127        float(args.target_ess)
128        if args.target_ess is not None
129        else float(2 * args.n_particles)
130    )
131
132    configure_logging()
133
134    rng = np.random.default_rng(args.seed)
135    log_prior, log_likelihood = make_likelihood_and_prior(
136        args.dimensions,
137        bound=args.bound,
138        scale=args.rosenbrock_scale,
139    )
140    initial_particles = rng.uniform(
141        -args.bound,
142        args.bound,
143        size=(args.n_particles, args.dimensions),
144    )
145    initial_state = ParticleState(initial_particles)
146
147    bounds = {
148        index: (-args.bound, args.bound) for index in range(args.dimensions)
149    }
150    bounded_transform = BoundedToUnboundedTransform(bounds=bounds)
151    if args.flow:
152        try:
153            from sequince.transforms import FlowTransform
154            from sequince_density.flows import ZukoFlow
155        except ImportError as exc:
156            parser.error(
157                "--flow requires the optional sequince-density dependency "
158                f"({exc})"
159            )
160        flow = ZukoFlow(
161            dims=args.dimensions,
162            flow_class="MAF",
163            hidden_features=[64, 64],
164            particle_transform=bounded_transform,
165        )
166        particle_transform = FlowTransform(
167            flow,
168            fit_kwargs={
169                "n_epochs": 100,
170                "lr": 1e-3,
171                "batch_size": 500,
172            },
173        )
174        transform_refit = "per-context"
175    else:
176        particle_transform = bounded_transform
177        transform_refit = "once"
178
179    sampler, posterior, metadata = make_persistent_bayesian_smc(
180        log_prior=log_prior,
181        log_likelihood=log_likelihood,
182        n_particles=args.n_particles,
183        target_ess=target_ess,
184        target_total_ess=args.target_total_ess,
185        mutation_n_steps=args.mutation_steps,
186        mutation_kwargs={"transform_refit": transform_refit},
187        particle_transform=particle_transform,
188    )
189    result = sampler.run(
190        initial_state,
191        rng=rng,
192        metadata=metadata,
193    )
194
195    final_state = result.final_state
196    if not isinstance(final_state, PersistentParticleState):
197        raise TypeError("persistent SMC must return a PersistentParticleState")
198    particles = np.asarray(final_state.particles)
199    weights = np.asarray(final_state.weights)
200    potentials = rosenbrock(particles)
201    best = particles[int(np.argmin(potentials))]
202    weighted_mean = np.average(particles, axis=0, weights=weights)
203    weighted_mean_potential = np.average(potentials, weights=weights)
204    print()
205    print(f"dimensions: {args.dimensions}")
206    print(f"iterations: {len(result)}")
207    print(f"final beta: {float(result.context.metadata['beta']):.6f}")
208    print(f"final archive size: {final_state.n_particles}")
209    print(f"final persistent ESS: {final_state.effective_sample_size():.1f}")
210    print(
211        "weighted mean: "
212        f"{np.array2string(weighted_mean, precision=3, threshold=12)}"
213    )
214    print(f"weighted mean Rosenbrock potential: {weighted_mean_potential:.6g}")
215    print(f"best Rosenbrock potential: {float(np.min(potentials)):.6g}")
216    print(f"best particle: {np.array2string(best, precision=3)}")
217    print(f"final log evidence: {float(final_state.log_evidence):+.6f}")
218
219    if documentation_output is not None:
220        summary = {
221            "log_evidence": float(final_state.log_evidence),
222            "weighted_mean": weighted_mean.tolist(),
223            "archive_size": int(final_state.n_particles),
224            "effective_sample_size": float(
225                final_state.effective_sample_size()
226            ),
227            "weighted_mean_potential": float(weighted_mean_potential),
228        }
229        (documentation_output / "summary.json").write_text(
230            json.dumps(summary, allow_nan=False) + "\n"
231        )
232        if args.output is None:
233            args.output = documentation_output / "corner.png"
234
235    if not args.no_plot:
236        output = args.output or Path(__file__).with_name(
237            f"rosenbrock_{args.dimensions}d.png"
238        )
239        output.parent.mkdir(parents=True, exist_ok=True)
240        posterior_samples, weights = result.posterior_samples(
241            method="rejection", rng=rng
242        )
243        plot_corner(
244            posterior_samples,
245            weights=weights,
246            parameter_labels=[
247                rf"$x_{{{index + 1}}}$" for index in range(args.dimensions)
248            ],
249            filename=output,
250            truths=np.ones(args.dimensions),
251        )
252        print(f"saved plot: {output}")
253
254
255if __name__ == "__main__":
256    main()
Contents