Implement implicit solvers using OrdinaryDiffEq.jl

[charleskawczynski]

Over the course of our development, we've moved to fully explicit time stepping in order to

• Unify the spatial operators
• Simplify timestepping to ease the transition to using OrdinaryDiffEq.jl
• Simplify/unify tendency collection

The relavent PRs (and their slow downs, pending) are:

• Make env diffusion explicit #486
• Make environment advection explicit #485
• Make env entr-detr and mixing length src terms explicit #484
• solve grid mean equations explicitly  #343
• Set timestep to 3 sec for all simulations  #115
• Remove substepping and unify timesteps in the model #111
  • ⛅ Bomex 3m 28s -> 4m 12s
  • ⛅ life_cycle_Tan2018 3m 27s -> 4m 15s
  • ⛅ Soares 3m 37s -> 4m 35s
  • ⛅ Rico 4m 58s -> 9m 22s
  • ⛅ Nieuwstadt 3m 39s -> 4m 26s
  • ✂️ TRMM_LBA 4m 2s -> 6m 55s
  • ⛅ ARM_SGP 4m 19s -> 5m 58s
  • ⛅ DYCOMS_RF01 4m 10s -> 5m 8s
  • ⛅ GABLS 3m 53s -> 3m 35s
  • ⛅ SP 3m 49s -> 4m 8s
  • 💭 DryBubble 3m 9s -> 2m 59s

We should now effectively revert this PRs, by using OrdinaryDiffEq.jl and treating specific terms (or all terms) implicitly in time.

[charleskawczynski]

When first pulling apart the implicit solver, I wrote this brief summary:

    N_a terms = 1 (diffusion)
    N_diagonal terms = 6 (unsteady+upwind_advection+2diffusion+entr_detr+dissipation)
    N_c terms = 2 (diffusion+upwind_advection)
    ∂_t ρa*tke                              ✓ unsteady
      + ∂_z(ρaw*tke) =                      ✓ advection
      + ρaK*(∂²_z(u)+∂²_z(v)+∂²_z(w̄))
      + ρawΣᵢ(εᵢ(wⱼ-w₀)²-δ₀*tke)            ✓ entr_detr
      + ∂_z(ρa₀K ∂_z(tke))                  ✓ diffusion
      + ρa₀*w̅₀b̅₀
      - a₀(u⋅∇p)
      + ρa₀D                                ✓ dissipation

[Zhengyu-Huang]

The advection and diffusion term can be approximated as

∂_z(ρaw*tke)  = ∂_z(ρaw^{n}*tke^{n+1}) 

∂_z(ρa₀K ∂_z(tke))  = ∂_z(ρa₀K^{n}∂_z(tke)^{n+1})      

superscripts indicate time, then this will give you a tridiagonal matrix about tke^{n+1}

The source term is a diagonal matrix ?

[charleskawczynski]

Some useful lines:

