Generalize Unconformity type to Constraint

OpenThermochronology · Nov 4, 2024 · 1f0dde9 · 1f0dde9 · brenhinkeller · Nov 4, 2024
1 parent 789abf2
commit 1f0dde9
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 59 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Thermochron"
 uuid = "b9a8379e-ee1a-4388-b7ca-7852af1cef08"
 authors = ["C. Brenhin Keller and collaborators"]
-version = "0.9.3"
+version = "0.9.4"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

diff --git a/examples/ZrnHeInversionVartCryst.jl b/examples/ZrnHeInversionVartCryst.jl
@@ -154,18 +154,16 @@
         ΔT = Float64[model.ΔTnow, model.ΔTinit],
     )
 
-    # Default: No unconformity is imposed
-    unconf = Unconformity();
+    # Default: No constraints are imposed
+    constraint = Constraint()
 
     # # Uncomment this section if you wish to impose an unconformity at any point in the record
     # # Uniform distributions from Age₀ to Age₀+ΔAge, T₀ to T₀+ΔT,
-    # unconf = Unconformity(
-    #     agepoints = Float64[550.0,],  # Ma
-    #     Tpoints = Float64[20.0,],     # Degrees C
-    #     Age₀ = Float64[500,],
-    #     ΔAge = Float64[80,],
-    #     T₀ = Float64[0,],
-    #     ΔT = Float64[40,],
+    # constraint = Constraint(
+    #     Age₀ = Float64[500,],         # [Ma] Minimum age
+    #     ΔAge = Float64[80,],          # [Ma] Duration of age range
+    #     T₀ = Float64[0,],             # [C] Minimum temperature
+    #     ΔT = Float64[40,],            # [C] Width of temperature range
     # )
     # name *= "_unconf"
 
@@ -182,8 +180,8 @@
     Tpoints[1:npoints] .= Tr  # Degrees C
 
     # Run Markov Chain
