From d867a7a8050a78b12bb699362ab16663c241bdf8 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 | 94 ++++---- .buildkite/Project.toml | 1 + Project.toml | 4 +- docs/Manifest-v1.11.toml | 16 +- docs/Manifest.toml | 29 ++- 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 | 17 +- ext/NeuralSnowExt.jl | 2 + ext/neural_snow/NeuralSnow.jl | 215 ++++++++++++++++++ src/Artifacts.jl | 3 + src/integrated/soil_snow_model.jl | 2 +- src/standalone/Snow/Snow.jl | 84 +++++-- src/standalone/Snow/snow_parameterizations.jl | 124 +++++++--- test/standalone/Snow/parameterizations.jl | 18 +- test/standalone/Snow/snow.jl | 10 +- test/standalone/Snow/tool_tests.jl | 123 ++++++++++ 19 files changed, 631 insertions(+), 122 deletions(-) create mode 100644 ext/neural_snow/NeuralSnow.jl diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index 6b2ebd3c40..ff0aba5b59 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 = "14f223707cbd1a2cb15b66a90e7be9a546a0f205" [[deps.ADTypes]] git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f" @@ -152,9 +152,9 @@ version = "1.11.0" [[deps.Atomix]] deps = ["UnsafeAtomics"] -git-tree-sha1 = "14e254ef74e44cd6ed27fbb751d4e1f9bbf085cc" +git-tree-sha1 = "c3b238aa28c1bebd4b5ea4988bebf27e9a01b72b" uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458" -version = "1.0.0" +version = "1.0.1" [deps.Atomix.extensions] AtomixCUDAExt = "CUDA" @@ -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" @@ -340,9 +346,9 @@ version = "1.1.1" [[deps.CairoMakie]] deps = ["CRC32c", "Cairo", "Cairo_jll", "Colors", "FileIO", "FreeType", "GeometryBasics", "LinearAlgebra", "Makie", "PrecompileTools"] -git-tree-sha1 = "c3161fbfe99d9d7ee121cf2017d49966b152857c" +git-tree-sha1 = "076fca013bc2c73b2038f7d27ea4fa3a624fcd9d" uuid = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" -version = "0.12.16" +version = "0.12.17" [[deps.Cairo_jll]] deps = ["Artifacts", "Bzip2_jll", "CompilerSupportLibraries_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "JLLWrappers", "LZO_jll", "Libdl", "Pixman_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Zlib_jll", "libpng_jll"] @@ -405,30 +411,21 @@ uuid = "908f55d8-4145-4867-9c14-5dad1a479e4d" version = "0.4.6" [[deps.ClimaDiagnostics]] -deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "SciMLBase"] -git-tree-sha1 = "ace174fe5e4ae04c50a7835683baecd92e49c0d4" +deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "OrderedCollections", "SciMLBase"] +git-tree-sha1 = "e9ac94af815dcae2a2ab24e54b53e76cca6258b7" uuid = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f" -version = "0.2.10" +version = "0.2.11" [[deps.ClimaLand]] deps = ["ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaUtilities", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"] 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"] @@ -986,9 +983,9 @@ version = "0.14.25" [[deps.Fontconfig_jll]] deps = ["Artifacts", "Bzip2_jll", "Expat_jll", "FreeType2_jll", "JLLWrappers", "Libdl", "Libuuid_jll", "Zlib_jll"] -git-tree-sha1 = "db16beca600632c95fc8aca29890d83788dd8b23" +git-tree-sha1 = "21fac3c77d7b5a9fc03b0ec503aa1a6392c34d2b" uuid = "a3f928ae-7b40-5064-980b-68af3947d34b" -version = "2.13.96+0" +version = "2.15.0+0" [[deps.Format]] git-tree-sha1 = "9c68794ef81b08086aeb32eeaf33531668d5f5fc" @@ -1230,15 +1227,15 @@ 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 = "ae350b8225575cc3ea385d4131c81594f86dfe4f" +git-tree-sha1 = "6c22309e9a356ac1ebc5c8a217045f9bae6f8d9a" uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3" -version = "1.10.12" +version = "1.10.13" [[deps.HarfBuzz_jll]] deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "Graphite2_jll", "JLLWrappers", "Libdl", "Libffi_jll"] -git-tree-sha1 = "401e4f3f30f43af2c8478fc008da50096ea5240f" +git-tree-sha1 = "55c53be97790242c29031e5cd45e8ac296dadda3" uuid = "2e76f6c2-a576-52d4-95c1-20adfe4de566" -version = "8.3.1+0" +version = "8.5.0+0" [[deps.Hwloc_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -1473,10 +1470,10 @@ uuid = "b14d175d-62b4-44ba-8fb7-3064adc8c3ec" version = "0.2.4" [[deps.KernelAbstractions]] -deps = ["Adapt", "Atomix", "InteractiveUtils", "MacroTools", "PrecompileTools", "Requires", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"] -git-tree-sha1 = "27044736be7c5727d35fc4318d7949dee33c37b4" +deps = ["Adapt", "Atomix", "InteractiveUtils", "MacroTools", "PrecompileTools", "Requires", "StaticArrays", "UUIDs"] +git-tree-sha1 = "b9a838cd3028785ac23822cded5126b3da394d1a" uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" -version = "0.9.30" +version = "0.9.31" weakdeps = ["EnzymeCore", "LinearAlgebra", "SparseArrays"] [deps.KernelAbstractions.extensions] @@ -1840,15 +1837,15 @@ version = "0.5.13" [[deps.Makie]] deps = ["Animations", "Base64", "CRC32c", "ColorBrewer", "ColorSchemes", "ColorTypes", "Colors", "Contour", "Dates", "DelaunayTriangulation", "Distributions", "DocStringExtensions", "Downloads", "FFMPEG_jll", "FileIO", "FilePaths", "FixedPointNumbers", "Format", "FreeType", "FreeTypeAbstraction", "GeometryBasics", "GridLayoutBase", "ImageBase", "ImageIO", "InteractiveUtils", "Interpolations", "IntervalSets", "InverseFunctions", "Isoband", "KernelDensity", "LaTeXStrings", "LinearAlgebra", "MacroTools", "MakieCore", "Markdown", "MathTeXEngine", "Observables", "OffsetArrays", "Packing", "PlotUtils", "PolygonOps", "PrecompileTools", "Printf", "REPL", "Random", "RelocatableFolders", "Scratch", "ShaderAbstractions", "Showoff", "SignedDistanceFields", "SparseArrays", "Statistics", "StatsBase", "StatsFuns", "StructArrays", "TriplotBase", "UnicodeFun", "Unitful"] -git-tree-sha1 = "5e4e0e027642293da251bf35dac408d692ccba8b" +git-tree-sha1 = "260d6e1ac8abcebd939029e6eedeba4e3870f13a" uuid = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" -version = "0.21.16" +version = "0.21.17" [[deps.MakieCore]] deps = ["ColorTypes", "GeometryBasics", "IntervalSets", "Observables"] -git-tree-sha1 = "ae4dbe0fcf1594ed98594e5f4ee685295a2a6f74" +git-tree-sha1 = "b774d0563bc332f64d136d50d0420a195d9bdcc6" uuid = "20f20a25-4f0e-4fdf-b5d1-57303727442b" -version = "0.8.10" +version = "0.8.11" [[deps.ManualMemory]] git-tree-sha1 = "bcaef4fc7a0cfe2cba636d84cda54b5e4e4ca3cd" @@ -2031,9 +2028,9 @@ uuid = "510215fc-4207-5dde-b226-833fc4488ee2" version = "0.5.5" [[deps.OffsetArrays]] -git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e" +git-tree-sha1 = "39d000d9c33706b8364817d8894fae1548f40295" uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" -version = "1.14.1" +version = "1.14.2" weakdeps = ["Adapt"] [deps.OffsetArrays.extensions] @@ -2070,9 +2067,9 @@ version = "3.2.4+0" [[deps.OpenJpeg_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "libpng_jll"] -git-tree-sha1 = "f4cb457ffac5f5cf695699f82c537073958a6a6c" +git-tree-sha1 = "0a41c2d8e204a3ad713242139628e01a29556967" uuid = "643b3616-a352-519d-856d-80112ee9badc" -version = "2.5.2+0" +version = "2.5.3+0" [[deps.OpenLibm_jll]] deps = ["Artifacts", "Libdl"] @@ -2527,9 +2524,9 @@ version = "0.1.0" [[deps.SQLite_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"] -git-tree-sha1 = "004fffbe2711abdc7263a980bbb1af9620781dd9" +git-tree-sha1 = "e84fab7d16107342d7638fbd519151d9a0d80720" uuid = "76ed43ae-9a5d-5a62-8c75-30186b810ce8" -version = "3.45.3+0" +version = "3.47.2+0" [[deps.SafeTestsets]] git-tree-sha1 = "81ec49d645af090901120a1542e67ecbbe044db3" @@ -2538,9 +2535,9 @@ version = "0.1.0" [[deps.SciMLBase]] deps = ["ADTypes", "Accessors", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "Expronicon", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "SciMLStructures", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface"] -git-tree-sha1 = "899468ac3c2fa6b87151cb3fa29367329017d365" +git-tree-sha1 = "87e054302a94a2d087f918ad50b0290977db67e6" uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -version = "2.66.0" +version = "2.67.0" [deps.SciMLBase.extensions] SciMLBaseChainRulesCoreExt = "ChainRulesCore" @@ -2984,12 +2981,6 @@ git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278" uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f" version = "0.2.1" -[[deps.UnsafeAtomicsLLVM]] -deps = ["LLVM", "UnsafeAtomics"] -git-tree-sha1 = "de4287a6569bcf3a8d6201d387991a8dda25c954" -uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249" -version = "0.2.2" - [[deps.Unzip]] git-tree-sha1 = "ca0969166a028236229f63514992fc073799bb78" uuid = "41fe7b60-77ed-43a1-b4f0-825fd5a5650d" @@ -3019,6 +3010,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" @@ -3031,6 +3028,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 7f298bb8a0..b0418cb60a 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 6fc944008c..85b87a6cf3 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 5140afaf63..c0a8c631f8 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 = "a5dc443666e7b6f001f45671e08cb0a420c01a66" +project_hash = "44cdd8eaaeb3c23ae85e35918e303ced7ff57a4f" [[deps.ADTypes]] git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f" @@ -217,6 +217,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" @@ -441,11 +446,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"] @@ -471,9 +476,14 @@ weakdeps = ["BenchmarkTools", "CUDA", "OrderedCollections", "PrettyTables", "Sta [[deps.ClimaUtilities]] deps = ["Artifacts", "ClimaComms", "Dates"] +<<<<<<< HEAD +<<<<<<< HEAD git-tree-sha1 = "cd699551a7c9f721363bbf92b2c61c436b1a3959" uuid = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513" version = "0.1.20" +git-tree-sha1 = "ae6e4d3223ecb29d3f3d9c5f82b622bae12f3bff" +uuid = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513" +version = "0.1.19" [deps.ClimaUtilities.extensions] ClimaUtilitiesCUDAExt = "CUDA" diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 9ba4846fb4..1ba6c9cabf 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.7" manifest_format = "2.0" -project_hash = "0271f9782f48e923fcefcee2c42babb40262588c" +project_hash = "17901eb92a5662b053574516009767eae929e671" [[deps.ADTypes]] git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f" @@ -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" @@ -1463,6 +1473,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" @@ -2735,6 +2756,12 @@ weakdeps = ["Adapt", "GPUArraysCore", "SparseArrays", "StaticArrays"] StructArraysSparseArraysExt = "SparseArrays" StructArraysStaticArraysExt = "StaticArrays" +[[deps.StyledStrings]] +deps = ["PrecompileTools", "TOML"] +git-tree-sha1 = "711c9650010c95814911c2005ea04e70e70e65ed" +uuid = "f489334b-da3d-4c2e-b8f0-e476e12c162b" +version = "1.0.3" + [[deps.SuiteSparse]] deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" 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..68cd4a1ba8 100644 --- a/experiments/standalone/Snow/snowmip_simulation.jl +++ b/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,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 +69,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..06fcdb48fe --- /dev/null +++ b/ext/neural_snow/NeuralSnow.jl @@ -0,0 +1,215 @@ +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 + +#Todo: Write tests for these functions and a few in Snow.jl or snow_parameterizations.jl diff --git a/src/Artifacts.jl b/src/Artifacts.jl index fd8abe052f..fafb666e02 100644 --- a/src/Artifacts.jl +++ b/src/Artifacts.jl @@ -336,4 +336,7 @@ function bareground_albedo_dataset_path(; context = nothing) ) end +neural_snow_znetwork_link() = + "https://caltech.box.com/shared/static/ay7cv0rhuiytrqbongpeq2y7m3cimhm4.bson" + end 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..eb6447e684 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, @@ -254,16 +280,16 @@ auxiliary_domain_names(::SnowModel) = ( :surface, ) - 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 +298,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 +339,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 +350,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..f511ce3d03 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}, @@ -34,13 +48,13 @@ Returns the surface height of the `Snow` model; surface (ground) elevation and ignores snow depth (CLM). Once topography or land ice is incorporated, this will need to change to -z_sfc + land_ice_depth. Note that land ice can +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. +This assumes that the surface elevation is zero. """ function ClimaLand.surface_height(model::SnowModel{FT}, Y, p) where {FT} return FT(0) @@ -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..ec05602764 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))) == _ρ_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