Merge pull request #234 from ReactiveBayes/lazylabels

Implementation of VariableRef aka lazylabel

Author: bvdmitri

Project.toml

@@ -1,7 +1,7 @@
 name = "GraphPPL"
 uuid = "b3f8163a-e979-4e85-b43e-1f63d8c8b42c"
 authors = ["Wouter Nuijten <[email protected]>", "Dmitry Bagaev <[email protected]>"]
-version = "4.1.0"
+version = "4.2.0"
 
 [deps]
 BitSetTuples = "0f2f92aa-23a3-4d05-b791-88071d064721"

docs/src/developers_guide.md

@@ -119,17 +119,25 @@ GraphPPL.NodeData
 GraphPPL.NodeLabel
 GraphPPL.EdgeLabel
 GraphPPL.ProxyLabel
+GraphPPL.indexed_last
+GraphPPL.lift_index
+GraphPPL.datalabel
 GraphPPL.StaticInterfaces
-GraphPPL.LazyIndex
+GraphPPL.VariableRef
+GraphPPL.makevarref
 GraphPPL.MissingCollection
 GraphPPL.VariableNodeProperties
 GraphPPL.FactorNodeProperties
 GraphPPL.VarDict
 GraphPPL.AnonymousVariable
 GraphPPL.NodeDataExtraKey
-GraphPPL.LazyNodeLabel
 GraphPPL.IndexedVariable
 
+GraphPPL.VariableKindRandom
+GraphPPL.VariableKindData
+GraphPPL.VariableKindConstant
+GraphPPL.VariableKindUnknown
+
 GraphPPL.Deterministic
 GraphPPL.Stochastic
 GraphPPL.Atomic
@@ -151,7 +159,6 @@ GraphPPL.add_variable_node!
 GraphPPL.add_composite_factor_node!
 GraphPPL.copy_markov_blanket_to_child_context
 GraphPPL.generate_nodelabel
-GraphPPL.check_variate_compatability
 
 GraphPPL.FunctionalIndex
 GraphPPL.FunctionalRange

docs/src/getting_started.md

@@ -128,15 +128,15 @@ end
 
 ## Instantiating the model 
 
