Skip to content

Commit 5be90de

Browse files
Merge pull request #3853 from CliMA/ck/slim_cache2
Remove more dependency on specific quantities
2 parents 4607ec2 + 01dafa1 commit 5be90de

File tree

10 files changed

+73
-80
lines changed

10 files changed

+73
-80
lines changed

perf/jet_test_nfailures.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ using Test
3838
# inference. By increasing this counter, we acknowledge that
3939
# we have introduced an inference failure. We hope to drive
4040
# this number down to 0.
41-
n_allowed_failures = 249
41+
n_allowed_failures = 253
4242
@show n
4343
if n < n_allowed_failures
4444
@info "Please update the n-failures to $n"

src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -961,7 +961,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
961961
(; params) = p
962962
(; dt) = p
963963
(; ᶜp, ᶜu, ᶜts) = p.precomputed
964-
(; q_tot) = p.precomputed.ᶜspecific
965964
(; ustar, obukhov_length) = p.precomputed.sfc_conditions
966965
(; ᶜtke⁰) = p.precomputed
967966
(;
@@ -1079,14 +1078,13 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
10791078
microphys_0m_params = CAP.microphysics_0m_params(p.params)
10801079
(; dt) = p
10811080
(; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
1082-
(; q_tot) = p.precomputed.ᶜspecific
10831081

10841082
# Environment precipitation sources (to be applied to grid mean)
10851083
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
10861084
thermo_params,
10871085
microphys_0m_params,
10881086
dt,
1089-
q_tot,
1087+
specific(Y.c.ρq_tot, Y.c.ρ),
10901088
ᶜts,
10911089
)
10921090
return nothing
@@ -1102,7 +1100,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
11021100
#microphys_1m_params = CAP.microphysics_1m_params(p.params)
11031101

11041102
#(; ᶜts, ᶜSqₜᵖ⁰, ᶜSeₜᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰) = p.precomputed
1105-
#(; q_tot) = p.precomputed.ᶜspecific
11061103
#(; ᶜqᵣ, ᶜqₛ) = p.precomputed
11071104

11081105
#ᶜSᵖ = p.scratch.ᶜtemp_scalar

src/cache/precipitation_precomputed_quantities.jl

Lines changed: 13 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -453,39 +453,29 @@ function set_precipitation_surface_fluxes!(
453453
)
454454
(; surface_rain_flux, surface_snow_flux) = p.precomputed
455455
(; col_integrated_precip_energy_tendency,) = p.conservation_check
456-
(; ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ, ᶜspecific) = p.precomputed
456+
(; ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ) = p.precomputed
457457
ᶜJ = Fields.local_geometry_field(Y.c).J
458458
ᶠJ = Fields.local_geometry_field(Y.f).J
459459
sfc_J = Fields.level(ᶠJ, Fields.half)
460460
sfc_space = axes(sfc_J)
461461

462462
# Jacobian-weighted extrapolation from interior to surface, consistent with
463463
# the reconstruction of density on cell faces, ᶠρ = ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ
464-
int_J = Fields.Field(Fields.field_values(Fields.level(ᶜJ, 1)), sfc_space)
465-
int_ρ = Fields.Field(Fields.field_values(Fields.level(Y.c.ρ, 1)), sfc_space)
464+
sfc_lev(x) =
465+
Fields.Field(Fields.field_values(Fields.level(x, 1)), sfc_space)
466+
int_J = sfc_lev(ᶜJ)
467+
int_ρ = sfc_lev(Y.c.ρ)
466468
sfc_ρ = @. lazy(int_ρ * int_J / sfc_J)
467469

468470
# Constant extrapolation to surface, consistent with simple downwinding
469-
sfc_qᵣ = Fields.Field(
470-
Fields.field_values(Fields.level(ᶜspecific.q_rai, 1)),
471-
sfc_space,
472-
)
473-
sfc_qₛ = Fields.Field(
474-
Fields.field_values(Fields.level(ᶜspecific.q_sno, 1)),
475-
sfc_space,
476-
)
477-
sfc_qₗ = Fields.Field(
478-
Fields.field_values(Fields.level(ᶜspecific.q_liq, 1)),
479-
sfc_space,
480-
)
481-
sfc_qᵢ = Fields.Field(
482-
Fields.field_values(Fields.level(ᶜspecific.q_ice, 1)),
483-
sfc_space,
484-
)
485-
sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space)
486-
sfc_wₛ = Fields.Field(Fields.field_values(Fields.level(ᶜwₛ, 1)), sfc_space)
487-
sfc_wₗ = Fields.Field(Fields.field_values(Fields.level(ᶜwₗ, 1)), sfc_space)
488-
sfc_wᵢ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵢ, 1)), sfc_space)
471+
sfc_wᵣ = sfc_lev(ᶜwᵣ)
472+
sfc_wₛ = sfc_lev(ᶜwₛ)
473+
sfc_wₗ = sfc_lev(ᶜwₗ)
474+
sfc_wᵢ = sfc_lev(ᶜwᵢ)
475+
sfc_qᵣ = lazy.(specific.(sfc_lev(Y.c.ρq_rai), sfc_ρ))
476+
sfc_qₛ = lazy.(specific.(sfc_lev(Y.c.ρq_sno), sfc_ρ))
477+
sfc_qₗ = lazy.(specific.(sfc_lev(Y.c.ρq_liq), sfc_ρ))
478+
sfc_qᵢ = lazy.(specific.(sfc_lev(Y.c.ρq_ice), sfc_ρ))
489479

