Make BilinearExponential into a Distribution

brenhinkeller · Jun 21, 2024 · cfae382 · cfae382
1 parent f0dc9b4
commit cfae382
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 26 deletions.
diff --git a/src/StratMetropolis.jl b/src/StratMetropolis.jl
@@ -446,19 +446,6 @@
     end
 
     adjust!(ages::AbstractVector, chronometer, systematic::Nothing) = ages
-    function adjust!(ages::AbstractVector{BilinearExponential{T}}, chronometer, systematic::SystematicUncertainty) where T
-        systUPb = randn()*systematic.UPb
-        systArAr = randn()*systematic.ArAr
-        @assert eachindex(ages)==eachindex(chronometer)
-        @inbounds for i ∈ eachindex(ages)
-            age = ages[i]
-            μ = age.μ
-            chronometer[i] === :UPb && (μ += systUPb)
-            chronometer[i] === :ArAr && (μ += systArAr)
-            ages[i] = BilinearExponential{T}(age.A, μ, age.σ, age.shp, age.skw)
-        end
-        return ages
-    end
     function adjust!(ages::AbstractVector, chronometer, systematic::SystematicUncertainty) where T
         systUPb = randn()*systematic.UPb
         systArAr = randn()*systematic.ArAr

diff --git a/src/Utilities.jl b/src/Utilities.jl
@@ -1,18 +1,52 @@
 ## --- Fitting non-Gaussian distributions
 
-    struct BilinearExponential{T}
+    struct BilinearExponential{T<:Real} <: ContinuousUnivariateDistribution
         A::T
         μ::T
         σ::T
-        sharpness::T
-        skew::T
+        shp::T
+        skw::T
     end
-    function BilinearExponential(p::AbstractVector)
+    function BilinearExponential(p::AbstractVector{T}) where {T}
         @assert length(p) == 5
         @assert all(x->!(x<0), Iterators.drop(p,1))
-        BilinearExponential(p...)
+        BilinearExponential{T}(p...)
+    end
+    function BilinearExponential(p::NTuple{5,T}) where {T}
+        @assert all(x->!(x<0), Iterators.drop(p,1))
+        BilinearExponential{T}(p...)
+    end
+
+    ## Conversions
+    Base.convert(::Type{BilinearExponential{T}}, d::Normal) where {T<:Real} = BilinearExponential{T}(T(d.A), T(d.μ), T(d.σ), T(d.shp), T(d.skw))
+    Base.convert(::Type{BilinearExponential{T}}, d::Normal{T}) where {T<:Real} = d
+
+    ## Parameters
+    Distributions.params(d::BilinearExponential) = (d.A, d.μ, d.σ, d.shp, d.skw)
+    @inline Distributions.partype(d::BilinearExponential{T}) where {T<:Real} = T
+
+    Distributions.location(d::BilinearExponential) = d.μ
+    Distributions.scale(d::BilinearExponential) = d.σ
+
+    Base.eltype(::Type{BilinearExponential{T}}) where {T} = T
+
+    ## Evaluation
+    function Distributions.pdf(d::BilinearExponential, x::Real)
+        xs = (x - d.μ)/d.σ # X scaled by mean and variance
+        v = 1/2 - atan(xs)/3.141592653589793 # Sigmoid (positive on LHS)
+        return exp(d.A + d.shp*d.skw*xs*v - d.shp/d.skw*xs*(1-v))
     end
 
