API Reference

This page documents simplesvgd's public API -- everything listed in simplesvgd.__all__.

Core

simplesvgd.update

Core SVGD update function with optional preconditioning and hierarchical sigma.

update

update(
    x0: NDArray[FloatDType],
    gradient_fn: GradientFn[FloatDType]
    | MinibatchGradientFn[FloatDType],
    config: SVGDConfig[FloatDType] | None = None,
) -> SVGDState[FloatDType]

Update a collection of samples using Stein Variational Gradient Descent.

Parameters:

Name	Type	Description	Default
x0	ndarray	Initial particle positions, shape (n_particles, n_dims). Its dtype (e.g. float32 or float64) is preserved throughout the run -- pass float32 particles to run the whole optimization in single precision. gradient_fn should return gradients in the same dtype; if it doesn't, they are cast to x0's dtype before use, so a careless gradient_fn can't silently upcast the run. Ignored when config.resume_from is set -- the run continues from resume_from.particles instead.	required
gradient_fn	callable	Computes gradients of the negative log-probability. Accepts particles of shape (n_particles, n_dims) and returns gradients of the same shape. When config.sigma.value is set or config.sigma.estimate is True, must return (gradients, misfits) where misfits has shape (n_particles,). When config.minibatch_sampler is set, called instead as gradient_fn(particles, batch_indices) -- see :class:SVGDConfig.	required
config	SVGDConfig or None	Every other tunable of the run (iteration count, step size, preconditioner, kernel, bounds, callback, resume, animation, ...). None uses SVGDConfig()'s defaults. See :class:SVGDConfig for the full list of fields.	None

Returns:

Type	Description
SVGDState	Final optimizer state. Access .particles for the particle array.

Source code in src/simplesvgd/update.py

