reduce allocations in snow and neural snow (#980)

Author: kmdeck

ext/neural_snow/NeuralSnow.jl

@@ -9,7 +9,7 @@ import ClimaLand.Snow:
     density_prog_names,
     update_density!,
     update_density_prog!,
-    snow_depth
+    snow_depth!
 import ClimaLand.Parameters as LP
 using Thermodynamics
 
@@ -114,14 +114,34 @@ ClimaLand.Snow.density_prog_names(m::NeuralDepthModel) =
 #Extend/Define the appropriate functions needed for this parameterization:
 
 """
-    snow_depth(m::NeuralDepthModel, Y, p, params)
+    snow_depth!(z_snow, m::NeuralDepthModel{FT}, Y, p, params) where {FT}
 
-An extension of the `snow_depth` function to the NeuralDepthModel density parameterization, which includes the prognostic 
+An extension of the `snow_depth!` function to the NeuralDepthModel density parameterization, which includes the prognostic 
 depth variable and thus does not need to derive snow depth from SWE and density.
-This is sufficient to enable dynamics of the auxillary variable `ρ_snow` without extension of update_density!, and avoids
-redundant computations in the computation of runoff.
+
+This function clips the snow depth to be between 0 and SWE.
 """
-ClimaLand.Snow.snow_depth(m::NeuralDepthModel, Y, p, params) = Y.snow.Z
+function snow_depth!(z_snow, m::NeuralDepthModel{FT}, Y, p, params) where {FT}
+    z_snow .= min(Y.snow.Z, Y.snow.S) # z cannot be larger than SWE
+    z_snow .= max(z_snow, eps(FT)) # z must be positive
+    return nothing
+end
+
+"""
+    update_density!(ρ_snow, density::NeuralDepthModel, Y, p, params::SnowParameters,)
+
+Updates the snow density in place given the current model state. Default for all model types,
+can be extended for alternative density parameterizations.
+"""
+function update_density!(
+    ρ_snow,
+    density::NeuralDepthModel,
+    Y,
+    p,
+    params::SnowParameters,
+)
+    @. ρ_snow = snow_bulk_density(Y.snow.S, p.snow.z_snow, params)
+end
 
 
 """
@@ -144,13 +164,13 @@ function eval_nn(
 end
 
 """
-    dzdt(density::NeuralDepthModel, model::SnowModel{FT}, Y, p, t) where {FT}
-Returns the change in snow depth (rate) given the current model state and the `NeuralDepthModel`
-density paramterization, passing the approximate average of the forcings over the last 24 hours instead of
-the instantaneous value.
+    update_dzdt!(density::NeuralDepthModel, model::SnowModel, Y, p, t)
+
+Updates the dY.snow.Z field in places with the predicted change in snow depth (rate) given the model state `Y` and the `NeuralDepthModel`
+density paramterization.
 """
-function dzdt(density::NeuralDepthModel, Y)
-    return eval_nn.(
+function update_dzdt!(dzdt, density::NeuralDepthModel, Y)
+    dzdt .= eval_nn(
         Ref(density.z_model),
         Y.snow.Z,
         Y.snow.S, # When snow-cover-fraction variable is implemented, make sure this value changes to the right input
@@ -164,6 +184,7 @@ end
 
 """
     clip_dZdt(S::FT, Z::FT, dSdt::FT, dZdt::FT, Δt::FT)::FT
+
 A helper function which clips the tendency of Z such that
 its behavior is consistent with that of S: if all snow melts
 within a timestep, we clip the tendency of S so that it does
@@ -187,6 +208,7 @@ end
 
 """
     update_density_prog!(density::NeuralDepthModel, model::SnowModel, Y, p)
+
 Updates all prognostic variables associated with density/depth given the current model state and the `NeuralDepthModel`
 density paramterization.
 """
@@ -197,15 +219,16 @@ function update_density_prog!(
     Y,
     p,
 )
+    update_dzdt!(dY.snow.Z, density, Y)
 
-    dY.snow.Z .=
-        clip_dZdt.(
-            Y.snow.S,
-            Y.snow.Z,
-            dY.snow.S, #assumes dY.snow.S is updated (and clipped) before dY.snow.Z
-            dzdt(density, Y), # Note that the `dzdt` call below allocates a field. Return to this in the future.
-            model.parameters.Δt,
-        )
+    # Now we clip the tendency so that Z stays within approximately physical bounds.
+    @. dY.snow.Z = clip_dZdt(
+        Y.snow.S,
+        Y.snow.Z,
+        dY.snow.S, #assumes dY.snow.S is updated (and clipped) before dY.snow.Z
+        dY.snow.Z,
+        model.parameters.Δt,
+    )
 
     @. dY.snow.P_avg = density.α * (abs(p.drivers.P_snow) - Y.snow.P_avg)
     @. dY.snow.T_avg =

src/integrated/soil_snow_model.jl

@@ -252,8 +252,7 @@ function update_soil_snow_ground_heat_flux!(
     κ_soil = ClimaLand.Domains.top_center_to_surface(p.soil.κ)
 
     # Depth of snow and soil layers interacting thermally at interface
-    # Note that the `snow_depth` call below allocates a field. Return to this in the future.
-    Δz_snow = Snow.snow_depth(snow_params.density, Y, p, snow_params) # Snow depth
+    Δz_snow = p.snow.z_snow # Snow depth
     Δz_soil = p.effective_soil_sfc_depth
     (; ρc_ds, earth_param_set) = soil_params
 

src/standalone/Snow/Snow.jl

@@ -249,7 +249,8 @@ Returns the auxiliary variable names for the snow model. These
 include the specific humidity at the surface of the snow `(`q_sfc`, unitless),
 the mass fraction in liquid water (`q_l`, unitless),
 the thermal conductivity (`κ`, W/m/K),
-the bulk temperature (`T`, K), the surface temperature (`T_sfc`, K), the bulk snow density (`ρ_snow`, kg/m^3)
+the bulk temperature (`T`, K), the surface temperature (`T_sfc`, K), the snow depth (`z_snow`, m), 
+the bulk snow density (`ρ_snow`, kg/m^3)
 the SHF, LHF, and vapor flux (`turbulent_fluxes.shf`, etc),
 the net radiation (`R_n, J/m^2/s)`, the energy flux in liquid water runoff
 (`energy_runoff`, J/m^2/s), the water volume in runoff (`water_runoff`, m/s), and the total energy and water fluxes applied to the snowpack.
@@ -265,6 +266,7 @@ auxiliary_vars(::SnowModel) = (
     :κ,
     :T,
     :T_sfc,
+    :z_snow,
     :ρ_snow,
     :turbulent_fluxes,
     :R_n,
@@ -284,6 +286,7 @@ auxiliary_types(::SnowModel{FT}) where {FT} = (
     FT,
     FT,
     FT,
+    FT,
     NamedTuple{(:lhf, :shf, :vapor_flux, :r_ae), Tuple{FT, FT, FT, FT}},
     FT,
     FT,
@@ -311,6 +314,7 @@ auxiliary_domain_names(::SnowModel) = (
     :surface,
     :surface,
     :surface,
+    :surface,
 )
 
 
@@ -319,8 +323,8 @@ ClimaLand.name(::SnowModel) = :snow
 function ClimaLand.make_update_aux(model::SnowModel{FT}) where {FT}
     function update_aux!(p, Y, t)
         parameters = model.parameters
-
-        update_density!(parameters.density, parameters, Y, p)
+        snow_depth!(p.snow.z_snow, model.parameters.density, Y, p, parameters)
+        update_density!(p.snow.ρ_snow, parameters.density, Y, p, parameters)
 
         @. p.snow.κ = snow_thermal_conductivity(p.snow.ρ_snow, parameters)
 
@@ -338,16 +342,14 @@ function ClimaLand.make_update_aux(model::SnowModel{FT}) where {FT}
             p.drivers.thermal_state,
             parameters,
         )
-
-        p.snow.water_runoff .=
-            compute_water_runoff.(
-                Y.snow.S,
-                p.snow.q_l,
-                p.snow.T,
-                p.snow.ρ_snow,
-                snow_depth(model.parameters.density, Y, p, parameters), # Note that the `snow_depth` call below allocates a field. Return to this in the future.
-                parameters,
-            )
+        @. p.snow.water_runoff = compute_water_runoff(
+            Y.snow.S,
+            p.snow.q_l,
+            p.snow.T,
+            p.snow.ρ_snow,
+            p.snow.z_snow,
+            parameters,
+        )
 
         @. p.snow.energy_runoff =
             p.snow.water_runoff * volumetric_internal_energy_liq(FT, parameters)

src/standalone/Snow/snow_parameterizations.jl

@@ -1,5 +1,5 @@
 export snow_surface_temperature,
-    snow_depth,
+    snow_depth!,
     specific_heat_capacity,
     snow_thermal_conductivity,
     snow_bulk_temperature,
@@ -24,19 +24,6 @@ function snow_cover_fraction(x::FT; α = FT(1e-3))::FT where {FT}
     return heaviside(x - α)
 end
 
-"""
-    snow_depth(model::AbstractDensityModel{FT}, Y, p, params) where {FT}
-
-Returns the snow depth given SWE, snow density ρ_snow, and
-the density of liquid water ρ_l.
-This can be extended for additional types of parameterizations.
-"""
-function snow_depth(density::AbstractDensityModel{FT}, Y, p, params) where {FT}
-    ρ_l = FT(LP.ρ_cloud_liq(params.earth_param_set))
-    return @. ρ_l * Y.snow.S / p.snow.ρ_snow
-end
-
-
 """
     ClimaLand.surface_height(
         model::SnowModel{FT},
@@ -393,41 +380,43 @@ function energy_from_T_and_swe(S::FT, T::FT, parameters) where {FT}
 
 end
 
+
 """
-    update_density!(density::AbstractDensityModel, params::SnowParameters, Y, p)
-Updates the snow density given the current model state. Default for all model types,
-can be extended for alternative density paramterizations.
+    snow_depth!(model::ConstantDensityModel, Y, p, params)
+
+Returns the snow depth given SWE, snow density ρ_snow, and
+the density of liquid water ρ_l for a constant density model.
 """
-function update_density!(
-    density::AbstractDensityModel,
-    params::SnowParameters,
-    Y,
-    p,
-)
-    p.snow.ρ_snow .=
-        snow_bulk_density.(Y.snow.S, snow_depth(density, Y, p, params), params)
+function snow_depth!(z_snow, density::ConstantDensityModel, Y, p, params)
+    ρ_l = LP.ρ_cloud_liq(params.earth_param_set)
+    @. z_snow = ρ_l * Y.snow.S / density.ρ_snow
+    return nothing
 end
 
 """
-    update_density!(density::ConstantDensityModel, params::SnowParameters, Y, p)
-Extends the update_density! function for the ConstantDensityModel type.
+    update_density!(ρ_snow, density::ConstantDensityModel, Y, p, params::SnowParameters)
+
+Extends the update_density! function for the ConstantDensityModel type; updates the snow density in place.
 """
 function update_density!(
+    ρ_snow,
     density::ConstantDensityModel,
-    params::SnowParameters,
     Y,
     p,
+    params::SnowParameters,
 )
-    p.snow.ρ_snow .= density.ρ_snow
+    ρ_snow .= density.ρ_snow
 end
 
 """
     update_density_prog!(density::AbstractDensityModel{FT}, model::SnowModel{FT}, Y, p) where {FT}
+
 Updates all prognostic variables associated with density/depth given the current model state.
-This is the default method for all density model types, which can be extended for alternative paramterizations.
+This is the default method for the constant density model,
+ which has no prognostic variables.
 """
 function update_density_prog!(
-    density::AbstractDensityModel,
+    density::ConstantDensityModel,
     model::SnowModel,
     dY,
     Y,

test/standalone/Snow/snow.jl

@@ -58,6 +58,7 @@ import ClimaLand.Parameters as LP
         :κ,
         :T,
         :T_sfc,
+        :z_snow,
         :ρ_snow,
         :turbulent_fluxes,
         :R_n,
@@ -139,12 +140,17 @@ import ClimaLand.Parameters as LP
     @test turb_fluxes_copy.lhf == p.snow.turbulent_fluxes.lhf
     @test turb_fluxes_copy.vapor_flux == p.snow.turbulent_fluxes.vapor_flux
     old_ρ = deepcopy(p.snow.ρ_snow)
-    Snow.update_density!(model.parameters.density, model.parameters, Y, p)
+    Snow.update_density!(
+        p.snow.ρ_snow,
+        model.parameters.density,
+        Y,
+        p,
+        model.parameters,
+    )
     @test p.snow.ρ_snow == old_ρ
-    old_z = similar(Y.snow.S)
-    old_z .= FT(0.5)
-    z = snow_depth(model.parameters.density, Y, p, parameters)
-    @test z == old_z
+    old_z = deepcopy(p.snow.z_snow)
+    snow_depth!(p.snow.z_snow, model.parameters.density, Y, p, parameters)
+    @test p.snow.z_snow == old_z
 
     # Now compute tendencies and make sure they operate correctly.
     dY = similar(Y)