• TurbulenceConvection.jl/src/Turbulence_PrognosticTKE.jl

  Lines 469 to 522 in e041f03

  	# ----------- TODO: move to compute_tendencies
  	implicit_eqs = self.implicit_eqs
  	# Matrix is the same for all variables that use the same eddy diffusivity, we can construct once and reuse
  	implicit_eqs.A_θq_gm .= construct_tridiag_diffusion_gm(grid, TS.dt, self.rho_ae_KH, ref_state.rho0_half, self.ae)
  	implicit_eqs.A_uv_gm .= construct_tridiag_diffusion_gm(grid, TS.dt, self.rho_ae_KM, ref_state.rho0_half, self.ae)
  	Δzi = grid.Δzi
  	@inbounds for k in real_center_indices(grid)
  	implicit_eqs.b_u_gm[k] = GMV.U.values[k]
  	implicit_eqs.b_v_gm[k] = GMV.V.values[k]
  	implicit_eqs.b_q_tot_gm[k] = self.EnvVar.QT.values[k]
  	implicit_eqs.b_θ_liq_ice_gm[k] = self.EnvVar.H.values[k]
  	if is_surface_center(grid, k)
  	implicit_eqs.b_u_gm[k] += TS.dt * Case.Sur.rho_uflux * Δzi * ref_state.alpha0_half[k] / self.ae[k]
  	implicit_eqs.b_v_gm[k] += TS.dt * Case.Sur.rho_vflux * Δzi * ref_state.alpha0_half[k] / self.ae[k]
  	implicit_eqs.b_q_tot_gm[k] += TS.dt * Case.Sur.rho_qtflux * Δzi * ref_state.alpha0_half[k] / self.ae[k]
  	implicit_eqs.b_θ_liq_ice_gm[k] += TS.dt * Case.Sur.rho_hflux * Δzi * ref_state.alpha0_half[k] / self.ae[k]
  	end
  	end
  	gm = GMV
  	up = self.UpdVar
  	en = self.EnvVar
  	KM = diffusivity_m(self).values
  	KH = diffusivity_h(self).values
  	common_args = (
  	grid,
  	param_set,
  	ref_state,
  	TS,
  	KM,
  	KH,
  	up.Area.bulkvalues,
  	up.Area.values,
  	up.W.values,
  	en.W.values,
  	en.TKE.values,
  	self.n_updrafts,
  	self.minimum_area,
  	self.pressure_plume_spacing,
  	self.frac_turb_entr,
  	self.entr_sc,
  	self.mixing_length,
  	)
  	implicit_eqs.A_TKE .= construct_tridiag_diffusion_en(common_args..., true)
  	implicit_eqs.A_Hvar .= construct_tridiag_diffusion_en(common_args..., false)
  	implicit_eqs.A_QTvar .= construct_tridiag_diffusion_en(common_args..., false)
  	implicit_eqs.A_HQTcov .= construct_tridiag_diffusion_en(common_args..., false)
  	implicit_eqs.b_TKE .= en_diffusion_tendencies(grid, ref_state, TS, up.Area.values, en.TKE)
  	implicit_eqs.b_Hvar .= en_diffusion_tendencies(grid, ref_state, TS, up.Area.values, en.Hvar)
  	implicit_eqs.b_QTvar .= en_diffusion_tendencies(grid, ref_state, TS, up.Area.values, en.QTvar)
  	implicit_eqs.b_HQTcov .= en_diffusion_tendencies(grid, ref_state, TS, up.Area.values, en.HQTcov)
  	# -----------

• TurbulenceConvection.jl/src/Turbulence_PrognosticTKE.jl

  Lines 1500 to 1526 in e041f03

  	function en_diffusion_tendencies(grid::Grid, ref_state::ReferenceState, TS, a_up, covar)
  	dti = TS.dti
  	b = center_field(grid)
  	ρ0_c = ref_state.rho0_half
  	ae = 1 .- up_sum(a_up)
  	kc_surf = kc_surface(grid)
  	covar_surf = covar.values[kc_surf]
  	@inbounds for k in real_center_indices(grid)
  	if is_surface_center(grid, k)
  	b[k] = covar_surf
  	else
  	b[k] = (
  	ρ0_c[k] * ae[k] * covar.values[k] * dti +
  	covar.press[k] +
  	covar.buoy[k] +
  	covar.shear[k] +
  	covar.entr_gain[k] +
  	covar.rain_src[k]
  	)
  	end
  	end
  	return b
  	end