def update(  # noqa: C901, PLR0912, PLR0915 -- single orchestration point for the whole run loop; splitting it further would scatter closely-coupled loop state across more helper signatures than it'd save in readability
    x0: npt.NDArray[FloatDType],
    gradient_fn: "GradientFn[FloatDType] | MinibatchGradientFn[FloatDType]",
    config: SVGDConfig[FloatDType] | None = None,
) -> SVGDState[FloatDType]:
    """Update a collection of samples using Stein Variational Gradient Descent.

    Parameters
    ----------
    x0 : np.ndarray
        Initial particle positions, shape ``(n_particles, n_dims)``. Its
        dtype (e.g. ``float32`` or ``float64``) is preserved throughout the
        run -- pass ``float32`` particles to run the whole optimization in
        single precision. ``gradient_fn`` should return gradients in the
        same dtype; if it doesn't, they are cast to ``x0``'s dtype before
        use, so a careless ``gradient_fn`` can't silently upcast the run.
        Ignored when ``config.resume_from`` is set -- the run continues from
        ``resume_from.particles`` instead.
    gradient_fn : callable
        Computes gradients of the negative log-probability. Accepts particles
        of shape ``(n_particles, n_dims)`` and returns gradients of the same
        shape. When ``config.sigma.value`` is set or ``config.sigma.estimate``
        is ``True``, must return ``(gradients, misfits)`` where misfits has
        shape ``(n_particles,)``. When ``config.minibatch_sampler`` is set,
        called instead as ``gradient_fn(particles, batch_indices)`` -- see
        :class:`SVGDConfig`.
    config : SVGDConfig or None
        Every other tunable of the run (iteration count, step size,
        preconditioner, kernel, bounds, callback, resume, animation, ...).
        ``None`` uses ``SVGDConfig()``'s defaults. See :class:`SVGDConfig`
        for the full list of fields.

    Returns
    -------
    SVGDState
        Final optimizer state. Access ``.particles`` for the particle array.

    """
    if x0 is None or gradient_fn is None:
        raise ValueError("x0 and gradient_fn cannot be None!")
    if config is None:
        config = SVGDConfig()
    if config.minibatch_sampler is not None and config.sigma.n_data_samples is None:
        raise ValueError("n_data_samples is required when minibatch_sampler is set")
    use_svn = config.preconditioner == "svn"
    if use_svn and config.hessian_vector_product is None:
        raise ValueError("hessian_vector_product is required when preconditioner='svn'")

    kernel_fn = _resolve_kernel_fn(config.kernel)
    needs_misfits = config.sigma.value is not None or config.sigma.estimate
    step_schedule = _resolve_step_schedule(config.step_schedule, config.preconditioner)
    if use_svn and step_schedule == "adagrad":
        raise ValueError(
            "step_schedule='adagrad' is not supported with preconditioner='svn' -- "
            "AdaGrad recomputes its own step from the raw kernel/gradient terms and "
            "would silently ignore the Newton-preconditioned direction"
        )
    temperature_fn = _resolve_temperature_fn(config.temperature_schedule, config.n_iter)
    use_lbfgs = config.preconditioner == "lbfgs"

    run = _init_run_state(x0, config, use_lbfgs=use_lbfgs, step_schedule=step_schedule)
    sigma_prior_beta = _resolve_sigma_prior_beta(config, run.current_sigma)
    anim = _setup_animation(config, run.particles)
    rerun_session = _setup_rerun(config)

    progress = make_progress(disable=config.disable_progressbar)
    task_id = progress.add_task("SVGD", total=config.n_iter, stats="")
    progress.start()
    try:
        for loop_iter in range(config.n_iter):
            iteration = run.start_iter + loop_iter

            batch_indices = (
                config.minibatch_sampler(iteration)
                if config.minibatch_sampler is not None
                else None
            )
            all_grads, misfits = _evaluate_gradients(
                gradient_fn, run.particles, needs_misfits=needs_misfits, batch_indices=batch_indices
            )
            all_grads, misfits = _rescale_for_minibatch(
                all_grads, misfits, batch_indices, config.sigma.n_data_samples
            )
            _lbfgs_curvature_update(run, all_grads, use_lbfgs=use_lbfgs)

            mean_misfit = _record_misfits(run, misfits)
            _update_sigma(run, config, misfits, sigma_prior_beta)
            _record_sigma_history(run)
            _record_particle_variance(run)

            # This iteration's repulsion ratio hasn't been computed yet at
            # this point in the loop (it's derived from the displacement
            # step, further down) -- stats.repulsion_ratio is the previous
            # iteration's value, one iteration stale.
            stats = RunStats(
                mean_misfit=mean_misfit,
                current_sigma=run.current_sigma,
                particle_variance=run.particle_variance_history[-1],
                repulsion_ratio=(
                    run.repulsion_ratio_history[-1] if run.repulsion_ratio_history else None
                ),
            )
            _report_progress(progress, task_id, stats)
            maybe_log_iteration(rerun_session, iteration, run.particles, stats)

            if config.callback is not None:
                config.callback(iteration, _snapshot(run, iteration))

            # Last iteration: don't update particles
            if loop_iter == config.n_iter - 1:
                break

            precond_grads = _precondition_gradients(run, all_grads, use_lbfgs=use_lbfgs)
            kernel_matrix, kernel_grad = kernel_fn(run.particles, config.bandwidth)
            temperature = temperature_fn(iteration)
            grads_scaled = _scale_grads_by_sigma(precond_grads, run.current_sigma, temperature)
            # SVGD update direction:
            #   phi = -(K @ grads_scaled - nabla_K) / n_particles
            attractive = np.matmul(kernel_matrix, grads_scaled)
            phi = -(attractive - kernel_grad) / run.n_particles
            _record_repulsion_ratio(run, kernel_grad, attractive)

            if use_svn:
                assert config.hessian_vector_product is not None  # noqa: S101 -- checked above
                phi = svn_direction(config.hessian_vector_product, run.particles, phi, config.svn)

            inputs = _StepInputs(
                kernel_matrix=kernel_matrix,
                kernel_grad=kernel_grad,
                phi=phi,
                attractive=attractive,
                precond_grads=precond_grads,
            )
            displacement = _compute_displacement(step_schedule, run, config, inputs, iteration)
            _store_lbfgs_prev(run, all_grads, use_lbfgs=use_lbfgs)

            run.particles = run.particles + displacement

            if config.bounds is not None:
                run.particles = np.clip(run.particles, config.bounds[0], config.bounds[1])

            if anim is not None:
                draw_frame(anim, run.particles, config.animation.dimensions_to_plot)

    except KeyboardInterrupt:
        pass
    finally:
        progress.stop()

    return SVGDState(
        particles=run.particles,
        iteration=run.start_iter + config.n_iter,
        lbfgs_states=run.lbfgs_states,
        historical_grad=run.historical_grad,
        data_sigma=run.current_sigma,
        sigma_history=run.sigma_history,
        misfit_history=run.misfit_history,
        particle_misfit_history=run.particle_misfit_history,
        particle_variance_history=run.particle_variance_history,
        repulsion_ratio_history=run.repulsion_ratio_history,
        prev_particles=run.prev_particles if use_lbfgs else None,
        prev_grads=run.prev_grads if use_lbfgs else None,
    )