+    function Distributions.logpdf(d::BilinearExponential, x::Real)
+        xs = (x - d.μ)/d.σ # X scaled by mean and variance
+        v = 1/2 - atan(xs)/3.141592653589793 # Sigmoid (positive on LHS)
+        return d.A + d.shp*d.skw*xs*v - d.shp/d.skw*xs*(1-v)
+    end
+
+    ## Affine transformations
+    Base.:+(d::BilinearExponential{T}, c::Real) where {T} = BilinearExponential{T}(d.A, d.μ + c, d.σ, d.shp, d.skw)
+    Base.:*(d::BilinearExponential{T}, c::Real) where {T} = BilinearExponential{T}(d.A, d.μ * c, d.σ * abs(c), d.shp, d.skw)
+
     """
     ```Julia
     bilinear_exponential(x::Number, p::AbstractVector)
@@ -95,7 +129,7 @@
             skw = abs(p[5,i])
             ll += shp*skw*xs*v - shp/skw*xs*(1-v)
         end
-        return ll
+        return lles
     end
     function bilinear_exponential_ll(x::AbstractVector, ages::AbstractVector{BilinearExponential{T}}) where T
         ll = zero(T)
@@ -104,7 +138,7 @@
             isnan(pᵢ.μ) && continue
             xs = (x[i]-pᵢ.μ)/(pᵢ.σ) # X scaled by mean and variance
             v = 1/2 - atan(xs)/3.141592653589793 # Sigmoid (positive on LHS)
-            shp, skw = pᵢ.sharpness, pᵢ.skew
+            shp, skw = pᵢ.shp, pᵢ.skw
             ll += shp*skw*xs*v - shp/skw*xs*(1-v)
         end
         return ll
@@ -146,12 +180,12 @@
 ## --- Biquadratic distribution (like bilinear, with quadratic sides)
 
 
-    struct BiquadraticExponential{T}
+    struct BiquadraticExponential{T<:Real}
         A::T
         μ::T
         σ::T
-        sharpness::T
-        skew::T
+        shp::T
+        skw::T
     end
     function BiquadraticExponential(p::AbstractVector)
         @assert length(p) == 5
@@ -191,8 +225,8 @@
             pᵢ = ages[i]
             xs = (x[i]-pᵢ.μ)/(pᵢ.σ) # X scaled by mean and variance
             v = 1/2 - atan(xs)/3.141592653589793 # Sigmoid (positive on LHS)
-            shp = pᵢ.sharpness
-            skw = pᵢ.skew
+            shp = pᵢ.shp
+            skw = pᵢ.skw
             ll -= (shp*skw*xs*v - shp/skw*xs*(1-v))^2
         end
         return ll

diff --git a/test/Project.toml b/test/Project.toml
@@ -1,3 +1,4 @@
 [deps]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,5 +1,5 @@
 using Chron
-using Test, Statistics
+using Test, Statistics, Distributions
 
 const make_plots = get(ENV, "MAKE_PLOTS", false) == "true"
 @testset "Utilities" begin include("testUtilities.jl") end

diff --git a/test/testUtilities.jl b/test/testUtilities.jl
@@ -4,11 +4,29 @@
     @test bilinear_exponential([66.02, 66.08, 66.15], [-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505]) ≈ [0.00010243952455548608, 1.0372066408907184e-13, 1.0569878258022094e-28]
     @test bilinear_exponential_ll(66.02, [-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505]) ≈ -3.251727600957773
 
+    d = Chron.BilinearExponential([-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505])
+    @test d isa Chron.BilinearExponential{Float64}
+    @test pdf(d, 66.02) ≈ 0.00010243952455548608
+    @test logpdf(d, 66.02) ≈ (-3.251727600957773 -5.9345103368858085)
+
+    @test eltype(d) === partype(d) === Float64
+    @test params(d) == (-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505)
+    @test d + 1 isa Chron.BilinearExponential{Float64}
+    @test location(d + 1) == 66.00606672179812 + 1
+    @test d * 2 isa Chron.BilinearExponential{Float64}
+    @test location(d * 2) == 66.00606672179812 * 2
+    @test scale(d * 2) == 0.17739474265630253 * 2
+
 ## --- Alternatively: biquadratic exponential
 
     @test Chron.biquadratic_exponential(66.02, [-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505]) ≈ 1.3384991282229901e-5
     @test Chron.biquadratic_exponential([66.02, 66.08, 66.15], [-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505]) ≈ [1.3384991282229901e-5, 5.441899612287734e-128, 0.0]
     @test Chron.biquadratic_exponential_ll(66.02, [-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505]) ≈ -10.573732390830594
 
 
+## --- Other utility functions for log likelihoods
+
+    @test Chron.strat_ll([0.0, 0.0], [Normal(0,1), Normal(0,1)]) ≈ 0
+    @test Chron.strat_ll([0.0, 0.5], [Normal(0,1), Uniform(0,1)]) ≈ -0.9189385332046728
+
 ## ---