remove more theta and fix

CliMA · Nov 13, 2023 · 34f6617 · 34f6617
1 parent a34f583
commit 34f6617
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 46 deletions.
diff --git a/post_processing/remap/remap_helpers.jl b/post_processing/remap/remap_helpers.jl
@@ -110,14 +110,7 @@ function remap2latlon(filein, data_dir, remap_tmpdir, weightfile, nlat, nlon)
     nc_time = def_time_coord(nc)
     # define variables for the prognostic states
     nc_rho = defVar(nc, "rho", FT, cspace, ("time",))
-    thermo_var = if :ρe_tot in propertynames(Y.c)
-        "e_tot"
-    elseif :ρθ in propertynames(Y.c)
-        "theta"
-    else
-        error("Unfound thermodynamic variable")
-    end
-    nc_thermo = defVar(nc, thermo_var, FT, cspace, ("time",))
+    nc_thermo = defVar(nc, "e_tot", FT, cspace, ("time",))
     nc_u = defVar(nc, "u", FT, cspace, ("time",))
     nc_v = defVar(nc, "v", FT, cspace, ("time",))
     nc_w = defVar(nc, "w", FT, cspace, ("time",))
@@ -199,11 +192,7 @@ function remap2latlon(filein, data_dir, remap_tmpdir, weightfile, nlat, nlon)
     # density
     nc_rho[:, 1] = Y.c.ρ
     # thermodynamics
-    if :ρe_tot in propertynames(Y.c)
-        nc_thermo[:, 1] = Y.c.ρe_tot ./ Y.c.ρ
-    elseif :ρθ in propertynames(Y.c)
-        nc_thermo[:, 1] = Y.c.ρθ ./ Y.c.ρ
-    end
+    nc_thermo[:, 1] = Y.c.ρe_tot ./ Y.c.ρ
     # physical horizontal velocity
     uh_phy = Geometry.transform.(tuple(Geometry.UVAxis()), Y.c.uₕ)
     nc_u[:, 1] = uh_phy.components.data.:1
@@ -279,7 +268,7 @@ function remap2latlon(filein, data_dir, remap_tmpdir, weightfile, nlat, nlon)
         joinpath(out_dir, first(splitext(basename(filein))) * ".nc")
     dry_variables = [
         "rho",
-        thermo_var,
+        "e_tot",
         "u",
         "v",
         "w",

diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -33,10 +33,10 @@ TODO: Rename `ᶜK` to `ᶜκ`.
 """
 function precomputed_quantities(Y, atmos)
     FT = eltype(Y)
-    @assert !(atmos.moisture_model isa DryModel)
-        || !(atmos.turbconv_model isa DiagnosticEDMFX)
-    @assert !(atmos.moisture_model isa DryModel)
-        || !(atmos.turbconv_model isa PrognosticEDMFX)
+    @assert !(atmos.moisture_model isa DryModel) ||
+            !(atmos.turbconv_model isa DiagnosticEDMFX)
+    @assert !(atmos.moisture_model isa DryModel) ||
+            !(atmos.turbconv_model isa PrognosticEDMFX)
     @assert !(atmos.edmfx_detr_model isa ConstantAreaDetrainment) ||
             !(atmos.turbconv_model isa DiagnosticEDMFX)
     TST = thermo_state_type(atmos.moisture_model, FT)
@@ -255,19 +255,17 @@ function thermo_vars(moisture_model, specific, K, Φ)
     return (; energy_var..., moisture_var...)
 end
 
-ts_gs(thermo_params, moisture_model, specific, K, Φ, ρ) =
-    thermo_state(
-        thermo_params;
-        thermo_vars(moisture_model, specific, K, Φ)...,
-        ρ,
-    )
+ts_gs(thermo_params, moisture_model, specific, K, Φ, ρ) = thermo_state(
+    thermo_params;
+    thermo_vars(moisture_model, specific, K, Φ)...,
+    ρ,
+)
 
-ts_sgs(thermo_params, moisture_model, specific, K, Φ, p) =
-    thermo_state(
-        thermo_params;
-        thermo_vars(moisture_model, specific, K, Φ)...,
-        p,
-    )
+ts_sgs(thermo_params, moisture_model, specific, K, Φ, p) = thermo_state(
+    thermo_params;
+    thermo_vars(moisture_model, specific, K, Φ)...,
+    p,
+)
 
 """
     set_precomputed_quantities!(Y, p, t)
@@ -322,8 +320,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
 
     (; ᶜh_tot) = p.precomputed
-    @. ᶜh_tot =
-        TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot)
+    @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot)
 
     if !isnothing(p.sfc_setup)
         SurfaceConditions.update_surface_conditions!(Y, p, t)

diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
@@ -328,8 +328,7 @@ NVTX.@annotate function compute_diagnostics(integrator)
 
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(Y.f), Fields.half)
-    surface_ct3_unit =
-        CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
     (; ρ_flux_uₕ, ρ_flux_h_tot) = p.precomputed.sfc_conditions
     sfc_flux_momentum =
         Geometry.UVVector.(

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
@@ -443,10 +443,8 @@ function compute_tau!(out, state, cache, component)
     return
 end
 
-compute_tauu!(out, state, cache, time) =
-    compute_tau!(out, state, cache, :1)
-compute_tauv!(out, state, cache, time) =
-    compute_tau!(out, state, cache, :2)
+compute_tauu!(out, state, cache, time) = compute_tau!(out, state, cache, :1)
+compute_tauv!(out, state, cache, time) = compute_tau!(out, state, cache, :2)
 
 add_diagnostic_variable!(
     short_name = "tauu",

diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl
@@ -461,17 +461,14 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
             ) ⋅ ᶠinterp_matrix()
         )
     @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] =
-        dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅
-        ᶠgradᵥ_matrix() * R_d / cv_d
+        dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() *
+        R_d / cv_d
     if p.atmos.rayleigh_sponge isa RayleighSponge
         @. ∂ᶠu₃_err_∂ᶠu₃[colidx] =
             dtγ * (
-                DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅
-                ᶠgradᵥ_matrix() ⋅
-                DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅
-                ∂ᶜK_∂ᶠu₃[colidx] + DiagonalMatrixRow(
-                    -p.ᶠβ_rayleigh_w[colidx] * (one_C3xACT3,),
-                )
+                DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅
+                DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ ∂ᶜK_∂ᶠu₃[colidx] +
+                DiagonalMatrixRow(-p.ᶠβ_rayleigh_w[colidx] * (one_C3xACT3,))
             ) - (I_u₃,)
     else
         @. ∂ᶠu₃_err_∂ᶠu₃[colidx] =