simplesvgd.SVGDConfig dataclass

SVGDConfig(
    n_iter: int = 1000,
    stepsize: float = 0.001,
    bandwidth: float = -1,
    preconditioner: str | None = None,
    lbfgs: LBFGSConfig = LBFGSConfig(),
    svn: SVNConfig = SVNConfig(),
    hessian_vector_product: HessianVectorProductFn[
        FloatDType
    ]
    | None = None,
    step_schedule: str | None = None,
    temperature_schedule: str
    | Callable[[int], float]
    | None = None,
    sigma: SigmaConfig = SigmaConfig(),
    minibatch_sampler: Callable[[int], BatchIndices]
    | None = None,
    kernel: str | KernelFn[FloatDType] | None = None,
    bounds: tuple[float, float] | None = None,
    callback: Callable[[int, SVGDState[FloatDType]], None]
    | None = None,
    disable_progressbar: bool = False,
    resume_from: SVGDState[FloatDType] | None = None,
    animation: AnimationConfig[
        FloatDType
    ] = AnimationConfig(),
    rerun: RerunConfig[FloatDType] = RerunConfig(),
)

Bases: Generic[FloatDType]

Every tunable of :func:update, grouped into one object.

All fields have defaults, so SVGDConfig() reproduces update()'s previous behavior with no arguments. Pass a partially-filled instance to override just what you need, e.g. SVGDConfig(n_iter=500, bandwidth=2.0). Related tunables are grouped into sub-objects -- L-BFGS settings under lbfgs, hierarchical noise estimation under sigma, and the legacy animation under animation -- constructed the same way, e.g. SVGDConfig(sigma=SigmaConfig(value=0.1, estimate=True)).

Attributes:

Name	Type	Description
n_iter	int	Number of iterations.
stepsize	float	Base step size (interpretation depends on step_schedule).
bandwidth	float	RBF kernel bandwidth. -1 for automatic (median heuristic).
preconditioner	str or None	None for AdaGrad+momentum (legacy). "lbfgs" for per-particle L-BFGS preconditioning with Robbins-Monro step decay. "svn" for mean-field Stein Variational Newton preconditioning (requires hessian_vector_product; see :mod:simplesvgd.svn) -- unlike "lbfgs", which preconditions each particle's raw gradient before kernel combination, this Newton-preconditions the full, already kernel-combined SVGD direction using one shared curvature operator for all particles.
lbfgs	LBFGSConfig	L-BFGS preconditioning parameters (only used for "lbfgs").
svn	SVNConfig	Mean-field SVN preconditioning parameters (only used for "svn").
hessian_vector_product	callable or None	Computes a Hessian-vector product, called as hessian_vector_product(particles, vectors) with both arguments of shape (n_particles, n_dims), returning an array of the same shape (row i is the Hessian-vector product at particles[i] applied to vectors[i]). Required when preconditioner == "svn"; ignored otherwise.
step_schedule	str or None	None uses the default for the preconditioner (AdaGrad for None, Robbins-Monro for "lbfgs", constant for "svn" -- "svn" already Newton-preconditions the full displacement, so stepsize there acts as a damped-Newton step fraction rather than needing the Robbins-Monro decay a noisier direction would). "robbins-monro" uses stepsize / (|attractive|_max * sqrt(1+t)). "constant" uses fixed stepsize. "adagrad" uses AdaGrad+momentum -- not supported together with preconditioner="svn" (AdaGrad recomputes its own step from the raw kernel/gradient terms and would silently ignore the Newton preconditioning).
temperature_schedule	str, callable, or None	Anneals the data-misfit gradient contribution: the (preconditioned, sigma-scaled) gradient is multiplied by a temperature in (0, 1] before the SVGD attractive term is formed, so a temperature below 1 reads as an inflated apparent noise level. Useful when the posterior is sharply peaked relative to the initial particle spread (e.g. full-waveform-inversion-style problems), where applying the full likelihood from iteration 0 risks collapsing particles onto a subset of modes before repulsion has had a chance to spread them out. None applies temperature 1.0 throughout (current behavior, unchanged). "linear" ramps linearly from a small floor to 1.0 over the run. "geometric" ramps geometrically (log-linear) from the same floor to 1.0. A callable is called as temperature_schedule(iteration) and must return a float in (0, 1].
sigma	SigmaConfig	Hierarchical likelihood-noise estimation parameters.
minibatch_sampler	callable or None	Called as minibatch_sampler(iteration), returning an array of indices into the data dimension (e.g. sources/receivers) to use for that iteration. When set, gradient_fn is called as gradient_fn(particles, batch_indices) instead of gradient_fn(particles), and must accept that second argument. The returned gradients (and misfits, if sigma.value is set or sigma.estimate is True) are treated as sums over just the sampled subset and rescaled by sigma.n_data_samples / len(batch_indices) -- an unbiased estimate of the full-dataset sum -- before use, so the SVGD attractive term and the hierarchical sigma estimate stay calibrated as if the full dataset had been evaluated. Requires sigma.n_data_samples to be set. None (default) evaluates the full dataset every iteration, unchanged.
kernel	str, KernelFn, or None	Kernel type. None or "rbf" for standard RBF. "rbf_normalized" for per-dimension normalized RBF, recommended for high-dimensional parameter spaces (d > ~100) where standard RBF repulsion vanishes. Can also be a custom KernelFn, e.g. one built with make_mass_weighted_kernel(weights) for particles representing a field discretized on a mesh or grid -- weighting distances by cell volume/quadrature weight keeps the kernel (and the resulting posterior statistics) stable as the mesh is refined, unlike plain Euclidean RBF.
bounds	tuple[float, float] or None	(lower, upper) bounds for particle clipping.
callback	callable or None	Called as callback(iteration, state) after each gradient evaluation.
disable_progressbar	bool	Suppress the live progress display (bar, ETA, and a live misfit/sigma/particle-variance/repulsion-ratio stats line).
resume_from	SVGDState or None	Resume from a previous run's state. When set, the run continues from resume_from.particles rather than x0.
animation	AnimationConfig	Legacy live-scatter animation parameters.
rerun	RerunConfig	Live particle + diagnostics visualization via the Rerun viewer, with a scrubbable timeline (see RerunConfig).