490480
@. surface_rain_flux = sfc_ρ * (sfc_qᵣ * (-sfc_wᵣ) + sfc_qₗ * (-sfc_wₗ))
491481
@. surface_snow_flux = sfc_ρ * (sfc_qₛ * (-sfc_wₛ) + sfc_qᵢ * (-sfc_wᵢ))

src/parameterized_tendencies/radiation/radiation.jl

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -425,12 +425,16 @@ end
425425

426426
function radiation_model_cache(Y, radiation_mode::RadiationDYCOMS)
427427
FT = Spaces.undertype(axes(Y.c))
428-
NT = NamedTuple{(:z, , :q_tot), NTuple{3, FT}}
428+
# The NT type is needed for the `column_reduce!` call below because
429+
# `column_reduce!` computes the output, `isoline_z_ρ_ρq`, in-place. `Y.c.ρ`
430+
# and `Y.c.ρq_tot` are input arguments, and `nt1` / `nt2` are output entry
431+
# values at each point along the vertical column.
432+
NT = NamedTuple{(:z, , :ρq_tot), NTuple{3, FT}}
429433
return (;
430434
ᶜκρq = similar(Y.c, FT),
431435
∫_0_∞_κρq = similar(Spaces.level(Y.c, 1), FT),
432436
ᶠ∫_0_z_κρq = similar(Y.f, FT),
433-
isoline_z_ρ_q = similar(Spaces.level(Y.c, 1), NT),
437+
isoline_z_ρ_ρq = similar(Spaces.level(Y.c, 1), NT),
434438
ᶠradiation_flux = similar(Y.f, Geometry.WVector{FT}),
435439
net_energy_flux_toa = [Geometry.WVector(FT(0))],
436440
net_energy_flux_sfc = [Geometry.WVector(FT(0))],
@@ -440,13 +444,13 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS)
440444
@assert !(p.atmos.moisture_model isa DryModel)
441445

442446
(; params) = p
443-
(; ᶜspecific, ᶜts) = p.precomputed
444-
(; ᶜκρq, ∫_0_∞_κρq, ᶠ∫_0_z_κρq, isoline_z_ρ_q, ᶠradiation_flux) =
447+
(; ᶜκρq, ∫_0_∞_κρq, ᶠ∫_0_z_κρq, isoline_z_ρ_ρq, ᶠradiation_flux) =
445448
p.radiation
449+
(; ᶜts) = p.precomputed
446450
thermo_params = CAP.thermodynamics_params(params)
447451
cp_d = CAP.cp_d(params)
448452
FT = Spaces.undertype(axes(Y.c))
449-
NT = NamedTuple{(:z, , :q_tot), NTuple{3, FT}}
453+
NT = NamedTuple{(:z, , :ρq_tot), NTuple{3, FT}}
450454
ᶜz = Fields.coordinate_field(Y.c).z
451455
ᶠz = Fields.coordinate_field(Y.f).z
452456

