From 4b38583c0f3e69b8c0c0a6828cd1cacd9a4975b0 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 12 Sep 2024 17:07:04 +0100 Subject: [PATCH 01/33] Add initial implementation SpeedyWeather integration --- SpeedyWeather/Manifest.toml | 1316 +++++++++++++++++++++++++++++++++++ SpeedyWeather/Project.toml | 3 + SpeedyWeather/src/speedy.jl | 202 ++++++ 3 files changed, 1521 insertions(+) create mode 100644 SpeedyWeather/Manifest.toml create mode 100644 SpeedyWeather/Project.toml create mode 100644 SpeedyWeather/src/speedy.jl diff --git a/SpeedyWeather/Manifest.toml b/SpeedyWeather/Manifest.toml new file mode 100644 index 0000000..05f8d07 --- /dev/null +++ b/SpeedyWeather/Manifest.toml @@ -0,0 +1,1316 @@ +# This file is machine-generated - editing it directly is not advised + +julia_version = "1.10.4" +manifest_format = "2.0" +project_hash = "cf0fb01158045c321488fa4aacf601601dfc4784" + +[[deps.AbstractFFTs]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef" +uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c" +version = "1.5.0" + + [deps.AbstractFFTs.extensions] + AbstractFFTsChainRulesCoreExt = "ChainRulesCore" + AbstractFFTsTestExt = "Test" + + [deps.AbstractFFTs.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[[deps.Accessors]] +deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown", "Test"] +git-tree-sha1 = "f61b15be1d76846c0ce31d3fcfac5380ae53db6a" +uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +version = "0.1.37" + + [deps.Accessors.extensions] + AccessorsAxisKeysExt = "AxisKeys" + AccessorsIntervalSetsExt = "IntervalSets" + AccessorsStaticArraysExt = "StaticArrays" + AccessorsStructArraysExt = "StructArrays" + AccessorsUnitfulExt = "Unitful" + + [deps.Accessors.weakdeps] + AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" + Requires = "ae029012-a4dd-5104-9daa-d747884805df" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" + Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" + +[[deps.Adapt]] +deps = ["LinearAlgebra", "Requires"] +git-tree-sha1 = "6a55b747d1812e699320963ffde36f1ebdda4099" +uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +version = "4.0.4" +weakdeps = ["StaticArrays"] + + [deps.Adapt.extensions] + AdaptStaticArraysExt = "StaticArrays" + +[[deps.AliasTables]] +deps = ["PtrArrays", "Random"] +git-tree-sha1 = "9876e1e164b144ca45e9e3198d0b689cadfed9ff" +uuid = "66dad0bd-aa9a-41b7-9441-69ab47430ed8" +version = "1.1.3" + +[[deps.ArgCheck]] +git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4" +uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197" +version = "2.3.0" + +[[deps.ArgTools]] +uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" +version = "1.1.1" + +[[deps.Artifacts]] +uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" + +[[deps.AssociatedLegendrePolynomials]] +git-tree-sha1 = "3204d769e06c5678b23cf928d850f2f4ad5ec8a5" +uuid = "2119f1ac-fb78-50f5-8cc0-dda848ebdb19" +version = "1.0.1" + +[[deps.Atomix]] +deps = ["UnsafeAtomics"] +git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be" +uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458" +version = "0.1.0" + +[[deps.BFloat16s]] +deps = ["LinearAlgebra", "Printf", "Random", "Test"] +git-tree-sha1 = "2c7cc21e8678eff479978a0a2ef5ce2f51b63dff" +uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b" +version = "0.5.0" + +[[deps.BangBang]] +deps = ["Accessors", "ConstructionBase", "InitialValues", "LinearAlgebra", "Requires"] +git-tree-sha1 = "e2144b631226d9eeab2d746ca8880b7ccff504ae" +uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" +version = "0.4.3" + + [deps.BangBang.extensions] + BangBangChainRulesCoreExt = "ChainRulesCore" + BangBangDataFramesExt = "DataFrames" + BangBangStaticArraysExt = "StaticArrays" + BangBangStructArraysExt = "StructArrays" + BangBangTablesExt = "Tables" + BangBangTypedTablesExt = "TypedTables" + + [deps.BangBang.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" + Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" + TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9" + +[[deps.Base64]] +uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" + +[[deps.Baselet]] +git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e" +uuid = "9718e550-a3fa-408a-8086-8db961cd8217" +version = "0.1.1" + +[[deps.BitInformation]] +deps = ["Distributions", "Random", "StatsBase"] +git-tree-sha1 = "2cf994e66a0886a91ba108cb3cfc044363f0bb01" +uuid = "de688a37-743e-4ac2-a6f0-bd62414d1aa7" +version = "0.6.3" + +[[deps.Blosc_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Lz4_jll", "Zlib_jll", "Zstd_jll"] +git-tree-sha1 = "19b98ee7e3db3b4eff74c5c9c72bf32144e24f10" +uuid = "0b7ba130-8d10-5ba8-a3d6-c5182647fed9" +version = "1.21.5+0" + +[[deps.Bzip2_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "9e2a6b69137e6969bab0152632dcb3bc108c8bdd" +uuid = "6e34b625-4abd-537c-b88f-471c36dfa7a0" +version = "1.0.8+1" + +[[deps.CEnum]] +git-tree-sha1 = "389ad5c84de1ae7cf0e28e381131c98ea87d54fc" +uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82" +version = "0.5.0" + +[[deps.CFTime]] +deps = ["Dates", "Printf"] +git-tree-sha1 = "5afb5c5ba2688ca43a9ad2e5a91cbb93921ccfa1" +uuid = "179af706-886a-5703-950a-314cd64e0468" +version = "0.1.3" + +[[deps.CUDA]] +deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"] +git-tree-sha1 = "fdd9dfb67dfefd548f51000cc400bb51003de247" +uuid = "052768ef-5323-5732-b1bb-66c8b64840ba" +version = "5.4.3" + + [deps.CUDA.extensions] + ChainRulesCoreExt = "ChainRulesCore" + EnzymeCoreExt = "EnzymeCore" + SpecialFunctionsExt = "SpecialFunctions" + + [deps.CUDA.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" + SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" + +[[deps.CUDA_Driver_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "325058b426c2b421e3d2df3d5fa646d72d2e3e7e" +uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc" +version = "0.9.2+0" + +[[deps.CUDA_Runtime_Discovery]] +deps = ["Libdl"] +git-tree-sha1 = "33576c7c1b2500f8e7e6baa082e04563203b3a45" +uuid = "1af6417a-86b4-443c-805f-a4643ffb695f" +version = "0.3.5" + +[[deps.CUDA_Runtime_jll]] +deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] +git-tree-sha1 = "afea94249b821dc754a8ca6695d3daed851e1f5a" +uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" +version = "0.14.1+0" + +[[deps.ChunkSplitters]] +deps = ["Compat", "TestItems"] +git-tree-sha1 = "783507c1f2371c8f2d321f41c3057ecd42cafa83" +uuid = "ae650224-84b6-46f8-82ea-d812ca08434e" +version = "2.4.5" + +[[deps.CodecZlib]] +deps = ["TranscodingStreams", "Zlib_jll"] +git-tree-sha1 = "bce6804e5e6044c6daab27bb533d1295e4a2e759" +uuid = "944b1d66-785c-5afd-91f1-9de20f533193" +version = "0.7.6" + +[[deps.ColorSchemes]] +deps = ["ColorTypes", "ColorVectorSpace", "Colors", "FixedPointNumbers", "PrecompileTools", "Random"] +git-tree-sha1 = "b5278586822443594ff615963b0c09755771b3e0" +uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4" +version = "3.26.0" + +[[deps.ColorTypes]] +deps = ["FixedPointNumbers", "Random"] +git-tree-sha1 = "b10d0b65641d57b8b4d5e234446582de5047050d" +uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f" +version = "0.11.5" + +[[deps.ColorVectorSpace]] +deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "Requires", "Statistics", "TensorCore"] +git-tree-sha1 = "a1f44953f2382ebb937d60dafbe2deea4bd23249" +uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4" +version = "0.10.0" +weakdeps = ["SpecialFunctions"] + + [deps.ColorVectorSpace.extensions] + SpecialFunctionsExt = "SpecialFunctions" + +[[deps.Colors]] +deps = ["ColorTypes", "FixedPointNumbers", "Reexport"] +git-tree-sha1 = "362a287c3aa50601b0bc359053d5c2468f0e7ce0" +uuid = "5ae59095-9a9b-59fe-a467-6f913c188581" +version = "0.12.11" + +[[deps.CommonDataModel]] +deps = ["CFTime", "DataStructures", "Dates", "Preferences", "Printf", "Statistics"] +git-tree-sha1 = "d6fb5bf939a2753c74984b11434ea25d6c397a58" +uuid = "1fbeeb36-5f17-413c-809b-666fb144f157" +version = "0.3.6" + +[[deps.Compat]] +deps = ["TOML", "UUIDs"] +git-tree-sha1 = "8ae8d32e09f0dcf42a36b90d4e17f5dd2e4c4215" +uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" +version = "4.16.0" +weakdeps = ["Dates", "LinearAlgebra"] + + [deps.Compat.extensions] + CompatLinearAlgebraExt = "LinearAlgebra" + +[[deps.CompilerSupportLibraries_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" +version = "1.1.1+0" + +[[deps.CompositionsBase]] +git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad" +uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b" +version = "0.1.2" +weakdeps = ["InverseFunctions"] + + [deps.CompositionsBase.extensions] + CompositionsBaseInverseFunctionsExt = "InverseFunctions" + +[[deps.ConstructionBase]] +git-tree-sha1 = "76219f1ed5771adbb096743bff43fb5fdd4c1157" +uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" +version = "1.5.8" + + [deps.ConstructionBase.extensions] + ConstructionBaseIntervalSetsExt = "IntervalSets" + ConstructionBaseLinearAlgebraExt = "LinearAlgebra" + ConstructionBaseStaticArraysExt = "StaticArrays" + + [deps.ConstructionBase.weakdeps] + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" + LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + +[[deps.ContextVariablesX]] +deps = ["Compat", "Logging", "UUIDs"] +git-tree-sha1 = "25cc3803f1030ab855e383129dcd3dc294e322cc" +uuid = "6add18c4-b38d-439d-96f6-d6bc489c04c5" +version = "0.1.3" + +[[deps.Contour]] +git-tree-sha1 = "439e35b0b36e2e5881738abc8857bd92ad6ff9a8" +uuid = "d38c429a-6771-53c6-b99e-75d170b6e991" +version = "0.6.3" + +[[deps.Crayons]] +git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15" +uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f" +version = "4.1.1" + +[[deps.DataAPI]] +git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe" +uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" +version = "1.16.0" + +[[deps.DataFrames]] +deps = ["Compat", "DataAPI", "DataStructures", "Future", "InlineStrings", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrecompileTools", "PrettyTables", "Printf", "REPL", "Random", "Reexport", "SentinelArrays", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"] +git-tree-sha1 = "04c738083f29f86e62c8afc341f0967d8717bdb8" +uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +version = "1.6.1" + +[[deps.DataStructures]] +deps = ["Compat", "InteractiveUtils", "OrderedCollections"] +git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82" +uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" +version = "0.18.20" + +[[deps.DataValueInterfaces]] +git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6" +uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464" +version = "1.0.0" + +[[deps.Dates]] +deps = ["Printf"] +uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" + +[[deps.DefineSingletons]] +git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c" +uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52" +version = "0.1.2" + +[[deps.DiskArrays]] +deps = ["LRUCache", "OffsetArrays"] +git-tree-sha1 = "ef25c513cad08d7ebbed158c91768ae32f308336" +uuid = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3" +version = "0.3.23" + +[[deps.Distributed]] +deps = ["Random", "Serialization", "Sockets"] +uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" + +[[deps.Distributions]] +deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"] +git-tree-sha1 = "e6c693a0e4394f8fda0e51a5bdf5aef26f8235e9" +uuid = "31c24e10-a181-5473-b8eb-7969acd0382f" +version = "0.25.111" + + [deps.Distributions.extensions] + DistributionsChainRulesCoreExt = "ChainRulesCore" + DistributionsDensityInterfaceExt = "DensityInterface" + DistributionsTestExt = "Test" + + [deps.Distributions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" + Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[[deps.DocStringExtensions]] +deps = ["LibGit2"] +git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d" +uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +version = "0.9.3" + +[[deps.Downloads]] +deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] +uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +version = "1.6.0" + +[[deps.ExprTools]] +git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" +uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" +version = "0.1.10" + +[[deps.FFTW]] +deps = ["AbstractFFTs", "FFTW_jll", "LinearAlgebra", "MKL_jll", "Preferences", "Reexport"] +git-tree-sha1 = "4820348781ae578893311153d69049a93d05f39d" +uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +version = "1.8.0" + +[[deps.FFTW_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "c6033cc3892d0ef5bb9cd29b7f2f0331ea5184ea" +uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a" +version = "3.3.10+0" + +[[deps.FLoops]] +deps = ["BangBang", "Compat", "FLoopsBase", "InitialValues", "JuliaVariables", "MLStyle", "Serialization", "Setfield", "Transducers"] +git-tree-sha1 = "0a2e5873e9a5f54abb06418d57a8df689336a660" +uuid = "cc61a311-1640-44b5-9fba-1b764f453329" +version = "0.2.2" + +[[deps.FLoopsBase]] +deps = ["ContextVariablesX"] +git-tree-sha1 = "656f7a6859be8673bf1f35da5670246b923964f7" +uuid = "b9860ae5-e623-471e-878b-f6a53c775ea6" +version = "0.1.1" + +[[deps.FastGaussQuadrature]] +deps = ["LinearAlgebra", "SpecialFunctions", "StaticArrays"] +git-tree-sha1 = "fd923962364b645f3719855c88f7074413a6ad92" +uuid = "442a2c76-b920-505d-bb47-c5924d526838" +version = "1.0.2" + +[[deps.FileIO]] +deps = ["Pkg", "Requires", "UUIDs"] +git-tree-sha1 = "82d8afa92ecf4b52d78d869f038ebfb881267322" +uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" +version = "1.16.3" + +[[deps.FileWatching]] +uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" + +[[deps.FillArrays]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "6a70198746448456524cb442b8af316927ff3e1a" +uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" +version = "1.13.0" +weakdeps = ["PDMats", "SparseArrays", "Statistics"] + + [deps.FillArrays.extensions] + FillArraysPDMatsExt = "PDMats" + FillArraysSparseArraysExt = "SparseArrays" + FillArraysStatisticsExt = "Statistics" + +[[deps.FixedPointNumbers]] +deps = ["Statistics"] +git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172" +uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" +version = "0.8.5" + +[[deps.Future]] +deps = ["Random"] +uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" + +[[deps.GMP_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "781609d7-10c4-51f6-84f2-b8444358ff6d" +version = "6.2.1+6" + +[[deps.GPUArrays]] +deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"] +git-tree-sha1 = "62ee71528cca49be797076a76bdc654a170a523e" +uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +version = "10.3.1" + +[[deps.GPUArraysCore]] +deps = ["Adapt"] +git-tree-sha1 = "ec632f177c0d990e64d955ccc1b8c04c485a0950" +uuid = "46192b85-c4d5-4398-a991-12ede77f4527" +version = "0.1.6" + +[[deps.GPUCompiler]] +deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Preferences", "Scratch", "Serialization", "TOML", "TimerOutputs", "UUIDs"] +git-tree-sha1 = "ab29216184312f99ff957b32cd63c2fe9c928b91" +uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" +version = "0.26.7" + +[[deps.GenericFFT]] +deps = ["AbstractFFTs", "FFTW", "LinearAlgebra", "Reexport"] +git-tree-sha1 = "1bc01f2ea9a0226a60723794ff86b8017739f5d9" +uuid = "a8297547-1b15-4a5a-a998-a2ac5f1cef28" +version = "0.1.6" + +[[deps.GnuTLS_jll]] +deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Nettle_jll", "P11Kit_jll", "Zlib_jll"] +git-tree-sha1 = "383db7d3f900f4c1f47a8a04115b053c095e48d3" +uuid = "0951126a-58fd-58f1-b5b3-b08c7c4a876d" +version = "3.8.4+0" + +[[deps.HDF5]] +deps = ["Compat", "HDF5_jll", "Libdl", "MPIPreferences", "Mmap", "Preferences", "Printf", "Random", "Requires", "UUIDs"] +git-tree-sha1 = "e856eef26cf5bf2b0f95f8f4fc37553c72c8641c" +uuid = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" +version = "0.17.2" +weakdeps = ["MPI"] + + [deps.HDF5.extensions] + MPIExt = "MPI" + +[[deps.HDF5_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LLVMOpenMP_jll", "LazyArtifacts", "LibCURL_jll", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "OpenSSL_jll", "TOML", "Zlib_jll", "libaec_jll"] +git-tree-sha1 = "38c8874692d48d5440d5752d6c74b0c6b0b60739" +uuid = "0234f1f7-429e-5d53-9886-15a909be8d59" +version = "1.14.2+1" + +[[deps.Hwloc_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "5e19e1e4fa3e71b774ce746274364aef0234634e" +uuid = "e33a78d0-f292-5ffc-b300-72abe9b543c8" +version = "2.11.1+0" + +[[deps.HypergeometricFunctions]] +deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"] +git-tree-sha1 = "7c4195be1649ae622304031ed46a2f4df989f1eb" +uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a" +version = "0.3.24" + +[[deps.InitialValues]] +git-tree-sha1 = "4da0f88e9a39111c2fa3add390ab15f3a44f3ca3" +uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c" +version = "0.3.1" + +[[deps.InlineStrings]] +git-tree-sha1 = "45521d31238e87ee9f9732561bfee12d4eebd52d" +uuid = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48" +version = "1.4.2" + + [deps.InlineStrings.extensions] + ArrowTypesExt = "ArrowTypes" + ParsersExt = "Parsers" + + [deps.InlineStrings.weakdeps] + ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" + Parsers = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" + +[[deps.IntegerMathUtils]] +git-tree-sha1 = "b8ffb903da9f7b8cf695a8bead8e01814aa24b30" +uuid = "18e54dd8-cb9d-406c-a71d-865a43cbb235" +version = "0.1.2" + +[[deps.IntelOpenMP_jll]] +deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"] +git-tree-sha1 = "10bd689145d2c3b2a9844005d01087cc1194e79e" +uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0" +version = "2024.2.1+0" + +[[deps.InteractiveUtils]] +deps = ["Markdown"] +uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" + +[[deps.InverseFunctions]] +git-tree-sha1 = "2787db24f4e03daf859c6509ff87764e4182f7d1" +uuid = "3587e190-3f89-42d0-90ee-14403ec27112" +version = "0.1.16" +weakdeps = ["Dates", "Test"] + + [deps.InverseFunctions.extensions] + InverseFunctionsDatesExt = "Dates" + InverseFunctionsTestExt = "Test" + +[[deps.InvertedIndices]] +git-tree-sha1 = "0dc7b50b8d436461be01300fd8cd45aa0274b038" +uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" +version = "1.3.0" + +[[deps.IrrationalConstants]] +git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" +uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" +version = "0.2.2" + +[[deps.IteratorInterfaceExtensions]] +git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856" +uuid = "82899510-4779-5014-852e-03e436cf321d" +version = "1.0.0" + +[[deps.JLD2]] +deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"] +git-tree-sha1 = "a0746c21bdc986d0dc293efa6b1faee112c37c28" +uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +version = "0.4.53" + +[[deps.JLLWrappers]] +deps = ["Artifacts", "Preferences"] +git-tree-sha1 = "f389674c99bfcde17dc57454011aa44d5a260a40" +uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" +version = "1.6.0" + +[[deps.JuliaNVTXCallbacks_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "af433a10f3942e882d3c671aacb203e006a5808f" +uuid = "9c1d0b0a-7046-5b2e-a33f-ea22f176ac7e" +version = "0.2.1+0" + +[[deps.JuliaVariables]] +deps = ["MLStyle", "NameResolution"] +git-tree-sha1 = "49fb3cb53362ddadb4415e9b73926d6b40709e70" +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 = "cb1cff88ef2f3a157cbad75bbe6b229e1975e498" +uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +version = "0.9.25" + + [deps.KernelAbstractions.extensions] + EnzymeExt = "EnzymeCore" + LinearAlgebraExt = "LinearAlgebra" + SparseArraysExt = "SparseArrays" + + [deps.KernelAbstractions.weakdeps] + EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" + LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" + +[[deps.LLVM]] +deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"] +git-tree-sha1 = "2470e69781ddd70b8878491233cd09bc1bd7fc96" +uuid = "929cbde3-209d-540e-8aea-75f648917ca0" +version = "8.1.0" +weakdeps = ["BFloat16s"] + + [deps.LLVM.extensions] + BFloat16sExt = "BFloat16s" + +[[deps.LLVMExtra_jll]] +deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] +git-tree-sha1 = "597d1c758c9ae5d985ba4202386a607c675ee700" +uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" +version = "0.0.31+0" + +[[deps.LLVMLoopInfo]] +git-tree-sha1 = "2e5c102cfc41f48ae4740c7eca7743cc7e7b75ea" +uuid = "8b046642-f1f6-4319-8d3c-209ddc03c586" +version = "1.0.0" + +[[deps.LLVMOpenMP_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "78211fb6cbc872f77cad3fc0b6cf647d923f4929" +uuid = "1d63c593-3942-5779-bab2-d838dc0a180e" +version = "18.1.7+0" + +[[deps.LRUCache]] +git-tree-sha1 = "b3cc6698599b10e652832c2f23db3cab99d51b59" +uuid = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637" +version = "1.6.1" +weakdeps = ["Serialization"] + + [deps.LRUCache.extensions] + SerializationExt = ["Serialization"] + +[[deps.LaTeXStrings]] +git-tree-sha1 = "50901ebc375ed41dbf8058da26f9de442febbbec" +uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" +version = "1.3.1" + +[[deps.LazyArtifacts]] +deps = ["Artifacts", "Pkg"] +uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3" + +[[deps.LibCURL]] +deps = ["LibCURL_jll", "MozillaCACerts_jll"] +uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" +version = "0.6.4" + +[[deps.LibCURL_jll]] +deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"] +uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" +version = "8.4.0+0" + +[[deps.LibGit2]] +deps = ["Base64", "LibGit2_jll", "NetworkOptions", "Printf", "SHA"] +uuid = "76f85450-5226-5b5a-8eaa-529ad045b433" + +[[deps.LibGit2_jll]] +deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll"] +uuid = "e37daf67-58a4-590a-8e99-b0245dd2ffc5" +version = "1.6.4+0" + +[[deps.LibSSH2_jll]] +deps = ["Artifacts", "Libdl", "MbedTLS_jll"] +uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8" +version = "1.11.0+1" + +[[deps.Libdl]] +uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb" + +[[deps.Libiconv_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "f9557a255370125b405568f9767d6d195822a175" +uuid = "94ce4f54-9a6c-5748-9c1c-f9c7231a4531" +version = "1.17.0+0" + +[[deps.LinearAlgebra]] +deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] +uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + +[[deps.LogExpFunctions]] +deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] +git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea" +uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +version = "0.3.28" + + [deps.LogExpFunctions.extensions] + LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" + LogExpFunctionsChangesOfVariablesExt = "ChangesOfVariables" + LogExpFunctionsInverseFunctionsExt = "InverseFunctions" + + [deps.LogExpFunctions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" + InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" + +[[deps.Logging]] +uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" + +[[deps.Lz4_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "7f26c8fc5229e68484e0b3447312c98e16207d11" +uuid = "5ced341a-0733-55b8-9ab6-a4889d929147" +version = "1.10.0+0" + +[[deps.MKL_jll]] +deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"] +git-tree-sha1 = "f046ccd0c6db2832a9f639e2c669c6fe867e5f4f" +uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7" +version = "2024.2.0+0" + +[[deps.MLStyle]] +git-tree-sha1 = "bc38dff0548128765760c79eb7388a4b37fae2c8" +uuid = "d8e11817-5142-5d16-987a-aa16d5891078" +version = "0.4.17" + +[[deps.MPI]] +deps = ["Distributed", "DocStringExtensions", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "PkgVersion", "PrecompileTools", "Requires", "Serialization", "Sockets"] +git-tree-sha1 = "b4d8707e42b693720b54f0b3434abee6dd4d947a" +uuid = "da04e1cc-30fd-572f-bb4f-1f8673147195" +version = "0.20.16" + + [deps.MPI.extensions] + AMDGPUExt = "AMDGPU" + CUDAExt = "CUDA" + + [deps.MPI.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + +[[deps.MPICH_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"] +git-tree-sha1 = "19d4bd098928a3263693991500d05d74dbdc2004" +uuid = "7cb0a576-ebde-5e09-9194-50597f1243b4" +version = "4.2.2+0" + +[[deps.MPIPreferences]] +deps = ["Libdl", "Preferences"] +git-tree-sha1 = "c105fe467859e7f6e9a852cb15cb4301126fac07" +uuid = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" +version = "0.1.11" + +[[deps.MPItrampoline_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"] +git-tree-sha1 = "8c35d5420193841b2f367e658540e8d9e0601ed0" +uuid = "f1f71cc9-e9ae-5b93-9b94-4fe0e1ad3748" +version = "5.4.0+0" + +[[deps.MacroTools]] +deps = ["Markdown", "Random"] +git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df" +uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" +version = "0.5.13" + +[[deps.MarchingCubes]] +deps = ["PrecompileTools", "StaticArrays"] +git-tree-sha1 = "27d162f37cc29de047b527dab11a826dd3a650ad" +uuid = "299715c1-40a9-479a-aaf9-4a633d36f717" +version = "0.1.9" + +[[deps.Markdown]] +deps = ["Base64"] +uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" + +[[deps.MbedTLS_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1" +version = "2.28.2+1" + +[[deps.MicroCollections]] +deps = ["Accessors", "BangBang", "InitialValues"] +git-tree-sha1 = "44d32db644e84c75dab479f1bc15ee76a1a3618f" +uuid = "128add7d-3638-4c79-886c-908ea0c25c34" +version = "0.2.0" + +[[deps.MicrosoftMPI_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "f12a29c4400ba812841c6ace3f4efbb6dbb3ba01" +uuid = "9237b28f-5490-5468-be7b-bb81f5f5e6cf" +version = "10.1.4+2" + +[[deps.Missings]] +deps = ["DataAPI"] +git-tree-sha1 = "ec4f7fbeab05d7747bdf98eb74d130a2a2ed298d" +uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28" +version = "1.2.0" + +[[deps.Mmap]] +uuid = "a63ad114-7e13-5084-954f-fe012c677804" + +[[deps.MozillaCACerts_jll]] +uuid = "14a3606d-f60d-562e-9121-12d972cd8159" +version = "2023.1.10" + +[[deps.NCDatasets]] +deps = ["CFTime", "CommonDataModel", "DataStructures", "Dates", "DiskArrays", "NetCDF_jll", "NetworkOptions", "Printf"] +git-tree-sha1 = "77df6d3708ec0eb3441551e1f20f7503b37c2393" +uuid = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" +version = "0.14.5" + +[[deps.NVTX]] +deps = ["Colors", "JuliaNVTXCallbacks_jll", "Libdl", "NVTX_jll"] +git-tree-sha1 = "53046f0483375e3ed78e49190f1154fa0a4083a1" +uuid = "5da4648a-3479-48b8-97b9-01cb529c0a1f" +version = "0.3.4" + +[[deps.NVTX_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "ce3269ed42816bf18d500c9f63418d4b0d9f5a3b" +uuid = "e98f9f5b-d649-5603-91fd-7774390e6439" +version = "3.1.0+2" + +[[deps.NaNMath]] +deps = ["OpenLibm_jll"] +git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" +uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +version = "1.0.2" + +[[deps.NameResolution]] +deps = ["PrettyPrint"] +git-tree-sha1 = "1a0fa0e9613f46c9b8c11eee38ebb4f590013c5e" +uuid = "71a1bf82-56d0-4bbc-8a3c-48b961074391" +version = "0.1.5" + +[[deps.NetCDF_jll]] +deps = ["Artifacts", "Blosc_jll", "Bzip2_jll", "HDF5_jll", "JLLWrappers", "LibCURL_jll", "Libdl", "OpenMPI_jll", "XML2_jll", "Zlib_jll", "Zstd_jll", "libzip_jll"] +git-tree-sha1 = "a8af1798e4eb9ff768ce7fdefc0e957097793f15" +uuid = "7243133f-43d8-5620-bbf4-c2c921802cf3" +version = "400.902.209+0" + +[[deps.Nettle_jll]] +deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "eca63e3847dad608cfa6a3329b95ef674c7160b4" +uuid = "4c82536e-c426-54e4-b420-14f461c4ed8b" +version = "3.7.2+0" + +[[deps.NetworkOptions]] +uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" +version = "1.2.0" + +[[deps.OffsetArrays]] +git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e" +uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" +version = "1.14.1" +weakdeps = ["Adapt"] + + [deps.OffsetArrays.extensions] + OffsetArraysAdaptExt = "Adapt" + +[[deps.OpenBLAS_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] +uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" +version = "0.3.23+4" + +[[deps.OpenLibm_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "05823500-19ac-5b8b-9628-191a04bc5112" +version = "0.8.1+2" + +[[deps.OpenMPI_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"] +git-tree-sha1 = "bfce6d523861a6c562721b262c0d1aaeead2647f" +uuid = "fe0851c0-eecd-5654-98d4-656369965a5c" +version = "5.0.5+0" + +[[deps.OpenSSL_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "1b35263570443fdd9e76c76b7062116e2f374ab8" +uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95" +version = "3.0.15+0" + +[[deps.OpenSpecFun_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1" +uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" +version = "0.5.5+0" + +[[deps.OrderedCollections]] +git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5" +uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" +version = "1.6.3" + +[[deps.P11Kit_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "2cd396108e178f3ae8dedbd8e938a18726ab2fbf" +uuid = "c2071276-7c44-58a7-b746-946036e04d0a" +version = "0.24.1+0" + +[[deps.PDMats]] +deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"] +git-tree-sha1 = "949347156c25054de2db3b166c52ac4728cbad65" +uuid = "90014a1f-27ba-587c-ab20-58faa44d9150" +version = "0.11.31" + +[[deps.ParticleDA]] +deps = ["ChunkSplitters", "HDF5", "LinearAlgebra", "MPI", "PDMats", "Random", "Statistics", "StructArrays", "TimerOutputs", "YAML"] +git-tree-sha1 = "445106e67dc21688691ed762446d393965e1d982" +uuid = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" +version = "1.1.0" + +[[deps.Pkg]] +deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] +uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +version = "1.10.0" + +[[deps.PkgVersion]] +deps = ["Pkg"] +git-tree-sha1 = "f9501cc0430a26bc3d156ae1b5b0c1b47af4d6da" +uuid = "eebad327-c553-4316-9ea0-9fa01ccd7688" +version = "0.3.3" + +[[deps.PooledArrays]] +deps = ["DataAPI", "Future"] +git-tree-sha1 = "36d8b4b899628fb92c2749eb488d884a926614d3" +uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" +version = "1.4.3" + +[[deps.PrecompileTools]] +deps = ["Preferences"] +git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f" +uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +version = "1.2.1" + +[[deps.Preferences]] +deps = ["TOML"] +git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6" +uuid = "21216c6a-2e73-6563-6e65-726566657250" +version = "1.4.3" + +[[deps.PrettyPrint]] +git-tree-sha1 = "632eb4abab3449ab30c5e1afaa874f0b98b586e4" +uuid = "8162dcfd-2161-5ef2-ae6c-7681170c5f98" +version = "0.2.0" + +[[deps.PrettyTables]] +deps = ["Crayons", "LaTeXStrings", "Markdown", "PrecompileTools", "Printf", "Reexport", "StringManipulation", "Tables"] +git-tree-sha1 = "66b20dd35966a748321d3b2537c4584cf40387c7" +uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" +version = "2.3.2" + +[[deps.Primes]] +deps = ["IntegerMathUtils"] +git-tree-sha1 = "cb420f77dc474d23ee47ca8d14c90810cafe69e7" +uuid = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae" +version = "0.5.6" + +[[deps.Printf]] +deps = ["Unicode"] +uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" + +[[deps.ProgressMeter]] +deps = ["Distributed", "Printf"] +git-tree-sha1 = "8f6bc219586aef8baf0ff9a5fe16ee9c70cb65e4" +uuid = "92933f4c-e287-5a05-a399-4b506db050ca" +version = "1.10.2" + +[[deps.PtrArrays]] +git-tree-sha1 = "77a42d78b6a92df47ab37e177b2deac405e1c88f" +uuid = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d" +version = "1.2.1" + +[[deps.QuadGK]] +deps = ["DataStructures", "LinearAlgebra"] +git-tree-sha1 = "1d587203cf851a51bf1ea31ad7ff89eff8d625ea" +uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" +version = "2.11.0" + + [deps.QuadGK.extensions] + QuadGKEnzymeExt = "Enzyme" + + [deps.QuadGK.weakdeps] + Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" + +[[deps.REPL]] +deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] +uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" + +[[deps.Random]] +deps = ["SHA"] +uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" + +[[deps.Random123]] +deps = ["Random", "RandomNumbers"] +git-tree-sha1 = "4743b43e5a9c4a2ede372de7061eed81795b12e7" +uuid = "74087812-796a-5b5d-8853-05524746bad3" +version = "1.7.0" + +[[deps.RandomNumbers]] +deps = ["Random"] +git-tree-sha1 = "c6ec94d2aaba1ab2ff983052cf6a606ca5985902" +uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143" +version = "1.6.0" + +[[deps.Reexport]] +git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b" +uuid = "189a3867-3050-52da-a836-e630ba90ab69" +version = "1.2.2" + +[[deps.Requires]] +deps = ["UUIDs"] +git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7" +uuid = "ae029012-a4dd-5104-9daa-d747884805df" +version = "1.3.0" + +[[deps.Rmath]] +deps = ["Random", "Rmath_jll"] +git-tree-sha1 = "f65dcb5fa46aee0cf9ed6274ccbd597adc49aa7b" +uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa" +version = "0.7.1" + +[[deps.Rmath_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "e60724fd3beea548353984dc61c943ecddb0e29a" +uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f" +version = "0.4.3+0" + +[[deps.SHA]] +uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" +version = "0.7.0" + +[[deps.Scratch]] +deps = ["Dates"] +git-tree-sha1 = "3bac05bc7e74a75fd9cba4295cde4045d9fe2386" +uuid = "6c6a2e73-6563-6170-7368-637461726353" +version = "1.2.1" + +[[deps.SentinelArrays]] +deps = ["Dates", "Random"] +git-tree-sha1 = "ff11acffdb082493657550959d4feb4b6149e73a" +uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c" +version = "1.4.5" + +[[deps.Serialization]] +uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" + +[[deps.Setfield]] +deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"] +git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac" +uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46" +version = "1.1.1" + +[[deps.Sockets]] +uuid = "6462fe0b-24de-5631-8697-dd941f90decc" + +[[deps.SortingAlgorithms]] +deps = ["DataStructures"] +git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085" +uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c" +version = "1.2.1" + +[[deps.SparseArrays]] +deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] +uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +version = "1.10.0" + +[[deps.SpecialFunctions]] +deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] +git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14" +uuid = "276daf66-3868-5448-9aa4-cd146d93841b" +version = "2.4.0" + + [deps.SpecialFunctions.extensions] + SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" + + [deps.SpecialFunctions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + +[[deps.SpeedyWeather]] +deps = ["AbstractFFTs", "Adapt", "AssociatedLegendrePolynomials", "BitInformation", "CUDA", "CodecZlib", "Dates", "DocStringExtensions", "FFTW", "FLoops", "FastGaussQuadrature", "GPUArrays", "GenericFFT", "JLD2", "KernelAbstractions", "LinearAlgebra", "NCDatasets", "Primes", "Printf", "ProgressMeter", "Random", "Statistics", "TOML", "UnicodePlots"] +git-tree-sha1 = "7af375091ba736521edae7d839a085354a3dad4e" +uuid = "9e226e20-d153-4fed-8a5b-493def4f21a9" +version = "0.10.0" + + [deps.SpeedyWeather.extensions] + SpeedyWeatherJLArraysExt = "JLArrays" + SpeedyWeatherMakieExt = "Makie" + + [deps.SpeedyWeather.weakdeps] + JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" + Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" + +[[deps.SplittablesBase]] +deps = ["Setfield", "Test"] +git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5" +uuid = "171d559e-b47b-412a-8079-5efa626c420e" +version = "0.1.15" + +[[deps.StaticArrays]] +deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] +git-tree-sha1 = "eeafab08ae20c62c44c8399ccb9354a04b80db50" +uuid = "90137ffa-7385-5640-81b9-e52037218182" +version = "1.9.7" + + [deps.StaticArrays.extensions] + StaticArraysChainRulesCoreExt = "ChainRulesCore" + StaticArraysStatisticsExt = "Statistics" + + [deps.StaticArrays.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" + +[[deps.StaticArraysCore]] +git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682" +uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +version = "1.4.3" + +[[deps.Statistics]] +deps = ["LinearAlgebra", "SparseArrays"] +uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +version = "1.10.0" + +[[deps.StatsAPI]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed" +uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0" +version = "1.7.0" + +[[deps.StatsBase]] +deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"] +git-tree-sha1 = "5cf7606d6cef84b543b483848d4ae08ad9832b21" +uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +version = "0.34.3" + +[[deps.StatsFuns]] +deps = ["HypergeometricFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"] +git-tree-sha1 = "cef0472124fab0695b58ca35a77c6fb942fdab8a" +uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c" +version = "1.3.1" + + [deps.StatsFuns.extensions] + StatsFunsChainRulesCoreExt = "ChainRulesCore" + StatsFunsInverseFunctionsExt = "InverseFunctions" + + [deps.StatsFuns.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" + +[[deps.StringEncodings]] +deps = ["Libiconv_jll"] +git-tree-sha1 = "b765e46ba27ecf6b44faf70df40c57aa3a547dcb" +uuid = "69024149-9ee7-55f6-a4c4-859efe599b68" +version = "0.3.7" + +[[deps.StringManipulation]] +deps = ["PrecompileTools"] +git-tree-sha1 = "a04cabe79c5f01f4d723cc6704070ada0b9d46d5" +uuid = "892a3eda-7b42-436c-8928-eab12a02cf0e" +version = "0.3.4" + +[[deps.StructArrays]] +deps = ["ConstructionBase", "DataAPI", "Tables"] +git-tree-sha1 = "f4dc295e983502292c4c3f951dbb4e985e35b3be" +uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" +version = "0.6.18" +weakdeps = ["Adapt", "GPUArraysCore", "SparseArrays", "StaticArrays"] + + [deps.StructArrays.extensions] + StructArraysAdaptExt = "Adapt" + StructArraysGPUArraysCoreExt = "GPUArraysCore" + StructArraysSparseArraysExt = "SparseArrays" + StructArraysStaticArraysExt = "StaticArrays" + +[[deps.SuiteSparse]] +deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] +uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" + +[[deps.SuiteSparse_jll]] +deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] +uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" +version = "7.2.1+1" + +[[deps.TOML]] +deps = ["Dates"] +uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76" +version = "1.0.3" + +[[deps.TableTraits]] +deps = ["IteratorInterfaceExtensions"] +git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39" +uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c" +version = "1.0.1" + +[[deps.Tables]] +deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"] +git-tree-sha1 = "598cd7c1f68d1e205689b1c2fe65a9f85846f297" +uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" +version = "1.12.0" + +[[deps.Tar]] +deps = ["ArgTools", "SHA"] +uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" +version = "1.10.0" + +[[deps.TensorCore]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6" +uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50" +version = "0.1.1" + +[[deps.Test]] +deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] +uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[[deps.TestItems]] +git-tree-sha1 = "42fd9023fef18b9b78c8343a4e2f3813ffbcefcb" +uuid = "1c621080-faea-4a02-84b6-bbd5e436b8fe" +version = "1.0.0" + +[[deps.TimerOutputs]] +deps = ["ExprTools", "Printf"] +git-tree-sha1 = "5a13ae8a41237cff5ecf34f73eb1b8f42fff6531" +uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" +version = "0.5.24" + +[[deps.TranscodingStreams]] +git-tree-sha1 = "e84b3a11b9bece70d14cce63406bbc79ed3464d2" +uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" +version = "0.11.2" + +[[deps.Transducers]] +deps = ["Accessors", "Adapt", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "ConstructionBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "SplittablesBase", "Tables"] +git-tree-sha1 = "5215a069867476fc8e3469602006b9670e68da23" +uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999" +version = "0.4.82" + + [deps.Transducers.extensions] + TransducersBlockArraysExt = "BlockArrays" + TransducersDataFramesExt = "DataFrames" + TransducersLazyArraysExt = "LazyArrays" + TransducersOnlineStatsBaseExt = "OnlineStatsBase" + TransducersReferenceablesExt = "Referenceables" + + [deps.Transducers.weakdeps] + BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" + DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" + LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" + OnlineStatsBase = "925886fa-5bf2-5e8e-b522-a9147a512338" + Referenceables = "42d2dcc6-99eb-4e98-b66c-637b7d73030e" + +[[deps.UUIDs]] +deps = ["Random", "SHA"] +uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" + +[[deps.Unicode]] +uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" + +[[deps.UnicodePlots]] +deps = ["ColorSchemes", "ColorTypes", "Contour", "Crayons", "Dates", "LinearAlgebra", "MarchingCubes", "NaNMath", "PrecompileTools", "Printf", "Requires", "SparseArrays", "StaticArrays", "StatsBase"] +git-tree-sha1 = "30646456e889c18fb3c23e58b2fc5da23644f752" +uuid = "b8865327-cd53-5732-bb35-84acbb429228" +version = "3.6.4" + + [deps.UnicodePlots.extensions] + FreeTypeExt = ["FileIO", "FreeType"] + ImageInTerminalExt = "ImageInTerminal" + IntervalSetsExt = "IntervalSets" + TermExt = "Term" + UnitfulExt = "Unitful" + + [deps.UnicodePlots.weakdeps] + FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" + FreeType = "b38be410-82b0-50bf-ab77-7b57e271db43" + ImageInTerminal = "d8c32880-2388-543b-8c61-d9f865259254" + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" + Term = "22787eb5-b846-44ae-b979-8e399b8463ab" + Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" + +[[deps.UnsafeAtomics]] +git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278" +uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f" +version = "0.2.1" + +[[deps.UnsafeAtomicsLLVM]] +deps = ["LLVM", "UnsafeAtomics"] +git-tree-sha1 = "2d17fabcd17e67d7625ce9c531fb9f40b7c42ce4" +uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249" +version = "0.2.1" + +[[deps.XML2_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"] +git-tree-sha1 = "1165b0443d0eca63ac1e32b8c0eb69ed2f4f8127" +uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a" +version = "2.13.3+0" + +[[deps.XZ_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "ac88fb95ae6447c8dda6a5503f3bafd496ae8632" +uuid = "ffd25f8a-64ca-5728-b0f7-c24cf3aae800" +version = "5.4.6+0" + +[[deps.YAML]] +deps = ["Base64", "Dates", "Printf", "StringEncodings"] +git-tree-sha1 = "dea63ff72079443240fbd013ba006bcbc8a9ac00" +uuid = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" +version = "0.4.12" + +[[deps.Zlib_jll]] +deps = ["Libdl"] +uuid = "83775a58-1f1d-513f-b197-d71354ab007a" +version = "1.2.13+1" + +[[deps.Zstd_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "e678132f07ddb5bfa46857f0d7620fb9be675d3b" +uuid = "3161d3a3-bdf6-5164-811a-617609db77b4" +version = "1.5.6+0" + +[[deps.libaec_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "46bf7be2917b59b761247be3f317ddf75e50e997" +uuid = "477f73a3-ac25-53e9-8cc3-50b2fa2566f0" +version = "1.1.2+0" + +[[deps.libblastrampoline_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" +version = "5.8.0+1" + +[[deps.libzip_jll]] +deps = ["Artifacts", "Bzip2_jll", "GnuTLS_jll", "JLLWrappers", "Libdl", "XZ_jll", "Zlib_jll", "Zstd_jll"] +git-tree-sha1 = "3282b7d16ae7ac3e57ec2f3fa8fafb564d8f9f7f" +uuid = "337d8026-41b4-5cde-a456-74a10e5b31d1" +version = "1.10.1+0" + +[[deps.nghttp2_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d" +version = "1.52.0+1" + +[[deps.oneTBB_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "7d0ea0f4895ef2f5cb83645fa689e52cb55cf493" +uuid = "1317d2d5-d96f-522e-a858-c73665f53c3e" +version = "2021.12.0+0" + +[[deps.p7zip_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" +version = "17.4.0+2" diff --git a/SpeedyWeather/Project.toml b/SpeedyWeather/Project.toml new file mode 100644 index 0000000..5dad286 --- /dev/null +++ b/SpeedyWeather/Project.toml @@ -0,0 +1,3 @@ +[deps] +ParticleDA = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl new file mode 100644 index 0000000..d6ee645 --- /dev/null +++ b/SpeedyWeather/src/speedy.jl @@ -0,0 +1,202 @@ +using ParticleDA +using SpeedyWeather + + +module SpeedyWeatherModel + +LAYERED_VARIABLES = (:vor, :div, :temp, :humid) +SURFACE_VARIABLES = (:pres) + +Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:ModelSetup} + spectral_truncation::Int = 31 + n_level::Int = 8 + n_days::Float64 = 0.25 + model_type::Type{M} = PrimitiveWetModel +end + +struct SpeedyModel{ + T<:AbstractFloat, + G<:AbstractSpectralGrid, + M<:ModelSetup, +} + parameters::SpeedyParameters{T, M} + spectral_grid::G + model_setup::M + prognostic_variables::PrognosticVariables + diagnostic_variables::DiagnosticVariables + n_layered_variables::Int + n_surface_variables::Int +end + + +function init(parameters::SpeedyParameters{T, M}) where {T<:AbstractFloat, M<:ModelSetup} + spectral_grid = SpectralGrid( + NF=T, trunc=parameters.spectral_truncation, nlev=parameters.n_level + ) + model = M(; spectral_grid) + simulation = initialize!(model) + (; prognostic_variables, diagnostic_variables) = simulation + n_layered_variables = count( + SpeedyWeather.has(model, var) for var in LAYERED_VARIABLES + ) + n_surface_variables = count( + SpeedyWeather.has(model, var) for var in SURFACE_VARIABLES + ) + return SpeedyModel( + parameters, + spectral_grid, + model, + prognostic_variables, + diagnostic_variables, + n_layered_variables, + n_surface_variables, + ) +end + +function ParticleDA.get_state_dimension(model::SpeedyModel) + (model.parameters.spectral_truncation + 1)^2 * ( + model.parameters.n_level * model.n_layered_variables + model.n_surface_variables + ) +end + +function ParticleDA.get_observation_dimension(model::SpeedyModel) + +end + +function update_spectral_coefficients_from_vector!( + spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, + vector::AbstractVector{T} +) where {T <: AbstractFloat} + m, n = size(spectral_coefficients) + spectral_coefficients[1:m - 1, 1] = vector[1:m - 1] + j = m - 1 + for i in 2:n + spectral_coefficients[i:m - 1, i] = reinterpret( + Complex{T}, vector[j+1:j+(m - i) * 2] + ) + j = j + (m - i) * 2 + end +end + +function update_prognostic_variables_from_state_vector!( + prognostic_variables::PrognosticVariables{T}, + state::AbstractVector{T}, +) where {T <: AbstractFloat} + start_index = 1 + dim_spectral = (prognostic_variables.trunc + 1)^2 + for name in LAYERED_VARIABLES + if SpeedyWeather.has(prognostic_variables, name) + layer_spectral_coefficients = SpeedyWeather.get_var( + prognostic_variables, name; lf=1 + ) + for spectral_coefficients in layer_spectral_coefficients + end_index = start_index + dim_spectral - 1 + update_spectral_coefficients_from_vector!( + spectral_coefficients, + state[start_index:end_index] + ) + start_index = end_index + 1 + end + end + end + if SpeedyWeather.has(prognostic_variables, :pres) + update_spectral_coefficients_from_vector!( + SpeedyWeather.get_pressure(prognostic_variables; lf=1), + state[start_index:end] + ) + end +end + +function update_vector_from_spectral_coefficients!( + vector::AbstractVector{T}, + spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, +) where {T <: AbstractFloat} + m, n = size(spectral_coefficients) + vector[1:m - 1] = spectral_coefficients[1:m - 1, 1] + j = m - 1 + for i in 2:n + vector[j+1:j+(m - i) * 2] = reinterpret(T, spectral_coefficients[i:m - 1, i]) + j = j + (m - i) * 2 + end +end + +function update_state_vector_from_prognostic_variables!( + state::AbstractVector{T}, + prognostic_variables::PrognosticVariables{T} +) where {T <: AbstractFloat} + start_index = 1 + dim_spectral = (prognostic_variables.lmax + 1) * (prognostic_variables.mmax + 1) + for name in (:vor, :div, :temp, :humid) + if SpeedyWeather.has(prognostic_variables, name) + layer_spectral_coefficients = SpeedyWeather.get_var( + prognostic_variables, name; lf=1 + ) + for spectral_coefficients in layer_spectral_coefficients + end_index = start_index + dim_spectral - 1 + update_vector_from_spectral_coefficients!( + state[start_index:end_index], + spectral_coefficients + ) + start_index = end_index + 1 + end + end + end + if SpeedyWeather.has(prognostic_variables, :pres) + update_vector_from_spectral_coefficients!( + state[start_index:end], + SpeedyWeather.get_pressure(prognostic_variables; lf=1) + ) + end +end + +ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T + +ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T + +function ParticleDA.sample_initial_state!( + state::V{T}, model::SpeedyModel{T}, rng::G +) where {T<:AbstractFloat, V<:AbstractVector, G<:AbstractRNG} + initial_conditions = model.model.initial_conditions + SpeedyWeather.initialize!( + model.prognostic_variables, initial_conditions, model.model_setup + ) + update_state_vector_from_prognostic_variables!(state, model.prognostic_variables) +end + +function ParticleDA.update_state_deterministic!( + state::V{T}, model::SpeedyModel{T}, time_index::Int +) where {T<:AbstractFloat, V<:AbstractVector} + update_prognostic_variables_from_state_vector!(model.prognostic_variables, state) + (; clock) = model.prognostic_variables + (; time_stepping) = model.model_setup + SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameter.n_days)) + SpeedyWeather.initialize!(clock, time_stepping) + SpeedyWeather.time_stepping!( + model.prognostic_variables, model.diagnostic_variables, model.model_setup + ) + clock.time += clock.n_timesteps * clock.Δt * (time_index - 1) + clock.timestep_counter += clock.n_timesteps * (time_index - 1) + update_state_vector_from_prognostic_variables!(state, model.prognostic_variables) +end + +function ParticleDA.update_state_stochastic!( + state::V{T}, model::SpeedyModel{T}, rng::G +) where {T<:AbstractFloat, V<:AbstractVector, G<:AbstractRNG} + +end + +function ParticleDA.sample_observation_given_state!( + observation::V{T}, state::V{T}, model::SpeedyModel{T}, rng::G +) where {T<:AbstractFloat, V<:AbstractVector, G<:AbstractRNG} + # TODO: figure out how to update diagnostic variables given prognostic variables + update_prognostic_variables_from_state_vector!(model.prognostic_variables, state) +end + +function ParticleDA.get_log_density_observation_given_state(observation, state, model) + +end + +function ParticleDA.write_model_metadata(file, model) + +end + From b224811f22dedcc8022fb4ddc9c8440b6950745b Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 29 Nov 2024 16:22:16 +0000 Subject: [PATCH 02/33] Update SpeedyWeather version --- SpeedyWeather/Manifest.toml | 322 ++---------------------------------- SpeedyWeather/src/speedy.jl | 83 +++++----- 2 files changed, 56 insertions(+), 349 deletions(-) diff --git a/SpeedyWeather/Manifest.toml b/SpeedyWeather/Manifest.toml index 05f8d07..f79edab 100644 --- a/SpeedyWeather/Manifest.toml +++ b/SpeedyWeather/Manifest.toml @@ -18,27 +18,6 @@ version = "1.5.0" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -[[deps.Accessors]] -deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown", "Test"] -git-tree-sha1 = "f61b15be1d76846c0ce31d3fcfac5380ae53db6a" -uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" -version = "0.1.37" - - [deps.Accessors.extensions] - AccessorsAxisKeysExt = "AxisKeys" - AccessorsIntervalSetsExt = "IntervalSets" - AccessorsStaticArraysExt = "StaticArrays" - AccessorsStructArraysExt = "StructArrays" - AccessorsUnitfulExt = "Unitful" - - [deps.Accessors.weakdeps] - AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" - IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" - Requires = "ae029012-a4dd-5104-9daa-d747884805df" - StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" - StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" - Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" - [[deps.Adapt]] deps = ["LinearAlgebra", "Requires"] git-tree-sha1 = "6a55b747d1812e699320963ffde36f1ebdda4099" @@ -55,11 +34,6 @@ git-tree-sha1 = "9876e1e164b144ca45e9e3198d0b689cadfed9ff" uuid = "66dad0bd-aa9a-41b7-9441-69ab47430ed8" version = "1.1.3" -[[deps.ArgCheck]] -git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4" -uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197" -version = "2.3.0" - [[deps.ArgTools]] uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" version = "1.1.1" @@ -78,42 +52,9 @@ git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be" uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458" version = "0.1.0" -[[deps.BFloat16s]] -deps = ["LinearAlgebra", "Printf", "Random", "Test"] -git-tree-sha1 = "2c7cc21e8678eff479978a0a2ef5ce2f51b63dff" -uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b" -version = "0.5.0" - -[[deps.BangBang]] -deps = ["Accessors", "ConstructionBase", "InitialValues", "LinearAlgebra", "Requires"] -git-tree-sha1 = "e2144b631226d9eeab2d746ca8880b7ccff504ae" -uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" -version = "0.4.3" - - [deps.BangBang.extensions] - BangBangChainRulesCoreExt = "ChainRulesCore" - BangBangDataFramesExt = "DataFrames" - BangBangStaticArraysExt = "StaticArrays" - BangBangStructArraysExt = "StructArrays" - BangBangTablesExt = "Tables" - BangBangTypedTablesExt = "TypedTables" - - [deps.BangBang.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" - StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" - StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" - Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" - TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9" - [[deps.Base64]] uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" -[[deps.Baselet]] -git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e" -uuid = "9718e550-a3fa-408a-8086-8db961cd8217" -version = "0.1.1" - [[deps.BitInformation]] deps = ["Distributions", "Random", "StatsBase"] git-tree-sha1 = "2cf994e66a0886a91ba108cb3cfc044363f0bb01" @@ -143,40 +84,6 @@ git-tree-sha1 = "5afb5c5ba2688ca43a9ad2e5a91cbb93921ccfa1" uuid = "179af706-886a-5703-950a-314cd64e0468" version = "0.1.3" -[[deps.CUDA]] -deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"] -git-tree-sha1 = "fdd9dfb67dfefd548f51000cc400bb51003de247" -uuid = "052768ef-5323-5732-b1bb-66c8b64840ba" -version = "5.4.3" - - [deps.CUDA.extensions] - ChainRulesCoreExt = "ChainRulesCore" - EnzymeCoreExt = "EnzymeCore" - SpecialFunctionsExt = "SpecialFunctions" - - [deps.CUDA.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" - SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" - -[[deps.CUDA_Driver_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "325058b426c2b421e3d2df3d5fa646d72d2e3e7e" -uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc" -version = "0.9.2+0" - -[[deps.CUDA_Runtime_Discovery]] -deps = ["Libdl"] -git-tree-sha1 = "33576c7c1b2500f8e7e6baa082e04563203b3a45" -uuid = "1af6417a-86b4-443c-805f-a4643ffb695f" -version = "0.3.5" - -[[deps.CUDA_Runtime_jll]] -deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "afea94249b821dc754a8ca6695d3daed851e1f5a" -uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" -version = "0.14.1+0" - [[deps.ChunkSplitters]] deps = ["Compat", "TestItems"] git-tree-sha1 = "783507c1f2371c8f2d321f41c3057ecd42cafa83" @@ -238,15 +145,6 @@ deps = ["Artifacts", "Libdl"] uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" version = "1.1.1+0" -[[deps.CompositionsBase]] -git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad" -uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b" -version = "0.1.2" -weakdeps = ["InverseFunctions"] - - [deps.CompositionsBase.extensions] - CompositionsBaseInverseFunctionsExt = "InverseFunctions" - [[deps.ConstructionBase]] git-tree-sha1 = "76219f1ed5771adbb096743bff43fb5fdd4c1157" uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" @@ -262,12 +160,6 @@ version = "1.5.8" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" -[[deps.ContextVariablesX]] -deps = ["Compat", "Logging", "UUIDs"] -git-tree-sha1 = "25cc3803f1030ab855e383129dcd3dc294e322cc" -uuid = "6add18c4-b38d-439d-96f6-d6bc489c04c5" -version = "0.1.3" - [[deps.Contour]] git-tree-sha1 = "439e35b0b36e2e5881738abc8857bd92ad6ff9a8" uuid = "d38c429a-6771-53c6-b99e-75d170b6e991" @@ -283,12 +175,6 @@ git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe" uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" version = "1.16.0" -[[deps.DataFrames]] -deps = ["Compat", "DataAPI", "DataStructures", "Future", "InlineStrings", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrecompileTools", "PrettyTables", "Printf", "REPL", "Random", "Reexport", "SentinelArrays", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"] -git-tree-sha1 = "04c738083f29f86e62c8afc341f0967d8717bdb8" -uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -version = "1.6.1" - [[deps.DataStructures]] deps = ["Compat", "InteractiveUtils", "OrderedCollections"] git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82" @@ -304,11 +190,6 @@ version = "1.0.0" deps = ["Printf"] uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" -[[deps.DefineSingletons]] -git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c" -uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52" -version = "0.1.2" - [[deps.DiskArrays]] deps = ["LRUCache", "OffsetArrays"] git-tree-sha1 = "ef25c513cad08d7ebbed158c91768ae32f308336" @@ -363,18 +244,6 @@ git-tree-sha1 = "c6033cc3892d0ef5bb9cd29b7f2f0331ea5184ea" uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a" version = "3.3.10+0" -[[deps.FLoops]] -deps = ["BangBang", "Compat", "FLoopsBase", "InitialValues", "JuliaVariables", "MLStyle", "Serialization", "Setfield", "Transducers"] -git-tree-sha1 = "0a2e5873e9a5f54abb06418d57a8df689336a660" -uuid = "cc61a311-1640-44b5-9fba-1b764f453329" -version = "0.2.2" - -[[deps.FLoopsBase]] -deps = ["ContextVariablesX"] -git-tree-sha1 = "656f7a6859be8673bf1f35da5670246b923964f7" -uuid = "b9860ae5-e623-471e-878b-f6a53c775ea6" -version = "0.1.1" - [[deps.FastGaussQuadrature]] deps = ["LinearAlgebra", "SpecialFunctions", "StaticArrays"] git-tree-sha1 = "fd923962364b645f3719855c88f7074413a6ad92" @@ -408,10 +277,6 @@ git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172" uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" version = "0.8.5" -[[deps.Future]] -deps = ["Random"] -uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" - [[deps.GMP_jll]] deps = ["Artifacts", "Libdl"] uuid = "781609d7-10c4-51f6-84f2-b8444358ff6d" @@ -429,12 +294,6 @@ git-tree-sha1 = "ec632f177c0d990e64d955ccc1b8c04c485a0950" uuid = "46192b85-c4d5-4398-a991-12ede77f4527" version = "0.1.6" -[[deps.GPUCompiler]] -deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Preferences", "Scratch", "Serialization", "TOML", "TimerOutputs", "UUIDs"] -git-tree-sha1 = "ab29216184312f99ff957b32cd63c2fe9c928b91" -uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "0.26.7" - [[deps.GenericFFT]] deps = ["AbstractFFTs", "FFTW", "LinearAlgebra", "Reexport"] git-tree-sha1 = "1bc01f2ea9a0226a60723794ff86b8017739f5d9" @@ -475,24 +334,6 @@ git-tree-sha1 = "7c4195be1649ae622304031ed46a2f4df989f1eb" uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a" version = "0.3.24" -[[deps.InitialValues]] -git-tree-sha1 = "4da0f88e9a39111c2fa3add390ab15f3a44f3ca3" -uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c" -version = "0.3.1" - -[[deps.InlineStrings]] -git-tree-sha1 = "45521d31238e87ee9f9732561bfee12d4eebd52d" -uuid = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48" -version = "1.4.2" - - [deps.InlineStrings.extensions] - ArrowTypesExt = "ArrowTypes" - ParsersExt = "Parsers" - - [deps.InlineStrings.weakdeps] - ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" - Parsers = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" - [[deps.IntegerMathUtils]] git-tree-sha1 = "b8ffb903da9f7b8cf695a8bead8e01814aa24b30" uuid = "18e54dd8-cb9d-406c-a71d-865a43cbb235" @@ -508,21 +349,6 @@ version = "2024.2.1+0" deps = ["Markdown"] uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" -[[deps.InverseFunctions]] -git-tree-sha1 = "2787db24f4e03daf859c6509ff87764e4182f7d1" -uuid = "3587e190-3f89-42d0-90ee-14403ec27112" -version = "0.1.16" -weakdeps = ["Dates", "Test"] - - [deps.InverseFunctions.extensions] - InverseFunctionsDatesExt = "Dates" - InverseFunctionsTestExt = "Test" - -[[deps.InvertedIndices]] -git-tree-sha1 = "0dc7b50b8d436461be01300fd8cd45aa0274b038" -uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" -version = "1.3.0" - [[deps.IrrationalConstants]] git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" @@ -545,18 +371,6 @@ git-tree-sha1 = "f389674c99bfcde17dc57454011aa44d5a260a40" uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" version = "1.6.0" -[[deps.JuliaNVTXCallbacks_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "af433a10f3942e882d3c671aacb203e006a5808f" -uuid = "9c1d0b0a-7046-5b2e-a33f-ea22f176ac7e" -version = "0.2.1+0" - -[[deps.JuliaVariables]] -deps = ["MLStyle", "NameResolution"] -git-tree-sha1 = "49fb3cb53362ddadb4415e9b73926d6b40709e70" -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 = "cb1cff88ef2f3a157cbad75bbe6b229e1975e498" @@ -578,22 +392,19 @@ deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", git-tree-sha1 = "2470e69781ddd70b8878491233cd09bc1bd7fc96" uuid = "929cbde3-209d-540e-8aea-75f648917ca0" version = "8.1.0" -weakdeps = ["BFloat16s"] [deps.LLVM.extensions] BFloat16sExt = "BFloat16s" + [deps.LLVM.weakdeps] + BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b" + [[deps.LLVMExtra_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] git-tree-sha1 = "597d1c758c9ae5d985ba4202386a607c675ee700" uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" version = "0.0.31+0" -[[deps.LLVMLoopInfo]] -git-tree-sha1 = "2e5c102cfc41f48ae4740c7eca7743cc7e7b75ea" -uuid = "8b046642-f1f6-4319-8d3c-209ddc03c586" -version = "1.0.0" - [[deps.LLVMOpenMP_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] git-tree-sha1 = "78211fb6cbc872f77cad3fc0b6cf647d923f4929" @@ -609,11 +420,6 @@ weakdeps = ["Serialization"] [deps.LRUCache.extensions] SerializationExt = ["Serialization"] -[[deps.LaTeXStrings]] -git-tree-sha1 = "50901ebc375ed41dbf8058da26f9de442febbbec" -uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" -version = "1.3.1" - [[deps.LazyArtifacts]] deps = ["Artifacts", "Pkg"] uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3" @@ -686,11 +492,6 @@ git-tree-sha1 = "f046ccd0c6db2832a9f639e2c669c6fe867e5f4f" uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7" version = "2024.2.0+0" -[[deps.MLStyle]] -git-tree-sha1 = "bc38dff0548128765760c79eb7388a4b37fae2c8" -uuid = "d8e11817-5142-5d16-987a-aa16d5891078" -version = "0.4.17" - [[deps.MPI]] deps = ["Distributed", "DocStringExtensions", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "PkgVersion", "PrecompileTools", "Requires", "Serialization", "Sockets"] git-tree-sha1 = "b4d8707e42b693720b54f0b3434abee6dd4d947a" @@ -744,12 +545,6 @@ deps = ["Artifacts", "Libdl"] uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1" version = "2.28.2+1" -[[deps.MicroCollections]] -deps = ["Accessors", "BangBang", "InitialValues"] -git-tree-sha1 = "44d32db644e84c75dab479f1bc15ee76a1a3618f" -uuid = "128add7d-3638-4c79-886c-908ea0c25c34" -version = "0.2.0" - [[deps.MicrosoftMPI_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "f12a29c4400ba812841c6ace3f4efbb6dbb3ba01" @@ -775,30 +570,12 @@ git-tree-sha1 = "77df6d3708ec0eb3441551e1f20f7503b37c2393" uuid = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" version = "0.14.5" -[[deps.NVTX]] -deps = ["Colors", "JuliaNVTXCallbacks_jll", "Libdl", "NVTX_jll"] -git-tree-sha1 = "53046f0483375e3ed78e49190f1154fa0a4083a1" -uuid = "5da4648a-3479-48b8-97b9-01cb529c0a1f" -version = "0.3.4" - -[[deps.NVTX_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "ce3269ed42816bf18d500c9f63418d4b0d9f5a3b" -uuid = "e98f9f5b-d649-5603-91fd-7774390e6439" -version = "3.1.0+2" - [[deps.NaNMath]] deps = ["OpenLibm_jll"] git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" version = "1.0.2" -[[deps.NameResolution]] -deps = ["PrettyPrint"] -git-tree-sha1 = "1a0fa0e9613f46c9b8c11eee38ebb4f590013c5e" -uuid = "71a1bf82-56d0-4bbc-8a3c-48b961074391" -version = "0.1.5" - [[deps.NetCDF_jll]] deps = ["Artifacts", "Blosc_jll", "Bzip2_jll", "HDF5_jll", "JLLWrappers", "LibCURL_jll", "Libdl", "OpenMPI_jll", "XML2_jll", "Zlib_jll", "Zstd_jll", "libzip_jll"] git-tree-sha1 = "a8af1798e4eb9ff768ce7fdefc0e957097793f15" @@ -886,12 +663,6 @@ git-tree-sha1 = "f9501cc0430a26bc3d156ae1b5b0c1b47af4d6da" uuid = "eebad327-c553-4316-9ea0-9fa01ccd7688" version = "0.3.3" -[[deps.PooledArrays]] -deps = ["DataAPI", "Future"] -git-tree-sha1 = "36d8b4b899628fb92c2749eb488d884a926614d3" -uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" -version = "1.4.3" - [[deps.PrecompileTools]] deps = ["Preferences"] git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f" @@ -904,17 +675,6 @@ git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6" uuid = "21216c6a-2e73-6563-6e65-726566657250" version = "1.4.3" -[[deps.PrettyPrint]] -git-tree-sha1 = "632eb4abab3449ab30c5e1afaa874f0b98b586e4" -uuid = "8162dcfd-2161-5ef2-ae6c-7681170c5f98" -version = "0.2.0" - -[[deps.PrettyTables]] -deps = ["Crayons", "LaTeXStrings", "Markdown", "PrecompileTools", "Printf", "Reexport", "StringManipulation", "Tables"] -git-tree-sha1 = "66b20dd35966a748321d3b2537c4584cf40387c7" -uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" -version = "2.3.2" - [[deps.Primes]] deps = ["IntegerMathUtils"] git-tree-sha1 = "cb420f77dc474d23ee47ca8d14c90810cafe69e7" @@ -956,18 +716,6 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" deps = ["SHA"] uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -[[deps.Random123]] -deps = ["Random", "RandomNumbers"] -git-tree-sha1 = "4743b43e5a9c4a2ede372de7061eed81795b12e7" -uuid = "74087812-796a-5b5d-8853-05524746bad3" -version = "1.7.0" - -[[deps.RandomNumbers]] -deps = ["Random"] -git-tree-sha1 = "c6ec94d2aaba1ab2ff983052cf6a606ca5985902" -uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143" -version = "1.6.0" - [[deps.Reexport]] git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b" uuid = "189a3867-3050-52da-a836-e630ba90ab69" @@ -995,27 +743,9 @@ version = "0.4.3+0" uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" version = "0.7.0" -[[deps.Scratch]] -deps = ["Dates"] -git-tree-sha1 = "3bac05bc7e74a75fd9cba4295cde4045d9fe2386" -uuid = "6c6a2e73-6563-6170-7368-637461726353" -version = "1.2.1" - -[[deps.SentinelArrays]] -deps = ["Dates", "Random"] -git-tree-sha1 = "ff11acffdb082493657550959d4feb4b6149e73a" -uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c" -version = "1.4.5" - [[deps.Serialization]] uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" -[[deps.Setfield]] -deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"] -git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac" -uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46" -version = "1.1.1" - [[deps.Sockets]] uuid = "6462fe0b-24de-5631-8697-dd941f90decc" @@ -1043,25 +773,23 @@ version = "2.4.0" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" [[deps.SpeedyWeather]] -deps = ["AbstractFFTs", "Adapt", "AssociatedLegendrePolynomials", "BitInformation", "CUDA", "CodecZlib", "Dates", "DocStringExtensions", "FFTW", "FLoops", "FastGaussQuadrature", "GPUArrays", "GenericFFT", "JLD2", "KernelAbstractions", "LinearAlgebra", "NCDatasets", "Primes", "Printf", "ProgressMeter", "Random", "Statistics", "TOML", "UnicodePlots"] -git-tree-sha1 = "7af375091ba736521edae7d839a085354a3dad4e" +deps = ["AbstractFFTs", "Adapt", "AssociatedLegendrePolynomials", "BitInformation", "CodecZlib", "Dates", "DocStringExtensions", "FFTW", "FastGaussQuadrature", "GPUArrays", "GenericFFT", "JLD2", "KernelAbstractions", "LinearAlgebra", "NCDatasets", "Primes", "Printf", "ProgressMeter", "Random", "Statistics", "TOML", "UnicodePlots"] +git-tree-sha1 = "74218651ed7254a03c5cec32df8b3acd32835780" uuid = "9e226e20-d153-4fed-8a5b-493def4f21a9" -version = "0.10.0" +version = "0.12.1" [deps.SpeedyWeather.extensions] + SpeedyWeatherCUDAExt = "CUDA" + SpeedyWeatherGeoMakieExt = "GeoMakie" SpeedyWeatherJLArraysExt = "JLArrays" SpeedyWeatherMakieExt = "Makie" [deps.SpeedyWeather.weakdeps] + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" -[[deps.SplittablesBase]] -deps = ["Setfield", "Test"] -git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5" -uuid = "171d559e-b47b-412a-8079-5efa626c420e" -version = "0.1.15" - [[deps.StaticArrays]] deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] git-tree-sha1 = "eeafab08ae20c62c44c8399ccb9354a04b80db50" @@ -1118,12 +846,6 @@ git-tree-sha1 = "b765e46ba27ecf6b44faf70df40c57aa3a547dcb" uuid = "69024149-9ee7-55f6-a4c4-859efe599b68" version = "0.3.7" -[[deps.StringManipulation]] -deps = ["PrecompileTools"] -git-tree-sha1 = "a04cabe79c5f01f4d723cc6704070ada0b9d46d5" -uuid = "892a3eda-7b42-436c-8928-eab12a02cf0e" -version = "0.3.4" - [[deps.StructArrays]] deps = ["ConstructionBase", "DataAPI", "Tables"] git-tree-sha1 = "f4dc295e983502292c4c3f951dbb4e985e35b3be" @@ -1174,10 +896,6 @@ git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6" uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50" version = "0.1.1" -[[deps.Test]] -deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] -uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - [[deps.TestItems]] git-tree-sha1 = "42fd9023fef18b9b78c8343a4e2f3813ffbcefcb" uuid = "1c621080-faea-4a02-84b6-bbd5e436b8fe" @@ -1194,26 +912,6 @@ git-tree-sha1 = "e84b3a11b9bece70d14cce63406bbc79ed3464d2" uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" version = "0.11.2" -[[deps.Transducers]] -deps = ["Accessors", "Adapt", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "ConstructionBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "SplittablesBase", "Tables"] -git-tree-sha1 = "5215a069867476fc8e3469602006b9670e68da23" -uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999" -version = "0.4.82" - - [deps.Transducers.extensions] - TransducersBlockArraysExt = "BlockArrays" - TransducersDataFramesExt = "DataFrames" - TransducersLazyArraysExt = "LazyArrays" - TransducersOnlineStatsBaseExt = "OnlineStatsBase" - TransducersReferenceablesExt = "Referenceables" - - [deps.Transducers.weakdeps] - BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" - DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" - LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" - OnlineStatsBase = "925886fa-5bf2-5e8e-b522-a9147a512338" - Referenceables = "42d2dcc6-99eb-4e98-b66c-637b7d73030e" - [[deps.UUIDs]] deps = ["Random", "SHA"] uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index d6ee645..f2cb05c 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -1,37 +1,39 @@ +module SpeedyWeatherSSM + using ParticleDA +using Random using SpeedyWeather - -module SpeedyWeatherModel - LAYERED_VARIABLES = (:vor, :div, :temp, :humid) -SURFACE_VARIABLES = (:pres) +SURFACE_VARIABLES = (:pres,) -Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:ModelSetup} +Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractModel} spectral_truncation::Int = 31 - n_level::Int = 8 - n_days::Float64 = 0.25 + n_layers::Int = 8 + n_days::T = 0.25 model_type::Type{M} = PrimitiveWetModel end struct SpeedyModel{ T<:AbstractFloat, - G<:AbstractSpectralGrid, - M<:ModelSetup, + G<:SpeedyWeather.AbstractSpectralGrid, + M<:SpeedyWeather.AbstractModel, } parameters::SpeedyParameters{T, M} spectral_grid::G - model_setup::M - prognostic_variables::PrognosticVariables - diagnostic_variables::DiagnosticVariables + model::M + prognostic_variables::PrognosticVariables{T} + diagnostic_variables::DiagnosticVariables{T} n_layered_variables::Int n_surface_variables::Int end -function init(parameters::SpeedyParameters{T, M}) where {T<:AbstractFloat, M<:ModelSetup} - spectral_grid = SpectralGrid( - NF=T, trunc=parameters.spectral_truncation, nlev=parameters.n_level +function init(parameters::SpeedyParameters{T, M}) where { + T<:AbstractFloat, M<:SpeedyWeather.AbstractModel +} + spectral_grid = SpectralGrid(; + NF=T, trunc=parameters.spectral_truncation, nlayers=parameters.n_layers ) model = M(; spectral_grid) simulation = initialize!(model) @@ -55,7 +57,7 @@ end function ParticleDA.get_state_dimension(model::SpeedyModel) (model.parameters.spectral_truncation + 1)^2 * ( - model.parameters.n_level * model.n_layered_variables + model.n_surface_variables + model.parameters.n_layers * model.n_layered_variables + model.n_surface_variables ) end @@ -65,7 +67,7 @@ end function update_spectral_coefficients_from_vector!( spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, - vector::AbstractVector{T} + vector::Vector{T} ) where {T <: AbstractFloat} m, n = size(spectral_coefficients) spectral_coefficients[1:m - 1, 1] = vector[1:m - 1] @@ -80,7 +82,7 @@ end function update_prognostic_variables_from_state_vector!( prognostic_variables::PrognosticVariables{T}, - state::AbstractVector{T}, + state::Vector{T}, ) where {T <: AbstractFloat} start_index = 1 dim_spectral = (prognostic_variables.trunc + 1)^2 @@ -108,10 +110,10 @@ function update_prognostic_variables_from_state_vector!( end function update_vector_from_spectral_coefficients!( - vector::AbstractVector{T}, + vector::Vector{T}, spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, ) where {T <: AbstractFloat} - m, n = size(spectral_coefficients) + m, n = size(spectral_coefficients, as=Matrix) vector[1:m - 1] = spectral_coefficients[1:m - 1, 1] j = m - 1 for i in 2:n @@ -121,11 +123,11 @@ function update_vector_from_spectral_coefficients!( end function update_state_vector_from_prognostic_variables!( - state::AbstractVector{T}, + state::Vector{T}, prognostic_variables::PrognosticVariables{T} ) where {T <: AbstractFloat} start_index = 1 - dim_spectral = (prognostic_variables.lmax + 1) * (prognostic_variables.mmax + 1) + dim_spectral = (prognostic_variables.trunc + 1)^2 for name in (:vor, :div, :temp, :humid) if SpeedyWeather.has(prognostic_variables, name) layer_spectral_coefficients = SpeedyWeather.get_var( @@ -154,42 +156,48 @@ ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T function ParticleDA.sample_initial_state!( - state::V{T}, model::SpeedyModel{T}, rng::G -) where {T<:AbstractFloat, V<:AbstractVector, G<:AbstractRNG} + state::Vector{T}, model::SpeedyModel{T}, rng::R +) where {T<:AbstractFloat, R<:AbstractRNG} initial_conditions = model.model.initial_conditions SpeedyWeather.initialize!( - model.prognostic_variables, initial_conditions, model.model_setup + model.prognostic_variables, initial_conditions, model.model ) update_state_vector_from_prognostic_variables!(state, model.prognostic_variables) end function ParticleDA.update_state_deterministic!( - state::V{T}, model::SpeedyModel{T}, time_index::Int -) where {T<:AbstractFloat, V<:AbstractVector} + state::Vector{T}, model::SpeedyModel{T}, time_index::Int +) where {T<:AbstractFloat} update_prognostic_variables_from_state_vector!(model.prognostic_variables, state) (; clock) = model.prognostic_variables - (; time_stepping) = model.model_setup + (; time_stepping) = model.model SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameter.n_days)) SpeedyWeather.initialize!(clock, time_stepping) - SpeedyWeather.time_stepping!( - model.prognostic_variables, model.diagnostic_variables, model.model_setup - ) clock.time += clock.n_timesteps * clock.Δt * (time_index - 1) clock.timestep_counter += clock.n_timesteps * (time_index - 1) + SpeedyWeather.time_stepping!( + model.prognostic_variables, model.diagnostic_variables, model.model + ) update_state_vector_from_prognostic_variables!(state, model.prognostic_variables) end function ParticleDA.update_state_stochastic!( - state::V{T}, model::SpeedyModel{T}, rng::G -) where {T<:AbstractFloat, V<:AbstractVector, G<:AbstractRNG} + state::Vector{T}, model::SpeedyModel{T}, rng::G +) where {T<:AbstractFloat, G<:AbstractRNG} end function ParticleDA.sample_observation_given_state!( - observation::V{T}, state::V{T}, model::SpeedyModel{T}, rng::G -) where {T<:AbstractFloat, V<:AbstractVector, G<:AbstractRNG} - # TODO: figure out how to update diagnostic variables given prognostic variables + observation::Vector{T}, state::Vector{T}, model::SpeedyModel{T}, rng::G +) where {T<:AbstractFloat, G<:AbstractRNG} update_prognostic_variables_from_state_vector!(model.prognostic_variables, state) + SpeedyWeather.transform!( + model.diagnostic_variables, + model.prognostic_variables, + 1, + model.model, + true + ) end function ParticleDA.get_log_density_observation_given_state(observation, state, model) @@ -197,6 +205,7 @@ function ParticleDA.get_log_density_observation_given_state(observation, state, end function ParticleDA.write_model_metadata(file, model) - + end +end From bf2848d88db42bdd20aef4ceae1da666af66fb10 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 2 Dec 2024 13:59:31 +0000 Subject: [PATCH 03/33] Fix mapping state vector to/from spectral coeffs --- SpeedyWeather/src/speedy.jl | 128 +++++++++++++++++++++++------------- 1 file changed, 81 insertions(+), 47 deletions(-) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index f2cb05c..b238b86 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -67,86 +67,120 @@ end function update_spectral_coefficients_from_vector!( spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, - vector::Vector{T} + vector::AbstractVector{T} ) where {T <: AbstractFloat} - m, n = size(spectral_coefficients) - spectral_coefficients[1:m - 1, 1] = vector[1:m - 1] - j = m - 1 - for i in 2:n - spectral_coefficients[i:m - 1, i] = reinterpret( - Complex{T}, vector[j+1:j+(m - i) * 2] + n_row, n_col = size(spectral_coefficients, as=Matrix) + # First column of spectral_coefficients (order = m = 0) are real-valued and we skip + # last row (degree = l = n_row - 1) as used only for computing meridional derivative + # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing + # skipping zero upper-triangular entries + spectral_coefficients[1:n_row - 1] .= vector[1:n_row - 1] + # vector index is i, spectral coefficient (flat) index is j + i = n_row - 1 + j = n_row + for col_index in 2:n_col + # Slice corresponding to column has non-zero entries from col_index row and we + # ignore last row as used only for computing meridional derivative for vector + # valued fields + slice_size = n_row - col_index + # Reinterpret real valued state coefficients to complex spectral coefficients + spectral_coefficients[j + 1:j + slice_size] .= reinterpret( + Complex{T}, vector[i + 1:i + 2 * slice_size] ) - j = j + (m - i) * 2 + # Zero entry corresponding to last row (degree = l = n_row - 1) as not used + # for scalar fields + spectral_coefficients[j + slice_size + 1] = 0 + # Update vector and spectral coefficient indices, adding 1 offset + # to latter to skip entries corresponding to last row + i = i + 2 * slice_size + j = j + 1 + slice_size end end function update_prognostic_variables_from_state_vector!( - prognostic_variables::PrognosticVariables{T}, - state::Vector{T}, + model::SpeedyModel{T}, state::Vector{T} ) where {T <: AbstractFloat} start_index = 1 - dim_spectral = (prognostic_variables.trunc + 1)^2 + dim_spectral = (model.parameters.spectral_truncation + 1)^2 + (; prognostic_variables) = model for name in LAYERED_VARIABLES - if SpeedyWeather.has(prognostic_variables, name) - layer_spectral_coefficients = SpeedyWeather.get_var( - prognostic_variables, name; lf=1 - ) - for spectral_coefficients in layer_spectral_coefficients + if SpeedyWeather.has(model.model, name) + # We only consider spectral coefficients for first leapfrog step (lf=1) to + # define state + layered_spectral_coefficients = getproperty(prognostic_variables, name)[1] + for layer_index in 1:model.parameters.n_layers end_index = start_index + dim_spectral - 1 update_spectral_coefficients_from_vector!( - spectral_coefficients, - state[start_index:end_index] + layered_spectral_coefficients[:, layer_index], + view(state, start_index:end_index) ) start_index = end_index + 1 end end end - if SpeedyWeather.has(prognostic_variables, :pres) + if SpeedyWeather.has(model.model, :pres) update_spectral_coefficients_from_vector!( - SpeedyWeather.get_pressure(prognostic_variables; lf=1), - state[start_index:end] + prognostic_variables.pres[1], + view(state, start_index:start_index + dim_spectral - 1) ) end end function update_vector_from_spectral_coefficients!( - vector::Vector{T}, + vector::AbstractVector{T}, spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, ) where {T <: AbstractFloat} - m, n = size(spectral_coefficients, as=Matrix) - vector[1:m - 1] = spectral_coefficients[1:m - 1, 1] - j = m - 1 - for i in 2:n - vector[j+1:j+(m - i) * 2] = reinterpret(T, spectral_coefficients[i:m - 1, i]) - j = j + (m - i) * 2 + n_row, n_col = size(spectral_coefficients, as=Matrix) + # First column of spectral_coefficients (order = m = 0) are real-valued and we skip + # last row (degree = l = n_row - 1) as used only for computing meridional derivative + # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing + # skipping zero upper-triangular entries + vector[1:n_row - 1] .= real(spectral_coefficients[1:n_row - 1]) + # vector index is i, spectral coefficient (flat) index is j + i = n_row - 1 + j = n_row + for col_index in 2:n_col + # Slice corresponding to column has non-zero entries from col_index row and we + # ignore last row as used only for computing meridional derivative for vector + # valued fields + slice_size = n_row - col_index + # Reinterpret complex valued spectral coefficients to extract both real and + # imaginary components + vector[i + 1:i + 2 * slice_size] .= reinterpret( + T, spectral_coefficients[j + 1:j + slice_size] + ) + # Update vector and spectral coefficient indices, adding 1 offset + # to latter to skip entries corresponding to last row + i = i + 2 * slice_size + j = j + 1 + slice_size end end function update_state_vector_from_prognostic_variables!( - state::Vector{T}, - prognostic_variables::PrognosticVariables{T} + state::Vector{T}, model::SpeedyModel{T}, ) where {T <: AbstractFloat} start_index = 1 - dim_spectral = (prognostic_variables.trunc + 1)^2 - for name in (:vor, :div, :temp, :humid) - if SpeedyWeather.has(prognostic_variables, name) - layer_spectral_coefficients = SpeedyWeather.get_var( - prognostic_variables, name; lf=1 - ) - for spectral_coefficients in layer_spectral_coefficients + dim_spectral = (model.parameters.spectral_truncation + 1)^2 + (; prognostic_variables) = model + for name in LAYERED_VARIABLES + if SpeedyWeather.has(model.model, name) + # We only consider spectral coefficients for first leapfrog step (lf=1) to + # define state + layered_spectral_coefficients = getproperty(prognostic_variables, name)[1] + for layer_index in 1:model.parameters.n_layers end_index = start_index + dim_spectral - 1 update_vector_from_spectral_coefficients!( - state[start_index:end_index], - spectral_coefficients + view(state, start_index:end_index), + layered_spectral_coefficients[:, layer_index] ) start_index = end_index + 1 end end end - if SpeedyWeather.has(prognostic_variables, :pres) + if SpeedyWeather.has(model.model, :pres) update_vector_from_spectral_coefficients!( - state[start_index:end], - SpeedyWeather.get_pressure(prognostic_variables; lf=1) + view(state, start_index:start_index + dim_spectral - 1), + prognostic_variables.pres[1] ) end end @@ -162,13 +196,13 @@ function ParticleDA.sample_initial_state!( SpeedyWeather.initialize!( model.prognostic_variables, initial_conditions, model.model ) - update_state_vector_from_prognostic_variables!(state, model.prognostic_variables) + update_state_vector_from_prognostic_variables!(state, model) end function ParticleDA.update_state_deterministic!( state::Vector{T}, model::SpeedyModel{T}, time_index::Int ) where {T<:AbstractFloat} - update_prognostic_variables_from_state_vector!(model.prognostic_variables, state) + update_prognostic_variables_from_state_vector!(model, state) (; clock) = model.prognostic_variables (; time_stepping) = model.model SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameter.n_days)) @@ -178,7 +212,7 @@ function ParticleDA.update_state_deterministic!( SpeedyWeather.time_stepping!( model.prognostic_variables, model.diagnostic_variables, model.model ) - update_state_vector_from_prognostic_variables!(state, model.prognostic_variables) + update_state_vector_from_prognostic_variables!(state, model) end function ParticleDA.update_state_stochastic!( @@ -190,7 +224,7 @@ end function ParticleDA.sample_observation_given_state!( observation::Vector{T}, state::Vector{T}, model::SpeedyModel{T}, rng::G ) where {T<:AbstractFloat, G<:AbstractRNG} - update_prognostic_variables_from_state_vector!(model.prognostic_variables, state) + update_prognostic_variables_from_state_vector!(model, state) SpeedyWeather.transform!( model.diagnostic_variables, model.prognostic_variables, @@ -205,7 +239,7 @@ function ParticleDA.get_log_density_observation_given_state(observation, state, end function ParticleDA.write_model_metadata(file, model) - + end end From 3fc0cc3a38ada0f0409a7e9dcb3277aa27ac1826 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 2 Dec 2024 14:00:06 +0000 Subject: [PATCH 04/33] Update SpeedyWeather/Project.toml MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Mosè Giordano <765740+giordano@users.noreply.github.com> --- SpeedyWeather/Project.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/SpeedyWeather/Project.toml b/SpeedyWeather/Project.toml index 5dad286..ffde9f2 100644 --- a/SpeedyWeather/Project.toml +++ b/SpeedyWeather/Project.toml @@ -1,3 +1,7 @@ [deps] ParticleDA = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" + +[compat] +ParticleDA = "1.1" +SpeedyWeather = "0.12" From 1960393e6a8f72c4216f5899aa2a0d160ea7e36e Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 17:38:25 +0000 Subject: [PATCH 05/33] Add basic observation model implementation --- SpeedyWeather/src/speedy.jl | 94 ++++++++++++++++++++++++++++++++----- 1 file changed, 82 insertions(+), 12 deletions(-) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index b238b86..3279274 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -3,21 +3,36 @@ module SpeedyWeatherSSM using ParticleDA using Random using SpeedyWeather +using SpeedyWeather.RingGrids LAYERED_VARIABLES = (:vor, :div, :temp, :humid) SURFACE_VARIABLES = (:pres,) +function equispaced_lat_lon_grid(T, n_lat, n_lon) + lat_interval = 180 / n_lat + lon_interval = 360 / n_lon + lat_range = (-90 + lat_interval/2):lat_interval:90 + lon_range = (lon_interval/2):lon_interval:360 + lat_lon_pairs = [(T(lat), T(lon)) for lat in lat_range for lon in lon_range] + collect(reshape(reinterpret(T, lat_lon_pairs), (2, :))) +end + Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractModel} spectral_truncation::Int = 31 n_layers::Int = 8 n_days::T = 0.25 + float_type::Type{T} = Float64 model_type::Type{M} = PrimitiveWetModel + observed_variable::Tuple{Symbol, Symbol} = (:physics, :precip_large_scale) + observed_coordinates::Matrix{T} = equispaced_lat_lon_grid(float_type, 6, 12) + observation_noise_std::T = 0.1 end struct SpeedyModel{ T<:AbstractFloat, G<:SpeedyWeather.AbstractSpectralGrid, M<:SpeedyWeather.AbstractModel, + I<:SpeedyWeather.RingGrids.AbstractInterpolator } parameters::SpeedyParameters{T, M} spectral_grid::G @@ -26,9 +41,10 @@ struct SpeedyModel{ diagnostic_variables::DiagnosticVariables{T} n_layered_variables::Int n_surface_variables::Int + n_observed_points::Int + observation_interpolator::I end - function init(parameters::SpeedyParameters{T, M}) where { T<:AbstractFloat, M<:SpeedyWeather.AbstractModel } @@ -44,6 +60,15 @@ function init(parameters::SpeedyParameters{T, M}) where { n_surface_variables = count( SpeedyWeather.has(model, var) for var in SURFACE_VARIABLES ) + n_observed_points = size(parameters.observed_coordinates, 2) + observation_interpolator = SpeedyWeather.AnvilInterpolator( + T, spectral_grid.Grid, spectral_grid.nlat_half, n_observed_points + ) + SpeedyWeather.RingGrids.update_locator!( + observation_interpolator, + parameters.observed_coordinates[1, :], + parameters.observed_coordinates[2, :] + ) return SpeedyModel( parameters, spectral_grid, @@ -52,6 +77,8 @@ function init(parameters::SpeedyParameters{T, M}) where { diagnostic_variables, n_layered_variables, n_surface_variables, + n_observed_points, + observation_interpolator ) end @@ -62,7 +89,7 @@ function ParticleDA.get_state_dimension(model::SpeedyModel) end function ParticleDA.get_observation_dimension(model::SpeedyModel) - + model.n_observed_points end function update_spectral_coefficients_from_vector!( @@ -98,7 +125,7 @@ function update_spectral_coefficients_from_vector!( end function update_prognostic_variables_from_state_vector!( - model::SpeedyModel{T}, state::Vector{T} + model::SpeedyModel{T}, state::AbstractVector{T} ) where {T <: AbstractFloat} start_index = 1 dim_spectral = (model.parameters.spectral_truncation + 1)^2 @@ -157,7 +184,7 @@ function update_vector_from_spectral_coefficients!( end function update_state_vector_from_prognostic_variables!( - state::Vector{T}, model::SpeedyModel{T}, + state::AbstractVector{T}, model::SpeedyModel{T}, ) where {T <: AbstractFloat} start_index = 1 dim_spectral = (model.parameters.spectral_truncation + 1)^2 @@ -190,7 +217,7 @@ ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T function ParticleDA.sample_initial_state!( - state::Vector{T}, model::SpeedyModel{T}, rng::R + state::AbstractVector{T}, model::SpeedyModel{T}, rng::R, task_index::Int=1 ) where {T<:AbstractFloat, R<:AbstractRNG} initial_conditions = model.model.initial_conditions SpeedyWeather.initialize!( @@ -200,7 +227,7 @@ function ParticleDA.sample_initial_state!( end function ParticleDA.update_state_deterministic!( - state::Vector{T}, model::SpeedyModel{T}, time_index::Int + state::AbstractVector{T}, model::SpeedyModel{T}, time_index::Int, task_index::Int=1 ) where {T<:AbstractFloat} update_prognostic_variables_from_state_vector!(model, state) (; clock) = model.prognostic_variables @@ -216,25 +243,68 @@ function ParticleDA.update_state_deterministic!( end function ParticleDA.update_state_stochastic!( - state::Vector{T}, model::SpeedyModel{T}, rng::G + state::AbstractVector{T}, model::SpeedyModel{T}, rng::G, task_index::Int=1 ) where {T<:AbstractFloat, G<:AbstractRNG} end -function ParticleDA.sample_observation_given_state!( - observation::Vector{T}, state::Vector{T}, model::SpeedyModel{T}, rng::G -) where {T<:AbstractFloat, G<:AbstractRNG} +function get_observed_variable_field( + diagnostic_variables::DiagnosticVariables, model::SpeedyModel +) + observed_outer, observed_inner = model.parameters.observed_variable + getfield( + getfield(diagnostic_variables, observed_outer), observed_inner + ) +end + +function ParticleDA.get_observation_mean_given_state!( + observation_mean::AbstractVector{T}, + state::AbstractVector{T}, + model::SpeedyModel{T}, + task_index::Int=1 +) where {T<:AbstractFloat} update_prognostic_variables_from_state_vector!(model, state) SpeedyWeather.transform!( model.diagnostic_variables, model.prognostic_variables, 1, model.model, - true + initialize=true + ) + observed_field_grid = get_observed_variable_field(model.diagnostic_variables, model) + SpeedyWeather.interpolate!( + observation_mean, observed_field_grid, model.observation_interpolator + ) +end + +function ParticleDA.sample_observation_given_state!( + observation::AbstractVector{T}, + state::AbstractVector{T}, + model::SpeedyModel{T}, + rng::G, + task_index::Int=1 +) where {T<:AbstractFloat, G<:AbstractRNG} + ParticleDA.get_observation_mean_given_state!(observation, state, model, task_index) + observation .+= ( + model.parameters.observation_noise_std + * randn(rng, T, ParticleDA.get_observation_dimension(model)) ) end -function ParticleDA.get_log_density_observation_given_state(observation, state, model) +function ParticleDA.get_log_density_observation_given_state( + observation::AbstractVector{T}, + state::AbstractVector{T}, + model::SpeedyModel{T}, + task_index::Int=1 +) where {T<:AbstractFloat} + observation_mean = Vector{T}(undef, ParticleDA.get_observation_dimension(model)) + ParticleDA.get_observation_mean_given_state!( + observation_mean, state, model, task_index + ) + return ( + -sum((observation - observation_mean).^2) + / (2 * model.parameters.observation_noise_std^2) + ) end From 73f751a9737587d36c4ae6cf44e51febf3f7ff48 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 17:38:41 +0000 Subject: [PATCH 06/33] Fix typo in field name --- SpeedyWeather/src/speedy.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index 3279274..2676d9d 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -232,7 +232,7 @@ function ParticleDA.update_state_deterministic!( update_prognostic_variables_from_state_vector!(model, state) (; clock) = model.prognostic_variables (; time_stepping) = model.model - SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameter.n_days)) + SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameters.n_days)) SpeedyWeather.initialize!(clock, time_stepping) clock.time += clock.n_timesteps * clock.Δt * (time_index - 1) clock.timestep_counter += clock.n_timesteps * (time_index - 1) From 342c9d5a22cfcb5ed873865f6d404f2d477d99b8 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 17:39:02 +0000 Subject: [PATCH 07/33] Add method for recording model metadata --- SpeedyWeather/Manifest.toml | 2 +- SpeedyWeather/Project.toml | 1 + SpeedyWeather/src/speedy.jl | 16 ++++++++++++++-- 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/SpeedyWeather/Manifest.toml b/SpeedyWeather/Manifest.toml index f79edab..b5ca649 100644 --- a/SpeedyWeather/Manifest.toml +++ b/SpeedyWeather/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.4" manifest_format = "2.0" -project_hash = "cf0fb01158045c321488fa4aacf601601dfc4784" +project_hash = "a7d89adc97aede5038df2403d7399e97c4b555a0" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] diff --git a/SpeedyWeather/Project.toml b/SpeedyWeather/Project.toml index ffde9f2..87061a5 100644 --- a/SpeedyWeather/Project.toml +++ b/SpeedyWeather/Project.toml @@ -1,4 +1,5 @@ [deps] +HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ParticleDA = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index 2676d9d..bc0a05f 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -1,5 +1,6 @@ module SpeedyWeatherSSM +using HDF5 using ParticleDA using Random using SpeedyWeather @@ -308,8 +309,19 @@ function ParticleDA.get_log_density_observation_given_state( end -function ParticleDA.write_model_metadata(file, model) - +function ParticleDA.write_model_metadata(file::HDF5.File, model::SpeedyModel) + group_name = "parameters" + if !haskey(file, group_name) + group = create_group(file, group_name) + for field in fieldnames(typeof(model.parameters)) + value = getfield(model.parameters, field) + isa(value, Type) && (value = string(nameof(value))) + isa(value, Tuple{Symbol, Symbol}) && (value = join(map(String, value), ".")) + HDF5.attributes(group)[string(field)] = value + end + else + @warn "Write failed, group $group_name already exists in $(file.filename)!" + end end end From 3ccaa22dca194e58255907979009097a5c2c1067 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 18:25:41 +0000 Subject: [PATCH 08/33] Disable progress meter --- SpeedyWeather/src/speedy.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index bc0a05f..8bad7d8 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -52,7 +52,9 @@ function init(parameters::SpeedyParameters{T, M}) where { spectral_grid = SpectralGrid(; NF=T, trunc=parameters.spectral_truncation, nlayers=parameters.n_layers ) - model = M(; spectral_grid) + feedback = SpeedyWeather.Feedback(; verbose=false) + model = M(; spectral_grid, feedback) + model.output.active = false simulation = initialize!(model) (; prognostic_variables, diagnostic_variables) = simulation n_layered_variables = count( From da9d89210f6f71155d817421a6429f2f45173a7f Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 18:26:07 +0000 Subject: [PATCH 09/33] Use views in state updates to avoid copy --- SpeedyWeather/src/speedy.jl | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index 8bad7d8..6822148 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -96,10 +96,11 @@ function ParticleDA.get_observation_dimension(model::SpeedyModel) end function update_spectral_coefficients_from_vector!( - spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, - vector::AbstractVector{T} + spectral_coefficients::AbstractVector{Complex{T}}, + vector::AbstractVector{T}, + spectral_truncation::Int ) where {T <: AbstractFloat} - n_row, n_col = size(spectral_coefficients, as=Matrix) + n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 # First column of spectral_coefficients (order = m = 0) are real-valued and we skip # last row (degree = l = n_row - 1) as used only for computing meridional derivative # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing @@ -141,8 +142,9 @@ function update_prognostic_variables_from_state_vector!( for layer_index in 1:model.parameters.n_layers end_index = start_index + dim_spectral - 1 update_spectral_coefficients_from_vector!( - layered_spectral_coefficients[:, layer_index], - view(state, start_index:end_index) + view(layered_spectral_coefficients, :, layer_index), + view(state, start_index:end_index), + model.parameters.spectral_truncation ) start_index = end_index + 1 end @@ -151,16 +153,18 @@ function update_prognostic_variables_from_state_vector!( if SpeedyWeather.has(model.model, :pres) update_spectral_coefficients_from_vector!( prognostic_variables.pres[1], - view(state, start_index:start_index + dim_spectral - 1) + view(state, start_index:start_index + dim_spectral - 1), + model.parameters.spectral_truncation ) end end function update_vector_from_spectral_coefficients!( vector::AbstractVector{T}, - spectral_coefficients::SpeedyWeather.LowerTriangularMatrix{Complex{T}}, + spectral_coefficients::AbstractVector{Complex{T}}, + spectral_truncation::Int ) where {T <: AbstractFloat} - n_row, n_col = size(spectral_coefficients, as=Matrix) + n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 # First column of spectral_coefficients (order = m = 0) are real-valued and we skip # last row (degree = l = n_row - 1) as used only for computing meridional derivative # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing @@ -201,7 +205,8 @@ function update_state_vector_from_prognostic_variables!( end_index = start_index + dim_spectral - 1 update_vector_from_spectral_coefficients!( view(state, start_index:end_index), - layered_spectral_coefficients[:, layer_index] + view(layered_spectral_coefficients, :, layer_index), + model.parameters.spectral_truncation ) start_index = end_index + 1 end @@ -210,7 +215,8 @@ function update_state_vector_from_prognostic_variables!( if SpeedyWeather.has(model.model, :pres) update_vector_from_spectral_coefficients!( view(state, start_index:start_index + dim_spectral - 1), - prognostic_variables.pres[1] + prognostic_variables.pres[1], + model.parameters.spectral_truncation ) end end From 8bac55f66332984510a5261158041f51c10dbb82 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 18:55:11 +0000 Subject: [PATCH 10/33] Explicitly set start date --- SpeedyWeather/src/speedy.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index 6822148..aaedc47 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -22,6 +22,7 @@ Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractM spectral_truncation::Int = 31 n_layers::Int = 8 n_days::T = 0.25 + start_date::DateTime = DateTime(2000, 1, 1) float_type::Type{T} = Float64 model_type::Type{M} = PrimitiveWetModel observed_variable::Tuple{Symbol, Symbol} = (:physics, :precip_large_scale) @@ -55,7 +56,7 @@ function init(parameters::SpeedyParameters{T, M}) where { feedback = SpeedyWeather.Feedback(; verbose=false) model = M(; spectral_grid, feedback) model.output.active = false - simulation = initialize!(model) + simulation = initialize!(model; time=parameters.start_date) (; prognostic_variables, diagnostic_variables) = simulation n_layered_variables = count( SpeedyWeather.has(model, var) for var in LAYERED_VARIABLES @@ -243,8 +244,9 @@ function ParticleDA.update_state_deterministic!( (; time_stepping) = model.model SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameters.n_days)) SpeedyWeather.initialize!(clock, time_stepping) - clock.time += clock.n_timesteps * clock.Δt * (time_index - 1) - clock.timestep_counter += clock.n_timesteps * (time_index - 1) + clock.start = model.parameters.start_date + clock.time = clock.start + clock.n_timesteps * clock.Δt * (time_index - 1) + clock.timestep_counter = clock.n_timesteps * (time_index - 1) SpeedyWeather.time_stepping!( model.prognostic_variables, model.diagnostic_variables, model.model ) From 86da18c7d24e6d8a035c7a1fbb5e641a7c1ea0a3 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 19:05:45 +0000 Subject: [PATCH 11/33] Convert date to string when recording parameters --- SpeedyWeather/src/speedy.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeather/src/speedy.jl index aaedc47..53212a0 100644 --- a/SpeedyWeather/src/speedy.jl +++ b/SpeedyWeather/src/speedy.jl @@ -327,6 +327,7 @@ function ParticleDA.write_model_metadata(file::HDF5.File, model::SpeedyModel) value = getfield(model.parameters, field) isa(value, Type) && (value = string(nameof(value))) isa(value, Tuple{Symbol, Symbol}) && (value = join(map(String, value), ".")) + isa(value, DateTime) && (value = string(value)) HDF5.attributes(group)[string(field)] = value end else From 5a39ef916a46c5f0220668694c9f43f3363aa96f Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 19:15:46 +0000 Subject: [PATCH 12/33] Renaming to avoid name clash --- {SpeedyWeather => SpeedyWeatherSSM}/Manifest.toml | 0 {SpeedyWeather => SpeedyWeatherSSM}/Project.toml | 0 {SpeedyWeather => SpeedyWeatherSSM}/src/speedy.jl | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename {SpeedyWeather => SpeedyWeatherSSM}/Manifest.toml (100%) rename {SpeedyWeather => SpeedyWeatherSSM}/Project.toml (100%) rename {SpeedyWeather => SpeedyWeatherSSM}/src/speedy.jl (100%) diff --git a/SpeedyWeather/Manifest.toml b/SpeedyWeatherSSM/Manifest.toml similarity index 100% rename from SpeedyWeather/Manifest.toml rename to SpeedyWeatherSSM/Manifest.toml diff --git a/SpeedyWeather/Project.toml b/SpeedyWeatherSSM/Project.toml similarity index 100% rename from SpeedyWeather/Project.toml rename to SpeedyWeatherSSM/Project.toml diff --git a/SpeedyWeather/src/speedy.jl b/SpeedyWeatherSSM/src/speedy.jl similarity index 100% rename from SpeedyWeather/src/speedy.jl rename to SpeedyWeatherSSM/src/speedy.jl From 56e502ddfe45d9b1eeb4889d50ec9113ed24045b Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 16 Dec 2024 19:53:51 +0000 Subject: [PATCH 13/33] Add name to project metadata --- SpeedyWeatherSSM/Manifest.toml | 2 +- SpeedyWeatherSSM/Project.toml | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/SpeedyWeatherSSM/Manifest.toml b/SpeedyWeatherSSM/Manifest.toml index b5ca649..f51dd8b 100644 --- a/SpeedyWeatherSSM/Manifest.toml +++ b/SpeedyWeatherSSM/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.4" manifest_format = "2.0" -project_hash = "a7d89adc97aede5038df2403d7399e97c4b555a0" +project_hash = "230975858e6ef4728d5bf461050c838f2c75017c" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] diff --git a/SpeedyWeatherSSM/Project.toml b/SpeedyWeatherSSM/Project.toml index 87061a5..8a58a32 100644 --- a/SpeedyWeatherSSM/Project.toml +++ b/SpeedyWeatherSSM/Project.toml @@ -1,6 +1,12 @@ +name = "SpeedyWeatherSSM" +uuid = "0d1ea686-ee24-4a74-b4ed-044400b22b75" +authors = ["Matt Graham"] +version = "0.1.0" + [deps] HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ParticleDA = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" [compat] From 026ff9837edeeda6e192b8f8efef7afbbf6fa48c Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 18 Dec 2024 16:50:52 +0000 Subject: [PATCH 14/33] Correct longitude interval --- SpeedyWeatherSSM/src/speedy.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/SpeedyWeatherSSM/src/speedy.jl b/SpeedyWeatherSSM/src/speedy.jl index 53212a0..241df5d 100644 --- a/SpeedyWeatherSSM/src/speedy.jl +++ b/SpeedyWeatherSSM/src/speedy.jl @@ -13,7 +13,7 @@ function equispaced_lat_lon_grid(T, n_lat, n_lon) lat_interval = 180 / n_lat lon_interval = 360 / n_lon lat_range = (-90 + lat_interval/2):lat_interval:90 - lon_range = (lon_interval/2):lon_interval:360 + lon_range = (-180 + lon_interval/2):lon_interval:180 lat_lon_pairs = [(T(lat), T(lon)) for lat in lat_range for lon in lon_range] collect(reshape(reinterpret(T, lat_lon_pairs), (2, :))) end From e4f1be5160954f7d3d8a83dd83d3d9c885029616 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 18 Dec 2024 16:51:30 +0000 Subject: [PATCH 15/33] Zero last row first column coefficient when updating --- SpeedyWeatherSSM/src/speedy.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/SpeedyWeatherSSM/src/speedy.jl b/SpeedyWeatherSSM/src/speedy.jl index 241df5d..c20427f 100644 --- a/SpeedyWeatherSSM/src/speedy.jl +++ b/SpeedyWeatherSSM/src/speedy.jl @@ -107,6 +107,8 @@ function update_spectral_coefficients_from_vector!( # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing # skipping zero upper-triangular entries spectral_coefficients[1:n_row - 1] .= vector[1:n_row - 1] + # Zero entry corresponding to last row as not used for scalar fields + spectral_coefficients[n_row] = 0 # vector index is i, spectral coefficient (flat) index is j i = n_row - 1 j = n_row @@ -119,8 +121,7 @@ function update_spectral_coefficients_from_vector!( spectral_coefficients[j + 1:j + slice_size] .= reinterpret( Complex{T}, vector[i + 1:i + 2 * slice_size] ) - # Zero entry corresponding to last row (degree = l = n_row - 1) as not used - # for scalar fields + # Zero entry corresponding to last row as not use for scalar fields spectral_coefficients[j + slice_size + 1] = 0 # Update vector and spectral coefficient indices, adding 1 offset # to latter to skip entries corresponding to last row From 8a27be532a0ee1d4c264f1c234ea7dc37c7447c9 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 18 Dec 2024 16:52:11 +0000 Subject: [PATCH 16/33] Renaming top-level package source file --- SpeedyWeatherSSM/src/{speedy.jl => SpeedyWeatherSSM.jl} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename SpeedyWeatherSSM/src/{speedy.jl => SpeedyWeatherSSM.jl} (100%) diff --git a/SpeedyWeatherSSM/src/speedy.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl similarity index 100% rename from SpeedyWeatherSSM/src/speedy.jl rename to SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl From 371f47875d1c537ed3a44738be30adaae0e5af4d Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 18 Dec 2024 16:53:11 +0000 Subject: [PATCH 17/33] Add initial tests --- SpeedyWeatherSSM/Manifest.toml | 277 +++++++++++++++--------------- SpeedyWeatherSSM/Project.toml | 9 +- SpeedyWeatherSSM/test/runtests.jl | 97 +++++++++++ 3 files changed, 242 insertions(+), 141 deletions(-) create mode 100644 SpeedyWeatherSSM/test/runtests.jl diff --git a/SpeedyWeatherSSM/Manifest.toml b/SpeedyWeatherSSM/Manifest.toml index f51dd8b..e40b62c 100644 --- a/SpeedyWeatherSSM/Manifest.toml +++ b/SpeedyWeatherSSM/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.4" manifest_format = "2.0" -project_hash = "230975858e6ef4728d5bf461050c838f2c75017c" +project_hash = "8a48103e70f2b517a2b98bb799150c3116af6c64" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] @@ -20,9 +20,9 @@ version = "1.5.0" [[deps.Adapt]] deps = ["LinearAlgebra", "Requires"] -git-tree-sha1 = "6a55b747d1812e699320963ffde36f1ebdda4099" +git-tree-sha1 = "50c3c56a52972d78e8be9fd135bfb91c9574c140" uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -version = "4.0.4" +version = "4.1.1" weakdeps = ["StaticArrays"] [deps.Adapt.extensions] @@ -48,9 +48,19 @@ version = "1.0.1" [[deps.Atomix]] deps = ["UnsafeAtomics"] -git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be" +git-tree-sha1 = "c3b238aa28c1bebd4b5ea4988bebf27e9a01b72b" uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458" -version = "0.1.0" +version = "1.0.1" + + [deps.Atomix.extensions] + AtomixCUDAExt = "CUDA" + AtomixMetalExt = "Metal" + AtomixoneAPIExt = "oneAPI" + + [deps.Atomix.weakdeps] + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + Metal = "dde4c033-4e86-420c-a63e-0dd931031962" + oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [[deps.Base64]] uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" @@ -63,15 +73,15 @@ version = "0.6.3" [[deps.Blosc_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Lz4_jll", "Zlib_jll", "Zstd_jll"] -git-tree-sha1 = "19b98ee7e3db3b4eff74c5c9c72bf32144e24f10" +git-tree-sha1 = "ef12cdd1c7fb7e1dfd6fa8fd60d4db6bc61d2f23" uuid = "0b7ba130-8d10-5ba8-a3d6-c5182647fed9" -version = "1.21.5+0" +version = "1.21.6+0" [[deps.Bzip2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "9e2a6b69137e6969bab0152632dcb3bc108c8bdd" +git-tree-sha1 = "8873e196c2eb87962a2048b3b8e08946535864a1" uuid = "6e34b625-4abd-537c-b88f-471c36dfa7a0" -version = "1.0.8+1" +version = "1.0.8+2" [[deps.CEnum]] git-tree-sha1 = "389ad5c84de1ae7cf0e28e381131c98ea87d54fc" @@ -85,10 +95,9 @@ uuid = "179af706-886a-5703-950a-314cd64e0468" version = "0.1.3" [[deps.ChunkSplitters]] -deps = ["Compat", "TestItems"] -git-tree-sha1 = "783507c1f2371c8f2d321f41c3057ecd42cafa83" +git-tree-sha1 = "397b871ff701290cc122cca06af61c5bdf9f5605" uuid = "ae650224-84b6-46f8-82ea-d812ca08434e" -version = "2.4.5" +version = "3.1.0" [[deps.CodecZlib]] deps = ["TranscodingStreams", "Zlib_jll"] @@ -98,21 +107,27 @@ version = "0.7.6" [[deps.ColorSchemes]] deps = ["ColorTypes", "ColorVectorSpace", "Colors", "FixedPointNumbers", "PrecompileTools", "Random"] -git-tree-sha1 = "b5278586822443594ff615963b0c09755771b3e0" +git-tree-sha1 = "c785dfb1b3bfddd1da557e861b919819b82bbe5b" uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4" -version = "3.26.0" +version = "3.27.1" [[deps.ColorTypes]] deps = ["FixedPointNumbers", "Random"] -git-tree-sha1 = "b10d0b65641d57b8b4d5e234446582de5047050d" +git-tree-sha1 = "c7acce7a7e1078a20a285211dd73cd3941a871d6" uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f" -version = "0.11.5" +version = "0.12.0" + + [deps.ColorTypes.extensions] + StyledStringsExt = "StyledStrings" + + [deps.ColorTypes.weakdeps] + StyledStrings = "f489334b-da3d-4c2e-b8f0-e476e12c162b" [[deps.ColorVectorSpace]] deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "Requires", "Statistics", "TensorCore"] -git-tree-sha1 = "a1f44953f2382ebb937d60dafbe2deea4bd23249" +git-tree-sha1 = "8b3b6f87ce8f65a2b4f857528fd8d70086cd72b1" uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4" -version = "0.10.0" +version = "0.11.0" weakdeps = ["SpecialFunctions"] [deps.ColorVectorSpace.extensions] @@ -120,15 +135,15 @@ weakdeps = ["SpecialFunctions"] [[deps.Colors]] deps = ["ColorTypes", "FixedPointNumbers", "Reexport"] -git-tree-sha1 = "362a287c3aa50601b0bc359053d5c2468f0e7ce0" +git-tree-sha1 = "64e15186f0aa277e174aa81798f7eb8598e0157e" uuid = "5ae59095-9a9b-59fe-a467-6f913c188581" -version = "0.12.11" +version = "0.13.0" [[deps.CommonDataModel]] deps = ["CFTime", "DataStructures", "Dates", "Preferences", "Printf", "Statistics"] -git-tree-sha1 = "d6fb5bf939a2753c74984b11434ea25d6c397a58" +git-tree-sha1 = "98d64d5b9e5263884276656a43c45424b3a645c2" uuid = "1fbeeb36-5f17-413c-809b-666fb144f157" -version = "0.3.6" +version = "0.3.7" [[deps.Compat]] deps = ["TOML", "UUIDs"] @@ -192,9 +207,9 @@ uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" [[deps.DiskArrays]] deps = ["LRUCache", "OffsetArrays"] -git-tree-sha1 = "ef25c513cad08d7ebbed158c91768ae32f308336" +git-tree-sha1 = "e0e89a60637a62d13aa2107f0acd169b9b9b77e7" uuid = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3" -version = "0.3.23" +version = "0.4.6" [[deps.Distributed]] deps = ["Random", "Serialization", "Sockets"] @@ -202,9 +217,9 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" [[deps.Distributions]] deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"] -git-tree-sha1 = "e6c693a0e4394f8fda0e51a5bdf5aef26f8235e9" +git-tree-sha1 = "4b138e4643b577ccf355377c2bc70fa975af25de" uuid = "31c24e10-a181-5473-b8eb-7969acd0382f" -version = "0.25.111" +version = "0.25.115" [deps.Distributions.extensions] DistributionsChainRulesCoreExt = "ChainRulesCore" @@ -240,9 +255,9 @@ version = "1.8.0" [[deps.FFTW_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "c6033cc3892d0ef5bb9cd29b7f2f0331ea5184ea" +git-tree-sha1 = "4d81ed14783ec49ce9f2e168208a12ce1815aa25" uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a" -version = "3.3.10+0" +version = "3.3.10+1" [[deps.FastGaussQuadrature]] deps = ["LinearAlgebra", "SpecialFunctions", "StaticArrays"] @@ -252,9 +267,15 @@ version = "1.0.2" [[deps.FileIO]] deps = ["Pkg", "Requires", "UUIDs"] -git-tree-sha1 = "82d8afa92ecf4b52d78d869f038ebfb881267322" +git-tree-sha1 = "2dd20384bf8c6d411b5c7370865b1e9b26cb2ea3" uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" -version = "1.16.3" +version = "1.16.6" + + [deps.FileIO.extensions] + HTTPExt = "HTTP" + + [deps.FileIO.weakdeps] + HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3" [[deps.FileWatching]] uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" @@ -277,11 +298,6 @@ git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172" uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" version = "0.8.5" -[[deps.GMP_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "781609d7-10c4-51f6-84f2-b8444358ff6d" -version = "6.2.1+6" - [[deps.GPUArrays]] deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"] git-tree-sha1 = "62ee71528cca49be797076a76bdc654a170a523e" @@ -300,12 +316,6 @@ git-tree-sha1 = "1bc01f2ea9a0226a60723794ff86b8017739f5d9" uuid = "a8297547-1b15-4a5a-a998-a2ac5f1cef28" version = "0.1.6" -[[deps.GnuTLS_jll]] -deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Nettle_jll", "P11Kit_jll", "Zlib_jll"] -git-tree-sha1 = "383db7d3f900f4c1f47a8a04115b053c095e48d3" -uuid = "0951126a-58fd-58f1-b5b3-b08c7c4a876d" -version = "3.8.4+0" - [[deps.HDF5]] deps = ["Compat", "HDF5_jll", "Libdl", "MPIPreferences", "Mmap", "Preferences", "Printf", "Random", "Requires", "UUIDs"] git-tree-sha1 = "e856eef26cf5bf2b0f95f8f4fc37553c72c8641c" @@ -324,15 +334,15 @@ version = "1.14.2+1" [[deps.Hwloc_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "5e19e1e4fa3e71b774ce746274364aef0234634e" +git-tree-sha1 = "50aedf345a709ab75872f80a2779568dc0bb461b" uuid = "e33a78d0-f292-5ffc-b300-72abe9b543c8" -version = "2.11.1+0" +version = "2.11.2+1" [[deps.HypergeometricFunctions]] deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"] -git-tree-sha1 = "7c4195be1649ae622304031ed46a2f4df989f1eb" +git-tree-sha1 = "b1c2585431c382e3fe5805874bda6aea90a95de9" uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a" -version = "0.3.24" +version = "0.3.25" [[deps.IntegerMathUtils]] git-tree-sha1 = "b8ffb903da9f7b8cf695a8bead8e01814aa24b30" @@ -361,21 +371,21 @@ version = "1.0.0" [[deps.JLD2]] deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"] -git-tree-sha1 = "a0746c21bdc986d0dc293efa6b1faee112c37c28" +git-tree-sha1 = "f1a1c1037af2a4541ea186b26b0c0e7eeaad232b" uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819" -version = "0.4.53" +version = "0.5.10" [[deps.JLLWrappers]] deps = ["Artifacts", "Preferences"] -git-tree-sha1 = "f389674c99bfcde17dc57454011aa44d5a260a40" +git-tree-sha1 = "be3dc50a92e5a386872a493a10050136d4703f9b" uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" -version = "1.6.0" +version = "1.6.1" [[deps.KernelAbstractions]] -deps = ["Adapt", "Atomix", "InteractiveUtils", "MacroTools", "PrecompileTools", "Requires", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"] -git-tree-sha1 = "cb1cff88ef2f3a157cbad75bbe6b229e1975e498" +deps = ["Adapt", "Atomix", "InteractiveUtils", "MacroTools", "PrecompileTools", "Requires", "StaticArrays", "UUIDs"] +git-tree-sha1 = "b9a838cd3028785ac23822cded5126b3da394d1a" uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" -version = "0.9.25" +version = "0.9.31" [deps.KernelAbstractions.extensions] EnzymeExt = "EnzymeCore" @@ -388,10 +398,10 @@ version = "0.9.25" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [[deps.LLVM]] -deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"] -git-tree-sha1 = "2470e69781ddd70b8878491233cd09bc1bd7fc96" +deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Unicode"] +git-tree-sha1 = "d422dfd9707bec6617335dc2ea3c5172a87d5908" uuid = "929cbde3-209d-540e-8aea-75f648917ca0" -version = "8.1.0" +version = "9.1.3" [deps.LLVM.extensions] BFloat16sExt = "BFloat16s" @@ -401,9 +411,9 @@ version = "8.1.0" [[deps.LLVMExtra_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "597d1c758c9ae5d985ba4202386a607c675ee700" +git-tree-sha1 = "05a8bd5a42309a9ec82f700876903abce1017dd3" uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" -version = "0.0.31+0" +version = "0.0.34+0" [[deps.LLVMOpenMP_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -453,9 +463,9 @@ uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb" [[deps.Libiconv_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "f9557a255370125b405568f9767d6d195822a175" +git-tree-sha1 = "61dfdba58e585066d8bce214c5a51eaa0539f269" uuid = "94ce4f54-9a6c-5748-9c1c-f9c7231a4531" -version = "1.17.0+0" +version = "1.17.0+1" [[deps.LinearAlgebra]] deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] @@ -463,9 +473,9 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [[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" @@ -482,9 +492,9 @@ uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" [[deps.Lz4_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "7f26c8fc5229e68484e0b3447312c98e16207d11" +git-tree-sha1 = "abf88ff67f4fd89839efcae2f4c39cbc4ecd0846" uuid = "5ced341a-0733-55b8-9ab6-a4889d929147" -version = "1.10.0+0" +version = "1.10.0+1" [[deps.MKL_jll]] deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"] @@ -494,9 +504,9 @@ version = "2024.2.0+0" [[deps.MPI]] deps = ["Distributed", "DocStringExtensions", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "PkgVersion", "PrecompileTools", "Requires", "Serialization", "Sockets"] -git-tree-sha1 = "b4d8707e42b693720b54f0b3434abee6dd4d947a" +git-tree-sha1 = "892676019c58f34e38743bc989b0eca5bce5edc5" uuid = "da04e1cc-30fd-572f-bb4f-1f8673147195" -version = "0.20.16" +version = "0.20.22" [deps.MPI.extensions] AMDGPUExt = "AMDGPU" @@ -508,9 +518,9 @@ version = "0.20.16" [[deps.MPICH_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"] -git-tree-sha1 = "19d4bd098928a3263693991500d05d74dbdc2004" +git-tree-sha1 = "7715e65c47ba3941c502bffb7f266a41a7f54423" uuid = "7cb0a576-ebde-5e09-9194-50597f1243b4" -version = "4.2.2+0" +version = "4.2.3+0" [[deps.MPIPreferences]] deps = ["Libdl", "Preferences"] @@ -520,9 +530,9 @@ version = "0.1.11" [[deps.MPItrampoline_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"] -git-tree-sha1 = "8c35d5420193841b2f367e658540e8d9e0601ed0" +git-tree-sha1 = "70e830dab5d0775183c99fc75e4c24c614ed7142" uuid = "f1f71cc9-e9ae-5b93-9b94-4fe0e1ad3748" -version = "5.4.0+0" +version = "5.5.1+0" [[deps.MacroTools]] deps = ["Markdown", "Random"] @@ -532,9 +542,9 @@ version = "0.5.13" [[deps.MarchingCubes]] deps = ["PrecompileTools", "StaticArrays"] -git-tree-sha1 = "27d162f37cc29de047b527dab11a826dd3a650ad" +git-tree-sha1 = "301345b808264ae42e60d10a519e55c5d992969b" uuid = "299715c1-40a9-479a-aaf9-4a633d36f717" -version = "0.1.9" +version = "0.1.10" [[deps.Markdown]] deps = ["Base64"] @@ -547,9 +557,9 @@ version = "2.28.2+1" [[deps.MicrosoftMPI_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "f12a29c4400ba812841c6ace3f4efbb6dbb3ba01" +git-tree-sha1 = "bc95bf4149bf535c09602e3acdf950d9b4376227" uuid = "9237b28f-5490-5468-be7b-bb81f5f5e6cf" -version = "10.1.4+2" +version = "10.1.4+3" [[deps.Missings]] deps = ["DataAPI"] @@ -566,9 +576,9 @@ version = "2023.1.10" [[deps.NCDatasets]] deps = ["CFTime", "CommonDataModel", "DataStructures", "Dates", "DiskArrays", "NetCDF_jll", "NetworkOptions", "Printf"] -git-tree-sha1 = "77df6d3708ec0eb3441551e1f20f7503b37c2393" +git-tree-sha1 = "2c9dc92001ac06d432f363f37ff5552954d9947c" uuid = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" -version = "0.14.5" +version = "0.14.6" [[deps.NaNMath]] deps = ["OpenLibm_jll"] @@ -582,20 +592,14 @@ git-tree-sha1 = "a8af1798e4eb9ff768ce7fdefc0e957097793f15" uuid = "7243133f-43d8-5620-bbf4-c2c921802cf3" version = "400.902.209+0" -[[deps.Nettle_jll]] -deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "eca63e3847dad608cfa6a3329b95ef674c7160b4" -uuid = "4c82536e-c426-54e4-b420-14f461c4ed8b" -version = "3.7.2+0" - [[deps.NetworkOptions]] uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" version = "1.2.0" [[deps.OffsetArrays]] -git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e" +git-tree-sha1 = "5e1897147d1ff8d98883cda2be2187dcf57d8f0c" uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" -version = "1.14.1" +version = "1.15.0" weakdeps = ["Adapt"] [deps.OffsetArrays.extensions] @@ -613,15 +617,15 @@ version = "0.8.1+2" [[deps.OpenMPI_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"] -git-tree-sha1 = "bfce6d523861a6c562721b262c0d1aaeead2647f" +git-tree-sha1 = "2dace87e14256edb1dd0724ab7ba831c779b96bd" uuid = "fe0851c0-eecd-5654-98d4-656369965a5c" -version = "5.0.5+0" +version = "5.0.6+0" [[deps.OpenSSL_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "1b35263570443fdd9e76c76b7062116e2f374ab8" +git-tree-sha1 = "7493f61f55a6cce7325f197443aa80d32554ba10" uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95" -version = "3.0.15+0" +version = "3.0.15+1" [[deps.OpenSpecFun_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] @@ -630,15 +634,9 @@ uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" version = "0.5.5+0" [[deps.OrderedCollections]] -git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5" +git-tree-sha1 = "12f1439c4f986bb868acda6ea33ebc78e19b95ad" uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" -version = "1.6.3" - -[[deps.P11Kit_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "2cd396108e178f3ae8dedbd8e938a18726ab2fbf" -uuid = "c2071276-7c44-58a7-b746-946036e04d0a" -version = "0.24.1+0" +version = "1.7.0" [[deps.PDMats]] deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"] @@ -647,10 +645,12 @@ uuid = "90014a1f-27ba-587c-ab20-58faa44d9150" version = "0.11.31" [[deps.ParticleDA]] -deps = ["ChunkSplitters", "HDF5", "LinearAlgebra", "MPI", "PDMats", "Random", "Statistics", "StructArrays", "TimerOutputs", "YAML"] -git-tree-sha1 = "445106e67dc21688691ed762446d393965e1d982" +deps = ["ChunkSplitters", "HDF5", "LinearAlgebra", "MPI", "PDMats", "Random", "Statistics", "StructArrays", "Test", "TimerOutputs", "YAML"] +git-tree-sha1 = "8799e295a625c885b8f1e7ce3711ced50c07afd1" +repo-rev = "238af83" +repo-url = "https://github.com/Team-RADDISH/ParticleDA.jl" uuid = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" -version = "1.1.0" +version = "1.2.0" [[deps.Pkg]] deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] @@ -698,9 +698,9 @@ version = "1.2.1" [[deps.QuadGK]] deps = ["DataStructures", "LinearAlgebra"] -git-tree-sha1 = "1d587203cf851a51bf1ea31ad7ff89eff8d625ea" +git-tree-sha1 = "cda3b045cf9ef07a08ad46731f5a3165e56cf3da" uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" -version = "2.11.0" +version = "2.11.1" [deps.QuadGK.extensions] QuadGKEnzymeExt = "Enzyme" @@ -729,15 +729,15 @@ version = "1.3.0" [[deps.Rmath]] deps = ["Random", "Rmath_jll"] -git-tree-sha1 = "f65dcb5fa46aee0cf9ed6274ccbd597adc49aa7b" +git-tree-sha1 = "852bd0f55565a9e973fcfee83a84413270224dc4" uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa" -version = "0.7.1" +version = "0.8.0" [[deps.Rmath_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "e60724fd3beea548353984dc61c943ecddb0e29a" +git-tree-sha1 = "58cdd8fb2201a6267e1db87ff148dd6c1dbd8ad8" uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f" -version = "0.4.3+0" +version = "0.5.1+0" [[deps.SHA]] uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" @@ -762,9 +762,9 @@ version = "1.10.0" [[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" [deps.SpecialFunctions.extensions] SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" @@ -792,9 +792,9 @@ version = "0.12.1" [[deps.StaticArrays]] deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "eeafab08ae20c62c44c8399ccb9354a04b80db50" +git-tree-sha1 = "777657803913ffc7e8cc20f0fd04b634f871af8f" uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.7" +version = "1.9.8" [deps.StaticArrays.extensions] StaticArraysChainRulesCoreExt = "ChainRulesCore" @@ -821,16 +821,16 @@ uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0" version = "1.7.0" [[deps.StatsBase]] -deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"] -git-tree-sha1 = "5cf7606d6cef84b543b483848d4ae08ad9832b21" +deps = ["AliasTables", "DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"] +git-tree-sha1 = "29321314c920c26684834965ec2ce0dacc9cf8e5" uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" -version = "0.34.3" +version = "0.34.4" [[deps.StatsFuns]] deps = ["HypergeometricFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"] -git-tree-sha1 = "cef0472124fab0695b58ca35a77c6fb942fdab8a" +git-tree-sha1 = "b423576adc27097764a90e163157bcfc9acf0f46" uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c" -version = "1.3.1" +version = "1.3.2" [deps.StatsFuns.extensions] StatsFunsChainRulesCoreExt = "ChainRulesCore" @@ -896,21 +896,20 @@ git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6" uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50" version = "0.1.1" -[[deps.TestItems]] -git-tree-sha1 = "42fd9023fef18b9b78c8343a4e2f3813ffbcefcb" -uuid = "1c621080-faea-4a02-84b6-bbd5e436b8fe" -version = "1.0.0" +[[deps.Test]] +deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] +uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [[deps.TimerOutputs]] deps = ["ExprTools", "Printf"] -git-tree-sha1 = "5a13ae8a41237cff5ecf34f73eb1b8f42fff6531" +git-tree-sha1 = "d7298ebdfa1654583468a487e8e83fae9d72dac3" uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" -version = "0.5.24" +version = "0.5.26" [[deps.TranscodingStreams]] -git-tree-sha1 = "e84b3a11b9bece70d14cce63406bbc79ed3464d2" +git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742" uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" -version = "0.11.2" +version = "0.11.3" [[deps.UUIDs]] deps = ["Random", "SHA"] @@ -920,10 +919,10 @@ uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" [[deps.UnicodePlots]] -deps = ["ColorSchemes", "ColorTypes", "Contour", "Crayons", "Dates", "LinearAlgebra", "MarchingCubes", "NaNMath", "PrecompileTools", "Printf", "Requires", "SparseArrays", "StaticArrays", "StatsBase"] -git-tree-sha1 = "30646456e889c18fb3c23e58b2fc5da23644f752" +deps = ["ColorSchemes", "ColorTypes", "Contour", "Crayons", "Dates", "LinearAlgebra", "MarchingCubes", "NaNMath", "PrecompileTools", "Printf", "SparseArrays", "StaticArrays", "StatsBase"] +git-tree-sha1 = "f18128aa9e5cf059426a91bdc750b1f63a2fdcd9" uuid = "b8865327-cd53-5732-bb35-84acbb429228" -version = "3.6.4" +version = "3.7.1" [deps.UnicodePlots.extensions] FreeTypeExt = ["FileIO", "FreeType"] @@ -941,27 +940,25 @@ version = "3.6.4" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [[deps.UnsafeAtomics]] -git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278" +git-tree-sha1 = "b13c4edda90890e5b04ba24e20a310fbe6f249ff" uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f" -version = "0.2.1" +version = "0.3.0" +weakdeps = ["LLVM"] -[[deps.UnsafeAtomicsLLVM]] -deps = ["LLVM", "UnsafeAtomics"] -git-tree-sha1 = "2d17fabcd17e67d7625ce9c531fb9f40b7c42ce4" -uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249" -version = "0.2.1" + [deps.UnsafeAtomics.extensions] + UnsafeAtomicsLLVM = ["LLVM"] [[deps.XML2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"] -git-tree-sha1 = "1165b0443d0eca63ac1e32b8c0eb69ed2f4f8127" +git-tree-sha1 = "a2fccc6559132927d4c5dc183e3e01048c6dcbd6" uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a" -version = "2.13.3+0" +version = "2.13.5+0" [[deps.XZ_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "ac88fb95ae6447c8dda6a5503f3bafd496ae8632" +git-tree-sha1 = "15e637a697345f6743674f1322beefbc5dcd5cfc" uuid = "ffd25f8a-64ca-5728-b0f7-c24cf3aae800" -version = "5.4.6+0" +version = "5.6.3+0" [[deps.YAML]] deps = ["Base64", "Dates", "Printf", "StringEncodings"] @@ -976,9 +973,9 @@ version = "1.2.13+1" [[deps.Zstd_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "e678132f07ddb5bfa46857f0d7620fb9be675d3b" +git-tree-sha1 = "555d1076590a6cc2fdee2ef1469451f872d8b41b" uuid = "3161d3a3-bdf6-5164-811a-617609db77b4" -version = "1.5.6+0" +version = "1.5.6+1" [[deps.libaec_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -992,10 +989,10 @@ uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" version = "5.8.0+1" [[deps.libzip_jll]] -deps = ["Artifacts", "Bzip2_jll", "GnuTLS_jll", "JLLWrappers", "Libdl", "XZ_jll", "Zlib_jll", "Zstd_jll"] -git-tree-sha1 = "3282b7d16ae7ac3e57ec2f3fa8fafb564d8f9f7f" +deps = ["Artifacts", "Bzip2_jll", "JLLWrappers", "Libdl", "OpenSSL_jll", "XZ_jll", "Zlib_jll", "Zstd_jll"] +git-tree-sha1 = "e797fa066eba69f4c0585ffbd81bc780b5118ce2" uuid = "337d8026-41b4-5cde-a456-74a10e5b31d1" -version = "1.10.1+0" +version = "1.11.2+0" [[deps.nghttp2_jll]] deps = ["Artifacts", "Libdl"] diff --git a/SpeedyWeatherSSM/Project.toml b/SpeedyWeatherSSM/Project.toml index 8a58a32..1a0a570 100644 --- a/SpeedyWeatherSSM/Project.toml +++ b/SpeedyWeatherSSM/Project.toml @@ -10,5 +10,12 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" [compat] -ParticleDA = "1.1" +ParticleDA = "1.2" SpeedyWeather = "0.12" + +[extras] +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +ReadOnlyArrays = "988b38a3-91fc-5605-94a2-ee2116b3bd83" + +[targets] +test = ["Test", "ReadOnlyArrays"] diff --git a/SpeedyWeatherSSM/test/runtests.jl b/SpeedyWeatherSSM/test/runtests.jl new file mode 100644 index 0000000..3d54dc7 --- /dev/null +++ b/SpeedyWeatherSSM/test/runtests.jl @@ -0,0 +1,97 @@ +using ParticleDA +using Random +using Test +using ReadOnlyArrays +using SpeedyWeather +using SpeedyWeatherSSM + + +@testset ( + "equispaced_lat_lon_grid with T=$T n_lat=$n_lat, n_lon=$n_lon" + ) for T in (Float32, Float64), n_lat in (1, 3), n_lon in (1, 2) + lat_lon_pairs = SpeedyWeatherSSM.equispaced_lat_lon_grid(T, n_lat, n_lon) + @test lat_lon_pairs isa Matrix{T} + @test size(lat_lon_pairs) == (2, n_lat * n_lon) + @test all(-90 .<= lat_lon_pairs[1, :] .<= 90) + @test all(-180 .<= lat_lon_pairs[2, :] .<= 180) +end + +function check_consistency_vector_and_spectral_coefficients( + vector::Vector{T}, spectral_coefficients::LowerTriangularMatrix{Complex{T}} +) where {T <: AbstractFloat} + n_row, n_col = size(spectral_coefficients, as=Matrix) + # Test first (n_row - 1) entries corresponding to real component of first column + @test all(vector[1:n_row - 1] == real(spectral_coefficients[1:n_row - 1])) + # Test remaining entries correspond to interleaved real and imaginary components + # scanning across columns + vector_index = n_row + for col_index in 2:n_col + for row_index in col_index:(n_row - 1) + @test vector[vector_index] == real( + spectral_coefficients[row_index, col_index] + ) + @test vector[vector_index + 1] == imag( + spectral_coefficients[row_index, col_index] + ) + vector_index += 2 + end + end +end + +@testset ( + "update_spectral_coefficients_from_vector with T = $T, trunc = $spectral_truncation" +) for T in (Float32, Float64), spectral_truncation in (1, 7, 15) + rng = Random.Xoshiro(1234) + vector = randn(T, (spectral_truncation + 1)^2) + n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 + spectral_coefficients = SpeedyWeather.LowerTriangularMatrix{Complex{T}}( + undef, n_row, n_col + ) + # Set coefficients in last row to NaN - these should be replaced with zero + for col_index in 1:n_col + spectral_coefficients[n_row, col_index] = NaN + end + # Wrap vector as read-only to check not being mutated + SpeedyWeatherSSM.update_spectral_coefficients_from_vector!( + spectral_coefficients, ReadOnlyVector(vector), spectral_truncation + ) + # Check last row entries all zero + for col_index in 1:n_col + @test spectral_coefficients[n_row, col_index] == 0 + end + check_consistency_vector_and_spectral_coefficients(vector, spectral_coefficients) +end + +@testset ( + "update_vector_from_spectral_coefficients with T = $T, trunc = $spectral_truncation" +) for T in (Float32, Float64), spectral_truncation in (1, 7, 15) + rng = Random.Xoshiro(1234) + vector = Vector{T}(undef, (spectral_truncation + 1)^2) + n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 + spectral_coefficients = SpeedyWeather.LowerTriangularMatrix{Complex{T}}( + undef, n_row, n_col + ) + # First column (order = m = 0) coefficients real-valued + spectral_coefficients[1:n_row] = randn(rng, T, n_row) + # Remainder of coefficients complex valued + spectral_coefficients[n_row + 1:end] = randn( + rng, Complex{T}, length(spectral_coefficients) - n_row + ) + # Set coefficients in last row to NaN - these should be ignored + for col_index in 1:n_col + spectral_coefficients[n_row, col_index] = NaN + end + # Wrap spectral_coefficient as read-only to check not being mutated + SpeedyWeatherSSM.update_vector_from_spectral_coefficients!( + vector, ReadOnlyVector(spectral_coefficients), spectral_truncation + ) + # Check no NaNs from last row ended up in vector + @test !any(isnan.(vector)) + check_consistency_vector_and_spectral_coefficients(vector, spectral_coefficients) +end + +@testset "Generic model interface unit tests" begin + seed = 1234 + model = SpeedyWeatherSSM.init(SpeedyWeatherSSM.SpeedyParameters()) + ParticleDA.run_unit_tests_for_generic_model_interface(model, seed) +end From a93b02aafc5cc03400e46fc29178e85f10b268df Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 18 Dec 2024 18:26:49 +0000 Subject: [PATCH 18/33] Decouple update functions from SpeedyModel type --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 56 +++++++++++++++++------- 1 file changed, 40 insertions(+), 16 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index c20427f..ab55974 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -131,36 +131,48 @@ function update_spectral_coefficients_from_vector!( end function update_prognostic_variables_from_state_vector!( - model::SpeedyModel{T}, state::AbstractVector{T} + prognostic_variables::SpeedyWeather.PrognosticVariables{T}, + state::AbstractVector{T}, + variable_names::Tuple ) where {T <: AbstractFloat} start_index = 1 - dim_spectral = (model.parameters.spectral_truncation + 1)^2 - (; prognostic_variables) = model + spectral_truncation = prognostic_variables.trunc + dim_spectral = (spectral_truncation + 1)^2 for name in LAYERED_VARIABLES - if SpeedyWeather.has(model.model, name) + if name in variable_names # We only consider spectral coefficients for first leapfrog step (lf=1) to # define state layered_spectral_coefficients = getproperty(prognostic_variables, name)[1] - for layer_index in 1:model.parameters.n_layers + for layer_index in 1:prognostic_variables.nlayers end_index = start_index + dim_spectral - 1 update_spectral_coefficients_from_vector!( view(layered_spectral_coefficients, :, layer_index), view(state, start_index:end_index), - model.parameters.spectral_truncation + spectral_truncation ) start_index = end_index + 1 end end end - if SpeedyWeather.has(model.model, :pres) + if :pres in variable_names update_spectral_coefficients_from_vector!( prognostic_variables.pres[1], view(state, start_index:start_index + dim_spectral - 1), - model.parameters.spectral_truncation + spectral_truncation ) end end +function update_prognostic_variables_from_state_vector!( + model::SpeedyModel{T}, state::AbstractVector{T} +) where {T <: AbstractFloat} + update_prognostic_variables_from_state_vector!( + model.prognostic_variables, + state, + SpeedyWeather.prognostic_variables(model.model) + ) +end + function update_vector_from_spectral_coefficients!( vector::AbstractVector{T}, spectral_coefficients::AbstractVector{Complex{T}}, @@ -193,36 +205,48 @@ function update_vector_from_spectral_coefficients!( end function update_state_vector_from_prognostic_variables!( - state::AbstractVector{T}, model::SpeedyModel{T}, + state::AbstractVector{T}, + prognostic_variables::SpeedyWeather.PrognosticVariables{T}, + variable_names::Tuple ) where {T <: AbstractFloat} start_index = 1 - dim_spectral = (model.parameters.spectral_truncation + 1)^2 - (; prognostic_variables) = model + spectral_truncation = prognostic_variables.trunc + dim_spectral = (spectral_truncation + 1)^2 for name in LAYERED_VARIABLES - if SpeedyWeather.has(model.model, name) + if name in variable_names # We only consider spectral coefficients for first leapfrog step (lf=1) to # define state layered_spectral_coefficients = getproperty(prognostic_variables, name)[1] - for layer_index in 1:model.parameters.n_layers + for layer_index in 1:prognostic_variables.nlayers end_index = start_index + dim_spectral - 1 update_vector_from_spectral_coefficients!( view(state, start_index:end_index), view(layered_spectral_coefficients, :, layer_index), - model.parameters.spectral_truncation + spectral_truncation ) start_index = end_index + 1 end end end - if SpeedyWeather.has(model.model, :pres) + if :pres in variable_names update_vector_from_spectral_coefficients!( view(state, start_index:start_index + dim_spectral - 1), prognostic_variables.pres[1], - model.parameters.spectral_truncation + spectral_truncation ) end end +function update_state_vector_from_prognostic_variables!( + state::AbstractVector{T}, model::SpeedyModel{T} +) where {T <: AbstractFloat} + update_state_vector_from_prognostic_variables!( + state, + model.prognostic_variables, + SpeedyWeather.prognostic_variables(model.model) + ) +end + ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T From c3e8bf8ae43dc7f89170852dc203e712547f9360 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 18 Dec 2024 18:27:08 +0000 Subject: [PATCH 19/33] Add tests for state update functions --- SpeedyWeatherSSM/test/runtests.jl | 47 +++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/SpeedyWeatherSSM/test/runtests.jl b/SpeedyWeatherSSM/test/runtests.jl index 3d54dc7..2becb45 100644 --- a/SpeedyWeatherSSM/test/runtests.jl +++ b/SpeedyWeatherSSM/test/runtests.jl @@ -89,6 +89,53 @@ end @test !any(isnan.(vector)) check_consistency_vector_and_spectral_coefficients(vector, spectral_coefficients) end + +@testset ( + "Updating prognostic variables to/from state vector with \ + T=$T, spectral_truncation=$spectral_truncation, n_layers=$n_layers" + ) for T in (Float32, Float64), spectral_truncation in (7, 15), n_layers in (1, 8) + rng = Random.Xoshiro(1234) + spectral_grid = SpectralGrid(; + NF=T, trunc=spectral_truncation, nlayers=n_layers + ) + layered_variables = (:vor, :div, :temp, :humid) + surface_variables = (:pres,) + all_variables = (layered_variables..., surface_variables...) + state_dimension = (spectral_truncation + 1)^2 * ( + n_layers * length(layered_variables) + length(surface_variables) + ) + state = rand(rng, T, state_dimension) + prognostic_variables_1 = SpeedyWeather.PrognosticVariables(spectral_grid) + prognostic_variables_1_init = SpeedyWeather.PrognosticVariables(spectral_grid) + SpeedyWeather.copy!(prognostic_variables_1_init, prognostic_variables_1) + prognostic_variables_2 = SpeedyWeather.PrognosticVariables(spectral_grid) + SpeedyWeatherSSM.update_prognostic_variables_from_state_vector!( + prognostic_variables_1, state, all_variables + ) + SpeedyWeatherSSM.update_prognostic_variables_from_state_vector!( + prognostic_variables_2, state, all_variables + ) + # Check all variables changed from initial values + for variable_name in all_variables + @test all( + getproperty(prognostic_variables_1, variable_name)[1] + != getproperty(prognostic_variables_1_init, variable_name)[1] + ) + end + # Check variables updates from same state vector match + for variable_name in all_variables + @test all( + getproperty(prognostic_variables_1, variable_name)[1] + == getproperty(prognostic_variables_2, variable_name)[1] + ) + end + # Check state vector reconstructed from prognostic variables matches original + state_2 = Vector{T}(undef, state_dimension) + SpeedyWeatherSSM.update_state_vector_from_prognostic_variables!( + state_2, prognostic_variables_1, all_variables + ) + @test all(state == state_2) +end @testset "Generic model interface unit tests" begin seed = 1234 From c07a9000b90ef5a76d9caa4aa724fe6900ee9274 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 19 Dec 2024 09:21:36 +0000 Subject: [PATCH 20/33] Expose leapfrog step to update as argument --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index ab55974..1872cc0 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -133,16 +133,17 @@ end function update_prognostic_variables_from_state_vector!( prognostic_variables::SpeedyWeather.PrognosticVariables{T}, state::AbstractVector{T}, - variable_names::Tuple + variable_names::Tuple; + leapfrog_step::Int = 1 ) where {T <: AbstractFloat} start_index = 1 spectral_truncation = prognostic_variables.trunc dim_spectral = (spectral_truncation + 1)^2 for name in LAYERED_VARIABLES if name in variable_names - # We only consider spectral coefficients for first leapfrog step (lf=1) to - # define state - layered_spectral_coefficients = getproperty(prognostic_variables, name)[1] + layered_spectral_coefficients = getproperty( + prognostic_variables, name + )[leapfrog_step] for layer_index in 1:prognostic_variables.nlayers end_index = start_index + dim_spectral - 1 update_spectral_coefficients_from_vector!( @@ -156,7 +157,7 @@ function update_prognostic_variables_from_state_vector!( end if :pres in variable_names update_spectral_coefficients_from_vector!( - prognostic_variables.pres[1], + prognostic_variables.pres[leapfrog_step], view(state, start_index:start_index + dim_spectral - 1), spectral_truncation ) @@ -207,7 +208,8 @@ end function update_state_vector_from_prognostic_variables!( state::AbstractVector{T}, prognostic_variables::SpeedyWeather.PrognosticVariables{T}, - variable_names::Tuple + variable_names::Tuple; + leapfrog_step::Int = 1 ) where {T <: AbstractFloat} start_index = 1 spectral_truncation = prognostic_variables.trunc @@ -216,7 +218,9 @@ function update_state_vector_from_prognostic_variables!( if name in variable_names # We only consider spectral coefficients for first leapfrog step (lf=1) to # define state - layered_spectral_coefficients = getproperty(prognostic_variables, name)[1] + layered_spectral_coefficients = getproperty( + prognostic_variables, name + )[leapfrog_step] for layer_index in 1:prognostic_variables.nlayers end_index = start_index + dim_spectral - 1 update_vector_from_spectral_coefficients!( @@ -231,7 +235,7 @@ function update_state_vector_from_prognostic_variables!( if :pres in variable_names update_vector_from_spectral_coefficients!( view(state, start_index:start_index + dim_spectral - 1), - prognostic_variables.pres[1], + prognostic_variables.pres[leapfrog_step], spectral_truncation ) end From 94023e8bdd84eb6c3388204fc3a5bffbab31937c Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 19 Dec 2024 09:56:57 +0000 Subject: [PATCH 21/33] Generalise treatment of surface variables --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 36 +++++++++++++++--------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 1872cc0..a929b44 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -155,12 +155,17 @@ function update_prognostic_variables_from_state_vector!( end end end - if :pres in variable_names - update_spectral_coefficients_from_vector!( - prognostic_variables.pres[leapfrog_step], - view(state, start_index:start_index + dim_spectral - 1), - spectral_truncation - ) + for name in SURFACE_VARIABLES + if name in variable_names + spectral_coefficients = getproperty( + prognostic_variables, name + )[leapfrog_step] + update_spectral_coefficients_from_vector!( + spectral_coefficients, + view(state, start_index:start_index + dim_spectral - 1), + spectral_truncation + ) + end end end @@ -216,8 +221,6 @@ function update_state_vector_from_prognostic_variables!( dim_spectral = (spectral_truncation + 1)^2 for name in LAYERED_VARIABLES if name in variable_names - # We only consider spectral coefficients for first leapfrog step (lf=1) to - # define state layered_spectral_coefficients = getproperty( prognostic_variables, name )[leapfrog_step] @@ -232,12 +235,17 @@ function update_state_vector_from_prognostic_variables!( end end end - if :pres in variable_names - update_vector_from_spectral_coefficients!( - view(state, start_index:start_index + dim_spectral - 1), - prognostic_variables.pres[leapfrog_step], - spectral_truncation - ) + for name in SURFACE_VARIABLES + if name in variable_names + spectral_coefficients = getproperty( + prognostic_variables, name + )[leapfrog_step] + update_vector_from_spectral_coefficients!( + view(state, start_index:start_index + dim_spectral - 1), + spectral_coefficients, + spectral_truncation + ) + end end end From 6442e67f78c6d80e11b99e4efe1e57cad7986830 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 19 Dec 2024 13:47:05 +0000 Subject: [PATCH 22/33] Factor out clock update function --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index a929b44..2eb97d8 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -73,6 +73,10 @@ function init(parameters::SpeedyParameters{T, M}) where { parameters.observed_coordinates[1, :], parameters.observed_coordinates[2, :] ) + SpeedyWeather.set_period!( + prognostic_variables.clock, SpeedyWeather.Day(parameters.n_days) + ) + SpeedyWeather.initialize!(prognostic_variables.clock, model.time_stepping) return SpeedyModel( parameters, spectral_grid, @@ -259,6 +263,15 @@ function update_state_vector_from_prognostic_variables!( ) end +function update_clock_from_time_index!(clock::SpeedyWeather.Clock,time_index::Int) + clock.time = clock.start + clock.n_timesteps * clock.Δt * (time_index - 1) + clock.timestep_counter = clock.n_timesteps * (time_index - 1) +end + +function update_clock_from_time_index!(model::SpeedyModel, time_index::Int) + update_clock_from_time_index!(model.prognostic_variables.clock, time_index) +end + ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T @@ -277,13 +290,7 @@ function ParticleDA.update_state_deterministic!( state::AbstractVector{T}, model::SpeedyModel{T}, time_index::Int, task_index::Int=1 ) where {T<:AbstractFloat} update_prognostic_variables_from_state_vector!(model, state) - (; clock) = model.prognostic_variables - (; time_stepping) = model.model - SpeedyWeather.set_period!(clock, SpeedyWeather.Day(model.parameters.n_days)) - SpeedyWeather.initialize!(clock, time_stepping) - clock.start = model.parameters.start_date - clock.time = clock.start + clock.n_timesteps * clock.Δt * (time_index - 1) - clock.timestep_counter = clock.n_timesteps * (time_index - 1) + update_clock_from_time_index!(model, time_index) SpeedyWeather.time_stepping!( model.prognostic_variables, model.diagnostic_variables, model.model ) From 21248a09774f5c8603d2c8377ca9830ad21853cc Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 19 Dec 2024 14:01:03 +0000 Subject: [PATCH 23/33] Factor out function to update diagnostic variables from state --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 38 +++++++++++++----------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 2eb97d8..b981b6a 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -272,6 +272,26 @@ function update_clock_from_time_index!(model::SpeedyModel, time_index::Int) update_clock_from_time_index!(model.prognostic_variables.clock, time_index) end +function get_observed_variable_field( + diagnostic_variables::DiagnosticVariables, model::SpeedyModel +) + observed_outer, observed_inner = model.parameters.observed_variable + getfield(getfield(diagnostic_variables, observed_outer), observed_inner) +end + +function update_prognostic_and_diagnostic_variables_from_state_vector!( + model::SpeedyModel{T}, state::AbstractVector{T} +) where {T <: AbstractFloat} + update_prognostic_variables_from_state_vector!(model, state) + SpeedyWeather.transform!( + model.diagnostic_variables, + model.prognostic_variables, + 1, + model.model, + initialize=true + ) +end + ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T @@ -303,29 +323,13 @@ function ParticleDA.update_state_stochastic!( end -function get_observed_variable_field( - diagnostic_variables::DiagnosticVariables, model::SpeedyModel -) - observed_outer, observed_inner = model.parameters.observed_variable - getfield( - getfield(diagnostic_variables, observed_outer), observed_inner - ) -end - function ParticleDA.get_observation_mean_given_state!( observation_mean::AbstractVector{T}, state::AbstractVector{T}, model::SpeedyModel{T}, task_index::Int=1 ) where {T<:AbstractFloat} - update_prognostic_variables_from_state_vector!(model, state) - SpeedyWeather.transform!( - model.diagnostic_variables, - model.prognostic_variables, - 1, - model.model, - initialize=true - ) + update_prognostic_and_diagnostic_variables_from_state_vector!(model, state) observed_field_grid = get_observed_variable_field(model.diagnostic_variables, model) SpeedyWeather.interpolate!( observation_mean, observed_field_grid, model.observation_interpolator From bbecbb752a6d04feb18eb4e00f3e8cf7b4a18f9d Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 19 Dec 2024 18:31:23 +0000 Subject: [PATCH 24/33] Zero second leapfrog step coefficients when updating prognostic variables --- SpeedyWeatherSSM/Manifest.toml | 2 +- SpeedyWeatherSSM/Project.toml | 3 ++- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 11 ++++++++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/SpeedyWeatherSSM/Manifest.toml b/SpeedyWeatherSSM/Manifest.toml index e40b62c..f7455a0 100644 --- a/SpeedyWeatherSSM/Manifest.toml +++ b/SpeedyWeatherSSM/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.4" manifest_format = "2.0" -project_hash = "8a48103e70f2b517a2b98bb799150c3116af6c64" +project_hash = "4d7cc7338286b6f128a00d1e215bbc942db42e27" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] diff --git a/SpeedyWeatherSSM/Project.toml b/SpeedyWeatherSSM/Project.toml index 1a0a570..670face 100644 --- a/SpeedyWeatherSSM/Project.toml +++ b/SpeedyWeatherSSM/Project.toml @@ -4,6 +4,7 @@ authors = ["Matt Graham"] version = "0.1.0" [deps] +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ParticleDA = "61cd1fb4-f4c4-4bc8-80c6-ea5639a6ca2e" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -14,8 +15,8 @@ ParticleDA = "1.2" SpeedyWeather = "0.12" [extras] -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" ReadOnlyArrays = "988b38a3-91fc-5605-94a2-ee2116b3bd83" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Test", "ReadOnlyArrays"] diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index b981b6a..ace8304 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -1,5 +1,6 @@ module SpeedyWeatherSSM +using FillArrays using HDF5 using ParticleDA using Random @@ -179,7 +180,15 @@ function update_prognostic_variables_from_state_vector!( update_prognostic_variables_from_state_vector!( model.prognostic_variables, state, - SpeedyWeather.prognostic_variables(model.model) + SpeedyWeather.prognostic_variables(model.model); + leapfrog_step=1 + ) + # Zero cofficients for second leapfrog step (corresponding to initial state) + update_prognostic_variables_from_state_vector!( + model.prognostic_variables, + Zeros(ParticleDA.get_state_dimension(model)), + SpeedyWeather.prognostic_variables(model.model); + leapfrog_step=2 ) end From e226ea2e58a561041ef3d62677cf8b064cf416f3 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 20 Dec 2024 13:55:12 +0000 Subject: [PATCH 25/33] Factor out function for mapping over spectral coefficients and state vector slices --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 73 ++++++++++-------------- 1 file changed, 31 insertions(+), 42 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index ace8304..3cbe0cf 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -135,11 +135,12 @@ function update_spectral_coefficients_from_vector!( end end -function update_prognostic_variables_from_state_vector!( +function map_over_spectral_coefficients_and_state_vector_slices( + map_function::Function, prognostic_variables::SpeedyWeather.PrognosticVariables{T}, state::AbstractVector{T}, - variable_names::Tuple; - leapfrog_step::Int = 1 + variable_names::Tuple, + leapfrog_step::Int ) where {T <: AbstractFloat} start_index = 1 spectral_truncation = prognostic_variables.trunc @@ -151,10 +152,9 @@ function update_prognostic_variables_from_state_vector!( )[leapfrog_step] for layer_index in 1:prognostic_variables.nlayers end_index = start_index + dim_spectral - 1 - update_spectral_coefficients_from_vector!( + map_function( view(layered_spectral_coefficients, :, layer_index), - view(state, start_index:end_index), - spectral_truncation + view(state, start_index:end_index) ) start_index = end_index + 1 end @@ -162,18 +162,30 @@ function update_prognostic_variables_from_state_vector!( end for name in SURFACE_VARIABLES if name in variable_names - spectral_coefficients = getproperty( - prognostic_variables, name - )[leapfrog_step] - update_spectral_coefficients_from_vector!( - spectral_coefficients, + map_function( + getproperty(prognostic_variables, name)[leapfrog_step], view(state, start_index:start_index + dim_spectral - 1), - spectral_truncation ) end end end +function update_prognostic_variables_from_state_vector!( + prognostic_variables::SpeedyWeather.PrognosticVariables{T}, + state::AbstractVector{T}, + variable_names::Tuple; + leapfrog_step::Int = 1 +) where {T <: AbstractFloat} + spectral_truncation = prognostic_variables.trunc + map_over_spectral_coefficients_and_state_vector_slices( + (c, v) -> update_spectral_coefficients_from_vector!(c, v, spectral_truncation), + prognostic_variables, + state, + variable_names, + leapfrog_step + ) +end + function update_prognostic_variables_from_state_vector!( model::SpeedyModel{T}, state::AbstractVector{T} ) where {T <: AbstractFloat} @@ -229,37 +241,14 @@ function update_state_vector_from_prognostic_variables!( variable_names::Tuple; leapfrog_step::Int = 1 ) where {T <: AbstractFloat} - start_index = 1 spectral_truncation = prognostic_variables.trunc - dim_spectral = (spectral_truncation + 1)^2 - for name in LAYERED_VARIABLES - if name in variable_names - layered_spectral_coefficients = getproperty( - prognostic_variables, name - )[leapfrog_step] - for layer_index in 1:prognostic_variables.nlayers - end_index = start_index + dim_spectral - 1 - update_vector_from_spectral_coefficients!( - view(state, start_index:end_index), - view(layered_spectral_coefficients, :, layer_index), - spectral_truncation - ) - start_index = end_index + 1 - end - end - end - for name in SURFACE_VARIABLES - if name in variable_names - spectral_coefficients = getproperty( - prognostic_variables, name - )[leapfrog_step] - update_vector_from_spectral_coefficients!( - view(state, start_index:start_index + dim_spectral - 1), - spectral_coefficients, - spectral_truncation - ) - end - end + map_over_spectral_coefficients_and_state_vector_slices( + (c, v) -> update_vector_from_spectral_coefficients!(v, c, spectral_truncation), + prognostic_variables, + state, + variable_names, + leapfrog_step + ) end function update_state_vector_from_prognostic_variables!( From 27a6bfd273f4041355c692caad4ee63701271406 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 5 Feb 2025 18:21:19 +0000 Subject: [PATCH 26/33] Further generalise function for mapping over slices --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 43 ++++++++++++++++++------ 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 3cbe0cf..da87629 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -135,26 +135,23 @@ function update_spectral_coefficients_from_vector!( end end -function map_over_spectral_coefficients_and_state_vector_slices( +function map_over_state_vector_slices( map_function::Function, - prognostic_variables::SpeedyWeather.PrognosticVariables{T}, state::AbstractVector{T}, variable_names::Tuple, - leapfrog_step::Int + spectral_truncation::Int, + n_layers::Int, ) where {T <: AbstractFloat} start_index = 1 - spectral_truncation = prognostic_variables.trunc dim_spectral = (spectral_truncation + 1)^2 for name in LAYERED_VARIABLES if name in variable_names - layered_spectral_coefficients = getproperty( - prognostic_variables, name - )[leapfrog_step] - for layer_index in 1:prognostic_variables.nlayers + for layer_index in 1:n_layers end_index = start_index + dim_spectral - 1 map_function( - view(layered_spectral_coefficients, :, layer_index), - view(state, start_index:end_index) + view(state, start_index:end_index), + name, + layer_index, ) start_index = end_index + 1 end @@ -163,13 +160,37 @@ function map_over_spectral_coefficients_and_state_vector_slices( for name in SURFACE_VARIABLES if name in variable_names map_function( - getproperty(prognostic_variables, name)[leapfrog_step], view(state, start_index:start_index + dim_spectral - 1), + name, + nothing, ) end end end +function map_over_spectral_coefficients_and_state_vector_slices( + map_function::Function, + prognostic_variables::SpeedyWeather.PrognosticVariables{T}, + state::AbstractVector{T}, + variable_names::Tuple, + leapfrog_step::Int +) where {T <: AbstractFloat} + function outer_map_function(state_slice, name, layer_index) + spectral_coefficients = getproperty(prognostic_variables, name)[leapfrog_step] + if !isnothing(layer_index) + spectral_coefficients = view(spectral_coefficients, :, layer_index) + end + map_function(spectral_coefficients, state_slice) + end + map_over_state_vector_slices( + outer_map_function, + state, + variable_names, + prognostic_variables.trunc, + prognostic_variables.nlayers, + ) +end + function update_prognostic_variables_from_state_vector!( prognostic_variables::SpeedyWeather.PrognosticVariables{T}, state::AbstractVector{T}, From 7db36bff81ec7ad518f07af4ad2a04f9386d8fdf Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 5 Feb 2025 18:21:46 +0000 Subject: [PATCH 27/33] Add stochastic state updates and initialisation --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 166 +++++++++++++++++++++-- 1 file changed, 154 insertions(+), 12 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index da87629..3b50921 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -19,6 +19,11 @@ function equispaced_lat_lon_grid(T, n_lat, n_lon) collect(reshape(reinterpret(T, lat_lon_pairs), (2, :))) end +Base.@kwdef struct GaussianRandomFieldParameters{T<:AbstractFloat} + output_scale::T = 1. + length_scale::T = 10. +end + Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractModel} spectral_truncation::Int = 31 n_layers::Int = 8 @@ -29,23 +34,36 @@ Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractM observed_variable::Tuple{Symbol, Symbol} = (:physics, :precip_large_scale) observed_coordinates::Matrix{T} = equispaced_lat_lon_grid(float_type, 6, 12) observation_noise_std::T = 0.1 + # TODO: Set more sensible per variable scales + initial_state_grf_parameters::Dict{Symbol, GaussianRandomFieldParameters{T}} = Dict( + + name => GaussianRandomFieldParameters(; output_scale=1e-6) + for name in (LAYERED_VARIABLES..., SURFACE_VARIABLES...) + ) + state_noise_grf_parameters::Dict{Symbol, GaussianRandomFieldParameters{T}} = Dict( + name => GaussianRandomFieldParameters(; output_scale=1e-7) + for name in (LAYERED_VARIABLES..., SURFACE_VARIABLES...) + ) end struct SpeedyModel{ T<:AbstractFloat, G<:SpeedyWeather.AbstractSpectralGrid, M<:SpeedyWeather.AbstractModel, - I<:SpeedyWeather.RingGrids.AbstractInterpolator + I<:SpeedyWeather.RingGrids.AbstractInterpolator, } parameters::SpeedyParameters{T, M} spectral_grid::G model::M prognostic_variables::PrognosticVariables{T} diagnostic_variables::DiagnosticVariables{T} + variable_names::Tuple n_layered_variables::Int n_surface_variables::Int n_observed_points::Int observation_interpolator::I + initial_state_grf_scale_factors::Dict{Symbol, Vector{T}} + state_noise_grf_scale_factors::Dict{Symbol, Vector{T}} end function init(parameters::SpeedyParameters{T, M}) where { @@ -59,6 +77,7 @@ function init(parameters::SpeedyParameters{T, M}) where { model.output.active = false simulation = initialize!(model; time=parameters.start_date) (; prognostic_variables, diagnostic_variables) = simulation + variable_names = SpeedyWeather.prognostic_variables(model) n_layered_variables = count( SpeedyWeather.has(model, var) for var in LAYERED_VARIABLES ) @@ -78,16 +97,35 @@ function init(parameters::SpeedyParameters{T, M}) where { prognostic_variables.clock, SpeedyWeather.Day(parameters.n_days) ) SpeedyWeather.initialize!(prognostic_variables.clock, model.time_stepping) + initial_state_grf_scale_factors = Dict( + name => get_grf_coefficient_scale_factors( + parameters.initial_state_grf_parameters[name], + parameters.spectral_truncation, + model.spectral_transform.norm_sphere + ) + for name in variable_names + ) + state_noise_grf_scale_factors = Dict( + name => get_grf_coefficient_scale_factors( + parameters.state_noise_grf_parameters[name], + parameters.spectral_truncation, + model.spectral_transform.norm_sphere + ) + for name in variable_names + ) return SpeedyModel( parameters, spectral_grid, model, prognostic_variables, diagnostic_variables, + variable_names, n_layered_variables, n_surface_variables, n_observed_points, - observation_interpolator + observation_interpolator, + initial_state_grf_scale_factors, + state_noise_grf_scale_factors ) end @@ -213,14 +251,14 @@ function update_prognostic_variables_from_state_vector!( update_prognostic_variables_from_state_vector!( model.prognostic_variables, state, - SpeedyWeather.prognostic_variables(model.model); + model.variable_names; leapfrog_step=1 ) # Zero cofficients for second leapfrog step (corresponding to initial state) update_prognostic_variables_from_state_vector!( model.prognostic_variables, Zeros(ParticleDA.get_state_dimension(model)), - SpeedyWeather.prognostic_variables(model.model); + model.variable_names; leapfrog_step=2 ) end @@ -228,14 +266,20 @@ end function update_vector_from_spectral_coefficients!( vector::AbstractVector{T}, spectral_coefficients::AbstractVector{Complex{T}}, - spectral_truncation::Int + spectral_truncation::Int; + increment::Bool = false ) where {T <: AbstractFloat} n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 # First column of spectral_coefficients (order = m = 0) are real-valued and we skip # last row (degree = l = n_row - 1) as used only for computing meridional derivative # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing # skipping zero upper-triangular entries - vector[1:n_row - 1] .= real(spectral_coefficients[1:n_row - 1]) + # TODO: Figure out how to do this without repetition here and below + if increment + vector[1:n_row - 1] .+= real(spectral_coefficients[1:n_row - 1]) + else + vector[1:n_row - 1] .= real(spectral_coefficients[1:n_row - 1]) + end # vector index is i, spectral coefficient (flat) index is j i = n_row - 1 j = n_row @@ -246,9 +290,15 @@ function update_vector_from_spectral_coefficients!( slice_size = n_row - col_index # Reinterpret complex valued spectral coefficients to extract both real and # imaginary components - vector[i + 1:i + 2 * slice_size] .= reinterpret( - T, spectral_coefficients[j + 1:j + slice_size] - ) + if increment + vector[i + 1:i + 2 * slice_size] .+= reinterpret( + T, spectral_coefficients[j + 1:j + slice_size] + ) + else + vector[i + 1:i + 2 * slice_size] .= reinterpret( + T, spectral_coefficients[j + 1:j + slice_size] + ) + end # Update vector and spectral coefficient indices, adding 1 offset # to latter to skip entries corresponding to last row i = i + 2 * slice_size @@ -278,7 +328,7 @@ function update_state_vector_from_prognostic_variables!( update_state_vector_from_prognostic_variables!( state, model.prognostic_variables, - SpeedyWeather.prognostic_variables(model.model) + model.variable_names ) end @@ -311,6 +361,75 @@ function update_prognostic_and_diagnostic_variables_from_state_vector!( ) end +function add_noise_to_state_vector!( + state::AbstractVector{T}, + spectral_truncation::Int, + n_layers::Int, + variable_names::Tuple, + grf_scale_factors::Dict{Symbol, Vector{T}}, + rng::AbstractRNG +) where {T <: AbstractFloat} + n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 + spectral_coefficients = SpeedyWeather.LowerTriangularMatrix{Complex{T}}( + undef, n_row, n_col + ) + function map_function(state_slice, name, layer_index) + generate_random_spectral_coefficients!( + spectral_coefficients, + spectral_truncation, + grf_scale_factors[name], + rng + ) + update_vector_from_spectral_coefficients!( + state_slice, spectral_coefficients, spectral_truncation; increment=true + ) + end + map_over_state_vector_slices( + map_function, + state, + variable_names, + spectral_truncation, + n_layers, + ) +end + +function get_grf_coefficient_scale_factors( + parameters::GaussianRandomFieldParameters{T}, + spectral_truncation::Int, + norm_sphere::T +) where {T <: AbstractFloat} + ell_max = spectral_truncation + 2 + scale_factors = zeros(T, ell_max) + denominator = 2 * sum( + [ + (2 * ell + 1) * exp(-ell * (ell + 1) / parameters.length_scale^2) + for ell in 1:spectral_truncation + ] + ) + common_scale = norm_sphere * sqrt(2 * parameters.output_scale^2 / denominator) + for ell in 2:ell_max + scale_factors[ell] = common_scale * exp( + -ell * (ell - 1) / (2 * parameters.length_scale^2) + ) + end + return scale_factors +end + +function generate_random_spectral_coefficients!( + spectral_coefficients::AbstractVector{Complex{T}}, + spectral_truncation::Int, + scale_factors::AbstractVector{T}, + rng::AbstractRNG +) where {T <: AbstractFloat} + ell_m = 0 + @inbounds for m in 1:spectral_truncation + 1 + for ell in m:spectral_truncation + 2 + ell_m += 1 + spectral_coefficients[ell_m] = scale_factors[ell] * randn(rng, Complex{T}) + end + end +end + ParticleDA.get_state_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat} = T @@ -323,6 +442,14 @@ function ParticleDA.sample_initial_state!( model.prognostic_variables, initial_conditions, model.model ) update_state_vector_from_prognostic_variables!(state, model) + add_noise_to_state_vector!( + state, + model.parameters.spectral_truncation, + model.parameters.n_layers, + model.variable_names, + model.initial_state_grf_scale_factors, + rng + ) end function ParticleDA.update_state_deterministic!( @@ -339,7 +466,14 @@ end function ParticleDA.update_state_stochastic!( state::AbstractVector{T}, model::SpeedyModel{T}, rng::G, task_index::Int=1 ) where {T<:AbstractFloat, G<:AbstractRNG} - + add_noise_to_state_vector!( + state, + model.parameters.spectral_truncation, + model.parameters.n_layers, + model.variable_names, + model.state_noise_grf_scale_factors, + rng + ) end function ParticleDA.get_observation_mean_given_state!( @@ -395,7 +529,15 @@ function ParticleDA.write_model_metadata(file::HDF5.File, model::SpeedyModel) isa(value, Type) && (value = string(nameof(value))) isa(value, Tuple{Symbol, Symbol}) && (value = join(map(String, value), ".")) isa(value, DateTime) && (value = string(value)) - HDF5.attributes(group)[string(field)] = value + if isa(value, Dict) + subgroup = create_group(group, string(field)) + for (key, val) in value + # TODO: Write struct val in a nicer way + HDF5.attributes(subgroup)[string(key)] = string(val) + end + else + HDF5.attributes(group)[string(field)] = value + end end else @warn "Write failed, group $group_name already exists in $(file.filename)!" From 0da51243f6810969dd9a0b30d9dcd045528e4c4c Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 6 Feb 2025 08:53:18 +0000 Subject: [PATCH 28/33] Adjust default noise scales to be per variable --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 3b50921..c5b9399 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -21,7 +21,7 @@ end Base.@kwdef struct GaussianRandomFieldParameters{T<:AbstractFloat} output_scale::T = 1. - length_scale::T = 10. + length_scale::T = 1. end Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractModel} @@ -34,15 +34,19 @@ Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractM observed_variable::Tuple{Symbol, Symbol} = (:physics, :precip_large_scale) observed_coordinates::Matrix{T} = equispaced_lat_lon_grid(float_type, 6, 12) observation_noise_std::T = 0.1 - # TODO: Set more sensible per variable scales initial_state_grf_parameters::Dict{Symbol, GaussianRandomFieldParameters{T}} = Dict( - - name => GaussianRandomFieldParameters(; output_scale=1e-6) - for name in (LAYERED_VARIABLES..., SURFACE_VARIABLES...) + :vor => GaussianRandomFieldParameters(; output_scale=5e-7), + :div => GaussianRandomFieldParameters(; output_scale=5e-7), + :temp => GaussianRandomFieldParameters(; output_scale=2e0), + :humid => GaussianRandomFieldParameters(; output_scale=1e-5), + :pres => GaussianRandomFieldParameters(; output_scale=2e-3), ) state_noise_grf_parameters::Dict{Symbol, GaussianRandomFieldParameters{T}} = Dict( - name => GaussianRandomFieldParameters(; output_scale=1e-7) - for name in (LAYERED_VARIABLES..., SURFACE_VARIABLES...) + :vor => GaussianRandomFieldParameters(; output_scale=5e-8), + :div => GaussianRandomFieldParameters(; output_scale=5e-8), + :temp => GaussianRandomFieldParameters(; output_scale=2e-1), + :humid => GaussianRandomFieldParameters(; output_scale=1e-6), + :pres => GaussianRandomFieldParameters(; output_scale=2e-4), ) end From f2665560f04b09d1e5db683c601ee76922a921d1 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 6 Feb 2025 10:00:35 +0000 Subject: [PATCH 29/33] Refactor model metadata write to dispatch on val type --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 61 ++++++++++++++++-------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index c5b9399..f55f5b0 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -524,28 +524,47 @@ function ParticleDA.get_log_density_observation_given_state( end -function ParticleDA.write_model_metadata(file::HDF5.File, model::SpeedyModel) - group_name = "parameters" - if !haskey(file, group_name) - group = create_group(file, group_name) - for field in fieldnames(typeof(model.parameters)) - value = getfield(model.parameters, field) - isa(value, Type) && (value = string(nameof(value))) - isa(value, Tuple{Symbol, Symbol}) && (value = join(map(String, value), ".")) - isa(value, DateTime) && (value = string(value)) - if isa(value, Dict) - subgroup = create_group(group, string(field)) - for (key, val) in value - # TODO: Write struct val in a nicer way - HDF5.attributes(subgroup)[string(key)] = string(val) - end - else - HDF5.attributes(group)[string(field)] = value - end - end - else - @warn "Write failed, group $group_name already exists in $(file.filename)!" +const HDF5FileOrGroup = Union{HDF5.File, HDF5.Group} + +function write!( + group::HDF5FileOrGroup, key::String, value::Union{Number, String, Array} +) + attributes(group)[key] = value +end + +function write!(group::HDF5FileOrGroup, key::String, value::Union{Symbol, DateTime}) + attributes(group)[key] = string(value) +end + +function write!(group::HDF5FileOrGroup, key::String, value::Type) + attributes(group)[key] = string(nameof(value)) +end + +function write!( + group::HDF5FileOrGroup, key::String, value::NTuple{N, T} +) where {N, T <: Union{Number, String, Symbol, Type}} + subgroup = create_group(group, key) + for (index, val) in enumerate(value) + write!(subgroup, string(index), val) + end +end + +function write!(group::HDF5FileOrGroup, key::String, value::Dict) + subgroup = create_group(group, key) + for (k, v) in value + write!(subgroup, string(k), v) + end +end + +function write!(group::HDF5FileOrGroup, key::String, value) + subgroup = create_group(group, key) + for name in fieldnames(typeof(value)) + write!(subgroup, string(name), getfield(value, name)) end end +function ParticleDA.write_model_metadata(file::HDF5.File, model::SpeedyModel) + write!(file, "parameters", model.parameters) +end + end From 1dd4f0328b8abd6ae8ae0e3da0ee8255089cbe09 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 6 Feb 2025 11:09:26 +0000 Subject: [PATCH 30/33] Fix comment typo and whitespacing --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index f55f5b0..3b43235 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -258,7 +258,7 @@ function update_prognostic_variables_from_state_vector!( model.variable_names; leapfrog_step=1 ) - # Zero cofficients for second leapfrog step (corresponding to initial state) + # Zero coefficients for second leapfrog step (corresponding to initial state) update_prognostic_variables_from_state_vector!( model.prognostic_variables, Zeros(ParticleDA.get_state_dimension(model)), @@ -336,7 +336,7 @@ function update_state_vector_from_prognostic_variables!( ) end -function update_clock_from_time_index!(clock::SpeedyWeather.Clock,time_index::Int) +function update_clock_from_time_index!(clock::SpeedyWeather.Clock, time_index::Int) clock.time = clock.start + clock.n_timesteps * clock.Δt * (time_index - 1) clock.timestep_counter = clock.n_timesteps * (time_index - 1) end From 40e2a7592ac278d4df0666fe490be0783b0698a6 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 6 Feb 2025 14:17:25 +0000 Subject: [PATCH 31/33] Avoid repetition in vector updates --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 3b43235..6af394e 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -274,16 +274,12 @@ function update_vector_from_spectral_coefficients!( increment::Bool = false ) where {T <: AbstractFloat} n_row, n_col = spectral_truncation + 2, spectral_truncation + 1 + update! = increment ? (lhs, rhs) -> (lhs .+= rhs) : (lhs, rhs) -> (lhs .= rhs) # First column of spectral_coefficients (order = m = 0) are real-valued and we skip # last row (degree = l = n_row - 1) as used only for computing meridional derivative # for vector valued fields. LowerTriangularMatrix allows vector (flat) indexing # skipping zero upper-triangular entries - # TODO: Figure out how to do this without repetition here and below - if increment - vector[1:n_row - 1] .+= real(spectral_coefficients[1:n_row - 1]) - else - vector[1:n_row - 1] .= real(spectral_coefficients[1:n_row - 1]) - end + @views update!(vector[1:n_row - 1], real(spectral_coefficients[1:n_row - 1])) # vector index is i, spectral coefficient (flat) index is j i = n_row - 1 j = n_row @@ -294,15 +290,10 @@ function update_vector_from_spectral_coefficients!( slice_size = n_row - col_index # Reinterpret complex valued spectral coefficients to extract both real and # imaginary components - if increment - vector[i + 1:i + 2 * slice_size] .+= reinterpret( - T, spectral_coefficients[j + 1:j + slice_size] - ) - else - vector[i + 1:i + 2 * slice_size] .= reinterpret( - T, spectral_coefficients[j + 1:j + slice_size] - ) - end + @views update!( + vector[i + 1:i + 2 * slice_size], + reinterpret(T, spectral_coefficients[j + 1:j + slice_size]) + ) # Update vector and spectral coefficient indices, adding 1 offset # to latter to skip entries corresponding to last row i = i + 2 * slice_size From 436818686a7c4f4fe0a85d8c40044cf37d4423be Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 6 Feb 2025 17:54:58 +0000 Subject: [PATCH 32/33] Correct GRF parameterization --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 34 +++++++++++------------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 6af394e..50b726b 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -21,7 +21,7 @@ end Base.@kwdef struct GaussianRandomFieldParameters{T<:AbstractFloat} output_scale::T = 1. - length_scale::T = 1. + length_scale::T = 0.1 end Base.@kwdef struct SpeedyParameters{T<:AbstractFloat, M<:SpeedyWeather.AbstractModel} @@ -393,19 +393,15 @@ function get_grf_coefficient_scale_factors( spectral_truncation::Int, norm_sphere::T ) where {T <: AbstractFloat} - ell_max = spectral_truncation + 2 - scale_factors = zeros(T, ell_max) - denominator = 2 * sum( - [ - (2 * ell + 1) * exp(-ell * (ell + 1) / parameters.length_scale^2) - for ell in 1:spectral_truncation - ] - ) - common_scale = norm_sphere * sqrt(2 * parameters.output_scale^2 / denominator) - for ell in 2:ell_max - scale_factors[ell] = common_scale * exp( - -ell * (ell - 1) / (2 * parameters.length_scale^2) - ) + el_max = spectral_truncation + 1 + (; output_scale, length_scale) = parameters + norm_factor = norm_sphere * output_scale / sqrt( + sum((2 * el + 1) * exp(-2 * length_scale^2 * el * (el + 1)) for el in 1:el_max) + ) + scale_factors = zeros(T, el_max) + # el = 1 case corresponds to mean - assume zero-mean + for el in 2:el_max + scale_factors[el] = norm_factor * exp(-length_scale^2 * el * (el + 1)) end return scale_factors end @@ -416,11 +412,13 @@ function generate_random_spectral_coefficients!( scale_factors::AbstractVector{T}, rng::AbstractRNG ) where {T <: AbstractFloat} - ell_m = 0 + el_m = 0 @inbounds for m in 1:spectral_truncation + 1 - for ell in m:spectral_truncation + 2 - ell_m += 1 - spectral_coefficients[ell_m] = scale_factors[ell] * randn(rng, Complex{T}) + for el in m:spectral_truncation + 2 + el_m += 1 + # Don't generate coefficients in last row (used only for meridional deriv.) + (el == spectral_truncation + 2) && continue + spectral_coefficients[el_m] = scale_factors[el] * randn(rng, Complex{T}) end end end From 2cc1d65362011422c88ae64e9b5f93f23565699d Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Mon, 10 Feb 2025 19:33:03 +0000 Subject: [PATCH 33/33] Allowing use of multiple tasks --- SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl | 82 +++++++++++++++--------- 1 file changed, 53 insertions(+), 29 deletions(-) diff --git a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl index 50b726b..c0cc1ae 100644 --- a/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl +++ b/SpeedyWeatherSSM/src/SpeedyWeatherSSM.jl @@ -55,12 +55,14 @@ struct SpeedyModel{ G<:SpeedyWeather.AbstractSpectralGrid, M<:SpeedyWeather.AbstractModel, I<:SpeedyWeather.RingGrids.AbstractInterpolator, + P<:SpeedyWeather.AbstractPrognosticVariables, + D<:SpeedyWeather.AbstractDiagnosticVariables } parameters::SpeedyParameters{T, M} spectral_grid::G model::M - prognostic_variables::PrognosticVariables{T} - diagnostic_variables::DiagnosticVariables{T} + prognostic_variables::Vector{P} + diagnostic_variables::Vector{D} variable_names::Tuple n_layered_variables::Int n_surface_variables::Int @@ -70,7 +72,7 @@ struct SpeedyModel{ state_noise_grf_scale_factors::Dict{Symbol, Vector{T}} end -function init(parameters::SpeedyParameters{T, M}) where { +function init(parameters::SpeedyParameters{T, M}, n_tasks::Int=1) where { T<:AbstractFloat, M<:SpeedyWeather.AbstractModel } spectral_grid = SpectralGrid(; @@ -81,6 +83,21 @@ function init(parameters::SpeedyParameters{T, M}) where { model.output.active = false simulation = initialize!(model; time=parameters.start_date) (; prognostic_variables, diagnostic_variables) = simulation + SpeedyWeather.set_period!( + prognostic_variables.clock, SpeedyWeather.Day(parameters.n_days) + ) + SpeedyWeather.initialize!(prognostic_variables.clock, model.time_stepping) + # We need separate copies of prognostic and diagnostic variables for each task to + # allow independent parallel read-write access + per_task_prognostic_variables = Vector{typeof(prognostic_variables)}(undef, n_tasks) + per_task_diagnostic_variables = Vector{typeof(diagnostic_variables)}(undef, n_tasks) + per_task_prognostic_variables[1] = prognostic_variables + per_task_diagnostic_variables[1] = diagnostic_variables + for t in 2:n_tasks + per_task_prognostic_variables[t] = PrognosticVariables(spectral_grid, model) + copy!(per_task_prognostic_variables[t], prognostic_variables) + per_task_diagnostic_variables[t] = DiagnosticVariables(spectral_grid) + end variable_names = SpeedyWeather.prognostic_variables(model) n_layered_variables = count( SpeedyWeather.has(model, var) for var in LAYERED_VARIABLES @@ -97,10 +114,6 @@ function init(parameters::SpeedyParameters{T, M}) where { parameters.observed_coordinates[1, :], parameters.observed_coordinates[2, :] ) - SpeedyWeather.set_period!( - prognostic_variables.clock, SpeedyWeather.Day(parameters.n_days) - ) - SpeedyWeather.initialize!(prognostic_variables.clock, model.time_stepping) initial_state_grf_scale_factors = Dict( name => get_grf_coefficient_scale_factors( parameters.initial_state_grf_parameters[name], @@ -121,8 +134,8 @@ function init(parameters::SpeedyParameters{T, M}) where { parameters, spectral_grid, model, - prognostic_variables, - diagnostic_variables, + per_task_prognostic_variables, + per_task_diagnostic_variables, variable_names, n_layered_variables, n_surface_variables, @@ -250,17 +263,17 @@ function update_prognostic_variables_from_state_vector!( end function update_prognostic_variables_from_state_vector!( - model::SpeedyModel{T}, state::AbstractVector{T} + model::SpeedyModel{T}, state::AbstractVector{T}, task_index::Int ) where {T <: AbstractFloat} update_prognostic_variables_from_state_vector!( - model.prognostic_variables, + model.prognostic_variables[task_index], state, model.variable_names; leapfrog_step=1 ) # Zero coefficients for second leapfrog step (corresponding to initial state) update_prognostic_variables_from_state_vector!( - model.prognostic_variables, + model.prognostic_variables[task_index], Zeros(ParticleDA.get_state_dimension(model)), model.variable_names; leapfrog_step=2 @@ -318,11 +331,11 @@ function update_state_vector_from_prognostic_variables!( end function update_state_vector_from_prognostic_variables!( - state::AbstractVector{T}, model::SpeedyModel{T} + state::AbstractVector{T}, model::SpeedyModel{T}, task_index::Int ) where {T <: AbstractFloat} update_state_vector_from_prognostic_variables!( state, - model.prognostic_variables, + model.prognostic_variables[task_index], model.variable_names ) end @@ -332,8 +345,12 @@ function update_clock_from_time_index!(clock::SpeedyWeather.Clock, time_index::I clock.timestep_counter = clock.n_timesteps * (time_index - 1) end -function update_clock_from_time_index!(model::SpeedyModel, time_index::Int) - update_clock_from_time_index!(model.prognostic_variables.clock, time_index) +function update_clock_from_time_index!( + model::SpeedyModel, time_index::Int, task_index::Int +) + update_clock_from_time_index!( + model.prognostic_variables[task_index].clock, time_index + ) end function get_observed_variable_field( @@ -344,12 +361,12 @@ function get_observed_variable_field( end function update_prognostic_and_diagnostic_variables_from_state_vector!( - model::SpeedyModel{T}, state::AbstractVector{T} + model::SpeedyModel{T}, state::AbstractVector{T}, task_index::Int ) where {T <: AbstractFloat} - update_prognostic_variables_from_state_vector!(model, state) + update_prognostic_variables_from_state_vector!(model, state, task_index) SpeedyWeather.transform!( - model.diagnostic_variables, - model.prognostic_variables, + model.diagnostic_variables[task_index], + model.prognostic_variables[task_index], 1, model.model, initialize=true @@ -430,11 +447,12 @@ ParticleDA.get_observation_eltype(model::SpeedyModel{T}) where {T<:AbstractFloat function ParticleDA.sample_initial_state!( state::AbstractVector{T}, model::SpeedyModel{T}, rng::R, task_index::Int=1 ) where {T<:AbstractFloat, R<:AbstractRNG} - initial_conditions = model.model.initial_conditions SpeedyWeather.initialize!( - model.prognostic_variables, initial_conditions, model.model + model.prognostic_variables[task_index], + model.model.initial_conditions, + model.model ) - update_state_vector_from_prognostic_variables!(state, model) + update_state_vector_from_prognostic_variables!(state, model, task_index) add_noise_to_state_vector!( state, model.parameters.spectral_truncation, @@ -448,12 +466,14 @@ end function ParticleDA.update_state_deterministic!( state::AbstractVector{T}, model::SpeedyModel{T}, time_index::Int, task_index::Int=1 ) where {T<:AbstractFloat} - update_prognostic_variables_from_state_vector!(model, state) - update_clock_from_time_index!(model, time_index) + update_prognostic_variables_from_state_vector!(model, state, task_index) + update_clock_from_time_index!(model, time_index, task_index) SpeedyWeather.time_stepping!( - model.prognostic_variables, model.diagnostic_variables, model.model + model.prognostic_variables[task_index], + model.diagnostic_variables[task_index], + model.model ) - update_state_vector_from_prognostic_variables!(state, model) + update_state_vector_from_prognostic_variables!(state, model, task_index) end function ParticleDA.update_state_stochastic!( @@ -475,8 +495,12 @@ function ParticleDA.get_observation_mean_given_state!( model::SpeedyModel{T}, task_index::Int=1 ) where {T<:AbstractFloat} - update_prognostic_and_diagnostic_variables_from_state_vector!(model, state) - observed_field_grid = get_observed_variable_field(model.diagnostic_variables, model) + update_prognostic_and_diagnostic_variables_from_state_vector!( + model, state, task_index + ) + observed_field_grid = get_observed_variable_field( + model.diagnostic_variables[task_index], model + ) SpeedyWeather.interpolate!( observation_mean, observed_field_grid, model.observation_interpolator )