Move factorization constraints to its own plugin

Author: bvdmitri

.JuliaFormatter.toml

@@ -1,6 +1,6 @@
 style = "blue"
 indent = 4
-margin = 180
+margin = 140
 always_for_in = true
 whitespace_typedefs = true
 whitespace_ops_in_indices = true

benchmark/model_zoo.jl

@@ -22,12 +22,11 @@ function create_hgf(n::Int)
     ω = GraphPPL.getorcreate!(model, ctx, :ω, nothing)
     θ = GraphPPL.getorcreate!(model, ctx, :θ, nothing)
     x_begin = GraphPPL.getorcreate!(model, ctx, :x_begin, nothing)
-    GraphPPL.add_terminated_submodel!(
+    GraphPPL.add_toplevel_model!(
         model,
         ctx,
         hgf,
-        (κ = κ, ω = ω, θ = θ, x_begin = x_begin, depth = n);
-        __debug__ = false,
+        (κ = κ, ω = ω, θ = θ, x_begin = x_begin, depth = n)
     )
     return model
 end

docs/src/developers_guide.md

@@ -57,12 +57,11 @@ x = GraphPPL.getorcreate!(model, context, :x, nothing; options = GraphPPL.Variab
 y = GraphPPL.getorcreate!(model, context, :y, nothing; options = GraphPPL.VariableNodeOptions(datavar=true))
 
 # Add the gcv model
-GraphPPL.add_terminated_submodel!(
+GraphPPL.add_toplevel_model!(
         model,
         context,
         gcv,
         (κ = κ, ω = ω, z = z, x = x, y = y);
-        __debug__ = false,
     )
 
 # Apply the constraints

src/graph_engine.jl

@@ -8,7 +8,7 @@ struct NodeCreatedByPlugin end
 
 const EmptyCreatedBy = Expr(:line, 0)
 
-GraphPPL.plugin_type(::NodeCreatedByPlugin) = FactorNodePlugin() 
+plugin_type(::NodeCreatedByPlugin) = FactorNodePlugin() 
 
 # The `created_by` field is used to track the expression that created the node.
 # The field can be a lambda function in which case it must be evaluated to get the expression.
@@ -21,7 +21,7 @@ Base.show(io::IO, createdby::CreatedBy) = show_createdby(io, createdby.created_b
 show_createdby(io::IO, created_by::Expr) = print(io, created_by)
 show_createdby(io::IO, created_by::Function) = show_createdby(io, created_by())
 
-function process_plugin(::NodeCreatedByPlugin, model::Model, context::Context, label::NodeLabel, nodedata::NodeData, options::NodeCreationOptions)
+function preprocess_plugin(::NodeCreatedByPlugin, model::Model, context::Context, label::NodeLabel, nodedata::NodeData, options::NodeCreationOptions)
     created_by = get(options, :created_by, EmptyCreatedBy)
     setextra!(nodedata, :created_by, CreatedBy(created_by))
     return label, nodedata

src/plugins/node_created_by.jl

@@ -15,3 +15,43 @@ struct FullFactorization end
 include("variational_constraints_macro.jl")
 include("variational_constraints_engine.jl")
 
+"""
+    VariationalConstraintsPlugin(constraints)
+
+A plugin that adds a VI related properties to the factor node for the variational inference procedure.
+"""
+struct VariationalConstraintsPlugin{C}
+    constraints::C
+end
+
+GraphPPL.plugin_type(::VariationalConstraintsPlugin) = FactorAndVariableNodesPlugin() 
+
+function preprocess_plugin(plugin::VariationalConstraintsPlugin, model::Model, context::Context, label::NodeLabel, nodedata::NodeData, options::NodeCreationOptions)
+    preprocess_vi_plugin!(plugin, nodedata, getproperties(nodedata))
+    return label, nodedata
+end
+
+function preprocess_vi_plugin!(::VariationalConstraintsPlugin, nodedata::NodeData, nodeproperties::FactorNodeProperties)
+    if hasextra(nodedata, :factorization_constraints)
+        error("Factorizatiom constraints has been already defined for the node ", nodedata, ".")
+    end
+    return nothing
+end
+
+function preprocess_vi_plugin!(::VariationalConstraintsPlugin, nodedata::NodeData, nodeproperties::VariableNodeProperties)
+    # TODO bvdmitri: todo, add functional form constraints and messages constraints here
+    return nothing
+end
+
+## Applies the constraints in `constraints` to `model`. This function materializes the constraints in `constraints` and applies them to `model`.
+function postprocess_plugin(plugin::VariationalConstraintsPlugin, model::Model)
+    # Attach `BitSetTuples` according to the number of neighbours of the factor node
+    foreach(factor_nodes(model)) do flabel 
+        nodedata = model[flabel]
+        nodeproperties = getproperties(nodedata)
+        number_of_neighbours = length(neighbors(nodeproperties))
+        setextra!(nodedata, :factorization_constraint_bitset, BitSetTuple(number_of_neighbours))
+    end
+    apply_constraints!(model, GraphPPL.get_principal_submodel(model), plugin.constraints, ConstraintStack())
+    materialize_constraints!(model)
+end

src/plugins/variational_constraints/variational_constraints.jl

@@ -51,9 +51,17 @@ function Base.filter(::UnknownPluginType, collection::PluginsCollection)
 end
 
 function Base.filter(trait::AbstractPluginTraitType, collection::PluginsCollection)
-    return PluginsCollection(filter(plugin -> plugin_type(plugin) === trait, collection.collection))
+    return PluginsCollection(filter(plugin -> isequal(plugin_type(plugin), trait), collection.collection))
 end
 
+struct UnionPluginType{T, U} <: AbstractPluginTraitType 
+    trait1::T
+    trait2::U
+end
+
+function Base.isequal(type::AbstractPluginTraitType, union::UnionPluginType) 
+    return isequal(type, union.trait1) || isequal(type, union.trait2)
+end
 
 
 # # OLD STUFF BELOW

src/plugins/variational_constraints/variational_constraints_engine.jl

@@ -2,12 +2,12 @@
 # We don't use models from the `model_zoo.jl` file because they are subject to change
 # These tests are meant to be stable and not change often
 
-@testitem "Simple model" begin
+@testitem "Simple model #1" begin
     using Distributions
 
-    import GraphPPL: create_model, getcontext, add_terminated_submodel!, factor_nodes, variable_nodes, is_constant, getproperties, as_node, as_variable
+    import GraphPPL: create_model, getcontext, add_toplevel_model!, factor_nodes, variable_nodes, is_constant, getproperties, as_node, as_variable
 
-    @model function simple_model()
+    @model function simple_model_1()
         x ~ Normal(0, 1)
         y ~ Gamma(1, 1)
         z ~ Normal(x, y)
@@ -16,7 +16,7 @@
     model = create_model()
     context = getcontext(model)
 
-    add_terminated_submodel!(model, context, simple_model, NamedTuple())
+    add_toplevel_model!(model, simple_model_1, NamedTuple())
 
     flabels = collect(factor_nodes(model))
     vlabels = collect(variable_nodes(model))
@@ -33,4 +33,29 @@
     @test length(collect(filter(as_variable(:x), model))) === 1
     @test length(collect(filter(as_variable(:y), model))) === 1
     @test length(collect(filter(as_variable(:z), model))) === 1
+end
+
+@testitem "Simple model #2" begin 
+    using Distributions
+    using GraphPPL: create_model, getcontext, getorcreate!, add_toplevel_model!, as_node, NodeCreationOptions, prune!
+
+    @model function simple_model_2(a, b, c)
+        x ~ Gamma(α = b, θ = sqrt(c))
+        a ~ Normal(μ = x, τ = 1)
+    end
+
+    model = create_model()
+    context = getcontext(model)
+
+    a = getorcreate!(model, context, NodeCreationOptions(datavar = true), :a, nothing)
+    b = getorcreate!(model, context, NodeCreationOptions(datavar = true), :b, nothing)
+    c = 1.0
+
+    add_toplevel_model!(model, simple_model_2, (a = a, b = b, c = c))
+
+    prune!(model)
+
+    @test length(collect(filter(as_node(Gamma), model))) === 1
+    @test length(collect(filter(as_node(Normal), model))) === 1
+    @test length(collect(filter(as_node(sqrt), model))) === 0 # should be compiled out, c is a constant
 end

src/plugins_collection.jl

@@ -271,7 +271,7 @@ end
 
     count = Ref(0)
 
-    function GraphPPL.process_plugin(::AnArbitraryPluginForTestUniqeness, model, context, label, nodedata, options)
+    function GraphPPL.preprocess_plugin(::AnArbitraryPluginForTestUniqeness, model, context, label, nodedata, options)
         setextra!(nodedata, :count, count[])
         count[] = count[] + 1
         return label, nodedata
@@ -298,7 +298,7 @@ end
 
     GraphPPL.plugin_type(::AnArbitraryPluginForChangingOptions) = GraphPPL.VariableNodePlugin()
 
-    function GraphPPL.process_plugin(::AnArbitraryPluginForChangingOptions, model, context, label, nodedata, options)
+    function GraphPPL.preprocess_plugin(::AnArbitraryPluginForChangingOptions, model, context, label, nodedata, options)
         # Here we replace the original options entirely
         return label, NodeData(context, convert(VariableNodeProperties, :x, nothing, NodeCreationOptions(constant = true, value = 1.0)))
     end
@@ -479,15 +479,18 @@ end
 
     model = create_model()
     ctx = getcontext(model)
+    @test isempty(model)
     @test nv(model) == 0
     @test ne(model) == 0
 
     model[NodeLabel(:a, 1)] = NodeData(ctx, VariableNodeProperties(name = :a, index = nothing))
     model[NodeLabel(:b, 2)] = NodeData(ctx, VariableNodeProperties(name = :b, index = nothing))
+    @test !isempty(model)
     @test nv(model) == 2
     @test ne(model) == 0
 
     model[NodeLabel(:a, 1), NodeLabel(:b, 2)] = EdgeLabel(:edge, 1)
+    @test !isempty(model)
     @test nv(model) == 2
     @test ne(model) == 1
 end

test/graph_construction_tests.jl

@@ -46,7 +46,7 @@ GraphPPL.factor_alias(::Type{Gamma}, ::Val{(:α, :θ)}) = GammaShapeScale
 function create_terminated_model(fform; plugins = GraphPPL.PluginsCollection())
     __model__ = GraphPPL.create_model(; fform = fform, plugins = plugins)
     __context__ = GraphPPL.getcontext(__model__)
-    GraphPPL.add_terminated_submodel!(__model__, __context__, fform, NamedTuple())
+    GraphPPL.add_toplevel_model!(__model__, __context__, fform, NamedTuple())
     return __model__
 end
 

test/graph_engine_tests.jl

@@ -65,7 +65,7 @@ end
 @testitem "Usage with the actual model" begin
     using Distributions
 
-    import GraphPPL: create_model, getcontext, add_terminated_submodel!, factor_nodes, as_node, hasextra, PluginsCollection, NodeCreatedByPlugin, getextra
+    import GraphPPL: create_model, getcontext, add_toplevel_model!, factor_nodes, as_node, hasextra, PluginsCollection, NodeCreatedByPlugin, getextra
 
     @model function simple_model()
         x ~ Normal(0, 1)
@@ -76,7 +76,7 @@ end
     model = create_model(plugins = PluginsCollection(NodeCreatedByPlugin()))
     context = getcontext(model)
 
-    add_terminated_submodel!(model, context, simple_model, NamedTuple())
+    add_toplevel_model!(model, simple_model, NamedTuple())
 
     fnormal = map(label -> model[label], filter(as_node(Normal), model))
     fgamma = map(label -> model[label], filter(as_node(Gamma), model))