Initial integration with ExponentialFamilyProjection

Project.toml

@@ -35,11 +35,13 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 Unrolled = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 
 [weakdeps]
+ExponentialFamilyProjection = "17f509fa-9a96-44ba-99b2-1c5f01f0931b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
 [extensions]
 ReactiveMPOptimisersExt = "Optimisers"
+ReactiveMPProjectionExt = "ExponentialFamilyProjection"
 ReactiveMPRequiresExt = "Requires"
 
 [compat]
@@ -50,6 +52,7 @@ Distributions = "0.24, 0.25"
 DomainIntegrals = "0.3.2, 0.4"
 DomainSets = "0.5.2, 0.6, 0.7"
 ExponentialFamily = "1.2.0"
+ExponentialFamilyProjection = "1.1"
 FastCholesky = "1.3.0"
 FastGaussQuadrature = "0.4, 0.5"
 FixedArguments = "0.1"
@@ -85,6 +88,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+ExponentialFamilyProjection = "17f509fa-9a96-44ba-99b2-1c5f01f0931b"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -96,4 +100,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestSetExtensions = "98d24dd4-01ad-11ea-1b02-c9a08f80db04"
 
 [targets]
-test = ["Aqua", "CpuId", "ReTestItems", "Test", "Pkg", "Logging", "InteractiveUtils", "TestSetExtensions", "Coverage", "Dates", "Distributed", "Documenter", "BenchmarkCI", "BenchmarkTools", "PkgBenchmark", "StableRNGs", "Optimisers", "DiffResults"]
+test = ["Aqua", "CpuId", "ReTestItems", "Test", "Pkg", "Logging", "InteractiveUtils", "TestSetExtensions", "Coverage", "Dates", "Distributed", "Documenter", "BenchmarkCI", "BenchmarkTools", "PkgBenchmark", "StableRNGs", "Optimisers", "DiffResults", "ExponentialFamilyProjection"]

docs/src/lib/nodes/delta.md

@@ -15,6 +15,7 @@ ReactiveMP.UT
 ReactiveMP.UnscentedTransform
 ReactiveMP.ProdCVI
 ReactiveMP.CVI
+ReactiveMP.CVIProjection
 ReactiveMP.cvi_setup!
 ReactiveMP.cvi_update!
 ReactiveMP.DeltaFnDefaultRuleLayout

docs/src/lib/rules.md

