Add hasvalue for (some) distributions

Author: penelopeysm

Project.toml

@@ -13,10 +13,17 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
+[weakdeps]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+
+[extensions]
+AbstractPPLDistributionsExt = ["Distributions"]
+
 [compat]
 AbstractMCMC = "2, 3, 4, 5"
 Accessors = "0.1"
 DensityInterface = "0.4"
+Distributions = "0.25"
 JSON = "0.19 - 0.21"
 Random = "1.6"
 StatsBase = "0.32, 0.33, 0.34"

ext/AbstractPPLDistributionsExt.jl

@@ -0,0 +1,51 @@
+module AbstractPPLDistributionsExt
+
+if isdefined(Base, :get_extension)
+    using AbstractPPL: AbstractPPL, VarName, Accessors
+    using Distributions: Distributions
+else
+    using ..AbstractPPL: AbstractPPL, VarName, Accessors
+    using ..Distributions: Distributions
+end
+
+# TODO(penelopeysm): Figure out tuple / namedtuple distributions, and LKJCholesky (grr)
+function AbstractPPL.hasvalue(
+    vals::AbstractDict, vn::VarName, dist::Distributions.Distribution
+)
+    @warn "`hasvalue(vals, vn, dist)` is not implemented for $(typeof(dist)); falling back to `hasvalue(vals, vn)`."
+    return AbstractPPL.hasvalue(vals, vn)
+end
+function AbstractPPL.hasvalue(
+    vals::AbstractDict, vn::VarName, ::Distributions.UnivariateDistribution
+)
+    return AbstractPPL.hasvalue(vals, vn)
+end
+function AbstractPPL.hasvalue(
+    vals::AbstractDict{<:VarName},
+    vn::VarName{sym},
+    dist::Union{Distributions.MultivariateDistribution,Distributions.MatrixDistribution},
+) where {sym}
+    # If `vn` is present as-is, then we are good
+    AbstractPPL.hasvalue(vals, vn) && return true
+    # If not, then we need to check inside `vals` to see if a subset of
+    # `vals` is enough to reconstruct `vn`. For example, if `vals` contains
+    # `x[1]` and `x[2]`, and `dist` is `MvNormal(zeros(2), I)`, then we
+    # can reconstruct `x`. If `dist` is `MvNormal(zeros(3), I)`, then we
+    # can't.
+    # To do this, we get the size of the distribution and iterate over all
+    # possible indices. If every index can be found in `subsumed_keys`, then we
+    # can return true.
+    sz = size(dist)
+    for idx in Iterators.product(map(Base.OneTo, sz)...)
+        new_optic = if AbstractPPL.getoptic(vn) === identity
+            Accessors.IndexLens(idx)
+        else
+            Accessors.IndexLens(idx) ∘ AbstractPPL.getoptic(vn)
+        end
+        new_vn = VarName{sym}(new_optic)
+        AbstractPPL.hasvalue(vals, new_vn) || return false
+    end
+    return true
+end
+
+end

src/hasvalue.jl

@@ -57,7 +57,7 @@ julia> getvalue(vals, @varname(x[1])) # different from `getindex`
 1.0
 
 julia> getvalue(vals, @varname(x[2]))
-ERROR: BoundsError: attempt to access 1-element Vector{Float64} at index [2]
+ERROR: getvalue: x[2] was not found in the values provided
 [...]
 ```
 
@@ -74,7 +74,7 @@ julia> getvalue(vals, @varname(x[1])) # different from `getindex`
 1.0
 
 julia> getvalue(vals, @varname(x[2]))
-ERROR: BoundsError: attempt to access 1-element Vector{Float64} at index [2]
+ERROR: getvalue: x[2] was not found in the values provided
 [...]
 ```
 
@@ -91,16 +91,53 @@ julia> getvalue(vals, @varname(x[1][1])) # different from `getindex`
 1.0
 
 julia> getvalue(vals, @varname(x[1][2]))
-ERROR: BoundsError: attempt to access 1-element Vector{Float64} at index [2]
+ERROR: getvalue: x[1][2] was not found in the values provided
 [...]
 
 julia> getvalue(vals, @varname(x[2][1]))
