docs menu

Module Norn.Ensemble

Ensemble sampling by the stretch move, without derivatives.

Walkers lie on the chain axis in independent ensembles of equal size, consecutive on the axis. A transition splits each ensemble at random into two halves, which mixes faster than fixed halves, and updates them in turn: a walker x of the moving half draws a walker w of the other half and a stretch z of density proportional to 1 / sqrt z on [1 / 2, 2], and moves to w + z (x - w) with probability min (1, z^(d - 1) p(y) / p(x)), d its float elements (Goodman and Weare 2010). A walker's move depends only on the half that does not move, as detailed balance requires, and on no scale: the move is unchanged by an affine map of the coordinates, so a correlated posterior costs what an isotropic one does. Nothing is tuned.

The density is evaluated once per transition on the moving half of every ensemble, a call of walkers / 2 rows when ensembles have an even size.

The walkers move along lines through each other, so they explore the region their spread spans. Walkers that start far from a narrow, correlated posterior, spread over a region of another shape, come back to it only over many transitions: start them where a warmed sampler or the posterior's approximation leaves them.

Posteriors it suits. The move suits a posterior that one affine map brings close to a Gaussian. Where the posterior's shape changes from region to region, as when a hyperparameter sets the scale of other elements in a hierarchical model, no such map exists. A move along the line between two walkers then stays where the density is high only when it is short, and the elements whose value shapes the rest, such as the hyperparameters, move slowest. Their effective sizes (Diag.ess_bulk) grow far slower than the draws, Summary.v reports Ess_low for them, and their nested R-hat may exceed 1.01. Nuts follows the gradient, so its trajectories bend with the posterior. Use it when the density has a gradient.

Keys. Transition n of a run has the key k = Nx.Rng.fold_in run n, and walker i row i of Nx.Rng.split_batch ~n:walkers k; an ensemble's split draws from its first walker's key. So with ensembles of a fixed size a walker's randomness does not depend on the number of ensembles. warmup and sample fold the state's draw counter into their key and advance it, so a + b draws are a draws then b.

#type ('f, 'd) stats = {

#lp : (float, 'f, 'd) Nx.t;
(*
The log density at the transition's end.
*)
#acceptance : (float, 'f, 'd) Nx.t;
(*
The probability with which the walker's move was accepted.
*)

}

The type for statistics of transitions, one element per walker.

#type ('u, 'f, 'd) state = private {

#position : 'u;
(*
Each walker's position.
*)
#lp : (float, 'f, 'd) Nx.t;
(*
The log density there, shape [walker].
*)
#stats : ('f, 'd) stats;
(*
The last transition's statistics.
*)
#draw : 'd Nx.int32_t;
(*
The number of transitions taken.
*)
#ensembles : int;
(*
The number of ensembles.
*)

}

The type for states of walkers over positions 'u.

#val ptree : ('u, 'd) Nx.Ptree.t -> (('u, 'f, 'd) state, 'd) Nx.Ptree.t

ptree u is the structure of states over u. It reports ensembles.

#val stats : (float, 'f) Nx.dtype -> (('f, 'd) stats, 'd) Nx.Ptree.t

stats dtype is the structure of statistics at dtype.

val init : 
  ('u, 'd) Nx.Ptree.t ->
  ?ensembles:int ->
  ('u -> (float, 'f, 'd) Nx.t) ->
  'u ->
  ('u, 'f, 'd) state

init u ?ensembles lp start is the state of walkers at start, a position, in ensembles ensembles, by default 1.

Raises Invalid_argument if ensembles < 1, if the walkers do not split into ensembles ensembles of equal size, if an ensemble has fewer walkers than twice a walker's float elements, or if lp does not return one log density per walker, each read from its own row.

val step : 
  ('u, 'd) Nx.Ptree.t ->
  ('u -> (float, 'f, 'd) Nx.t) ->
  'd Nx.Rng.key ->
  ('u, 'f, 'd) state ->
  ('u, 'f, 'd) state

step u lp k s is one transition of every walker from s.

Raises Invalid_argument if lp is NaN or +inf at a walker.

val warmup : 
  ('u, 'd) Nx.Ptree.t ->
  ('u -> (float, 'f, 'd) Nx.t) ->
  'd Nx.Rng.key ->
  steps:int ->
  ('u, 'f, 'd) state ->
  ('u, 'f, 'd) state

warmup u lp k ~steps s is s after steps transitions, whose draws are discarded. Transition n of the run, counted by s.draw, has the key Nx.Rng.fold_in k n, as in sample.

Raises Invalid_argument if steps < 0.

val sample : 
  ('u, 'd) Nx.Ptree.t ->
  ('u -> (float, 'f, 'd) Nx.t) ->
  'd Nx.Rng.key ->
  draws:int ->
  ('u, 'f, 'd) state ->
  ('u, 'f, 'd) state * 'u Draws.t * ('f, 'd) stats Draws.t

sample u lp k ~draws s is the state after draws transitions, the walkers' positions and their statistics, each walker a chain of the draws. A walker's moves read the other walkers of its ensemble, so the chains of one ensemble are dependent and those of two independent: their R-hat is Diag.nested_rhat with a superchain per ensemble (Summary.v's superchains at s.ensembles). Transition n of the run, counted by s.draw, has the key Nx.Rng.fold_in k n.

Raises Invalid_argument if draws < 1.