simplesvgd.SVGDState dataclass

SVGDState(
    particles: NDArray[FloatDType],
    iteration: int = 0,
    lbfgs_states: list[LBFGSState[FloatDType]]
    | None = None,
    historical_grad: NDArray[FloatDType] | None = None,
    data_sigma: float | None = None,
    sigma_history: list[float] = list(),
    misfit_history: list[float] = list(),
    particle_misfit_history: list[list[float]] = list(),
    particle_variance_history: list[float] = list(),
    repulsion_ratio_history: list[float] = list(),
    prev_particles: NDArray[FloatDType] | None = None,
    prev_grads: NDArray[FloatDType] | None = None,
)

Bases: Generic[FloatDType]

Complete state of an SVGD run.

This object is returned by :func:simplesvgd.update and can be passed back via SVGDConfig(resume_from=...) to continue optimization.

Attributes: particles: Current particle positions, shape (n_particles, n_dims). iteration: Total number of completed iterations. lbfgs_states: Per-particle L-BFGS states (None when using AdaGrad). historical_grad: AdaGrad accumulator (None when using L-BFGS). data_sigma: Current likelihood noise standard deviation. sigma_history: data_sigma at each iteration. misfit_history: Mean misfit across particles at each iteration. particle_misfit_history: Per-particle misfits at each iteration. particle_variance_history: Total particle-ensemble variance (trace of the empirical covariance, i.e. sum of per-dimension variances) at each iteration -- a cheap proxy for ensemble spread. A value that shrinks steadily over the run, well below what the target distribution's actual variance should be, is a variance-collapse warning sign (see Ba et al., "Understanding the Variance Collapse of SVGD in High Dimensions", ICLR 2022). repulsion_ratio_history: Ratio of the repulsive kernel-gradient term's norm to the attractive term's norm, at each iteration a displacement is computed (shorter than the other histories -- the final iteration of any given update() call only records state, it doesn't step, so each resume_from boundary drops one more entry than the other histories accumulate). Read this early in a run, not as a monotonic trend over the whole run -- the attractive term shrinks toward zero near any converged mode regardless of collapse, which swamps the ratio's trend late in a run. A ratio far below 1 in the first few iterations (while particles are still diffuse and attraction hasn't decayed yet) is the direct, cheap signature of the collapse mechanism in the paper above: repulsion already overwhelmed by attraction before it's had any chance to spread the ensemble out. prev_particles: Previous particle positions (for deferred L-BFGS update). prev_grads: Previous gradients (for deferred L-BFGS update).