@@ -30,6 +30,13 @@ ReactiveMP.@test_rules
 ReactiveMP.@test_marginalrules
 ```
 
+## [Rule fallbacks](@id lib-rules-fallbacks)
+
+```@docs
+ReactiveMP.NodeFunctionRuleFallback
+ReactiveMP.nodefunction
+```
+
 ## [Table of available update rules](@id lib-rules-table)
 
 !!! note

ext/ReactiveMPProjectionExt/ReactiveMPProjectionExt.jl

@@ -0,0 +1,35 @@
+module ReactiveMPProjectionExt
+
+using ReactiveMP, ExponentialFamily, Distributions, ExponentialFamilyProjection, BayesBase, Random, LinearAlgebra, FastCholesky
+
+struct DivisionOf{A, B}
+    numerator::A
+    denumerator::B
+end
+
+BayesBase.insupport(d::DivisionOf, p) = insupport(d.numerator, p) && insupport(d.denumerator, p)
+BayesBase.logpdf(d::DivisionOf, p) = logpdf(d.numerator, p) - logpdf(d.denumerator, p)
+
+function BayesBase.prod(::GenericProd, something, division::DivisionOf) 
+    return prod(GenericProd(), division, something)
+end
+
+function BayesBase.prod(::GenericProd, division::DivisionOf, something)
+    if division.denumerator == something
+        return division.numerator
+    else
+        return ProductOf(division, something)
+    end
+end
+
+include("layout/cvi_projection.jl")
+include("rules/in.jl")
+include("rules/out.jl")
+include("rules/marginals.jl")
+
+# This will enable the extension and make `CVIProjection` compatible with delta nodes 
+# Otherwise it should throw an error suggesting users to install `ExponentialFamilyProjection`
+# See `approximations/cvi_projection.jl`
+ReactiveMP.is_delta_node_compatible(::ReactiveMP.CVIProjection) = Val(true)
+
+end

ext/ReactiveMPProjectionExt/layout/cvi_projection.jl

@@ -0,0 +1,85 @@
+
+using Rocket
+import ReactiveMP:
+    deltafn_rule_layout,
+    deltafn_apply_layout,
+    AbstractDeltaNodeDependenciesLayout,
+    DeltaFnDefaultRuleLayout,
+    DeltaFnNode,
+    getmarginal,
+    functionalform,
+    tag,
+    Marginalisation,
+    MessageMapping,
+    DefferedMessage,
+    with_statics,
+    apply_pipeline_stage,
+    messageout,
+    messagein,
+    connect!
+
+"""
+    CVIProjectionApproximationDeltaFnRuleLayout
+
+Custom rule layout for the Delta node in case of the CVI projection approximation method:
+
+# Layout 
+
+In order to compute:
+
+- `q_out`: mirrors the posterior marginal on the `out` edge
+- `q_ins`: uses inbound message on the `out` edge and all inbound messages on the `ins` edges
+- `m_out`: uses the posterior over `out`, message from `out` and the joint over the `ins` edges
+- `m_in_k`: uses the inbound message on the `in_k` edge and `q_ins`
+"""
+struct CVIProjectionApproximationDeltaFnRuleLayout <: AbstractDeltaNodeDependenciesLayout end
+
+deltafn_rule_layout(::DeltaFnNode, ::CVIProjection, inverse::Nothing) = CVIProjectionApproximationDeltaFnRuleLayout()
+
+function deltafn_rule_layout(::DeltaFnNode, ::CVIProjection, inverse::Any)
+    @warn "CVI projection approximation does not accept the inverse function. Ignoring the provided inverse."
+    return CVIProjectionApproximationDeltaFnRuleLayout()
+end
+
+# This function declares how to compute `q_out` locally around `DeltaFn`
+function deltafn_apply_layout(::CVIProjectionApproximationDeltaFnRuleLayout, ::Val{:q_out}, factornode::DeltaFnNode, meta, pipeline_stages, scheduler, addons, rulefallback)
+    return deltafn_apply_layout(DeltaFnDefaultRuleLayout(), Val(:q_out), factornode, meta, pipeline_stages, scheduler, addons, rulefallback)
+end
+
+# This function declares how to compute `q_ins` locally around `DeltaFn`
+function deltafn_apply_layout(::CVIProjectionApproximationDeltaFnRuleLayout, ::Val{:q_ins}, factornode::DeltaFnNode, meta, pipeline_stages, scheduler, addons, rulefallback)
+    return deltafn_apply_layout(DeltaFnDefaultRuleLayout(), Val(:q_ins), factornode, meta, pipeline_stages, scheduler, addons, rulefallback)
+end
+
+# This function declares how to compute `m_out` 
+function deltafn_apply_layout(::CVIProjectionApproximationDeltaFnRuleLayout, ::Val{:m_out}, factornode::DeltaFnNode, meta, pipeline_stages, scheduler, addons, rulefallback)
+    let interface = factornode.out
+        msgs_names      = Val{(:out,)}()
+        msgs_observable = combineLatestUpdates((messagein(factornode.out),), PushNew())
+
+        marginal_names       = Val{(:out, :ins)}()
+        marginals_observable = combineLatestUpdates((getmarginal(factornode.localmarginals.marginals[1]), getmarginal(factornode.localmarginals.marginals[2])), PushNew())
+
+        fform       = functionalform(factornode)
+        vtag        = tag(interface)
+        vconstraint = Marginalisation()
+
+        vmessageout = combineLatest((msgs_observable, marginals_observable), PushNew())
+
+        mapping = let messagemap = MessageMapping(fform, vtag, vconstraint, msgs_names, marginal_names, meta, addons, factornode, rulefallback)
+            (dependencies) -> DefferedMessage(dependencies[1], dependencies[2], messagemap)
+        end
+
+        vmessageout = with_statics(factornode, vmessageout)
+        vmessageout = vmessageout |> map(AbstractMessage, mapping)
+        vmessageout = apply_pipeline_stage(pipeline_stages, factornode, vtag, vmessageout)
+        vmessageout = vmessageout |> schedule_on(scheduler)
+
+        connect!(messageout(interface), vmessageout)
+    end
+end
+
+# This function declares how to compute `m_in` for each `k` 
+function deltafn_apply_layout(::CVIProjectionApproximationDeltaFnRuleLayout, ::Val{:m_in}, factornode::DeltaFnNode, meta, pipeline_stages, scheduler, addons, rulefallback)
+    return deltafn_apply_layout(DeltaFnDefaultRuleLayout(), Val(:m_in), factornode, meta, pipeline_stages, scheduler, addons, rulefallback)
+end

ext/ReactiveMPProjectionExt/rules/in.jl

@@ -0,0 +1,5 @@
+
+@rule DeltaFn((:in, k), Marginalisation) (q_ins::FactorizedJoint, m_in::Any, meta::DeltaMeta{M}) where {M <: CVIProjection} = begin
+    q_ins_k = component(q_ins, k)
+    return DivisionOf(q_ins_k, m_in)
+end

ext/ReactiveMPProjectionExt/rules/marginals.jl

@@ -0,0 +1,52 @@
+using TupleTools
+
+import Distributions: Distribution
+import BayesBase: AbstractContinuousGenericLogPdf
+
+@marginalrule DeltaFn(:ins) (m_out::Any, m_ins::ManyOf{1, Any}, meta::DeltaMeta{M}) where {M <: CVIProjection} = begin
+    method = ReactiveMP.getmethod(meta)
+    g = getnodefn(meta, Val(:out))
+
+    m_in = first(m_ins)
+    # Create an `AbstractContinuousGenericLogPdf` with an unspecified domain and the transformed `logpdf` function
+    F = promote_variate_type(variate_form(typeof(m_in)), BayesBase.AbstractContinuousGenericLogPdf)
+    f = convert(F, UnspecifiedDomain(), (z) -> logpdf(m_out, g(z)))
+
+    T = ExponentialFamily.exponential_family_typetag(m_in)
+    prj = ProjectedTo(T, size(m_in)...; parameters = something(method.prjparams, ExponentialFamilyProjection.DefaultProjectionParameters()))
+    q = project_to(prj, f, first(m_ins))
+
+    return FactorizedJoint((q,))
+end
+
+@marginalrule DeltaFn(:ins) (m_out::Any, m_ins::ManyOf{N, Any}, meta::DeltaMeta{M}) where {N, M <: CVIProjection} = begin
+    method = ReactiveMP.getmethod(meta)
+    rng = method.rng
+    pre_samples = zip(map(m_in_k -> ReactiveMP.cvilinearize(rand(rng, m_in_k, method.marginalsamples)), m_ins)...)
+
+    logp_nc_drop_index = let g = getnodefn(meta, Val(:out)), pre_samples = pre_samples
+        (z, i, pre_samples) -> begin
+            samples = map(ttuple -> ReactiveMP.TupleTools.insertat(ttuple, i, (z,)), pre_samples)
+            t_samples = map(s -> g(s...), samples)
+            logpdfs = map(out -> logpdf(m_out, out), t_samples)
+            return mean(logpdfs)
+        end
+    end
+
+    optimize_natural_parameters = let m_ins = m_ins, logp_nc_drop_index = logp_nc_drop_index
+        (i, pre_samples) -> begin
+            # Create an `AbstractContinuousGenericLogPdf` with an unspecified domain and the transformed `logpdf` function
+            df = let i = i, pre_samples = pre_samples, logp_nc_drop_index = logp_nc_drop_index
+                (z) -> logp_nc_drop_index(z, i, pre_samples)
+            end
+            logp = convert(promote_variate_type(variate_form(typeof(first(m_ins))), BayesBase.AbstractContinuousGenericLogPdf), UnspecifiedDomain(), df)
+
+            T = ExponentialFamily.exponential_family_typetag(m_ins[i])
+            prj = ProjectedTo(T, size(m_ins[i])...; parameters = something(method.prjparams, ExponentialFamilyProjection.DefaultProjectionParameters()))
+
+            return project_to(prj, logp, m_ins[i])
+        end
+    end
+
+    return FactorizedJoint(ntuple(i -> optimize_natural_parameters(i, pre_samples), length(m_ins)))
+end

ext/ReactiveMPProjectionExt/rules/out.jl

@@ -0,0 +1,63 @@
+using ForwardDiff
+
+# cost function
+function targetfn(M, p, data)
+    ef = convert(ExponentialFamilyDistribution, M, p)
+    return -sum((d) -> logpdf(ef, d), data)
+end
+
+# # gradient function
+## I think this is wrong. This is not a gradient on the manifolds. It is just Euclidean gradient.
+function grad_targetfn(M, p, data)
+    ef = convert(ExponentialFamilyDistribution, M, p)
+    ifisher = cholinv(Hermitian(fisherinformation(ef)))
+    X = ExponentialFamilyProjection.ExponentialFamilyManifolds.partition_point(M, ifisher * ForwardDiff.gradient((p) -> targetfn(M, p, data), p))
+    X = ExponentialFamilyProjection.ExponentialFamilyManifolds.ManifoldsBase.project(M, p, X)
+    return X
+end
+
+@rule DeltaFn(:out, Marginalisation) (m_out::Any, q_out::Any, q_ins::FactorizedJoint, meta::DeltaMeta{U}) where {U <: CVIProjection} = begin
+    node_function         = getnodefn(meta, Val(:out))
+    method                = ReactiveMP.getmethod(meta)
+    rng                   = method.rng
+    q_ins_components      = components(q_ins)
+    dimensions            = map(size, q_ins_components)
+    q_ins_sample_friendly = map(q_in -> sampling_optimized(q_in), q_ins_components)
+    ## Option 1
+    # samples               = map(i -> collect(map(q -> rand(rng, q), q_ins_sample_friendly)), 1:method.out_samples_no)
+    # q_out_samples         = map(sample -> node_function(ReactiveMP.__splitjoin(sample, dimensions)...), samples)
+
+    ## Option 2
+    samples       = map(ReactiveMP.cvilinearize, map(q_in -> rand(rng, q_in, method.outsamples), q_ins_sample_friendly))
+    q_out_samples = map(x -> node_function(x...), zip(samples...))
+
+    ## Option 3
+    # T = ExponentialFamily.exponential_family_typetag(q_out)
+    # s = sampling_optimized(q_out)
+    # d = fit_mle(typeof(s), q_out_samples)
+    # m = DivisionOf(d, m_out)
+    # r = project_to(ProjectedTo(T, size(q_out)...; parameters = method.prjparams), (x) -> logpdf(m, x))
+    # return r
+
+    T           = ExponentialFamily.exponential_family_typetag(q_out)
+    q_out_ef    = convert(ExponentialFamilyDistribution, q_out)
+    conditioner = getconditioner(q_out_ef)
+    manifold    = ExponentialFamilyProjection.ExponentialFamilyManifolds.get_natural_manifold(T, size(mean(q_out_ef)), conditioner)
+    nat_params  = ExponentialFamilyProjection.ExponentialFamilyManifolds.partition_point(manifold, getnaturalparameters(q_out_ef))
+
+    f = (M, p) -> targetfn(M, p, q_out_samples)
+    g = (M, p) -> grad_targetfn(M, p, q_out_samples)
+
+    est = convert(
+        ExponentialFamilyDistribution,
+        manifold,
+        ExponentialFamilyProjection.Manopt.gradient_descent(
+            manifold, f, g, nat_params; 
+            stepsize = ExponentialFamilyProjection.Manopt.ConstantStepsize(0.1), 
+            direction = ExponentialFamilyProjection.BoundedNormUpdateRule(1),
+            debug = missing
+        )
+    )
+    # return x -> logpdf(est, x) - logpdf(m_out, x)
+    return DivisionOf(est, m_out)
+end

src/ReactiveMP.jl

@@ -71,6 +71,7 @@ include("approximations/rts.jl")
 include("approximations/linearization.jl")
 include("approximations/unscented.jl")
 include("approximations/cvi.jl")
+include("approximations/cvi_projection.jl")
 
 # Equality node is a special case and needs to be included before random variable implementation
 include("nodes/equality.jl")
@@ -96,6 +97,7 @@ include("score/node.jl")
 
 include("nodes/predefined.jl")
 include("rules/predefined.jl")
+include("rules/fallbacks.jl")
 
 # This symbol is only defined on Julia versions that support extensions
 @static if !isdefined(Base, :get_extension)
@@ -109,6 +111,7 @@ function __init__()
     @static if !isdefined(Base, :get_extension)
         @require Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" include("../ext/ReactiveMPOptimisersExt/ReactiveMPOptimisersExt.jl")
         @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("../ext/ReactiveMPZygoteExt/ReactiveMPZygoteExt.jl")
+        @require ExponentialFamilyProjection = "17f509fa-9a96-44ba-99b2-1c5f01f0931b" include("../ext/ReactiveMPProjectionExt/ReactiveMPProjectionExt.jl")
     end
 end
 

src/approximations/cvi.jl

@@ -33,6 +33,9 @@ cvilinearize(matrix::AbstractMatrix) = eachcol(matrix)
 The `ProdCVI` structure defines the approximation method hyperparameters of the `prod(approximation::CVI, logp::F, dist)`.
 This method performs an approximation of the product of the `dist` and `logp` with Stochastic Variational message passing (SVMP-CVI) (See [`Probabilistic programming with stochastic variational message passing`](https://biaslab.github.io/publication/probabilistic_programming_with_stochastic_variational_message_passing/)).
 
+!!! note
+    `ProdCVI` is deprecated in favor of `CVIProjection`.
+
 Arguments
  - `rng`: random number generator
  - `n_samples`: number of samples to use for statistics approximation
@@ -74,6 +77,8 @@ end
 """Alias for the `ProdCVI` method. See help for [`ProdCVI`](@ref)"""
 const CVI = ProdCVI
 
+is_delta_node_compatible(::ProdCVI) = Val(true)
+
 #---------------------------
 # CVI implementations
 #---------------------------

src/approximations/cvi_projection.jl

@@ -0,0 +1,32 @@
+export CVIProjection
+
+"""
+    CVIProjection(; parameters...)
+
+A structure representing the parameters for the Conjugate Variational Inference (CVI) projection method. 
+This structure is a subtype of `AbstractApproximationMethod` and is used to configure the settings for CVI.
+
+!!! note
+    The `CVIProjection` method requires `ExponentialFamilyProjection` package installed in the current environment.
+
+# Parameters
+
+- `rng::R`: The random number generator used for sampling. Default is `Random.MersenneTwister(42)`.
+- `marginalsamples::S`: The number of samples used for approximating marginal distributions. Default is `10`.
+- `outsamples::S`: The number of samples used for approximating output message distributions. Default is `100`.
+- `prjparams::P`: Parameters for the exponential family projection. Default is `nothing`, in which case it will use `ExponentialFamilyProjection.DefaultProjectionParameters()`.
+
+!!! note
+    The `CVIProjection` method is an experimental enhancement of the now-deprecated `CVI`, offering better stability and improved accuracy. 
+    Note that the parameters of this structure, as well as their defaults, are subject to change during the experimentation phase.
+"""
+Base.@kwdef struct CVIProjection{R, S, P} <: AbstractApproximationMethod
+    rng::R = Random.MersenneTwister(42)
+    marginalsamples::S = 10
+    outsamples::S = 100
+    prjparams::P = nothing # ExponentialFamilyProjection.DefaultProjectionParameters()
+end
+
+# This method should only be invoked if a user did not install `ExponentialFamilyProjection`
+# in the current Julia session
+check_delta_node_compatibility(::Val{false}, ::CVIProjection) = error("CVI projection requires `using ExponentialFamilyProjection` in the current session.")

src/approximations/linearization.jl

@@ -15,6 +15,8 @@ The `Linearization` structure is used inside the `DeltaMeta` or `FlowMeta` struc
 """
 struct Linearization <: AbstractApproximationMethod end
 
+is_delta_node_compatible(::Linearization) = Val(true)
+
 # ported from ForneyLab.jl
 using ForwardDiff
 

src/approximations/unscented.jl

@@ -62,6 +62,8 @@ const UT = Unscented
 """An alias for the [`Unscented`](@ref) approximation method."""
 const UnscentedTransform = Unscented
 
+is_delta_node_compatible(::Unscented) = Val(true)
+
 # get-functions for the Unscented structure
 
 getα(approximation::Unscented) = approximation.α