Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found

Target

Select target project
  • cryogrid/cryogridjulia
  • jnitzbon/cryogridjulia
  • bgroenke/cryogridjulia
3 results
Show changes
Commits on Source (56)
Showing
with 409 additions and 346 deletions
name = "CryoGrid" name = "CryoGrid"
uuid = "a535b82e-5f3d-4d97-8b0b-d6483f5bebd5" uuid = "a535b82e-5f3d-4d97-8b0b-d6483f5bebd5"
authors = ["Brian Groenke <brian.groenke@awi.de>", "Moritz Langer <moritz.langer@awi.de>"] authors = ["Brian Groenke <brian.groenke@awi.de>", "Moritz Langer <moritz.langer@awi.de>"]
version = "0.9.0" version = "0.10.3"
[deps] [deps]
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
...@@ -29,7 +29,6 @@ PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" ...@@ -29,7 +29,6 @@ PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
Reexport = "189a3867-3050-52da-a836-e630ba90ab69" Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
SimulationLogs = "e3c6736c-93d9-479d-93f3-96ff5e8836d5" SimulationLogs = "e3c6736c-93d9-479d-93f3-96ff5e8836d5"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
...@@ -40,6 +39,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" ...@@ -40,6 +39,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
[compat] [compat]
julia = ">= 1.6" julia = ">= 1.6"
PreallocationTools = ">= 0.3.1"
[extras] [extras]
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
......
...@@ -10,7 +10,7 @@ const modules = [ ...@@ -10,7 +10,7 @@ const modules = [
CryoGrid.Physics, CryoGrid.Physics,
CryoGrid.Boundaries, CryoGrid.Boundaries,
CryoGrid.HeatConduction, CryoGrid.HeatConduction,
CryoGrid.WaterBalance, CryoGrid.Hydrology,
CryoGrid.Soils, CryoGrid.Soils,
CryoGrid.SEB, CryoGrid.SEB,
CryoGrid.Strat, CryoGrid.Strat,
...@@ -35,7 +35,7 @@ makedocs(modules=modules, ...@@ -35,7 +35,7 @@ makedocs(modules=modules,
"Utilities" => "api/utils.md", "Utilities" => "api/utils.md",
"Physics" => [ "Physics" => [
"Heat Conduction" => "api/heat_conduction.md", "Heat Conduction" => "api/heat_conduction.md",
"Water Balance" => "api/water_balance.md", "Hydrology" => "api/hydrology.md",
"Soils" => "api/soils.md", "Soils" => "api/soils.md",
"Surface Energy Balance" => "api/seb.md", "Surface Energy Balance" => "api/seb.md",
], ],
......
# Water balance # Hydrology
```@autodocs ```@autodocs
Modules = [WaterBalance] Modules = [Hydrology]
Private = false Private = false
Order = [:type, :function, :macro] Order = [:type, :function, :macro]
``` ```
\ No newline at end of file
using CryoGrid
using Plots
# forcings = loadforcings(CryoGrid.Presets.Forcings.Samoylov_ERA5_fitted_daily_1979_2020, :Tair => u"°C", :swe => u"m", :ρsn => u"kg/m^3"; spec=JsonSpec{2});
forcings = loadforcings(CryoGrid.Presets.Forcings.Samoylov_ERA_obs_fitted_1979_2014_spinup_extended_2044, :Tair => u"°C", :snowfall => u"mm/d"; spec=JsonSpec{1});
# use air temperature as upper boundary forcing;
tair = TimeSeriesForcing(ustrip.(forcings.data.Tair), forcings.timestamps, :Tair);
# swe = TimeSeriesForcing(ustrip.(forcings.data.swe), forcings.timestamps, :swe);
# ρsn = TimeSeriesForcing(ustrip.(forcings.data.ρsn), forcings.timestamps, :ρsn);
snowfall = TimeSeriesForcing(convert.(Float64, ustrip.(u"m/s", forcings.data.snowfall)), forcings.timestamps, :snowfall)
# use default profiles for samoylov
soilprofile, tempprofile = CryoGrid.Presets.SamoylovDefault
# "simple" heat conduction model w/ 5 cm grid spacing (defaults to free water freezing scheme)
modelgrid = Grid(vcat([-1.0u"m"], CryoGrid.Presets.DefaultGrid_5cm))
initT = initializer(:T, tempprofile)
z_top = -2.0u"m"
z_sub = map(knot -> knot.depth, soilprofile)
z_bot = modelgrid[end]
snowmass = SnowMassBalance(
para = Snow.Dynamic(
ablation = Snow.DegreeDayMelt(factor=5.0u"mm/K/d")
)
)
strat = @Stratigraphy(
z_top => top(TemperatureGradient(tair), Snowfall(snowfall)),
# prescribed snow
# z_top => subsurface(:snowpack, Snowpack(para=Snow.Bulk()), SnowMassBalance(para=Snow.Prescribed(swe=swe, ρsn=ρsn)), Heat(:H)),
# "dynamic" snow (i.e. modeled snow accumulation and ablation)
z_top => subsurface(:snowpack, Snowpack(para=Snow.Bulk(thresh=2.0u"cm")), snowmass, Heat(:H)),
z_sub[1] => subsurface(:topsoil1, Soil(para=soilprofile[1].value), Heat(:H)),
z_sub[2] => subsurface(:topsoil2, Soil(para=soilprofile[2].value), Heat(:H)),
z_sub[3] => subsurface(:sediment1, Soil(para=soilprofile[3].value), Heat(:H)),
z_sub[4] => subsurface(:sediment2, Soil(para=soilprofile[4].value), Heat(:H)),
z_sub[5] => subsurface(:sediment3, Soil(para=soilprofile[5].value), Heat(:H)),
z_bot => bottom(GeothermalHeatFlux(0.053u"J/s/m^2"))
);
tile = Tile(strat, modelgrid, initT)
# define time span
tspan = (DateTime(2010,9,30),DateTime(2012,9,30))
p = parameters(tile)
u0, du0 = initialcondition!(tile, tspan, p)
prob = CryoGridProblem(tile,u0,tspan,p,step_limiter=nothing,savevars=(:T,:snowpack => (:dsn,:T_ub)))
sol = @time solve(prob, SSPRK22(), dt=300.0, saveat=24*3600.0, progress=true);
out = CryoGridOutput(sol)
# Plot it!
zs = [1,10,20,30,50,100,200,500,1000]u"cm"
cg = Plots.cgrad(:copper,rev=true);
plot(ustrip(out.T[Z(Near(zs))]), color=cg[LinRange(0.0,1.0,length(zs))]', ylabel="Temperature (°C)", leg=false, dpi=150)
plt1 = plot!(ustrip(out.snowpack.T_ub), color=:skyblue, linestyle=:dash, alpha=0.7, leg=false, dpi=150)
plot(ustrip(out.swe), ylabel="Depth (m)", label="Snow water equivalent", dpi=150)
plt2 = plot!(ustrip(out.snowpack.dsn), label="Snow depth", legendtitle="", dpi=150)
plot(plt1, plt2, size=(1200,400), margins=5*Plots.Measures.mm)
...@@ -7,11 +7,11 @@ const gridvals = vcat([0:0.02:2...,2.05:0.05:4.0..., ...@@ -7,11 +7,11 @@ const gridvals = vcat([0:0.02:2...,2.05:0.05:4.0...,
35:5:50...,60:10:100...,200:100:1000...]...)u"m" 35:5:50...,60:10:100...,200:100:1000...]...)u"m"
# soil profile: depth => (excess ice, natural porosity, saturation, organic fraction) # soil profile: depth => (excess ice, natural porosity, saturation, organic fraction)
soilprofile = SoilProfile( soilprofile = SoilProfile(
0.0u"m" => CharacteristicFractions(xic=0.0,por=0.80,sat=1.0,org=0.75), #(θw=0.80,θm=0.05,θo=0.15,ϕ=0.80), 0.0u"m" => soilparameters(xic=0.0,por=0.80,sat=1.0,org=0.75), #(θwi=0.80,θm=0.05,θo=0.15,ϕ=0.80),
0.1u"m" => CharacteristicFractions(xic=0.0,por=0.80,sat=1.0,org=0.25), #(θw=0.80,θm=0.15,θo=0.05,ϕ=0.80), 0.1u"m" => soilparameters(xic=0.0,por=0.80,sat=1.0,org=0.25), #(θwi=0.80,θm=0.15,θo=0.05,ϕ=0.80),
0.4u"m" => CharacteristicFractions(xic=0.30,por=0.55,sat=1.0,org=0.25), #(θw=0.80,θm=0.15,θo=0.05,ϕ=0.55), 0.4u"m" => soilparameters(xic=0.30,por=0.55,sat=1.0,org=0.25), #(θwi=0.80,θm=0.15,θo=0.05,ϕ=0.55),
3.0u"m" => CharacteristicFractions(xic=0.0,por=0.50,sat=1.0,org=0.0), #(θw=0.50,θm=0.50,θo=0.0,ϕ=0.50), 3.0u"m" => soilparameters(xic=0.0,por=0.50,sat=1.0,org=0.0), #(θwi=0.50,θm=0.50,θo=0.0,ϕ=0.50),
10.0u"m" => CharacteristicFractions(xic=0.0,por=0.30,sat=1.0,org=0.0), #(θw=0.30,θm=0.70,θo=0.0,ϕ=0.30), 10.0u"m" => soilparameters(xic=0.0,por=0.30,sat=1.0,org=0.0), #(θwi=0.30,θm=0.70,θo=0.0,ϕ=0.30),
), ),
# mid-winter temperature profile # mid-winter temperature profile
tempprofile = TemperatureProfile( tempprofile = TemperatureProfile(
...@@ -49,7 +49,7 @@ u0, du0 = initialcondition!(model, tspan, p) ...@@ -49,7 +49,7 @@ u0, du0 = initialcondition!(model, tspan, p)
# CryoGrid front-end for ODEProblem # CryoGrid front-end for ODEProblem
prob = CryoGridProblem(model,u0,tspan,p,savevars=(:T,)) prob = CryoGridProblem(model,u0,tspan,p,savevars=(:T,))
# solve with forward Euler (w/ CFL) and construct CryoGridOutput from solution # solve with forward Euler (w/ CFL) and construct CryoGridOutput from solution
out = @time solve(prob, Euler(), dt=2*60.0, callback=CFLStepLimiter(model), saveat=24*3600.0, progress=true) |> CryoGridOutput; out = @time solve(prob, Euler(), dt=2*60.0, saveat=24*3600.0, progress=true) |> CryoGridOutput;
# Plot it! # Plot it!
zs = [1:10...,20:10:100...] zs = [1:10...,20:10:100...]
cg = Plots.cgrad(:copper,rev=true); cg = Plots.cgrad(:copper,rev=true);
......
...@@ -11,33 +11,35 @@ using Reexport ...@@ -11,33 +11,35 @@ using Reexport
# Common types and methods # Common types and methods
export Layer, SubSurface, Top, Bottom export Layer, SubSurface, Top, Bottom
export Process, SubSurfaceProcess, BoundaryProcess, CoupledProcesses, Coupled export Process, SubSurfaceProcess, BoundaryProcess, CoupledProcesses
export Coupled, Coupled2, Coupled3, Coupled4
include("types.jl") include("types.jl")
export variables, initialcondition!, diagnosticstep!, prognosticstep!, interact! export variables, initialcondition!, diagnosticstep!, prognosticstep!, interact!, timestep
export boundaryflux, boundaryvalue, criterion, affect!, observe export boundaryflux, boundaryvalue, criterion, trigger!, observe
include("methods.jl") include("methods.jl")
# Submodules # Submodules
include("Utils/Utils.jl") include("Utils/Utils.jl")
using .Utils
export convert_t, convert_tspan, pstrip, @pstrip, @sym_str
include("Numerics/Numerics.jl") include("Numerics/Numerics.jl")
include("IO/InputOutput.jl")
include("Physics/Physics.jl")
include("Strat/Strat.jl")
include("Diagnostics/Diagnostics.jl")
using .Numerics using .Numerics
export Grid, cells, edges, subgridinds, Δ, volume, area, initializer, getvar export Grid, cells, edges, subgridinds, Δ, volume, area, initializer, getvar
using .Utils include("IO/InputOutput.jl")
export convert_t, convert_tspan, pstrip, @pstrip, @sym_str @reexport using .InputOutput
# Re-exported submodules include("Physics/Physics.jl")
@reexport using .Physics @reexport using .Physics
include("Strat/Strat.jl")
@reexport using .Strat @reexport using .Strat
@reexport using .InputOutput include("Diagnostics/Diagnostics.jl")
@reexport using .Diagnostics @reexport using .Diagnostics
# Coupling
include("coupling.jl")
# Re-exported packages # Re-exported packages
@reexport using Dates: Dates, DateTime @reexport using Dates: Dates, DateTime
@reexport using DimensionalData: DimArray, Z, Dim, dims @reexport using DimensionalData
@reexport using IntervalSets @reexport using IntervalSets
@reexport using Unitful @reexport using Unitful
......
...@@ -3,15 +3,16 @@ Driver module SciML diffeq solvers. ...@@ -3,15 +3,16 @@ Driver module SciML diffeq solvers.
""" """
module DiffEq module DiffEq
using CryoGrid
using CryoGrid: Event, ContinuousEvent, DiscreteEvent, ContinuousTrigger, Increasing, Decreasing
using CryoGrid.Drivers using CryoGrid.Drivers
using CryoGrid: Strat, SubSurface, CoupledProcesses, Callback, CallbackStyle, Discrete, Continuous
using CryoGrid.InputOutput using CryoGrid.InputOutput
using CryoGrid.Numerics using CryoGrid.Numerics
using CryoGrid.Physics: Heat using CryoGrid.Physics: Heat
using CryoGrid.Strat: Stratigraphy, StratComponent
using CryoGrid.Utils using CryoGrid.Utils
import CryoGrid: variables, callbacks, criterion, affect! import CryoGrid.Strat
import CryoGrid.Strat: Tile, Stratigraphy, StratComponent
using ComponentArrays using ComponentArrays
using Dates using Dates
...@@ -29,124 +30,28 @@ using DiffEqBase ...@@ -29,124 +30,28 @@ using DiffEqBase
using DiffEqBase.SciMLBase using DiffEqBase.SciMLBase
using DiffEqCallbacks using DiffEqCallbacks
import DiffEqCallbacks
@reexport using OrdinaryDiffEq @reexport using OrdinaryDiffEq
@reexport using DiffEqBase: solve, init, ODEProblem, SciMLBase @reexport using DiffEqBase: solve, init, ODEProblem, SciMLBase
export CryoGridProblem
export CFLStepLimiter
include("steplimiters.jl")
export TDMASolver export TDMASolver
include("solvers.jl") include("solvers.jl")
include("callbacks.jl")
export CryoGridProblem
include("problem.jl")
include("output.jl")
# Add method dispatches for other CryoGrid methods on DEIntegrator type
""" """
Specialized problem type for CryoGrid `ODEProblem`s. Tile(integrator::SciMLBase.DEIntegrator)
"""
struct CryoGridODEProblem end
Constructs a `Tile` from a `SciMLBase` integrator.
"""
function Strat.Tile(integrator::SciMLBase.DEIntegrator) function Strat.Tile(integrator::SciMLBase.DEIntegrator)
tile = integrator.sol.prob.f.f tile = integrator.sol.prob.f.f
return Strat.updateparams(tile, Strat.withaxes(integrator.u, tile), integrator.p, integrator.t) return Strat.updateparams(tile, Strat.withaxes(integrator.u, tile), integrator.p, integrator.t)
end end
"""
CryoGridProblem(setup::Tile, tspan::NTuple{2,Float64}, p=nothing;kwargs...)
CryoGrid specialized constructor for ODEProblem that automatically generates the initial
condition and necessary callbacks.
"""
function CryoGridProblem(
tile::Tile,
u0::ComponentVector,
tspan::NTuple{2,Float64},
p=nothing;
saveat=3600.0,
savevars=(),
save_everystep=false,
save_start=true,
save_end=true,
callback=nothing,
kwargs...
)
# workaround for bug in DiffEqCallbacks; see https://github.com/SciML/DifferentialEquations.jl/issues/326
# we have to manually expand single-number `saveat` (i.e. time interval for saving) to a step-range.
expandtstep(tstep::Number) = tspan[1]:tstep:tspan[end]
expandtstep(tstep::AbstractVector) = tstep
getsavestate(tile::Tile, u, du) = deepcopy(Strat.getvars(tile.state, Strat.withaxes(u, tile), Strat.withaxes(du, tile), savevars...))
savefunc(u, t, integrator) = getsavestate(Tile(integrator), Strat.withaxes(u, Tile(integrator)), get_du(integrator))
# remove units
tile = stripunits(tile)
# collect parameters
p = isnothing(p) ? dustrip.(collect(Model(tile)[:val])) : collect(p)
du0 = zero(u0)
# set up saving callback
stateproto = getsavestate(tile, u0, du0)
savevals = SavedValues(Float64, typeof(stateproto))
savingcallback = SavingCallback(savefunc, savevals; saveat=expandtstep(saveat), save_start=save_start, save_end=save_end, save_everystep=save_everystep)
layercallbacks = _makecallbacks(tile)
usercallbacks = isnothing(callback) ? () : callback
callbacks = CallbackSet(savingcallback, layercallbacks..., usercallbacks...)
# note that this implicitly discards any existing saved values in the model setup's state history
tile.hist.vals = savevals
# set up default mass matrix, M:
# M⋅∂u∂t = f(u)
M_diag = similar(tile.state.uproto)
M_idxmap = ComponentArrays.indexmap(getaxes(M_diag)[1])
allvars = Flatten.flatten(tile.state.vars, Flatten.flattenable, Var)
progvars = map(varname, filter(isprognostic, allvars))
algvars = map(varname, filter(isalgebraic, allvars))
for name in keys(M_idxmap)
M_diag_var = @view M_diag[name]
if name progvars
M_diag_var .= one(eltype(M_diag))
elseif name algvars
M_diag_var .= zero(eltype(M_diag))
end
end
# if no algebraic variables are present, use identity matrix
num_algebraic = length(M_diag) - sum(M_diag)
M = num_algebraic > 0 ? Diagonal(M_diag) : I
func = odefunction(tile, M, u0, p, tspan; kwargs...)
ODEProblem(func, u0, tspan, p, CryoGridODEProblem(); callback=callbacks, kwargs...)
end
"""
CryoGridProblem(setup::Tile, tspan::NTuple{2,DateTime}, args...;kwargs...)
"""
CryoGridProblem(setup::Tile, u0::ComponentVector, tspan::NTuple{2,DateTime}, args...;kwargs...) = CryoGridProblem(setup,u0,convert_tspan(tspan),args...;kwargs...)
"""
odefunction(setup::Tile, M, u0, p, tspan; kwargs...)
Constructs a SciML `ODEFunction` given the model setup, mass matrix M, initial state u0, parameters p, and tspan.
Can (and should) be overridden by users to provide customized ODEFunction configurations for specific problem setups, e.g:
```
tile = Tile(strat,grid)
function CryoGrid.Setup.odefunction(::DefaultJac, setup::typeof(tile), M, u0, p, tspan)
...
# make sure to return an instance of ODEFunction
end
...
prob = CryoGridProblem(tile, tspan, p)
```
`JacobianStyle` can also be extended to create custom traits which can then be applied to compatible `Tile`s.
"""
odefunction(setup::TSetup, M, u0, p, tspan; kwargs...) where {TSetup<:Tile} = odefunction(JacobianStyle(TSetup), setup, M, u0, p, tspan; kwargs...)
odefunction(::DefaultJac, setup::TSetup, M, u0, p, tspan; kwargs...) where {TSetup<:Tile} = ODEFunction(setup, mass_matrix=M; kwargs...)
function odefunction(::TridiagJac, setup::Tile, M, u0, p, tspan; kwargs...)
if :jac_prototype in keys(kwargs)
@warn "using user specified jac_prorotype instead of tridiagonal"
ODEFunction(setup, mass_matrix=M; kwargs...)
else
N = length(u0)
J = Tridiagonal(
similar(u0, eltype(p), N-1) |> Vector,
similar(u0, eltype(p), N) |> Vector,
similar(u0, eltype(p), N-1) |> Vector
)
ODEFunction(setup, mass_matrix=M, jac_prototype=J, kwargs...)
end
end
""" """
getstate(integrator::SciMLBase.DEIntegrator) getstate(integrator::SciMLBase.DEIntegrator)
getstate(layername::Symbol, integrator::SciMLBase.DEIntegrator) getstate(layername::Symbol, integrator::SciMLBase.DEIntegrator)
...@@ -160,119 +65,5 @@ Strat.getstate(::Val{layername}, integrator::SciMLBase.DEIntegrator) where {laye ...@@ -160,119 +65,5 @@ Strat.getstate(::Val{layername}, integrator::SciMLBase.DEIntegrator) where {laye
getvar(var::Symbol, integrator::SciMLBase.DEIntegrator) getvar(var::Symbol, integrator::SciMLBase.DEIntegrator)
""" """
Numerics.getvar(var::Symbol, integrator::SciMLBase.DEIntegrator) = Numerics.getvar(Val{var}(), Tile(integrator), integrator.u) Numerics.getvar(var::Symbol, integrator::SciMLBase.DEIntegrator) = Numerics.getvar(Val{var}(), Tile(integrator), integrator.u)
"""
CryoGridOutput(sol::TSol, tspan::NTuple{2,Float64}=(-Inf,Inf)) where {TSol<:SciMLBase.AbstractODESolution}
Constructs a `CryoGridOutput` from the given `ODESolution`. Optional argument `tspan` restricts the time span of the output.
"""
InputOutput.CryoGridOutput(sol::TSol, tspan::NTuple{2,DateTime}) where {TSol<:SciMLBase.AbstractODESolution} = CryoGridOutput(sol, convert_tspan(tspan))
function InputOutput.CryoGridOutput(sol::TSol, tspan::NTuple{2,Float64}=(-Inf,Inf)) where {TSol<:SciMLBase.AbstractODESolution}
# Helper functions for mapping variables to appropriate DimArrays by grid/shape.
withdims(var::Var{name,<:OnGrid{Cells}}, arr, grid, ts) where {name} = DimArray(arr*one(vartype(var))*varunits(var), (Z(round.(typeof(1.0u"m"), cells(grid), digits=5)),Ti(ts)))
withdims(var::Var{name,<:OnGrid{Edges}}, arr, grid, ts) where {name} = DimArray(arr*one(vartype(var))*varunits(var), (Z(round.(typeof(1.0u"m"), edges(grid), digits=5)),Ti(ts)))
withdims(var::Var{name}, arr, zs, ts) where {name} = DimArray(arr*one(vartype(var))*varunits(var), (Ti(ts),))
save_interval = ClosedInterval(tspan...)
model = sol.prob.f.f # Tile
ts = model.hist.vals.t # use save callback time points
t_mask = map((save_interval), ts) # indices within t interval
u_all = reduce(hcat, sol.(ts[t_mask])) # build prognostic state from continuous solution
pax = ComponentArrays.indexmap(getaxes(model.state.uproto)[1])
# get saved diagnostic states and timestamps only in given interval
savedstates = model.hist.vals.saveval[t_mask]
ts_datetime = Dates.epochms2datetime.(round.(ts[t_mask]*1000.0))
allvars = variables(model)
progvars = tuplejoin(filter(isprognostic, allvars), filter(isalgebraic, allvars))
diagvars = filter(isdiagnostic, allvars)
fluxvars = filter(isflux, allvars)
outputs = Dict{Symbol,Any}()
for var in progvars
name = varname(var)
outputs[name] = withdims(var, u_all[pax[name],:], model.grid, ts_datetime)
end
for var in filter(isongrid, tuplejoin(diagvars, fluxvars))
name = varname(var)
states = collect(skipmissing([name keys(state) ? state[name] : missing for state in savedstates]))
if length(states) == length(ts_datetime)
arr = reduce(hcat, states)
outputs[name] = withdims(var, arr, model.grid, ts_datetime)
end
end
# loop over remaining (user defined) log variables
# uservars = Dict()
# for varname in logvarnames
# var_log = log[varname]
# if eltype(var_log) <: AbstractVector
# vardata = reduce(hcat, var_log)
# uservars[varname] = DimArray(vardata, (Z(1:size(vardata,1)),Ti(ts)))
# else
# uservars[varname] = DimArray(vardata, (Ti(ts),))
# end
# end
# if length(logvarnames) > 0
# nt = NamedTuple{tuple(keys(uservars)...)}(tuple(values(uservars)...))
# layerstates = merge(layerstates, (user=nt,))
# end
CryoGridOutput(ts_datetime, sol, (;outputs...))
end
"""
Evaluates the continuous solution at time `t`.
"""
(out::CryoGridOutput{<:ODESolution})(t::Real) = withaxes(out.res(t), out.res.prob.f.f)
(out::CryoGridOutput{<:ODESolution})(t::DateTime) = out(Dates.datetime2epochms(t)/1000.0)
# callback building functions
function _criterionfunc(::Val{name}, i_layer, i_proc, i_cb) where name
function _condition(u,t,integrator)
let tile = Tile(integrator),
comp = tile.strat[i_layer],
layer = comp.layer,
process = comp.processes[i_proc],
cb = tile.callbacks[name][i_cb],
u = Strat.withaxes(u, tile),
du = Strat.withaxes(get_du(integrator), tile),
t = t;
criterion(cb, layer, process, Strat.getstate(Val{name}(), tile, u, du, t))
end
end
end
function _affectfunc(::Val{name}, i_layer, i_proc, i_cb) where name
function _affect!(integrator)
let tile=Tile(integrator),
comp = tile.strat[i_layer],
layer = comp.layer,
process = comp.processes[i_proc],
cb = tile.callbacks[name][i_cb],
u = Strat.withaxes(integrator.u, tile),
du = Strat.withaxes(get_du(integrator), tile),
t = integrator.t;
affect!(cb, layer, process, Strat.getstate(Val{name}(), tile, u, du, t))
end
end
end
_diffeqcallback(::Discrete, ::Val{name}, i_layer, i_proc, i_cb) where name = DiffEqCallbacks.DiscreteCallback(
_criterionfunc(Val{name}(), i_layer, i_proc, i_cb),
_affectfunc(Val{name}(), i_layer, i_proc, i_cb),
# todo: initialize and finalize?
)
_diffeqcallback(::Continuous, ::Val{name}, i_layer, i_proc, i_cb) where name = DiffEqCallbacks.ContinuousCallback(
_criterionfunc(Val{name}(), i_layer, i_proc, i_cb),
_affectfunc(Val{name}(), i_layer, i_proc, i_cb),
# todo: initialize and finalize?
)
function _makecallbacks(component::StratComponent{L,P,name}, i_layer) where {L,P,name}
cbs = []
i_cb = 1
for (i_proc, proc) in enumerate(component.processes)
for callback in callbacks(component.layer, proc)
push!(cbs, _diffeqcallback(CallbackStyle(callback), Val{name}(), i_layer, i_proc, i_cb))
i_cb += 1
end
end
return Tuple(cbs)
end
function _makecallbacks(tile::Tile)
diffeq_callbacks = tuplejoin((_makecallbacks(comp, i) for (i,comp) in enumerate(tile.strat))...)
return diffeq_callbacks
end
end end
# callback building functions
function makecallbacks(component::StratComponent{L,P,name}, i_layer) where {L,P,name}
cbs = []
i_ev = 1
for ev in CryoGrid.events(component.layer, component.process)
push!(cbs, _diffeqcallback(ev, Val{name}(), i_layer, i_ev))
i_ev += 1
end
return Tuple(cbs)
end
function makecallbacks(tile::Tile)
diffeq_callbacks = tuplejoin((makecallbacks(comp, i) for (i,comp) in enumerate(tile.strat))...)
return diffeq_callbacks
end
function _criterionfunc(::Val{name}, i_layer, i_ev) where name
function _condition(u,t,integrator)
let tile = Tile(integrator),
comp = tile.strat[i_layer],
layer = comp.layer,
process = comp.process,
ev = tile.events[name][i_ev],
u = Strat.withaxes(u, tile),
du = Strat.withaxes(get_du(integrator), tile),
t = t;
criterion(ev, layer, process, Strat.getstate(Val{name}(), tile, u, du, t, integrator.dt))
end
end
end
function _triggerfunc(::Val{name}, trig, i_layer, i_ev) where name
_invoke_trigger!(ev, ::Nothing, layer, process, state) = trigger!(ev, layer, process, state)
_invoke_trigger!(ev, trig::ContinuousTrigger, layer, process, state) = trigger!(ev, trig, layer, process, state)
function _trigger!(integrator)
let tile=Tile(integrator),
comp = tile.strat[i_layer],
layer = comp.layer,
process = comp.process,
ev = tile.events[name][i_ev],
u = Strat.withaxes(integrator.u, tile),
du = Strat.withaxes(get_du(integrator), tile),
t = integrator.t,
state = Strat.getstate(Val{name}(), tile, u, du, t, integrator.dt);
_invoke_trigger!(ev, trig, layer, process, state)
end
end
end
_diffeqcallback(::DiscreteEvent, ::Val{name}, i_layer, i_ev) where name = DiffEqCallbacks.DiscreteCallback(
_criterionfunc(Val{name}(), i_layer, i_ev),
_triggerfunc(Val{name}(), nothing, i_layer, i_ev);
# todo: initialize and finalize?
)
_diffeqcallback(::ContinuousEvent, ::Val{name}, i_layer, i_ev) where name = DiffEqCallbacks.ContinuousCallback(
_criterionfunc(Val{name}(), i_layer, i_ev),
_triggerfunc(Val{name}(), Increasing(), i_layer, i_ev);
affect_neg! = _triggerfunc(Val{name}(), Decreasing(), i_layer, i_ev),
# todo: initialize and finalize?
)
"""
(p::DiffEqCallbacks.StepsizeLimiterAffect{typeof(CryoGrid.timestep)})(integrator)
Custom implementation of `StepsizeLimiterAffect` function for `CryoGrid.timestep` that invokes the
`timestep` function with `tile,du,u,p,t` as arguments.
"""
function (p::DiffEqCallbacks.StepsizeLimiterAffect{typeof(CryoGrid.timestep)})(integrator)
integrator.opts.dtmax = p.safety_factor*p.dtFE(integrator.sol.prob.f.f, get_du(integrator), integrator.u, integrator.p, integrator.t)
# This part is copied from the implementation in DiffEqCallbacks
if !integrator.opts.adaptive
if integrator.opts.dtmax < integrator.dtcache
integrator.dtcache = integrator.opts.dtmax
elseif p.cached_dtcache <= integrator.opts.dtmax
integrator.dtcache = p.cached_dtcache
end
end
if p.max_step && isfinite(integrator.opts.dtmax)
set_proposed_dt!(integrator,integrator.opts.dtmax)
integrator.dtcache = integrator.opts.dtmax
end
u_modified!(integrator, false)
end
\ No newline at end of file
"""
CryoGridOutput(sol::TSol, tspan::NTuple{2,Float64}=(-Inf,Inf)) where {TSol<:SciMLBase.AbstractODESolution}
Constructs a `CryoGridOutput` from the given `ODESolution`. Optional argument `tspan` restricts the time span of the output.
"""
InputOutput.CryoGridOutput(sol::TSol, tspan::NTuple{2,DateTime}) where {TSol<:SciMLBase.AbstractODESolution} = CryoGridOutput(sol, convert_tspan(tspan))
function InputOutput.CryoGridOutput(sol::TSol, tspan::NTuple{2,Float64}=(-Inf,Inf)) where {TSol<:SciMLBase.AbstractODESolution}
# Helper functions for mapping variables to appropriate DimArrays by grid/shape.
withdims(var::Var{name,<:OnGrid{Cells}}, arr, grid, ts) where {name} = DimArray(arr*one(vartype(var))*varunits(var), (Z(round.(typeof(1.0u"m"), cells(grid), digits=5)),Ti(ts)))
withdims(var::Var{name,<:OnGrid{Edges}}, arr, grid, ts) where {name} = DimArray(arr*one(vartype(var))*varunits(var), (Z(round.(typeof(1.0u"m"), edges(grid), digits=5)),Ti(ts)))
withdims(var::Var{name}, arr, zs, ts) where {name} = DimArray(arr*one(vartype(var))*varunits(var), (Dim{name}(1:size(arr,1)),Ti(ts)))
save_interval = ClosedInterval(tspan...)
tile = sol.prob.f.f # Tile
ts = tile.hist.vals.t # use save callback time points
t_mask = map((save_interval), ts) # indices within t interval
u_all = reduce(hcat, sol.(ts[t_mask])) # build prognostic state from continuous solution
pax = ComponentArrays.indexmap(getaxes(tile.state.uproto)[1])
# get saved diagnostic states and timestamps only in given interval
savedstates = tile.hist.vals.saveval[t_mask]
ts_datetime = Dates.epochms2datetime.(round.(ts[t_mask]*1000.0))
allvars = variables(tile)
progvars = tuplejoin(filter(isprognostic, allvars), filter(isalgebraic, allvars))
diagvars = filter(isdiagnostic, allvars)
fluxvars = filter(isflux, allvars)
outputs = Dict{Symbol,Any}()
for var in progvars
name = varname(var)
outputs[name] = withdims(var, u_all[pax[name],:], tile.grid, ts_datetime)
end
for var in filter(isongrid, tuplejoin(diagvars, fluxvars))
name = varname(var)
states = collect(skipmissing([name keys(state) ? state[name] : missing for state in savedstates]))
if length(states) == length(ts_datetime)
arr = reduce(hcat, states)
outputs[name] = withdims(var, arr, tile.grid, ts_datetime)
end
end
for layer in Strat.componentnames(tile.strat)
# if layer name appears in saved states, then add these variables to the output.
if haskey(savedstates[1], layer)
layerouts = map(savedstates) do state
layerstate = state[layer]
layerout = Dict()
for var in keys(layerstate)
layerout[var] = layerstate[var]
end
(;layerout...)
end
layerouts_combined = reduce(layerouts[2:end]; init=layerouts[1]) do out1, out2
map(vcat, out1, out2)
end
# for each variable in the named tuple, find the corresponding diagnostic variable in `diagvars`
layervars = (; map(name -> name => first(filter(var -> varname(var) == name, diagvars)), keys(layerouts_combined))...)
# map each output to a variable and call withdims to wrap in a DimArray
outputs[layer] = map((var,out) -> withdims(var, reshape(out,1,:), nothing, ts_datetime), layervars, layerouts_combined)
end
end
return CryoGridOutput(ts_datetime, sol, (;outputs...))
end
"""
Evaluates the continuous solution at time `t`.
"""
(out::CryoGridOutput{<:ODESolution})(t::Real) = withaxes(out.res(t), out.res.prob.f.f)
(out::CryoGridOutput{<:ODESolution})(t::DateTime) = out(Dates.datetime2epochms(t)/1000.0)
"""
Specialized problem type for CryoGrid `ODEProblem`s.
"""
struct CryoGridODEProblem end
"""
CryoGridProblem(setup::Tile, tspan::NTuple{2,Float64}, p=nothing;kwargs...)
CryoGrid specialized constructor for ODEProblem that automatically generates the initial
condition and necessary callbacks.
"""
function CryoGridProblem(
tile::Tile,
u0::ComponentVector,
tspan::NTuple{2,Float64},
p=nothing;
saveat=3600.0,
savevars=(),
save_everystep=false,
save_start=true,
save_end=true,
step_limiter=CryoGrid.timestep,
safety_factor=9//10,
max_step=true,
callback=nothing,
isoutofdomain=Strat.domain(tile),
kwargs...
)
# workaround for bug in DiffEqCallbacks; see https://github.com/SciML/DifferentialEquations.jl/issues/326
# we have to manually expand single-number `saveat` (i.e. time interval for saving) to a step-range.
expandtstep(tstep::Number) = tspan[1]:tstep:tspan[end]
expandtstep(tstep::AbstractVector) = tstep
getsavestate(tile::Tile, u, du) = deepcopy(Strat.getvars(tile.state, Strat.withaxes(u, tile), Strat.withaxes(du, tile), savevars...))
savefunc(u, t, integrator) = getsavestate(Tile(integrator), Strat.withaxes(u, Tile(integrator)), get_du(integrator))
# remove units
model_tile = Model(tile)
model_tile[:val] = p
tile = parent(model_tile)
# collect parameters
p = isnothing(p) ? dustrip.(collect(model_tile[:val])) : collect(p)
du0 = zero(u0)
# remove units
tile = stripunits(tile)
# set up saving callback
stateproto = getsavestate(tile, u0, du0)
savevals = SavedValues(Float64, typeof(stateproto))
savingcallback = SavingCallback(savefunc, savevals; saveat=expandtstep(saveat), save_start=save_start, save_end=save_end, save_everystep=save_everystep)
# add step limiter to default callbacks, if defined
defaultcallbacks = isnothing(step_limiter) ? (savingcallback,) : (savingcallback, StepsizeLimiter(step_limiter; safety_factor, max_step))
# build layer callbacks
layercallbacks = DiffEq.makecallbacks(tile)
# add user callbacks
usercallbacks = isnothing(callback) ? () : callback
callbacks = CallbackSet(defaultcallbacks..., layercallbacks..., usercallbacks...)
# note that this implicitly discards any existing saved values in the model setup's state history
tile.hist.vals = savevals
# set up default mass matrix, M:
# M⋅∂u∂t = f(u)
M_diag = similar(tile.state.uproto)
M_idxmap = ComponentArrays.indexmap(getaxes(M_diag)[1])
allvars = Flatten.flatten(tile.state.vars, Flatten.flattenable, Var)
progvars = map(varname, filter(isprognostic, allvars))
algvars = map(varname, filter(isalgebraic, allvars))
for name in keys(M_idxmap)
M_diag_var = @view M_diag[name]
if name progvars
M_diag_var .= one(eltype(M_diag))
elseif name algvars
M_diag_var .= zero(eltype(M_diag))
end
end
# if no algebraic variables are present, use identity matrix
num_algebraic = length(M_diag) - sum(M_diag)
M = num_algebraic > 0 ? Diagonal(M_diag) : I
func = odefunction(tile, M, u0, p, tspan; kwargs...)
ODEProblem(func, u0, tspan, p, CryoGridODEProblem(); callback=callbacks, isoutofdomain, kwargs...)
end
"""
CryoGridProblem(setup::Tile, tspan::NTuple{2,DateTime}, args...;kwargs...)
"""
CryoGridProblem(setup::Tile, u0::ComponentVector, tspan::NTuple{2,DateTime}, args...;kwargs...) = CryoGridProblem(setup,u0,convert_tspan(tspan),args...;kwargs...)
"""
odefunction(setup::Tile, M, u0, p, tspan; kwargs...)
Constructs a SciML `ODEFunction` given the model setup, mass matrix M, initial state u0, parameters p, and tspan.
Can (and should) be overridden by users to provide customized ODEFunction configurations for specific problem setups, e.g:
```
tile = Tile(strat,grid)
function CryoGrid.Setup.odefunction(::DefaultJac, setup::typeof(tile), M, u0, p, tspan)
...
# make sure to return an instance of ODEFunction
end
...
prob = CryoGridProblem(tile, tspan, p)
```
`JacobianStyle` can also be extended to create custom traits which can then be applied to compatible `Tile`s.
"""
odefunction(setup::TSetup, M, u0, p, tspan; kwargs...) where {TSetup<:Tile} = odefunction(JacobianStyle(TSetup), setup, M, u0, p, tspan; kwargs...)
odefunction(::DefaultJac, setup::TSetup, M, u0, p, tspan; kwargs...) where {TSetup<:Tile} = ODEFunction(setup, mass_matrix=M; kwargs...)
function odefunction(::TridiagJac, setup::Tile, M, u0, p, tspan; kwargs...)
if :jac_prototype in keys(kwargs)
@warn "using user specified jac_prorotype instead of tridiagonal"
ODEFunction(setup, mass_matrix=M; kwargs...)
else
N = length(u0)
J = Tridiagonal(
similar(u0, eltype(p), N-1) |> Vector,
similar(u0, eltype(p), N) |> Vector,
similar(u0, eltype(p), N-1) |> Vector
)
ODEFunction(setup, mass_matrix=M, jac_prototype=J, kwargs...)
end
end
struct CFLStepLimiter{TTile,A}
tile::TTile
Δt::A
default_dt::Float64
end
function (cfl::CFLStepLimiter{<:Tile,<:AbstractArray})(u,p,t)
let Δt = cfl.Δt,
cfl! = Drivers.cfl!(Heat),
Δx = dustrip(Δ(cfl.tile.grid)),
dHdT = getvar(:dHdT, cfl.tile, u), # apparent heat capacity
kc = getvar(:kc, cfl.tile, u); # thermal cond. at grid centers
cfl!(Δt, Δx, dHdT, kc)
Δt_min = minimum(Δt)
IfElse.ifelse(isfinite(Δt_min) && Δt_min > 0, Δt_min, cfl.default_dt)
end
end
function (cfl::CFLStepLimiter{<:Tile,Nothing})(u,p,t)
let cfl = Drivers.cfl(Heat),
Δx = dustrip(Δ(cfl.tile.grid)),
dHdT = getvar(:dHdT, cfl.tile, u), # apparent heat capacity
kc = getvar(:kc, cfl.tile, u); # thermal cond. at grid centers
Δt = cfl(Δt, Δx, dHdT, kc)
Δt_min = minimum(Δt)
IfElse.ifelse(isfinite(Δt_min) && Δt_min > 0, Δt_min, cfl.default_dt)
end
end
function CFLStepLimiter(tile::HeatOnlyTile; courant_number=1//2, default_dt=60.0, iip=true)
cfl = iip ? CFLStepLimiter(tile, zero(dustrip(Δ(tile.grid))), default_dt) : CFLStepLimiter(tile, nothing, default_dt)
StepsizeLimiter(cfl; safety_factor=courant_number, max_step=true)
end
...@@ -27,21 +27,9 @@ JacobianStyle(::Type{<:Tile}) = DefaultJac() ...@@ -27,21 +27,9 @@ JacobianStyle(::Type{<:Tile}) = DefaultJac()
# Auto-detect Jacobian sparsity for problems with one or more heat-only layers. # Auto-detect Jacobian sparsity for problems with one or more heat-only layers.
# Note: This assumes that the processes/forcings on the boundary layers do not violate the tridiagonal structure! # Note: This assumes that the processes/forcings on the boundary layers do not violate the tridiagonal structure!
# Unfortunately, the Stratigraphy type signature is a bit nasty to work with :( # Unfortunately, the Stratigraphy type signature is a bit nasty to work with :(
const HeatOnlyTile = Tile{<:Stratigraphy{N,<:Tuple{TTop,Vararg{<:Union{<:StratComponent{<:SubSurface, <:CoupledProcesses{<:Tuple{<:Heat}}},TBot}}}}} where {N,TTop,TBot} const HeatOnlyTile = Tile{<:Stratigraphy{N,<:Tuple{TTop,Vararg{<:Union{<:StratComponent{<:SubSurface, <:Heat},TBot}}}}} where {N,TTop,TBot}
JacobianStyle(::Type{<:HeatOnlyTile}) = TridiagJac() JacobianStyle(::Type{<:HeatOnlyTile}) = TridiagJac()
# CFL conditions
"""
cfl(::Type{<:SubSurfaceProcess})
cfl!(::Type{<:SubSurfaceProcess})
Returns a function of the form (Δx, args...) -> Δt (or in-place, (Δt, Δx, args...) -> Δt)
which comptues the CFL condition with process-specific parameters `args`.
"""
cfl(::Type{<:SubSurfaceProcess}) = error("not implemented")
cfl(::Type{<:Heat}) = (Δx, dHdT, kc) -> Utils.adstrip(Δx^2 * dHdT / kc)
cfl!(::Type{T}) where {T<:SubSurfaceProcess} = (Δt, Δx, dHdT, kc) -> @. Δt = cfl(T)(Δx, dHdT, kc)
# DiffEq/SciML driver (possibly should be a soft dependency with Requires.jl) # DiffEq/SciML driver (possibly should be a soft dependency with Requires.jl)
include("DiffEq/DiffEq.jl") include("DiffEq/DiffEq.jl")
@reexport using .DiffEq @reexport using .DiffEq
......
...@@ -18,14 +18,15 @@ end ...@@ -18,14 +18,15 @@ end
function Base.show(io::IO, out::CryoGridOutput) function Base.show(io::IO, out::CryoGridOutput)
countvars(x) = 1 countvars(x) = 1
countvars(nt::NamedTuple) = sum(map(countvars, nt)) countvars(nt::NamedTuple) = sum(map(countvars, nt))
format(res::Tuple) = join(res, ", ")
describe(key, val) = "$key => $(typeof(val).name.wrapper) of $(eltype(val)) with dimensions $(size(val))" describe(key, val) = "$key => $(typeof(val).name.wrapper) of $(eltype(val)) with dimensions $(size(val))"
describe(key, val::NamedTuple) = "$key => $(map(describe, keys(val), values(val)))" describe(key, val::NamedTuple) = "$key => $(format(map(describe, keys(val), values(val))))"
data = out.data data = out.data
nvars = countvars(data) nvars = countvars(data)
println(io, "CryoGridOutput with $(length(out.ts)) time steps and $(nvars != 1 ? "$nvars variables" : "1 variable"):") println(io, "CryoGridOutput with $(length(out.ts)) time steps and $(nvars != 1 ? "$nvars variables" : "1 variable"):")
strs = map(describe, keys(data), values(data)) strs = map(describe, keys(data), values(data))
for r in strs for r in strs
println(io, " $r") println(io, " $r")
end end
end end
stack(out::CryoGridOutput, var::Symbol, vars::Symbol...) = DimStack((;map(n -> n => getproperty(out, n), tuple(var, vars...))...)) stack(out::CryoGridOutput, var::Symbol, vars::Symbol...) = DimStack((;map(n -> n => getproperty(out, n), tuple(var, vars...))...))
......
...@@ -17,12 +17,9 @@ using Interpolations ...@@ -17,12 +17,9 @@ using Interpolations
using IntervalSets using IntervalSets
using LinearAlgebra using LinearAlgebra
using LoopVectorization using LoopVectorization
using RuntimeGeneratedFunctions
using Unitful using Unitful
using StructTypes using StructTypes
RuntimeGeneratedFunctions.init(@__MODULE__)
export GridSpec, Edges, Cells export GridSpec, Edges, Cells
export Profile, ProfileKnot export Profile, ProfileKnot
...@@ -52,10 +49,6 @@ struct Profile{N,TKnots} ...@@ -52,10 +49,6 @@ struct Profile{N,TKnots}
Profile(pairs::Tuple{Vararg{<:Pair}}) where D = Profile(map(Base.splat(ProfileKnot), pairs)) Profile(pairs::Tuple{Vararg{<:Pair}}) where D = Profile(map(Base.splat(ProfileKnot), pairs))
Profile(pairs::Pair...) = Profile(pairs) Profile(pairs::Pair...) = Profile(pairs)
end end
Flatten.flattenable(::Type{<:ProfileKnot}, ::Type{Val{:depth}}) = false
Flatten.flattenable(::Type{ProfileKnot{D,T}}, ::Type{Val{:value}}) where {D,T<:Number} = false
Flatten.flattenable(::Type{ProfileKnot{D,T}}, ::Type{Val{:depth}}) where {D<:ModelParameters.Param,T} = true
Flatten.flattenable(::Type{ProfileKnot{D,T}}, ::Type{Val{:value}}) where {D,T<:ModelParameters.Param} = true
Base.length(::Profile{N}) where N = N Base.length(::Profile{N}) where N = N
Base.iterate(profile::Profile) = iterate(profile.knots) Base.iterate(profile::Profile) = iterate(profile.knots)
Base.iterate(profile::Profile, state) = iterate(profile.knots, state) Base.iterate(profile::Profile, state) = iterate(profile.knots, state)
...@@ -72,7 +65,7 @@ export Grid, cells, edges, subgridinds, Δ, volume, area, updategrid! ...@@ -72,7 +65,7 @@ export Grid, cells, edges, subgridinds, Δ, volume, area, updategrid!
include("grid.jl") include("grid.jl")
export Var, Prognostic, Algebraic, Diagnostic, VarDim, OnGrid, Shape, Scalar export Var, Prognostic, Algebraic, Diagnostic, VarDim, OnGrid, Shape, Scalar
export varname, vartype, vardims, varunits, isprognostic, isalgebraic, isflux, isdiagnostic, isongrid, dimlength export varname, vartype, vardims, varunits, vardomain, isprognostic, isalgebraic, isflux, isdiagnostic, isongrid, dimlength
include("variables.jl") include("variables.jl")
export VarStates, DiffCache, retrieve, getvar, getvars export VarStates, DiffCache, retrieve, getvar, getvars
......
...@@ -18,10 +18,10 @@ function discretize(::Type{A}, grid::Grid, pvars::Union{<:Prognostic,<:Algebraic ...@@ -18,10 +18,10 @@ function discretize(::Type{A}, grid::Grid, pvars::Union{<:Prognostic,<:Algebraic
Np = length(pointvar_ns) > 0 ? sum(pointvar_ns) : 0 Np = length(pointvar_ns) > 0 ? sum(pointvar_ns) : 0
# build axis indices; # build axis indices;
# non-grid prognostic variables get collected at the top of the vector, in the order provided # non-grid prognostic variables get collected at the top of the vector, in the order provided
pointvar_ax = (;(varname(p) => i:(i+n) for (p,n,i) in zip(pointvars, pointvar_ns, cumsum(vcat([1],collect(pointvar_ns[1:end-1])))))...) pointvar_ax = (;(varname(p) => i:(i+n-1) for (p,n,i) in zip(pointvars, pointvar_ns, cumsum(vcat([1],collect(pointvar_ns[1:end-1])))))...)
# grid variables get interlaced throughout the rest of the vector; i.e. for variable i, its grid points are: # grid variables get interlaced throughout the rest of the vector; i.e. for variable i, its grid points are:
# i:k:kn where k is the number of grid variables and n is the length of the grid. # i:k:kn where k is the number of grid variables and n is the length of the grid.
gridvar_ax = (;(varname(p) => st:length(gridvars):Ng for (p,st) in zip(gridvars, (Np+1):(1+Np+length(gridvars))))...) gridvar_ax = (;(varname(p) => st:length(gridvars):(Ng+Np) for (p,st) in zip(gridvars, (Np+1):(1+Np+length(gridvars))))...)
# allocate component array; assumes all variables have (and should!) have the same type # allocate component array; assumes all variables have (and should!) have the same type
u = zero(similar(A{vartype(first(pvars))}, Ng+Np)) u = zero(similar(A{vartype(first(pvars))}, Ng+Np))
ComponentVector(u, (Axis(merge(pointvar_ax, gridvar_ax)),)) ComponentVector(u, (Axis(merge(pointvar_ax, gridvar_ax)),))
......
...@@ -74,13 +74,14 @@ end ...@@ -74,13 +74,14 @@ end
@propagate_inbounds Base.getindex(grid::Grid, i::AbstractRange) = grid[grid[first(i)]..grid[last(i)]] @propagate_inbounds Base.getindex(grid::Grid, i::AbstractRange) = grid[grid[first(i)]..grid[last(i)]]
@propagate_inbounds Base.getindex(grid::Grid{S,G,Q,A}, interval::Interval{L,R,Q}) where {S,G,Q,A,L,R} = grid[subgridinds(grid, interval)] @propagate_inbounds Base.getindex(grid::Grid{S,G,Q,A}, interval::Interval{L,R,Q}) where {S,G,Q,A,L,R} = grid[subgridinds(grid, interval)]
@propagate_inbounds Base.getindex(grid::Grid, interval::Interval{L,R,Int}) where {L,R} = Grid(grid, first(grid.bounds)+interval.left-1..first(grid.bounds)+interval.right-1) @propagate_inbounds Base.getindex(grid::Grid, interval::Interval{L,R,Int}) where {L,R} = Grid(grid, first(grid.bounds)+interval.left-1..first(grid.bounds)+interval.right-1)
Base.setindex!(::Grid, args...) = error("setindex! is not allowed for Grid types") Base.setindex!(grid::Grid{Edges}, val, i...) = setindex!(values(grid), val, i...)
Base.setindex!(::Grid{Cells}, args...) = error("setindex! is permitted only for edge grids; use `edges(grid)` and call `updategrid!` directly after.")
""" """
updategrid!(grid::Grid{Edges,G,Q}, vals::Q) where {G,Q} updategrid!(grid::Grid{Edges,G,Q}, vals::Q=grid) where {G,Q}
Overwrites `grid` edges with `vals`, and recomputes grid centers/deltas to be consistent with the new grid. Overwrites `grid` edges with `vals`, and recomputes grid centers/deltas to be consistent with the new grid.
""" """
function updategrid!(grid::Grid{Edges,G,Q}, vals::AbstractVector{Q}) where {G,Q} function updategrid!(grid::Grid{Edges,G,Q}, vals::AbstractVector{Q}=grid) where {G,Q}
z_edges = values(grid) z_edges = values(grid)
z_cells = values(cells(grid)) z_cells = values(cells(grid))
Δz_edges = Δ(grid) Δz_edges = Δ(grid)
......
...@@ -10,7 +10,7 @@ struct FunctionInitializer{varname,F} <: VarInitializer{varname} ...@@ -10,7 +10,7 @@ struct FunctionInitializer{varname,F} <: VarInitializer{varname}
f::F f::F
FunctionInitializer(varname::Symbol, f::F) where {F} = new{varname,F}(f) FunctionInitializer(varname::Symbol, f::F) where {F} = new{varname,F}(f)
end end
(init::FunctionInitializer)(args...) = f(args...) (init::FunctionInitializer)(args...) = init.f(args...)
Base.getindex(init::FunctionInitializer, itrv::Interval) = init Base.getindex(init::FunctionInitializer, itrv::Interval) = init
""" """
InterpInitializer{varname,P,I,E} <: VarInitializer{varname} InterpInitializer{varname,P,I,E} <: VarInitializer{varname}
...@@ -44,10 +44,6 @@ function (init::InterpInitializer{var})(state) where var ...@@ -44,10 +44,6 @@ function (init::InterpInitializer{var})(state) where var
end end
# automatic partitioning of profile based on interval # automatic partitioning of profile based on interval
Base.getindex(init::InterpInitializer{var}, itrv::Interval) where var = InterpInitializer(var, init.profile[itrv], init.interp, init.extrap) Base.getindex(init::InterpInitializer{var}, itrv::Interval) where var = InterpInitializer(var, init.profile[itrv], init.interp, init.extrap)
# constructor for constant initializer function
function _constantinit(::Val{varname}, x) where {varname}
return state -> state[varname] .= x
end
""" """
initializer(varname::Symbol, x::Number) => FunctionInitializer w/ constant initializer(varname::Symbol, x::Number) => FunctionInitializer w/ constant
initializer(varname::Symbol, f::Function) => FunctionInitializer initializer(varname::Symbol, f::Function) => FunctionInitializer
...@@ -55,6 +51,6 @@ end ...@@ -55,6 +51,6 @@ end
Convenience constructor for `VarInitializer` that selects the appropriate initializer type based on the arguments. Convenience constructor for `VarInitializer` that selects the appropriate initializer type based on the arguments.
""" """
initializer(varname::Symbol, x::Number) = FunctionInitializer(varname, _constantinit(Val{varname}(), x)) initializer(varname::Symbol, x::Number) = FunctionInitializer(varname, state -> getproperty(state, varname) .= x)
initializer(varname::Symbol, f::Function) = FunctionInitializer(varname, f) initializer(varname::Symbol, f::Function) = FunctionInitializer(varname, f)
initializer(varname::Symbol, profile::Profile, interp=Linear(), extrap=Flat()) = InterpInitializer(varname, profile, interp, extrap) initializer(varname::Symbol, profile::Profile, interp=Linear(), extrap=Flat()) = InterpInitializer(varname, profile, interp, extrap)
...@@ -12,62 +12,64 @@ dimlength(::Shape{dims}, grid::Grid) where dims = prod(dims) ...@@ -12,62 +12,64 @@ dimlength(::Shape{dims}, grid::Grid) where dims = prod(dims)
dimlength(d::OnGrid{Cells}, grid::Grid) = d.f(length(cells(grid))) dimlength(d::OnGrid{Cells}, grid::Grid) = d.f(length(cells(grid)))
dimlength(d::OnGrid{Edges}, grid::Grid) = d.f(length(edges(grid))) dimlength(d::OnGrid{Edges}, grid::Grid) = d.f(length(edges(grid)))
abstract type Var{name,S<:VarDim,T,units} end abstract type Var{name,S<:VarDim,T,units,domain} end
""" """
Prognostic{name,S,T,units} <: Var{name,S,T,units} Prognostic{name,S,T,units,domain} <: Var{name,S,T,units,domain}
Defines a prognostic (time-integrated) state variable. Defines a prognostic (time-integrated) state variable.
""" """
struct Prognostic{name,S,T,units} <: Var{name,S,T,units} struct Prognostic{name,S,T,units,domain} <: Var{name,S,T,units,domain}
dim::S dim::S
Prognostic(name::Symbol, dims::Union{<:Shape,OnGrid{Cells,typeof(identity)}}, units=NoUnits, ::Type{T}=Float64) where {T} = new{name,typeof(dims),T,units}(dims) Prognostic(name::Symbol, dims::Union{<:Shape,OnGrid{Cells,typeof(identity)}}, units=NoUnits, ::Type{T}=Float64; domain=-Inf..Inf) where {T} = new{name,typeof(dims),T,units,domain}(dims)
Prognostic(::Symbol, dims::OnGrid, args...) where {T} = error("Off-cell prognostic/algebraic spatial variables are not currently supported.") Prognostic(::Symbol, dims::OnGrid, args...) = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
end end
""" """
Algebraic{name,S,T,units} <: Var{name,S,T,units} Algebraic{name,S,T,units,domain} <: Var{name,S,T,units,domain}
Defines an algebraic (implicit) state variable. Defines an algebraic (implicit) state variable.
""" """
struct Algebraic{name,S,T,units} <: Var{name,S,T,units} struct Algebraic{name,S,T,units,domain} <: Var{name,S,T,units,domain}
dim::S dim::S
# maybe a mass matrix init function? # maybe a mass matrix init function?
Algebraic(name::Symbol, dims::Union{<:Shape,OnGrid{Cells,typeof(identity)}}, units=NoUnits, ::Type{T}=Float64) where {T} = new{name,typeof(dims),T,units}(dims) Algebraic(name::Symbol, dims::Union{<:Shape,OnGrid{Cells,typeof(identity)}}, units=NoUnits, ::Type{T}=Float64; domain=-Inf..Inf) where {T} = new{name,typeof(dims),T,units,domain}(dims)
Algebraic(::Symbol, dims::OnGrid, args...) where {T} = error("Off-cell prognostic/algebraic spatial variables are not currently supported.") Algebraic(::Symbol, dims::OnGrid, args...) = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
end end
""" """
Delta{dname,name,S,T,units} <: Var{dname,S,T,units} Delta{dname,name,S,T,units,domain} <: Var{dname,S,T,units,domain}
Defines a "delta" term `du` for variable `u`, which is the time-derivative or flux for prognostic variables and Defines a "delta" term `du` for variable `u`, which is the time-derivative or flux for prognostic variables and
the residual for algebraic variables. the residual for algebraic variables.
""" """
struct Delta{dname,name,S,T,units} <: Var{dname,S,T,units} struct Delta{dname,name,S,T,units,domain} <: Var{dname,S,T,units,domain}
dim::S dim::S
Delta(::Symbol, dims::OnGrid, args...) where {T} = error("Off-cell prognostic/algebraic spatial variables are not currently supported.") Delta(::Symbol, dims::OnGrid, args...) = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
Delta(dname::Symbol, name::Symbol, dims::Union{<:Shape,OnGrid{Cells,typeof(identity)}}, units=NoUnits, ::Type{T}=Float64) where {T} = new{dname,name,typeof(dims),T,units}(dims) Delta(dname::Symbol, name::Symbol, dims::Union{<:Shape,OnGrid{Cells,typeof(identity)}}, units=NoUnits, ::Type{T}=Float64; domain=-Inf..Inf) where {T} = new{dname,name,typeof(dims),T,units,domain}(dims)
Delta(var::Prognostic{name,S,T,units}) where {name,S,T,units} = let dims=vardims(var); new{Symbol(:d,name),name,typeof(dims),T,upreferred(units)/u"s"}(dims) end Delta(var::Prognostic{name,S,T,units,domain}) where {name,S,T,units,domain} = let dims=vardims(var); new{Symbol(:d,name),name,typeof(dims),T,upreferred(units)/u"s",domain}(dims) end
Delta(var::Algebraic{name,S,T,units}) where {name,S,T,units} = let dims=vardims(var); new{Symbol(:d,name),name,typeof(dims),T,units}(dims) end Delta(var::Algebraic{name,S,T,units,domain}) where {name,S,T,units,domain} = let dims=vardims(var); new{Symbol(:d,name),name,typeof(dims),T,units,domain}(dims) end
end end
""" """
Diagnostic{name,S,T,units} <: Var{name,S,T,units} Diagnostic{name,S,T,units,domain} <: Var{name,S,T,units,domain}
Defines a diagnostic variable which is allocated and cached per timestep but not integrated over time. Defines a diagnostic variable which is allocated and cached per timestep but not integrated over time.
""" """
struct Diagnostic{name,S,T,units} <: Var{name,S,T,units} struct Diagnostic{name,S,T,units,domain} <: Var{name,S,T,units,domain}
dim::S dim::S
Diagnostic(name::Symbol, dims::VarDim, units=NoUnits, ::Type{T}=Float64) where {T} = new{name,typeof(dims),T,units}(dims) Diagnostic(name::Symbol, dims::VarDim, units=NoUnits, ::Type{T}=Float64; domain=-Inf..Inf) where {T} = new{name,typeof(dims),T,units,domain}(dims)
end end
# constructors for Flatten.reconstruct # constructors for Flatten.reconstruct
ConstructionBase.constructorof(::Type{Prognostic{name,S,T,units}}) where {name,S,T,units} = s -> Prognostic(name, s, T, units) ConstructionBase.constructorof(::Type{Prognostic{name,S,T,units,domain}}) where {name,S,T,units,domain} = s -> Prognostic(name, s, units, T; domain)
ConstructionBase.constructorof(::Type{Diagnostic{name,S,T,units}}) where {name,S,T,units} = s -> Diagnostic(name, s, T, units) ConstructionBase.constructorof(::Type{Diagnostic{name,S,T,units,domain}}) where {name,S,T,units,domain} = s -> Diagnostic(name, s, units, T; domain)
ConstructionBase.constructorof(::Type{Algebraic{name,S,T,units}}) where {name,S,T,units} = s -> Algebraic(name, s, T, units) ConstructionBase.constructorof(::Type{Algebraic{name,S,T,units,domain}}) where {name,S,T,units,domain} = s -> Algebraic(name, s, units, T; domain)
ConstructionBase.constructorof(::Type{Delta{dname,name,S,T,units}}) where {dname,name,S,T,units} = s -> Delta(dname, name, s, T, units) ConstructionBase.constructorof(::Type{Delta{dname,name,S,T,units,domain}}) where {dname,name,S,T,units,domain} = s -> Delta(dname, name, s, units, T; domain)
==(var1::Var{N1,S1,T1,u1},var2::Var{N2,S2,T2,u2}) where {N1,N2,S1,S2,T1,T2,u1,u2} = (N1==N2) && (S1==S2) && (T1==T2) && (u1 == u2) ==(::Var{N1,S1,T1,u1,d1},::Var{N2,S2,T2,u2,d2}) where {N1,N2,S1,S2,T1,T2,u1,u2,d1,d2} = (N1 == N2) && (S1 == S2) && (T1 == T2) && (u1 == u2) && (d1 == d2)
varname(::Var{name}) where {name} = name varname(::Var{name}) where {name} = name
varname(::Type{<:Var{name}}) where {name} = name varname(::Type{<:Var{name}}) where {name} = name
vartype(::Var{name,S,T}) where {name,S,T} = T vartype(::Var{name,S,T}) where {name,S,T} = T
vartype(::Type{<:Var{name,S,T}}) where {name,S,T} = T vartype(::Type{<:Var{name,S,T}}) where {name,S,T} = T
varunits(::Var{name,S,T,units}) where {name,S,T,units} = units varunits(::Var{name,S,T,units}) where {name,S,T,units} = units
varunits(::Type{<:Var{name,S,T,units}}) where {name,S,T,units} = units varunits(::Type{<:Var{name,S,T,units}}) where {name,S,T,units} = units
vardomain(::Var{name,S,T,units,domain}) where {name,S,T,units,domain} = domain
vardomain(::Type{<:Var{name,S,T,units,domain}}) where {name,S,T,units,domain} = domain
vardims(var::Var) = var.dim vardims(var::Var) = var.dim
isprognostic(::T) where {T<:Var} = T <: Prognostic isprognostic(::T) where {T<:Var} = T <: Prognostic
isalgebraic(::T) where {T<:Var} = T <: Algebraic isalgebraic(::T) where {T<:Var} = T <: Algebraic
......
...@@ -32,14 +32,22 @@ end ...@@ -32,14 +32,22 @@ end
end end
end end
function getvars(vs::VarStates{layers,gridvars,TU}, u::ComponentVector, du::ComponentVector, vals::Union{Symbol,<:Pair{Symbol}}...) where {layers,gridvars,T,A,pax,TU<:ComponentVector{T,A,Tuple{Axis{pax}}}} function getvars(vs::VarStates{layers,gridvars,TU}, u::ComponentVector, du::ComponentVector, vals::Union{Symbol,<:Pair{Symbol}}...) where {layers,gridvars,T,A,pax,TU<:ComponentVector{T,A,Tuple{Axis{pax}}}}
vars = map(filter(n -> n keys(pax), vals)) do val # map over given variable names, ignoring prognostic variables # case 1: grid variable (no layer specified)
isprognostic(name::Symbol) = name keys(pax)
# case 2: non-grid variable on specific layer; ignore here, defer until handled below
isprognostic(other) = false
symbols(name::Symbol) = tuple(name)
symbols(names::NTuple{N,Symbol}) where N = names
vars = map(filter(!(isprognostic), vals)) do val # map over given variable names, ignoring prognostic variables
if val gridvars if val gridvars
val => getvar(Val{val}(), vs, u, du) val => getvar(Val{val}(), vs, u, du)
elseif val map(n -> Symbol(:d,n), keys(pax)) elseif val map(n -> Symbol(:d,n), keys(pax))
val => du[val] val => du[val]
else else
nestedvals = isa(vals[val], Tuple) ? vals[val] : tuple(vals[val]) layername = val[1]
first(val) => (;map(n -> n => retrieve(getproperty(vs.diag[first(val)], n)), nestedvals)...) # handle either a single variable name or multiple, also filtering out prognostic variables
layervars = filter(!(isprognostic), symbols(val[2]))
layername => (;map(n -> n => retrieve(getproperty(vs.diag[layername], n)), layervars)...)
end end
end end
return (;vars...) return (;vars...)
...@@ -63,9 +71,9 @@ Base.show(io::IO, mime::MIME{Symbol("text/plain")}, cache::DiffCache) = show(io, ...@@ -63,9 +71,9 @@ Base.show(io::IO, mime::MIME{Symbol("text/plain")}, cache::DiffCache) = show(io,
function Prealloc.get_tmp(dc::Prealloc.DiffCache, ::Type{T}) where {T<:ForwardDiff.Dual} function Prealloc.get_tmp(dc::Prealloc.DiffCache, ::Type{T}) where {T<:ForwardDiff.Dual}
# this part is copied from PreallocationTools source code # this part is copied from PreallocationTools source code
nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*length(dc.du) nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*length(dc.du)
Prealloc.ArrayInterface.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) Prealloc.ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
end end
retrieve(dc::DiffCache) = dc.du retrieve(dc::DiffCache) = dc.cache.du
# for matching chunk sizes, retrieve from cache # for matching chunk sizes, retrieve from cache
retrieve(dc::DiffCache{N}, ::Type{T}) where {tag,U,N,T<:ForwardDiff.Dual{tag,U,N}} = Prealloc.get_tmp(dc.cache, T) retrieve(dc::DiffCache{N}, ::Type{T}) where {tag,U,N,T<:ForwardDiff.Dual{tag,U,N}} = Prealloc.get_tmp(dc.cache, T)
# otherwise, create new DiffCache on demand # otherwise, create new DiffCache on demand
......
...@@ -13,20 +13,21 @@ boundaryvalue(bc::ConstantBC,l1,p2,l2,s1,s2) = bc.value ...@@ -13,20 +13,21 @@ boundaryvalue(bc::ConstantBC,l1,p2,l2,s1,s2) = bc.value
BoundaryStyle(::Type{<:ConstantBC{P,S}}) where {P,S} = S() BoundaryStyle(::Type{<:ConstantBC{P,S}}) where {P,S} = S()
""" """
PeriodicBC{P,S,T1,T2,T3} <: BoundaryProcess{P} PeriodicBC{P,S,T1,T2,T3,T4} <: BoundaryProcess{P}
Periodic boundary condition (of any type/unit) specified by `period`, `amplitude`, and `phaseshift`. Periodic boundary condition (of any type/unit) specified by `period`, `amplitude`, and `phaseshift`.
""" """
struct PeriodicBC{P,S,T1,T2,T3} <: BoundaryProcess{P} struct PeriodicBC{P,S,T1,T2,T3,T4} <: BoundaryProcess{P}
period::T1 period::T1
amplitude::T2 amplitude::T2
phaseshift::T3 phaseshift::T3
PeriodicBC(::Type{P}, ::Type{S}, period::T1, amplitude::T2=one(T), phaseshift::T3=one(T)) where level::T4
{P<:SubSurfaceProcess,S<:BoundaryStyle,T1<:TimeQuantity,T2,T3} = PeriodicBC(::Type{P}, ::Type{S}, period::T1=1.0, amplitude::T2=1.0, phaseshift::T3=0.0, level::T4=0.0) where
new{P,S,T1,T2,T3}(period, amplitude, phaseshift) {P<:SubSurfaceProcess,S<:BoundaryStyle,T1<:TimeQuantity,T2,T3,T4} =
new{P,S,T1,T2,T3,T4}(period, amplitude, phaseshift, level)
end end
ConstructionBase.constructorof(::Type{<:PeriodicBC{P,S}}) where {P,S} = (args...) -> PeriodicBC(P, S, args...) ConstructionBase.constructorof(::Type{<:PeriodicBC{P,S}}) where {P,S} = (args...) -> PeriodicBC(P, S, args...)
@inline boundaryvalue(bc::PeriodicBC,l1,p2,l2,s1,s2) = bc.amplitude*sin(π*(1/bc.period)*t + bc.phaseshift) @inline boundaryvalue(bc::PeriodicBC,l1,p2,l2,s1,s2) = bc.amplitude*sin(π*(1/bc.period)*t + bc.phaseshift) + bc.level
BoundaryStyle(::Type{<:PeriodicBC{P,S}}) where {P,S} = S() BoundaryStyle(::Type{<:PeriodicBC{P,S}}) where {P,S} = S()
......