PyTorch bridge

simplesvgd.update_torch

update_torch(
    x0: NDArray[FloatDType],
    gradient_fn: Callable[
        [NDArray[FloatDType]], NDArray[FloatDType]
    ],
    optimizer_class: type[_TorchOptimizerLike],
    optimizer_parameters: dict[str, Any] | None = None,
    schedulers: list[_TorchSchedulerLike] | None = None,
    *,
    n_iter: int = 1000,
    animate: bool = False,
    figure: Figure | None = None,
    dimensions_to_plot: list[int] | None = None,
    background: Background[FloatDType] | None = None,
    disable_progressbar: bool = False,
) -> npt.NDArray[FloatDType]

Update samples using SVGD with a PyTorch optimizer.

Parameters:

Name	Type	Description	Default
x0	ndarray	Initial samples, shape (n_samples, dimensionality).	required
gradient_fn	callable	Computes gradients of the negative log-probability.	required
optimizer_class	torch.optim.Optimizer subclass	PyTorch optimizer to use.	required
optimizer_parameters	dict or None	Keyword arguments forwarded to the optimizer constructor.	None
schedulers	list or None	Learning rate schedulers to step after each iteration.	None
n_iter	int	Number of iterations.	1000
animate	bool	Enable 2D scatter animation.	False
figure	matplotlib Figure or None	Figure to draw the animation on; created if not given.	None
dimensions_to_plot	list of int or None	Which two particle dimensions to animate. Defaults to [0, 1].	None
background	tuple or None	(x1s, x2s, background_image) contour data to draw behind the animation.	None
disable_progressbar	bool	Suppress the tqdm progress bar.	False

Source code in src/simplesvgd/__init__.py

def update_torch(  # noqa: PLR0913 -- legacy PyTorch bridge; not covered by update()'s SVGDConfig redesign
    x0: npt.NDArray[FloatDType],
    gradient_fn: Callable[[npt.NDArray[FloatDType]], npt.NDArray[FloatDType]],
    optimizer_class: type[_TorchOptimizerLike],
    optimizer_parameters: dict[str, Any] | None = None,
    schedulers: list[_TorchSchedulerLike] | None = None,
    *,
    n_iter: int = 1000,
    animate: bool = False,
    figure: "Figure | None" = None,
    dimensions_to_plot: list[int] | None = None,
    background: Background[FloatDType] | None = None,
    disable_progressbar: bool = False,
) -> npt.NDArray[FloatDType]:
    """Update samples using SVGD with a PyTorch optimizer.

    Parameters
    ----------
    x0 : np.ndarray
        Initial samples, shape ``(n_samples, dimensionality)``.
    gradient_fn : callable
        Computes gradients of the negative log-probability.
    optimizer_class : torch.optim.Optimizer subclass
        PyTorch optimizer to use.
    optimizer_parameters : dict or None
        Keyword arguments forwarded to the optimizer constructor.
    schedulers : list or None
        Learning rate schedulers to step after each iteration.
    n_iter : int
        Number of iterations.
    animate : bool
        Enable 2D scatter animation.
    figure : matplotlib Figure or None
        Figure to draw the animation on; created if not given.
    dimensions_to_plot : list of int or None
        Which two particle dimensions to animate. Defaults to ``[0, 1]``.
    background : tuple or None
        ``(x1s, x2s, background_image)`` contour data to draw behind the
        animation.
    disable_progressbar : bool
        Suppress the tqdm progress bar.

    """
    import torch  # noqa: PLC0415 -- torch is an optional extra  # ty: ignore[unresolved-import]

    from .helpers import torch_wrapper  # noqa: PLC0415 -- only needed for this optional path

    if x0 is None or gradient_fn is None:
        raise ValueError("x0 or gradient_fn cannot be None!")

    if optimizer_parameters is None:
        optimizer_parameters = {}
    if schedulers is None:
        schedulers = []
    if dimensions_to_plot is None:
        dimensions_to_plot = [0, 1]

    x0_updated = np.copy(x0)

    anim = (
        setup_animation(
            figure=figure,
            background=background,
            particles=x0_updated,
            dimensions_to_plot=dimensions_to_plot,
        )
        if animate
        else None
    )

    x0_updated_tensor = torch.tensor(x0_updated, requires_grad=True)
    total = torch_wrapper(gradient_fn, rbf_kernel)
    optimizer = optimizer_class([x0_updated_tensor], **optimizer_parameters)

    try:
        for _ in tqdm_auto.trange(n_iter, disable=disable_progressbar):
            def closure() -> Any:  # noqa: ANN401 -- torch.Tensor loss, unavailable without installing torch for type-checking
                optimizer.zero_grad()
                loss = total(x0_updated_tensor).mean()
                loss.backward()
                return loss

            optimizer.step(closure)

            for scheduler in schedulers:
                scheduler.step()

            if anim is not None:
                draw_frame(anim, x0_updated_tensor.detach().numpy(), dimensions_to_plot)

    except KeyboardInterrupt:
        sleep(0.5)

    return x0_updated_tensor.detach().numpy()

