[Merged by Bors] - Fix for #371

[torfjelde]

This PR adds a method called resolve_varnames(varname, dist) and adds an additional generated variable for each ~ which now holds the RHS of ~.

It does address #371 but uncertain if this is the best way, so wouldn't recommend merging this just yet. But putting it here so we can colab on it.

[yebai]

It looks sensible to me!

[devmotion]

uncertain if this is the best way

Another solution would be to remove NamedDist 😛 Maybe someone can remind me, what's the main motivation for NamedDist and what problem does it fix?

[torfjelde]

Another solution would be to remove NamedDist stuck_out_tongue Maybe someone can remind me, what's the main motivation for NamedDist and what problem does it fix?

Allows one to use a different variable-name, e.g. x ~ NamedDist(..., :z) means that this will now show up as z in VarInfo (and the resulting chain, etc.).
Never use it myself, but there are people using it 🤷

[bgroenks96]

@devmotion Is there another way besides NamedDist to programmatically define variable names in a Turing model, without using internal APIs like VarInfo?

I am surprised that this is such an uncommon use-case. For large systems with dozens or more parameters, it's a pain to write and maintain Turing models entirely by hand. It's only natural to want to automate this and build models programmatically.

[devmotion]

Sure, I know this use case, eg from DiffEqBayes. But to me NamedDist seems like a hacky workaround but not a "nice" solution to this problem. Maybe we should just support VarNames on the LHS of ~? Or maybe (an improved version of) the interpolation support is sufficient?

[yebai]

Maybe we should just support VarNames on the LHS of ~? Or maybe (an improved version of) the interpolation support is sufficient?

It sounds like a good idea!

[yebai]

For large systems with dozens or more parameters, it's a pain to write and maintain Turing models entirely by hand. It's only natural to want to automate this and build models programmatically.

I’m curious; what models are you working with? Some examples would help us to improve DynamicPPL / AbstractPPL.

[bgroenks96]

Maybe we should just support VarNames on the LHS of ~? Or maybe (an improved version of) the interpolation support is sufficient?

That would also be good, maybe even preferable. I actually only learned about NamedDist from @torfjelde , and I was doing something far hackier before via the internal APIs 😅

I’m curious; what models are you working with? Some examples would help us to improve DynamicPPL / AbstractPPL.

I am working on a permafrost model (at the moment basically a 1-D heat PDE with phase change) which uses a kind of "modular" design philosophy; i.e. one builds a model by stacking various components that may include one or more physical processes and each component may have an arbitrary number of parameters associated with it. I am interested in solving the Bayesian inverse problem (and also obtaining predictive uncertainty) for a subset of these parameters.

I actually don't use Turing at the moment for posterior sampling because I found HMC and NUTS to be far too costly for a model with such an expensive likelihood function (I do fairly long simulations of the PDE), and the particle based samplers like SMC in AdvancedPS don't really support proper parallelization. Instead, I use Turing to build and sample from the prior and then I use Ensemble Kalman Sampling for (approximte) posterior inference. In the future, I hope to be able to use Turing for the whole thing :)

I can give you a code snippet of how I specify the Turing model at present:

function makepriors(ps::CryoGridParams)
    return map(ps[:component], ps[:fieldname], ps[:val]) do comp, name, value
        prior(comp, Val{name}(), value)
    end
end
# `NamedDist` comes from DynamicPPL.jl. (this part comes from Tor)
@model priormodel(shared::NamedTuple, name::Symbol, prior::Distribution) = x ~ NamedDist(prior, name)
@model priormodel(shared::NamedTuple, name::Symbol, nt::NamedTuple) = @submodel prefix=$name priormodel(shared, nt)
@model priormodel(shared::NamedTuple, name::Symbol, val::Symbol) = shared[val]
@model priormodel(shared::NamedTuple, name::Symbol, val) = val
@model function priormodel(shared, priors::NamedTuple)
    # NOTE: The below will update `__varinfo__` for every `k`.
    vals = map(keys(priors)) do k
        @submodel priormodel(shared, k, priors[k])
    end
    return NamedTuple{keys(priors)}(vals)
