adding Model type back in

Author: PavanChaggar

src/AbstractPPL.jl

@@ -7,7 +7,7 @@ export VarName, getsym, getlens, inspace, subsumes, varname, vsym, @varname, @vs
 # Abstract model functions
 export AbstractProbabilisticProgram, condition, decondition, logdensityof, densityof
 
-# SimplePPL
+# GraphInfo
 export GraphInfo, Model, dag, nodes
 
 # Abstract traces
@@ -18,5 +18,5 @@ include("varname.jl")
 include("abstractmodeltrace.jl")
 include("abstractprobprog.jl")
 include("deprecations.jl")
-include("simpleppl.jl")
+include("graphinfo.jl")
 end # module

src/graphinfo.jl

@@ -54,12 +54,17 @@ s2 = (value = 0.0, input = (), eval = var"#5#8"(), kind = :Stochastic)
 y = (value = 0.0, input = (:μ, :s2), eval = var"#7#10"(), kind = :Stochastic)
 ```
 """
+
+struct Model
+    g::GraphInfo
+end
+
 function Model(;kwargs...)
     vals = getvals(NamedTuple(kwargs))
     args = [argnames(f) for f in vals[2]]
     A, sorted_vertices = DAG(NamedTuple{keys(kwargs)}(args))    
     modelinputs = NamedTuple{Tuple(sorted_vertices)}.([vals[1], Tuple.(args), vals[2], vals[3]])
-    GraphInfo(modelinputs..., A, sorted_vertices)
+    Model(GraphInfo(modelinputs..., A, sorted_vertices))
 end
 
 
@@ -134,8 +139,6 @@ function adjacency_matrix(inputs::NamedTuple{nodes}) where {nodes}
     return A
 end
 
-adjacency_matrix(m::GraphInfo) = adjacency_matrix(m.input)
-
 function outneighbors(A::SparseMatrixCSC, u::T) where T <: Int
     #adapted from Graph.jl https://github.com/JuliaGraphs/Graphs.jl/blob/06669054ed470bcfe4b2ad90ed974f2e65c84bb6/src/interface.jl#L302
     inds, _ = findnz(A[:, u])
@@ -185,54 +188,67 @@ Index a Model with a `VarName{p}` lens. Retrieves the `value``, `input`,
 # Examples
 
 ```jl-doctest 
-# add a model 
+julia> using AbstractPPL
+
+julia> m = Model( s2 = (0.0, () -> InverseGamma(2.0,3.0), :Stochastic), 
+                   μ = (1.0, () -> 1.0, :Logical), 
+                   y = (0.0, (μ, s2) -> MvNormal(μ, sqrt(s2)), :Stochastic))
+(s2 = Symbol[], μ = Symbol[], y = [:μ, :s2])
+Nodes: 
+μ = (value = 0.0, input = (), eval = var"#43#46"(), kind = :Stochastic)
+s2 = (value = 1.0, input = (), eval = var"#44#47"(), kind = :Logical)
+y = (value = 0.0, input = (:μ, :s2), eval = var"#45#48"(), kind = :Stochastic)
 
-julia> m[@varname y]
-(value = 0.0, input = (:μ, :s2), eval = var"#35#38"(), kind = :Stochastic)
 
+julia> m[@varname y]
+(value = 0.0, input = (:μ, :s2), eval = var"#45#48"(), kind = :Stochastic)
 ```
 """
-@generated function Base.getindex(m::GraphInfo, vn::VarName{p}) where {p}
+@generated function Base.getindex(g::GraphInfo, vn::VarName{p}) where {p}
     fns = fieldnames(GraphInfo)[1:4]
     name_lens = Setfield.PropertyLens{p}()
     field_lenses = [Setfield.PropertyLens{f}() for f in fns]
-    values = [:(get(m, Setfield.compose($l, $name_lens, getlens(vn)))) for l in field_lenses]
+    values = [:(get(g, Setfield.compose($l, $name_lens, getlens(vn)))) for l in field_lenses]
     return :(NamedTuple{$(fns)}(($(values...),)))
 end
 
-function Base.show(io::IO, m::GraphInfo)
+function Base.getindex(m::Model, vn::VarName)
+    return m.g[vn]
+end
+
+function Base.show(io::IO, m::Model)
     print(io, "Nodes: \n")
     for node in nodes(m)
         print(io, "$node = ", m[VarName{node}()], "\n")
     end
 end
 
 
-function Base.iterate(m::GraphInfo, state=1)
-    state > length(nodes(m)) ? nothing : (m[VarName{m.sorted_vertices[state]}()], state+1)
+function Base.iterate(m::Model, state=1)
+    state > length(nodes(m)) ? nothing : (m[VarName{m.g.sorted_vertices[state]}()], state+1)
 end
 
-Base.eltype(m::GraphInfo) = NamedTuple{fieldnames(GraphInfo)[1:4]}
-Base.IteratorEltype(m::GraphInfo) = HasEltype()
+Base.eltype(m::Model) = NamedTuple{fieldnames(GraphInfo)[1:4]}
+Base.IteratorEltype(m::Model) = HasEltype()
 
-Base.keys(m::GraphInfo) = (VarName{n}() for n in m.sorted_vertices)
-Base.values(m::GraphInfo) = Base.Generator(identity, m)
-Base.length(m::GraphInfo) = length(nodes(m))
-Base.keytype(m::GraphInfo) = eltype(keys(m))
-Base.valtype(m::GraphInfo) = eltype(m)
+Base.keys(m::Model) = (VarName{n}() for n in m.g.sorted_vertices)
+Base.values(m::Model) = Base.Generator(identity, m)
+Base.length(m::Model) = length(nodes(m))
+Base.keytype(m::Model) = eltype(keys(m))
+Base.valtype(m::Model) = eltype(m)
 
 
 """
     dag(m::Model)
 
 Returns the adjacency matrix of the model as a SparseArray.
 """
-dag(m::GraphInfo) = m.A
+dag(m::Model) = m.g.A
 
 """
     nodes(m::Model)
 
 Returns a `Vector{Symbol}` containing the sorted vertices 
 of the DAG. 
 """
-nodes(m::GraphInfo) = m.sorted_vertices
+nodes(m::Model) = m.g.sorted_vertices

test/graphinfo.jl

@@ -21,8 +21,9 @@ model = (
 m = Model(; zip(keys(model), values(model))...) # uses Model(; kwargs...) constructor
 
 # test the type of the model is correct
-@test typeof(m) <: GraphInfo <: AbstractModelTrace
-@test typeof(m) == GraphInfo{(:s2, :xmat, :β, :μ, :y)}
+@test typeof(m) == Model
+@test typeof(m.g) <: GraphInfo <: AbstractModelTrace
+@test typeof(m.g) == GraphInfo{(:s2, :xmat, :β, :μ, :y)}
 
 # test the dag is correct
 A = sparse([0 0 0 0 0; 0 0 0 0 0; 0 0 0 0 0; 0 1 1 0 0; 1 0 0 1 0])
@@ -34,11 +35,11 @@ A = sparse([0 0 0 0 0; 0 0 0 0 0; 0 0 0 0 0; 0 1 1 0 0; 1 0 0 1 0])
 # check the values from the NamedTuple match the values in the fields of GraphInfo
 vals = AbstractPPL.getvals(model)
 for (i, field) in enumerate([:value, :eval, :kind])
-    @test eval( :(values(m.$field) == vals[$i]) )
+    @test eval( :(values(m.g.$field) == vals[$i]) )
 end
 
 # test the right inputs have been inferred 
-@test m.input == (s2 = (), xmat = (), β = (), μ = (:xmat, :β), y = (:μ, :s2))
+@test m.g.input == (s2 = (), xmat = (), β = (), μ = (:xmat, :β), y = (:μ, :s2))
 
 # test keys are VarNames
 for key in keys(m)
@@ -47,9 +48,9 @@ end
 
 
 # test Model constructor for model with single parent node
-@test typeof(
-        Model(μ = (1.0, () -> 3, :Logical), y = (1.0, (μ) -> MvNormal(μ, sqrt(1)), :Stochastic))
-    ) == GraphInfo{(:μ, :y)}
+single_parent_m = Model(μ = (1.0, () -> 3, :Logical), y = (1.0, (μ) -> MvNormal(μ, sqrt(1)), :Stochastic))
+@test typeof(single_parent_m) == Model
+@test typeof(single_parent_m.g) == GraphInfo{(:μ, :y)}
 
 # test ErrorException for parent node not found
 @test_throws ErrorException Model( μ = (1.0, (β) -> 3, :Logical), y = (1.0, (μ) -> MvNormal(μ, sqrt(1)), :Stochastic))

test/runtests.jl

@@ -12,7 +12,7 @@ using Test
 
 @testset "AbstractPPL.jl" begin
     include("deprecations.jl")
-    include("simpleppl.jl")
+    include("graphinfo.jl")
     @testset "doctests" begin
         DocMeta.setdocmeta!(
             AbstractPPL,