Kernels

simplesvgd.rbf_kernel_normalized

rbf_kernel_normalized(
    particles: NDArray[FloatDType], h: float = -1
) -> tuple[
    npt.NDArray[FloatDType], npt.NDArray[FloatDType]
]

RBF kernel with per-dimension normalization for high-dimensional spaces.

In high dimensions (d >> 1), the standard median heuristic produces bandwidth h^2 ~ O(d), which causes the repulsive gradient per dimension to scale as O(1/d) while the attractive gradient stays O(1). This kills particle diversity as d grows.

This variant normalizes each dimension to unit variance before computing pairwise distances and the bandwidth, then maps the kernel gradient back to the original space. The effective bandwidth is dimension-independent, preserving repulsion in spaces with thousands of dimensions (e.g. FWI parameter vectors).

Source code in src/simplesvgd/kernels.py

def rbf_kernel_normalized(
    particles: npt.NDArray[FloatDType], h: float = -1
) -> tuple[npt.NDArray[FloatDType], npt.NDArray[FloatDType]]:
    """RBF kernel with per-dimension normalization for high-dimensional spaces.

    In high dimensions (d >> 1), the standard median heuristic produces
    bandwidth h^2 ~ O(d), which causes the repulsive gradient per dimension
    to scale as O(1/d) while the attractive gradient stays O(1).  This kills
    particle diversity as d grows.

    This variant normalizes each dimension to unit variance before computing
    pairwise distances and the bandwidth, then maps the kernel gradient back
    to the original space.  The effective bandwidth is dimension-independent,
    preserving repulsion in spaces with thousands of dimensions (e.g. FWI
    parameter vectors).
    """
    n_particles, n_dims = particles.shape

    # Per-dimension normalization
    std = np.std(particles, axis=0)
    # Avoid division by zero for constant dimensions
    std = np.where(std < _DEGENERATE_SCALE_THRESHOLD, 1.0, std)
    particles_normalized = particles / std

    # Compute kernel in normalized space
    pairwise_dists = _pairwise_sq_dists(particles_normalized)
    bandwidth: FloatDType | float
    if h < 0:
        bandwidth = np.median(pairwise_dists)
        bandwidth = np.sqrt(0.5 * bandwidth / np.log(n_particles + 1))
        bandwidth = particles.dtype.type(bandwidth)
    else:
        bandwidth = h

    if bandwidth < _DEGENERATE_SCALE_THRESHOLD:
        return np.ones((n_particles, n_particles), dtype=particles.dtype), np.zeros_like(particles)

    kernel_matrix = np.exp(-pairwise_dists / bandwidth ** 2 / 2)

    # Kernel gradient in normalized space
    kernel_grad_normalized = -np.matmul(kernel_matrix, particles_normalized)
    kernel_row_sums = np.sum(kernel_matrix, axis=1)
    for i in range(n_dims):
        kernel_grad_normalized[:, i] = kernel_grad_normalized[:, i] + np.multiply(
            particles_normalized[:, i], kernel_row_sums
        )
    kernel_grad_normalized = kernel_grad_normalized / (bandwidth ** 2)

    # Map gradient back to original space: d/dx_k = (1/std_k) * d/dx_n_k
    kernel_grad = kernel_grad_normalized / std

    return (kernel_matrix, kernel_grad)

L-BFGS preconditioning

simplesvgd.LBFGSState dataclass

LBFGSState(
    s_history: NDArray[FloatDType],
    y_history: NDArray[FloatDType],
    cursor: int = 0,
    count: int = 0,
)

Bases: Generic[FloatDType]

Circular buffer storing L-BFGS curvature pairs (s, y).

