Allow empty subsets of VarInfos (#692)

* Allow empty subsets of VarInfos * Run JuliaFormatter Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
TuringLang · Oct 17, 2024 · 54691bf · 54691bf · mhauru · Oct 18, 2024
1 parent 27ba772
commit 54691bf
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 17 deletions.
diff --git a/src/simple_varinfo.jl b/src/simple_varinfo.jl
@@ -439,22 +439,17 @@ function subset(varinfo::SimpleVarInfo, vns::AbstractVector{<:VarName})
     return Accessors.@set varinfo.values = _subset(varinfo.values, vns)
 end
 
-function _subset(x::AbstractDict, vns)
+function _subset(x::AbstractDict, vns::AbstractVector{VN}) where {VN<:VarName}
     vns_present = collect(keys(x))
-    vns_found = mapreduce(vcat, vns) do vn
+    vns_found = mapreduce(vcat, vns; init=VN[]) do vn
         return filter(Base.Fix1(subsumes, vn), vns_present)
     end
-
-    # NOTE: This `vns` to be subsume varnames explicitly present in `x`.
+    C = ConstructionBase.constructorof(typeof(x))
     if isempty(vns_found)
-        throw(
-            ArgumentError(
-                "Cannot subset `AbstractDict` with `VarName` which does not subsume any keys.",
-            ),
-        )
+        return C()
+    else
+        return C(vn => x[vn] for vn in vns_found)
     end
-    C = ConstructionBase.constructorof(typeof(x))
-    return C(vn => x[vn] for vn in vns_found)
 end
 
 function _subset(x::NamedTuple, vns)

diff --git a/src/varinfo.jl b/src/varinfo.jl
@@ -368,20 +368,24 @@ function subset(varinfo::TypedVarInfo, vns::AbstractVector{<:VarName})
     )
 end
 
-function subset(metadata::Metadata, vns_given::AbstractVector{<:VarName})
+function subset(metadata::Metadata, vns_given::AbstractVector{VN}) where {VN<:VarName}
     # TODO: Should we error if `vns` contains a variable that is not in `metadata`?
     # For each `vn` in `vns`, get the variables subsumed by `vn`.
-    vns = mapreduce(vcat, vns_given) do vn
+    vns = mapreduce(vcat, vns_given; init=VN[]) do vn
         filter(Base.Fix1(subsumes, vn), metadata.vns)
     end
     indices_for_vns = map(Base.Fix1(getindex, metadata.idcs), vns)
-    indices = Dict(vn => i for (i, vn) in enumerate(vns))
+    indices = if isempty(vns)
+        Dict{VarName,Int}()
+    else
+        Dict(vn => i for (i, vn) in enumerate(vns))
+    end
     # Construct new `vals` and `ranges`.
     vals_original = metadata.vals
     ranges_original = metadata.ranges
     # Allocate the new `vals`. and `ranges`.
-    vals = similar(metadata.vals, sum(length, ranges_original[indices_for_vns]))
-    ranges = similar(ranges_original)
+    vals = similar(metadata.vals, sum(length, ranges_original[indices_for_vns]; init=0))
+    ranges = similar(ranges_original, length(vns))
     # The new range `r` for `vns[i]` is offset by `offset` and
     # has the same length as the original range `r_original`.
     # The new `indices` (from above) ensures ordering according to `vns`.
@@ -415,7 +419,7 @@ function subset(metadata::Metadata, vns_given::AbstractVector{<:VarName})
         ranges,
         vals,
         metadata.dists[indices_for_vns],
-        metadata.gids,
+        metadata.gids[indices_for_vns],
         metadata.orders[indices_for_vns],
         flags,
     )

diff --git a/test/varinfo.jl b/test/varinfo.jl
@@ -578,6 +578,13 @@ DynamicPPL.getspace(::DynamicPPL.Sampler{MySAlg}) = (:s,)
                 else
                     vns_supported_standard
                 end
+
+            @testset ("$(convert(Vector{VarName}, vns_subset)) empty") for vns_subset in
+                                                                           vns_supported
+                varinfo_subset = subset(varinfo, VarName[])
+                @test isempty(varinfo_subset)
+            end
+
             @testset "$(convert(Vector{VarName}, vns_subset))" for vns_subset in
                                                                    vns_supported
                 varinfo_subset = subset(varinfo, vns_subset)