Merge pull request #50 from phipsgabler/phg/varinfo-indices

Refactor VarName

Author: yebai

src/DynamicPPL.jl

@@ -6,11 +6,9 @@ using Bijectors
 using MacroTools
 import ZygoteRules
 
-import Base: string,
-             Symbol,
+import Base: Symbol,
              ==,
              hash,
-             in,
              getindex,
              setindex!,
              push!,
@@ -23,8 +21,7 @@ import Base: string,
              haskey
 
 # VarInfo
-export  VarName,
-        AbstractVarInfo,
+export  AbstractVarInfo,
         VarInfo,
         UntypedVarInfo,
         getlogp,
@@ -45,6 +42,10 @@ export  VarName,
         link!,
         invlink!,
         tonamedtuple,
+#VarName
+        VarName,
+        inspace,
+        subsumes,
 # Compiler
         ModelGen,
         @model,

src/compiler.jl

@@ -70,7 +70,7 @@ end
 To generate a `Model`, call `model_generator(x_value)`.
 """
 macro model(input_expr)
-    build_model_info(input_expr) |> replace_tilde! |> replace_vi! |> 
+    Base.replace_ref_end!(input_expr) |> build_model_info |> replace_tilde! |> replace_vi! |> 
         replace_logpdf! |> replace_sampler! |> build_output
 end
 

src/context_implementations.jl

@@ -52,19 +52,7 @@ function _tilde(sampler, right, vn::VarName, vi)
     return assume(sampler, right, vn, vi)
 end
 function _tilde(sampler, right::NamedDist, vn::VarName, vi)
-    name = right.name
-    if name isa String
-        sym_str, inds = split_var_str(name, String)
-        sym = Symbol(sym_str)
-        vn = VarName{sym}(inds)
-    elseif name isa Symbol
-        vn = VarName{name}("")
-    elseif name isa VarName
-        vn = name
-    else
-        throw("Unsupported variable name. Please use either a string, symbol or VarName.")
-    end
-    return _tilde(sampler, right.dist, vn, vi)
+    return _tilde(sampler, right.dist, right.name, vi)
 end
 
 # observe
@@ -214,30 +202,19 @@ end
 
 
 function get_vns_and_dist(dist::NamedDist, var, vn::VarName)
-    name = dist.name
-    if name isa String
-        sym_str, inds = split_var_str(name, String)
-        sym = Symbol(sym_str)
-        vn = VarName{sym}(inds)
-    elseif name isa Symbol
-        vn = VarName{name}("")
-    elseif name isa VarName
-        vn = name
-    else
-        throw("Unsupported variable name. Please use either a string, symbol or VarName.")
-    end
-    return get_vns_and_dist(dist.dist, var, vn)
+    return get_vns_and_dist(dist.dist, var, dist.name)
 end
 function get_vns_and_dist(dist::MultivariateDistribution, var::AbstractMatrix, vn::VarName)
-    getvn = i -> VarName(vn, vn.indexing * "[Colon(),$i]")
+    getvn = i -> VarName(vn, (vn.indexing..., (Colon(), i)))
     return getvn.(1:size(var, 2)), dist
+    
 end
 function get_vns_and_dist(
     dist::Union{Distribution, AbstractArray{<:Distribution}}, 
     var::AbstractArray, 
     vn::VarName
 )
-    getvn = ind -> VarName(vn, vn.indexing * "[" * join(Tuple(ind), ",") * "]")
+    getvn = ind -> VarName(vn, (vn.indexing..., Tuple(ind)))
     return getvn.(CartesianIndices(var)), dist
 end
 

src/distribution_wrappers.jl

@@ -11,13 +11,15 @@ A named distribution that carries the name of the random variable with it.
 struct NamedDist{
     variate, 
     support, 
-    Td <: Distribution{variate, support}, 
-    Tn
+    Td <: Distribution{variate, support},
+    Tv <: VarName
 } <: Distribution{variate, support}
     dist::Td
-    name::Tn
+    name::Tv
 end
 
+NamedDist(dist::Distribution, name::Symbol) = NamedDist(dist, VarName(name))
+
 
 struct NoDist{
     variate, 

src/varinfo.jl

@@ -125,7 +125,7 @@ end
     offset = :(0)
     for f in names
         mdf = :(metadata.$f)
-        if f in space || length(space) == 0
+        if inspace(f, space) || length(space) == 0
             len = :(length($mdf.vals))
             push!(exprs, :($f = Metadata($mdf.idcs,
                                         $mdf.vns,
@@ -330,13 +330,6 @@ setall!(vi::TypedVarInfo, val) = _setall!(vi.metadata, val)
     return expr
 end
 
-"""
-    getsym(vn::VarName)
-
-Return the symbol of the Julia variable used to generate `vn`.
-"""
-getsym(vn::VarName{sym}) where sym = sym
-
 """
     getgid(vi::VarInfo, vn::VarName)
 