@@ -467,14 +471,14 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS)
467471
q_tot_isoline = FT(0.008)
468472
Operators.column_reduce!(
469473
(nt1, nt2) ->
470-
abs(nt1.q_tot - q_tot_isoline) < abs(nt2.q_tot - q_tot_isoline) ?
471-
nt1 : nt2,
472-
isoline_z_ρ_q,
473-
Base.broadcasted(NT tuple, ᶜz, Y.c.ρ, ᶜspecific.q_tot),
474+
abs(specific(nt1.ρq_tot, nt1.ρ) - q_tot_isoline) <
475+
abs(specific(nt2.ρq_tot, nt2.ρ) - q_tot_isoline) ? nt1 : nt2,
476+
isoline_z_ρ_ρq,
477+
Base.broadcasted(NT tuple, ᶜz, Y.c.ρ, Y.c.ρq_tot),
474478
)
475479

476-
zi = isoline_z_ρ_q.z
477-
ρi = isoline_z_ρ_q.ρ
480+
zi = isoline_z_ρ_ρq.z
481+
ρi = isoline_z_ρ_ρq.ρ
478482

479483
# TODO: According to the paper, we should remove the ifelse condition that
480484
# clips the third term to 0 below zi, and we should also replace cp_d with

src/prognostic_equations/advection.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
228228
end
229229

230230
if !(p.atmos.moisture_model isa DryModel) && tracer_upwinding != Val(:none)
231-
ᶜq_tot = ᶜspecific.q_tot
231+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
232232
vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), tracer_upwinding)
233233
vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), Val(:none))
234234
@. Yₜ.c.ρq_tot += vtt - vtt_central