Attributes: s_history: Array of shape (m, n) storing s_k = x_{k+1} - x_k vectors. y_history: Array of shape (m, n) storing y_k = g_{k+1} - g_k vectors. cursor: Index where the next pair will be written. count: Number of pairs stored so far (up to m).

simplesvgd.make_lbfgs_state

make_lbfgs_state(
    n: int, m: int = 10
) -> LBFGSState[np.float64]

make_lbfgs_state(
    n: int, m: int = 10, *, dtype: dtype[FloatDType]
) -> LBFGSState[FloatDType]

make_lbfgs_state(
    n: int,
    m: int = 10,
    *,
    dtype: dtype[FloatDType] | type[float64] = np.float64,
) -> LBFGSState[FloatDType] | LBFGSState[np.float64]

Create an empty L-BFGS state with history size m for vectors of length n.

dtype should match the dtype of the gradients/particles this state will be used with (e.g. particles.dtype), so the two-loop recursion in :func:lbfgs_direction doesn't get upcast by a mismatched buffer dtype. Defaults to float64 when omitted.

Source code in src/simplesvgd/lbfgs.py

def make_lbfgs_state(
    n: int, m: int = 10, *, dtype: np.dtype[FloatDType] | type[np.float64] = np.float64
) -> LBFGSState[FloatDType] | LBFGSState[np.float64]:
    """Create an empty L-BFGS state with history size *m* for vectors of length *n*.

    *dtype* should match the dtype of the gradients/particles this state
    will be used with (e.g. ``particles.dtype``), so the two-loop recursion
    in :func:`lbfgs_direction` doesn't get upcast by a mismatched buffer
    dtype. Defaults to float64 when omitted.
    """
    return LBFGSState(
        s_history=np.zeros((m, n), dtype=dtype),
        y_history=np.zeros((m, n), dtype=dtype),
        cursor=0,
        count=0,
    )

simplesvgd.lbfgs_direction

lbfgs_direction(
    state: LBFGSState[FloatDType], grad: NDArray[FloatDType]
) -> npt.NDArray[FloatDType]

Compute the L-BFGS search direction via two-loop recursion.

Returns -H_k @ grad where H_k is the L-BFGS approximation to the inverse Hessian. Falls back to -grad when the history is empty.

Source code in src/simplesvgd/lbfgs.py

def lbfgs_direction(
    state: LBFGSState[FloatDType], grad: npt.NDArray[FloatDType]
) -> npt.NDArray[FloatDType]:
    """Compute the L-BFGS search direction via two-loop recursion.

    Returns ``-H_k @ grad`` where H_k is the L-BFGS approximation to the
    inverse Hessian.  Falls back to ``-grad`` when the history is empty.
    """
    k = state.count
    if k == 0:
        return -grad.copy()

    m = state.s_history.shape[0]
    # Indices from newest to oldest
    indices = [(state.cursor - 1 - i) % m for i in range(k)]

    q = grad.copy()
    # dtype=grad.dtype matters: a plain np.zeros(k) defaults to float64, and
    # a float64 numpy scalar pulled from it would upcast every float32 array
    # it later touches (unlike a plain Python float, which stays "weak" and
    # doesn't force a promotion).
    alphas = np.zeros(k, dtype=grad.dtype)
    rhos = np.zeros(k, dtype=grad.dtype)

    # Forward pass (newest to oldest)
    for j, idx in enumerate(indices):
        s_j = state.s_history[idx]
        y_j = state.y_history[idx]
        rho_j = 1.0 / np.dot(y_j, s_j)
        rhos[j] = rho_j
        alpha_j = rho_j * np.dot(s_j, q)
        alphas[j] = alpha_j
        q = q - alpha_j * y_j

    # Initial Hessian scaling: H0 = (y_k . s_k) / (y_k . y_k) * I
    newest = indices[0]
    s_newest = state.s_history[newest]
    y_newest = state.y_history[newest]
    gamma = np.dot(y_newest, s_newest) / np.dot(y_newest, y_newest)
    r = gamma * q

    # Backward pass (oldest to newest)
    for j in reversed(range(k)):
        idx = indices[j]
        y_j = state.y_history[idx]
        beta = rhos[j] * np.dot(y_j, r)
        r = r + (alphas[j] - beta) * state.s_history[idx]

    return -r

simplesvgd.lbfgs_update

lbfgs_update(
    state: LBFGSState[FloatDType],
    s: NDArray[FloatDType],
    y: NDArray[FloatDType],
) -> None

Push a new (s, y) pair into the circular buffer.