end
@model function cryogridparams(ps::CryoGridParams, sharedpriors::NamedTuple=(;), ::Type{T}=Float64) where T
    shared = @submodel priormodel((;), sharedpriors)
    pvals = @submodel priormodel(shared, (; map(Pair, ps[:pid], ps[:prior])...))
    # create copy of parameter array with compatible type (for autodiff)
    p = copyto!(similar(collect(values(ps)), T), pvals)
    return p
end

where prior is a method that has dispatches for various parameters in the model and CryoGridParams wraps a tabular interface to the parameters provided by ModelParameters.jl.

I hope that helps!

Edit: I should note that I just changed this very recently; the parameters used to be organized into a hierarchy/tree and then priormodel built submodels recursively, hence the priormodel(shared::NamedTuple, name::Symbol, nt::NamedTuple) dispatch. I found this to have some limitations, so now it just uses a table-like structure, and I kind of did the bare minimum to make this code work again... so yeah, that's why it's a bit weird 😅.

[yebai]

One main advantage of NamedDist is it allows interpolation in VarName at runtime (via an RHS named distribution). TuringLang/AbstractPPL.jl#54 can probably replace this PR since users can create VarNames with interpolation from the LHS.

@phipsgabler can probably confirm or falsify this claim

[phipsgabler]

I think VarName interpolation is not sufficient for replacing that. The intent of it was to make manual fiddling with @varname easier, not to have much effect on DPPL. The main work to be done would be in the @model macro, since it has control over interpolation (and we don't use @varname inside of it at all). However, the same kind of logic can maybe be reduce in @model to expand

$name[i]

to

(VarName){name}((Setfield.compose)((Setfield.IndexLens)((i,))))

---

OK, actually, the LHS is transformed as

$vn = $(AbstractPPL.drop_escape(varname(left)))

so it might actually work... but I haven't tried it!

[yebai]

@torfjelde a gentle reminder on this

[torfjelde]

As @phipsgabler originally said, @varname interpolation isn't sufficient.

This is because we explicitly "unwrap" the interpolated expressions in the model macro:

DynamicPPL.jl/src/compiler.jl

Lines 321 to 322 in d222316

	# Do not touch interpolated expressions
	expr.head === :$ && return expr.args[1]

As a result, even if @varname $name.a works, if you do $name ~ Normal(), the LHS will be parsed as VarName{:name}() and subsequent LHSs in the expanded code will just be name.

If we keep the $ around rather than "unwrapping" it as we do in the lines referenced above, we correctly get VarName{name}() but now subsequent usages of left will be a Expr(:$, name) which aren't interpolated. I've forgotten exactly what's the reason for this, but it probably has something to do with how we escape everything which in turn messes up the macro-hygiene.

[yebai]

bors try

[bors]

try

Build succeeded:

• test (1, ubuntu-latest, x64, 1)
• test (1, ubuntu-latest, x64, 2)
• test (1, ubuntu-latest, x86, 2)
• test (1.3, ubuntu-latest, x64, 1)
• test (1.3, ubuntu-latest, x64, 2)
• test (1, macOS-latest, x64, 2)
• test (1, windows-latest, x64, 2)
• Turing.jl
• format
• docs

[bgroenks96]

It would be great to get this merged, if everything is working ok! I currently have my project pointed at this PR branch.

[yebai]

@torfjelde pls, feel free to go ahead and merge it if you think it is ready.

[torfjelde]

bors r+

[torfjelde]

bors cancel

[bors]

Canceled.

[torfjelde]

bors r+

[bors]

Pull request successfully merged into master.

Build succeeded:

• docs
• format
• test (1.3, ubuntu-latest, x64, 1)
• test (1.3, ubuntu-latest, x64, 2)
• test (1, macOS-latest, x64, 2)
• test (1, ubuntu-latest, x64, 1)
• test (1, ubuntu-latest, x64, 2)
• test (1, ubuntu-latest, x86, 2)
• test (1, windows-latest, x64, 2)
• Turing.jl