Modify variables structures for predictions functionality

Project.toml

@@ -3,16 +3,6 @@ uuid = "a194aa59-28ba-4574-a09c-4a745416d6e3"
 authors = ["Dmitry Bagaev <d.v.bagaev@tue.nl>", "Albert Podusenko <a.podusenko@tue.nl>", "Bart van Erp <b.v.erp@tue.nl>", "Ismail Senoz <i.senoz@tue.nl>"]
 version = "3.9.3"
 
-[weakdeps]
-Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
-
-[extensions]
-ReactiveMPOptimisersExt = "Optimisers"
-ReactiveMPZygoteExt = "Zygote"
-ReactiveMPRequiresExt = "Requires"
-
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -38,6 +28,16 @@ TinyHugeNumbers = "783c9a47-75a3-44ac-a16b-f1ab7b3acf04"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 Unrolled = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 
+[weakdeps]
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[extensions]
+ReactiveMPOptimisersExt = "Optimisers"
+ReactiveMPRequiresExt = "Requires"
+ReactiveMPZygoteExt = "Zygote"
+
 [compat]
 DataStructures = "0.17, 0.18"
 Distributions = "0.24, 0.25"
@@ -73,9 +73,9 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"

src/constraints/specifications/constraints.jl

@@ -104,7 +104,10 @@
     foreach(constraints.factorisation) do spec
         specnames = getnames(spec)
         foreach(specnames) do specname
-            if warn && (hasdatavar(variables, specname) || hasconstvar(variables, specname))
+            if hasdatavar(variables, specname) && allows_missings(variables[specname])
+                # skip, because it is fine to have a datavar in the factorization constraint, which allows missings
+                nothing
+            elseif warn && (hasdatavar(variables, specname) || hasconstvar(variables, specname))
                 @warn "Constraints specification has factorisation constraint for `q($(join(specnames, ", ")))`, but `$(specname)` is not a random variable. Data variables and constants in the model are forced to be factorized by default such that `q($(join(specnames, ", "))) = q($(specname))q(...)` . Use `warn = false` option during constraints specification to suppress this warning."
             elseif warn && !hasrandomvar(variables, specname)
                 @warn "Constraints specification has factorisation constraint for `q($(join(specnames, ", ")))`, but variables collection has no random variable named `$(specname)`. Use `warn = false` option during constraints specification to suppress this warning."

src/constraints/specifications/factorisation.jl

@@ -187,19 +187,56 @@
 # Preoptimised dispatch rule for unspecified constraints and a deterministic node with any number of inputs
 resolve_factorisation(::UnspecifiedConstraints, ::Deterministic, allvariables, fform, variables) = FullFactorisation()
 
-# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 2 inputs
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2}) where {V1 <: RandomVariable, V2 <: RandomVariable}                         = ((1, 2),)
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2}) where {V1 <: Union{<:ConstVariable, <:DataVariable}, V2 <: RandomVariable} = ((1,), (2,))
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2}) where {V1 <: RandomVariable, V2 <: Union{<:ConstVariable, <:DataVariable}} = ((1,), (2,))
-
-# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 3 inputs
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: RandomVariable, V3 <: RandomVariable}                                                 = ((1, 2, 3),)
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: Union{<:ConstVariable, <:DataVariable}, V2 <: RandomVariable, V3 <: RandomVariable}                         = ((1,), (2, 3))
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: Union{<:ConstVariable, <:DataVariable}, V3 <: RandomVariable}                         = ((1, 3), (2,))
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: RandomVariable, V3 <: Union{<:ConstVariable, <:DataVariable}}                         = ((1, 2), (3,))
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: Union{<:ConstVariable, <:DataVariable}, V3 <: Union{<:ConstVariable, <:DataVariable}} = ((1,), (2,), (3,))
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: Union{<:ConstVariable, <:DataVariable}, V2 <: RandomVariable, V3 <: Union{<:ConstVariable, <:DataVariable}} = ((1,), (2,), (3,))
-resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: Union{<:ConstVariable, <:DataVariable}, V2 <: Union{<:ConstVariable, <:DataVariable}, V3 <: RandomVariable} = ((1,), (2,), (3,))
+# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 2 inputs, random variable & constant variable
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2}) where {V1 <: RandomVariable, V2 <: RandomVariable} = ((1, 2),)
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2}) where {V1 <: ConstVariable, V2 <: RandomVariable} = ((1,), (2,))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2}) where {V1 <: RandomVariable, V2 <: ConstVariable} = ((1,), (2,))
+
+# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 2 inputs, random variable & data variable
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2}) where {V1 <: DataVariable, V2 <: RandomVariable} =
+    allows_missings(variables[1]) ? ((1, 2),) : ((1,), (2,))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2}) where {V1 <: RandomVariable, V2 <: DataVariable} =
+    allows_missings(variables[2]) ? ((1, 2),) : ((1,), (2,))
+
+# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 3 inputs, random variable & constant variables
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: RandomVariable, V3 <: RandomVariable} = ((1, 2, 3),)
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: ConstVariable, V2 <: RandomVariable, V3 <: RandomVariable}  = ((1,), (2, 3))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: ConstVariable, V3 <: RandomVariable}  = ((1, 3), (2,))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: RandomVariable, V3 <: ConstVariable}  = ((1, 2), (3,))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: RandomVariable, V2 <: ConstVariable, V3 <: ConstVariable}   = ((1,), (2,), (3,))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: ConstVariable, V2 <: RandomVariable, V3 <: ConstVariable}   = ((1,), (2,), (3,))
+resolve_factorisation(::UnspecifiedConstraints, ::Stochastic, allvariables, fform, ::Tuple{V1, V2, V3}) where {V1 <: ConstVariable, V2 <: ConstVariable, V3 <: RandomVariable}   = ((1,), (2,), (3,))
+
+# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 3 inputs, random variable & data variable
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: DataVariable, V2 <: RandomVariable, V3 <: RandomVariable} = allows_missings(variables[1]) ? ((1, 2, 3),) : ((1,), (2, 3))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: RandomVariable, V2 <: DataVariable, V3 <: RandomVariable} = allows_missings(variables[2]) ? ((1, 2, 3),) : ((1, 3), (2,))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: RandomVariable, V2 <: RandomVariable, V3 <: DataVariable} = allows_missings(variables[3]) ? ((1, 2, 3),) : ((1, 2), (3,))
+
+# Preoptimised dispatch rules for unspecified constraints and a stochastic node with 3 inputs, random variable & data variable & const variable
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: DataVariable, V2 <: ConstVariable, V3 <: RandomVariable} = allows_missings(variables[1]) ? ((1, 3), (2,)) : ((1,), (2,), (3,))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: DataVariable, V2 <: RandomVariable, V3 <: ConstVariable} = allows_missings(variables[1]) ? ((1, 2), (3,)) : ((1,), (2,), (3,))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: ConstVariable, V2 <: DataVariable, V3 <: RandomVariable} = allows_missings(variables[2]) ? ((1,), (2, 3)) : ((1,), (3,), (2,))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: RandomVariable, V2 <: DataVariable, V3 <: ConstVariable} = allows_missings(variables[2]) ? ((1, 2), (3,)) : ((1,), (2,), (3,))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: ConstVariable, V2 <: RandomVariable, V3 <: DataVariable} = allows_missings(variables[3]) ? ((1,), (2, 3)) : ((1,), (2,), (3,))
+resolve_factorisation(
+    ::UnspecifiedConstraints, ::Stochastic, allvariables, fform, variables::Tuple{V1, V2, V3}
+) where {V1 <: RandomVariable, V2 <: ConstVariable, V3 <: DataVariable} = allows_missings(variables[3]) ? ((1, 3), (2,)) : ((1,), (2,), (3,))
 
 """
     resolve_factorisation(constraints, allvariables, fform, variables) 
@@ -419,8 +456,11 @@
     index::Int = 1
     shift::Int = 0
     for varref in var_refs
-        if israndom(varref[3])
+        if israndom(varref[3]) || (isdata(varref[3]) && allows_missings(varref[3]))
             # We process everything as usual if varref is a random variable
+            # or if the variable is data variable and it allows missing 
+            # We probably should change the logic from "allows missings" to "used as prediction"
+            # For now we assume that if data variable allows missing input it is indeed "used as prediction"
             __process_factorisation_entry!(varref[1], varref[2], shift)
         else
             # We filter out varref from all clusters if it is not random

src/marginal.jl

@@ -195,7 +195,13 @@
     # Marginal is initial if it is not clamped and all of the inputs are either clamped or initial
     is_marginal_initial = !is_marginal_clamped && (__check_all(is_clamped_or_initial, messages) && __check_all(is_clamped_or_initial, marginals))
 
-    marginal = marginalrule(marginal_mapping_fform(mapping), mapping.vtag, mapping.msgs_names, messages, mapping.marginals_names, marginals, mapping.meta, mapping.factornode)
+    marginal = if !isnothing(messages) && any(ismissing, TupleTools.flatten(getdata.(messages)))
+        missing
+    elseif !isnothing(marginals) && any(ismissing, TupleTools.flatten(getdata.(marginals)))
+        missing
+    else
+        marginalrule(marginal_mapping_fform(mapping), mapping.vtag, mapping.msgs_names, messages, mapping.marginals_names, marginals, mapping.meta, mapping.factornode)
+    end
 
     return Marginal(marginal, is_marginal_clamped, is_marginal_initial, nothing)
 end

src/message.jl

@@ -324,18 +324,24 @@
     # Message is initial if it is not clamped and all of the inputs are either clamped or initial
     is_message_initial = !is_message_clamped && (__check_all(is_clamped_or_initial, messages) && __check_all(is_clamped_or_initial, marginals))
 
-    result, addons = rule(
-        message_mapping_fform(mapping),
-        mapping.vtag,
-        mapping.vconstraint,
-        mapping.msgs_names,
-        messages,
-        mapping.marginals_names,
-        marginals,
-        mapping.meta,
-        mapping.addons,
-        mapping.factornode
-    )
+    result, addons = if !isnothing(messages) && any(ismissing, TupleTools.flatten(getdata.(messages)))
+        missing, mapping.addons
+    elseif !isnothing(marginals) && any(ismissing, TupleTools.flatten(getdata.(marginals)))
+        missing, mapping.addons
+    else
+        rule(
+            message_mapping_fform(mapping),
+            mapping.vtag,
+            mapping.vconstraint,
+            mapping.msgs_names,
+            messages,
+            mapping.marginals_names,
+            marginals,
+            mapping.meta,
+            mapping.addons,
+            mapping.factornode
+        )
+    end
 
     # Inject extra addons after the rule has been executed
     addons = message_mapping_addons(mapping, getdata(messages), getdata(marginals), result, addons)

src/node.jl

@@ -1149,7 +1149,7 @@ macro node(fformtype, sdtype, interfaces_list)
             missingclustererr = "Cannot find the cluster for the variable connected to the `$(name)` interface around the `$fformtype` node."
             quote
                 # If a variable `$name` is a constvar or a datavar
-                if ReactiveMP.isconst($(name)) || ReactiveMP.isdata($(name))
+                if ReactiveMP.isconst($(name)) || (ReactiveMP.isdata($(name)) && !ReactiveMP.allows_missings($(name)))
                     local __factorisation = ReactiveMP.factorisation(node)
                     # Find the factorization cluster associated with the constvar `$name`
                     local __index  = ReactiveMP.interface_get_index(Val{$(QuoteNode(fbottomtype))}, Val{$(QuoteNode(name))})

src/variables/data.jl

@@ -5,6 +5,8 @@
 mutable struct DataVariable{D, S} <: AbstractVariable
     name            :: Symbol
     collection_type :: AbstractVariableCollectionType
+    prediction      :: MarginalObservable
+    input_messages  :: Vector{MessageObservable{AbstractMessage}}
     messageout      :: S
     nconnected      :: Int
     isproxy         :: Bool
@@ -74,12 +76,16 @@
 datavar(name::Symbol, ::Type{D}, dims::Vararg{Int}) where {D}                                                      = datavar(DataVariableCreationOptions(D), name, D, dims)
 
 datavar(options::DataVariableCreationOptions{S}, name::Symbol, ::Type{D}, collection_type::AbstractVariableCollectionType = VariableIndividual()) where {S, D} =
-    DataVariable{D, S}(name, collection_type, options.subject, 0, options.isproxy, options.isused)
+    DataVariable{D, S}(name, collection_type, MarginalObservable(), Vector{MessageObservable{AbstractMessage}}(), options.subject, 0, options.isproxy, options.isused)
 
 function datavar(options::DataVariableCreationOptions, name::Symbol, ::Type{D}, length::Int) where {D}
     return map(i -> datavar(similar(options), name, D, VariableVector(i)), 1:length)
 end
 
+function datavar(options::DataVariableCreationOptions, name::Symbol, ::Type{D}, dim1::Int, extra_dims::Vararg{Int}) where {D}
+    return datavar(options, name, D, (dim1, extra_dims...))
+end
+
 function datavar(options::DataVariableCreationOptions, name::Symbol, ::Type{D}, dims::Tuple) where {D}
     indices = CartesianIndices(dims)
     size    = axes(indices)
@@ -106,11 +112,17 @@
 isconst(::DataVariable)                   = false
 isconst(::AbstractArray{<:DataVariable})  = false
 
+allows_missings(datavar::DataVariable) = allows_missings(datavar, eltype(datavar.messageout))
+
+allows_missings(datavars::AbstractArray{<:DataVariable}) = all(allows_missings, datavars)
+allows_missings(datavar::DataVariable, ::Type{Message{D}}) where {D} = false
+allows_missings(datavar::DataVariable, ::Type{Union{Message{Missing}, Message{D}}} where {D}) = true
+
 function Base.getindex(datavar::DataVariable, i...)
     error("Variable $(indexed_name(datavar)) has been indexed with `[$(join(i, ','))]`. Direct indexing of `data` variables is not allowed.")
 end
 
-getlastindex(::DataVariable) = 1
+getlastindex(datavar::DataVariable) = degree(datavar) + 1
 
 messageout(datavar::DataVariable, ::Int) = datavar.messageout
 messagein(datavar::DataVariable, ::Int)  = error("It is not possible to get a reference for inbound message for datavar")
@@ -163,8 +175,27 @@
 
 setanonymous!(::DataVariable, ::Bool) = nothing
 
-function setmessagein!(datavar::DataVariable, ::Int, messagein)
-    datavar.nconnected += 1
-    datavar.isused = true
+function setmessagein!(datavar::DataVariable, index::Int, messagein)
+    if index === (degree(datavar) + 1)
+        push!(datavar.input_messages, messagein)
+        datavar.nconnected += 1
+        datavar.isused = true
+    else
+        error(
+            "Inconsistent state in setmessagein! function for data variable $(datavar). `index` should be equal to `degree(datavar) + 1 = $(degree(datavar) + 1)`, $(index) is given instead"
+        )
+    end
+    return nothing
+end
+
+marginal_prod_fn(datavar::DataVariable) = marginal_prod_fn(FoldLeftProdStrategy(), ProdAnalytical(), UnspecifiedFormConstraint(), FormConstraintCheckLast())
+
+_getprediction(datavar::DataVariable)              = datavar.prediction
+_setprediction!(datavar::DataVariable, observable) = connect!(_getprediction(datavar), observable)
+_makeprediction(datavar::DataVariable)             = collectLatest(AbstractMessage, Marginal, datavar.input_messages, marginal_prod_fn(datavar))
+
+function activate!(datavar::DataVariable, options)
+    _setprediction!(datavar, _makeprediction(datavar))
+
     return nothing
 end

src/variables/variable.jl

@@ -1,7 +1,7 @@
 export AbstractVariable, degree
 export is_clamped, is_marginalisation, is_moment_matching
 export FoldLeftProdStrategy, FoldRightProdStrategy, CustomProdStrategy
-export getmarginal, getmarginals, setmarginal!, setmarginals!, name, as_variable
+export getprediction, getpredictions, getmarginal, getmarginals, setmarginal!, setmarginals!, name, as_variable
 export setmessage!, setmessages!
 
 using Rocket
@@ -80,6 +80,9 @@
 # Helper functions
 # Getters
 
+getprediction(variable::AbstractVariable)                    = _getprediction(variable)
+getpredictions(variables::AbstractArray{<:AbstractVariable}) = collectLatest(map(v -> getprediction(v), variables))
+
 getmarginal(variable::AbstractVariable)                                      = getmarginal(variable, SkipInitial())
 getmarginal(variable::AbstractVariable, skip_strategy::MarginalSkipStrategy) = apply_skip_filter(_getmarginal(variable), skip_strategy)