From c9ee3a4f07c43feba1a18e0d765ba9efdb6155fb Mon Sep 17 00:00:00 2001
From: "C. Brenhin Keller" <cbkeller@dartmouth.edu>
Date: Mon, 24 Jun 2024 16:28:53 -0400
Subject: [PATCH] Convert `Radiocarbon` struct to be a `Distribution`

---
 Project.toml          |  2 +-
 src/Utilities.jl      | 55 +++++++++++++++++++++++++++++++++++++++----
 test/testUtilities.jl | 15 ++++++++++++
 3 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index a3bff65..be7fa64 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Chron"
 uuid = "68885b1f-77b5-52a7-b2e7-6a8014c56b98"
 authors = ["C. Brenhin Keller <cbkeller@dartmouth.edu>"]
-version = "0.5.2"
+version = "0.5.3"
 
 [deps]
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
diff --git a/src/Utilities.jl b/src/Utilities.jl
index 25b3ddc..2825421 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -86,25 +86,70 @@
 
 ## --- Radiocarbon distribution type
 
-    struct Radiocarbon{T}
+    struct Radiocarbon{T<:Real} <: ContinuousUnivariateDistribution
         μ::T
         σ::T
         dist::Vector{T}
+        ldist::Vector{T}
     end
 
+    function Radiocarbon(μ::T, σ::T, ldist::Vector{T}) where {T<:Real}
+        dist = exp.(ldist)
+        dist ./= sum(dist) * 1 # Normalize
+        return Radiocarbon{T}(μ, σ, dist, ldist)
+    end
+
+    function Radiocarbon(Age_14C::Real, Age_14C_sigma::Real, calibration::NamedTuple=intcal13)
+        @assert calibration.Age_Calendar == 1:1:length(calibration.Age_14C)
+        @assert step(calibration.Age_Calendar) == calibration.dt == 1
+
+        ldist = normproduct_ll.(Age_14C, Age_14C_sigma, calibration.Age_14C, calibration.Age_sigma)
+
+        dist = exp.(ldist)
+        dist ./= sum(dist) * calibration.dt # Normalize
+        samples = draw_from_distribution(dist, 10^6) .* maximum(calibration.Age_Calendar)
+        μ, σ = nanmean(samples), nanstd(samples)
+
+        return Radiocarbon(μ, σ, ldist)
+    end
+
+    ## Conversions
+    Base.convert(::Type{Radiocarbon{T}}, d::Normal) where {T<:Real} = Radiocarbon{T}(T(d.μ), T(d.σ), T.(d.ldist))
+    Base.convert(::Type{Radiocarbon{T}}, d::Normal{T}) where {T<:Real} = d
+
+    ## Parameters
+    Distributions.params(d::Radiocarbon) = (d.dist, d.ldist)
+    @inline Distributions.partype(d::Radiocarbon{T}) where {T<:Real} = T
+
+    Distributions.location(d::Radiocarbon) = d.μ
+    Distributions.scale(d::Radiocarbon) = d.σ
+
+    Base.eltype(::Type{Radiocarbon{T}}) where {T} = T
+
+    ## Evaluation
+    @inline function Distributions.pdf(d::Radiocarbon{T}, x::Real) where {T}
+        return linterp_at_index(d.dist, x, zero(T))
+    end
+
+    @inline function Distributions.logpdf(d::Radiocarbon{T}, x::Real) where {T}
+        return linterp_at_index(d.ldist, x, -maxintfloat(T))
+    end
+
+## --- 
+
     # Interpolate log likelihood from an array
-    function interpolate_ll(x::AbstractVector,p::AbstractMatrix{T}) where {T<:Number}
+    function interpolate_ll(x::AbstractVector,p::AbstractMatrix{T}) where {T<:Real}
         ll = zero(T)
         @inbounds for i ∈ eachindex(x)
-            ll += linterp_at_index(view(p,:,i), x[i], -1e9)
+            ll += linterp_at_index(view(p,:,i), x[i], -maxintfloat(T))
         end
         return ll
     end
     function interpolate_ll(x::AbstractVector,ages::AbstractVector{Radiocarbon{T}}) where T
         ll = zero(T)
         @inbounds for i ∈ eachindex(x,ages)
-            dist = ages[i].dist
-            ll += linterp_at_index(dist, x[i], -1e9)
+            ldist = ages[i].ldist
+            ll += linterp_at_index(ldist, x[i], -maxintfloat(T))
         end
         return ll
     end
diff --git a/test/testUtilities.jl b/test/testUtilities.jl
index 9bfaa2d..ed83db1 100644
--- a/test/testUtilities.jl
+++ b/test/testUtilities.jl
@@ -14,12 +14,27 @@
     @test location(d * 2) == 66.00606672179812 * 2
     @test scale(d * 2) == 0.17739474265630253 * 2
 
+
+## --- Test Radiocarbon distribution
+
+    d = Radiocarbon(1000, 10, intcal13)
+    @test d isa Radiocarbon{Float64}
+    @test pdf(d, 950) ≈ 0.0036210273644940918
+    @test logpdf(d, 950) ≈ -3.4446721311475414 
+    @test pdf(d, [900, 950, 1000]) ≈ [3.372067127114498e-7, 0.0036210273644940918, 1.5582187464154278e-12]
+
+    @test eltype(d) === partype(d) === Float64
+    @test location(d) ≈ 926.3094740428785 atol=0.1
+    @test scale(d) ≈ 7.411299532288234 atol=0.1
+
 ## --- Other utility functions for log likelihoods
 
     @test strat_ll(66.02, BilinearExponential(-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505)) ≈  (-3.251727600957773 -5.9345103368858085)
+    @test strat_ll(950, Radiocarbon(1000, 10, intcal13)) ≈ -3.4446721311475414 
 
     @test strat_ll([0.0, 0.0], [Normal(0,1), Normal(0,1)]) ≈ 0
     @test strat_ll([0.0, 0.5], [Normal(0,1), Uniform(0,1)]) ≈ -0.9189385332046728
     @test strat_ll([0.0, 0.5, 66.02], [Normal(0,1), Uniform(0,1), BilinearExponential(-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505)]) ≈ (-0.9189385332046728 -3.251727600957773 -5.9345103368858085)
+    @test strat_ll([0.0, 0.5, 66.02, 900], [Normal(0,1), Uniform(0,1), BilinearExponential(-5.9345103368858085, 66.00606672179812, 0.17739474265630253, 70.57882331291309, 0.6017142541555505), Radiocarbon(1000, 10, intcal13)]) ≈ -22.831420814939655
 
 ## ---