Merge pull request #24 from boriskaus/bk/zircon_growth

boriskaus · web-flow · commit de386ae43b2e · 2026-04-12T14:28:03.000+02:00
ZirconGrowth interaction
diff --git a/Project.toml b/Project.toml
@@ -45,6 +45,12 @@ TimerOutputs = "0.5"
 WriteVTK = "1"
 julia = "1.10"
 
+[weakdeps]
+ZirconGrowth = "5085b5e1-578a-4bbf-9dfc-2862bf7461c3"
+
+[extensions]
+ZirconGrowthExt = "ZirconGrowth"
+
 [extras]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
diff --git a/docs/make.jl b/docs/make.jl
@@ -112,6 +112,7 @@ makedocs(;
                 "MTK_GMG 3D Examples" => "man/example_mtk_gmg_3d.md"
             ],
             "Numerics and Physics" => "man/numerics.md",
+            "ZirconGrowth Integration" => "man/zircon_growth.md",
             "Ongoing Development" => "man/development.md",
         ],
         "API" => Any[
diff --git a/docs/src/man/zircon_growth.md b/docs/src/man/zircon_growth.md
@@ -0,0 +1,122 @@
+# ZirconGrowth Integration
+
+MagmaThermoKinematics has an optional integration with
+[ZirconGrowth.jl](https://github.com/JuliaGeodynamics/ZirconGrowth.jl), a package that
+simulates zircon crystal growth along a magma Tt-path using an implicit finite-difference
+scheme on a 1-D spherical grid. The integration is implemented as a
+[Julia package extension](https://pkgdocs.julialang.org/v1/creating-packages/#Conditional-loading-of-code-in-packages-(Extensions))
+and is loaded automatically when you do `using ZirconGrowth` in your session.
+
+## Installation
+
+ZirconGrowth.jl is a separate package:
+
+```julia
+] add https://github.com/JuliaGeodynamics/ZirconGrowth.jl   # while unregistered
+] add ZirconGrowth                                           # once registered
+```
+
+## Workflow
+
+A typical MagmaThermoKinematics run records the temperature–time history of each passive
+tracer particle.  After the simulation, these Tt-paths can be passed directly to the
+ZirconGrowth model:
+
+```julia
+using MagmaThermoKinematics
+using ZirconGrowth          # activates the extension
+
+# Option 1 — point at an output directory
+age_years, zircon_radius_um = simulate_zircon_growth_from_tracers("MyRun/")
+
+# Option 2 — pass tracers already loaded in memory
+using JLD2
+Tracers = JLD2.load("MyRun/Tracers_SimParams.jld2", "Tracers")
+age_years, zircon_radius_um = simulate_zircon_growth_from_tracers(Tracers)
+
+# Option 3 — single tracer
+result = simulate_zircon_growth_from_tracers(Tracers[1])
+```
+
+## Return values
+
+`simulate_zircon_growth_from_tracers` returns a `NamedTuple` with one entry per
+successfully simulated tracer (tracers with fewer than 2 time steps are skipped):
+
+| Field               | Type              | Description |
+|:--------------------|:------------------|:------------|
+| `age_years`         | `Vector{Float64}` | Volume-averaged crystallisation age in years before the end of the simulation |
+| `zircon_radius_um`  | `Vector{Float64}` | Final crystal radius in µm |
+
+### Volume-averaged age
+
+The volume-averaged age weights each concentric shell of the crystal by its volume.
+The shell between radii $r_i$ and $r_{i+1}$ crystallised at simulation time $t_i$ and has
+volume proportional to $r_{i+1}^3 - r_i^3$.  Its age (years before the end of the
+simulation) is $t_\text{end} - \tfrac{1}{2}(t_i + t_{i+1})$.  The weighted mean is:
+
+$$\bar{t} = \frac{\sum_i (t_\text{end} - \bar{t}_i)\,(r_{i+1}^3 - r_i^3)}{\sum_i (r_{i+1}^3 - r_i^3)}$$
+
+This can also be called directly on a `SimulationResult`:
+
+```julia
+age = volume_averaged_age(results[1])      # single crystal → Float64
+ages = volume_averaged_age(results)        # vector → Vector{Float64}
+```
+
+## Keyword arguments
+
+```julia
+simulate_zircon_growth_from_tracers(Tracers;
+    params         = nothing,          # ZirconGrowth.GrowthParams (one per tracer if nothing)
+    elements       = ZirconGrowth.default_element_data(),  # trace elements to track
+    filename       = nothing,          # save age_years & zircon_radius_um to this JLD2 file
+    return_results = false)            # also return Vector{SimulationResult}
+```
+
+### `return_results`
+
+Set `return_results = true` to retrieve the full `ZirconGrowth.SimulationResult` objects.
+This uses more memory but gives access to the complete crystal profile, trace-element
+concentrations, growth rates, etc.:
+
+```julia
+age_years, zircon_radius_um, results = simulate_zircon_growth_from_tracers(
+    Tracers; return_results = true)
+
+# example: inspect the radius–time history of the first crystal
+results[1].time_years
+results[1].zircon_radius_um
+```
+
+### `filename`
+
+When a `filename` is provided the summary vectors are saved to a JLD2 file:
+
+```julia
+# explicit filename
+simulate_zircon_growth_from_tracers(Tracers; filename = "MyRun/ZirconGrowth.jld2")
+
+# directory form saves to dirname/ZirconGrowth.jld2 automatically
+simulate_zircon_growth_from_tracers("MyRun/")
+
+# suppress saving when using the directory form
+simulate_zircon_growth_from_tracers("MyRun/"; filename = nothing)
+```
+
+## Performance
+
+Each tracer is independent, so the loop is parallelised with `Threads.@threads`.  Start
+Julia with multiple threads for a proportional speedup:
+
+```bash
+julia --threads auto
+```
+
+or set the environment variable permanently:
+
+```bash
+export JULIA_NUM_THREADS=auto
+```
+
+Check the thread count at runtime with `Threads.nthreads()`.
diff --git a/ext/ZirconGrowthExt.jl b/ext/ZirconGrowthExt.jl
@@ -0,0 +1,191 @@
+module ZirconGrowthExt
+
+using MagmaThermoKinematics
+using MagmaThermoKinematics: JLD2
+using ZirconGrowth
+
+"""
+    simulate_zircon_growth_from_tracers(Tracers;
+                                        params         = nothing,
+                                        elements       = ZirconGrowth.default_element_data(),
+                                        filename       = nothing,
+                                        return_results = false)
+    simulate_zircon_growth_from_tracers(tr::Tracer; kwargs...)
+    simulate_zircon_growth_from_tracers(dirname::AbstractString; kwargs...)
+
+Run the [ZirconGrowth.jl](https://github.com/JuliaGeodynamics/ZirconGrowth.jl) crystal-growth
+model on tracer Tt-paths recorded during a MagmaThermoKinematics simulation.
+
+Three call forms are supported:
+- Pass a `StructArray` of `Tracer` objects directly.
+- Pass a single `Tracer`.
+- Pass a directory path; `dirname/Tracers_SimParams.jld2` is loaded automatically
+  (saves to `dirname/ZirconGrowth.jld2` by default).
+
+Tracers with fewer than 2 time steps are skipped.
+`Tracer.time_vec` must be in **Myr** and `Tracer.T_vec` in **°C**.
+The loop runs on all available Julia threads (`julia --threads auto`).
+
+## Return value
+
+By default returns a `NamedTuple` with two `Vector{Float64}` fields, one entry per
+successfully simulated tracer:
+
+| Field               | Description |
+|:--------------------|:------------|
+| `age_years`         | Volume-averaged crystallisation age in years before the end of the simulation (see [`volume_averaged_age`](@ref)) |
+| `zircon_radius_um`  | Final crystal radius in µm |
+
+When `return_results = true` a third field `results` is included, containing a
+`Vector{SimulationResult}` with the full ZirconGrowth output for each tracer.
+
+## Arguments
+- `params`         : `ZirconGrowth.GrowthParams`; if `nothing` (default) a fresh instance
+                     is constructed from each tracer's Tt-path individually.
+- `nx`             : number of 1-D spatial grid points for the ZirconGrowth solver
+                     (default 100).  Ignored when `params` is provided explicitly.
+- `elements`       : `ZirconGrowth.ElementData` selecting which trace elements are tracked;
+                     defaults to `ZirconGrowth.default_element_data()`.
+- `filename`       : path to a JLD2 file; if provided the `age_years` and
+                     `zircon_radius_um` vectors (and `results` when requested) are saved
+                     there.  Default `nothing` (no saving) for the `Tracers` form;
+                     `dirname/ZirconGrowth.jld2` for the `dirname` form.
+- `return_results` : set to `true` to include the full `Vector{SimulationResult}` in the
+                     return value.  Default `false` to save memory.
+
+## Examples
+
+```julia
+using MagmaThermoKinematics
+using ZirconGrowth          # triggers the extension
+
+# from a directory produced by a simulation run
+age_years, zircon_radius_um = simulate_zircon_growth_from_tracers("MyRun/")
+
+# from tracers already in memory
+using JLD2
+Tracers = JLD2.load("MyRun/Tracers_SimParams.jld2", "Tracers")
+age_years, zircon_radius_um = simulate_zircon_growth_from_tracers(Tracers)
+
+# single tracer
+res = simulate_zircon_growth_from_tracers(Tracers[1])
+
+# also retrieve full SimulationResult objects
+age_years, zircon_radius_um, results = simulate_zircon_growth_from_tracers(
+    Tracers; return_results = true)
+```
+"""
+function MagmaThermoKinematics.simulate_zircon_growth_from_tracers(Tracers;
+        params::Union{Nothing, ZirconGrowth.GrowthParams} = nothing,
+        nx::Int = 100,
+        elements::ZirconGrowth.ElementData = ZirconGrowth.default_element_data(),
+        filename::Union{Nothing, AbstractString} = nothing,
+        return_results::Bool = false)
+
+    n                = length(Tracers)
+    age_years        = Vector{Union{Nothing, Float64}}(nothing, n)
+    zircon_radius_um = Vector{Union{Nothing, Float64}}(nothing, n)
+    _results         = return_results ? Vector{Union{Nothing, ZirconGrowth.SimulationResult}}(nothing, n) : nothing
+    done             = Threads.Atomic{Int}(0)
+    interval         = max(1, n ÷ 20)
+
+    println("Simulating zircon growth for $n tracers on $(Threads.nthreads()) thread(s)...")
+    Threads.@threads for i in eachindex(Tracers)
+        tr = Tracers[i]
+        isempty(tr.time_vec) && continue
+
+        time_Myr = Float64.(tr.time_vec)   # already in Myr
+        T_C      = Float64.(tr.T_vec)
+
+        length(time_Myr) < 2 && continue
+
+        p = isnothing(params) ? ZirconGrowth.GrowthParams(time_Myr, T_C; nx=nx) : params
+
+        res = ZirconGrowth.simulate_from_cooling_path(time_Myr, T_C;
+                  params = p, elements = elements)
+
+        age_years[i]        = MagmaThermoKinematics.volume_averaged_age(res)
+        zircon_radius_um[i] = res.zircon_radius_um[end]
+        return_results && (_results[i] = res)
+
+        k = Threads.atomic_add!(done, 1) + 1
+        k % interval == 0 && print("\r  $k / $n done...")
+    end
+
+    age_years        = Float64[v for v in age_years        if !isnothing(v)]
+    zircon_radius_um = Float64[v for v in zircon_radius_um if !isnothing(v)]
+    println("\rDone: $(length(age_years)) / $n tracers simulated.          ")
+
+    if !isnothing(filename)
+        JLD2.jldsave(filename; age_years, zircon_radius_um)
+        println("Saved ZirconGrowth results to: $filename  ($(length(age_years)) tracers)")
+    end
+
+    if return_results
+        results = ZirconGrowth.SimulationResult[r for r in _results if !isnothing(r)]
+        return (; age_years, zircon_radius_um, results)
+    end
+
+    return (; age_years, zircon_radius_um)
+end
+
+function MagmaThermoKinematics.simulate_zircon_growth_from_tracers(tr::MagmaThermoKinematics.Tracer; kwargs...)
+    return MagmaThermoKinematics.simulate_zircon_growth_from_tracers((tr,); kwargs...)
+end
+
+"""
+    simulate_zircon_growth_from_tracers(dirname::AbstractString; kwargs...)
+
+Load tracers from `dirname/Tracers_SimParams.jld2` and call
+`simulate_zircon_growth_from_tracers(Tracers; kwargs...)`.
+"""
+function MagmaThermoKinematics.simulate_zircon_growth_from_tracers(dirname::AbstractString;
+        filename::Union{Nothing, AbstractString} = joinpath(dirname, "ZirconGrowth.jld2"),
+        kwargs...)
+
+    Tracers = JLD2.load(joinpath(dirname, "Tracers_SimParams.jld2"), "Tracers")
+    return MagmaThermoKinematics.simulate_zircon_growth_from_tracers(Tracers;
+               filename = filename, kwargs...)
+end
+
+"""
+    volume_averaged_age(result::SimulationResult) -> Float64
+
+Compute the volume-averaged crystallisation age (in years) of a single zircon crystal.
+
+Each concentric shell crystallised at a different time. The shell between radii
+`r[i]` and `r[i+1]` has volume ∝ r[i+1]³ − r[i]³ and an age of
+`time_years[end] − time_years[i]` (measured back from the end of the simulation).
+Shells with zero or negative growth are excluded.
+
+Returns the volume-weighted mean age in **years**.
+"""
+function MagmaThermoKinematics.volume_averaged_age(result::ZirconGrowth.SimulationResult)
+    t = result.time_years
+    r = result.zircon_radius_um
+
+    age_sum = 0.0
+    vol_sum = 0.0
+    t_end   = t[end]
+
+    for i in 1:(length(r) - 1)
+        dV = r[i+1]^3 - r[i]^3   # proportional to shell volume; 4π/3 cancels
+        dV <= 0 && continue
+        age_mid  = t_end - 0.5*(t[i] + t[i+1])   # age of shell midpoint
+        age_sum += age_mid * dV
+        vol_sum += dV
+    end
+
+    return vol_sum > 0 ? age_sum / vol_sum : 0.0
+end
+
+"""
+    volume_averaged_age(results::AbstractVector) -> Vector{Float64}
+
+Apply `volume_averaged_age` to each element of `results` and return a `Vector{Float64}`.
+"""
+function MagmaThermoKinematics.volume_averaged_age(results::AbstractVector)
+    return [MagmaThermoKinematics.volume_averaged_age(r) for r in results]
+end
+
+end # module ZirconGrowthExt
diff --git a/src/MagmaThermoKinematics.jl b/src/MagmaThermoKinematics.jl
@@ -230,7 +230,7 @@ export AdvectTemperature, Interpolate!, CorrectBounds, evaluate_interp_2D, evalu
 # Shared utility routines and MTK/GMG structs are loaded at module initialization
 # to avoid defining new global bindings inside environment!.
 include("Utils.jl")
-export Process_ZirconAges, copy_arrays_GPU2CPU!, copy_arrays_CPU2GPU!
+export Process_ZirconAges, simulate_zircon_growth_from_tracers, volume_averaged_age, copy_arrays_GPU2CPU!, copy_arrays_CPU2GPU!
 
 include("MTK_GMG_structs.jl")
 export NumParam, DikeParam, TimeDepProps
diff --git a/src/Utils.jl b/src/Utils.jl
@@ -57,6 +57,21 @@ function Process_ZirconAges(dirname; ZirconData=ZirconAgeData())
     return nothing
 end
 
+"""
+    simulate_zircon_growth_from_tracers(dirname; params, elements)
+
+Requires `ZirconGrowth.jl` to be loaded.  Load the extension with `using ZirconGrowth`.
+"""
+function simulate_zircon_growth_from_tracers end
+
+"""
+    volume_averaged_age(result)
+    volume_averaged_age(results)
+
+Requires `ZirconGrowth.jl` to be loaded.  Load the extension with `using ZirconGrowth`.
+"""
+function volume_averaged_age end
+
 function copy_arrays_GPU2CPU!(T_CPU::AbstractArray,  ϕ_CPU::AbstractArray, T_GPU::AbstractArray, ϕ_GPU::AbstractArray)
 
     T_CPU  .= Array(T_GPU)
