From 4f62e77cd061c1e682d6400e99828c20f7f65d3d Mon Sep 17 00:00:00 2001 From: Andy Charbonneau Date: Thu, 31 Oct 2024 20:23:01 -0400 Subject: [PATCH] Rebasing and integrating NeuralSnow module --- .buildkite/Manifest-v1.11.toml | 77 ++++--- .buildkite/Project.toml | 1 + Project.toml | 4 +- docs/Manifest-v1.11.toml | 31 +-- docs/Project.toml | 1 + .../standalone/Snow/base_tutorial.jl | 4 +- .../standalone/Snow/data_tutorial.jl | 4 +- .../fluxnet/snow_soil/simulation.jl | 2 +- .../standalone/Snow/snowmip_simulation.jl | 16 +- ext/NeuralSnowExt.jl | 2 + ext/neural_snow/NeuralSnow.jl | 213 ++++++++++++++++++ src/Artifacts.jl | 3 + src/integrated/soil_snow_model.jl | 2 +- src/standalone/Snow/Snow.jl | 83 +++++-- src/standalone/Snow/snow_parameterizations.jl | 120 +++++++--- test/standalone/Snow/parameterizations.jl | 18 +- test/standalone/Snow/snow.jl | 10 +- test/standalone/Snow/tool_tests.jl | 123 ++++++++++ 18 files changed, 600 insertions(+), 114 deletions(-) create mode 100644 ext/neural_snow/NeuralSnow.jl diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index fd55e053a2..00ab515d8b 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.buildkite/Manifest-v1.11.toml @@ -2,7 +2,7 @@ julia_version = "1.11.2" manifest_format = "2.0" -project_hash = "bf1216ed2bc76713a01437d07332219f6bcd0c54" +project_hash = "ed73b8f9b8b31b7bd82adc044809f42aee092974" [[deps.ADTypes]] git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f" @@ -296,6 +296,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", "demumble_jll"] git-tree-sha1 = "e0725a467822697171af4dae15cec10b4fc19053" @@ -415,26 +421,17 @@ deps = ["ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "Clim path = ".." uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532" version = "0.15.6" +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"] -git-tree-sha1 = "489c5655993c62fb34293908a6b0877e32f183ee" +git-tree-sha1 = "1e17ed5997da08f1ca186df1882ef10010f70481" uuid = "5c42b081-d73a-476f-9059-fd94b934656c" -version = "0.10.16" +version = "0.10.18" [[deps.ClimaTimeSteppers]] deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"] @@ -963,10 +960,10 @@ uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" version = "0.8.5" [[deps.Flux]] -deps = ["Adapt", "ChainRulesCore", "Compat", "Functors", "LinearAlgebra", "MLDataDevices", "MLUtils", "MacroTools", "NNlib", "OneHotArrays", "Optimisers", "Preferences", "ProgressLogging", "Random", "Reexport", "Setfield", "SparseArrays", "SpecialFunctions", "Statistics", "Zygote"] -git-tree-sha1 = "df520a0727f843576801a0294f5be1a94be28e23" +deps = ["Adapt", "ChainRulesCore", "Compat", "EnzymeCore", "Functors", "LinearAlgebra", "MLDataDevices", "MLUtils", "MacroTools", "NNlib", "OneHotArrays", "Optimisers", "Preferences", "ProgressLogging", "Random", "Reexport", "Setfield", "SparseArrays", "SpecialFunctions", "Statistics", "Zygote"] +git-tree-sha1 = "40c77b726f127356110a3f0aa6e3ecd3ac14159b" uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c" -version = "0.14.25" +version = "0.15.2" [deps.Flux.extensions] FluxAMDGPUExt = "AMDGPU" @@ -1047,10 +1044,10 @@ uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" version = "0.1.3" [[deps.Functors]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "64d8e93700c7a3f28f717d265382d52fac9fa1c1" +deps = ["Compat", "ConstructionBase", "LinearAlgebra", "Random"] +git-tree-sha1 = "60a0339f28a233601cb74468032b5c302d5067de" uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196" -version = "0.4.12" +version = "0.5.2" [[deps.Future]] deps = ["Random"] @@ -1230,9 +1227,9 @@ version = "1.14.2+1" [[deps.HTTP]] deps = ["Base64", "CodecZlib", "ConcurrentUtilities", "Dates", "ExceptionUnwrapping", "Logging", "LoggingExtras", "MbedTLS", "NetworkOptions", "OpenSSL", "PrecompileTools", "Random", "SimpleBufferStream", "Sockets", "URIs", "UUIDs"] -git-tree-sha1 = "6c22309e9a356ac1ebc5c8a217045f9bae6f8d9a" +git-tree-sha1 = "627fcacdb7cb51dc67f557e1598cdffe4dda386d" uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3" -version = "1.10.13" +version = "1.10.14" [[deps.HarfBuzz_jll]] deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "Graphite2_jll", "JLLWrappers", "Libdl", "Libffi_jll"] @@ -1393,9 +1390,9 @@ weakdeps = ["Dates", "Test"] InverseFunctionsTestExt = "Test" [[deps.InvertedIndices]] -git-tree-sha1 = "0dc7b50b8d436461be01300fd8cd45aa0274b038" +git-tree-sha1 = "6da3c4316095de0f5ee2ebd875df8721e7e0bdbe" uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" -version = "1.3.0" +version = "1.3.1" [[deps.IrrationalConstants]] git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" @@ -1709,9 +1706,9 @@ version = "2.16.0+0" [[deps.LogExpFunctions]] deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] -git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea" +git-tree-sha1 = "13ca9e2586b89836fd20cccf56e57e2b9ae7f38f" uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -version = "0.3.28" +version = "0.3.29" [deps.LogExpFunctions.extensions] LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" @@ -1747,9 +1744,9 @@ version = "2024.2.0+0" [[deps.MLDataDevices]] deps = ["Adapt", "Compat", "Functors", "Preferences", "Random"] -git-tree-sha1 = "85b47bc5a8bf0c886286638585df3bec7c9f8269" +git-tree-sha1 = "2580925570e9906bb812896c7940345b68cdb6a3" uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" -version = "1.5.3" +version = "1.6.3" [deps.MLDataDevices.extensions] MLDataDevicesAMDGPUExt = "AMDGPU" @@ -2105,9 +2102,14 @@ version = "0.5.5+0" [[deps.Optimisers]] deps = ["ChainRulesCore", "Functors", "LinearAlgebra", "Random", "Statistics"] -git-tree-sha1 = "c9ff5c686240c31eb8570b662dd1f66f4b183116" +git-tree-sha1 = "c5feff34a5cf6bdc6ca06de0c5b7d6847199f1c0" uuid = "3bd65402-5787-11e9-1adc-39752487f4e2" -version = "0.3.4" +version = "0.4.2" +weakdeps = ["Adapt", "EnzymeCore"] + + [deps.Optimisers.extensions] + OptimisersAdaptExt = ["Adapt"] + OptimisersEnzymeCoreExt = "EnzymeCore" [[deps.Opus_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -2674,9 +2676,9 @@ version = "0.1.2" [[deps.SpecialFunctions]] deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] -git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14" +git-tree-sha1 = "64cca0c26b4f31ba18f13f6c12af7c85f478cfde" uuid = "276daf66-3868-5448-9aa4-cd146d93841b" -version = "2.4.0" +version = "2.5.0" weakdeps = ["ChainRulesCore"] [deps.SpecialFunctions.extensions] @@ -3013,6 +3015,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" @@ -3025,6 +3033,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" diff --git a/.buildkite/Project.toml b/.buildkite/Project.toml index 3e15187523..2de2bfe678 100644 --- a/.buildkite/Project.toml +++ b/.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" diff --git a/Project.toml b/Project.toml index 9c61b07e58..d4f5e0eaaa 100644 --- a/Project.toml +++ b/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" diff --git a/docs/Manifest-v1.11.toml b/docs/Manifest-v1.11.toml index 74cbe53b67..62dd07a083 100644 --- a/docs/Manifest-v1.11.toml +++ b/docs/Manifest-v1.11.toml @@ -2,7 +2,7 @@ julia_version = "1.11.2" manifest_format = "2.0" -project_hash = "375d00e829a5d8c49c86dce3a9f254044ec00898" +project_hash = "04bb15e7c257352c3c06bc57a69afee68a5405e7" [[deps.ADTypes]] git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f" @@ -207,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 = "2a81cc8adf470ac6bd87eef3ca3d194d08a8754c" @@ -431,11 +436,11 @@ deps = ["ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "Clim path = ".." uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532" version = "0.15.6" -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"] @@ -445,9 +450,9 @@ version = "0.1.0" [[deps.ClimaParams]] deps = ["TOML"] -git-tree-sha1 = "489c5655993c62fb34293908a6b0877e32f183ee" +git-tree-sha1 = "1e17ed5997da08f1ca186df1882ef10010f70481" uuid = "5c42b081-d73a-476f-9059-fd94b934656c" -version = "0.10.16" +version = "0.10.18" [[deps.ClimaTimeSteppers]] deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"] @@ -1248,9 +1253,9 @@ version = "1.12.2+2" [[deps.HTTP]] deps = ["Base64", "CodecZlib", "ConcurrentUtilities", "Dates", "ExceptionUnwrapping", "Logging", "LoggingExtras", "MbedTLS", "NetworkOptions", "OpenSSL", "PrecompileTools", "Random", "SimpleBufferStream", "Sockets", "URIs", "UUIDs"] -git-tree-sha1 = "6c22309e9a356ac1ebc5c8a217045f9bae6f8d9a" +git-tree-sha1 = "627fcacdb7cb51dc67f557e1598cdffe4dda386d" uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3" -version = "1.10.13" +version = "1.10.14" [[deps.HarfBuzz_jll]] deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "Graphite2_jll", "JLLWrappers", "Libdl", "Libffi_jll"] @@ -1410,9 +1415,9 @@ weakdeps = ["Dates", "Test"] InverseFunctionsTestExt = "Test" [[deps.InvertedIndices]] -git-tree-sha1 = "0dc7b50b8d436461be01300fd8cd45aa0274b038" +git-tree-sha1 = "6da3c4316095de0f5ee2ebd875df8721e7e0bdbe" uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" -version = "1.3.0" +version = "1.3.1" [[deps.IrrationalConstants]] git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" @@ -1721,9 +1726,9 @@ version = "2.20.1" [[deps.LogExpFunctions]] deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] -git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea" +git-tree-sha1 = "13ca9e2586b89836fd20cccf56e57e2b9ae7f38f" uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -version = "0.3.28" +version = "0.3.29" [deps.LogExpFunctions.extensions] LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" @@ -2668,9 +2673,9 @@ version = "0.1.2" [[deps.SpecialFunctions]] deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] -git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14" +git-tree-sha1 = "64cca0c26b4f31ba18f13f6c12af7c85f478cfde" uuid = "276daf66-3868-5448-9aa4-cd146d93841b" -version = "2.4.0" +version = "2.5.0" weakdeps = ["ChainRulesCore"] [deps.SpecialFunctions.extensions] diff --git a/docs/Project.toml b/docs/Project.toml index 83032395f1..c46132fcfb 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,6 @@ [deps] 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" diff --git a/docs/tutorials/standalone/Snow/base_tutorial.jl b/docs/tutorials/standalone/Snow/base_tutorial.jl index a1f954a4fe..b0328c9b38 100644 --- a/docs/tutorials/standalone/Snow/base_tutorial.jl +++ b/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; diff --git a/docs/tutorials/standalone/Snow/data_tutorial.jl b/docs/tutorials/standalone/Snow/data_tutorial.jl index 1ce7828eb9..e817dcd187 100644 --- a/docs/tutorials/standalone/Snow/data_tutorial.jl +++ b/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; diff --git a/experiments/integrated/fluxnet/snow_soil/simulation.jl b/experiments/integrated/fluxnet/snow_soil/simulation.jl index a9ea1b8cdf..2811b1d592 100644 --- a/experiments/integrated/fluxnet/snow_soil/simulation.jl +++ b/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); diff --git a/experiments/standalone/Snow/snowmip_simulation.jl b/experiments/standalone/Snow/snowmip_simulation.jl index e9a1f11fb9..ef986ef3d7 100644 --- a/experiments/standalone/Snow/snowmip_simulation.jl +++ b/experiments/standalone/Snow/snowmip_simulation.jl @@ -20,6 +20,9 @@ using Statistics using Insolation using DelimitedFiles +using CSV, HTTP, Flux, cuDNN, BSON, StatsBase +NeuralSnow = Base.get_extension(ClimaLand, :NeuralSnowExt).NeuralSnow; + # Site-specific quantities # Error if no site argument is provided if length(ARGS) < 1 @@ -46,8 +49,16 @@ ndays = (tf - t0) / 3600 / 24 domain = ClimaLand.Domains.Point(; z_sfc = FT(0)) -parameters = - SnowParameters{FT}(Δt; α_snow = α, ρ_snow = ρ, earth_param_set = param_set) +#density_model = NeuralSnow.NeuralDepthModel(FT) +density_model = Snow.ConstantDensityModel(ρ) +depths = z[snow_data_avail] + +parameters = SnowParameters{FT}( + Δt; + α_snow = α, + density = density_model, + earth_param_set = param_set, +) model = ClimaLand.Snow.SnowModel( parameters = parameters, domain = domain, @@ -57,6 +68,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 instead of ConstantDensityModel Y.snow.U .= ClimaLand.Snow.energy_from_q_l_and_swe(FT(SWE[1]), FT(0), parameters) # with q_l = 0 diff --git a/ext/NeuralSnowExt.jl b/ext/NeuralSnowExt.jl index 2e2073bddf..4d2c12de8b 100644 --- a/ext/NeuralSnowExt.jl +++ b/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 diff --git a/ext/neural_snow/NeuralSnow.jl b/ext/neural_snow/NeuralSnow.jl new file mode 100644 index 0000000000..ec8b587af9 --- /dev/null +++ b/ext/neural_snow/NeuralSnow.jl @@ -0,0 +1,213 @@ +module NeuralSnow + +using Flux +using ClimaLand +using ClimaLand.Snow: AbstractDensityModel, SnowModel, SnowParameters +import ClimaLand.Snow: + density_prog_vars, + density_prog_types, + density_prog_names, + update_density!, + update_density_prog!, + snow_depth +import ClimaLand.Parameters as LP +using Thermodynamics + +using HTTP, Flux, BSON +include("./ModelTools.jl") +using .ModelTools + +export NeuralDepthModel + +""" + NeuralDepthModel{FT <: AbstractFloat} <: AbstractDensityModel{FT} +Establishes the density parameterization where snow density +is calculated from a neural network determining the rate of change of snow depth, dzdt. +""" +struct NeuralDepthModel{FT} <: AbstractDensityModel{FT} + z_model::Flux.Chain + α::FT +end + +""" + get_znetwork() +Return the snow-depth neural network from the paper. +""" +function get_znetwork() + download_link = ClimaLand.Artifacts.neural_snow_znetwork_link() + z_idx = 1 + p_idx = 7 + nfeatures = 7 + zmodel = ModelTools.make_model(nfeatures, 4, z_idx, p_idx) + zmodel_state = BSON.load(IOBuffer(HTTP.get(download_link).body))[:zstate] + Flux.loadmodel!(zmodel, zmodel_state) + ModelTools.settimescale!(zmodel, 86400.0) + return zmodel +end + +""" + converted_model_type!(model::Flux.Chain, FT::Type) +A wrapper function to aid conversion of the utilized Flux network weights to the type determined by the simulation. +Currently supports `Float32` and `Float64` types. +""" +function converted_model_type(model::Flux.Chain, FT::DataType) + if FT == Float32 + return Flux.f32(model) + elseif FT == Float64 + return Flux.f64(model) + else + error( + "Conversion of Flux Chain weights to the desired float-type is not supported. Please implement the desired conversion + in NeuralSnow.convert_model_type!() or a custom constructor for the NeuralDepthModel with the desired type.", + ) + end +end + +""" + 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. +Can choose to use custom models via the `model` argument, or a custom exponential moving average weight for the network inputs. +The optional Δt argument (model time step, in seconds) can be used to set the lower boundary scaling on the network (when not using a custom model) +in accordance with the paper upon initialization, equivalent to ModelTools.settimescale!(model, Δt, dtype = FT). +If a weight for the exponential moving average is not provided, the default is determined as 1/43200 s⁻¹. This +weight will be unstable for any Δt greater than 43200 seconds, and an alternative value must be supplied. +""" +function NeuralDepthModel( + FT::DataType; + Δt::Union{AbstractFloat, Nothing} = nothing, + model::Flux.Chain = get_znetwork(), + α::Union{AbstractFloat, Nothing} = nothing, +) + usemodel = converted_model_type(model, FT) + weight = !isnothing(α) ? FT(α) : FT(2 / 86400) + if !isnothing(Δt) + ModelTools.settimescale!(usemodel, Δt, dtype = FT) #assumes Δt is provided in seconds + if (Δt > 43200) & (isnothing(α)) + error("Please supply a weight for the + exponential moving average, the + default value is unstable in this case.") + end + end + return NeuralDepthModel{FT}(usemodel, weight) +end + +#Define the additional prognostic variables needed for using this parameterization: +ClimaLand.Snow.density_prog_vars(m::NeuralDepthModel) = + (:Z, :P_avg, :T_avg, :R_avg, :Qrel_avg, :u_avg) +ClimaLand.Snow.density_prog_types(m::NeuralDepthModel{FT}) where {FT} = + (FT, FT, FT, FT, FT, FT) +ClimaLand.Snow.density_prog_names(m::NeuralDepthModel) = + (:surface, :surface, :surface, :surface, :surface, :surface) + +#Extend/Define the appropriate functions needed for this parameterization: + +""" + snow_depth(m::NeuralDepthModel, Y, p, params) + +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. +""" +ClimaLand.Snow.snow_depth(m::NeuralDepthModel, Y, p, params) = Y.snow.Z + + +""" + eval_nn(vmodel, z::FT, swe::FT, P::FT, T::FT, R::FT, qrel::FT, u::FT)::FT where {FT} +Helper function for evaluating the neural network in a pointwise manner over a `ClimaCore.Field` +and returning the output in a broadcastable way. +""" +function eval_nn( + model, + z::FT, + swe::FT, + P::FT, + T::FT, + R::FT, + qrel::FT, + u::FT, +)::FT where {FT} + #model() of a Vector{FT} returns a 1-element Matrix, return the internal value: + return model([z, swe, qrel, R, u, T, P])[1] +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. +""" +function dzdt(density::NeuralDepthModel, Y) + return 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 + Y.snow.P_avg, + Y.snow.T_avg, + Y.snow.R_avg, + Y.snow.Qrel_avg, + Y.snow.u_avg, + ) +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 +not become negative, and here we also clip the tendency of Z +so that depth does not become negative. Additionally, if the +tendencies of Z and S are such that we would encounter Z < S +(rho_snow > rho_liq), we also clip the tendency. +""" +function clip_dZdt(S::FT, Z::FT, dSdt::FT, dZdt::FT, Δt::FT)::FT where {FT} + #Case if S is set to zero: + if (S + dSdt * Δt) <= 0 + return -Z / Δt + #Case if Z would have been set to Z < S: + elseif (Z + dZdt * Δt) < (S + dSdt * Δt) + #more stable form for very small Z, S + return ((dSdt * Δt + S) - Z) / Δt + else + return dZdt + end +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. +""" +function update_density_prog!( + density::NeuralDepthModel, + model::SnowModel, + dY, + Y, + p, +) + + 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), + model.parameters.Δt, + ) + + @. dY.snow.P_avg = density.α * (abs(p.drivers.P_snow) - Y.snow.P_avg) + @. dY.snow.T_avg = + density.α * + (p.drivers.T - model.parameters.earth_param_set.T_freeze - Y.snow.T_avg) + @. dY.snow.R_avg = density.α * (p.drivers.SW_d - Y.snow.R_avg) + dY.snow.Qrel_avg .= + density.α .* ( + Thermodynamics.relative_humidity.( + model.boundary_conditions.atmos.thermo_params, + p.drivers.thermal_state, + ) .- Y.snow.Qrel_avg + ) + @. dY.snow.u_avg = density.α * (p.drivers.u - Y.snow.u_avg) +end + +end diff --git a/src/Artifacts.jl b/src/Artifacts.jl index 6827d30516..9fde8d5459 100644 --- a/src/Artifacts.jl +++ b/src/Artifacts.jl @@ -336,6 +336,9 @@ function bareground_albedo_dataset_path(; context = nothing) ) end +neural_snow_znetwork_link() = + "https://caltech.box.com/shared/static/ay7cv0rhuiytrqbongpeq2y7m3cimhm4.bson" + """ ilamb_dataset_path(; context = nothing) diff --git a/src/integrated/soil_snow_model.jl b/src/integrated/soil_snow_model.jl index 1e9e0910c5..76e04d9899 100644 --- a/src/integrated/soil_snow_model.jl +++ b/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 Δz_soil = p.effective_soil_sfc_depth (; ρc_ds, earth_param_set) = soil_params diff --git a/src/standalone/Snow/Snow.jl b/src/standalone/Snow/Snow.jl index 54cd915ecc..2377eb4293 100644 --- a/src/standalone/Snow/Snow.jl +++ b/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 +for use within an `AbstractSnowModel` type. Current examples include the +`ConstantDensityModel` models. +""" +abstract type AbstractDensityModel{FT <: AbstractFloat} end + +""" + ConstantDensityModel{FT <: AbstractFloat} <: AbstractDensityModel{FT} +Establishes the density parameterization where snow density +is always treated as a constant (type FT), with the provided value in units of kg/m³. +""" +struct ConstantDensityModel{FT} <: AbstractDensityModel{FT} + ρ_snow::FT +end + """ SnowParameters{FT <: AbstractFloat, PSE} @@ -56,9 +73,13 @@ a parameter, and take the larger of the timestep and the physical timescale as the value used in the model. Future implementations will revisit this. $(DocStringExtensions.FIELDS) """ -Base.@kwdef struct SnowParameters{FT <: AbstractFloat, PSE} - "Density of snow (kg/m^3)" - ρ_snow::FT +Base.@kwdef struct SnowParameters{ + FT <: AbstractFloat, + DM <: AbstractDensityModel, + PSE, +} + "Choice of parameterization for snow density" + density::DM "Roughness length over snow for momentum (m)" z_0m::FT "Roughness length over snow for scalars (m)" @@ -83,7 +104,7 @@ end """ SnowParameters{FT}(Δt; - ρ_snow = FT(200), + density = ConstantDensityModel(200), z_0m = FT(0.0024), z_0b = FT(0.00024), α_snow = FT(0.8), @@ -99,7 +120,7 @@ all arguments but `earth_param_set`. """ function SnowParameters{FT}( Δt; - ρ_snow = FT(200), + density::DM = ConstantDensityModel(FT(200)), z_0m = FT(0.0024), z_0b = FT(0.00024), α_snow = FT(0.8), @@ -109,9 +130,9 @@ function SnowParameters{FT}( κ_ice = FT(2.21), ρcD_g = FT(3.553e5), earth_param_set::PSE, -) where {FT <: AbstractFloat, PSE} - return SnowParameters{FT, PSE}( - ρ_snow, +) where {FT <: AbstractFloat, DM <: AbstractDensityModel, PSE} + return SnowParameters{FT, DM, PSE}( + density, z_0m, z_0b, α_snow, @@ -148,10 +169,10 @@ struct SnowModel{FT, PS <: SnowParameters{FT}, BC, D} <: AbstractSnowModel{FT} end function SnowModel(; - parameters::SnowParameters{FT, PSE}, + parameters::SnowParameters{FT, DM, PSE}, domain::ClimaLand.Domains.AbstractDomain, boundary_conditions::BC, -) where {FT, PSE, BC} +) where {FT, DM, PSE, BC} args = (parameters, boundary_conditions, domain) SnowModel{FT, typeof.(args)...}(args...) end @@ -165,7 +186,9 @@ For this model, we track the snow water equivalent S in meters (liquid water volume per ground area) and the energy per unit ground area U [J/m^2] prognostically. """ -prognostic_vars(::SnowModel) = (:S, :U) +prognostic_vars(m::SnowModel) = + (:S, :U, density_prog_vars(m.parameters.density)...) +density_prog_vars(::AbstractDensityModel) = () """ prognostic_types(::SnowModel{FT}) @@ -174,7 +197,9 @@ Returns the prognostic variable types of the snow model; both snow water equivalent and energy per unit area are scalars. """ -prognostic_types(::SnowModel{FT}) where {FT} = (FT, FT) +prognostic_types(m::SnowModel{FT}) where {FT} = + (FT, FT, density_prog_types(m.parameters.density)...) +density_prog_types(::AbstractDensityModel{FT}) where {FT} = () """ prognostic_domain_names(::SnowModel) @@ -184,7 +209,9 @@ both snow water equivalent and energy per unit area are modeling only as a function of (x,y), and not as a function of depth. Therefore their domain name is ":surface". """ -prognostic_domain_names(::SnowModel) = (:surface, :surface) +prognostic_domain_names(m::SnowModel) = + (:surface, :surface, density_prog_names(m.parameters.density)...) +density_prog_names(::AbstractDensityModel) = () """ auxiliary_vars(::SnowModel) @@ -192,8 +219,7 @@ prognostic_domain_names(::SnowModel) = (:surface, :surface) Returns the auxiliary variable names for the snow model. These include 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 depth (`z`, m), +the bulk temperature (`T`, K), the surface temperature (`T_sfc`, K), 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. @@ -208,7 +234,7 @@ auxiliary_vars(::SnowModel) = ( :κ, :T, :T_sfc, - :z, + :ρ_snow, :turbulent_fluxes, :R_n, :energy_runoff, @@ -260,10 +286,11 @@ ClimaLand.name(::SnowModel) = :snow function ClimaLand.make_update_aux(model::SnowModel{FT}) where {FT} function update_aux!(p, Y, t) parameters = model.parameters - _ρ_liq = FT(LP.ρ_cloud_liq(parameters.earth_param_set)) - p.snow.κ .= snow_thermal_conductivity(parameters.ρ_snow, parameters) - @. p.snow.z = snow_depth(Y.snow.S, parameters.ρ_snow, _ρ_liq) + update_density!(parameters.density, parameters, Y, p) + + @. p.snow.κ = snow_thermal_conductivity(p.snow.ρ_snow, parameters) + @. p.snow.q_l = snow_liquid_mass_fraction(Y.snow.U, Y.snow.S, parameters) @@ -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 .= + compute_water_runoff.( + Y.snow.S, + p.snow.q_l, + p.snow.T, + p.snow.ρ_snow, + snow_depth(model.parameters.density, Y, p, parameters), + parameters, + ) @. p.snow.energy_runoff = p.snow.water_runoff * volumetric_internal_energy_liq(FT, parameters) @@ -306,6 +340,7 @@ function ClimaLand.make_compute_exp_tendency(model::SnowModel{FT}) where {FT} # positive fluxes are TOWARDS atmos; negative fluxes increase quantity in snow @. dY.snow.S = -p.snow.applied_water_flux @. dY.snow.U = -p.snow.applied_energy_flux + update_density_prog!(model.parameters.density, model, dY, Y, p) end return compute_exp_tendency! end @@ -316,7 +351,11 @@ end A helper function which clips the total water flux so that snow water equivalent S will not become negative in a timestep Δt. """ -function clip_total_snow_water_flux(S, total_water_flux, Δt) +function clip_total_snow_water_flux( + S::FT, + total_water_flux::FT, + Δt::FT, +) where {FT} if S - total_water_flux * Δt < 0 return S / Δt else diff --git a/src/standalone/Snow/snow_parameterizations.jl b/src/standalone/Snow/snow_parameterizations.jl index 94251fff2a..3fe5d42594 100644 --- a/src/standalone/Snow/snow_parameterizations.jl +++ b/src/standalone/Snow/snow_parameterizations.jl @@ -9,7 +9,8 @@ export snow_surface_temperature, compute_water_runoff, energy_from_q_l_and_swe, energy_from_T_and_swe, - snow_cover_fraction + snow_cover_fraction, + snow_bulk_density """ snow_cover_fraction(x::FT; α = FT(1e-3))::FT where {FT} @@ -23,6 +24,19 @@ 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}, @@ -37,7 +51,7 @@ Once topography or land ice is incorporated, this will need to change to z_sfc + land_ice_depth. Note that land ice can be ~1-3 km thick on Greenland/ -In order to compute surface fluxes, this cannot be large than the +In order to compute surface fluxes, this cannot be larger than the height of the atmosphere measurement location (z_atmos > z_land always). This assumes that the surface elevation is zero. @@ -77,7 +91,7 @@ end """ - ClimaLand.surface_specific_humidity(model::BucketModel, Y, p) + ClimaLand.surface_specific_humidity(model::SnowModel, Y, p) Computes and returns the specific humidity over snow as a weighted fraction of the saturated specific humidity over liquid and frozen @@ -121,19 +135,6 @@ as the bulk temperature T. snow_surface_temperature(T::FT) where {FT} = T -""" - snow_depth(SWE::FT, ρ_snow::FT, ρ_l::FT) where {FT} - -Returns the snow depth given SWE, snow density ρ_snow, and -the density of liquid water ρ_l. - -""" -function snow_depth(SWE::FT, ρ_snow::FT, ρ_l::FT)::FT where {FT} - return SWE * ρ_l / ρ_snow -end - - - """ specific_heat_capacity(q_l::FT, parameters::SnowParameters{FT} @@ -203,18 +204,32 @@ function snow_bulk_temperature( q_l::FT, parameters::SnowParameters{FT}, ) where {FT} - _ρ_i = FT(LP.ρ_cloud_ice(parameters.earth_param_set)) _ρ_l = FT(LP.ρ_cloud_liq(parameters.earth_param_set)) _T_ref = FT(LP.T_0(parameters.earth_param_set)) _LH_f0 = FT(LP.LH_f0(parameters.earth_param_set)) cp_s = specific_heat_capacity(q_l, parameters) - _T_freeze = FT(LP.T_freeze(parameters.earth_param_set)) _ρcD_g = parameters.ρcD_g return _T_ref + (U + (1 - q_l) * _LH_f0 * _ρ_l * SWE) / (_ρ_l * SWE * cp_s + _ρcD_g) end +""" + snow_bulk_density(SWE::FT, z::FT, parameters::SnowParameters{FT}) where {FT} +Returns the snow density given the current model state when depth and SWE are available. +""" +function snow_bulk_density( + SWE::FT, + z::FT, + parameters::SnowParameters{FT}, +)::FT where {FT} + ρ_l = FT(LP.ρ_cloud_liq(parameters.earth_param_set)) + ε = eps(FT) #for preventing dividing by zero + #return SWE/z * ρ_l but tend to ρ_l as SWE → 0 + #also handle instabilities when z, SWE both near machine precision + return max(SWE, ε) / max(z, SWE, ε) * ρ_l +end + """ snow_liquid_mass_fraction(U::FT, SWE::FT, parameters::SnowParameters{FT}) where {FT} @@ -226,7 +241,6 @@ function snow_liquid_mass_fraction( SWE::FT, parameters::SnowParameters{FT}, ) where {FT} - _ρ_i = FT(LP.ρ_cloud_ice(parameters.earth_param_set)) _ρ_l = FT(LP.ρ_cloud_liq(parameters.earth_param_set)) _T_ref = FT(LP.T_0(parameters.earth_param_set)) _T_freeze = FT(LP.T_freeze(parameters.earth_param_set)) @@ -234,19 +248,17 @@ function snow_liquid_mass_fraction( _cp_i = FT(LP.cp_i(parameters.earth_param_set)) _cp_l = FT(LP.cp_l(parameters.earth_param_set)) - ΔT = _T_freeze - _T_ref - _ρcD_g = parameters.ρcD_g Uminus = (_ρ_l * SWE * _cp_i + _ρcD_g) * (_T_freeze - _T_ref) - _ρ_l * SWE * _LH_f0 Uplus = (_ρ_l * SWE * _cp_l + _ρcD_g) * (_T_freeze - _T_ref) if U < Uminus - FT(0) + return FT(0) elseif U > Uplus - FT(1) + return FT(1) else - (U - Uminus) / max((Uplus - Uminus), eps(FT)) + return (U - Uminus) / max((Uplus - Uminus), eps(FT)) end end @@ -311,13 +323,15 @@ liquid water (runoff) from the snowpack. Runoff occurs as the snow melts and exceeds the water holding capacity. """ -function compute_water_runoff(S::FT, q_l::FT, T::FT, parameters) where {FT} - _ρ_l = FT(LP.ρ_cloud_liq(parameters.earth_param_set)) - ρ_snow::FT = parameters.ρ_snow - Ksat::FT = parameters.Ksat - Δt::FT = parameters.Δt - depth = snow_depth(S, ρ_snow, _ρ_l) - τ = runoff_timescale(depth, Ksat, Δt) +function compute_water_runoff( + S::FT, + q_l::FT, + T::FT, + ρ_snow::FT, + z::FT, + parameters, +) where {FT} + τ = runoff_timescale(z, parameters.Ksat, parameters.Δt) q_l_max::FT = maximum_liquid_mass_fraction(T, ρ_snow, parameters) return -(q_l - q_l_max) * heaviside(q_l - q_l_max) / τ * S / max(1 - q_l, FT(0.01)) @@ -356,7 +370,6 @@ If T = T_freeze, we return the energy as if q_l = 0. """ function energy_from_T_and_swe(S::FT, T::FT, parameters) where {FT} - _ρ_i = FT(LP.ρ_cloud_ice(parameters.earth_param_set)) _ρ_l = FT(LP.ρ_cloud_liq(parameters.earth_param_set)) _T_ref = FT(LP.T_0(parameters.earth_param_set)) _T_freeze = FT(LP.T_freeze(parameters.earth_param_set)) @@ -372,3 +385,46 @@ function energy_from_T_and_swe(S::FT, T::FT, parameters) where {FT} end 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. +""" +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) +end + +""" + update_density!(density::ConstantDensityModel, params::SnowParameters, Y, p) +Extends the update_density! function for the ConstantDensityModel type. +""" +function update_density!( + density::ConstantDensityModel, + params::SnowParameters, + Y, + p, +) + p.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. +""" +function update_density_prog!( + density::AbstractDensityModel, + model::SnowModel, + dY, + Y, + p, +) + return nothing +end diff --git a/test/standalone/Snow/parameterizations.jl b/test/standalone/Snow/parameterizations.jl index 8ab6aafb09..33fe751777 100644 --- a/test/standalone/Snow/parameterizations.jl +++ b/test/standalone/Snow/parameterizations.jl @@ -27,6 +27,7 @@ for FT in (Float32, Float64) κ_air = FT(LP.K_therm(param_set)) ρ_snow = FT(200) + densitymodel = Snow.ConstantDensityModel(ρ_snow) z_0m = FT(0.0024) α_snow = FT(0.8) # These values should match ClimaParams @@ -37,9 +38,13 @@ for FT in (Float32, Float64) κ_ice = FT(2.21) ρcD_g = FT(1700 * 2.09e3 * 0.1) Δt = Float64(180.0) - parameters = SnowParameters{FT}(Δt; earth_param_set = param_set) - @test parameters.ρ_snow == ρ_snow - @test typeof(parameters.ρ_snow) == FT + parameters = SnowParameters{FT}( + Δt; + density = densitymodel, + earth_param_set = param_set, + ) + @test parameters.density.ρ_snow == ρ_snow + @test typeof(parameters.density.ρ_snow) == FT @test parameters.ρcD_g == ρcD_g @test typeof(parameters.ρcD_g) == FT @test parameters.z_0m == z_0m @@ -67,8 +72,7 @@ for FT in (Float32, Float64) ) SWE = cat(FT.(rand(100) .+ 0.2), FT(0), dims = 1) - z = snow_depth.(SWE, ρ_snow, _ρ_l) - @test all(z .== SWE * _ρ_l ./ ρ_snow) + z = SWE * _ρ_l ./ ρ_snow @test specific_heat_capacity(FT(1.0), parameters) == _cp_l @test specific_heat_capacity(FT(0.0), parameters) == _cp_i @test snow_thermal_conductivity(ρ_snow, parameters) == @@ -76,6 +80,10 @@ for FT in (Float32, Float64) (FT(0.07) * (ρ_snow / _ρ_i) + FT(0.93) * (ρ_snow / _ρ_i)^2) * (κ_ice - κ_air) @test runoff_timescale.(z, Ksat, FT(Δt)) ≈ max.(Δt, z ./ Ksat) + ρ_calc = snow_bulk_density.(SWE, z, parameters) + @test prod(ρ_calc[1:(end - 1)] .≈ ρ_snow) + @test ρ_calc[end] == _ρ_l + @test snow_bulk_density(eps(FT(0)), 2 * eps(FT(0)), parameters) == _ρ_l U = cat(FT.(Array(LinRange(-1e8, 1e7, 100))), FT(0), dims = 1) diff --git a/test/standalone/Snow/snow.jl b/test/standalone/Snow/snow.jl index 9098aedaa7..575570b6a8 100644 --- a/test/standalone/Snow/snow.jl +++ b/test/standalone/Snow/snow.jl @@ -57,7 +57,7 @@ import ClimaLand.Parameters as LP :κ, :T, :T_sfc, - :z, + :ρ_snow, :turbulent_fluxes, :R_n, :energy_runoff, @@ -131,6 +131,11 @@ import ClimaLand.Parameters as LP @test turb_fluxes.shf == p.snow.turbulent_fluxes.shf @test turb_fluxes.lhf == p.snow.turbulent_fluxes.lhf @test turb_fluxes.vapor_flux == p.snow.turbulent_fluxes.vapor_flux + old_ρ = deepcopy(p.snow.ρ_snow) + Snow.update_density!(model.parameters.density, model.parameters, Y, p) + @test p.snow.ρ_snow == old_ρ + z = snow_depth(model.parameters.density, Y, p, parameters) + @test parent(z)[1] == FT(0.5) # Now compute tendencies and make sure they operate correctly. dY = similar(Y) @@ -146,6 +151,9 @@ import ClimaLand.Parameters as LP -p.snow.turbulent_fluxes.shf - p.snow.turbulent_fluxes.lhf - p.snow.R_n + p.snow.energy_runoff ) + @test isnothing( + Snow.update_density_prog!(model.parameters.density, model, dY, Y, p), + ) # Now try a step where the snow will melt diff --git a/test/standalone/Snow/tool_tests.jl b/test/standalone/Snow/tool_tests.jl index 6aae52b780..c139733642 100644 --- a/test/standalone/Snow/tool_tests.jl +++ b/test/standalone/Snow/tool_tests.jl @@ -3,6 +3,17 @@ using Test using BSON, Dates, HTTP using DataFrames, CSV, StatsBase, Flux, LinearAlgebra +import ClimaParams as CP +import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput +using Thermodynamics +using SurfaceFluxes +using StaticArrays +using Dates +using ClimaLand.Snow +using ClimaLand.Domains +import ClimaLand.Parameters as LP + + try import CUDA import cuDNN @@ -12,10 +23,12 @@ end DataToolsExt = Base.get_extension(ClimaLand, :NeuralSnowExt) ModelToolsExt = Base.get_extension(ClimaLand, :NeuralSnowExt) +NeuralSnowExt = Base.get_extension(ClimaLand, :NeuralSnowExt) if !isnothing(DataToolsExt) DataTools = DataToolsExt.DataTools ModelTools = ModelToolsExt.ModelTools + NeuralSnow = NeuralSnowExt.NeuralSnow @testset "Testing Data Utilities" begin start_date = "2015-01-01" end_date = "2015-02-01" @@ -322,4 +335,114 @@ if !isnothing(DataToolsExt) sqrt(sum((pred_series .- true_series) .^ 2) ./ length(pred_series)) @test series_err <= 0.2 end + + @testset "Testing NeuralSnow module" begin + #Model setup: + FT = Float32 + earth_param_set = LP.LandParameters(FT) + start_date = DateTime(2005) + param_set = LP.LandParameters(FT) + Δt = FT(180.0) + domain = Point(; z_sfc = FT(0)) + "Radiation" + SW_d = TimeVaryingInput((t) -> eltype(t)(20.0)) + LW_d = TimeVaryingInput((t) -> eltype(t)(20.0)) + rad = ClimaLand.PrescribedRadiativeFluxes(FT, SW_d, LW_d, start_date) + "Atmos" + precip = TimeVaryingInput((t) -> eltype(t)(0.01 / Δt)) + T_atmos = TimeVaryingInput((t) -> eltype(t)(278.0)) + u_atmos = TimeVaryingInput((t) -> eltype(t)(10.0)) + q_atmos = TimeVaryingInput((t) -> eltype(t)(0.03)) + h_atmos = FT(3) + P_atmos = TimeVaryingInput((t) -> eltype(t)(101325)) + atmos = ClimaLand.PrescribedAtmosphere( + precip, + precip, + T_atmos, + u_atmos, + q_atmos, + P_atmos, + start_date, + h_atmos, + earth_param_set, + ) + + #Test extension utilities + z_model = NeuralSnow.get_znetwork() + @test typeof(z_model) <: Flux.Chain + z32 = NeuralSnow.converted_model_type(z_model, FT) + @test eltype(z32[1].layers[1].weight) == FT + z64 = NeuralSnow.converted_model_type(z_model, Float64) + @test eltype(z64[1].layers[1].weight) == Float64 + dens_model1 = NeuralSnow.NeuralDepthModel(FT) + @test prod(Flux.params(dens_model1.z_model) .== Flux.params(z32)) + @test eltype(dens_model1.α) == FT + test_alph = 3 / 86400 + dens_model2 = NeuralSnow.NeuralDepthModel(FT, α = test_alph, Δt = Δt) + @test dens_model2.α == FT(test_alph) + @test dens_model2.z_model[:final_scale].weight[2, 2] == FT(1 / Δt) + + parameters = SnowParameters{FT}( + Δt; + earth_param_set = param_set, + density = dens_model2, + ) + model = ClimaLand.Snow.SnowModel( + parameters = parameters, + domain = domain, + boundary_conditions = ClimaLand.Snow.AtmosDrivenSnowBC(atmos, rad), + ) + + drivers = ClimaLand.get_drivers(model) + Y, p, coords = ClimaLand.initialize(model) + @test (Y.snow |> propertynames) == + (:S, :U, :Z, :P_avg, :T_avg, :R_avg, :Qrel_avg, :u_avg) + + Y.snow.S .= FT(0.1) + Y.snow.U .= + ClimaLand.Snow.energy_from_T_and_swe.( + Y.snow.S, + FT(273.0), + Ref(model.parameters), + ) + Y.snow.Z .= FT(0.2) + set_initial_cache! = ClimaLand.make_set_initial_cache(model) + t0 = FT(0.0) + set_initial_cache!(p, Y, t0) + @test snow_depth(model.parameters.density, Y, p, parameters) == Y.snow.Z + + output1 = NeuralSnow.eval_nn( + dens_model2.z_model, + FT.([0, 0, 0, 0, 0, 0, 0])..., + ) + + @test eltype(output1) == FT + @test output1 == 0.0f0 + + zerofield = similar(Y.snow.Z) + zerofield .= FT(0) + @test NeuralSnow.dzdt(dens_model2, Y) == zerofield + + Z = FT(0.5) + S = FT(0.1) + dzdt = FT(1 / Δt) + dsdt = FT(1 / Δt) + @test NeuralSnow.clip_dZdt(S, Z, dsdt, dzdt, Δt) == dzdt + + @test NeuralSnow.clip_dZdt(Z, S, dsdt, dzdt, Δt) ≈ FT(1.4 / Δt) + + @test NeuralSnow.clip_dZdt(S, Z, FT(-S / Δt), dzdt, Δt) ≈ FT(-Z / Δt) + + oldρ = p.snow.ρ_snow + Snow.update_density!(dens_model2, model.parameters, Y, p) + @test p.snow.ρ_snow == oldρ + + dY = similar(Y) + dswe_by_precip = 0.1 + Y.snow.P_avg .= FT(dswe_by_precip / Δt) + NeuralSnow.update_density_prog!(dens_model2, model, dY, Y, p) + @test parent(dY.snow.Z)[1] * Δt > dswe_by_precip + new_dYP = FT(test_alph) .* (p.drivers.P_snow .- Y.snow.P_avg) + @test dY.snow.P_avg == new_dYP + end end