GitHub

Annotations

Messages and marginals in ReactiveMP carry a probability distribution as their primary content. Annotations are an optional side-channel that can travel alongside a message, holding arbitrary extra information keyed by Symbol. Typical uses include tracking log-scale factors (see LogScaleAnnotations), recording which messages were used to compute a result, or attaching debugging information.

Annotations are designed to be zero-cost when unused: the underlying dictionary is only allocated on the first write.

AnnotationDict

Every message and marginal holds an ReactiveMP.AnnotationDict. The basic operations are:

ReactiveMP.AnnotationDict — Type
AnnotationDict()
AnnotationDict(other::AnnotationDict)

A mutable dictionary that associates Symbol keys with arbitrary annotation values. Supports lazy initialization — no memory is allocated until the first write.

The copy constructor creates an independent shallow copy of other.

source
ReactiveMP.annotate! — Function
annotate!(ann::AnnotationDict, key::Symbol, value)

Store value under key in ann. Always returns nothing.

source
ReactiveMP.get_annotation — Function
get_annotation(ann::AnnotationDict, key::Symbol)

Return the value stored under key. Throws KeyError if key is absent.

source
get_annotation(ann::AnnotationDict, ::Type{T}, key::Symbol) where {T}

Return the value stored under key, converted to type T. Throws KeyError if key is absent.

source
ReactiveMP.has_annotation — Function
has_annotation(ann::AnnotationDict, key::Symbol) -> Bool

Return true if ann contains an entry for key, false otherwise.

source

Annotation processors

Annotation processors are subtypes of ReactiveMP.AbstractAnnotations that define how annotations are written and merged. There are three integration points:

  • Before a rule executes — ReactiveMP.pre_rule_annotations! is called with the processor, the rule's AnnotationDict, the MessageMapping, the incoming messages and marginals. Use this to write annotations that does not depend on what the rule computed.
  • After a rule executes — ReactiveMP.post_rule_annotations! is called with the processor, the rule's AnnotationDict, the MessageMapping, the incoming messages and marginals, and the result distribution. Use this to write annotations that depend on what the rule computed.
  • During a message product — ReactiveMP.post_product_annotations! is called with the processor, a fresh merged AnnotationDict, and the left and right annotation dicts together with the distributions involved. Use this to merge annotations from the two incoming messages into the product message.
ReactiveMP.pre_rule_annotations! — Function
pre_rule_annotations!(processor::AbstractAnnotations, ann::AnnotationDict, mapping, messages, marginals)

Write annotations into ann before a rule has executed. Called once per processor inside the MessageMapping callable, before the rule returns its result distribution.

source
ReactiveMP.post_rule_annotations! — Function
post_rule_annotations!(processor::AbstractAnnotations, ann::AnnotationDict, mapping, messages, marginals, result)

Write annotations into ann after a rule has executed. Called once per processor inside the MessageMapping callable, after the rule returns its result distribution.

source
ReactiveMP.post_product_annotations! — Function
post_product_annotations!(processor::AbstractAnnotations, merged::AnnotationDict, left_ann::AnnotationDict, right_ann::AnnotationDict, new_dist, left_dist, right_dist)

Write annotations into merged based on left_ann, right_ann, and the distributions involved in the message product. Called once per processor inside compute_product_of_two_messages.

source

Implementing a custom annotation processor

To add a new kind of annotation, subtype AbstractAnnotations and implement the two callbacks:

using ReactiveMP

# not exported by default
import ReactiveMP: AbstractAnnotations, AnnotationDict, has_annotation, get_annotation, annotate!

struct CountAnnotations <: AbstractAnnotations end

# Called before each rule execution
function ReactiveMP.pre_rule_annotations!(::CountAnnotations, ann::AnnotationDict, mapping, messages, marginals)
    return nothing
end

# Called after each rule execution
function ReactiveMP.post_rule_annotations!(::CountAnnotations, ann::AnnotationDict, mapping, messages, marginals, result)
    prev = has_annotation(ann, :count) ? get_annotation(ann, Int, :count) : 0
    annotate!(ann, :count, prev + 1)
    return nothing
end

# Called when two messages are multiplied
function ReactiveMP.post_product_annotations!(::CountAnnotations, merged::AnnotationDict, left_ann::AnnotationDict, right_ann::AnnotationDict, new_dist, left_dist, right_dist)
    left_count  = has_annotation(left_ann,  :count) ? get_annotation(left_ann,  Int, :count) : 0
    right_count = has_annotation(right_ann, :count) ? get_annotation(right_ann, Int, :count) : 0
    annotate!(merged, :count, left_count + right_count)
    return nothing
end

Processors are passed to FactorNodeActivationOptions (for rule-time annotation) and ReactiveMP.MessageProductContext (for product-time merging) when building a model. Both sites must be configured — see the RxInfer documentation for how to set this up at the model level.

Built-in annotation processors