small typos/nitpick fixes and coderabbit suggestions implemented

Author: a-charbon

experiments/standalone/Snow/snowmip_simulation.jl

@@ -22,11 +22,14 @@ using DelimitedFiles
 
 # Site-specific quantities
 # Error if no site argument is provided
+#=
 if length(ARGS) < 1
     @error("Please provide a site name as command line argument")
 else
     SITE_NAME = ARGS[1]
 end
+=#
+SITE_NAME = "cdp"
 
 climaland_dir = pkgdir(ClimaLand)
 

ext/neural_snow/ModelTools.jl

@@ -13,7 +13,7 @@ export make_model,
     save_predictive_model_weights,
     load_model_weights!,
     freeze_fixed_layers!,
-    convert_model!
+    convert_model
 
 include("./WebTools.jl")
 
@@ -165,19 +165,19 @@ paper_lower_bound(
 
 
 """ 
-    convert_model!(model, FT)
+    convert_model(model, FT)
 Provides a utility funciton to set a model's float type by passing
 the type as an argument.
 """
-function convert_model!(model, FT::Type{<:AbstractFloat})
+function convert_model(model, FT::Type{<:AbstractFloat})
     if FT == Float32
         model = Flux.f32(model)
     elseif FT == Float64
         model = Flux.f64(model)
     else
         error(
             "Conversion of Flux Chain weights to the desired float-type is not supported. Please implement the desired conversion
-      in ModelTools.convert_model!() or a custom constructor for the desired type.",
+      in ModelTools.convert_model() or a custom constructor for the desired type.",
         )
     end
     return model
@@ -258,7 +258,7 @@ function make_model(
             apply_upper = Dense(get_min, false, relu),
             sumlayer = MulLayer(final_mul),
         )
-        return convert_model!(model, FT)
+        return convert_model(model, FT)
     end
 end
 
@@ -357,7 +357,7 @@ function make_model_paper(;
             apply_upper = Dense(get_min, false, relu),
             sumlayer = MulLayer(final_mul),
         )
-        return convert_model!(model, FT)
+        return convert_model(model, FT)
     end
 end
 
@@ -378,10 +378,10 @@ of the predictive component by the constant `scale`.
 # Arguments
 - `model::Chain`: the neural model to be scaled.
 - `scale::Real`: the scaling parameter to return data to applicable units.
-- `dtype::Type`: Sets type, consistent with neural model. Default is `Float32`.
 """
-function setoutscale!(model::Chain, scale::Real; dtype::Type = Float32)
-    model[:apply_relus].weight[end, 3] = dtype(-1 * scale)
+function setoutscale!(model::Chain, scale::Real)
+    FT = eltype(model[:apply_relus].weight)
+    model[:apply_relus].weight[end, 3] = FT(-1 * scale)
 end
 
 
@@ -408,21 +408,16 @@ directly by the constant `scale`.
 - `model::Chain`: the neural model to be scaled.
 - `bound`::Symbol: specifies the upper or lower bound with :upper or :lower.
 - `scale::Real`: the scaling parameter to return data to applicable units.
-- `dtype::Type`: Sets type, consistent with neural model. Default is `Float32`.
 """
-function setboundscale!(
-    model::Chain,
-    bound::Symbol,
-    scale::Real;
-    dtype::Type = Float32,
-)
+function setboundscale!(model::Chain, bound::Symbol, scale::Real;)
+    FT = eltype(model[:apply_relus].weight)
     if bound == :upper
-        model[:apply_relus].weight[1, 1] = dtype(1.0 * scale)
-        model[:apply_relus].weight[2, 1] = dtype(-1.0 * scale)
-        model[:apply_relus].weight[5, 1] = dtype(-1.0 * scale)
+        model[:apply_relus].weight[1, 1] = FT(1.0 * scale)
+        model[:apply_relus].weight[2, 1] = FT(-1.0 * scale)
+        model[:apply_relus].weight[5, 1] = FT(-1.0 * scale)
     elseif bound == :lower
-        model[:apply_relus].weight[3, 2] = dtype(1.0 * scale)
-        model[:apply_relus].weight[4, 2] = dtype(-1.0 * scale)
+        model[:apply_relus].weight[3, 2] = FT(1.0 * scale)
+        model[:apply_relus].weight[4, 2] = FT(-1.0 * scale)
     else
         error("provided `bound` must be `:upper` or `:lower`.")
     end
@@ -439,10 +434,10 @@ a custom form of `setboundscale!()` to use with the custom model type made with
 # Arguments
 - `model::Chain`: the neural model to be used.
 - `dt::Real`: the number of seconds per timestep for usage.
-- `dtype::Type`: Sets type, consistent with neural model. Default is `Float32`.
 """
-function settimescale!(model::Chain, dt::Real; dtype::Type = Float32)
-    model[:apply_relus].weight[2, 2] = dtype(1.0 / dt)
+function settimescale!(model::Chain, dt::Real)
+    FT = eltype(model[:apply_relus].weight)
+    model[:apply_relus].weight[2, 2] = FT(1.0 / dt)
 end
 
 """
@@ -840,8 +835,12 @@ and in a way that is stable to changing versions of the Flux package.
 # Arguments
 - `path`::String: The filepath (or URL) from where to get the text file.
 - `model`: the model for which to load in the predictive parameters.
+- `FT` : Optional arg for passing the right Float type if not inferred from passed model.
 """
 function load_model_weights!(filepath::String, model::Chain; FT = Float32)
+    T =
+        hasproperty(model.layers, :apply_relus) ?
+        eltype(model[:apply_relus].weight) : FT
     data = Dict()  # Dictionary to store the parsed data
     # Check if the filepath is a URL
     if startswith(filepath, "http://") || startswith(filepath, "https://")
@@ -869,7 +868,7 @@ function load_model_weights!(filepath::String, model::Chain; FT = Float32)
             elseif isnothing(current_name)
                 current_name = line
             else
-                push!(current_values, parse.(FT, split(line)))
+                push!(current_values, parse.(T, split(line)))
             end
         end
         # Handle the last item (if file doesn't end with a blank line)

ext/neural_snow/WebTools.jl

@@ -1,4 +1,4 @@
-using Downloads, DataFrames
+using Downloads, DataFrames, Dates
 export read_webpage_body, df_from_url
 
 const extension_dir_path = @__DIR__()

src/standalone/Snow/AndersonModel.jl

@@ -23,18 +23,15 @@ Base.@kwdef struct Anderson1976{FT} <: AbstractDensityModel{FT}
     c3::FT = FT(0.005)
     c4::FT = FT(0.10)
     c5::FT = FT(0)
-    cx::FT = FT(0.15)
-    ρ_d::FT = FT(23.0)
+    cx::FT = FT(23)
+    ρ_d::FT = FT(0.15)
 end
 
 
 #Define the additional prognostic variables needed for using this parameterization:
-density_prog_vars(m::Anderson1976) = 
-    (:Z,)
-density_prog_types(m::Anderson1976{FT}) where {FT} =
-    (FT,)
-density_prog_names(m::Anderson1976) =
-    (:surface,)
+density_prog_vars(m::Anderson1976) = (:Z,)
+density_prog_types(m::Anderson1976{FT}) where {FT} = (FT,)
+density_prog_names(m::Anderson1976) = (:surface,)
 
 
 """
@@ -116,7 +113,12 @@ Returns the change in snow depth (rate) given the current model state and the `A
 density paramterization, which estimates contributions from new snowfall, as well as compaction
 effects from the existing and newly fallen snow.
 """
-function get_dzdt(density::Anderson1976{FT}, model::SnowModel{FT}, Y, p) where {FT}
+function get_dzdt(
+    density::Anderson1976{FT},
+    model::SnowModel{FT},
+    Y,
+    p,
+) where {FT}
     Δt_t = model.parameters.Δt / 3600 #in hours (when Δt in SnowParameters is ::FT instead of ::Period)
 
     #Contribution from new snowfall: (uses parameterization of newsnow_temp, newsnow_density)
@@ -136,7 +138,8 @@ function get_dzdt(density::Anderson1976{FT}, model::SnowModel{FT}, Y, p) where {
     Δz = snowfall ./ ρ_newsnow
 
     #Estimate resulting change from compaction effects of old snow (and the incoming snow on top):
-    S_snowheight = swe_snow_area.(Y.snow.S, p.snow.snow_cover_fraction, Y.snow.Z)
+    S_snowheight =
+        swe_snow_area.(Y.snow.S, p.snow.snow_cover_fraction, Y.snow.Z)
     W_ice = (1 .- p.snow.q_l) .* S_snowheight #S_snowheight instead of Y.snow.S
     ρ_ice = W_ice ./ Y.snow.Z #unitless
     sliq = p.snow.q_l .* S_snowheight
@@ -150,6 +153,10 @@ function get_dzdt(density::Anderson1976{FT}, model::SnowModel{FT}, Y, p) where {
 end
 
 
+# Do we want to change the type assertion to ::Union{Anderson1976, NeuralDepthModel} in NeuralDepthModel
+# to avoid redundancy since the two are identical and use the same functionality, or is better to keep
+# them separate and visible with their own code file? I can always just comment this one out like I did
+# above/below with clip_dZdt() and swe_snow_area. 
 """
     update_density_and_depth!(ρ_snow, z_snow,density::Anderson1976, Y, p, params::SnowParameters,)
 
@@ -210,19 +217,8 @@ end
 Updates all prognostic variables associated with density/depth given the current model state and the `Anderson1976`
 density paramterization.
 """
-function update_density_prog!(
-    density::Anderson1976,
-    model::SnowModel,
-    dY,
-    Y,
-    p,
-)
-    dY.snow.Z .= get_dzdt(
-        density,
-        model,
-        Y,
-        p
-    )
+function update_density_prog!(density::Anderson1976, model::SnowModel, dY, Y, p)
+    dY.snow.Z .= get_dzdt(density, model, Y, p)
 
     # Now we clip the tendency so that Z stays within approximately physical bounds.
     @. dY.snow.Z = clip_dZdt(

src/standalone/Snow/NeuralDepthModel.jl

@@ -8,7 +8,6 @@ import ClimaLand.Snow:
     density_prog_names,
     update_density_and_depth!,
     update_density_prog!
-import ClimaLand.Parameters as LP
 using Thermodynamics
 
 export NeuralDepthModel
@@ -33,7 +32,7 @@ end
 Retrieves the network weights for the snow depth model formulated in Charbonneau et. al., 2025, 
 and stores/returns them within a dictionary to enable easy loading of the weights into the generated
 neural model.
-Note that because we are loading from a document of pre-trained weights,
+Note that because we are loading from an unchanging document of pre-trained weights,
 the weight sizes are hard-coded in the function below.
 
 # Arguments
@@ -42,7 +41,7 @@ the weight sizes are hard-coded in the function below.
 function get_network_weights(FT::DataType)
     weightfile =
         readlines(ClimaLand.Artifacts.neural_depth_model_weights_path())
-    @assert length(weightfile) == 56 "Might not be reading the right weight file, expected 56 lines and got $(length(weightfile))"
+    @assert length(weightfile) == 56 "Artifact path might need updating, expected 56 lines and got $(length(weightfile))"
     data = Dict()
     data["SCALING"] = parse.(FT, split(weightfile[2], " "))
     data["l1_WEIGHT"] =
@@ -129,7 +128,7 @@ end
 
 
 """
-    setoutscale!(model, scale; dtype)
+    setoutscale!(model, scale)
 Set the physical scaling parameter for model usage (i.e. rectifying scaling done on model input).
 This scaling constant is combined with a -1 factor that enforces the boundary for speed and memory purposes -
 for this reason, we do not recommend first `extracting`/storing the weight this function refers to,
@@ -140,24 +139,24 @@ of the predictive component by the constant `scale`.
 # Arguments
 - `model::Chain`: the neural model to be used.
 - `scale::Real`: the scaling parameter to return data to applicable units.
-- `dtype::Type`: Sets type, consistent with neural model. Default is `Float32`.
 """
-function setoutscale!(model, scale::Real; dtype::Type = Float32)
-    model[:apply_relus].weight[end, 3] = dtype(-1 * scale)
+function setoutscale!(model, scale::Real)
+    FT = eltype(model[:apply_relus].weight)
+    model[:apply_relus].weight[end, 3] = FT(-1 * scale)
 end
 
 
 """
-    settimescale!(model, dt; dtype)
+    settimescale!(model, dt)
 Set the timescale parameter for model usage
 
 # Arguments
 - `model::Chain`: the neural model to be used.
 - `dt::Real`: the number of seconds per timestep for usage.
-- `dtype::Type`: Sets type, consistent with neural model. Default is `Float32`.
 """
-function settimescale!(model, dt::Real; dtype::Type = Float32)
-    model[:apply_relus].weight[2, 2] = dtype(1.0 / dt)
+function settimescale!(model, dt::Real)
+    FT = eltype(model[:apply_relus].weight)
+    model[:apply_relus].weight[2, 2] = FT(1.0 / dt)
 end
 
 """
@@ -179,8 +178,8 @@ function get_znetwork(; FT = Float32)
     network[:pred].layers.layers[:l2].bias .= data["l2_BIAS"]
     network[:pred].layers.layers[:l3].weight .= data["l3_WEIGHT"]
     network[:pred].layers.layers[:l3].bias .= data["l3_BIAS"]
-    setoutscale!(network, data["FINALSCALE"], dtype = FT)
-    settimescale!(network, 86400, dtype = FT)
+    setoutscale!(network, data["FINALSCALE"])
+    settimescale!(network, 86400)
     return network
 end
 
@@ -204,7 +203,7 @@ function NeuralDepthModel(
     usemodel = get_znetwork(FT = FT)
     weight = !isnothing(α) ? FT(α) : FT(2 / 86400)
     if !isnothing(Δt)
-        settimescale!(usemodel, Δt, dtype = FT) #assumes Δt is provided in seconds
+        settimescale!(usemodel, Δt) #assumes Δt is provided in seconds
         if (Δt > 43200) && isnothing(α)
             error("Please supply a weight for the
                    exponential moving average, the
@@ -250,7 +249,7 @@ as the depth over the snow-covered portion and S is provided as the depth over t
 Used in computing the time derivative of snow depth.
 """
 function swe_snow_area(S::FT, scf::FT, z::FT)::FT where {FT}
-    min_scf = 0.05
+    min_scf = 0.05 #placeholder until Katherine and Andy discuss how to implement this
     return min(z, S / max(scf, min_scf)) #double check if this is the formula Katherine wants? see clip_dzdt() too.
 end
 
@@ -316,12 +315,12 @@ function clip_dZdt(
     scf::FT,
     Δt::FT,
 )::FT where {FT}
-    min_scf = 0.05
+    min_scf = 0.05 #placeholder until Katherine and Andy discuss this
     #Case if S is set to zero:
     if (S + dSdt * Δt) <= eps(FT)
         return -Z / Δt
         #Case if Z would have been set to Z < S:
-    elseif (Z + dZdt * Δt) * max(scf, min_scf) < (S + dSdt * Δt)
+    elseif (Z + dZdt * Δt) * max(scf, min_scf) < (S + dSdt * Δt) #technically, the scf used here should be the scf of the following timestep, so how to handle this?
         #more stable form for very small Z, S
         return ((dSdt * Δt + S) / max(scf, min_scf) - Z) / Δt
     else

test/standalone/Snow/tool_tests.jl

@@ -258,7 +258,7 @@ if !isnothing(DataToolsExt)
         end
 
         @test eltype(model[:pred].layers[:l1].weight) == Float32
-        model64 = ModelTools.convert_model!(deepcopy(model), Float64)
+        model64 = ModelTools.convert_model(deepcopy(model), Float64)
         @test eltype(model64[:pred].layers[:l1].weight) == Float64
         test_input = Matrix{Float32}(ones(nfeatures, 8))
         test_output = model(test_input)