• TurbulenceConvection.jl/src/Operators.jl

  Lines 527 to 658 in e041f03

  	function construct_tridiag_diffusion_gm(grid::Grid, dt, ρ_ae_K_m, ρ_0, ae)
  	a = center_field(grid) # for tridiag solver
  	b = center_field(grid) # for tridiag solver
  	c = center_field(grid) # for tridiag solver
  	Δzi = grid.Δzi
  	@inbounds for k in real_center_indices(grid)
  	X = ρ_0[k] * ae[k] / dt
  	Y = ρ_ae_K_m[k + 1] * Δzi * Δzi
  	Z = ρ_ae_K_m[k] * Δzi * Δzi
  	if is_surface_center(grid, k)
  	Z = 0.0
  	elseif is_toa_center(grid, k)
  	Y = 0.0
  	end
  	a[k] = -Z / X
  	b[k] = 1.0 + Y / X + Z / X
  	c[k] = -Y / X
  	end
  	A = LinearAlgebra.Tridiagonal(a[2:end], vec(b), c[1:(end - 1)])
  	return A
  	end
  	tridiag_solve(b_rhs, A) = A \ b_rhs
  	#= TODO: clean this up somehow!
  	construct_tridiag_diffusion_en configures the tridiagonal
  	matrix for solving 2nd order environment variables where
  	some terms are treated implicitly. Here is a list of all
  	the terms in the matrix (and the overall equation solved):
  	N_a terms = 1 (diffusion)
  	N_diagonal terms = 6 (unsteady+upwind_advection+2diffusion+entr_detr+dissipation)
  	N_c terms = 2 (diffusion+upwind_advection)
  	∂_t ρa*tke ✓ unsteady
  	+ ∂_z(ρaw*tke) = ✓ advection
  	+ ρaK*(∂²_z(u)+∂²_z(v)+∂²_z(w̄))
  	+ ρawΣᵢ(εᵢ(wⱼ-w₀)²-δ₀*tke) ✓ entr_detr
  	+ ∂_z(ρa₀K ∂_z(tke)) ✓ diffusion
  	+ ρa₀*w̅₀b̅₀
  	- a₀(u⋅∇p)
  	+ ρa₀D ✓ dissipation
  	=#
  	function construct_tridiag_diffusion_en(
  	grid::Grid,
  	param_set::APS,
  	ref_state::ReferenceState,
  	TS,
  	KM,
  	KH,
  	a_up_bulk,
  	a_up,
  	w_up,
  	w_en,
  	tke_en,
  	n_updrafts::Int,
  	minimum_area::Float64,
  	pressure_plume_spacing::Vector,
  	frac_turb_entr,
  	entr_sc,
  	mixing_length,
  	is_tke,
  	)
  	c_d = CPEDMF.c_d(param_set)
  	kc_surf = kc_surface(grid)
  	kc_toa = kc_top_of_atmos(grid)
  	Δzi = grid.Δzi
  	dti = TS.dti
  	a = center_field(grid)
  	b = center_field(grid)
  	c = center_field(grid)
  	ρ0_c = ref_state.rho0_half
  	ae = 1 .- a_up_bulk
  	rho_ae_K_m = face_field(grid)
  	w_en_c = center_field(grid)
  	D_env = 0.0
  	aeK = is_tke ? ae .* KM : ae .* KH
  	aeK_bcs = (; bottom = SetValue(aeK[kc_surf]), top = SetValue(aeK[kc_toa]))
  	@inbounds for k in real_face_indices(grid)
  	rho_ae_K_m[k] = interpc2f(aeK, grid, k; aeK_bcs...) * ref_state.rho0[k]
  	end
  	@inbounds for k in real_center_indices(grid)
  	w_en_c[k] = interpf2c(w_en, grid, k)
  	end
  	@inbounds for k in real_center_indices(grid)
  	D_env = 0.0
  	@inbounds for i in xrange(n_updrafts)
  	if a_up[i, k] > minimum_area
  	turb_entr = frac_turb_entr[i, k]
  	R_up = pressure_plume_spacing[i]
  	w_up_c = interpf2c(w_up, grid, k, i)
  	D_env += ρ0_c[k] * a_up[i, k] * w_up_c * (entr_sc[i, k] + turb_entr)
  	else
  	D_env = 0.0
  	end
  	end
  	# TODO: this tridiagonal matrix needs to be re-verified, as it's been pragmatically
  	# modified to not depend on ghost points, and these changes have not been
  	# carefully verified.
  	if is_surface_center(grid, k)
  	a[k] = 0.0
  	b[k] = 1.0
  	c[k] = 0.0
  	else
  	a[k] = (-rho_ae_K_m[k] * Δzi * Δzi)
  	b[k] = (
  	ρ0_c[k] * ae[k] * dti - ρ0_c[k] * ae[k] * w_en_c[k] * Δzi +
  	rho_ae_K_m[k + 1] * Δzi * Δzi +
  	rho_ae_K_m[k] * Δzi * Δzi +
  	D_env +
  	ρ0_c[k] * ae[k] * c_d * sqrt(max(tke_en[k], 0)) / max(mixing_length[k], 1)
  	)
  	if is_toa_center(grid, k)
  	c[k] = 0.0
  	else
  	c[k] = (ρ0_c[k + 1] * ae[k + 1] * w_en_c[k + 1] * Δzi - rho_ae_K_m[k + 1] * Δzi * Δzi)
  	end
  	end
  	end
  	A = LinearAlgebra.Tridiagonal(a[2:end], vec(b), c[1:(end - 1)])
  	return A
  	end

cc @Zhengyu-Huang

[yairchn]

The advection and diffusion term can be approximated as

∂_z(ρaw*tke)  = ∂_z(ρaw^{n}*tke^{n+1}) 

∂_z(ρa₀K ∂_z(tke))  = ∂_z(ρa₀K^{n}∂_z(tke)^{n+1})      

This is indeed what we did in SCAMPy, but be mindful that K is a. function of tke so that you are mixing time steps (it worked reasonably well). The same is true for the buoyancy gradient and shear sources where K is used