Remove Sampler and initialstep

Author: penelopeysm

HISTORY.md

@@ -1,5 +1,12 @@
 # DynamicPPL Changelog
 
+## 0.38.0
+
+TODO: write up InitContext changes
+
+`DynamicPPL.Sampler` and `DynamicPPL.initialstep` have also been removed entirely.
+If you were using these, you should unwrap your sampler and implement `AbstractMCMC.step` as usual for any other `AbstractMCMC.AbstractSampler`.
+
 ## 0.37.0
 
 DynamicPPL 0.37 comes with a substantial reworking of its internals.

docs/src/api.md

@@ -493,23 +493,21 @@ DynamicPPL.AbstractInitStrategy
 DynamicPPL.init
 ```
 
-### Samplers
+### Sampling
 
-In DynamicPPL a generic sampler for inference is implemented.
+`init_strategy(spl)` defines how parameters are to be initialised when performing MCMC sampling with `spl`.
 
 ```@docs
-Sampler
+DynamicPPL.init_strategy
 ```
 
-The default implementation of [`Sampler`](@ref) uses the following unexported functions.
+`loadstate` is used for resuming sampling from a previous chain.
 
 ```@docs
-DynamicPPL.initialstep
 DynamicPPL.loadstate
-DynamicPPL.init_strategy
 ```
 
-Finally, to specify which varinfo type a [`Sampler`](@ref) should use for a given [`Model`](@ref), this is specified by [`DynamicPPL.default_varinfo`](@ref) and can thus be overloaded for each  `model`-`sampler` combination. This can be useful in cases where one has explicit knowledge that one type of varinfo will be more performant for the given `model` and `sampler`.
+Finally, to specify which varinfo type a sampler should use for a given [`Model`](@ref), this is specified by [`DynamicPPL.default_varinfo`](@ref) and can thus be overloaded for each  `model`-`sampler` combination. This can be useful in cases where one has explicit knowledge that one type of varinfo will be more performant for the given `model` and `sampler`.
 
 ```@docs
 DynamicPPL.default_varinfo

src/DynamicPPL.jl

@@ -94,8 +94,6 @@ export AbstractVarInfo,
     getargnames,
     extract_priors,
     values_as_in_model,
-    # Samplers
-    Sampler,
     # LogDensityFunction
     LogDensityFunction,
     # Contexts

src/sampler.jl

@@ -1,23 +1,3 @@
-# TODO(mhauru) Could we get rid of Sampler now that it's just a wrapper around `alg`?
-# (Selector has been removed).
-"""
-    Sampler{T}
-
-Generic sampler type for inference algorithms of type `T` in DynamicPPL.
-
-`Sampler` should implement the AbstractMCMC interface, and in particular
-`AbstractMCMC.step`. A default implementation of the initial sampling step is
-provided that supports resuming sampling from a previous state and setting initial
-parameter values. It requires to overload [`loadstate`](@ref) and [`initialstep`](@ref)
-for loading previous states and actually performing the initial sampling step,
-respectively. Additionally, sometimes one might want to implement [`init_strategy`](@ref)
-that specifies how the initial parameter values are sampled if they are not provided.
-By default, values are sampled from the prior.
-"""
-struct Sampler{T} <: AbstractSampler
-    alg::T
-end
-
 """
     default_varinfo(rng, model, sampler)
 
@@ -39,49 +19,13 @@ function default_varinfo(::Random.AbstractRNG, ::Model, ::AbstractSampler)
     return typed_varinfo(VarInfo())
 end
 
-function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
-    model::Model,
-    sampler::Sampler,
-    N::Integer;
-    chain_type=default_chain_type(sampler),
-    resume_from=nothing,
-    initial_state=loadstate(resume_from),
-    kwargs...,
-)
-    return AbstractMCMC.mcmcsample(
-        rng, model, sampler, N; chain_type, initial_state, kwargs...
-    )
-end
-
 """
     init_strategy(sampler)
 
 Define the initialisation strategy used for generating initial values when
 sampling with `sampler`. Defaults to `PriorInit()`, but can be overridden.
 """
-init_strategy(::Sampler) = PriorInit()
-
-function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    model::Model,
-    spl::Sampler;
-    initial_params::AbstractInitStrategy=init_strategy(spl),
-    kwargs...,
-)
-    # Generate the default varinfo (usually this just makes an empty VarInfo
-    # with NamedTuple of Metadata).
-    vi = default_varinfo(rng, model, spl)
-
-    # Fill it with initial parameters. Note that, if `ParamsInit` is used, the
-    # parameters provided must be in unlinked space (when inserted into the
-    # varinfo, they will be adjusted to match the linking status of the
-    # varinfo).
-    _, vi = init!!(rng, model, vi, initial_params)
-
-    # Call the actual function that does the first step.
-    return initialstep(rng, model, spl, vi; initial_params, kwargs...)
-end
+init_strategy(::AbstractMCMC.AbstractSampler) = PriorInit()
 
 """
     loadstate(data)
@@ -97,14 +41,4 @@ loadstate(data) = data
 
 Default type of the chain of posterior samples from `sampler`.
 """
-default_chain_type(::Sampler) = Any
-
-"""
-    initialstep(rng, model, sampler, varinfo; kwargs...)
-
-Perform the initial sampling step of the `sampler` for the `model`.
-
-The `varinfo` contains the initial samples, which can be provided by the user or
-sampled randomly.
-"""
-function initialstep end
+default_chain_type(::AbstractMCMC.AbstractSampler) = Any

test/sampler.jl

@@ -4,28 +4,23 @@
         abstract type OnlyInitAlg end
         struct OnlyInitAlgDefault <: OnlyInitAlg end
         struct OnlyInitAlgUniform <: OnlyInitAlg end
-        function DynamicPPL.initialstep(
-            rng::Random.AbstractRNG,
-            model::Model,
-            ::Sampler{<:OnlyInitAlg},
-            vi::AbstractVarInfo;
-            kwargs...,
+        function AbstractMCMC.step(
+            rng::Random.AbstractRNG, model::Model, ::OnlyInitAlg, kwargs...
         )
             return vi, nothing
         end
 
         # initial samplers
-        DynamicPPL.init_strategy(::Sampler{OnlyInitAlgUniform}) = UniformInit()
-        @test DynamicPPL.init_strategy(Sampler(OnlyInitAlgDefault())) == PriorInit()
+        DynamicPPL.init_strategy(::OnlyInitAlgUniform) = UniformInit()
+        @test DynamicPPL.init_strategy(OnlyInitAlgDefault()) == PriorInit()
 
-        for alg in (OnlyInitAlgDefault(), OnlyInitAlgUniform())
+        for sampler in (OnlyInitAlgDefault(), OnlyInitAlgUniform())
             # model with one variable: initialization p = 0.2
             @model function coinflip()
                 p ~ Beta(1, 1)
                 return 10 ~ Binomial(25, p)
             end
             model = coinflip()
-            sampler = Sampler(alg)
             lptrue = logpdf(Binomial(25, 0.2), 10)
             let inits = ParamsInit((; p=0.2))
                 chain = sample(model, sampler, 1; initial_params=inits, progress=false)