src/prognostic_equations/edmfx_sgs_flux.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ function edmfx_sgs_mass_flux_tendency!(
4040

4141
n = n_mass_flux_subdomains(turbconv_model)
4242
(; edmfx_sgsflux_upwinding) = p.atmos.numerics
43-
(; ᶠu³, ᶜh_tot, ᶜspecific) = p.precomputed
43+
(; ᶠu³, ᶜh_tot) = p.precomputed
4444
(; ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = p.precomputed
4545
(; ᶜρa⁰, ᶜρ⁰, ᶠu³⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed
4646
if (
@@ -236,7 +236,7 @@ function edmfx_sgs_mass_flux_tendency!(
236236
a_max = CAP.max_area(turbconv_params)
237237
n = n_mass_flux_subdomains(turbconv_model)
238238
(; edmfx_sgsflux_upwinding) = p.atmos.numerics
239-
(; ᶠu³, ᶜh_tot, ᶜspecific) = p.precomputed
239+
(; ᶠu³, ᶜh_tot) = p.precomputed
240240
(; ᶜρaʲs, ᶜρʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs) = p.precomputed
241241
(; dt) = p
242242
ᶜJ = Fields.local_geometry_field(Y.c).J
@@ -450,7 +450,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
450450
(; dt, params) = p
451451
turbconv_params = CAP.turbconv_params(params)
452452
c_d = CAP.tke_diss_coeff(turbconv_params)
453-
(; ᶜu, ᶜh_tot, ᶜspecific, ᶜtke⁰, ᶜmixing_length) = p.precomputed
453+
(; ᶜu, ᶜh_tot, ᶜtke⁰, ᶜmixing_length) = p.precomputed
454454
(; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
455455
ᶠgradᵥ = Operators.GradientC2F()
456456

src/prognostic_equations/forcing/external_forcing.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -337,7 +337,7 @@ function external_forcing_tendency!(
337337
# horizontal advection, vertical fluctuation, nudging, subsidence (need to add),
338338
(; params) = p
339339
thermo_params = CAP.thermodynamics_params(params)
340-
(; ᶜspecific, ᶜts, ᶜh_tot) = p.precomputed
340+
(; ᶜts, ᶜh_tot) = p.precomputed
341341
(;
342342
ᶜdTdt_fluc,
343343
ᶜdqtdt_fluc,
@@ -398,11 +398,12 @@ function external_forcing_tendency!(
398398
ᶜh_tot,
399399
Val{:first_order}(),
400400
)
401+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
401402
subsidence!(
402403
Yₜ.c.ρq_tot,
403404
Y.c.ρ,
404405
ᶠls_subsidence³,
405-
ᶜspecific.q_tot,
406+
ᶜq_tot,
406407
Val{:first_order}(),
407408
)
408409

src/prognostic_equations/forcing/subsidence.jl

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ scalar quantity `χ`.
5555
Arguments:
5656
- `Yₜ`: The tendency state vector, modified in place.
5757
- `Y`: The current state vector (used for `Y.c.ρ`).
58-
- `p`: Cache containing parameters, precomputed fields (`ᶜh_tot`, `ᶜspecific`),
58+
- `p`: Cache containing parameters, precomputed fields (`ᶜh_tot`),
5959
atmospheric model configurations (`p.atmos.moisture_model`, `p.atmos.subsidence`),
6060
and scratch space.
6161
- `t`: Current simulation time (unused by this specific tendency calculation).
@@ -68,7 +68,7 @@ subsidence_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing # No subsidence
6868
function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence)
6969
(; moisture_model) = p.atmos
7070
subsidence_profile = p.atmos.subsidence.prof
71-
(; ᶜh_tot, ᶜspecific) = p.precomputed
71+
(; ᶜh_tot) = p.precomputed
7272

7373
ᶠz = Fields.coordinate_field(axes(Y.f)).z
7474
ᶠlg = Fields.local_geometry_field(Y.f)
@@ -78,26 +78,23 @@ function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence)
7878

7979
# Large-scale subsidence
8080
subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}())
81-
subsidence!(
82-
Yₜ.c.ρq_tot,
83-
Y.c.ρ,
84-
ᶠsubsidence³,
85-
ᶜspecific.q_tot,
86-
Val{:first_order}(),
87-
)
81+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
82+
subsidence!(Yₜ.c.ρq_tot, Y.c.ρ, ᶠsubsidence³, ᶜq_tot, Val{:first_order}())
8883
if moisture_model isa NonEquilMoistModel
84+
ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
8985
subsidence!(
9086
Yₜ.c.ρq_liq,
9187
Y.c.ρ,
9288
ᶠsubsidence³,
93-
ᶜspecific.q_liq,
89+
ᶜq_liq,
9490
Val{:first_order}(),
9591
)
92+
ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
9693
subsidence!(
9794
Yₜ.c.ρq_ice,
9895
Y.c.ρ,
9996
ᶠsubsidence³,
100-
ᶜspecific.q_ice,
97+
ᶜq_ice,
10198
Val{:first_order}(),
10299
)
103100
end

src/prognostic_equations/implicit/implicit_tendency.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,15 +139,16 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
139139
ᶜJ = Fields.local_geometry_field(Y.c).J
140140
ᶠJ = Fields.local_geometry_field(Y.f).J
141141
(; ᶠgradᵥ_ᶜΦ) = p.core
142-
(; ᶜh_tot, ᶜspecific, ᶠu³, ᶜp) = p.precomputed
142+
(; ᶜh_tot, ᶠu³, ᶜp) = p.precomputed
143143

144144
@. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠu³)
145145

146146
# Central vertical advection of active tracers (e_tot and q_tot)
147147
vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜh_tot, dt, Val(:none))
148148
@. Yₜ.c.ρe_tot += vtt
149149
if !(moisture_model isa DryModel)
150-
vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜspecific.q_tot, dt, Val(:none))
150+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
151+
vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜq_tot, dt, Val(:none))
151152
@. Yₜ.c.ρq_tot += vtt
152153
end
153154

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -399,7 +399,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
399399
(; matrix) = cache
400400
(; params) = p
401401
(; ᶜΦ, ᶠgradᵥ_ᶜΦ) = p.core
402-
(; ᶜspecific, ᶠu³, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed
402+
(; ᶠu³, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed
403403
(;
404404
∂ᶜK_∂ᶜuₕ,
405405
∂ᶜK_∂ᶠu₃,
@@ -477,20 +477,22 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
477477
∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)]
478478
@. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix DiagonalMatrixRow(g³³(ᶠgⁱʲ))
479479

480-
tracer_info = (
481-
(@name(c.ρe_tot), @name(ᶜh_tot)),
482-
(@name(c.ρq_tot), @name(ᶜspecific.q_tot)),
483-
)
484-
MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name)
480+
tracer_info = (@name(c.ρe_tot), @name(c.ρq_tot))
481+
MatrixFields.unrolled_foreach(tracer_info) do ρχ_name
485482
MatrixFields.has_field(Y, ρχ_name) || return
486-
ᶜχ = MatrixFields.get_field(p.precomputed, χ_name)
483+
ᶜχ = if ρχ_name === @name(c.ρe_tot)
484+
p.precomputed.ᶜh_tot
485+
else
486+
@. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
487+
end
487488
if use_derivative(topography_flag)
488489
∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)]
490+
@. ∂ᶜρχ_err_∂ᶜuₕ =
491+
dtγ * ᶜadvection_matrix DiagonalMatrixRow(ᶠinterp(ᶜχ))
492+
ᶠwinterp_matrix(ᶜJ * ᶜρ) DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ))
489493
end
494+
490495
∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)]
491-
use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ =
492-
dtγ * ᶜadvection_matrix DiagonalMatrixRow(ᶠinterp(ᶜχ))
493-
ᶠwinterp_matrix(ᶜJ * ᶜρ) DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ))
494496
@. ∂ᶜρχ_err_∂ᶠu₃ =
495497
dtγ * ᶜadvection_matrix DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ))
496498
end
@@ -536,12 +538,12 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
536538
end
537539

538540
tracer_info = (
539-
(@name(c.ρq_liq), @name(q_liq), @name(ᶜwₗ)),
540-
(@name(c.ρq_ice), @name(q_ice), @name(ᶜwᵢ)),
541-
(@name(c.ρq_rai), @name(q_rai), @name(ᶜwᵣ)),
542-
(@name(c.ρq_sno), @name(q_sno), @name(ᶜwₛ)),
543-
(@name(c.ρn_liq), @name(n_liq), @name(ᶜwnₗ)),
544-
(@name(c.ρn_rai), @name(n_rai), @name(ᶜwnᵣ)),
541+
(@name(c.ρq_liq), @name(ᶜwₗ)),
542+
(@name(c.ρq_ice), @name(ᶜwᵢ)),
543+
(@name(c.ρq_rai), @name(ᶜwᵣ)),
544+
(@name(c.ρq_sno), @name(ᶜwₛ)),
545+
(@name(c.ρn_liq), @name(ᶜwnₗ)),
546+
(@name(c.ρn_rai), @name(ᶜwnᵣ)),
545547
)
546548
if !(p.atmos.moisture_model isa DryModel) || use_derivative(diffusion_flag)
547549
∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)]
@@ -583,13 +585,13 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
583585
# DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅
584586
# DiagonalMatrixRow(
585587
# -1 / ᶜρ * ifelse(
586-
# ᶜspecific.q_tot == 0,
588+
# specific(Y.c.ρq_tot, Y.c.ρ) == 0,
587589
# (Geometry.WVector(FT(0)),),
588-
# p.precomputed.ᶜwₜqₜ / ᶜspecific.q_tot,
590+
# p.precomputed.ᶜwₜqₜ / specific(Y.c.ρq_tot, Y.c.ρ),
589591
# ),
590592
# ) - (I,)
591593

592-
MatrixFields.unrolled_foreach(tracer_info) do (ρχₚ_name, _, wₚ_name)
594+
MatrixFields.unrolled_foreach(tracer_info) do (ρχₚ_name, wₚ_name)
593595
MatrixFields.has_field(Y, ρχₚ_name) || return
594596
∂ᶜρχₚ_err_∂ᶜρχₚ = matrix[ρχₚ_name, ρχₚ_name]
595597
ᶜwₚ = MatrixFields.get_field(p.precomputed, wₚ_name)
@@ -651,13 +653,14 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
651653
dtγ * ᶜdiffusion_h_matrix DiagonalMatrixRow(1 / ᶜρ)
652654
end
653655

654-
MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name, _)
656+
MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, _)
655657
MatrixFields.has_field(Y, ρχ_name) || return
656-
ᶜχ = MatrixFields.get_field(ᶜspecific, χ_name)
658+
ᶜρχ = MatrixFields.get_field(Y, ρχ_name)
659+
ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
657660
∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)]
658661
∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name]
659662
ᶜtridiagonal_matrix_scalar = ifelse(
660-
χ_name in (@name(q_rai), @name(q_sno), @name(n_rai)),
663+
ρχ_name in (@name(c.ρq_rai), @name(c.ρq_sno), @name(c.ρn_rai)),
661664
ᶜdiffusion_h_matrix_scaled,
662665
ᶜdiffusion_h_matrix,
663666
)

0 commit comments

Comments
 (0)