k-dimensional OU process

src/continuous.jl

@@ -1,31 +1,40 @@
-#OU process
-struct OrnsteinUhlenbeckDiffusion{T <: Real} <: GaussianStateProcess
+# OU process
+struct OrnsteinUhlenbeckDiffusion{T} <: GaussianStateProcess
     mean::T
     volatility::T
     reversion::T
 end
 
-OrnsteinUhlenbeckDiffusion(mean::Real, volatility::Real, reversion::Real) = OrnsteinUhlenbeckDiffusion(float.(promote(mean, volatility, reversion))...)
+function OrnsteinUhlenbeckDiffusion(mean::Real, volatility::Real, reversion::Real)
+    μ, σ, θ = float.(promote(mean, volatility, reversion))
+    return OrnsteinUhlenbeckDiffusion{typeof(μ)}(μ, σ, θ)
+end
 
 OrnsteinUhlenbeckDiffusion(mean::T) where T <: Real = OrnsteinUhlenbeckDiffusion(mean,T(1.0),T(0.5))
 
 var(model::OrnsteinUhlenbeckDiffusion) = (model.volatility^2) / (2 * model.reversion)
 
 eq_dist(model::OrnsteinUhlenbeckDiffusion) = Normal(model.mean,sqrt(var(model)))
 
-function forward(process::OrnsteinUhlenbeckDiffusion, x_s::AbstractArray, s::Real, t::Real)
+# These are for nested broadcasting
+elmwisesqrt(x) = sqrt.(x)
+elmwiseinv(x) = inv.(x)
+elmwisemul(x, y) = x .* y
+elmwisediv(x, y) = x ./ y
+
+function forward(process::OrnsteinUhlenbeckDiffusion{T}, x_s::AbstractArray{T}, s::Real, t::Real) where T
     μ, σ, θ = process.mean, process.volatility, process.reversion
-    mean = @. exp(-(t - s) * θ) * (x_s - μ) + μ
+    mean = elmwisemul.((exp.(-(t - s) * θ),), x_s .- (μ,)) .+ (μ,)
     var = similar(mean)
-    var .= ((1 - exp(-2(t - s) * θ)) * σ^2) / 2θ
+    fill!(var, @. ((1 - exp(-2(t - s) * θ)) * σ^2) / 2θ)
     return GaussianVariables(mean, var)
 end
 
-function backward(process::OrnsteinUhlenbeckDiffusion, x_t::AbstractArray, s::Real, t::Real)
+function backward(process::OrnsteinUhlenbeckDiffusion{T}, x_t::AbstractArray{T}, s::Real, t::Real) where T
     μ, σ, θ = process.mean, process.volatility, process.reversion
-    mean = @. exp((t - s) * θ) * (x_t - μ) + μ
+    mean = elmwisemul.((exp.((t - s) * θ),), x_t .- (μ,)) .+ (μ,)
     var = similar(mean)
-    var .= -(σ^2 / 2θ) + (σ^2 * exp(2(t - s) * θ)) / 2θ
+    fill!(var, @. -(σ^2 / 2θ) + (σ^2 * exp(2(t - s) * θ)) / 2θ)
     return (μ = mean, σ² = var)
 end
 

src/randomvariable.jl

@@ -9,11 +9,12 @@ end
 
 Base.size(X::GaussianVariables) = size(X.μ)
 
-sample(rng::AbstractRNG, X::GaussianVariables{T}) where T = randn(rng, T, size(X)) .* .√X.σ² .+ X.μ
+sample(rng::AbstractRNG, X::GaussianVariables{T}) where T =
+    elmwisemul.(randn(rng, T, size(X)), elmwisesqrt.(X.σ²)) .+ X.μ
 
 function combine(X::GaussianVariables, lik)
-    σ² = @. inv(inv(X.σ²) + inv(lik.σ²))
-    μ = @. σ² * (X.μ / X.σ² + lik.μ / lik.σ²)
+    σ² = elmwiseinv.(elmwiseinv.(X.σ²) .+ elmwiseinv.(lik.σ²))
+    μ = elmwisemul.(σ², elmwisediv.(X.μ, X.σ²) .+ elmwisediv.(lik.μ, lik.σ²))
     return GaussianVariables(μ, σ²)
 end
 

test/runtests.jl

@@ -74,6 +74,16 @@ end
     x = one(QuatRotation{Float32})
     t = 0.29999998f0
     @test sampleforward(diffusion, t, [x]) isa Vector
+
+    # three-dimensional diffusion
+    μ = @SVector [0.0, 0.0, 0.0]
+    θ = @SVector [1.0, 1.0, 1.0]
+    σ = @SVector [0.5, 0.5, 0.5]
+    x_0 = [zero(μ), zero(μ)]
+    diffusion = OrnsteinUhlenbeckDiffusion(μ, σ, θ)
+    x_t = sampleforward(diffusion, 1.0, x_0)
+    @test x_t isa typeof(x_0)
+    @test size(x_t) == size(x_0)
 end
 
 @testset "Discrete Diffusions" begin
@@ -175,6 +185,15 @@ end
     x = samplebackward((x, t) -> x + randn(size(x)), process, [1/8, 1/4, 1/2, 1/1], x_t)
     @test size(x) == size(x_t)
     @test x isa Matrix
+
+    μ = @SVector [0.0, 0.0, 0.0]
+    θ = @SVector [1.0, 1.0, 1.0]
+    σ = @SVector [0.5, 0.5, 0.5]
+    x_t = randn(typeof(μ), 4, 10)
+    process = OrnsteinUhlenbeckDiffusion(μ, σ, θ)
+    x = samplebackward((x, t) -> x + randn(eltype(x), size(x)), process, [1/8, 1/4, 1/2, 1/1], x_t)
+    @test size(x) == size(x_t)
+    @test x isa Matrix
 end
 
 @testset "Masked Diffusion" begin