stripe / rainier

Bayesian inference in Scala.
https://rainier.fit
Apache License 2.0
433 stars 51 forks source link

Introducing SamplerConfig #487

Closed avibryant closed 4 years ago

avibryant commented 4 years ago

This is another try at refactoring Sampler to allow for more modular tuning. The approach here centers around a new SamplerConfig trait:

trait SamplerConfig {
  def iterations: Int
  def warmupIterations: Int

  def stepSizeTuner(): StepSizeTuner
  def metricTuner(): MetricTuner
  def sampler(): Sampler
}

The notable thing here is that it separately/orthogonally creates stateful Tuner and Sampler objects for each chain. There's now a Driver which takes care of threading things together, and Sampler's job is narrowed to just performing a single iteration of either the warmup or sampling phase, given a stepsize and metric. (Windowed tuning can be implemented within the tuners, and there's a mechanism for communicating a reset to the stepsize tuner when the metric gets replaced).

A backwards-compatible HMC() constructor is provided for an appropriate config. For the moment, EHMC and mass matrix tuning have been dropped to keep the PR small, and will be reintroduced.

avibryant commented 4 years ago

Ok, looks like I need to bring back EHMC so that the docs compile, before merging.