Skip to content

Commit

Permalink
Figure out the difference between the two Jacobian partitioning schemes
Browse files Browse the repository at this point in the history
  • Loading branch information
dennisYatunin committed Oct 6, 2023
1 parent a5f27d7 commit 722f800
Show file tree
Hide file tree
Showing 8 changed files with 165 additions and 34 deletions.
27 changes: 27 additions & 0 deletions .buildkite/pipeline.yml
Original file line number Diff line number Diff line change
Expand Up @@ -696,6 +696,33 @@ steps:
agents:
slurm_mem: 20GB

- label: ":genie: EDMFX Bomex in a box (Float64)"
command: >
julia --color=yes --project=examples examples/hybrid/driver.jl
--config_file $CONFIG_PATH/edmfx_bomex_box_ft64.yml
artifact_paths: "edmfx_bomex_box_ft64/*"
soft_fail: true
agents:
slurm_mem: 20GB

- label: ":genie: EDMFX Bomex in a box (alternative partitioning)"
command: >
julia --color=yes --project=examples examples/hybrid/driver.jl
--config_file $CONFIG_PATH/edmfx_bomex_box_alternative_partitioning.yml
artifact_paths: "edmfx_bomex_box_alternative_partitioning/*"
soft_fail: true
agents:
slurm_mem: 20GB

- label: ":genie: EDMFX Bomex in a box (alternative partitioning + Float64)"
command: >
julia --color=yes --project=examples examples/hybrid/driver.jl
--config_file $CONFIG_PATH/edmfx_bomex_box_alternative_partitioning_ft64.yml
artifact_paths: "edmfx_bomex_box_alternative_partitioning_ft64/*"
soft_fail: true
agents:
slurm_mem: 20GB

- label: ":genie: EDMFX Rico in a box"
command: >
julia --color=yes --project=examples examples/hybrid/driver.jl
Expand Down
3 changes: 3 additions & 0 deletions config/default_configs/default_config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,9 @@ smoothing_order:
help: "Define the surface smoothing kernel factor (integer) [`3 (default)`]"
value: 3
# ODE
partition_by_cell_face_values:
help: "Experimental flag for the Jacobian that switches between SchurComplementSolve(@name(f)) and SchurComplementSolve(@name(f.u₃))"
value: false
use_newton_rtol:
help: "Whether to check if the current iteration of Newton's method has an error within a relative tolerance, instead of always taking the maximum number of iterations (only for ClimaTimeSteppers.jl)"
value: false
Expand Down
31 changes: 31 additions & 0 deletions config/model_configs/edmfx_bomex_box_alternative_partitioning.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
job_id: "edmfx_bomex_box_alternative_partitioning"
initial_condition: "Bomex"
subsidence: "Bomex"
edmf_coriolis: "Bomex"
ls_adv: "Bomex"
surface_setup: "Bomex"
turbconv: "edmfx"
edmfx_upwinding: first_order
edmfx_entr_model: "ConstantCoefficient"
edmfx_detr_model: "ConstantCoefficient"
edmfx_sgs_mass_flux: true
edmfx_sgs_diffusive_flux: true
edmfx_nh_pressure: false
prognostic_tke: false
moist: "equil"
config: "box"
hyperdiff: "true"
kappa_4: 1.0e12
x_max: 1e5
y_max: 1e5
z_max: 3e3
x_elem: 2
y_elem: 2
z_elem: 60
z_stretch: false
perturb_initstate: false
partition_by_cell_face_values: true
dt: "1secs"
t_end: "10800secs"
dt_save_to_disk: "10secs"
toml: [toml/edmfx_box.toml]
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
job_id: "edmfx_bomex_box_alternative_partitioning_ft64"
initial_condition: "Bomex"
subsidence: "Bomex"
edmf_coriolis: "Bomex"
ls_adv: "Bomex"
surface_setup: "Bomex"
turbconv: "edmfx"
edmfx_upwinding: first_order
edmfx_entr_model: "ConstantCoefficient"
edmfx_detr_model: "ConstantCoefficient"
edmfx_sgs_mass_flux: true
edmfx_sgs_diffusive_flux: true
edmfx_nh_pressure: false
prognostic_tke: false
moist: "equil"
config: "box"
hyperdiff: "true"
kappa_4: 1.0e12
x_max: 1e5
y_max: 1e5
z_max: 3e3
x_elem: 2
y_elem: 2
z_elem: 60
z_stretch: false
perturb_initstate: false
partition_by_cell_face_values: true
FLOAT_TYPE: Float64
dt: "1secs"
t_end: "10800secs"
dt_save_to_disk: "10secs"
toml: [toml/edmfx_box.toml]
31 changes: 31 additions & 0 deletions config/model_configs/edmfx_bomex_box_ft64.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
job_id: "edmfx_bomex_box_ft64"
initial_condition: "Bomex"
subsidence: "Bomex"
edmf_coriolis: "Bomex"
ls_adv: "Bomex"
surface_setup: "Bomex"
turbconv: "edmfx"
edmfx_upwinding: first_order
edmfx_entr_model: "ConstantCoefficient"
edmfx_detr_model: "ConstantCoefficient"
edmfx_sgs_mass_flux: true
edmfx_sgs_diffusive_flux: true
edmfx_nh_pressure: false
prognostic_tke: false
moist: "equil"
config: "box"
hyperdiff: "true"
kappa_4: 1.0e12
x_max: 1e5
y_max: 1e5
z_max: 3e3
x_elem: 2
y_elem: 2
z_elem: 60
z_stretch: false
perturb_initstate: false
FLOAT_TYPE: Float64
dt: "1secs"
t_end: "10800secs"
dt_save_to_disk: "10secs"
toml: [toml/edmfx_box.toml]
4 changes: 2 additions & 2 deletions examples/Manifest.toml
Original file line number Diff line number Diff line change
Expand Up @@ -299,11 +299,11 @@ version = "0.5.4"