-    @time (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, unconf, detail)
-    # @time (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC_varkinetics(data, model, npoints, agepoints, Tpoints, boundary, unconf, detail)
+    @time (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, constraint, detail)
+    # @time (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC_varkinetics(data, model, npoints, agepoints, Tpoints, boundary, constraint, detail)
     @info """tpointdist & Tpointdist collected, size: $(size(Tpointdist))
     Mean log-likelihood: $(nanmean(view(lldist, model.burnin:model.nsteps)))
     Mean acceptance rate: $(nanmean(view(acceptancedist, model.burnin:model.nsteps)))

diff --git a/src/inversion.jl b/src/inversion.jl
@@ -1,18 +1,18 @@
 
     """
     ```julia
-    MCMC(data::NamedTuple, model::NamedTuple, npoints::Int, agepoints::Vector, Tpoints::Vector, unconf::Unconformity, boundary::Boundary, [detail::DetailInterval])
+    MCMC(data::NamedTuple, model::NamedTuple, npoints::Int, agepoints::Vector, Tpoints::Vector, constraint::Constraint, boundary::Boundary, [detail::DetailInterval])
     ```
     Markov chain Monte Carlo time-Temperature inversion of the data specified in `data` and model parameters specified by `model`.
 
     Returns a tuple of distributions `(tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist)`
 
     ## Examples
     ```julia
-    tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist = MCMC(data, model, npoints, agepoints, Tpoints, unconf, boundary)
+    tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist = MCMC(data, model, npoints, agepoints, Tpoints, constraint, boundary)
     ```
     """
-    function MCMC(data::NamedTuple, model::NamedTuple, npoints::Int, agepoints::DenseVector{T}, Tpoints::DenseVector{T}, boundary::Boundary{T}, unconf::Unconformity{T}=Unconformity(T), detail::DetailInterval{T}=DetailInterval(T)) where T <: AbstractFloat
+    function MCMC(data::NamedTuple, model::NamedTuple, npoints::Int, agepoints::DenseVector{T}, Tpoints::DenseVector{T}, boundary::Boundary{T}, constraint::Constraint{T}=Constraint(T), detail::DetailInterval{T}=DetailInterval(T)) where T <: AbstractFloat
         # Sanitize inputs
         @assert firstindex(agepoints) === 1
         @assert firstindex(Tpoints) === 1
@@ -23,10 +23,10 @@
         U = T.(data.U)::DenseVector{T}
         Th = T.(data.Th)::DenseVector{T}
         Sm = (haskey(data, :Sm) ? T.(data.Th) : zeros(T, size(U)))::DenseVector{T}
-        nsteps = model.nsteps::Int
-        maxpoints = model.maxpoints::Int
+        nsteps = (haskey(model, :nsteps) ? model.nsteps : 10^6)::Int
+        maxpoints = (haskey(model, :maxpoints) ? model.maxpoints : 50)::Int
         minpoints = (haskey(model, :minpoints) ? model.minpoints : 1)::Int
-        totalpoints = maxpoints + boundary.npoints + unconf.npoints::Int
+        totalpoints = maxpoints + boundary.npoints + constraint.npoints::Int
         simplified = (haskey(model, :simplified) ? model.simplified : false)::Bool
         boundarytype = (haskey(model, :boundarytype) ? model.boundarytype : :hard)::Symbol
         dynamicjumping = (haskey(model, :dynamicjumping) ? model.dynamicjumping : false)::Bool
@@ -50,8 +50,8 @@
         knot_index = similar(agesteps, Int)::DenseVector{Int}
 
         # Calculate model ages for initial proposal
-        ages = collectto!(agepointbuffer, view(agepoints, 1:npoints), boundary.agepoints, unconf.agepoints)::StridedVector{T}
-        temperatures = collectto!(Tpointbuffer, view(Tpoints, 1:npoints), boundary.Tpoints, unconf.Tpoints)::StridedVector{T}
+        ages = collectto!(agepointbuffer, view(agepoints, 1:npoints), boundary.agepoints, constraint.agepoints)::StridedVector{T}
+        temperatures = collectto!(Tpointbuffer, view(Tpoints, 1:npoints), boundary.Tpoints, constraint.Tpoints)::StridedVector{T}
         Tsteps = linterp1s(ages, temperatures, agesteps)::DenseVector{T}
         calcHeAges = similar(HeAge)::DenseVector{T}
 
@@ -149,8 +149,8 @@
             npointsₚ = npoints
             copyto!(agepointsₚ, agepoints)
             copyto!(Tpointsₚ, Tpoints)
-            copyto!(unconf.agepointsₚ, unconf.agepoints)
-            copyto!(unconf.Tpointsₚ, unconf.Tpoints)
+            copyto!(constraint.agepointsₚ, constraint.agepoints)
+            copyto!(constraint.Tpointsₚ, constraint.Tpoints)
             copyto!(boundary.Tpointsₚ, boundary.Tpoints)
 
             # Adjust the proposal
@@ -178,15 +178,15 @@
                 @. boundary.Tpointsₚ = boundary.T₀ + rand()*boundary.ΔT
 
                 # If there's an imposed unconformity, adjust within parameters
-                if unconf.npoints > 0
-                    @. unconf.agepointsₚ = unconf.Age₀ + rand()*unconf.ΔAge
-                    @. unconf.Tpointsₚ = unconf.T₀ + rand()*unconf.ΔT
+                if constraint.npoints > 0
+                    @. constraint.agepointsₚ = constraint.Age₀ + rand()*constraint.ΔAge
+                    @. constraint.Tpointsₚ = constraint.T₀ + rand()*constraint.ΔT
                 end
             end
 
             # Recalculate interpolated proposed t-T path
-            ages = collectto!(agepointbuffer, view(agepointsₚ, 1:npointsₚ), boundary.agepoints, unconf.agepointsₚ)::StridedVector{T}
-            temperatures = collectto!(Tpointbuffer, view(Tpointsₚ, 1:npointsₚ), boundary.Tpointsₚ, unconf.Tpointsₚ)::StridedVector{T}
+            ages = collectto!(agepointbuffer, view(agepointsₚ, 1:npointsₚ), boundary.agepoints, constraint.agepointsₚ)::StridedVector{T}
+            temperatures = collectto!(Tpointbuffer, view(Tpointsₚ, 1:npointsₚ), boundary.Tpointsₚ, constraint.Tpointsₚ)::StridedVector{T}
             linterp1s!(Tsteps, knot_index, ages, temperatures, agesteps)
 
             # Old version of imposing max reheating rate, for reference:
@@ -242,8 +242,8 @@
                 npoints = npointsₚ
                 copyto!(agepoints, 1, agepointsₚ, 1, npoints)
                 copyto!(Tpoints, 1, Tpointsₚ, 1, npoints)
-                copyto!(unconf.agepoints, unconf.agepointsₚ)
-                copyto!(unconf.Tpoints, unconf.Tpointsₚ)
+                copyto!(constraint.agepoints, constraint.agepointsₚ)
+                copyto!(constraint.Tpoints, constraint.Tpointsₚ)
                 copyto!(boundary.Tpoints, boundary.Tpointsₚ)
                 copyto!(calcHeAges, calcHeAgesₚ)
 
@@ -259,8 +259,8 @@
             HeAgedist[:,n] .= calcHeAges # distribution of He ages
 
             # This is the actual output we want -- the distribution of t-T paths (t path is always identical)
-            collectto!(view(tpointdist, :, n), view(agepoints, 1:npoints), boundary.agepoints, unconf.agepoints)
-            collectto!(view(Tpointdist, :, n), view(Tpoints, 1:npoints), boundary.Tpoints, unconf.Tpoints)
+            collectto!(view(tpointdist, :, n), view(agepoints, 1:npoints), boundary.agepoints, constraint.agepoints)
+            collectto!(view(Tpointdist, :, n), view(Tpoints, 1:npoints), boundary.Tpoints, constraint.Tpoints)
 
             # Update progress meter every `progress_interval` steps
             (mod(n, progress_interval) == 0) && update!(progress, n)
@@ -269,7 +269,7 @@
     end
     export MCMC
 
-    function MCMC_varkinetics(data::NamedTuple, model::NamedTuple, npoints::Int, agepoints::DenseVector{T}, Tpoints::DenseVector{T}, boundary::Boundary{T}, unconf::Unconformity{T}=Unconformity(T), detail::DetailInterval{T}=DetailInterval(T)) where T <: AbstractFloat
+    function MCMC_varkinetics(data::NamedTuple, model::NamedTuple, npoints::Int, agepoints::DenseVector{T}, Tpoints::DenseVector{T}, boundary::Boundary{T}, constraint::Constraint{T}=Constraint(T), detail::DetailInterval{T}=DetailInterval(T)) where T <: AbstractFloat
         # Sanitize inputs
         @assert firstindex(agepoints) === 1
         @assert firstindex(Tpoints) === 1
@@ -280,10 +280,10 @@
         U = T.(data.U)::DenseVector{T}
         Th = T.(data.Th)::DenseVector{T}
         Sm = (haskey(data, :Sm) ? T.(data.Th) : zeros(T, size(U)))::DenseVector{T}
-        nsteps = model.nsteps::Int
-        maxpoints = model.maxpoints::Int
+        nsteps = (haskey(model, :nsteps) ? model.nsteps : 10^6)::Int
+        maxpoints = (haskey(model, :maxpoints) ? model.maxpoints : 50)::Int
         minpoints = (haskey(model, :minpoints) ? model.minpoints : 1)::Int
-        totalpoints = maxpoints + boundary.npoints + unconf.npoints::Int
+        totalpoints = maxpoints + boundary.npoints + constraint.npoints::Int
         simplified = (haskey(model, :simplified) ? model.simplified : false)::Bool
         boundarytype = (haskey(model, :boundarytype) ? model.boundarytype : :hard)::Symbol
         dynamicjumping = (haskey(model, :dynamicjumping) ? model.dynamicjumping : false)::Bool
@@ -307,8 +307,8 @@
         knot_index = similar(agesteps, Int)::DenseVector{Int}
 
         # Calculate model ages for initial proposal
-        ages = collectto!(agepointbuffer, view(agepoints, 1:npoints), boundary.agepoints, unconf.agepoints)::StridedVector{T}
-        temperatures = collectto!(Tpointbuffer, view(Tpoints, 1:npoints), boundary.Tpoints, unconf.Tpoints)::StridedVector{T}
+        ages = collectto!(agepointbuffer, view(agepoints, 1:npoints), boundary.agepoints, constraint.agepoints)::StridedVector{T}
+        temperatures = collectto!(Tpointbuffer, view(Tpoints, 1:npoints), boundary.Tpoints, constraint.Tpoints)::StridedVector{T}
         Tsteps = linterp1s(ages, temperatures, agesteps)::DenseVector{T}
         calcHeAges = similar(HeAge)::DenseVector{T}
 
@@ -409,8 +409,8 @@
             npointsₚ = npoints
             copyto!(agepointsₚ, agepoints)
             copyto!(Tpointsₚ, Tpoints)
-            copyto!(unconf.agepointsₚ, unconf.agepoints)
-            copyto!(unconf.Tpointsₚ, unconf.Tpoints)
+            copyto!(constraint.agepointsₚ, constraint.agepoints)
+            copyto!(constraint.Tpointsₚ, constraint.Tpoints)
             copyto!(boundary.Tpointsₚ, boundary.Tpoints)
 
             # Adjust the proposal
@@ -438,9 +438,9 @@
                 @. boundary.Tpointsₚ = boundary.T₀ + rand()*boundary.ΔT
 
                 # If there's an imposed unconformity, adjust within parameters
-                if unconf.npoints > 0
-                    @. unconf.agepointsₚ = unconf.Age₀ + rand()*unconf.ΔAge
-                    @. unconf.Tpointsₚ = unconf.T₀ + rand()*unconf.ΔT
+                if constraint.npoints > 0
+                    @. constraint.agepointsₚ = constraint.Age₀ + rand()*constraint.ΔAge
+                    @. constraint.Tpointsₚ = constraint.T₀ + rand()*constraint.ΔT
                 end
 
             elseif (r < p_move+p_birth+p_death+p_bounds+p_kinetics)
@@ -451,8 +451,8 @@
             end
 
             # Recalculate interpolated proposed t-T path
-            ages = collectto!(agepointbuffer, view(agepointsₚ, 1:npointsₚ), boundary.agepoints, unconf.agepointsₚ)::StridedVector{T}
-            temperatures = collectto!(Tpointbuffer, view(Tpointsₚ, 1:npointsₚ), boundary.Tpointsₚ, unconf.Tpointsₚ)::StridedVector{T}
+            ages = collectto!(agepointbuffer, view(agepointsₚ, 1:npointsₚ), boundary.agepoints, constraint.agepointsₚ)::StridedVector{T}
+            temperatures = collectto!(Tpointbuffer, view(Tpointsₚ, 1:npointsₚ), boundary.Tpointsₚ, constraint.Tpointsₚ)::StridedVector{T}
             linterp1s!(Tsteps, knot_index, ages, temperatures, agesteps)
 
             # Old version of imposing max reheating rate, for reference:
@@ -512,8 +512,8 @@
                 npoints = npointsₚ
                 copyto!(agepoints, 1, agepointsₚ, 1, npoints)
                 copyto!(Tpoints, 1, Tpointsₚ, 1, npoints)
-                copyto!(unconf.agepoints, unconf.agepointsₚ)
-                copyto!(unconf.Tpoints, unconf.Tpointsₚ)
+                copyto!(constraint.agepoints, constraint.agepointsₚ)
+                copyto!(constraint.Tpoints, constraint.Tpointsₚ)
                 copyto!(boundary.Tpoints, boundary.Tpointsₚ)
                 copyto!(calcHeAges, calcHeAgesₚ)
 
@@ -529,8 +529,8 @@
             HeAgedist[:,n] .= calcHeAges # distribution of He ages
 
             # This is the actual output we want -- the distribution of t-T paths (t path is always identical)
-            collectto!(view(tpointdist, :, n), view(agepoints, 1:npoints), boundary.agepoints, unconf.agepoints)
-            collectto!(view(Tpointdist, :, n), view(Tpoints, 1:npoints), boundary.Tpoints, unconf.Tpoints)
+            collectto!(view(tpointdist, :, n), view(agepoints, 1:npoints), boundary.agepoints, constraint.agepoints)
+            collectto!(view(Tpointdist, :, n), view(Tpoints, 1:npoints), boundary.Tpoints, constraint.Tpoints)
 
             # Update progress meter every `progress_interval` steps
             (mod(n, progress_interval) == 0) && update!(progress, n)

diff --git a/src/types.jl b/src/types.jl
@@ -1,5 +1,7 @@
-# Define Damage model types
+## --- Define DamageModel types
 abstract type DamageModel end
+export DamageModel
+
 Base.@kwdef struct ZRDAAM{T<:AbstractFloat} <: DamageModel 
     DzD0::T = 193188.0          # Diffusivity [cm^2/sec], crystalline endmember
     DzD0_logsigma::T=log(2)     # log units (default = log(2) = a factor of 2)
@@ -47,6 +49,7 @@ Base.@kwdef struct RDAAM{T<:AbstractFloat} <: DamageModel
 end
 export RDAAM
 
+## --- Define Boundary type to specify the working area
 struct Boundary{T<:AbstractFloat}
     agepoints::Vector{T} # Ma
     Tpoints::Vector{T}   # Degrees C
@@ -67,7 +70,8 @@ function Boundary(T::Type=Float64; agepoints, Tpoints, T₀, ΔT)
 end
 export Boundary
 
-struct Unconformity{T<:AbstractFloat}
+## -- Define Constraint type to specify aribitrary t-T constraint boxes
+struct Constraint{T<:AbstractFloat}
     agepoints::Vector{T} # Ma
     Tpoints::Vector{T}   # Degrees C
     agepointsₚ::Vector{T} # Ma
@@ -78,11 +82,11 @@ struct Unconformity{T<:AbstractFloat}
     ΔT::Vector{T}
     npoints::Int
 end
-function Unconformity(T::Type=Float64; agepoints=Float64[], Tpoints=Float64[], Age₀=Float64[], ΔAge=Float64[], T₀=Float64[], ΔT=Float64[])
+function Constraint(T::Type=Float64; agepoints=Float64[], Tpoints=Float64[], Age₀=Float64[], ΔAge=Float64[], T₀=Float64[], ΔT=Float64[])
     agepoints = isempty(agepoints) ? Age₀ + ΔAge/2 : agepoints
     Tpoints = isempty(Tpoints) ? T₀ + ΔT/2 : Tpoints
     @assert length(agepoints) == length(Tpoints) == length(Age₀) == length(ΔAge) == length(T₀) == length(ΔT)
-    Unconformity{T}(T.(agepoints),
+    Constraint{T}(T.(agepoints),
         T.(Tpoints),
         T.(agepoints), 
         T.(Tpoints), 
@@ -93,8 +97,13 @@ function Unconformity(T::Type=Float64; agepoints=Float64[], Tpoints=Float64[], A
         length(agepoints),
     )
 end
+export Constraint
+
+# For backwards compatibility with old scripts
+const Unconformity = Constraint
 export Unconformity
 
+## --- Define DetailInterval type to specify a minumum number of t-T path nodes within a given time interval
 struct DetailInterval{T<:AbstractFloat}
     agemin::T
     agemax::T

diff --git a/test/testcomplete.jl b/test/testcomplete.jl
@@ -61,7 +61,7 @@ boundary = Boundary(
 )
 
 # Default: No unconformity is imposed
-unconf = Unconformity()
+unconf = Constraint()
 
 
 ## --- Invert for maximum likelihood t-T path

diff --git a/test/testcompletezrn.jl b/test/testcompletezrn.jl
@@ -59,8 +59,8 @@ boundary = Boundary(
     ΔT = Float64[model.ΔTnow, model.ΔTinit],
 )
 
-# Default: No unconformity is imposed
-unconf = Unconformity()
+# Default: No constraint is imposed
+constraint = Constraint()
 
 
 ## --- Invert for maximum likelihood t-T path
@@ -75,8 +75,8 @@ agepoints[1:npoints] .= (model.tinit/30,model.tinit/4,model.tinit/2,model.tinit-
 Tpoints[1:npoints] .= Tr  # Degrees C
 
 # Run Markov Chain
-@time "Compiling MCMC" MCMC(data, model, npoints, agepoints, Tpoints, boundary, unconf)
-@time "Running MCMC" (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, unconf)
+@time "Compiling MCMC" MCMC(data, model, npoints, agepoints, Tpoints, boundary, constraint)
+@time "Running MCMC" (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, constraint)
 
 @test isa(Tpointdist, AbstractMatrix)
 @test maximum(Tpointdist) <= model.Tinit
@@ -114,7 +114,7 @@ detail = DetailInterval(
     agemax = 541, # Oldest end of detail interval
     minpoints = 5, # Minimum number of points in detail interval
 )
-@time "MCMC with Detail interval" (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, unconf, detail)
+@time "MCMC with Detail interval" (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, constraint, detail)
 
 @test isa(Tpointdist, AbstractMatrix)
 @test maximum(Tpointdist) <= model.Tinit
@@ -150,7 +150,7 @@ llmean = mean(@view(lldist[model.burnin:end]))
 model = (model...,
     dynamicjumping=true
 )
-@time "MCMC with Detail interval & dynamicjumping" (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, unconf, detail)
+@time "MCMC with Detail interval & dynamicjumping" (tpointdist, Tpointdist, ndist, HeAgedist, lldist, acceptancedist, σⱼtdist, σⱼTdist) = MCMC(data, model, npoints, agepoints, Tpoints, boundary, constraint, detail)
 
 @test isa(Tpointdist, AbstractMatrix)
 @test maximum(Tpointdist) <= model.Tinit