From 3937f89c092aae9f76abff729fb9bf4b815b4784 Mon Sep 17 00:00:00 2001
From: Nikos Ignatiadis <nikos.ignatiadis01@gmail.com>
Date: Sun, 21 Aug 2022 10:24:45 +0300
Subject: [PATCH 1/3] analytical form for convolution and bind of two Normal
 distributions

---
 src/MeasureTheory.jl                        |  7 +++++-
 src/combinators/affine.jl                   | 15 ++++++++++-
 src/combinators/convolve.jl                 | 18 +++++++++++++
 src/parameterized/pairwise/normal_normal.jl | 13 ++++++++++
 test/runtests.jl                            | 28 +++++++++++++++++----
 5 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 src/combinators/convolve.jl
 create mode 100644 src/parameterized/pairwise/normal_normal.jl

diff --git a/src/MeasureTheory.jl b/src/MeasureTheory.jl
index 5be30634..c1769238 100644
--- a/src/MeasureTheory.jl
+++ b/src/MeasureTheory.jl
@@ -58,7 +58,8 @@ import MeasureBase:
     paramnames,
     ∫,
     𝒹,
-    ∫exp
+    ∫exp,
+    bind
 import MeasureBase: ≪
 using MeasureBase: BoundedInts, BoundedReals, CountingMeasure, IntegerDomain, IntegerNumbers
 using MeasureBase: weightedmeasure, restrict
@@ -127,6 +128,7 @@ include("parameterized.jl")
 
 include("macros.jl")
 include("combinators/affine.jl")
+include("combinators/convolve.jl")
 include("combinators/weighted.jl")
 include("combinators/product.jl")
 include("combinators/transforms.jl")
@@ -165,4 +167,7 @@ include("transforms/corrcholesky.jl")
 include("transforms/ordered.jl")
 
 include("distproxy.jl")
+
+include("parameterized/pairwise/normal_normal.jl")
+
 end # module
diff --git a/src/combinators/affine.jl b/src/combinators/affine.jl
index b5fe4312..0f902e98 100644
--- a/src/combinators/affine.jl
+++ b/src/combinators/affine.jl
@@ -36,7 +36,7 @@ import InverseFunctions: inverse
 @inline inverse(f::AffineTransform{(:λ,)}) = AffineTransform((σ = f.λ,))
 @inline inverse(f::AffineTransform{(:μ,)}) = AffineTransform((μ = -f.μ,))
 
-# `size(f) == (m,n)` means `f : ℝⁿ → ℝᵐ`  
+# `size(f) == (m,n)` means `f : ℝⁿ → ℝᵐ`
 Base.size(f::AffineTransform{(:μ, :σ)}) = size(f.σ)
 Base.size(f::AffineTransform{(:μ, :λ)}) = size(f.λ)
 Base.size(f::AffineTransform{(:σ,)}) = size(f.σ)
@@ -339,3 +339,16 @@ end
 @inline function Distributions.cdf(d::Affine, x)
     cdf(parent(d), inverse(d.f)(x))
 end
+
+function mean(d::Affine)
+    m = mean(parent(d))
+    f = getfield(d, :f)
+    return f(m)
+end
+
+# std only for univariate distributions
+std(d::Affine{(:μ,)}) = std(parent(d))
+std(d::Affine{(:σ,)}) = d.σ * std(parent(d))
+std(d::Affine{(:λ,)}) = d.λ \ std(parent(d))
+std(d::Affine{(:μ, :σ)}) = d.σ * std(parent(d))
+std(f::Affine{(:μ, :λ)}) = d.λ \ std(parent(d))
diff --git a/src/combinators/convolve.jl b/src/combinators/convolve.jl
new file mode 100644
index 00000000..7ea7d818
--- /dev/null
+++ b/src/combinators/convolve.jl
@@ -0,0 +1,18 @@
+struct Convolution{M,N} <: AbstractMeasure
+    μ::M
+    ν::N
+end
+
+export ↣
+
+"""
+If μ, ν are subtypes of `AbstractMeasure` or satisfy the Measure interface, r
+then `convolve(μ, ν)` is a measure, called the convolution of μ and ν.
+"""
+convolve(μ, ν) = Convolution(μ, ν)
+
+function Base.rand(rng::AbstractRNG, ::Type{T}, d::Convolution) where {T}
+    x = rand(rng, T, d.μ)
+    y = rand(rng, T, d.ν)
+    return x+y
+end
diff --git a/src/parameterized/pairwise/normal_normal.jl b/src/parameterized/pairwise/normal_normal.jl
new file mode 100644
index 00000000..e07147f6
--- /dev/null
+++ b/src/parameterized/pairwise/normal_normal.jl
@@ -0,0 +1,13 @@
+function convolve(μ::Normal, ν::Normal)
+    Normal(mean(μ) + mean(ν), hypot(std(μ), std(ν)))
+end
+
+function bind(μ::Normal,
+    ::ParameterizedTransitionKernel{Type{Normal{(:μ,)}}, typeof(identity), (:μ,), Tuple{typeof(identity)}})
+   convolve(μ, Normal())
+end
+
+function bind(μ::Normal,
+    k::ParameterizedTransitionKernel{Type{Normal{(:μ, :σ)}}, typeof(identity), (:μ, :σ), Tuple{typeof(identity), T}} where T<:Number)
+   convolve(μ, Normal(σ=k.param_maps.σ))
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index aaea8646..78763ac5 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -277,7 +277,7 @@ end
 
 @testset "Product of Diracs" begin
     x = randn(3)