@@ -407,7 +400,7 @@ end
         # If the varname is in the sampler space
         # or the sample space is empty (all variables)
         # then return the indices for that variable.
-        if f in space || length(space) == 0
+        if inspace(f, space) || length(space) == 0
             push!(exprs, :($f = findinds(metadata.$f, s, Val($space))))
         end
     end
@@ -418,7 +411,7 @@ end
     # Get all the idcs of the vns in `space` and that belong to the selector `s`
     return filter((i) ->
         (s in f_meta.gids[i] || isempty(f_meta.gids[i])) &&
-        (isempty(space) || in(f_meta.vns[i], space)), 1:length(f_meta.gids))
+        (isempty(space) || inspace(f_meta.vns[i], space)), 1:length(f_meta.gids))
 end
 @inline function findinds(f_meta)
     # Get all the idcs of the vns
@@ -488,69 +481,6 @@ end
 #### APIs for typed and untyped VarInfo
 ####
 
-# VarName
-
-"""
-    VarName(sym, indexing)
-    VarName{sym}(indexing::String)
-
-Construct a new instance of `VarName{sym}`
-"""
-VarName(sym, indexing) = VarName{sym}(indexing)
-
-"""
-    VarName(vn::VarName, indexing::String)
-
-Return a copy of `vn` with a new index `indexing`.
-"""
-function VarName(vn::VarName, indexing::String)
-    return VarName{getsym(vn)}(indexing)
-end
-
-"""
-    uid(vn::VarName)
-
-Return a unique tuple identifier for `vn`.
-"""
-uid(vn::VarName) = (getsym(vn), vn.indexing)
-
-hash(vn::VarName) = hash(uid(vn))
-
-==(x::VarName, y::VarName) = hash(uid(x)) == hash(uid(y))
-
-function string(vn::VarName)
-    return "$(getsym(vn))$(vn.indexing)"
-end
-function string(vns::Vector{<:VarName})
-    return replace(string(map(string, vns)), "String" => "")
-end
-
-"""
-    Symbol(vn::VarName)
-
-Return a `Symbol` represenation of the variable identifier `VarName`.
-"""
-Symbol(vn::VarName) = Symbol(string(vn))  # simplified symbol
-
-"""
-    in(vn::VarName, space::Set)
-
-Check whether `vn`'s symbol is in `space`.
-"""
-in(::VarName, ::Tuple{}) = true
-in(vn::VarName, space::Tuple)::Bool = getsym(vn) in space || _in(string(vn), space)
-
-_in(::String, ::Tuple{}) = false
-_in(vn_str::String, space::Tuple)::Bool = _in(vn_str, Base.tail(space))
-function _in(vn_str::String, space::Tuple{Expr,Vararg})::Bool
-    # Collect expressions from space
-    expr = first(space)
-    # Filter `(` and `)` out and get a string representation of `exprs`
-    expr_str = replace(string(expr), r"\(|\)" => "")
-    # Check if `vn_str` is in `expr_strs`
-    valid = occursin(expr_str, vn_str)
-    return valid || _in(vn_str, Base.tail(space))
-end
 
 # VarInfo
 
@@ -602,8 +532,7 @@ function TypedVarInfo(vi::UntypedVarInfo)
         sym_vals = foldl(vcat, _vals)
 
         push!(new_metas, Metadata(sym_idcs, sym_vns, sym_ranges, sym_vals,
-                                    sym_dists, sym_gids, sym_orders, sym_flags)
-            )
+                                  sym_dists, sym_gids, sym_orders, sym_flags))
     end
     logp = getlogp(vi)
     num_produce = get_num_produce(vi)
@@ -764,7 +693,7 @@ end
 @generated function _link!(metadata::NamedTuple{names}, vi, vns, ::Val{space}) where {names, space}
     expr = Expr(:block)
     for f in names
-        if f in space || length(space) == 0
+        if inspace(f, space) || length(space) == 0
             push!(expr.args, quote
                 f_vns = vi.metadata.$f.vns
                 if ~istrans(vi, f_vns[1])
@@ -810,7 +739,7 @@ end
 @generated function _invlink!(metadata::NamedTuple{names}, vi, vns, ::Val{space}) where {names, space}
     expr = Expr(:block)
     for f in names
-        if f in space || length(space) == 0
+        if inspace(f, space) || length(space) == 0
             push!(expr.args, quote
                 f_vns = vi.metadata.$f.vns
                 if istrans(vi, f_vns[1])
@@ -1173,7 +1102,7 @@ Set `vn`'s `gid` to `Set([spl.selector])`, if `vn` does not have a sampler selec
 and `vn`'s symbol is in the space of `spl`.
 """
 function updategid!(vi::AbstractVarInfo, vn::VarName, spl::Sampler)
-    if vn in getspace(spl)
+    if inspace(vn, getspace(spl))
         setgid!(vi, spl.selector, vn)
     end
 end