Skips the update if the curvature condition y.s > 0 is not satisfied.

Source code in src/simplesvgd/lbfgs.py

def lbfgs_update(
    state: LBFGSState[FloatDType], s: npt.NDArray[FloatDType], y: npt.NDArray[FloatDType]
) -> None:
    """Push a new (s, y) pair into the circular buffer.

    Skips the update if the curvature condition y.s > 0 is not satisfied.
    """
    ys = float(np.dot(s, y))
    if ys <= 0:
        return
    m = state.s_history.shape[0]
    idx = state.cursor % m
    state.s_history[idx] = s
    state.y_history[idx] = y
    state.cursor = (idx + 1) % m
    state.count = min(state.count + 1, m)

Mean-field SVN preconditioning

simplesvgd.SVNConfig dataclass

SVNConfig(cg_iters: int = 10, damping: float = 1e-06)

Mean-field Stein Variational Newton preconditioning parameters.

Only used when SVGDConfig.preconditioner == "svn", which also requires SVGDConfig.hessian_vector_product to be set. See :mod:simplesvgd.svn for what "mean-field" means here and how it differs from the full Stein Variational Newton algorithm.

Attributes:

Name	Type	Description
cg_iters	int	Conjugate-gradient iterations used to solve each iteration's shared Newton system.
damping	float	Tikhonov damping added to the shared curvature operator before solving, so it stays positive-definite (and the CG solve well-posed) even when hessian_vector_product returns a rank-deficient or near-singular Hessian-vector product.

simplesvgd.svn_direction

svn_direction(
    hessian_vector_product: HessianVectorProductFn[
        FloatDType
    ],
    particles: NDArray[FloatDType],
    phi: NDArray[FloatDType],
    svn_config: SVNConfig,
) -> npt.NDArray[FloatDType]

Newton-precondition phi with the shared, mean-field curvature operator.

Solves (H_mean + damping * I) @ v_i = phi_i for every particle i via batched CG, where H_mean is hessian_vector_product evaluated at the particle ensemble's mean position (the same operator for every particle). phi is the already kernel-combined SVGD direction (attraction and repulsion together), shape (n_particles, n_dims).

Source code in src/simplesvgd/svn.py

def svn_direction(
    hessian_vector_product: HessianVectorProductFn[FloatDType],
    particles: npt.NDArray[FloatDType],
    phi: npt.NDArray[FloatDType],
    svn_config: "SVNConfig",
) -> npt.NDArray[FloatDType]:
    """Newton-precondition *phi* with the shared, mean-field curvature operator.

    Solves ``(H_mean + damping * I) @ v_i = phi_i`` for every particle *i*
    via batched CG, where ``H_mean`` is *hessian_vector_product* evaluated at
    the particle ensemble's mean position (the same operator for every
    particle). *phi* is the already kernel-combined SVGD direction
    (attraction and repulsion together), shape ``(n_particles, n_dims)``.
    """
    n_particles = particles.shape[0]
    mean_particle = np.mean(particles, axis=0, keepdims=True)
    mean_particles = np.tile(mean_particle, (n_particles, 1)).astype(particles.dtype)
    damping = svn_config.damping

    def apply_operator(v: npt.NDArray[FloatDType]) -> npt.NDArray[FloatDType]:
        hv = np.asarray(hessian_vector_product(mean_particles, v), dtype=phi.dtype)
        return cast("npt.NDArray[FloatDType]", hv + damping * v)

    return _batched_cg(apply_operator, phi, svn_config.cg_iters)

Utilities

simplesvgd.gradient_vectorizer

gradient_vectorizer(
    non_vectorized_gradient: Callable[
        [NDArray[FloatDType]], NDArray[FloatDType]
    ],
) -> Callable[
    [npt.NDArray[FloatDType]], npt.NDArray[FloatDType]
]

Wrap a single-point gradient function to accept batched inputs.

Source code in src/simplesvgd/__init__.py

def gradient_vectorizer(
    non_vectorized_gradient: Callable[[npt.NDArray[FloatDType]], npt.NDArray[FloatDType]],
) -> Callable[[npt.NDArray[FloatDType]], npt.NDArray[FloatDType]]:
    """Wrap a single-point gradient function to accept batched inputs."""
    def grd(m: npt.NDArray[FloatDType]) -> npt.NDArray[FloatDType]:
        return np.hstack(
            [non_vectorized_gradient(m[idm, :, None]) for idm in range(m.shape[0])]
        ).T
    return grd