-    t = as(productmeasure(Dirac.(x))) 
+    t = as(productmeasure(Dirac.(x)))
     @test transform(t, []) == x
 end
 
@@ -297,7 +297,7 @@ end
 
 #     chain = Chain(kernel, μ)
 
-#     dyniterate(iter::TimeLift, ::Nothing) = dyniterate(iter, 0=>nothing) 
+#     dyniterate(iter::TimeLift, ::Nothing) = dyniterate(iter, 0=>nothing)
 #     tr1 = trace(TimeLift(chain), nothing, u -> u[1] > 15)
 #     tr2 = trace(TimeLift(rand(Random.GLOBAL_RNG, chain)), nothing, u -> u[1] > 15)
 #     collect(Iterators.take(chain, 10))
@@ -348,8 +348,8 @@ end
     # NOTE: The `test_broken` below are mostly because of the change to `Affine`.
     # For example, `Normal{(:μ,:σ)}` is now `Affine{(:μ,:σ), Normal{()}}`.
     # The problem is not really with these measures, but with the tests
-    # themselves. 
-    # 
+    # themselves.
+    #
     # We should instead probably be doing e.g.
     # `D = typeof(Normal(μ=0.3, σ=4.1))`
 
@@ -652,7 +652,7 @@ end
     end
 
     x = rand(d)
-    
+
     @test logdensityof(d, x) isa Real
 end
 
@@ -660,3 +660,21 @@ end
     @test cdf(Normal(0, 1), 0) == 0.5
     @test cdf.((Normal(0, 1),), [0, 0]) == [0.5, 0.5]
 end
+
+@testset "pairwise normal-normal" begin
+    n0 = Normal(0, 0)
+    n1 = Normal(0.0, 1.0)
+    n2 = Normal(1.0, 2.0)
+    @test MeasureTheory.convolve(n1, n0) == n1
+    @test MeasureTheory.convolve(n1, n2) == MeasureTheory.convolve(n2, n1)
+
+    standard_normal_kernel = MeasureTheory.kernel(Normal{(:μ,)}, μ=identity)
+    σs = [0.0, 1.0, 2.0]
+    normal_kernels = [MeasureTheory.kernel(Normal{(:μ,:σ)}, μ=identity, σ=σ) for σ in σs]
+
+    @test (n1 ↣ standard_normal_kernel) == Normal(0.0, sqrt(2))
+    @test (n0 ↣ standard_normal_kernel) == n1
+
+    @test (n2 ↣ normal_kernels[1]) == n2
+    @test (n2 ↣ normal_kernels[2]) == (n2 ↣ standard_normal_kernel)
+end

From a91b323575cc7b498f3286cbb762e0418cdf81ce Mon Sep 17 00:00:00 2001
From: Nikos Ignatiadis <nikos.ignatiadis01@gmail.com>
Date: Sun, 21 Aug 2022 10:29:57 +0300
Subject: [PATCH 2/3] remove erroneous export

---
 src/combinators/convolve.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/combinators/convolve.jl b/src/combinators/convolve.jl
index 7ea7d818..3e920cc7 100644
--- a/src/combinators/convolve.jl
+++ b/src/combinators/convolve.jl
@@ -3,8 +3,6 @@ struct Convolution{M,N} <: AbstractMeasure
     ν::N
 end
 
-export ↣
-
 """
 If μ, ν are subtypes of `AbstractMeasure` or satisfy the Measure interface, r
 then `convolve(μ, ν)` is a measure, called the convolution of μ and ν.

From 06b893320cdf834c5560e6584c58e923dd3a2b8c Mon Sep 17 00:00:00 2001
From: Nikos Ignatiadis <nikos.ignatiadis01@gmail.com>
Date: Sun, 21 Aug 2022 10:44:52 +0300
Subject: [PATCH 3/3] fix two typos

---
 src/combinators/affine.jl   | 2 +-
 src/combinators/convolve.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/combinators/affine.jl b/src/combinators/affine.jl
index 0f902e98..4c8f45ad 100644
--- a/src/combinators/affine.jl
+++ b/src/combinators/affine.jl
@@ -351,4 +351,4 @@ std(d::Affine{(:μ,)}) = std(parent(d))
 std(d::Affine{(:σ,)}) = d.σ * std(parent(d))
 std(d::Affine{(:λ,)}) = d.λ \ std(parent(d))
 std(d::Affine{(:μ, :σ)}) = d.σ * std(parent(d))
-std(f::Affine{(:μ, :λ)}) = d.λ \ std(parent(d))
+std(d::Affine{(:μ, :λ)}) = d.λ \ std(parent(d))
diff --git a/src/combinators/convolve.jl b/src/combinators/convolve.jl
index 3e920cc7..039fa983 100644
--- a/src/combinators/convolve.jl
+++ b/src/combinators/convolve.jl
@@ -4,7 +4,7 @@ struct Convolution{M,N} <: AbstractMeasure
 end
 
 """
-If μ, ν are subtypes of `AbstractMeasure` or satisfy the Measure interface, r
+If μ, ν are subtypes of `AbstractMeasure` or satisfy the Measure interface,
 then `convolve(μ, ν)` is a measure, called the convolution of μ and ν.
 """
 convolve(μ, ν) = Convolution(μ, ν)