-ERROR: KeyError: key x[2][1] not found
+ERROR: getvalue: x[2][1] was not found in the values provided
 [...]
 ```
 """
-getvalue(vals::NamedTuple, vn::VarName) = get(vals, vn)
-getvalue(vals::AbstractDict, vn::VarName) = nested_getindex(vals, vn)
+function getvalue(vals::NamedTuple, vn::VarName{sym}) where {sym}
+    optic = getoptic(vn)
+    if haskey(vals, sym) && canview(optic, getproperty(vals, sym))
+        return optic(vals[sym])
+    else
+        error("getvalue: $(vn) was not found in the values provided")
+    end
+end
+# For the Dict case, it is more complicated. There are two cases:
+# 1. `vn` itself is already a key of `vals` (the easy case)
+# 2. `vn` is not a key of `vals`, but some parent of `vn` is a key of `vals`
+#    (the harder case). For example, if `vn` is `x[1][2]`, then we need to
+#    check if either `x` or `x[1]` is a key of `vals`, and if so, whether
+#    we can index into the corresponding value.
+function getvalue(vals::AbstractDict{<:VarName}, vn::VarName{sym}) where {sym}
+    # First we check if `vn` is present as is.
+    haskey(vals, vn) && return vals[vn]
+
+    # Otherwise, we start by testing the bare `vn` (e.g., if `vn` is `x[1][2]`,
+    # we start by checking if `x` is present). We will then keep adding optics
+    # to `test_optic`, either until we find a key that is present, or until we
+    # run out of optics to test (which is determined by _inner(test_optic) ==
+    # identity).
+    test_vn = VarName{sym}()
+    test_optic = getoptic(vn)
+
+    while _inner(test_optic) != identity
+        if haskey(vals, test_vn) && canview(test_optic, vals[test_vn])
+            return test_optic(vals[test_vn])
+        else
+            # Move the innermost optic into test_vn
+            test_optic_outer = _outer(test_optic)
+            test_optic_inner = _inner(test_optic)
+            test_vn = VarName{sym}(test_optic_inner ∘ getoptic(test_vn))
+            test_optic = test_optic_outer
+        end
+    end
+    return error("getvalue: $(vn) was not found in the values provided")
+end
 
 """
     hasvalue(vals::NamedTuple, vn::VarName)
@@ -168,13 +205,6 @@ false
 function hasvalue(vals::NamedTuple, vn::VarName{sym}) where {sym}
     return haskey(vals, sym) && canview(getoptic(vn), getproperty(vals, sym))
 end
-
-# For the Dict case, it is more complicated. There are two cases:
-# 1. `vn` itself is already a key of `vals` (the easy case)
-# 2. `vn` is not a key of `vals`, but some parent of `vn` is a key of `vals`
-#    (the harder case). For example, if `vn` is `x[1][2]`, then we need to
-#    check if either `x` or `x[1]` is a key of `vals`, and if so, whether
-#    we can index into the corresponding value.
 function hasvalue(vals::AbstractDict{<:VarName}, vn::VarName{sym}) where {sym}
     # First we check if `vn` is present as is.
     haskey(vals, vn) && return true
@@ -186,12 +216,8 @@ function hasvalue(vals::AbstractDict{<:VarName}, vn::VarName{sym}) where {sym}
     # identity).
     test_vn = VarName{sym}()
     test_optic = getoptic(vn)
-    
+
     while _inner(test_optic) != identity
-        @show test_vn, test_optic
-        if haskey(vals, test_vn)
-            @show canview(test_optic, vals[test_vn])
-        end
         if haskey(vals, test_vn) && canview(test_optic, vals[test_vn])
             return true
         else
@@ -204,68 +230,3 @@ function hasvalue(vals::AbstractDict{<:VarName}, vn::VarName{sym}) where {sym}
     end
     return false
 end
-# TODO(penelopeysm): Figure out tuple / namedtuple distributions, and LKJCholesky (grr)
-# function hasvalue(vals::AbstractDict, vn::VarName, dist::Distribution)
-#     @warn "`hasvalue(vals, vn, dist)` is not implemented for $(typeof(dist)); falling back to `hasvalue(vals, vn)`."
-#     return hasvalue(vals, vn)
-# end
-# hasvalue(vals::AbstractDict, vn::VarName, ::UnivariateDistribution) = hasvalue(vals, vn)
-# function hasvalue(
-#     vals::AbstractDict{<:VarName},
-#     vn::VarName{sym},
-#     dist::Union{MultivariateDistribution,MatrixDistribution},
-# ) where {sym}
-#     # If `vn` is present as-is, then we are good
-#     hasvalue(vals, vn) && return true
-#     # If not, then we need to check inside `vals` to see if a subset of
-#     # `vals` is enough to reconstruct `vn`. For example, if `vals` contains
-#     # `x[1]` and `x[2]`, and `dist` is `MvNormal(zeros(2), I)`, then we
-#     # can reconstruct `x`. If `dist` is `MvNormal(zeros(3), I)`, then we
-#     # can't.
-#     # To do this, we get the size of the distribution and iterate over all
-#     # possible indices. If every index can be found in `subsumed_keys`, then we
-#     # can return true.
-#     sz = size(dist)
-#     for idx in Iterators.product(map(Base.OneTo, sz)...)
-#         new_optic = if getoptic(vn) === identity
-#             Accessors.IndexLens(idx)
-#         else
-#             Accessors.IndexLens(idx) ∘ getoptic(vn)
-#         end
-#         new_vn = VarName{sym}(new_optic)
-#         hasvalue(vals, new_vn) || return false
-#     end
-#     return true
-# end
-
-# """
-#     nested_getindex(values::AbstractDict, vn::VarName)
-#
-# Return value corresponding to `vn` in `values` by also looking
-# in the the actual values of the dict.
-# """
-# function nested_getindex(values::AbstractDict, vn::VarName)
-#     maybeval = get(values, vn, nothing)
-#     if maybeval !== nothing
-#         return maybeval
-#     end
-#
-#     # Split the optic into the key / `parent` and the extraction optic / `child`.
-#     parent, child, issuccess = splitoptic(getoptic(vn)) do optic
-#         o = optic === nothing ? identity : optic
-#         haskey(values, VarName(vn, o))
-#     end
-#     # When combined with `VarInfo`, `nothing` is equivalent to `identity`.
-#     keyoptic = parent === nothing ? identity : parent
-#
-#     # If we found a valid split, then we can extract the value.
-#     if !issuccess
-#         # At this point we just throw an error since the key could not be found.
-#         throw(KeyError(vn))
-#     end
-#
-#     # TODO: Should we also check that we `canview` the extracted `value`
-#     # rather than just let it fail upon `get` call?
-#     value = values[VarName(vn, keyoptic)]
-#     return child(value)
-# end

test/Project.toml

@@ -1,7 +1,9 @@
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/hasvalue.jl

@@ -1,17 +1,28 @@
-@testset "hasvalue" begin
+@testset "base getvalue + hasvalue" begin
     @testset "NamedTuple" begin
         nt = (a=[1], b=2, c=(x=3,), d=[1.0 0.5; 0.5 1.0])
         @test hasvalue(nt, @varname(a))
+        @test getvalue(nt, @varname(a)) == [1]
         @test hasvalue(nt, @varname(a[1]))
+        @test getvalue(nt, @varname(a[1])) == 1
         @test hasvalue(nt, @varname(b))
+        @test getvalue(nt, @varname(b)) == 2
         @test hasvalue(nt, @varname(c))
+        @test getvalue(nt, @varname(c)) == (x=3,)
         @test hasvalue(nt, @varname(c.x))
+        @test getvalue(nt, @varname(c.x)) == 3
         @test hasvalue(nt, @varname(d))
+        @test getvalue(nt, @varname(d)) == [1.0 0.5; 0.5 1.0]
         @test hasvalue(nt, @varname(d[1, 1]))
+        @test getvalue(nt, @varname(d[1, 1])) == 1.0
         @test hasvalue(nt, @varname(d[1, 2]))
+        @test getvalue(nt, @varname(d[1, 2])) == 0.5
         @test hasvalue(nt, @varname(d[2, 1]))
+        @test getvalue(nt, @varname(d[2, 1])) == 0.5
         @test hasvalue(nt, @varname(d[2, 2]))
+        @test getvalue(nt, @varname(d[2, 2])) == 1.0
         @test hasvalue(nt, @varname(d[3]))  # linear indexing works....
+        @test getvalue(nt, @varname(d[3])) == 0.5
         @test !hasvalue(nt, @varname(nope))
         @test !hasvalue(nt, @varname(a[2]))
         @test !hasvalue(nt, @varname(a[1][1]))
@@ -23,40 +34,77 @@
 
     @testset "Dict" begin
         # same tests as above
-        d = Dict(@varname(a) => [1],
+        d = Dict(
+            @varname(a) => [1],
             @varname(b) => 2,
             @varname(c) => (x=3,),
-            @varname(d) => [1.0 0.5; 0.5 1.0])
+            @varname(d) => [1.0 0.5; 0.5 1.0],
+        )
         @test hasvalue(d, @varname(a))
+        @test getvalue(d, @varname(a)) == [1]
         @test hasvalue(d, @varname(a[1]))
+        @test getvalue(d, @varname(a[1])) == 1
         @test hasvalue(d, @varname(b))
+        @test getvalue(d, @varname(b)) == 2
         @test hasvalue(d, @varname(c))
+        @test getvalue(d, @varname(c)) == (x=3,)
         @test hasvalue(d, @varname(c.x))
+        @test getvalue(d, @varname(c.x)) == 3
         @test hasvalue(d, @varname(d))
+        @test getvalue(d, @varname(d)) == [1.0 0.5; 0.5 1.0]
         @test hasvalue(d, @varname(d[1, 1]))
+        @test getvalue(d, @varname(d[1, 1])) == 1.0
         @test hasvalue(d, @varname(d[1, 2]))
+        @test getvalue(d, @varname(d[1, 2])) == 0.5
         @test hasvalue(d, @varname(d[2, 1]))
+        @test getvalue(d, @varname(d[2, 1])) == 0.5
         @test hasvalue(d, @varname(d[2, 2]))
-        @test hasvalue(d, @varname(d[3])) # linear indexing works....
+        @test getvalue(d, @varname(d[2, 2])) == 1.0
+        @test hasvalue(d, @varname(d[3]))  # linear indexing works....
+        @test getvalue(d, @varname(d[3])) == 0.5
         @test !hasvalue(d, @varname(nope))
         @test !hasvalue(d, @varname(a[2]))
         @test !hasvalue(d, @varname(a[1][1]))
         @test !hasvalue(d, @varname(c.x[1]))
         @test !hasvalue(d, @varname(c.y))
         @test !hasvalue(d, @varname(d[1, 3]))
-        @test !hasvalue(d, @varname(d[3]))
 
-        # extra ones since Dict can have weird key
-        d = Dict(@varname(a[1]) => [1.0, 2.0],
-                 @varname(b.x) => [3.0])
+        # extra ones since Dict can have weird keys
+        d = Dict(
+            @varname(a[1]) => [1.0, 2.0],
+            @varname(b.x) => [3.0],
+            @varname(c[2]) => (a=4.0, b=5.0),
+        )
         @test hasvalue(d, @varname(a[1]))
+        @test getvalue(d, @varname(a[1])) == [1.0, 2.0]
         @test hasvalue(d, @varname(a[1][1]))
+        @test getvalue(d, @varname(a[1][1])) == 1.0
         @test hasvalue(d, @varname(a[1][2]))
+        @test getvalue(d, @varname(a[1][2])) == 2.0
         @test hasvalue(d, @varname(b.x))
+        @test getvalue(d, @varname(b.x)) == [3.0]
         @test hasvalue(d, @varname(b.x[1]))
+        @test getvalue(d, @varname(b.x[1])) == 3.0
+        @test hasvalue(d, @varname(c[2]))
+        @test getvalue(d, @varname(c[2])) == (a=4.0, b=5.0)
+        @test hasvalue(d, @varname(c[2].a))
+        @test getvalue(d, @varname(c[2].a)) == 4.0
+        @test hasvalue(d, @varname(c[2].b))
+        @test getvalue(d, @varname(c[2].b)) == 5.0
         @test !hasvalue(d, @varname(a))
         @test !hasvalue(d, @varname(a[2]))
         @test !hasvalue(d, @varname(b.y))
         @test !hasvalue(d, @varname(b.x[2]))
+        @test !hasvalue(d, @varname(c[1]))
+        @test !hasvalue(d, @varname(c[2].x))
     end
 end
+
+@testset "with Distributions: getvalue + hasvalue" begin
+    using Distributions
+    using LinearAlgebra
+
+    d = Dict(@varname(x[1]) => 1.0, @varname(x[2]) => 2.0)
+    @test hasvalue(d, @varname(x), MvNormal(zeros(2), I))
+    @test !hasvalue(d, @varname(x), MvNormal(zeros(3), I))
+end