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)
Hide whitespace changes
Inline Side-by-side
Showing
with 409 additions and 346 deletions
Project.toml
View file @ 738aa942
name = "CryoGrid"
uuid = "a535b82e-5f3d-4d97-8b0b-d6483f5bebd5"
authors = ["Brian Groenke <brian.groenke@awi.de>", "Moritz Langer <moritz.langer@awi.de>"]
version = "0.9.0"
version = "0.10.3"
[deps]
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
......@@ -29,7 +29,6 @@ PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
SimulationLogs = "e3c6736c-93d9-479d-93f3-96ff5e8836d5"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
......@@ -40,6 +39,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
[compat]
julia = ">= 1.6"
PreallocationTools = ">= 0.3.1"
[extras]
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
......
docs/make.jl
View file @ 738aa942
......@@ -10,7 +10,7 @@ const modules = [
CryoGrid.Physics,
CryoGrid.Boundaries,
CryoGrid.HeatConduction,
CryoGrid.WaterBalance,
CryoGrid.Hydrology,
CryoGrid.Soils,
CryoGrid.SEB,
CryoGrid.Strat,
......@@ -35,7 +35,7 @@ makedocs(modules=modules,
"Utilities" => "api/utils.md",
"Physics" => [
"Heat Conduction" => "api/heat_conduction.md",
"Water Balance" => "api/water_balance.md",
"Hydrology" => "api/hydrology.md",
"Soils" => "api/soils.md",
"Surface Energy Balance" => "api/seb.md",
],
......
docs/src/api/water_balance.md → docs/src/api/hydrology.md
View file @ 738aa942
# Water balance
# Hydrology
```@autodocs
Modules = [WaterBalance]
Modules = [Hydrology]
Private = false
Order = [:type, :function, :macro]
```
\ No newline at end of file
examples/heat_freeW_samoylov_snow.jl 0 → 100644
View file @ 738aa942
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)
examples/heat_vgfc_seb_samoylov_custom.jl
View file @ 738aa942
......@@ -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"
# soil profile: depth => (excess ice, natural porosity, saturation, organic fraction)
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.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.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),
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),
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),
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" => 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" => 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" => 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" => 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
tempprofile = TemperatureProfile(
......@@ -49,7 +49,7 @@ u0, du0 = initialcondition!(model, tspan, p)
# CryoGrid front-end for ODEProblem
prob = CryoGridProblem(model,u0,tspan,p,savevars=(:T,))
# 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!
zs = [1:10...,20:10:100...]
cg = Plots.cgrad(:copper,rev=true);
......
src/CryoGrid.jl
View file @ 738aa942
......@@ -11,33 +11,35 @@ using Reexport
# Common types and methods
export Layer, SubSurface, Top, Bottom
export Process, SubSurfaceProcess, BoundaryProcess, CoupledProcesses, Coupled
export Process, SubSurfaceProcess, BoundaryProcess, CoupledProcesses
export Coupled, Coupled2, Coupled3, Coupled4
include("types.jl")
export variables, initialcondition!, diagnosticstep!, prognosticstep!, interact!
export boundaryflux, boundaryvalue, criterion, affect!, observe
export variables, initialcondition!, diagnosticstep!, prognosticstep!, interact!, timestep
export boundaryflux, boundaryvalue, criterion, trigger!, observe
include("methods.jl")
# Submodules
include("Utils/Utils.jl")
using .Utils
export convert_t, convert_tspan, pstrip, @pstrip, @sym_str
include("Numerics/Numerics.jl")
include("IO/InputOutput.jl")
include("Physics/Physics.jl")
include("Strat/Strat.jl")
include("Diagnostics/Diagnostics.jl")
using .Numerics
export Grid, cells, edges, subgridinds, Δ, volume, area, initializer, getvar
using .Utils
export convert_t, convert_tspan, pstrip, @pstrip, @sym_str
# Re-exported submodules
include("IO/InputOutput.jl")
@reexport using .InputOutput
include("Physics/Physics.jl")
@reexport using .Physics
include("Strat/Strat.jl")
@reexport using .Strat
@reexport using .InputOutput
include("Diagnostics/Diagnostics.jl")
@reexport using .Diagnostics
# Coupling
include("coupling.jl")
# Re-exported packages
@reexport using Dates: Dates, DateTime
@reexport using DimensionalData: DimArray, Z, Dim, dims
@reexport using DimensionalData
@reexport using IntervalSets
@reexport using Unitful
......
src/Drivers/DiffEq/DiffEq.jl
View file @ 738aa942
......@@ -3,15 +3,16 @@ Driver module SciML diffeq solvers.
"""
module DiffEq
using CryoGrid
using CryoGrid: Event, ContinuousEvent, DiscreteEvent, ContinuousTrigger, Increasing, Decreasing
using CryoGrid.Drivers
using CryoGrid: Strat, SubSurface, CoupledProcesses, Callback, CallbackStyle, Discrete, Continuous
using CryoGrid.InputOutput
using CryoGrid.Numerics
using CryoGrid.Physics: Heat
using CryoGrid.Strat: Stratigraphy, StratComponent
using CryoGrid.Utils
import CryoGrid: variables, callbacks, criterion, affect!
import CryoGrid.Strat: Tile, Stratigraphy, StratComponent
import CryoGrid.Strat
using ComponentArrays
using Dates
......@@ -29,124 +30,28 @@ using DiffEqBase
using DiffEqBase.SciMLBase
using DiffEqCallbacks
import DiffEqCallbacks
@reexport using OrdinaryDiffEq
@reexport using DiffEqBase: solve, init, ODEProblem, SciMLBase
export CryoGridProblem
export CFLStepLimiter
include("steplimiters.jl")
export TDMASolver
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.
"""
struct CryoGridODEProblem end
Tile(integrator::SciMLBase.DEIntegrator)
Constructs a `Tile` from a `SciMLBase` integrator.
"""
function Strat.Tile(integrator::SciMLBase.DEIntegrator)
tile = integrator.sol.prob.f.f
return Strat.updateparams(tile, Strat.withaxes(integrator.u, tile), integrator.p, integrator.t)
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(layername::Symbol, integrator::SciMLBase.DEIntegrator)
......@@ -160,119 +65,5 @@ Strat.getstate(::Val{layername}, integrator::SciMLBase.DEIntegrator) where {laye
getvar(var::Symbol, integrator::SciMLBase.DEIntegrator)
"""
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
src/Drivers/DiffEq/callbacks.jl 0 → 100644
View file @ 738aa942
# 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
src/Drivers/DiffEq/output.jl 0 → 100644
View file @ 738aa942
"""
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)
src/Drivers/DiffEq/problem.jl 0 → 100644
View file @ 738aa942
"""
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
src/Drivers/DiffEq/steplimiters.jl deleted 100644 → 0
View file @ 57f0b540
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
src/Drivers/Drivers.jl
View file @ 738aa942
......@@ -27,21 +27,9 @@ JacobianStyle(::Type{<:Tile}) = DefaultJac()
# 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!
# 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()
# 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)
include("DiffEq/DiffEq.jl")
@reexport using .DiffEq
......
src/IO/output.jl
View file @ 738aa942
......@@ -18,14 +18,15 @@ end
function Base.show(io::IO, out::CryoGridOutput)
countvars(x) = 1
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::NamedTuple) = "$key => $(map(describe, keys(val), values(val)))"
describe(key, val::NamedTuple) = "$key => $(format(map(describe, keys(val), values(val))))"
data = out.data
nvars = countvars(data)
println(io, "CryoGridOutput with $(length(out.ts)) time steps and $(nvars != 1 ? "$nvars variables" : "1 variable"):")
strs = map(describe, keys(data), values(data))
for r in strs
println(io, " $r")
println(io, " $r")
end
end
stack(out::CryoGridOutput, var::Symbol, vars::Symbol...) = DimStack((;map(n -> n => getproperty(out, n), tuple(var, vars...))...))
......
src/Numerics/Numerics.jl
View file @ 738aa942
......@@ -17,12 +17,9 @@ using Interpolations
using IntervalSets
using LinearAlgebra
using LoopVectorization
using RuntimeGeneratedFunctions
using Unitful
using StructTypes
RuntimeGeneratedFunctions.init(@__MODULE__)
export GridSpec, Edges, Cells
export Profile, ProfileKnot
......@@ -52,10 +49,6 @@ struct Profile{N,TKnots}
Profile(pairs::Tuple{Vararg{<:Pair}}) where D = Profile(map(Base.splat(ProfileKnot), pairs))
Profile(pairs::Pair...) = Profile(pairs)
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.iterate(profile::Profile) = iterate(profile.knots)
Base.iterate(profile::Profile, state) = iterate(profile.knots, state)
......@@ -72,7 +65,7 @@ export Grid, cells, edges, subgridinds, Δ, volume, area, updategrid!
include("grid.jl")
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")
export VarStates, DiffCache, retrieve, getvar, getvars
......
src/Numerics/discretize.jl
View file @ 738aa942
......@@ -18,10 +18,10 @@ function discretize(::Type{A}, grid::Grid, pvars::Union{<:Prognostic,<:Algebraic
Np = length(pointvar_ns) > 0 ? sum(pointvar_ns) : 0
# build axis indices;
# 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:
# 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
u = zero(similar(A{vartype(first(pvars))}, Ng+Np))
ComponentVector(u, (Axis(merge(pointvar_ax, gridvar_ax)),))
......
src/Numerics/grid.jl
View file @ 738aa942
......@@ -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{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)
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.
"""
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_cells = values(cells(grid))
Δz_edges = Δ(grid)
......
src/Numerics/init.jl
View file @ 738aa942
......@@ -10,7 +10,7 @@ struct FunctionInitializer{varname,F} <: VarInitializer{varname}
f::F
FunctionInitializer(varname::Symbol, f::F) where {F} = new{varname,F}(f)
end
(init::FunctionInitializer)(args...) = f(args...)
(init::FunctionInitializer)(args...) = init.f(args...)
Base.getindex(init::FunctionInitializer, itrv::Interval) = init
"""
InterpInitializer{varname,P,I,E} <: VarInitializer{varname}
......@@ -44,10 +44,6 @@ function (init::InterpInitializer{var})(state) where var
end
# 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)
# 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, f::Function) => FunctionInitializer
......@@ -55,6 +51,6 @@ end
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, profile::Profile, interp=Linear(), extrap=Flat()) = InterpInitializer(varname, profile, interp, extrap)
src/Numerics/variables.jl
View file @ 738aa942
......@@ -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{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.
"""
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
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(::Symbol, dims::OnGrid, args...) where {T} = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
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...) = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
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.
"""
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
# 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(::Symbol, dims::OnGrid, args...) where {T} = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
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...) = error("Off-cell prognostic/algebraic spatial variables are not currently supported.")
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
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
Delta(::Symbol, dims::OnGrid, args...) where {T} = 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(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::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(::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; domain=-Inf..Inf) where {T} = new{dname,name,typeof(dims),T,units,domain}(dims)
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,domain}) where {name,S,T,units,domain} = let dims=vardims(var); new{Symbol(:d,name),name,typeof(dims),T,units,domain}(dims) 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.
"""
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
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
# 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{Diagnostic{name,S,T,units}}) where {name,S,T,units} = s -> Diagnostic(name, s, T, units)
ConstructionBase.constructorof(::Type{Algebraic{name,S,T,units}}) where {name,S,T,units} = s -> Algebraic(name, s, T, units)
ConstructionBase.constructorof(::Type{Delta{dname,name,S,T,units}}) where {dname,name,S,T,units} = s -> Delta(dname, name, s, T, units)
==(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)
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,domain}}) where {name,S,T,units,domain} = s -> Diagnostic(name, s, units, T; domain)
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,domain}}) where {dname,name,S,T,units,domain} = s -> Delta(dname, name, s, units, T; domain)
==(::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(::Type{<:Var{name}}) where {name} = name
vartype(::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(::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
isprognostic(::T) where {T<:Var} = T <: Prognostic
isalgebraic(::T) where {T<:Var} = T <: Algebraic
......
src/Numerics/varstates.jl
View file @ 738aa942
......@@ -32,14 +32,22 @@ 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}}}}
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
val => getvar(Val{val}(), vs, u, du)
elseif val ∈ map(n -> Symbol(:d,n), keys(pax))
val => du[val]
else
nestedvals = isa(vals[val], Tuple) ? vals[val] : tuple(vals[val])
first(val) => (;map(n -> n => retrieve(getproperty(vs.diag[first(val)], n)), nestedvals)...)
layername = val[1]
# 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
return (;vars...)
......@@ -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}
# this part is copied from PreallocationTools source code
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
retrieve(dc::DiffCache) = dc.du
retrieve(dc::DiffCache) = dc.cache.du
# 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)
# otherwise, create new DiffCache on demand
......
src/Physics/Boundaries/bc.jl
View file @ 738aa942
......@@ -13,20 +13,21 @@ boundaryvalue(bc::ConstantBC,l1,p2,l2,s1,s2) = bc.value
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`.
"""
struct PeriodicBC{P,S,T1,T2,T3} <: BoundaryProcess{P}
struct PeriodicBC{P,S,T1,T2,T3,T4} <: BoundaryProcess{P}
period::T1
amplitude::T2
phaseshift::T3
PeriodicBC(::Type{P}, ::Type{S}, period::T1, amplitude::T2=one(T), phaseshift::T3=one(T)) where
{P<:SubSurfaceProcess,S<:BoundaryStyle,T1<:TimeQuantity,T2,T3} =
new{P,S,T1,T2,T3}(period, amplitude, phaseshift)
level::T4
PeriodicBC(::Type{P}, ::Type{S}, period::T1=1.0, amplitude::T2=1.0, phaseshift::T3=0.0, level::T4=0.0) where
{P<:SubSurfaceProcess,S<:BoundaryStyle,T1<:TimeQuantity,T2,T3,T4} =
new{P,S,T1,T2,T3,T4}(period, amplitude, phaseshift, level)
end
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()
......