@@ -4,36 +4,98 @@ function _check_model(model::DynamicPPL.Model)
44    new_model =  DynamicPPL. setleafcontext (model, DynamicPPL. InitContext ())
55    return  DynamicPPL. check_model (new_model, VarInfo (); error_on_failure= true )
66end 
7- function  _check_model (model:: DynamicPPL.Model , alg :: InferenceAlgorithm )
7+ function  _check_model (model:: DynamicPPL.Model , :: AbstractSampler )
88    return  _check_model (model)
99end 
1010
11+ """ 
12+     Turing.Inference.init_strategy(spl::AbstractSampler) 
13+ 
14+ Get the default initialization strategy for a given sampler `spl`, i.e. how initial 
15+ parameters for sampling are chosen if not specified by the user. By default, this is 
16+ `InitFromPrior()`, which samples initial parameters from the prior distribution. 
17+ """ 
18+ init_strategy (:: AbstractSampler ) =  DynamicPPL. InitFromPrior ()
19+ 
20+ """ 
21+     _convert_initial_params(initial_params) 
22+ 
23+ Convert `initial_params` to a `DynamicPPl.AbstractInitStrategy` if it is not already one, or 
24+ throw a useful error message. 
25+ """ 
26+ _convert_initial_params (initial_params:: DynamicPPL.AbstractInitStrategy ) =  initial_params
27+ function  _convert_initial_params (nt:: NamedTuple )
28+     @info  " Using a NamedTuple for `initial_params` will be deprecated in a future release. Please use `InitFromParams(namedtuple)` instead." 
29+     return  InitFromParams (nt)
30+ end 
31+ function  _convert_initial_params (d:: AbstractDict{<:VarName} )
32+     @info  " Using a Dict for `initial_params` will be deprecated in a future release. Please use `InitFromParams(dict)` instead." 
33+     return  InitFromParams (d)
34+ end 
35+ function  _convert_initial_params (:: AbstractVector{<:Real} )
36+     errmsg =  " `initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or ideally a `DynamicPPL.AbstractInitStrategy`. Using a vector of parameters for `initial_params` is no longer supported. Please see https://turinglang.org/docs/usage/sampling-options/#specifying-initial-parameters for details on how to update your code." 
37+     throw (ArgumentError (errmsg))
38+ end 
39+ function  _convert_initial_params (@nospecialize (_:: Any ))
40+     errmsg =  " `initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or a `DynamicPPL.AbstractInitStrategy`." 
41+     throw (ArgumentError (errmsg))
42+ end 
43+ 
44+ """ 
45+     default_varinfo(rng, model, sampler) 
46+ 
47+ Return a default varinfo object for the given `model` and `sampler`. 
48+ The default method for this returns a NTVarInfo (i.e. 'typed varinfo'). 
49+ """ 
50+ function  default_varinfo (
51+     rng:: Random.AbstractRNG , model:: DynamicPPL.Model , :: AbstractSampler 
52+ )
53+     #  Note that in `AbstractMCMC.step`, the values in the varinfo returned here are
54+     #  immediately overwritten by a subsequent call to `init!!`. The reason why we
55+     #  _do_ create a varinfo with parameters here (as opposed to simply returning
56+     #  an empty `typed_varinfo(VarInfo())`) is to avoid issues where pushing to an empty
57+     #  typed VarInfo would fail. This can happen if two VarNames have different types
58+     #  but share the same symbol (e.g. `x.a` and `x.b`).
59+     #  TODO (mhauru) Fix push!! to work with arbitrary lens types, and then remove the arguments
60+     #  and return an empty VarInfo instead.
61+     return  DynamicPPL. typed_varinfo (VarInfo (rng, model))
62+ end 
63+ 
1164# ########################################
1265#  Default definitions for the interface #
1366# ########################################
1467
1568function  AbstractMCMC. sample (
16-     model:: AbstractModel , alg :: InferenceAlgorithm , N:: Integer ; kwargs... 
69+     model:: DynamicPPL.Model , spl :: AbstractSampler , N:: Integer ; kwargs... 
1770)
18-     return  AbstractMCMC. sample (Random. default_rng (), model, alg , N; kwargs... )
71+     return  AbstractMCMC. sample (Random. default_rng (), model, spl , N; kwargs... )
1972end 
2073
2174function  AbstractMCMC. sample (
2275    rng:: AbstractRNG ,
23-     model:: AbstractModel ,
24-     alg :: InferenceAlgorithm ,
76+     model:: DynamicPPL.Model ,
77+     spl :: AbstractSampler ,
2578    N:: Integer ;
79+     initial_params= init_strategy (spl),
2680    check_model:: Bool = true ,
2781    chain_type= DEFAULT_CHAIN_TYPE,
2882    kwargs... ,
2983)
30-     check_model &&  _check_model (model, alg)
31-     return  AbstractMCMC. sample (rng, model, Sampler (alg), N; chain_type, kwargs... )
84+     check_model &&  _check_model (model, spl)
85+     return  AbstractMCMC. sample (
86+         rng,
87+         model,
88+         spl,
89+         N;
90+         initial_params= _convert_initial_params (initial_params),
91+         chain_type,
92+         kwargs... ,
93+     )
3294end 
3395
3496function  AbstractMCMC. sample (
35-     model:: AbstractModel ,
36-     alg:: InferenceAlgorithm ,
97+     model:: DynamicPPL.Model ,
98+     alg:: AbstractSampler ,
3799    ensemble:: AbstractMCMC.AbstractMCMCEnsemble ,
38100    N:: Integer ,
39101    n_chains:: Integer ;
@@ -45,18 +107,53 @@ function AbstractMCMC.sample(
45107end 
46108
47109function  AbstractMCMC. sample (
48-     rng:: AbstractRNG ,
49-     model:: AbstractModel ,
50-     alg :: InferenceAlgorithm ,
110+     rng:: Random. AbstractRNG
111+     model:: DynamicPPL.Model ,
112+     spl :: AbstractSampler ,
51113    ensemble:: AbstractMCMC.AbstractMCMCEnsemble ,
52114    N:: Integer ,
53115    n_chains:: Integer ;
54116    chain_type= DEFAULT_CHAIN_TYPE,
55117    check_model:: Bool = true ,
118+     initial_params= init_strategy (spl),
56119    kwargs... ,
57120)
58-     check_model &&  _check_model (model, alg )
121+     check_model &&  _check_model (model, spl )
59122    return  AbstractMCMC. sample (
60-         rng, model, Sampler (alg), ensemble, N, n_chains; chain_type, kwargs... 
123+         rng,
124+         model,
125+         spl,
126+         ensemble,
127+         N,
128+         n_chains;
129+         chain_type,
130+         initial_params= map (_convert_initial_params, initial_params),
131+         kwargs... ,
61132    )
62133end 
134+ 
135+ loadstate (c:: MCMCChains.Chains ) =  c. info. samplerstate
136+ 
137+ #  TODO (penelopeysm): Remove initialstep and generalise MCMC sampling procedures
138+ function  initialstep end 
139+ 
140+ function  AbstractMCMC. step (
141+     rng:: Random.AbstractRNG ,
142+     model:: DynamicPPL.Model ,
143+     spl:: AbstractSampler ;
144+     initial_params,
145+     kwargs... ,
146+ )
147+     #  Generate the default varinfo. Note that any parameters inside this varinfo
148+     #  will be immediately overwritten by the next call to `init!!`.
149+     vi =  default_varinfo (rng, model, spl)
150+ 
151+     #  Fill it with initial parameters. Note that, if `InitFromParams` is used, the
152+     #  parameters provided must be in unlinked space (when inserted into the
153+     #  varinfo, they will be adjusted to match the linking status of the
154+     #  varinfo).
155+     _, vi =  DynamicPPL. init!! (rng, model, vi, initial_params)
156+ 
157+     #  Call the actual function that does the first step.
158+     return  initialstep (rng, model, spl, vi; initial_params, kwargs... )
159+ end 
0 commit comments