[[deps.ClimaCore]]
deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"]
git-tree-sha1 = "d94b29e8db7e8f0b09bee51c0f6a8fe98a99bb4b"
git-tree-sha1 = "9e5f8f0ec1c93f6db8efc68959402b0ed26d1b5f"
repo-rev = "dy/field_matrix"
repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
version = "0.10.51"
version = "0.10.52"

[[deps.ClimaCoreMakie]]
deps = ["ClimaCore", "Makie"]
Expand Down
59 changes: 31 additions & 28 deletions src/prognostic_equations/implicit/implicit_solver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,18 @@ import LinearAlgebra: I, Adjoint, ldiv!
using ForwardDiff: Dual
using ClimaCore.MatrixFields

abstract type JacobianMode end
struct EnthalpyCorrection <: JacobianMode end
struct NoEnthalpyCorrection <: JacobianMode end
abstract type EnthalpyCorrectionFlag end
struct CorrectEnthalpy <: EnthalpyCorrectionFlag end
struct DoNotCorrectEnthalpy <: EnthalpyCorrectionFlag end

abstract type BlockPartitioningFlag end
struct PartitionByGridMeanVelocity <: BlockPartitioningFlag end
struct PartitionByCellFaceValues <: BlockPartitioningFlag end

struct ImplicitEquationJacobian{
W <: MatrixFields.FieldMatrix,
S <: MatrixFields.FieldMatrixSolver,
M <: JacobianMode,
E <: EnthalpyCorrectionFlag,
FT,
T <: Fields.FieldVector{FT},
}
Expand All @@ -20,7 +24,7 @@ struct ImplicitEquationJacobian{
solver::S

# whether to compute ∂(ᶜh_tot)/∂(ᶠu₃) or to approximate it as 0
mode::M
enthalpy_flag::E

# required by OrdinaryDiffEq.jl to run non-Rosenbrock integration algorithms
transform::Bool
Expand All @@ -33,7 +37,8 @@ end

function ImplicitEquationJacobian(
Y;
mode = NoEnthalpyCorrection(),
enthalpy_flag = DoNotCorrectEnthalpy(),
partition_flag = PartitionByGridMeanVelocity(),
transform = false,
)
FT = Spaces.undertype(axes(Y.c))
Expand All @@ -43,40 +48,39 @@ function ImplicitEquationJacobian(
Quaddiagonal_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
Tridiagonal_C3xACT3 = TridiagonalMatrixRow{typeof(one_C3xACT3)}

name_tree = MatrixFields.FieldNameTree(Y)
is_name_in_Y(name) = MatrixFields.has_subtree_at_name(name_tree, name)
is_in_Y(name) = MatrixFields.has_field(Y, name)

energy_name = is_name_in_Y(@name(c.ρe_tot)) ? @name(c.ρe_tot) : @name(c.ρθ)
possible_energy_names = (@name(c.ρe_tot), @name(c.ρθ))
energy_name = MatrixFields.unrolled_findonly(is_in_Y, possible_energy_names)
possible_tracer_names = (
@name(c.ρq_tot),
@name(c.ρq_liq),
@name(c.ρq_ice),
@name(c.ρq_rai),
@name(c.ρq_sno),
)
tracer_names =
MatrixFields.unrolled_filter(is_name_in_Y, possible_tracer_names)
possible_remaining_names = (
tracer_names = MatrixFields.unrolled_filter(is_in_Y, possible_tracer_names)
possible_other_names = (
@name(c.sgs⁰.ρatke),
@name(c.sgsʲs),
@name(f.sgsʲs),
@name(c.turbconv),
@name(f.turbconv),
@name(sfc),
)
remaining_names =
MatrixFields.unrolled_filter(is_name_in_Y, possible_remaining_names)
other_names = MatrixFields.unrolled_filter(is_in_Y, possible_other_names)

W = MatrixFields.FieldMatrix(
(@name(c.ρ), @name(c.ρ)) => FT(-1) * I,
(@name(c.uₕ), @name(c.uₕ)) => FT(-1) * I,
MatrixFields.unrolled_map(
name -> (name, name) => FT(-1) * I,
(energy_name, tracer_names..., remaining_names...),
(energy_name, tracer_names..., other_names...),
)...,
(@name(c.ρ), @name(f.u₃)) => similar(Y.c, Bidiagonal_ACT3),
(energy_name, @name(f.u₃)) =>
energy_name == @name(c.ρe_tot) && mode isa EnthalpyCorrection ?
energy_name == @name(c.ρe_tot) &&
enthalpy_flag isa CorrectEnthalpy ?
similar(Y.c, Quaddiagonal_ACT3) : similar(Y.c, Bidiagonal_ACT3),
MatrixFields.unrolled_map(
name -> (name, @name(f.u₃)) => similar(Y.c, Bidiagonal_ACT3),
Expand All @@ -87,12 +91,14 @@ function ImplicitEquationJacobian(
(@name(f.u₃), @name(f.u₃)) => similar(Y.f, Tridiagonal_C3xACT3),
)

alg = MatrixFields.SchurComplementSolve(@name(f.u₃))
partition_name =
partition_flag isa PartitionByCellFaceValues ? @name(f) : @name(f.u₃)
alg = MatrixFields.SchurComplementSolve(partition_name)

return ImplicitEquationJacobian(
W,
MatrixFields.FieldMatrixSolver(alg, W, Y),
mode,
enthalpy_flag,
transform,
Ref{FT}(),
similar(Y),
Expand Down Expand Up @@ -159,7 +165,7 @@ function Wfact!(A, Y, p, dtγ, t)
(rayleigh_sponge isa RayleighSponge ? (; p.ᶠβ_rayleigh_w) : (;))...,
)

# Convert dtγ from a Float64 to an FT.
# Convert dtγ from a Float64 to an FT and save it to dtγ_ref.
FT = Spaces.undertype(axes(Y.c))
dtγ′ = FT(dtγ)
A.dtγ_ref[] = dtγ′
Expand All @@ -171,7 +177,7 @@ function Wfact!(A, Y, p, dtγ, t)
end

function _Wfact!(A, Y, p, dtγ, colidx)
(; W, mode) = A
(; W, enthalpy_flag) = A
(; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK∂ᶠu₃) = p
(; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params) = p
(; energy_upwinding, tracer_upwinding, density_upwinding) = p
Expand All @@ -192,9 +198,6 @@ function _Wfact!(A, Y, p, dtγ, colidx)
ᶜJ = Fields.local_geometry_field(Y.c).J
ᶠg³³ = g³³_field(Y.f)

name_tree = MatrixFields.FieldNameTree(Y)
is_name_in_Y(name) = MatrixFields.has_subtree_at_name(name_tree, name)

# We can express the kinetic energy as
# ᶜK =
# adjoint(CT12(ᶜuₕ)) * ᶜuₕ / 2 +
Expand Down Expand Up @@ -268,7 +271,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
)
ρχ_and_χ_names =
MatrixFields.unrolled_filter(possible_ρχ_and_χ_names) do (ρχ_name, _)
is_name_in_Y(ρχ_name)
MatrixFields.has_field(Y, ρχ_name)
end
MatrixFields.unrolled_foreach(ρχ_and_χ_names) do (ρχ_name, χ_name)
∂ᶜρχₜ∂ᶠu₃ = W[ρχ_name, @name(f.u₃)]
Expand All @@ -289,7 +292,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
tracer_upwinding
end
if upwinding == Val(:none)
if ρχ_name == @name(c.ρe_tot) && mode isa EnthalpyCorrection
if ρχ_name == @name(c.ρe_tot) && enthalpy_flag isa CorrectEnthalpy
@. ∂ᶜρχₜ∂ᶠu₃[colidx] =
dtγ * (
-(ᶜadvdivᵥ_matrix())
Expand Down Expand Up @@ -319,7 +322,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
ᶠupwind_matrix =
upwinding == Val(:third_order) ? ᶠupwind3_matrix :
ᶠupwind1_matrix
if ρχ_name == @name(c.ρe_tot) && mode isa EnthalpyCorrection
if ρχ_name == @name(c.ρe_tot) && enthalpy_flag isa CorrectEnthalpy
@. ∂ᶜρχₜ∂ᶠu₃[colidx] =
dtγ * (
-(ᶜadvdivᵥ_matrix())
Expand Down Expand Up @@ -394,7 +397,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
∂ᶠu₃ₜ∂ᶜρ = W[@name(f.u₃), @name(c.ρ)]
∂ᶠu₃ₜ∂ᶠu₃ = W[@name(f.u₃), @name(f.u₃)]
minus_I_u₃ = -DiagonalMatrixRow(one_C3xACT3)
if is_name_in_Y(@name(c.ρθ))
if MatrixFields.has_field(Y, @name(c.ρθ))
ᶜρθ = Y.c.ρθ
∂ᶠu₃ₜ∂ᶜρθ = W[@name(f.u₃), @name(c.ρθ)]
@. ∂ᶠu₃ₜ∂ᶜρ[colidx] =
Expand Down Expand Up @@ -422,7 +425,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
else
@. ∂ᶠu₃ₜ∂ᶠu₃[colidx] = (minus_I_u₃,)
end
elseif is_name_in_Y(@name(c.ρe_tot))
elseif MatrixFields.has_field(Y, @name(c.ρe_tot))
∂ᶠu₃ₜ∂ᶜρe_tot = W[@name(f.u₃), @name(c.ρe_tot)]
@. ∂ᶠu₃ₜ∂ᶜρ[colidx] =
dtγ * (
Expand Down
12 changes: 8 additions & 4 deletions src/solver/type_getters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -382,10 +382,14 @@ additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (;
is_cts_algo(::SciMLBase.AbstractODEAlgorithm) = false
is_cts_algo(::CTS.DistributedODEAlgorithm) = true

function jac_kwargs(ode_algo, Y)
function jac_kwargs(parsed_args, ode_algo, Y)
if is_implicit(ode_algo)
A = ImplicitEquationJacobian(Y; transform = use_transform(ode_algo))
if use_transform(ode_algo)
partition_flag =
parsed_args["partition_by_cell_face_values"] ?
PartitionByCellFaceValues() : PartitionByGridMeanVelocity()
transform = use_transform(ode_algo)
A = ImplicitEquationJacobian(Y; partition_flag, transform)
if transform
return (; jac_prototype = A, Wfact_t = Wfact!)
else
return (; jac_prototype = A, Wfact = Wfact!)
Expand Down Expand Up @@ -708,7 +712,7 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback)
func = if parsed_args["split_ode"]
implicit_func = SciMLBase.ODEFunction(
implicit_tendency!;
jac_kwargs(ode_algo, Y)...,
jac_kwargs(parsed_args, ode_algo, Y)...,
tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0),
)
if is_cts_algo(ode_algo)
Expand Down

0 comments on commit 722f800

Please sign in to comment.