Improvements to VarNamedVector (#1098)

* Change VNV to use Dict rather than OrderedDict

* Change concretisation from map(identity, x) to a comprehension

* Improve tighten_types!! and loosen_types!!

* Fix use of set_transformed!!

* Fix push!! for VarInfos

* Change the default element types in VNV to be Union{}

* In untyped_vector_varinfo, don't rely on Metadata

* Code style

* Run formatter

* In VNV, use typejoin rather than promote_type

* Bump patch version to 0.38.4

Author: mhauru

HISTORY.md

@@ -1,5 +1,9 @@
 # DynamicPPL Changelog
 
+## 0.38.4
+
+Improve performance of VarNamedVector. It should now be very nearly on par with Metadata for all models we've benchmarked on.
+
 ## 0.38.3
 
 Add an implementation of `returned(::Model, ::AbstractDict{<:VarName})`.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.38.3"
+version = "0.38.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/api.md

@@ -414,7 +414,7 @@ DynamicPPL.reset!
 DynamicPPL.update!
 DynamicPPL.insert!
 DynamicPPL.loosen_types!!
-DynamicPPL.tighten_types
+DynamicPPL.tighten_types!!
 ```
 
 ```@docs

src/contexts/init.jl

@@ -180,7 +180,9 @@ function tilde_assume!!(
     end
     # Neither of these set the `trans` flag so we have to do it manually if
     # necessary.
-    insert_transformed_value && set_transformed!!(vi, true, vn)
+    if insert_transformed_value
+        vi = set_transformed!!(vi, true, vn)
+    end
     # `accumulate_assume!!` wants untransformed values as the second argument.
     vi = accumulate_assume!!(vi, x, logjac, vn, dist)
     # We always return the untransformed value here, as that will determine

src/debug_utils.jl

@@ -27,7 +27,7 @@ add_io_context(io::IO) = IOContext(io, :compact => true, :limit => true)
 show_varname(io::IO, varname::VarName) = print(io, varname)
 function show_varname(io::IO, varname::Array{<:VarName,N}) where {N}
     # Attempt to make the type concrete in case the symbol is shared.
-    return _show_varname(io, map(identity, varname))
+    return _show_varname(io, [vn for vn in varname])
 end
 function _show_varname(io::IO, varname::Array{<:VarName,N}) where {N}
     # Print the first and last element of the array.
@@ -407,7 +407,7 @@ julia> @model function demo_incorrect()
        end
 demo_incorrect (generic function with 2 methods)
 
-julia> # Notice that VarInfo(model_incorrect) evaluates the model, but doesn't actually 
+julia> # Notice that VarInfo(model_incorrect) evaluates the model, but doesn't actually
        # alert us to the issue of `x` being sampled twice.
        model = demo_incorrect(); varinfo = VarInfo(model);
 

src/logdensityfunction.jl

@@ -49,7 +49,7 @@ box:
 - [`getlogprior`](@ref): calculate the log prior in the model space, ignoring
   any effects of linking
 - [`getloglikelihood`](@ref): calculate the log likelihood (this is unaffected
-  by linking, since transforms are only applied to random variables) 
+  by linking, since transforms are only applied to random variables)
 
 !!! note
     By default, `LogDensityFunction` uses `getlogjoint_internal`, i.e., the
@@ -146,7 +146,7 @@ struct LogDensityFunction{
             is_supported(adtype) ||
                 @warn "The AD backend $adtype is not officially supported by DynamicPPL. Gradient calculations may still work, but compatibility is not guaranteed."
             # Get a set of dummy params to use for prep
-            x = map(identity, varinfo[:])
+            x = [val for val in varinfo[:]]
             if use_closure(adtype)
                 prep = DI.prepare_gradient(
                     LogDensityAt(model, getlogdensity, varinfo), adtype, x
@@ -282,7 +282,7 @@ function LogDensityProblems.logdensity_and_gradient(
 ) where {M,F,V,AD<:ADTypes.AbstractADType}
     f.prep === nothing &&
         error("Gradient preparation not available; this should not happen")
-    x = map(identity, x)  # Concretise type
+    x = [val for val in x]  # Concretise type
     # Make branching statically inferrable, i.e. type-stable (even if the two
     # branches happen to return different types)
     return if use_closure(f.adtype)

src/simple_varinfo.jl

@@ -484,6 +484,7 @@ function set_transformed!!(vi::SimpleOrThreadSafeSimple, trans::Bool, ::VarName)
             "Individual variables in SimpleVarInfo cannot have different `set_transformed` statuses.",
         )
     end
+    return vi
 end
 
 is_transformed(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation)

src/test_utils/ad.jl

@@ -197,7 +197,7 @@ Everything else is optional, and can be categorised into several groups:
 1. _How to specify the results to compare against._
 
    Once logp and its gradient has been calculated with the specified `adtype`,
-   it can optionally be tested for correctness. The exact way this is tested 
+   it can optionally be tested for correctness. The exact way this is tested
    is specified in the `test` parameter.
 
    There are several options for this:
@@ -260,7 +260,7 @@ function run_ad(
     if isnothing(params)
         params = varinfo[:]
     end
-    params = map(identity, params)  # Concretise
+    params = [p for p in params]  # Concretise
 
     # Calculate log-density and gradient with the backend of interest
     verbose && @info "Running AD on $(model.f) with $(adtype)\n"

src/varinfo.jl

@@ -315,7 +315,7 @@ function untyped_vector_varinfo(
     model::Model,
     init_strategy::AbstractInitStrategy=InitFromPrior(),
 )
-    return untyped_vector_varinfo(untyped_varinfo(rng, model, init_strategy))
+    return last(init!!(rng, model, VarInfo(VarNamedVector()), init_strategy))
 end
 function untyped_vector_varinfo(
     model::Model, init_strategy::AbstractInitStrategy=InitFromPrior()
@@ -789,18 +789,24 @@ function setval!(md::Metadata, val, vn::VarName)
     return md.vals[getrange(md, vn)] = tovec(val)
 end
 
+function set_transformed!!(vi::NTVarInfo, val::Bool, vn::VarName)
+    md = set_transformed!!(getmetadata(vi, vn), val, vn)
+    return Accessors.@set vi.metadata[getsym(vn)] = md
+end
+
 function set_transformed!!(vi::VarInfo, val::Bool, vn::VarName)
-    set_transformed!!(getmetadata(vi, vn), val, vn)
-    return vi
+    md = set_transformed!!(getmetadata(vi, vn), val, vn)
+    return VarInfo(md, vi.accs)
 end
+
 function set_transformed!!(metadata::Metadata, val::Bool, vn::VarName)
     metadata.is_transformed[getidx(metadata, vn)] = val
     return metadata
 end
 
 function set_transformed!!(vi::VarInfo, val::Bool)
     for vn in keys(vi)
-        set_transformed!!(vi, val, vn)
+        vi = set_transformed!!(vi, val, vn)
     end
 
     return vi
@@ -977,7 +983,7 @@ function filter_subsumed(filter_vns, filtered_vns)
 end
 
 @generated function _link!!(
-    ::NamedTuple{metadata_names}, vi, vns::NamedTuple{vns_names}
+    ::NamedTuple{metadata_names}, vi, varnames::NamedTuple{vns_names}
 ) where {metadata_names,vns_names}
     expr = Expr(:block)
     for f in metadata_names
@@ -988,7 +994,7 @@ end
             expr.args,
             quote
                 f_vns = vi.metadata.$f.vns
-                f_vns = filter_subsumed(vns.$f, f_vns)
+                f_vns = filter_subsumed(varnames.$f, f_vns)
                 if !isempty(f_vns)
                     if !is_transformed(vi, f_vns[1])
                         # Iterate over all `f_vns` and transform
@@ -1652,30 +1658,47 @@ end
 Push a new random variable `vn` with a sampled value `r` from a distribution `dist` to
 the `VarInfo` `vi`, mutating if it makes sense.
 """
-function BangBang.push!!(vi::VarInfo, vn::VarName, r, dist::Distribution)
-    if vi isa UntypedVarInfo
-        @assert ~(vn in keys(vi)) "[push!!] attempt to add an existing variable $(getsym(vn)) ($(vn)) to VarInfo (keys=$(keys(vi))) with dist=$dist"
-    elseif vi isa NTVarInfo
-        @assert ~(haskey(vi, vn)) "[push!!] attempt to add an existing variable $(getsym(vn)) ($(vn)) to NTVarInfo of syms $(syms(vi)) with dist=$dist"
-    end
+function BangBang.push!!(vi::VarInfo, vn::VarName, val, dist::Distribution)
+    @assert ~(vn in keys(vi)) "[push!!] attempt to add an existing variable $(getsym(vn)) ($(vn)) to VarInfo (keys=$(keys(vi))) with dist=$dist"
+    md = push!!(getmetadata(vi, vn), vn, val, dist)
+    return VarInfo(md, vi.accs)
+end
 
+function BangBang.push!!(vi::NTVarInfo, vn::VarName, val, dist::Distribution)
+    @assert ~(haskey(vi, vn)) "[push!!] attempt to add an existing variable $(getsym(vn)) ($(vn)) to NTVarInfo of syms $(syms(vi)) with dist=$dist"
     sym = getsym(vn)
-    if vi isa NTVarInfo && ~haskey(vi.metadata, sym)
+    meta = if ~haskey(vi.metadata, sym)
         # The NamedTuple doesn't have an entry for this variable, let's add one.
-        val = tovec(r)
-        md = Metadata(Dict(vn => 1), [vn], [1:length(val)], val, [dist], BitVector([false]))
-        vi = Accessors.@set vi.metadata[sym] = md
+        _new_submetadata(vi, vn, val, dist)
     else
-        meta = getmetadata(vi, vn)
-        push!(meta, vn, r, dist)
+        push!!(getmetadata(vi, vn), vn, val, dist)
     end
-
+    vi = Accessors.@set vi.metadata[sym] = meta
     return vi
 end
 
-function Base.push!(vi::UntypedVectorVarInfo, vn::VarName, val, args...)
-    push!(getmetadata(vi, vn), vn, val, args...)
-    return vi
+"""
+    _new_submetadata(vi::VarInfo{NamedTuple{Names,SubMetas}}, args...) where {Names,SubMetas}
+
+Create a new sub-metadata for an NTVarInfo. The type is chosen by the types of existing
+SubMetas.
+"""
+@generated function _new_submetadata(
+    vi::VarInfo{NamedTuple{Names,SubMetas}}, vn, r, dist
+) where {Names,SubMetas}
+    has_vnv = any(s -> s <: VarNamedVector, SubMetas.parameters)
+    return if has_vnv
+        :(return _new_vnv_submetadata(vn, r, dist))
+    else
+        :(return _new_metadata_submetadata(vn, r, dist))
+    end
+end
+
+_new_vnv_submetadata(vn, r, _) = VarNamedVector([vn], [r])
+
+function _new_metadata_submetadata(vn, r, dist)
+    val = tovec(r)
+    return Metadata(Dict(vn => 1), [vn], [1:length(val)], val, [dist], BitVector([false]))
 end
 
 function Base.push!(vi::UntypedVectorVarInfo, pair::Pair, args...)
@@ -1700,6 +1723,11 @@ function Base.push!(meta::Metadata, vn, r, dist)
     return meta
 end
 
+function BangBang.push!!(meta::Metadata, vn, r, dist)
+    push!(meta, vn, r, dist)
+    return meta
+end
+
 function Base.delete!(vi::VarInfo, vn::VarName)
     delete!(getmetadata(vi, vn), vn)
     return vi