Merge NeuralSnow module into main #946

a-charbon · 2024-12-10T19:54:43Z

Purpose

To-do

Content

I have read and checked the items on the review checklist.

kmdeck

thanks for your hard work on this! You can merge once the doc comments are addressed. 🎉

kmdeck · 2024-12-13T17:45:24Z

src/standalone/Snow/Snow.jl

@@ -42,6 +42,23 @@ and is used as a bulk snow model.
 """
 abstract type AbstractSnowModel{FT} <: ClimaLand.AbstractExpModel{FT} end

+"""
+    AbstractDensityModel{FT}
+Defines the model type for density and depth parameterizations


Suggested change

Defines the model type for density and depth parameterizations

Defines the model type for density and depth parameterizations

for use within an `AbstractSnowModel` type. Current examples include the

`ConstantDensityModel` models.

Since depth and bulk density are related via SWE, and SWE is a prognostic variable

of the snow model, only depth or density can be independently modeled at one time.

This is why they are treated as a single "density model" even if the parameterization is actually

a model for snow depth.

The snow depth/density model can be diagnostic or introduce additional prognostic variables.

kmdeck · 2024-12-13T17:47:11Z

src/standalone/Snow/Snow.jl

-prognostic_vars(::SnowModel) = (:S, :U)
+prognostic_vars(m::SnowModel) =
+    (:S, :U, density_prog_vars(m.parameters.density)...)
+density_prog_vars(::AbstractDensityModel) = ()


Suggested change

density_prog_vars(::AbstractDensityModel) = ()

"""

density_prog_vars(::AbstractDensityModel)

A default method for adding prognostic variables to the snow model as required

by the density model choice.

"""

density_prog_vars(::AbstractDensityModel) = ()

kmdeck · 2024-12-13T17:49:30Z

src/standalone/Snow/Snow.jl

@@ -272,8 +299,15 @@ function ClimaLand.make_update_aux(model::SnowModel{FT}) where {FT}

        @. p.snow.T_sfc = snow_surface_temperature(p.snow.T)

-        @. p.snow.water_runoff =
-            compute_water_runoff(Y.snow.S, p.snow.q_l, p.snow.T, parameters)
+        p.snow.water_runoff .=


Suggested change

p.snow.water_runoff .=

# Note that the `snow_depth` call below allocates a field. Return to this in the future.

p.snow.water_runoff .=

kmdeck · 2024-12-13T17:49:59Z

src/integrated/soil_snow_model.jl

@@ -253,7 +253,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
-    Δz_snow = p.snow.z # Snow depth
+    Δz_snow = Snow.snow_depth(snow_params.density, Y, p, snow_params) # Snow depth


Suggested change

Δz_snow = Snow.snow_depth(snow_params.density, Y, p, snow_params) # Snow depth

# 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

kmdeck · 2024-12-13T17:50:58Z

ext/neural_snow/NeuralSnow.jl

+is calculated from a neural network determining the rate of change of snow depth, dzdt.
+"""
+struct NeuralDepthModel{FT} <: AbstractDensityModel{FT}
+    z_model::Flux.Chain


Suggested change

z_model::Flux.Chain

" The Flux neural network for compute dz/dt"

z_model::Flux.Chain

kmdeck · 2024-12-13T17:51:28Z

ext/neural_snow/NeuralSnow.jl

+"""
+struct NeuralDepthModel{FT} <: AbstractDensityModel{FT}
+    z_model::Flux.Chain
+    α::FT


Suggested change

α::FT

"The inverse of the averaging window time (1/s)"

α::FT

kmdeck · 2024-12-13T17:52:46Z

ext/neural_snow/NeuralSnow.jl

+
+"""
+    get_znetwork()
+Return the snow-depth neural network from the paper.


Suggested change

Return the snow-depth neural network from the paper.

Return the snow-depth neural network from Charbonneau et al (2024; https://arxiv.org/abs/2412.06819), and sets the timescale [@a-charbon explain what this timescale is]

Note that because we are loading a pre-trained model, the number of features input, the size of the output, etc, are hard coded in the function below.

kmdeck · 2024-12-13T17:57:09Z

ext/neural_snow/NeuralSnow.jl

+
+"""
+    NeuralDepthModel{FT <: AbstractFloat} <: AbstractDensityModel{FT}
+Establishes the density parameterization where snow density