-To instantiate the model we need to pass the data `x`. We can do that with the [`GraphPPL.create_model`](@ref) function in combination with [`GraphPPL.LazyIndex`](@ref).
+To instantiate the model we need to pass the data `x`. We can do that with the [`GraphPPL.create_model`](@ref) function in combination with [`GraphPPL.datalabel`](@ref).
 
 ```@example getting-started
-xdata = [ 1.0, 0.0, 0.0, 1.0 ]
+data_for_x = [ 1.0, 0.0, 0.0, 1.0 ]
 
 model = GraphPPL.create_model(coin_toss()) do model, context
     return (; 
         # This expression creates data handle for `x` in the model using the `xdata` as the underlying collection
-        x = GraphPPL.getorcreate!(model, context, GraphPPL.NodeCreationOptions(kind = :data), :x, GraphPPL.LazyIndex(xdata))
+        x = GraphPPL.datalabel(model, context, GraphPPL.NodeCreationOptions(kind = GraphPPL.VariableKindData), :x, data_for_x)
     )
 end
 nothing #hide

ext/GraphPPLDistributionsExt.jl

@@ -13,7 +13,9 @@ end
 function distributions_ext_default_parametrization(
     t::Type{<:Distributions.Distribution}, ::GraphPPL.StaticInterfaces{interfaces}, interface_values
 ) where {interfaces}
-    @assert length(interface_values) == length(interfaces) "Distribution $t has $(length(interfaces)) fields $(interfaces) but $(length(interface_values)) values were provided."
+    if !(length(interface_values) == length(interfaces)) 
+        error(lazy"Distribution $t has $(length(interfaces)) fields $(interfaces) but $(length(interface_values)) values were provided.")
+    end
     return NamedTuple{interfaces}(interface_values)
 end
 

src/graph_engine.jl

@@ -57,11 +57,11 @@ GraphPPL.@model function beta_bernoulli(y, a, b)
     end
 end
 
-data_y = rand(Bernoulli(0.5), 100)
+data_for_y = rand(Bernoulli(0.5), 100)
 
 model = GraphPPL.create_model(beta_bernoulli(a = 1, b = 1)) do model, ctx 
     # Inject the data into the model
-    y = GraphPPL.getorcreate!(model, ctx, GraphPPL.NodeCreationOptions(kind = :data), :y, GraphPPL.LazyIndex(data_y))
+    y = GraphPPL.datalabel(model, ctx, GraphPPL.NodeCreationOptions(kind = GraphPPL.VariableKindData), :y, data_for_y)
     return (; y = y, )
 end
 

src/model_generator.jl

@@ -350,7 +350,7 @@ function add_get_or_create_expression(e::Expr)
     if @capture(e, (lhs_ ~ rhs_ where {options__}))
         @capture(lhs, (var_[index__]) | (var_))
         return quote
-            $(generate_get_or_create(var, index))
+            $(generate_get_or_create(var, index, rhs))
             $e
         end
     end
@@ -371,7 +371,7 @@ Generates code to get or create a variable in the graph. This function is used t
 # Returns
 A `quote` block with the code to get or create the variable in the graph.
 """
-generate_get_or_create(s::Symbol, index::Nothing) = generate_get_or_create(s, :((nothing,)))
+generate_get_or_create(s::Symbol, index::Nothing, rhs) = generate_get_or_create(s, :((nothing,)), rhs)
 
 """
     generate_get_or_create(s::Symbol, lhs::Expr, index::AbstractArray)
@@ -385,20 +385,15 @@ Generates code to get or create a variable in the graph. This function is used t
 # Returns
 A `quote` block with the code to get or create the variable in the graph.
 """
-generate_get_or_create(s::Symbol, index::AbstractArray) = generate_get_or_create(s, :(($(index...),)))
+generate_get_or_create(s::Symbol, index::AbstractArray, rhs) = generate_get_or_create(s, :(($(index...),)), rhs)
 
-function generate_get_or_create(s::Symbol, index::Expr)
+function generate_get_or_create(s::Symbol, index::Expr, rhs)
+    type = @capture(rhs, (f_()) | (f_(args__) | (f_(; kwargs__)) | (f_(args__; kwargs__)))) ? f : :(GraphPPL.Composite())
     return quote
         $s = if !@isdefined($s)
-            GraphPPL.getorcreate!(__model__, __context__, $(QuoteNode(s)), $(index)...)
+            GraphPPL.makevarref($type, __model__, __context__, GraphPPL.NodeCreationOptions(), $(QuoteNode(s)), $(index))
         else
-            (
-                if GraphPPL.check_variate_compatability($s, $(index)...)
-                    $s
-                else
-                    GraphPPL.getorcreate!(__model__, __context__, $(QuoteNode(s)), $(index)...)
-                end
-            )
+            $s
         end
     end
 end
@@ -445,7 +440,7 @@ Converts an expression into its proxied equivalent. Used to pass variables in su
 julia> x = GraphPPL.NodeLabel(:x, 1)
 x_1
 julia> GraphPPL.proxy_args(:(y = x))
-:(y = GraphPPL.proxylabel(:x, nothing, x))
+:(y = GraphPPL.proxylabel(:x, x, nothing, GraphPPL.False()))
 ```
 """
 function proxy_args end
@@ -464,29 +459,21 @@ function proxy_args(arg)
 end
 
 function proxy_args_lhs_eq_rhs(lhs, rhs)
-    @assert isa(lhs, Symbol) "Cannot wrap a ProxyLabel of `$lhs = $rhs` expression. The LHS must be a Symbol."
+    if !(isa(lhs, Symbol))
+        error(lazy"Cannot wrap a ProxyLabel of `$lhs = $rhs` expression. The LHS must be a Symbol.")
+    end
     return :($lhs = $(proxy_args_rhs(rhs)))
 end
 
 function proxy_args_rhs(rhs)
     if isa(rhs, Symbol)
-        return :(GraphPPL.proxylabel($(QuoteNode(rhs)), nothing, $rhs))
+        return :(GraphPPL.proxylabel($(QuoteNode(rhs)), $rhs, nothing, GraphPPL.False()))
     elseif @capture(rhs, rlabel_[index__])
-        return :(GraphPPL.proxylabel($(QuoteNode(rlabel)), $(Expr(:tuple, index...)), $rlabel))
+        return :(GraphPPL.proxylabel($(QuoteNode(rlabel)), $rlabel, $(Expr(:tuple, index...)), GraphPPL.False()))
     elseif @capture(rhs, new(rlabel_[index__]))
-        newrhs = gensym(:force_create)
-        errmsg = "Cannot force create a new label with the `new($rlabel[$(index...)])`. The label already exists."
-        return :(
-            let $newrhs = if isassigned($rlabel, $(index...))
-                    error($errmsg)
-                else
-                    GraphPPL.getorcreate!(__model__, __context__, $(QuoteNode(rlabel)), $(index...))
-                end
-                GraphPPL.proxylabel($(QuoteNode(rlabel)), $(Expr(:tuple, index...)), $newrhs)
-            end
-        )
+        return :(GraphPPL.proxylabel($(QuoteNode(rlabel)), $rlabel, $(Expr(:tuple, index...)), GraphPPL.True()))
     end
-    return rhs
+    return :(GraphPPL.proxylabel(:anonymous, $rhs, nothing, GraphPPL.False()))
 end
 
 """
@@ -553,10 +540,10 @@ end
 combine_broadcast_args(args::Nothing, kwargs::Nothing) = nothing
 
 generate_lhs_proxylabel(var, index::Nothing) = quote
-    GraphPPL.proxylabel($(QuoteNode(var)), nothing, $var)
+    GraphPPL.proxylabel($(QuoteNode(var)), $var, nothing, GraphPPL.True())
 end
 generate_lhs_proxylabel(var, index::AbstractArray) = quote
-    GraphPPL.proxylabel($(QuoteNode(var)), $(Expr(:tuple, index...)), $var)
+    GraphPPL.proxylabel($(QuoteNode(var)), $var, $(Expr(:tuple, index...)), GraphPPL.True())
 end
 
 __combine_axes() = Base.OneTo(1)
@@ -611,7 +598,7 @@ function convert_tilde_expression(e::Expr)
         combinable_args = kwargs === nothing ? args : vcat(args, [kwarg.args[2] for kwarg in kwargs])
         @capture(lhs, (var_[index__]) | (var_)) || error("Invalid left-hand side $(lhs). Must be in a `var` or `var[index]` form.")
         combinablesym = gensym()
-        getorcreate_lhs = generate_get_or_create(var, :(GraphPPL.__combine_axes($combinablesym...)))
+        getorcreate_lhs = generate_get_or_create(var, :(GraphPPL.__combine_axes($combinablesym...)), :(($fform)()))
         returnvalsym = gensym()
         return quote
             $combinablesym = ($(combinable_args...),)
@@ -710,7 +697,7 @@ function get_make_node_function(ms_body, ms_args, ms_name)
             __parent_context__::GraphPPL.Context,
             __options__::GraphPPL.NodeCreationOptions,
             ::typeof($ms_name),
-            __lhs_interface__::Union{GraphPPL.NodeLabel, GraphPPL.ProxyLabel},
+            __lhs_interface__::Union{GraphPPL.NodeLabel, GraphPPL.ProxyLabel, GraphPPL.VariableRef},
             __rhs_interfaces__::NamedTuple,
             __n_interfaces__::GraphPPL.StaticInt{$(length(ms_args))}
         )
@@ -722,7 +709,7 @@ function get_make_node_function(ms_body, ms_args, ms_name)
                 __model__, __context__, __options__, $ms_name, __interfaces__, __n_interfaces__
             )
             GraphPPL.returnval!(__context__, __returnval__)
-            return __context__, GraphPPL.unroll(__lhs_interface__)
+            return __context__, __lhs_interface__
         end
 
         function GraphPPL.add_terminated_submodel!(

src/model_macro.jl

@@ -677,6 +677,10 @@ edge case but don't know how we can resolve this, let alone efficiently. Please
     end
 end
 
+__resolve(model::Model, label::VariableRef) = __resolve(model, getifcreated(model, label.context, label))
+__resolve(model::Model, label::AbstractArray{T}) where {T <: VariableRef} =
+    __resolve(model, map(l -> getifcreated(model, l.context, l), label))
+
 function __resolve(model::Model, label::NodeLabel)
     data = model[label]
     return __resolve(model, data, getproperties(data), index(getproperties(data)))

src/plugins/variational_constraints/variational_constraints_engine.jl

@@ -68,7 +68,9 @@ function recursive_size(::Val{N}, vector::Vector{T}) where {N, T <: Vector}
 end
 
 function setindex!(array::ResizableArray{T, V, N}, value, index...) where {T, V, N}
-    @assert N === length(index) "Invalid index $(index) for $(array)"
+    if !(N === length(index))
+        error(lazy"Invalid index $(index) for $(array)")
+    end
     recursive_setindex!(Val(N), array.data, value, index...)
     return array
 end
@@ -128,7 +130,9 @@ function getindex(array::ResizableArray{T, V, N}, index::Vararg{UnitRange}) wher
 end
 
 function getindex(array::ResizableArray{T, V, N}, index::Vararg{Int}) where {T, V, N}
-    @assert N >= length(index) "Invalid index $(index) for $(array) of shape $(size(array)))"
+    if !(N >= length(index))
+        error(lazy"Invalid index $(index) for $(array) of shape $(size(array)))")
+    end
     return recursive_getindex(Val(length(index)), array.data, index...)
 end