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/factorisation.jl

@@ -187,19 +187,56 @@ resolve_factorisation(::UnspecifiedConstraints, any, allvariables, fform, variab
 # 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 @@ function resolve_factorisation(::Stochastic, constraints, allvariables, fform, _
  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 @@ function (mapping::MarginalMapping)(dependencies)
  # 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 @@ function materialize!(mapping::MessageMapping, messages, marginals)
  # 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 @@ import Base: show
 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::Tuple) where {D}
 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 @@ isdata(::AbstractArray{<:DataVariable}) = true
 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,28 @@ _makemarginal(datavar::DataVariable) = error("It is not possible to
 
 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))
+
+# options here must implement at least `Rocket.getscheduler`
+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 @@ add_pipeline_stage!(variable::AbstractVariable, stage) = error("Its not possible
 # 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)