integrate neural network into Snow.jl architecture (Rebased)

.buildkite/Manifest-v1.11.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.11.1"
 manifest_format = "2.0"
-project_hash = "bf1216ed2bc76713a01437d07332219f6bcd0c54"
+project_hash = "14f223707cbd1a2cb15b66a90e7be9a546a0f205"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "30bb95a372787af850addf28ac937f1be7b79173"
@@ -286,6 +286,12 @@ git-tree-sha1 = "e329286945d0cfc04456972ea732551869af1cfc"
 uuid = "4e9b3aee-d8a1-5a3d-ad8b-7d824db253f0"
 version = "1.0.1+0"
 
+[[deps.CSV]]
+deps = ["CodecZlib", "Dates", "FilePathsBase", "InlineStrings", "Mmap", "Parsers", "PooledArrays", "PrecompileTools", "SentinelArrays", "Tables", "Unicode", "WeakRefStrings", "WorkerUtilities"]
+git-tree-sha1 = "deddd8725e5e1cc49ee205a1964256043720a6c3"
+uuid = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+version = "0.10.15"
+
 [[deps.CUDA]]
 deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"]
 git-tree-sha1 = "cdbdca28f19c2c7fcf34ffb48bfdaca404dcd18a"
@@ -405,20 +411,11 @@ deps = ["ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "Clim
 path = ".."
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 version = "0.15.5"
+weakdeps = ["BSON", "CSV", "CUDA", "ClimaParams", "DataFrames", "Flux", "HTTP", "StatsBase", "cuDNN"]
 
     [deps.ClimaLand.extensions]
     CreateParametersExt = "ClimaParams"
-    NeuralSnowExt = ["CSV", "DataFrames", "HTTP", "Flux", "StatsBase", "cuDNN"]
-
-    [deps.ClimaLand.weakdeps]
-    CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
-    CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-    ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
-    DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-    Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-    HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
-    StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
-    cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
+    NeuralSnowExt = ["CSV", "DataFrames", "HTTP", "Flux", "StatsBase", "cuDNN", "BSON"]
 
 [[deps.ClimaParams]]
 deps = ["TOML"]
@@ -3005,6 +3002,12 @@ git-tree-sha1 = "93f43ab61b16ddfb2fd3bb13b3ce241cafb0e6c9"
 uuid = "2381bf8a-dfd0-557d-9999-79630e7b1b91"
 version = "1.31.0+0"
 
+[[deps.WeakRefStrings]]
+deps = ["DataAPI", "InlineStrings", "Parsers"]
+git-tree-sha1 = "b1be2855ed9ed8eac54e5caff2afcdb442d52c23"
+uuid = "ea10d353-3f73-51f8-a26c-33c1cb351aa5"
+version = "1.4.2"
+
 [[deps.WebP]]
 deps = ["CEnum", "ColorTypes", "FileIO", "FixedPointNumbers", "ImageCore", "libwebp_jll"]
 git-tree-sha1 = "aa1ca3c47f119fbdae8770c29820e5e6119b83f2"
@@ -3017,6 +3020,11 @@ git-tree-sha1 = "c1a7aa6219628fcd757dede0ca95e245c5cd9511"
 uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6"
 version = "1.0.0"
 
+[[deps.WorkerUtilities]]
+git-tree-sha1 = "cd1659ba0d57b71a464a29e64dbc67cfe83d54e7"
+uuid = "76eceee3-57b5-4d4a-8e66-0e911cebbf60"
+version = "1.6.1"
+
 [[deps.XML2_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"]
 git-tree-sha1 = "a2fccc6559132927d4c5dc183e3e01048c6dcbd6"

.buildkite/Project.toml

@@ -2,6 +2,7 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArtifactWrappers = "a14bc488-3040-4b00-9dc1-f6467924858a"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ClimaAnalysis = "29b5916a-a76c-4e73-9657-3c8fd22e65e6"

Project.toml

@@ -22,6 +22,7 @@ SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 
 [weakdeps]
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
@@ -33,10 +34,11 @@ cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [extensions]
 CreateParametersExt = "ClimaParams"
-NeuralSnowExt = ["CSV", "DataFrames", "HTTP", "Flux", "StatsBase", "cuDNN"]
+NeuralSnowExt = ["CSV", "DataFrames", "HTTP", "Flux", "StatsBase", "cuDNN", "BSON"]
 
 [compat]
 ArtifactWrappers = "0.2"
+BSON = "0.3.9"
 CSV = "0.10.14"
 CUDA = "5.5"
 ClimaComms = "0.6"

docs/Manifest-v1.11.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.11.1"
 manifest_format = "2.0"
-project_hash = "15fe7e962a3bea15852e19f92138be4db325cca0"
+project_hash = "8d95e673f3bf47d00566191d57db2b0549068c76"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "30bb95a372787af850addf28ac937f1be7b79173"
@@ -26,6 +26,16 @@ git-tree-sha1 = "574baf8110975760d391c710b6341da1afa48d8c"
 uuid = "a4c015fc-c6ff-483c-b24f-f7ea428134e9"
 version = "0.0.1"
 
+[[deps.About]]
+deps = ["InteractiveUtils", "JuliaSyntaxHighlighting", "PrecompileTools", "StyledStrings"]
+git-tree-sha1 = "efbf5b623b7ee2a41ce5aed6299aa62ab7f2d5b9"
+uuid = "69d22d85-9f48-4c46-bbbe-7ad8341ff72a"
+version = "1.0.1"
+weakdeps = ["Pkg"]
+
+    [deps.About.extensions]
+    PkgExt = "Pkg"
+
 [[deps.AbstractFFTs]]
 deps = ["LinearAlgebra"]
 git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef"
@@ -197,6 +207,11 @@ git-tree-sha1 = "2c7cc21e8678eff479978a0a2ef5ce2f51b63dff"
 uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 version = "0.5.0"
 
+[[deps.BSON]]
+git-tree-sha1 = "4c3e506685c527ac6a54ccc0c8c76fd6f91b42fb"
+uuid = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+version = "0.3.9"
+
 [[deps.BandedMatrices]]
 deps = ["ArrayLayouts", "FillArrays", "LinearAlgebra", "PrecompileTools"]
 git-tree-sha1 = "a2c85f53ddcb15b4099da59867868bd40f005579"
@@ -421,11 +436,11 @@ deps = ["ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "Clim
 path = ".."
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 version = "0.15.5"
-weakdeps = ["CSV", "CUDA", "ClimaParams", "DataFrames", "Flux", "HTTP", "StatsBase", "cuDNN"]
+weakdeps = ["BSON", "CSV", "CUDA", "ClimaParams", "DataFrames", "Flux", "HTTP", "StatsBase", "cuDNN"]
 
     [deps.ClimaLand.extensions]
     CreateParametersExt = "ClimaParams"
-    NeuralSnowExt = ["CSV", "DataFrames", "HTTP", "Flux", "StatsBase", "cuDNN"]
+    NeuralSnowExt = ["CSV", "DataFrames", "HTTP", "Flux", "StatsBase", "cuDNN", "BSON"]
 
 [[deps.ClimaLandSimulations]]
 deps = ["ArtifactWrappers", "Bonito", "CairoMakie", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaLand", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "DataFrames", "Dates", "DelimitedFiles", "Formatting", "HTTP", "Insolation", "Interpolations", "InverseFunctions", "JSON", "LaTeXStrings", "MutableArithmetics", "NLsolve", "PlotUtils", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "Unitful", "UnitfulMoles", "WGLMakie", "cuDNN"]
@@ -1459,6 +1474,17 @@ git-tree-sha1 = "af433a10f3942e882d3c671aacb203e006a5808f"
 uuid = "9c1d0b0a-7046-5b2e-a33f-ea22f176ac7e"
 version = "0.2.1+0"
 
+[[deps.JuliaSyntax]]
+git-tree-sha1 = "937da4713526b96ac9a178e2035019d3b78ead4a"
+uuid = "70703baa-626e-46a2-a12c-08ffd08c73b4"
+version = "0.4.10"
+
+[[deps.JuliaSyntaxHighlighting]]
+deps = ["JuliaSyntax", "StyledStrings"]
+git-tree-sha1 = "19ecee1ea81c60156486a92b062e443b6bba60b7"
+uuid = "ac6e5ff7-fb65-4e79-a425-ec3bc9c03011"
+version = "0.1.0"
+
 [[deps.JuliaVariables]]
 deps = ["MLStyle", "NameResolution"]
 git-tree-sha1 = "49fb3cb53362ddadb4415e9b73926d6b40709e70"

docs/Project.toml

@@ -1,5 +1,7 @@
 [deps]
+About = "69d22d85-9f48-4c46-bbbe-7ad8341ff72a"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"

docs/tutorials/standalone/Snow/base_tutorial.jl

@@ -43,10 +43,10 @@
 # We begin by importing the developed code to create and run the neural network,
 # as well as some preliminary packages:
 using ClimaLand
-using DataFrames, CSV, HTTP, Dates, Flux, StatsBase, cuDNN
+using DataFrames, CSV, HTTP, Dates, Flux, StatsBase, cuDNN, BSON
 
 # The code lives in an extenson that we have to manually load. The extension can
-# be loaded only if "CSV", "HTTP", "Flux", "StatsBase", "cuDNN" and "ClimaLand"
+# be loaded only if "DataFrames", "CSV", "HTTP", "Flux", "StatsBase", "cuDNN", "BSON", and "ClimaLand"
 # are loaded.
 DataTools = Base.get_extension(ClimaLand, :NeuralSnowExt).DataTools
 ModelTools = Base.get_extension(ClimaLand, :NeuralSnowExt).ModelTools;

docs/tutorials/standalone/Snow/data_tutorial.jl

@@ -11,10 +11,10 @@
 
 # We begin by importing all required packages:
 using ClimaLand
-using DataFrames, CSV, HTTP, Dates, Flux, StatsBase, cuDNN
+using DataFrames, CSV, HTTP, Dates, Flux, StatsBase, cuDNN, BSON
 
 # The code lives in an extenson that we have to manually load. The extension can
-# be loaded only if "CSV", "HTTP", "Flux", "StatsBase", "cuDNN" and "ClimaLand"
+# be loaded only if "DataFrames", "CSV", "HTTP", "Flux", "StatsBase", "cuDNN", "BSON", and "ClimaLand"
 # are loaded.
 DataTools = Base.get_extension(ClimaLand, :NeuralSnowExt).DataTools;
 

experiments/integrated/fluxnet/snow_soil/simulation.jl

@@ -88,7 +88,7 @@ soil_model_type = Soil.EnergyHydrology
 snow_parameters = SnowParameters{FT}(
     dt;
     α_snow = α,
-    ρ_snow = ρ,
+    density = Snow.ConstantDensityModel(ρ),
     earth_param_set = earth_param_set,
 );
 snow_args = (; parameters = snow_parameters);

experiments/standalone/Snow/snowmip_simulation.jl

@@ -20,6 +20,10 @@ using Statistics
 using Insolation
 using DelimitedFiles
 
+using CSV, HTTP, Flux, cuDNN, BSON, StatsBase
+ModelTools = Base.get_extension(ClimaLand, :NeuralSnowExt).ModelTools;
+NeuralSnow = Base.get_extension(ClimaLand, :NeuralSnowExt).NeuralSnow;
+
 # Site-specific quantities
 # Error if no site argument is provided
 if length(ARGS) < 1
@@ -46,8 +50,17 @@ ndays = (tf - t0) / 3600 / 24
 
 domain = ClimaLand.Domains.Point(; z_sfc = FT(0))
 
-parameters =
-    SnowParameters{FT}(Δt; α_snow = α, ρ_snow = ρ, earth_param_set = param_set)
+#dens = NeuralSnow.NeuralDepthModel(FT)
+#dens = Snow.Anderson1976{FT}()
+dens = Snow.ConstantDensityModel(ρ)
+depths = z[snow_data_avail]
+
+parameters = SnowParameters{FT}(
+    Δt;
+    α_snow = α,
+    density = dens,
+    earth_param_set = param_set,
+)
 model = ClimaLand.Snow.SnowModel(
     parameters = parameters,
     domain = domain,
@@ -57,6 +70,7 @@ Y, p, coords = ClimaLand.initialize(model)
 
 # Set initial conditions
 Y.snow.S .= FT(SWE[1]) # first data point
+#Y.snow.Z .= FT(depths[1]) #uncomment if using NeuralDepthModel/Anderson1976 instead of ConstantDensityModel
 Y.snow.U .=
     ClimaLand.Snow.energy_from_q_l_and_swe(FT(SWE[1]), FT(0), parameters) # with q_l = 0
 

ext/NeuralSnowExt.jl

@@ -4,5 +4,7 @@ include("neural_snow/DataTools.jl")
 using .DataTools
 include("neural_snow/ModelTools.jl")
 using .ModelTools
+include("neural_snow/NeuralSnow.jl")
+using .NeuralSnow
 
 end