Suggested change

Establishes the density parameterization where snow density

Establishes the density parameterization where snow density is calculated from a prognostic snow depth variable,

along with the prognostic SWE variable, using a neural network for the rate of change of snow depth, dz/dt.

The input to the network are temporally averaged, which we achieve using an exponentially moving average, with a rate of `\alpha`. Therefore, to achieve approximately daily averages, use `\alpha = 1/86400 s^{-1}`.

kmdeck · 2024-12-13T18:00:42Z

ext/neural_snow/NeuralSnow.jl

+
+"""
+    NeuralDepthModel(FT::DataType; Δt::Union{Nothing, FT}=nothing; model::Flux.Chain, α::Union{FT, Nothing})
+An outer constructor for the `NeuralDepthModel` density parameterization for usage in a snow model.


"An outer constructor for the NeuralDepthModel density parameterization for usage in a snow model.

By default, the neural network weights trained in Charbonneau et al 2015 are used, but one can also supply their own model by keyword argument. Since the neural network models are trained using temporally averaged input, we also must supply information about how to compute these averaged inputs [using an exponential-weighting scheme]. By default, an averaging window of 43200 seconds is used. "

Is Delta t really optional? I thought the lower constraint was baked into the model.

Can you explain the instability better?

oh, Delta t not used anymore maybe?

kmdeck · 2024-12-13T19:57:54Z

ext/neural_snow/NeuralSnow.jl

+    p,
+)
+
+    dY.snow.Z .=


Maybe we can comment that dzdt(density,Y) allocates a field here as a reminder

a-charbon requested a review from kmdeck December 10, 2024 19:54

a-charbon force-pushed the ac/snow_ema_PR branch from e6c9746 to d867a7a Compare December 10, 2024 20:01

kmdeck mentioned this pull request Dec 11, 2024

integrate neural network into Snow.jl architecture (Rebased) #941

Closed

1 task

a-charbon force-pushed the ac/snow_ema_PR branch 3 times, most recently from 4f62e77 to 6b9eb71 Compare December 12, 2024 17:47

Rebasing and integrating NeuralSnow module

adbc045

a-charbon force-pushed the ac/snow_ema_PR branch from 4be4bcc to adbc045 Compare December 13, 2024 01:36

kmdeck approved these changes Dec 13, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Merge NeuralSnow module into main #946

Merge NeuralSnow module into main #946

a-charbon commented Dec 10, 2024

kmdeck left a comment

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

kmdeck Dec 13, 2024

-Defines the model type for density and depth parameterizations
+Defines the model type for density and depth parameterizations
+for use within an `AbstractSnowModel` type. Current examples include the
+`ConstantDensityModel` models.
+Since depth and bulk density are related via SWE, and SWE is a prognostic variable
+of the snow model, only depth or density can be independently modeled at one time.
+This is why they are treated as a single "density model" even if the parameterization is actually
+a model for snow depth.
+The snow depth/density model can be diagnostic or introduce additional prognostic variables.

	p.snow.water_runoff .=
	# Note that the `snow_depth` call below allocates a field. Return to this in the future.
	p.snow.water_runoff .=

	z_model::Flux.Chain
	" The Flux neural network for compute dz/dt"
	z_model::Flux.Chain

	α::FT
	"The inverse of the averaging window time (1/s)"
	α::FT

-Return the snow-depth neural network from the paper.
+Return the snow-depth neural network from Charbonneau et al (2024; https://arxiv.org/abs/2412.06819), and sets the timescale [@a-charbon explain what this timescale is]
+Note that because we are loading a pre-trained model, the number of features input, the size of the output, etc, are hard coded in the function below.

-Establishes the density parameterization where snow density
+Establishes the density parameterization where snow density is calculated from a prognostic snow depth variable,
+along with the prognostic SWE variable, using a neural network for the rate of change of snow depth, dz/dt.
+The input to the network are temporally averaged, which we achieve using an exponentially moving average, with a rate of `\alpha`. Therefore, to achieve approximately daily averages, use `\alpha = 1/86400 s^{-1}`.

Merge NeuralSnow module into main #946

Are you sure you want to change the base?

Merge NeuralSnow module into main #946

Conversation

a-charbon commented Dec 10, 2024

Purpose

To-do

Content

kmdeck left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment