diff --git a/.gitignore b/.gitignore index f1cbd1f..9a70561 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,4 @@ docs/src/tutorials/intermediate/clustered_sites_probo_allsites.jld2 projects/clustered_sites/intermediate/* build/* +intermediate/* diff --git a/Project.toml b/Project.toml index e738ec1..284c5cc 100644 --- a/Project.toml +++ b/Project.toml @@ -16,6 +16,7 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" DistributionFits = "45214091-1ed4-4409-9bcf-fdb48a05e921" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e" @@ -29,7 +30,9 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" +PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" @@ -40,11 +43,10 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5" -DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" - [extensions] HybridVariationalInferenceCUDAExt = "CUDA" @@ -69,6 +71,7 @@ DistributionFits = "0.3.9" Distributions = "0.25.117" FillArrays = "1.13.0" Flux = "0.14, 0.15, 0.16" +ForwardDiff = "1" Functors = "0.4, 0.5" GPUArraysCore = "0.1, 0.2" IterTools = "1.10.0" @@ -83,7 +86,9 @@ NaNMath = "1.1.3" Optimisers = "0.4.6" Optimization = "5.3" PDMats = "0.11.37" +PreallocationTools = "1.2.1" Random = "1.10.0" +RecursiveArrayTools = "3.54.0" SimpleChains = "0.4.7" StableRNGs = "1.0.2" StaticArrays = "1.9.13" diff --git a/docs/src/tutorials/Manifest-v1.11.toml b/docs/src/tutorials/Manifest-v1.11.toml index 403e39b..19b8d34 100644 --- a/docs/src/tutorials/Manifest-v1.11.toml +++ b/docs/src/tutorials/Manifest-v1.11.toml @@ -5,9 +5,9 @@ manifest_format = "2.0" project_hash = "9c33017a188b1974d0c0c424bd1863131c3a2b25" [[deps.ADTypes]] -git-tree-sha1 = "f7304359109c768cf32dc5fa2d371565bb63b68a" +git-tree-sha1 = "5970c86505ae9c07bf5bc521ef2bbbb3849e8b7b" uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -version = "1.21.0" +version = "1.23.0" weakdeps = ["ChainRulesCore", "ConstructionBase", "EnzymeCore"] [deps.ADTypes.extensions] @@ -28,9 +28,9 @@ weakdeps = ["ChainRulesCore", "Test"] [[deps.AbstractMCMC]] deps = ["BangBang", "ConsoleProgressMonitor", "Dates", "Distributed", "LogDensityProblems", "Logging", "LoggingExtras", "ProgressLogging", "Random", "StatsBase", "TerminalLoggers", "UUIDs"] -git-tree-sha1 = "8ac6182431567907e0d5170bcac6dd48fa541f78" +git-tree-sha1 = "328c7d50f307c66308a915abb20d9889e5aab48b" uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001" -version = "5.15.1" +version = "5.16.0" [deps.AbstractMCMC.extensions] AbstractMCMCOnlineStatsExt = "OnlineStats" @@ -57,9 +57,9 @@ version = "0.4.5" [[deps.Accessors]] deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "MacroTools"] -git-tree-sha1 = "2eeb2c9bef11013efc6f8f97f32ee59b146b09fb" +git-tree-sha1 = "7063ad1083578215c7c4bf410368150abe8d5524" uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" -version = "0.1.44" +version = "0.1.45" [deps.Accessors.extensions] AxisKeysExt = "AxisKeys" @@ -80,10 +80,10 @@ version = "0.1.44" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [[deps.Adapt]] -deps = ["LinearAlgebra", "Requires"] -git-tree-sha1 = "35ea197a51ce46fcd01c4a44befce0578a1aaeca" +deps = ["LinearAlgebra"] +git-tree-sha1 = "daa72978cd7a624246e894a4f4f067706d4e17e2" uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -version = "4.5.0" +version = "4.7.0" weakdeps = ["SparseArrays", "StaticArrays"] [deps.Adapt.extensions] @@ -118,9 +118,9 @@ version = "1.1.2" [[deps.ArrayInterface]] deps = ["Adapt", "LinearAlgebra"] -git-tree-sha1 = "78b3a7a536b4b0a747a0f296ea77091ca0a9f9a3" +git-tree-sha1 = "13f3b228c230ef0b4ecafd73c8ca9e99987ca692" uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -version = "7.23.0" +version = "7.30.0" [deps.ArrayInterface.extensions] ArrayInterfaceAMDGPUExt = "AMDGPU" @@ -130,6 +130,7 @@ version = "7.23.0" ArrayInterfaceCUDSSExt = ["CUDSS", "CUDA"] ArrayInterfaceChainRulesCoreExt = "ChainRulesCore" ArrayInterfaceChainRulesExt = "ChainRules" + ArrayInterfaceFillArraysExt = "FillArrays" ArrayInterfaceGPUArraysCoreExt = "GPUArraysCore" ArrayInterfaceMetalExt = "Metal" ArrayInterfaceReverseDiffExt = "ReverseDiff" @@ -145,6 +146,7 @@ version = "7.23.0" CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" @@ -175,10 +177,10 @@ version = "1.1.3" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [[deps.Automa]] -deps = ["PrecompileTools", "SIMD", "TranscodingStreams"] -git-tree-sha1 = "a8f503e8e1a5f583fbef15a8440c8c7e32185df2" +deps = ["PrecompileTools", "TranscodingStreams"] +git-tree-sha1 = "94eab0b3ccdcac361188cc661daf69d4433c1818" uuid = "67c07d97-cdcb-5c2c-af73-a7f9c32a568b" -version = "1.1.0" +version = "1.2.0" [[deps.AxisAlgorithms]] deps = ["LinearAlgebra", "Random", "SparseArrays", "WoodburyMatrices"] @@ -225,20 +227,15 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" version = "1.11.0" [[deps.BaseDirs]] -git-tree-sha1 = "bca794632b8a9bbe159d56bf9e31c422671b35e0" +git-tree-sha1 = "8c290a1b223deaeea9aea44b235d24546da8eb98" uuid = "18cc8868-cbac-4acf-b575-c8ff214dc66f" -version = "1.3.2" - -[[deps.Baselet]] -git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e" -uuid = "9718e550-a3fa-408a-8086-8db961cd8217" -version = "0.1.1" +version = "1.4.0" [[deps.Bijectors]] deps = ["AbstractPPL", "ArgCheck", "ChainRulesCore", "ChangesOfVariables", "DifferentiationInterface", "Distributions", "DocStringExtensions", "EnzymeCore", "FillArrays", "Functors", "InverseFunctions", "IrrationalConstants", "LinearAlgebra", "LogExpFunctions", "MappedArrays", "Random", "Reexport", "Roots", "SparseArrays", "Statistics", "Test"] -git-tree-sha1 = "d6ee8f89dd20f933fbdad578a798e32babb617ee" +git-tree-sha1 = "b425418b2644f826823e8ebd4b68a076e8d8d2ec" uuid = "76274a88-744f-5084-9051-94815aaf08c4" -version = "0.15.20" +version = "0.15.24" [deps.Bijectors.extensions] BijectorsDistributionsADExt = "DistributionsAD" @@ -257,9 +254,9 @@ version = "0.15.20" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [[deps.BitFlags]] -git-tree-sha1 = "0691e34b3bb8be9307330f88d1a3c3f25466c24d" +git-tree-sha1 = "bbe1079eecf9c9fbb52765193ad2bae27ae09bc8" uuid = "d1d4a3ce-64b1-5f1a-9ba4-7e7e69966f35" -version = "0.1.9" +version = "0.1.10" [[deps.BitTwiddlingConvenienceFunctions]] deps = ["Static"] @@ -312,9 +309,9 @@ version = "1.0.1+0" [[deps.CUDA]] deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Compiler_jll", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "ExprTools", "GPUArrays", "GPUCompiler", "GPUToolbox", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "demumble_jll"] -git-tree-sha1 = "ea6a2ab8307059b6c9ea186ff7dfcd032a13b731" +git-tree-sha1 = "b267c611dcbbcb70d42e398192ee0af160358075" uuid = "052768ef-5323-5732-b1bb-66c8b64840ba" -version = "5.11.0" +version = "5.11.1" [deps.CUDA.extensions] ChainRulesCoreExt = "ChainRulesCore" @@ -330,15 +327,15 @@ version = "5.11.0" [[deps.CUDA_Compiler_jll]] deps = ["Artifacts", "CUDA_Driver_jll", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "8c19e97de5b7574672e4a7a3abd55714ad66d59a" +git-tree-sha1 = "c32d22f2f563ce192c88a44b09c2b569f1e7a980" uuid = "d1e2174e-dfdc-576e-b43e-73b79eb1aca8" -version = "0.4.2+0" +version = "0.4.4+1" [[deps.CUDA_Driver_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "TOML"] -git-tree-sha1 = "061f39cc84e99928830aa1005d79f7e99097ba28" +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "901cc2d895842b2ddef81715cb70fffc3b5028ae" uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc" -version = "13.2.0+0" +version = "13.3.1+0" [[deps.CUDA_Runtime_Discovery]] deps = ["Libdl"] @@ -348,15 +345,15 @@ version = "1.0.0" [[deps.CUDA_Runtime_jll]] deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "af17d37b5b8b4d7525f8902eba1ef6141a9a7d3b" +git-tree-sha1 = "c0314d9fb0ebd00e404feba4c3fbc04c9975abc1" uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" -version = "0.21.0+0" +version = "0.21.0+1" [[deps.CUDNN_jll]] deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "70dea6a7133d2100a143b515a00d6d887e208500" +git-tree-sha1 = "77b169898d4cdc234b1cd9afc1d0a8cac1017a24" uuid = "62b44479-cb7b-5706-934f-f13b2eb2e645" -version = "9.20.0+0" +version = "9.24.0+0" [[deps.Cairo]] deps = ["Cairo_jll", "Colors", "Glib_jll", "Graphics", "Libdl", "Pango_jll"] @@ -366,15 +363,15 @@ version = "1.1.1" [[deps.CairoMakie]] deps = ["CRC32c", "Cairo", "Cairo_jll", "Colors", "FileIO", "FreeType", "GeometryBasics", "LinearAlgebra", "Makie", "PrecompileTools"] -git-tree-sha1 = "fa072933899aae6dc61dde934febed8254e66c6a" +git-tree-sha1 = "47142129b1777e21da58cff265050b10d8560588" uuid = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" -version = "0.15.9" +version = "0.15.13" [[deps.Cairo_jll]] deps = ["Artifacts", "Bzip2_jll", "CompilerSupportLibraries_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "JLLWrappers", "Libdl", "Pixman_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Zlib_jll", "libpng_jll"] -git-tree-sha1 = "d0efe2c6fdcdaa1c161d206aa8b933788397ec71" +git-tree-sha1 = "1fa950ebc3e37eccd51c6a8fe1f92f7d86263522" uuid = "83423d85-b0ee-5818-9007-b63ccbeb887a" -version = "1.18.6+0" +version = "1.18.7+0" [[deps.ChainRules]] deps = ["Adapt", "ChainRulesCore", "Compat", "Distributed", "GPUArraysCore", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "SparseInverseSubset", "Statistics", "StructArrays", "SuiteSparse"] @@ -394,9 +391,9 @@ weakdeps = ["SparseArrays"] [[deps.ChangesOfVariables]] deps = ["LinearAlgebra"] -git-tree-sha1 = "3aa4bf1532aa2e14e0374c4fd72bed9a9d0d0f6c" +git-tree-sha1 = "83ee8183bd8c4a390ae178385e6c7b3aa4e468b2" uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" -version = "0.1.10" +version = "0.1.11" weakdeps = ["InverseFunctions", "Test"] [deps.ChangesOfVariables.extensions] @@ -410,9 +407,9 @@ version = "1.0.1" [[deps.ChunkCodecLibZlib]] deps = ["ChunkCodecCore", "Zlib_jll"] -git-tree-sha1 = "cee8104904c53d39eb94fd06cbe60cb5acde7177" +git-tree-sha1 = "d4101e848e8d3f585d61d244c2fe0c80a70e6b3b" uuid = "4c0bbee4-addc-4d73-81a0-b6caacae83c8" -version = "1.0.0" +version = "1.1.0" [[deps.ChunkCodecLibZstd]] deps = ["ChunkCodecCore", "Zstd_jll"] @@ -432,11 +429,23 @@ git-tree-sha1 = "3e22db924e2945282e70c33b75d4dde8bfa44c94" uuid = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5" version = "0.15.8" +[[deps.CodeTracking]] +deps = ["InteractiveUtils", "REPL", "UUIDs"] +git-tree-sha1 = "cfb7a2e89e245a9d5016b70323db412b3a7438d5" +uuid = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2" +version = "3.0.2" + [[deps.CodecZlib]] deps = ["TranscodingStreams", "Zlib_jll"] -git-tree-sha1 = "962834c22b66e32aa10f7611c08c8ca4e20749a9" +git-tree-sha1 = "970758a3d591a2a5c2a907c53f2e2f8c1b1d3537" uuid = "944b1d66-785c-5afd-91f1-9de20f533193" -version = "0.7.8" +version = "0.7.9" + +[[deps.CodecZstd]] +deps = ["TranscodingStreams", "Zstd_jll"] +git-tree-sha1 = "da54a6cd93c54950c15adf1d336cfd7d71f51a56" +uuid = "6b39b394-51ab-5f42-8807-6242bab2b4c2" +version = "0.8.7" [[deps.ColorBrewer]] deps = ["Colors", "JSON"] @@ -482,9 +491,9 @@ uuid = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" version = "1.1.0" [[deps.CommonSolve]] -git-tree-sha1 = "78ea4ddbcf9c241827e7035c3a03e2e456711470" +git-tree-sha1 = "cf963add2340ad9960e5eb22844e61ad8f931fe1" uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" -version = "0.2.6" +version = "0.2.13" [[deps.CommonSubexpressions]] deps = ["MacroTools"] @@ -493,9 +502,9 @@ uuid = "bbf7d656-a473-5ed7-a52c-81e309532950" version = "0.3.1" [[deps.CommonWorldInvalidations]] -git-tree-sha1 = "ae52d1c52048455e85a387fbee9be553ec2b68d0" +git-tree-sha1 = "ef2022bff55342a8c9846cdf218f62e475f0444d" uuid = "f70d9fcc-98c5-4d4a-abd7-e4cdeebd8ca8" -version = "1.0.0" +version = "1.1.2" [[deps.Compat]] deps = ["TOML", "UUIDs"] @@ -514,9 +523,9 @@ version = "1.1.1+0" [[deps.ComponentArrays]] deps = ["Adapt", "ArrayInterface", "ChainRulesCore", "ConstructionBase", "Functors", "LinearAlgebra", "StaticArrayInterface", "StaticArraysCore"] -git-tree-sha1 = "f9873e13c3f89808e518118951eb6b244f7af735" +git-tree-sha1 = "122b5c26469f63d50d1e76f6c63428aa7d2b244b" uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" -version = "0.15.33" +version = "0.15.44" [deps.ComponentArrays.extensions] ComponentArraysGPUArraysExt = "GPUArrays" @@ -526,7 +535,7 @@ version = "0.15.33" ComponentArraysReactantExt = "Reactant" ComponentArraysRecursiveArrayToolsExt = "RecursiveArrayTools" ComponentArraysReverseDiffExt = "ReverseDiff" - ComponentArraysSciMLBaseExt = "SciMLBase" + ComponentArraysSciMLBaseExt = ["SciMLBase", "SymbolicIndexingInterface"] ComponentArraysTrackerExt = "Tracker" ComponentArraysZygoteExt = "Zygote" @@ -539,6 +548,7 @@ version = "0.15.33" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" + SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" @@ -553,20 +563,20 @@ weakdeps = ["InverseFunctions"] [[deps.ComputePipeline]] deps = ["Observables", "Preferences"] -git-tree-sha1 = "3b4be73db165146d8a88e47924f464e55ab053cd" +git-tree-sha1 = "7bc84b769c1d384315e7b5c4ac03a6c303e6cf35" uuid = "95dc2771-c249-4cd0-9c9f-1f3b4330693c" -version = "0.1.7" +version = "0.1.8" [[deps.ConcreteStructs]] -git-tree-sha1 = "f749037478283d372048690eb3b5f92a79432b34" +git-tree-sha1 = "804fc3ca1cbfdd5aa52ae3149c0cb0555f875eec" uuid = "2569d6c7-a4a2-43d3-a901-331e8e4be471" -version = "0.2.3" +version = "0.2.7" [[deps.ConcurrentUtilities]] deps = ["Serialization", "Sockets"] -git-tree-sha1 = "21d088c496ea22914fe80906eb5bce65755e5ec8" +git-tree-sha1 = "3c9be947934c38475bafe822c6d61aaed17f0738" uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb" -version = "2.5.1" +version = "2.6.0" [[deps.ConsoleProgressMonitor]] deps = ["Logging", "ProgressMeter"] @@ -585,12 +595,6 @@ weakdeps = ["IntervalSets", "LinearAlgebra", "StaticArrays"] ConstructionBaseLinearAlgebraExt = "LinearAlgebra" ConstructionBaseStaticArraysExt = "StaticArrays" -[[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" @@ -615,9 +619,9 @@ uuid = "adafc99b-e345-5852-983c-f28acb93d879" version = "0.3.1" [[deps.Crayons]] -git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15" +git-tree-sha1 = "54b76cbb40d9a0f5368c880725b2f141da77c94f" uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f" -version = "4.1.1" +version = "4.2.0" [[deps.DataAPI]] git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe" @@ -626,9 +630,9 @@ version = "1.16.0" [[deps.DataStructures]] deps = ["OrderedCollections"] -git-tree-sha1 = "e86f4a2805f7f19bec5129bc9150c38208e5dc23" +git-tree-sha1 = "b0bc6d2cad1fed8b7fd59a1551a991cb3d2809e6" uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" -version = "0.19.4" +version = "0.19.6" [[deps.DataValueInterfaces]] git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6" @@ -646,11 +650,6 @@ git-tree-sha1 = "473e9afc9cf30814eb67ffa5f2db7df82c3ad9fd" uuid = "ee1fde0b-3d02-5ea6-8484-8dfef6360eab" version = "1.16.2+0" -[[deps.DefineSingletons]] -git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c" -uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52" -version = "0.1.2" - [[deps.DelaunayTriangulation]] deps = ["AdaptivePredicates", "EnumX", "ExactPredicates", "Random"] git-tree-sha1 = "c55f5a9fd67bdbc8e089b5a3111fe4292986a8e8" @@ -677,15 +676,15 @@ version = "1.1.0" [[deps.DiffRules]] deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"] -git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272" +git-tree-sha1 = "79a2aca180a85c690c58a020d47b426954b590f8" uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" -version = "1.15.1" +version = "1.16.0" [[deps.DifferentiationInterface]] deps = ["ADTypes", "LinearAlgebra"] -git-tree-sha1 = "7ae99144ea44715402c6c882bfef2adbeadbc4ce" +git-tree-sha1 = "0693d8b0a4608ff289d228ab4c598df5894845cd" uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" -version = "0.7.16" +version = "0.7.21" [deps.DifferentiationInterface.extensions] DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore" @@ -695,8 +694,9 @@ version = "0.7.16" DifferentiationInterfaceFiniteDiffExt = "FiniteDiff" DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences" DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"] - DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore" + DifferentiationInterfaceGPUArraysCoreExt = ["GPUArraysCore", "Adapt"] DifferentiationInterfaceGTPSAExt = "GTPSA" + DifferentiationInterfaceHyperHessiansExt = "HyperHessians" DifferentiationInterfaceMooncakeExt = "Mooncake" DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"] DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"] @@ -709,6 +709,7 @@ version = "0.7.16" DifferentiationInterfaceZygoteExt = ["Zygote", "ForwardDiff"] [deps.DifferentiationInterface.weakdeps] + Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c" @@ -720,6 +721,7 @@ version = "0.7.16" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" + HyperHessians = "06b494a0-c8e0-40cc-ad32-d99506a00a6c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" @@ -759,10 +761,10 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" version = "1.11.0" [[deps.DistributionFits]] -deps = ["Distributions", "FillArrays", "LinearAlgebra", "Reexport", "Requires", "StaticArrays", "Statistics", "StatsAPI", "StatsFuns"] -git-tree-sha1 = "a50a9fbdcccc1cb08ff2e1f40c37a268a74b8879" +deps = ["Distributions", "FillArrays", "LinearAlgebra", "Reexport", "StaticArrays", "Statistics", "StatsAPI", "StatsFuns"] +git-tree-sha1 = "1ebcc04db80b949d464046c21496c1738ed5aaa9" uuid = "45214091-1ed4-4409-9bcf-fdb48a05e921" -version = "0.3.9" +version = "0.3.10" [deps.DistributionFits.extensions] DistributionFitsOptimExt = "Optim" @@ -771,15 +773,16 @@ version = "0.3.9" Optim = "429524aa-4258-5aef-a3af-852621145aeb" [[deps.Distributions]] -deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"] -git-tree-sha1 = "fbcc7610f6d8348428f722ecbe0e6cfe22e672c6" +deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "Roots", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"] +git-tree-sha1 = "d2facc77c08c1c2bfb1a77c148edd05b3db5410b" uuid = "31c24e10-a181-5473-b8eb-7969acd0382f" -version = "0.25.123" -weakdeps = ["ChainRulesCore", "DensityInterface", "Test"] +version = "0.25.130" +weakdeps = ["ChainRulesCore", "DensityInterface", "SparseConnectivityTracer", "Test"] [deps.Distributions.extensions] DistributionsChainRulesCoreExt = "ChainRulesCore" DistributionsDensityInterfaceExt = "DensityInterface" + DistributionsSparseConnectivityTracerExt = "SparseConnectivityTracer" DistributionsTestExt = "Test" [[deps.DocStringExtensions]] @@ -804,9 +807,9 @@ uuid = "4e289a0a-7415-4d19-859d-a7e5c4648b56" version = "1.0.7" [[deps.EnzymeCore]] -git-tree-sha1 = "24bbb6fc8fb87eb71c1f8d00184a60fc22c63903" +git-tree-sha1 = "971d7831cc85f43bc9f51d615a3f7f21270c2f1d" uuid = "f151be2c-9106-41f4-ab19-57ee4f262869" -version = "0.8.19" +version = "0.8.21" weakdeps = ["Adapt", "ChainRulesCore"] [deps.EnzymeCore.extensions] @@ -833,30 +836,25 @@ version = "0.1.11" [[deps.Expat_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "27af30de8b5445644e8ffe3bcb0d72049c089cf1" +git-tree-sha1 = "f4d39eee89f1e58c26bf447f1d4156c0125d6838" uuid = "2e619515-83b5-522b-bb60-26c02a35a201" -version = "2.7.3+0" +version = "2.8.3+0" [[deps.ExprTools]] -git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" +git-tree-sha1 = "d2e49e7efd29719d6f28b891b0e0e159daa9d2b4" uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" -version = "0.1.10" +version = "0.1.11" [[deps.ExpressionExplorer]] -git-tree-sha1 = "5f1c005ed214356bbe41d442cc1ccd416e510b7e" +git-tree-sha1 = "678f8b7cd246ed441d29ca42a99f5f138be2ecac" uuid = "21656369-7473-754a-2065-74616d696c43" -version = "1.1.4" +version = "1.1.5" [[deps.ExproniconLite]] git-tree-sha1 = "c13f0b150373771b0fdc1713c97860f8df12e6c2" uuid = "55351af7-c7e9-48d6-89ff-24e801d99491" version = "0.10.14" -[[deps.Extents]] -git-tree-sha1 = "b309b36a9e02fe7be71270dd8c0fd873625332b4" -uuid = "411431e0-e8b7-467b-b5e0-f676ba4f2910" -version = "0.1.6" - [[deps.FFMPEG]] deps = ["FFMPEG_jll"] git-tree-sha1 = "95ecf07c2eea562b5adbd0696af6db62c0f52560" @@ -865,9 +863,9 @@ version = "0.4.5" [[deps.FFMPEG_jll]] deps = ["Artifacts", "Bzip2_jll", "FreeType2_jll", "FriBidi_jll", "JLLWrappers", "LAME_jll", "Libdl", "Ogg_jll", "OpenSSL_jll", "Opus_jll", "PCRE2_jll", "Zlib_jll", "libaom_jll", "libass_jll", "libfdk_aac_jll", "libva_jll", "libvorbis_jll", "x264_jll", "x265_jll"] -git-tree-sha1 = "66381d7059b5f3f6162f28831854008040a4e905" +git-tree-sha1 = "7a58e45171b63ed4782f2d36fdee8713a469e6e0" uuid = "b22a6f82-2f65-5046-a5b2-351ab43fb4e5" -version = "8.0.1+1" +version = "8.1.2+0" [[deps.FFTA]] deps = ["AbstractFFTs", "DocStringExtensions", "LinearAlgebra", "MuladdMacro", "Primes", "Random", "Reexport"] @@ -875,18 +873,6 @@ git-tree-sha1 = "65e55303b72f4a567a51b174dd2c47496efeb95a" uuid = "b86e33f2-c0db-4aa1-a6e0-ab43e668529e" version = "0.3.1" -[[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.FastClosures]] git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef" uuid = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" @@ -894,9 +880,9 @@ version = "0.3.2" [[deps.FileIO]] deps = ["Pkg", "Requires", "UUIDs"] -git-tree-sha1 = "6522cfb3b8fe97bec632252263057996cbd3de20" +git-tree-sha1 = "6621fef488e496356c9c9625d0562c12a6070819" uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" -version = "1.18.0" +version = "1.20.0" weakdeps = ["HTTP"] [deps.FileIO.extensions] @@ -935,9 +921,9 @@ version = "1.11.0" [[deps.FillArrays]] deps = ["LinearAlgebra"] -git-tree-sha1 = "2f979084d1e13948a3352cf64a25df6bd3b4dca3" +git-tree-sha1 = "5bad39456d9f0166184fce2248783dd9862645c1" uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" -version = "1.16.0" +version = "1.17.0" weakdeps = ["PDMats", "SparseArrays", "StaticArrays", "Statistics"] [deps.FillArrays.extensions] @@ -947,10 +933,10 @@ weakdeps = ["PDMats", "SparseArrays", "StaticArrays", "Statistics"] FillArraysStatisticsExt = "Statistics" [[deps.FixedPointNumbers]] -deps = ["Statistics"] -git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172" +deps = ["Random", "Statistics"] +git-tree-sha1 = "59af96b98217c6ef4ae0dfe065ac7c20831d1a84" uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" -version = "0.8.5" +version = "0.8.6" [[deps.Fontconfig_jll]] deps = ["Artifacts", "Bzip2_jll", "Expat_jll", "FreeType2_jll", "JLLWrappers", "Libdl", "Libuuid_jll", "Zlib_jll"] @@ -965,9 +951,9 @@ version = "1.3.7" [[deps.ForwardDiff]] deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions"] -git-tree-sha1 = "cddeab6487248a39dae1a960fff0ac17b2a28888" +git-tree-sha1 = "1b86cca764a61dcac4fef4c5e16e378e5ed6953c" uuid = "f6369f11-7733-5829-9624-2563aa707210" -version = "1.3.3" +version = "1.4.5" weakdeps = ["StaticArrays"] [deps.ForwardDiff.extensions] @@ -1003,23 +989,25 @@ uuid = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" version = "1.1.3" [[deps.FunctionWrappersWrappers]] -deps = ["FunctionWrappers", "PrecompileTools", "TruncatedStacktraces"] -git-tree-sha1 = "5201523536a43bf8aef3914b7f60b552b098ef8e" +deps = ["FunctionWrappers", "PrecompileTools", "SciMLPublic"] +git-tree-sha1 = "2bcce3ad6f6977d617928d7707fdc86ac83cce03" uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" -version = "1.1.0" +version = "1.13.0" [deps.FunctionWrappersWrappers.extensions] FunctionWrappersWrappersEnzymeExt = ["Enzyme", "EnzymeCore"] + FunctionWrappersWrappersMooncakeExt = "Mooncake" [deps.FunctionWrappersWrappers.weakdeps] Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" + Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" [[deps.Functors]] deps = ["Compat", "ConstructionBase", "LinearAlgebra", "Random"] -git-tree-sha1 = "60a0339f28a233601cb74468032b5c302d5067de" +git-tree-sha1 = "1ac2813982db52b974c9343124ca61adbf297316" uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196" -version = "0.5.2" +version = "0.5.3" [[deps.Future]] deps = ["Random"] @@ -1028,15 +1016,15 @@ version = "1.11.0" [[deps.GLFW_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libglvnd_jll", "Xorg_libXcursor_jll", "Xorg_libXi_jll", "Xorg_libXinerama_jll", "Xorg_libXrandr_jll", "libdecor_jll", "xkbcommon_jll"] -git-tree-sha1 = "b7bfd56fa66616138dfe5237da4dc13bbd83c67f" +git-tree-sha1 = "9e0fb9e54594c47f278d75063980e43066e26e20" uuid = "0656b61e-2033-5cc2-a64a-77c0f6c09b89" -version = "3.4.1+0" +version = "3.4.1+1" [[deps.GPUArrays]] deps = ["Adapt", "GPUArraysCore", "KernelAbstractions", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "ScopedValues", "Serialization", "SparseArrays", "Statistics"] -git-tree-sha1 = "6487601563e4a1d1dab796e88b4548bf5544209e" +git-tree-sha1 = "4939facbe63151c92b622dad9f9bb2b639b5bbde" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "11.4.1" +version = "11.5.12" weakdeps = ["JLD2"] [deps.GPUArrays.extensions] @@ -1050,9 +1038,12 @@ version = "0.2.0" [[deps.GPUCompiler]] deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "PrecompileTools", "Preferences", "Scratch", "Serialization", "TOML", "Tracy", "UUIDs"] -git-tree-sha1 = "fedfe5e7db7035271c3f58359007f971da1dde87" +git-tree-sha1 = "b863a2e71f89328e2af69069490aec2448106e59" uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "1.9.1" +version = "1.17.1" + + [deps.GPUCompiler.weakdeps] + LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444" [[deps.GPUToolbox]] deps = ["LLVM"] @@ -1062,9 +1053,9 @@ version = "1.1.1" [[deps.GR]] deps = ["Artifacts", "Base64", "DelimitedFiles", "Downloads", "GR_jll", "HTTP", "JSON", "Libdl", "LinearAlgebra", "Preferences", "Printf", "Qt6Wayland_jll", "Random", "Serialization", "Sockets", "TOML", "Tar", "Test", "p7zip_jll"] -git-tree-sha1 = "44716a1a667cb867ee0e9ec8edc31c3e4aa5afdc" +git-tree-sha1 = "f954322d5de03ec630d177cda203dcd92b6be399" uuid = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71" -version = "0.73.24" +version = "0.73.26" [deps.GR.extensions] IJuliaExt = "IJulia" @@ -1074,21 +1065,30 @@ version = "0.73.24" [[deps.GR_jll]] deps = ["Artifacts", "Bzip2_jll", "Cairo_jll", "FFMPEG_jll", "Fontconfig_jll", "FreeType2_jll", "GLFW_jll", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "Pixman_jll", "Qt6Base_jll", "Zlib_jll", "libpng_jll"] -git-tree-sha1 = "be8a1b8065959e24fdc1b51402f39f3b6f0f6653" +git-tree-sha1 = "6fada551286ab6ea4ca1628cb2de9f166a2ec966" uuid = "d2c73de3-f751-5644-a686-071e5b155ba9" -version = "0.73.24+0" +version = "0.73.26+0" + +[[deps.Gamma]] +git-tree-sha1 = "86f86b6168a016ed88e4ae4e64577b98c3b59e8e" +uuid = "a0844989-3bd2-4988-8bea-c9407ab0941b" +version = "1.1.0" [[deps.GeometryBasics]] -deps = ["EarCut_jll", "Extents", "IterTools", "LinearAlgebra", "PrecompileTools", "Random", "StaticArrays"] -git-tree-sha1 = "1f5a80f4ed9f5a4aada88fc2db456e637676414b" +deps = ["EarCut_jll", "LinearAlgebra", "PrecompileTools", "Random", "StaticArrays"] +git-tree-sha1 = "364685f5ffde25deb1bbcfd5bb278a5c6b7a9b37" uuid = "5c1252a2-5f33-56bf-86c9-59e7332b4326" -version = "0.5.10" +version = "0.5.11" [deps.GeometryBasics.extensions] + ExtentsExt = "Extents" GeometryBasicsGeoInterfaceExt = "GeoInterface" + IntervalSetsExt = "IntervalSets" [deps.GeometryBasics.weakdeps] + Extents = "411431e0-e8b7-467b-b5e0-f676ba4f2910" GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f" + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" [[deps.GettextRuntime_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Libiconv_jll"] @@ -1110,9 +1110,9 @@ version = "5.2.3+0" [[deps.Glib_jll]] deps = ["Artifacts", "GettextRuntime_jll", "JLLWrappers", "Libdl", "Libffi_jll", "Libiconv_jll", "Libmount_jll", "PCRE2_jll", "Zlib_jll"] -git-tree-sha1 = "24f6def62397474a297bfcec22384101609142ed" +git-tree-sha1 = "090526e65de8f69648ac156daae153de8b56df62" uuid = "7746bdde-850d-59dc-9ae8-88ece973131d" -version = "2.86.3+0" +version = "2.88.3+0" [[deps.Graphics]] deps = ["Colors", "LinearAlgebra", "NaNMath"] @@ -1122,9 +1122,9 @@ version = "1.1.3" [[deps.Graphite2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "8a6dbda1fd736d60cc477d99f2e7a042acfa46e8" +git-tree-sha1 = "69ffb934a5c5b7e086a0b4fee3427db2556fba6e" uuid = "3b182d85-2403-5c21-9c21-1e1f0cc25472" -version = "1.3.15+0" +version = "1.3.16+0" [[deps.GridLayoutBase]] deps = ["GeometryBasics", "InteractiveUtils", "Observables"] @@ -1161,34 +1161,36 @@ uuid = "3e5b6fbb-0976-4d2c-9146-d79de83f2fb0" version = "0.1.18" [[deps.HybridVariationalInference]] -deps = ["Bijectors", "BlockDiagonals", "ChainRulesCore", "Clustering", "Combinatorics", "CommonSolve", "ComponentArrays", "DifferentiationInterface", "Distances", "DistributionFits", "Distributions", "FillArrays", "Functors", "GPUArraysCore", "IterTools", "KernelAbstractions", "LinearAlgebra", "LogExpFunctions", "MLDataDevices", "MLUtils", "Missings", "NaNMath", "Optimisers", "Optimization", "PDMats", "Random", "StableRNGs", "StaticArrays", "StatsBase", "StatsFuns", "Test", "UnPack", "Zygote"] +deps = ["Bijectors", "BlockDiagonals", "ChainRulesCore", "Clustering", "Combinatorics", "CommonSolve", "ComponentArrays", "DifferentiationInterface", "Distances", "DistributionFits", "Distributions", "FillArrays", "ForwardDiff", "Functors", "GPUArraysCore", "IterTools", "KernelAbstractions", "LinearAlgebra", "LogExpFunctions", "MLDataDevices", "MLUtils", "Missings", "NaNMath", "Optimisers", "Optimization", "PDMats", "PreallocationTools", "Random", "RecursiveArrayTools", "StableRNGs", "StaticArrays", "StatsBase", "StatsFuns", "Test", "UnPack", "Zygote"] path = "../../.." uuid = "a108c475-a4e2-4021-9a84-cfa7df242f64" version = "0.2.0" [deps.HybridVariationalInference.extensions] HybridVariationalInferenceCUDAExt = "CUDA" + HybridVariationalInferenceDataFramesExt = "DataFrames" HybridVariationalInferenceFluxExt = "Flux" HybridVariationalInferenceLuxExt = "Lux" HybridVariationalInferenceSimpleChainsExt = "SimpleChains" [deps.HybridVariationalInference.weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5" [[deps.HypergeometricFunctions]] -deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"] -git-tree-sha1 = "68c173f4f449de5b438ee67ed0c9c748dc31a2ec" +deps = ["Gamma", "LinearAlgebra"] +git-tree-sha1 = "31bb6c92405c084617facc1d7ed9eb6c402d061e" uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a" -version = "0.3.28" +version = "0.3.30" [[deps.IRTools]] deps = ["InteractiveUtils", "MacroTools"] -git-tree-sha1 = "57e9ce6cf68d0abf5cb6b3b4abf9bedf05c939c0" +git-tree-sha1 = "88d07a6b68b8fffb13cacd49e23ea73c571859b2" uuid = "7869d1d1-7146-5819-86e3-90919afe41df" -version = "0.4.15" +version = "0.4.20" [[deps.IfElse]] git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1" @@ -1247,9 +1249,9 @@ uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c" version = "0.3.1" [[deps.IntegerMathUtils]] -git-tree-sha1 = "4c1acff2dc6b6967e7e750633c50bc3b8d83e617" +git-tree-sha1 = "c72458f1962faeb003bf23cbdb75164fe6280906" uuid = "18e54dd8-cb9d-406c-a71d-865a43cbb235" -version = "0.1.3" +version = "0.1.4" [[deps.InteractiveUtils]] deps = ["Markdown"] @@ -1258,9 +1260,9 @@ version = "1.11.0" [[deps.Interpolations]] deps = ["Adapt", "AxisAlgorithms", "ChainRulesCore", "LinearAlgebra", "OffsetArrays", "Random", "Ratios", "SharedArrays", "SparseArrays", "StaticArrays", "WoodburyMatrices"] -git-tree-sha1 = "65d505fa4c0d7072990d659ef3fc086eb6da8208" +git-tree-sha1 = "48922d06068130f87e43edef52382e6a94305ae6" uuid = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" -version = "0.16.2" +version = "0.16.3" weakdeps = ["ForwardDiff", "Unitful"] [deps.Interpolations.extensions] @@ -1269,9 +1271,9 @@ weakdeps = ["ForwardDiff", "Unitful"] [[deps.IntervalArithmetic]] deps = ["CRlibm", "CoreMath", "MacroTools", "OpenBLASConsistentFPCSR_jll", "Printf", "Random", "RoundingEmulator"] -git-tree-sha1 = "f1c42fcaca2d8034fe392f3e86c2e0809f75b2a1" +git-tree-sha1 = "ff294afb9a15d31d8d7422da138844641a73135f" uuid = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253" -version = "1.0.6" +version = "1.0.11" [deps.IntervalArithmetic.extensions] IntervalArithmeticArblibExt = "Arblib" @@ -1337,9 +1339,9 @@ version = "1.0.0" [[deps.JLD2]] deps = ["ChunkCodecLibZlib", "ChunkCodecLibZstd", "FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "ScopedValues"] -git-tree-sha1 = "941f87a0ae1b14d1ac2fa57245425b23a9d7a516" +git-tree-sha1 = "9ebadf3f8f0de07031359917549bbdadc23f5dc3" uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819" -version = "0.6.4" +version = "0.6.5" weakdeps = ["UnPack"] [deps.JLD2.extensions] @@ -1353,15 +1355,15 @@ version = "0.1.11" [[deps.JLLWrappers]] deps = ["Artifacts", "Preferences"] -git-tree-sha1 = "0533e564aae234aff59ab625543145446d8b6ec2" +git-tree-sha1 = "7204148362dafe5fe6a273f855b8ccbe4df8173e" uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" -version = "1.7.1" +version = "1.8.0" [[deps.JSON]] deps = ["Dates", "Logging", "Parsers", "PrecompileTools", "StructUtils", "UUIDs", "Unicode"] -git-tree-sha1 = "67c6f1f085cb2671c93fe34244c9cccde30f7a26" +git-tree-sha1 = "c7345ab1a7ca4dc8a02c9f6510da0d9857bbe513" uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" -version = "1.5.0" +version = "1.7.1" [deps.JSON.extensions] JSONArrowExt = ["ArrowTypes"] @@ -1383,9 +1385,9 @@ version = "0.1.6" [[deps.JpegTurbo_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "b6893345fd6658c8e475d40155789f4860ac3b21" +git-tree-sha1 = "037babc10853eeb8e585418922246cb97b8e5b74" uuid = "aacddb02-875f-59d6-b918-886e6ef4fbf8" -version = "3.1.4+0" +version = "3.2.0+1" [[deps.JuliaNVTXCallbacks_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] @@ -1393,17 +1395,11 @@ 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"] -git-tree-sha1 = "f2e76d3ced51a2a9e185abc0b97494c7273f649f" +git-tree-sha1 = "a5b87110fa95d711355af44832497745aa93fb52" uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" -version = "0.9.41" +version = "0.9.42" weakdeps = ["EnzymeCore", "LinearAlgebra", "SparseArrays"] [deps.KernelAbstractions.extensions] @@ -1413,9 +1409,9 @@ weakdeps = ["EnzymeCore", "LinearAlgebra", "SparseArrays"] [[deps.KernelDensity]] deps = ["Distributions", "DocStringExtensions", "FFTA", "Interpolations", "StatsBase"] -git-tree-sha1 = "4260cfc991b8885bf747801fb60dd4503250e478" +git-tree-sha1 = "9eda8292dd3268b3b7ec9df21bbfac24e177ec52" uuid = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b" -version = "0.6.11" +version = "0.6.12" [[deps.LAME_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -1425,15 +1421,15 @@ version = "3.100.3+0" [[deps.LERC_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "aaafe88dccbd957a8d82f7d05be9b69172e0cee3" +git-tree-sha1 = "17b94ecafcfa45e8360a4fc9ca6b583b049e4e37" uuid = "88015f11-f218-50d7-93a8-a6af411a945d" -version = "4.0.1+0" +version = "4.1.0+0" [[deps.LLVM]] -deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Unicode"] -git-tree-sha1 = "69e4739502b7ab5176117e97e1664ed181c35036" +deps = ["CEnum", "LLVMExtra_jll", "Libdl", "PrecompileTools", "Preferences", "Printf", "Unicode"] +git-tree-sha1 = "5f708df9df936ff026cc6232e187ddf93f343430" uuid = "929cbde3-209d-540e-8aea-75f648917ca0" -version = "9.4.6" +version = "9.13.0" weakdeps = ["BFloat16s"] [deps.LLVM.extensions] @@ -1441,9 +1437,9 @@ weakdeps = ["BFloat16s"] [[deps.LLVMExtra_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "8e76807afb59ebb833e9b131ebf1a8c006510f33" +git-tree-sha1 = "23caf34b74ef5d02b03f2c2853724df7cf6ec3ca" uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" -version = "0.0.38+0" +version = "0.0.46+0" [[deps.LLVMLoopInfo]] git-tree-sha1 = "2e5c102cfc41f48ae4740c7eca7743cc7e7b75ea" @@ -1452,20 +1448,20 @@ version = "1.0.0" [[deps.LLVMOpenMP_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "eb62a3deb62fc6d8822c0c4bef73e4412419c5d8" +git-tree-sha1 = "b7970cef8ae1c990ba0c09cd8bdc1145e006632f" uuid = "1d63c593-3942-5779-bab2-d838dc0a180e" -version = "18.1.8+0" +version = "22.1.7+0" [[deps.LaTeXStrings]] -git-tree-sha1 = "dda21b8cbd6a6c40d9d02a73230f9d70fed6918c" +git-tree-sha1 = "f88f3ccef05a6a72a0cf0ed417c8fd68530f4ab2" uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" -version = "1.4.0" +version = "1.4.1" [[deps.Latexify]] deps = ["Format", "Ghostscript_jll", "InteractiveUtils", "LaTeXStrings", "MacroTools", "Markdown", "OrderedCollections", "Requires"] -git-tree-sha1 = "44f93c47f9cd6c7e431f2f2091fcba8f01cd7e8f" +git-tree-sha1 = "df7566479bd64f20bd16b09960145e70160ffb3b" uuid = "23fbe1c1-3f47-55db-b15f-69d7ec21a316" -version = "0.16.10" +version = "0.16.12" [deps.Latexify.extensions] DataFramesExt = "DataFrames" @@ -1497,9 +1493,9 @@ version = "0.3.1" [[deps.LeftChildRightSiblingTrees]] deps = ["AbstractTrees"] -git-tree-sha1 = "95ba48564903b43b2462318aa243ee79d81135ff" +git-tree-sha1 = "d4816abce26971e3237b46a47b99991f306e4832" uuid = "1d6d02ad-be62-4b6b-8a6d-2f90e265016e" -version = "0.2.1" +version = "0.3.0" [[deps.LibCURL]] deps = ["LibCURL_jll", "MozillaCACerts_jll"] @@ -1556,21 +1552,21 @@ version = "1.18.0+0" [[deps.Libmount_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "97bbca976196f2a1eb9607131cb108c69ec3f8a6" +git-tree-sha1 = "cc3ad4faf30015a3e8094c9b5b7f19e85bdf2386" uuid = "4b2f31a3-9ecc-558c-b454-b3730dcb73e9" -version = "2.41.3+0" +version = "2.42.0+0" [[deps.Libtiff_jll]] deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "LERC_jll", "Libdl", "XZ_jll", "Zlib_jll", "Zstd_jll"] -git-tree-sha1 = "f04133fe05eff1667d2054c53d59f9122383fe05" +git-tree-sha1 = "aebd334d06cee9f24cea70bd19a39749daf73881" uuid = "89763e89-9b03-5906-acba-b20f662cd828" -version = "4.7.2+0" +version = "4.7.3+0" [[deps.Libuuid_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "d0205286d9eceadc518742860bf23f703779a3d6" +git-tree-sha1 = "d620582b1f0cbe2c72dd1d5bd195a9ce73370ab1" uuid = "38a345b3-de98-5d2b-a5d3-14cd9215e700" -version = "2.41.3+0" +version = "2.42.0+0" [[deps.LinearAlgebra]] deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] @@ -1607,9 +1603,9 @@ version = "1.2.0" [[deps.LoopVectorization]] deps = ["ArrayInterface", "CPUSummary", "CloseOpenIntervals", "DocStringExtensions", "HostCPUFeatures", "IfElse", "LayoutPointers", "LinearAlgebra", "OffsetArrays", "PolyesterWeave", "PrecompileTools", "SIMDTypes", "SLEEFPirates", "Static", "StaticArrayInterface", "ThreadingUtilities", "UnPack", "VectorizationBase"] -git-tree-sha1 = "a9fc7883eb9b5f04f46efb9a540833d1fad974b3" +git-tree-sha1 = "514e8475e33c6faf3155efee5f3c10d9e65a11ab" uuid = "bdcacae8-1622-11e9-2a5c-532679323890" -version = "0.12.173" +version = "0.12.174" weakdeps = ["ChainRulesCore", "ForwardDiff", "NNlib", "SpecialFunctions"] [deps.LoopVectorization.extensions] @@ -1619,9 +1615,9 @@ weakdeps = ["ChainRulesCore", "ForwardDiff", "NNlib", "SpecialFunctions"] [[deps.Lux]] deps = ["ADTypes", "Adapt", "ArrayInterface", "ChainRulesCore", "ConcreteStructs", "DiffResults", "DispatchDoctor", "EnzymeCore", "FastClosures", "ForwardDiff", "Functors", "GPUArraysCore", "LinearAlgebra", "LuxCore", "LuxLib", "MLDataDevices", "MacroTools", "Markdown", "NNlib", "Optimisers", "PrecompileTools", "Preferences", "Random", "ReactantCore", "Reexport", "SciMLPublic", "Setfield", "Static", "StaticArraysCore", "Statistics", "UUIDs", "WeightInitializers"] -git-tree-sha1 = "334de475ff414c8eb67f88f57f7b02d40cd8f320" +git-tree-sha1 = "b7654d9b1144792d7fa165add2e07434329e3193" uuid = "b2108857-7c20-44ae-9111-449ecde12c47" -version = "1.31.3" +version = "1.31.4" [deps.Lux.extensions] ComponentArraysExt = "ComponentArrays" @@ -1689,9 +1685,9 @@ version = "1.5.3" [[deps.LuxLib]] deps = ["ArrayInterface", "CPUSummary", "ChainRulesCore", "DispatchDoctor", "EnzymeCore", "FastClosures", "Functors", "KernelAbstractions", "LinearAlgebra", "LuxCore", "MLDataDevices", "Markdown", "NNlib", "Preferences", "Random", "Reexport", "SciMLPublic", "Static", "StaticArraysCore", "Statistics", "UUIDs"] -git-tree-sha1 = "77f3257b18e9fedd39b7b7990f0d3a0800a834ae" +git-tree-sha1 = "6a6453d556f7bc3870d797657636b1ad5f45fd27" uuid = "82251201-b29d-42c6-8e01-566dec8acb11" -version = "1.15.6" +version = "1.15.9" [deps.LuxLib.extensions] AppleAccelerateExt = "AppleAccelerate" @@ -1732,21 +1728,27 @@ version = "1.15.6" [[deps.MCMCDiagnosticTools]] deps = ["AbstractFFTs", "DataAPI", "DataStructures", "Distributions", "LinearAlgebra", "MLJModelInterface", "Random", "SpecialFunctions", "Statistics", "StatsBase", "StatsFuns", "Tables"] -git-tree-sha1 = "2f464b68e84673727b4e4216a6254fba7da5cf4e" +git-tree-sha1 = "345bedeaf7d650f673fefa042dd8384c63de6c68" uuid = "be115224-59cd-429b-ad48-344e309966f0" -version = "0.3.17" +version = "0.3.19" [[deps.MLCore]] deps = ["DataAPI", "SimpleTraits", "Tables"] -git-tree-sha1 = "73907695f35bc7ffd9f11f6c4f2ee8c1302084be" +git-tree-sha1 = "c4ab44fe709638fda6f2c0cbfea2c114932d6c2f" uuid = "c2834f40-e789-41da-a90e-33b280584a8c" -version = "1.0.0" +version = "1.1.0" + + [deps.MLCore.extensions] + MLCorePythonCallExt = "PythonCall" + + [deps.MLCore.weakdeps] + PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" [[deps.MLDataDevices]] deps = ["Adapt", "Functors", "Preferences", "Random", "SciMLPublic"] -git-tree-sha1 = "39a69ca451c3e78b9a6a2e42ef894fdf7505e629" +git-tree-sha1 = "29b00f22be6fd821a214760f0224329f21998a05" uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" -version = "1.17.5" +version = "1.17.10" [deps.MLDataDevices.extensions] AMDGPUExt = "AMDGPU" @@ -1796,16 +1798,11 @@ git-tree-sha1 = "c275fae2e693206b4527dd9d2382aa15359ef3ed" uuid = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" version = "1.12.1" -[[deps.MLStyle]] -git-tree-sha1 = "bc38dff0548128765760c79eb7388a4b37fae2c8" -uuid = "d8e11817-5142-5d16-987a-aa16d5891078" -version = "0.4.17" - [[deps.MLUtils]] -deps = ["ChainRulesCore", "Compat", "DataAPI", "DelimitedFiles", "FLoops", "MLCore", "NNlib", "Random", "ShowCases", "SimpleTraits", "Statistics", "StatsBase", "Tables", "Transducers"] -git-tree-sha1 = "a772d8d1987433538a5c226f79393324b55f7846" +deps = ["ChainRulesCore", "CodeTracking", "Compat", "DataAPI", "DelimitedFiles", "Distributed", "InteractiveUtils", "MLCore", "Mmap", "NNlib", "Random", "ShowCases", "SimpleTraits", "Statistics", "StatsBase", "Tables"] +git-tree-sha1 = "0a589dc0ada20d30b7e9ad13752cf25361875bf2" uuid = "f1d291b0-491e-4a28-83b9-f70985020b54" -version = "0.4.8" +version = "0.4.13" [[deps.MacroTools]] git-tree-sha1 = "1e0228a030642014fe5cfe68c2c0a818f9e3f522" @@ -1813,10 +1810,10 @@ uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" version = "0.5.16" [[deps.Makie]] -deps = ["Animations", "Base64", "CRC32c", "ColorBrewer", "ColorSchemes", "ColorTypes", "Colors", "ComputePipeline", "Contour", "Dates", "DelaunayTriangulation", "Distributions", "DocStringExtensions", "Downloads", "FFMPEG_jll", "FileIO", "FilePaths", "FixedPointNumbers", "Format", "FreeType", "FreeTypeAbstraction", "GeometryBasics", "GridLayoutBase", "ImageBase", "ImageIO", "InteractiveUtils", "Interpolations", "IntervalSets", "InverseFunctions", "Isoband", "KernelDensity", "LaTeXStrings", "LinearAlgebra", "MacroTools", "Markdown", "MathTeXEngine", "Observables", "OffsetArrays", "PNGFiles", "Packing", "Pkg", "PlotUtils", "PolygonOps", "PrecompileTools", "Printf", "REPL", "Random", "RelocatableFolders", "Scratch", "ShaderAbstractions", "Showoff", "SignedDistanceFields", "SparseArrays", "Statistics", "StatsBase", "StatsFuns", "StructArrays", "TriplotBase", "UnicodeFun", "Unitful"] -git-tree-sha1 = "68af66ec16af8b152309310251ecb4fbfe39869f" +deps = ["Animations", "Base64", "CRC32c", "ColorBrewer", "ColorSchemes", "ColorTypes", "Colors", "ComputePipeline", "Contour", "Dates", "DelaunayTriangulation", "Distributions", "DocStringExtensions", "Downloads", "FFMPEG_jll", "FileIO", "FilePaths", "FixedPointNumbers", "Format", "FreeType", "FreeTypeAbstraction", "GeometryBasics", "GridLayoutBase", "ImageBase", "ImageIO", "InteractiveUtils", "Interpolations", "IntervalSets", "InverseFunctions", "Isoband", "KernelDensity", "LaTeXStrings", "LinearAlgebra", "MacroTools", "Markdown", "MathTeXEngine", "Observables", "OffsetArrays", "PNGFiles", "Packing", "Pkg", "PlotUtils", "PolygonOps", "PrecompileTools", "Printf", "REPL", "Random", "RelocatableFolders", "Scratch", "ShaderAbstractions", "SignedDistanceFields", "SparseArrays", "Statistics", "StatsBase", "StatsFuns", "StructArrays", "TriplotBase", "UnicodeFun", "Unitful"] +git-tree-sha1 = "f2c8715d05bf10f9d4dc354e69dee30b6be53239" uuid = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" -version = "0.24.9" +version = "0.24.13" [deps.Makie.extensions] MakieDynamicQuantitiesExt = "DynamicQuantities" @@ -1841,9 +1838,9 @@ version = "1.11.0" [[deps.MathTeXEngine]] deps = ["AbstractTrees", "Automa", "DataStructures", "FreeTypeAbstraction", "GeometryBasics", "LaTeXStrings", "REPL", "RelocatableFolders", "UnicodeFun"] -git-tree-sha1 = "7eb8cdaa6f0e8081616367c10b31b9d9b34bb02a" +git-tree-sha1 = "aa1078778be5a8e5259ff04fbc3d258b3e78d464" uuid = "0a4f8689-d25c-4efe-a92b-7142dfc1aa53" -version = "0.6.7" +version = "0.6.9" [[deps.MbedTLS]] deps = ["Dates", "MbedTLS_jll", "MozillaCACerts_jll", "NetworkOptions", "Random", "Sockets"] @@ -1861,12 +1858,6 @@ git-tree-sha1 = "b513cedd20d9c914783d8ad83d08120702bf2c77" uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e" version = "0.3.3" -[[deps.MicroCollections]] -deps = ["Accessors", "BangBang", "InitialValues"] -git-tree-sha1 = "44d32db644e84c75dab479f1bc15ee76a1a3618f" -uuid = "128add7d-3638-4c79-886c-908ea0c25c34" -version = "0.2.0" - [[deps.Missings]] deps = ["DataAPI"] git-tree-sha1 = "ec4f7fbeab05d7747bdf98eb74d130a2a2ed298d" @@ -1885,33 +1876,36 @@ version = "0.3.4" [[deps.Moshi]] deps = ["ExproniconLite", "Jieko"] -git-tree-sha1 = "53f817d3e84537d84545e0ad749e483412dd6b2a" +git-tree-sha1 = "60beb0717782a3bbe0f7df56decad0ef89048c23" uuid = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d" -version = "0.3.7" +version = "0.3.12" [[deps.MozillaCACerts_jll]] uuid = "14a3606d-f60d-562e-9121-12d972cd8159" version = "2023.12.12" [[deps.MuladdMacro]] -git-tree-sha1 = "cac9cc5499c25554cba55cd3c30543cff5ca4fab" +deps = ["PrecompileTools"] +git-tree-sha1 = "283bf85d4a767481dd924dff0eee1735e95f449e" uuid = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" -version = "0.2.4" +version = "0.2.7" [[deps.NNlib]] -deps = ["Adapt", "Atomix", "ChainRulesCore", "GPUArraysCore", "KernelAbstractions", "LinearAlgebra", "Random", "ScopedValues", "Statistics"] -git-tree-sha1 = "6dc9ffc3a9931e6b988f913b49630d0fb986d0a8" +deps = ["Adapt", "Atomix", "BFloat16s", "ChainRulesCore", "GPUArraysCore", "KernelAbstractions", "LinearAlgebra", "Random", "ScopedValues", "Statistics"] +git-tree-sha1 = "c052421ac94583479b6b535f0884dbbfa27561b2" uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" -version = "0.9.33" +version = "0.9.44" [deps.NNlib.extensions] NNlibAMDGPUExt = "AMDGPU" NNlibCUDACUDNNExt = ["CUDA", "cuDNN"] NNlibCUDAExt = "CUDA" + NNlibEnzymeCoreCUDNNExt = ["EnzymeCore", "CUDA", "cuDNN"] NNlibEnzymeCoreExt = "EnzymeCore" NNlibFFTWExt = "FFTW" NNlibForwardDiffExt = "ForwardDiff" NNlibMetalExt = "Metal" + NNlibMooncakeCUDAExt = ["Mooncake", "CUDA"] NNlibSpecialFunctionsExt = "SpecialFunctions" [deps.NNlib.weakdeps] @@ -1921,6 +1915,7 @@ version = "0.9.33" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" + Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" @@ -1942,21 +1937,15 @@ version = "3.2.2+0" [[deps.NaNMath]] deps = ["OpenLibm_jll"] -git-tree-sha1 = "9b8215b1ee9e78a293f99797cd31375471b2bcae" +git-tree-sha1 = "dbd2e8cd2c1c27f0b584f6661b4309609c5a685e" uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" -version = "1.1.3" - -[[deps.NameResolution]] -deps = ["PrettyPrint"] -git-tree-sha1 = "1a0fa0e9613f46c9b8c11eee38ebb4f590013c5e" -uuid = "71a1bf82-56d0-4bbc-8a3c-48b961074391" -version = "0.1.5" +version = "1.1.4" [[deps.NearestNeighbors]] deps = ["AbstractTrees", "Distances", "StaticArrays"] -git-tree-sha1 = "e2c3bba08dd6dedfe17a17889131b885b8c082f0" +git-tree-sha1 = "576eb4656529c12e77a46b17c23103dfba9fa570" uuid = "b8a86587-4115-5ab1-83bc-aa920d37bbce" -version = "0.4.27" +version = "0.4.29" [[deps.Netpbm]] deps = ["FileIO", "ImageCore", "ImageMetadata"] @@ -1990,9 +1979,9 @@ version = "1.3.6+0" [[deps.OpenBLASConsistentFPCSR_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl"] -git-tree-sha1 = "f2b3b9e52a5eb6a3434c8cca67ad2dde011194f4" +git-tree-sha1 = "38a93f17e431141c6470bb67a88952a7c4f0e928" uuid = "6cdc7f73-28fd-5e50-80fb-958a8875b1af" -version = "0.3.30+0" +version = "0.3.34+0" [[deps.OpenBLAS_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] @@ -2007,9 +1996,9 @@ version = "0.3.3" [[deps.OpenEXR_jll]] deps = ["Artifacts", "Imath_jll", "JLLWrappers", "Libdl", "Zlib_jll"] -git-tree-sha1 = "135492b7e97fc86d9b132b96a54d2d3dd3e0c6a8" +git-tree-sha1 = "0d621a4beb5e48d195f907c3c5b0bea285d9ff9d" uuid = "18a262bb-aa17-5467-a713-aee519bc75cb" -version = "3.4.8+0" +version = "3.4.13+0" [[deps.OpenLibm_jll]] deps = ["Artifacts", "Libdl"] @@ -2024,9 +2013,9 @@ version = "1.6.1" [[deps.OpenSSL_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "2ac022577e5eac7da040de17776d51bb770cd895" +git-tree-sha1 = "d8cce34295c55f47be683580f44791716045b8fe" uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95" -version = "3.5.6+0" +version = "3.5.7+0" [[deps.OpenSpecFun_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl"] @@ -2035,32 +2024,34 @@ uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" version = "0.5.6+0" [[deps.Optimisers]] -deps = ["ChainRulesCore", "ConstructionBase", "Functors", "LinearAlgebra", "Random", "Statistics"] -git-tree-sha1 = "36b5d2b9dd06290cd65fcf5bdbc3a551ed133af5" +deps = ["ChainRulesCore", "Compat", "ConstructionBase", "Functors", "LinearAlgebra", "Random", "Statistics"] +git-tree-sha1 = "b6a586b581eccc60a181145ffd4382099e32e2df" uuid = "3bd65402-5787-11e9-1adc-39752487f4e2" -version = "0.4.7" +version = "0.4.9" [deps.Optimisers.extensions] OptimisersAdaptExt = ["Adapt"] OptimisersEnzymeCoreExt = "EnzymeCore" OptimisersReactantExt = "Reactant" + OptimisersReactantMLDataDevicesExt = ["Reactant", "MLDataDevices"] [deps.Optimisers.weakdeps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" + MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" [[deps.Optimization]] deps = ["ADTypes", "ArrayInterface", "ConsoleProgressMonitor", "DocStringExtensions", "LinearAlgebra", "Logging", "LoggingExtras", "OptimizationBase", "Printf", "Reexport", "SciMLBase", "SparseArrays", "TerminalLoggers"] -git-tree-sha1 = "2c409c814c2d745620fdd55391a66ee514561146" +git-tree-sha1 = "2a7b377ca40f17db759d079f5368571ffe10f6c5" uuid = "7f7a1694-90dd-40f0-9382-eb1efda571ba" -version = "5.5.0" +version = "5.7.0" [[deps.OptimizationBase]] -deps = ["ADTypes", "ArrayInterface", "DifferentiationInterface", "DocStringExtensions", "FastClosures", "LinearAlgebra", "PDMats", "PrecompileTools", "Reexport", "SciMLBase", "SciMLLogging", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "SymbolicIndexingInterface"] -git-tree-sha1 = "a3d7837832e515111c95a02df7dc55edbdf17d8a" +deps = ["ADTypes", "ArrayInterface", "DifferentiationInterface", "DocStringExtensions", "FastClosures", "LinearAlgebra", "PrecompileTools", "SciMLBase", "SciMLLogging", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "SymbolicIndexingInterface"] +git-tree-sha1 = "37e44c9a2b0fadf55a57de550c9a06c35d3a1793" uuid = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" -version = "5.1.0" +version = "5.3.0" [deps.OptimizationBase.extensions] OptimizationChainRulesCoreExt = "ChainRulesCore" @@ -2087,10 +2078,10 @@ version = "5.1.0" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [[deps.OptimizationOptimisers]] -deps = ["Logging", "Optimisers", "OptimizationBase", "Reexport", "SciMLBase"] -git-tree-sha1 = "7caf4c41e3ee6d348381228b6517decea28867e3" +deps = ["Logging", "Optimisers", "OptimizationBase", "Reexport", "SciMLBase", "SciMLLogging"] +git-tree-sha1 = "d7414b3eb9fec8ad9d77aed6730eaff6276e35a0" uuid = "42dfb2eb-d2b4-4451-abcd-913932933ac1" -version = "0.3.16" +version = "0.3.21" [[deps.Opus_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -2099,9 +2090,9 @@ uuid = "91d4177d-7536-5919-b921-800302f37372" version = "1.6.1+0" [[deps.OrderedCollections]] -git-tree-sha1 = "05868e21324cede2207c6f0f466b4bfef6d5e7ee" +git-tree-sha1 = "94ba93778373a53bfd5a0caaf7d809c445292ff4" uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" -version = "1.8.1" +version = "1.8.2" [[deps.PCRE2_jll]] deps = ["Artifacts", "Libdl"] @@ -2110,9 +2101,9 @@ version = "10.42.0+1" [[deps.PDMats]] deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "e4cff168707d441cd6bf3ff7e4832bdf34278e4a" +git-tree-sha1 = "123266c25174ef6c8d4718920abc206452cf8de6" uuid = "90014a1f-27ba-587c-ab20-58faa44d9150" -version = "0.11.37" +version = "0.11.41" weakdeps = ["StatsBase"] [deps.PDMats.extensions] @@ -2120,9 +2111,9 @@ weakdeps = ["StatsBase"] [[deps.PNGFiles]] deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"] -git-tree-sha1 = "cf181f0b1e6a18dfeb0ee8acc4a9d1672499626c" +git-tree-sha1 = "32b657a0d57c310a1a172bfc8c8cf68c5e674323" uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883" -version = "0.4.4" +version = "0.4.5" [[deps.Packing]] deps = ["GeometryBasics"] @@ -2138,9 +2129,9 @@ version = "0.5.12" [[deps.PairPlots]] deps = ["Contour", "Distributions", "KernelDensity", "LinearAlgebra", "MCMCDiagnosticTools", "Makie", "Measures", "Missings", "OrderedCollections", "PolygonOps", "PrecompileTools", "Printf", "Requires", "StaticArrays", "Statistics", "StatsBase", "TableOperations", "Tables"] -git-tree-sha1 = "ca8501a0912f6c1e6533904e28cc80a1f23e9247" +git-tree-sha1 = "e4bf0aceaf2a9443c7c72ff59b1c2b954f210a32" uuid = "43a3c2be-4208-490b-832a-a21dcd55d7da" -version = "3.0.3" +version = "3.0.8" [deps.PairPlots.extensions] MCMCChainsExt = "MCMCChains" @@ -2154,21 +2145,21 @@ version = "3.0.3" [[deps.Pango_jll]] deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "FriBidi_jll", "Glib_jll", "HarfBuzz_jll", "JLLWrappers", "Libdl"] -git-tree-sha1 = "0662b083e11420952f2e62e17eddae7fc07d5997" +git-tree-sha1 = "7126b66b721a605a2fec966a2874c5ed53258eb3" uuid = "36c8627f-9965-5494-a995-c6b170f724f3" -version = "1.57.0+0" +version = "1.58.0+0" [[deps.Parsers]] deps = ["Dates", "PrecompileTools", "UUIDs"] -git-tree-sha1 = "7d2f8f21da5db6a806faf7b9b292296da42b2810" +git-tree-sha1 = "3de8f5e6e90ebfa8d6d1f86997d6cdcd6a912ff3" uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" -version = "2.8.3" +version = "2.8.7" [[deps.Pixman_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LLVMOpenMP_jll", "Libdl"] -git-tree-sha1 = "db76b1ecd5e9715f3d043cec13b2ec93ce015d53" +git-tree-sha1 = "e4a6721aa89e62e5d4217c0b21bd714263779dda" uuid = "30392449-352a-5448-841d-b1acce4e97dc" -version = "0.44.2+0" +version = "0.46.4+0" [[deps.Pkg]] deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "Random", "SHA", "TOML", "Tar", "UUIDs", "p7zip_jll"] @@ -2235,20 +2226,20 @@ uuid = "647866c9-e3ac-4575-94e7-e3d426903924" version = "0.1.2" [[deps.PreallocationTools]] -deps = ["Adapt", "ArrayInterface", "PrecompileTools"] -git-tree-sha1 = "e16b73bf892c55d16d53c9c0dbd0fb31cb7e25da" +deps = ["Adapt", "ArrayInterface", "PrecompileTools", "SciMLPublic"] +git-tree-sha1 = "5e1c95fb1366c7f92c44839b22fc362257895a34" uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46" -version = "1.2.0" +version = "1.5.0" [deps.PreallocationTools.extensions] + PreallocationToolsEnzymeCoreExt = "EnzymeCore" PreallocationToolsForwardDiffExt = "ForwardDiff" PreallocationToolsReverseDiffExt = "ReverseDiff" - PreallocationToolsSparseConnectivityTracerExt = "SparseConnectivityTracer" [deps.PreallocationTools.weakdeps] + EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" - SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" [[deps.PrecompileTools]] deps = ["Preferences"] @@ -2262,22 +2253,19 @@ git-tree-sha1 = "8b770b60760d4451834fe79dd483e318eee709c4" uuid = "21216c6a-2e73-6563-6e65-726566657250" version = "1.5.2" -[[deps.PrettyPrint]] -git-tree-sha1 = "632eb4abab3449ab30c5e1afaa874f0b98b586e4" -uuid = "8162dcfd-2161-5ef2-ae6c-7681170c5f98" -version = "0.2.0" - [[deps.PrettyTables]] deps = ["Crayons", "LaTeXStrings", "Markdown", "PrecompileTools", "Printf", "REPL", "Reexport", "StringManipulation", "Tables"] -git-tree-sha1 = "624de6279ab7d94fc9f672f0068107eb6619732c" +git-tree-sha1 = "69a6c6b1ba14f89970d8d8838551b6d5efa6a685" uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" -version = "3.3.2" +version = "3.4.7" [deps.PrettyTables.extensions] + PrettyTablesExcelExt = "XLSX" PrettyTablesTypstryExt = "Typstry" [deps.PrettyTables.weakdeps] Typstry = "f0ed7684-a786-439e-b1e3-3b82803b501e" + XLSX = "fdbf4ff8-1666-58a4-91e7-1b58723a45e0" [[deps.Primes]] deps = ["IntegerMathUtils"] @@ -2315,15 +2303,15 @@ version = "1.0.2" [[deps.Qt6Base_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Fontconfig_jll", "Glib_jll", "JLLWrappers", "Libdl", "Libglvnd_jll", "OpenSSL_jll", "Vulkan_Loader_jll", "Xorg_libSM_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Xorg_libxcb_jll", "Xorg_xcb_util_cursor_jll", "Xorg_xcb_util_image_jll", "Xorg_xcb_util_keysyms_jll", "Xorg_xcb_util_renderutil_jll", "Xorg_xcb_util_wm_jll", "Zlib_jll", "libinput_jll", "xkbcommon_jll"] -git-tree-sha1 = "d7a4bff94f42208ce3cf6bc8e4e7d1d663e7ee8b" +git-tree-sha1 = "144895f6166994730ee7ff8113b981fc360638f1" uuid = "c0090381-4147-56d7-9ebc-da0b1113ec56" -version = "6.10.2+1" +version = "6.10.2+2" [[deps.Qt6Declarative_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Qt6Base_jll", "Qt6ShaderTools_jll", "Qt6Svg_jll"] -git-tree-sha1 = "d5b7dd0e226774cbd87e2790e34def09245c7eab" +git-tree-sha1 = "159d253ab126d5b29230cf53521899bea4ef4648" uuid = "629bc702-f1f5-5709-abd5-49b8460ea067" -version = "6.10.2+1" +version = "6.10.2+2" [[deps.Qt6ShaderTools_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Qt6Base_jll"] @@ -2394,9 +2382,9 @@ weakdeps = ["FixedPointNumbers"] [[deps.ReactantCore]] deps = ["ExpressionExplorer", "MacroTools"] -git-tree-sha1 = "5b9e0fe7fb2cf3794fd96ac32bf2732aa4bb9776" +git-tree-sha1 = "9ae01d91c9326873cea17b97ad198befdf231661" uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433" -version = "0.1.19" +version = "0.1.21" [[deps.RealDot]] deps = ["LinearAlgebra"] @@ -2477,15 +2465,15 @@ version = "0.9.0" [[deps.Rmath_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "58cdd8fb2201a6267e1db87ff148dd6c1dbd8ad8" +git-tree-sha1 = "6d40b2fe70437b01397d2a4d5b020008da4e7019" uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f" -version = "0.5.1+0" +version = "0.5.2+0" [[deps.Roots]] deps = ["Accessors", "CommonSolve", "Printf"] -git-tree-sha1 = "b2f70f34eb9973572d55c332933c6a04c911f549" +git-tree-sha1 = "7fb25a964849d90a0446366cdefca822e0e84900" uuid = "f2b01f46-fcfa-551c-844a-d8ac1e96c665" -version = "2.2.14" +version = "3.0.6" [deps.Roots.extensions] RootsChainRulesCoreExt = "ChainRulesCore" @@ -2510,9 +2498,9 @@ version = "0.2.1" [[deps.RuntimeGeneratedFunctions]] deps = ["ExprTools", "SHA", "Serialization"] -git-tree-sha1 = "7257165d5477fd1025f7cb656019dcb6b0512c38" +git-tree-sha1 = "65c9e1142f0372bfc16ba14b9edd57737fe0039f" uuid = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" -version = "0.5.17" +version = "0.5.24" [[deps.SHA]] uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" @@ -2531,15 +2519,15 @@ version = "0.1.0" [[deps.SLEEFPirates]] deps = ["IfElse", "Static", "VectorizationBase"] -git-tree-sha1 = "456f610ca2fbd1c14f5fcf31c6bfadc55e7d66e0" +git-tree-sha1 = "72312aa278823c0e99ce31186e22d917d2d11f99" uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa" -version = "0.6.43" +version = "0.6.46" [[deps.SciMLBase]] deps = ["ADTypes", "Accessors", "Adapt", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "Moshi", "PreallocationTools", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLLogging", "SciMLOperators", "SciMLPublic", "SciMLStructures", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface"] -git-tree-sha1 = "908c0bf271604d09393a21c142116ab26f66f67c" +git-tree-sha1 = "d56cb3b924cdf9297c171113e6a6ae9694e370eb" uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -version = "2.154.0" +version = "2.155.2" [deps.SciMLBase.extensions] SciMLBaseChainRulesCoreExt = "ChainRulesCore" @@ -2582,35 +2570,35 @@ version = "2.154.0" [[deps.SciMLLogging]] deps = ["Logging", "LoggingExtras", "Preferences"] -git-tree-sha1 = "0161be062570af4042cf6f69e3d5d0b0555b6927" +git-tree-sha1 = "4e1e21f14a284f892eb62923a356c70a2a0c68e1" uuid = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1" -version = "1.9.1" +version = "1.10.1" weakdeps = ["Tracy"] [deps.SciMLLogging.extensions] SciMLLoggingTracyExt = "Tracy" [[deps.SciMLOperators]] -deps = ["Accessors", "ArrayInterface", "DocStringExtensions", "LinearAlgebra"] -git-tree-sha1 = "234869cf9fee9258a95464b7a7065cc7be84db00" +deps = ["Accessors", "Adapt", "ArrayInterface", "DocStringExtensions", "LinearAlgebra", "SciMLPublic"] +git-tree-sha1 = "fdba76643b52cc34baeda645b7a0c9a7a89db01a" uuid = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" -version = "1.16.0" -weakdeps = ["SparseArrays", "StaticArraysCore"] +version = "1.28.0" +weakdeps = ["LoopVectorization", "SparseArrays"] [deps.SciMLOperators.extensions] + SciMLOperatorsLoopVectorizationExt = "LoopVectorization" SciMLOperatorsSparseArraysExt = "SparseArrays" - SciMLOperatorsStaticArraysCoreExt = "StaticArraysCore" [[deps.SciMLPublic]] -git-tree-sha1 = "0ba076dbdce87ba230fff48ca9bca62e1f345c9b" +git-tree-sha1 = "cf9aaf8b9ed5db993259ea8b24cf2b7ba9bd3b79" uuid = "431bcebd-1456-4ced-9d72-93c2757fff0b" -version = "1.0.1" +version = "1.2.4" [[deps.SciMLStructures]] deps = ["ArrayInterface", "PrecompileTools"] -git-tree-sha1 = "607f6867d0b0553e98fc7f725c9f9f13b4d01a32" +git-tree-sha1 = "53bf620cb2c3763d41495b2a145611c6ca400dcd" uuid = "53ae85a6-f571-4167-b2af-e1d143709226" -version = "1.10.0" +version = "1.10.4" [[deps.ScientificTypesBase]] deps = ["InteractiveUtils"] @@ -2620,9 +2608,9 @@ version = "3.1.0" [[deps.ScopedValues]] deps = ["HashArrayMappedTries", "Logging"] -git-tree-sha1 = "ac4b837d89a58c848e85e698e2a2514e9d59d8f6" +git-tree-sha1 = "67a144433c4ce877ee6d1ada69a124d6b1ecf7be" uuid = "7e506255-f358-4e82-b7e4-beb19740aa63" -version = "1.6.0" +version = "1.6.2" [[deps.Scratch]] deps = ["Dates"] @@ -2632,9 +2620,9 @@ version = "1.3.0" [[deps.SentinelArrays]] deps = ["Dates", "Random"] -git-tree-sha1 = "ebe7e59b37c400f694f52b58c93d26201387da70" +git-tree-sha1 = "084c47c7c5ce5cfecefa0a98dff69eb3646b5a80" uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c" -version = "1.4.9" +version = "1.4.10" [[deps.Serialization]] uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" @@ -2687,9 +2675,9 @@ version = "0.4.8" [[deps.SimpleTraits]] deps = ["InteractiveUtils", "MacroTools"] -git-tree-sha1 = "be8eeac05ec97d379347584fa9fe2f5f76795bcb" +git-tree-sha1 = "7ddb0b49c109481b046972c0e4ab02b2127d6a75" uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" -version = "0.9.5" +version = "0.9.6" [[deps.Sixel]] deps = ["Dates", "FileIO", "ImageCore", "IndirectArrays", "OffsetArrays", "REPL", "libsixel_jll"] @@ -2703,9 +2691,9 @@ version = "1.11.0" [[deps.SortingAlgorithms]] deps = ["DataStructures"] -git-tree-sha1 = "64d974c2e6fdf07f8155b5b2ca2ffa9069b608d9" +git-tree-sha1 = "13cd91cc9be159e3f4d95b857fa2aa383b53772a" uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c" -version = "1.2.2" +version = "1.2.3" [[deps.SparseArrays]] deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] @@ -2714,9 +2702,9 @@ version = "1.11.0" [[deps.SparseConnectivityTracer]] deps = ["ADTypes", "DocStringExtensions", "FillArrays", "LinearAlgebra", "Random", "SparseArrays"] -git-tree-sha1 = "590b72143436e443888124aaf4026a636049e3f5" +git-tree-sha1 = "ad4d1275eeb223cbd4d362563954a661fe12d2f7" uuid = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" -version = "1.2.1" +version = "1.2.2" weakdeps = ["ChainRulesCore", "LogExpFunctions", "NNlib", "NaNMath", "SpecialFunctions"] [deps.SparseConnectivityTracer.extensions] @@ -2728,9 +2716,9 @@ weakdeps = ["ChainRulesCore", "LogExpFunctions", "NNlib", "NaNMath", "SpecialFun [[deps.SparseInverseSubset]] deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "52962839426b75b3021296f7df242e40ecfc0852" +git-tree-sha1 = "eec446511ab8c3293dd846c61c15128392fceed5" uuid = "dc90abb0-5640-4711-901d-7e5b23a2fada" -version = "0.1.2" +version = "0.1.3" [[deps.SparseMatrixColorings]] deps = ["ADTypes", "DocStringExtensions", "LinearAlgebra", "PrecompileTools", "Random", "SparseArrays"] @@ -2753,20 +2741,14 @@ version = "0.4.26" [[deps.SpecialFunctions]] deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] -git-tree-sha1 = "2700b235561b0335d5bef7097a111dc513b8655e" +git-tree-sha1 = "429071b23f4c9a13fb6582f807cc2ef454082408" uuid = "276daf66-3868-5448-9aa4-cd146d93841b" -version = "2.7.2" +version = "2.9.0" weakdeps = ["ChainRulesCore"] [deps.SpecialFunctions.extensions] SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" -[[deps.SplittablesBase]] -deps = ["Setfield", "Test"] -git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5" -uuid = "171d559e-b47b-412a-8079-5efa626c420e" -version = "0.1.15" - [[deps.StableRNGs]] deps = ["Random"] git-tree-sha1 = "4f96c596b8c8258cc7d3b19797854d368f243ddc" @@ -2781,15 +2763,15 @@ version = "0.1.2" [[deps.Static]] deps = ["CommonWorldInvalidations", "IfElse", "PrecompileTools", "SciMLPublic"] -git-tree-sha1 = "49440414711eddc7227724ae6e570c7d5559a086" +git-tree-sha1 = "474a5283ad435618090122872eea6a8165ea6bcf" uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" -version = "1.3.1" +version = "1.4.6" [[deps.StaticArrayInterface]] deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "PrecompileTools", "SciMLPublic", "Static"] -git-tree-sha1 = "aa1ea41b3d45ac449d10477f65e2b40e3197a0d2" +git-tree-sha1 = "2a635e15d5035c53b345077c947f31ff91744078" uuid = "0d7ed370-da01-4f52-bd93-41d350b8b718" -version = "1.9.0" +version = "1.10.0" weakdeps = ["OffsetArrays", "StaticArrays"] [deps.StaticArrayInterface.extensions] @@ -2798,9 +2780,9 @@ weakdeps = ["OffsetArrays", "StaticArrays"] [[deps.StaticArrays]] deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "246a8bb2e6667f832eea063c3a56aef96429a3db" +git-tree-sha1 = "fac51faf3bb96e8bc0bf6f9f39ca4955652776bb" uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.18" +version = "1.9.19" weakdeps = ["ChainRulesCore", "Statistics"] [deps.StaticArrays.extensions] @@ -2836,9 +2818,9 @@ version = "1.8.0" [[deps.StatsBase]] deps = ["AliasTables", "DataAPI", "DataStructures", "IrrationalConstants", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"] -git-tree-sha1 = "aceda6f4e598d331548e04cc6b2124a6148138e3" +git-tree-sha1 = "e4d7a1a0edc20af42689ea6f4f3587a2175d50ee" uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" -version = "0.34.10" +version = "0.34.12" [[deps.StatsFuns]] deps = ["HypergeometricFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"] @@ -2853,15 +2835,15 @@ weakdeps = ["ChainRulesCore", "InverseFunctions"] [[deps.StrideArraysCore]] deps = ["ArrayInterface", "CloseOpenIntervals", "IfElse", "LayoutPointers", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface", "ThreadingUtilities"] -git-tree-sha1 = "83151ba8065a73f53ca2ae98bc7274d817aa30f2" +git-tree-sha1 = "5316097111523c9a970596a5b33cfea5f92e8581" uuid = "7792a7ef-975c-4747-a70f-980b88e8d1da" -version = "0.5.8" +version = "0.5.9" [[deps.StringManipulation]] deps = ["PrecompileTools"] -git-tree-sha1 = "d05693d339e37d6ab134c5ab53c29fce5ee5d7d5" +git-tree-sha1 = "773065c6e0e903924a9d838259be74338422aef2" uuid = "892a3eda-7b42-436c-8928-eab12a02cf0e" -version = "0.4.4" +version = "0.5.0" [[deps.StructArrays]] deps = ["ConstructionBase", "DataAPI", "Tables"] @@ -2879,9 +2861,9 @@ weakdeps = ["Adapt", "GPUArraysCore", "KernelAbstractions", "LinearAlgebra", "Sp [[deps.StructUtils]] deps = ["Dates", "UUIDs"] -git-tree-sha1 = "fa95b3b097bcef5845c142ea2e085f1b2591e92c" +git-tree-sha1 = "2d0fc55c61321ba245c47be599570d11bac50303" uuid = "ec057cc2-7a8d-4b58-b3b3-92acb9f63b42" -version = "2.7.1" +version = "2.8.5" [deps.StructUtils.extensions] StructUtilsMeasurementsExt = ["Measurements"] @@ -2908,9 +2890,9 @@ version = "7.7.0+0" [[deps.SymbolicIndexingInterface]] deps = ["Accessors", "ArrayInterface", "RuntimeGeneratedFunctions", "StaticArraysCore"] -git-tree-sha1 = "94c58884e013efff548002e8dc2fdd1cb74dfce5" +git-tree-sha1 = "2167b9913f3013a1485bdc9bb249123eb8b53cb0" uuid = "2efcf032-c050-4f8e-a9bb-153293bab1f5" -version = "0.3.46" +version = "0.3.54" weakdeps = ["PrettyTables"] [deps.SymbolicIndexingInterface.extensions] @@ -2935,9 +2917,9 @@ version = "1.0.1" [[deps.Tables]] deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"] -git-tree-sha1 = "f2c1efbc8f3a609aadf318094f8fc5204bdaf344" +git-tree-sha1 = "0f38a06c83f0007bbab3cf911262841c9a0f07e0" uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -version = "1.12.1" +version = "1.13.0" [[deps.Tar]] deps = ["ArgTools", "SHA"] @@ -2952,9 +2934,9 @@ version = "0.1.1" [[deps.TerminalLoggers]] deps = ["LeftChildRightSiblingTrees", "Logging", "Markdown", "Printf", "ProgressLogging", "UUIDs"] -git-tree-sha1 = "f133fab380933d042f6796eda4e130272ba520ca" +git-tree-sha1 = "81c9b4137edfe56a56efcdcb35d721b2ce3e2416" uuid = "5d786b92-1e48-4d6f-9151-6b4477ca9bed" -version = "0.1.7" +version = "0.1.8" [[deps.Test]] deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] @@ -2963,15 +2945,15 @@ version = "1.11.0" [[deps.ThreadingUtilities]] deps = ["ManualMemory"] -git-tree-sha1 = "d969183d3d244b6c33796b5ed01ab97328f2db85" +git-tree-sha1 = "7c73336785b21f723f5b143f6e99cf6c43b37dc1" uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5" -version = "0.5.5" +version = "0.5.6" [[deps.TiffImages]] -deps = ["ColorTypes", "DataStructures", "DocStringExtensions", "FileIO", "FixedPointNumbers", "IndirectArrays", "Inflate", "Mmap", "OffsetArrays", "PkgVersion", "PrecompileTools", "ProgressMeter", "SIMD", "UUIDs"] -git-tree-sha1 = "08c10bc34f4e7743f530793d0985bf3c254e193d" +deps = ["CodecZstd", "ColorTypes", "DataStructures", "DocStringExtensions", "FileIO", "FixedPointNumbers", "IndirectArrays", "Inflate", "Mmap", "OffsetArrays", "PkgVersion", "PrecompileTools", "ProgressMeter", "SIMD", "UUIDs"] +git-tree-sha1 = "9ca5f1f2d42f80df4b8c9f6ab5a64f438bbd9976" uuid = "731e570b-9d59-4bfa-96dc-6df516fadf69" -version = "0.11.8" +version = "0.11.9" [[deps.Tracy]] deps = ["ExprTools", "LibTracyClient_jll", "Libdl"] @@ -2990,43 +2972,15 @@ git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742" uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" version = "0.11.3" -[[deps.Transducers]] -deps = ["Accessors", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "ConstructionBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "SplittablesBase", "Tables"] -git-tree-sha1 = "4aa1fdf6c1da74661f6f5d3edfd96648321dade9" -uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999" -version = "0.4.85" - - [deps.Transducers.extensions] - TransducersAdaptExt = "Adapt" - TransducersBlockArraysExt = "BlockArrays" - TransducersDataFramesExt = "DataFrames" - TransducersLazyArraysExt = "LazyArrays" - TransducersOnlineStatsBaseExt = "OnlineStatsBase" - TransducersReferenceablesExt = "Referenceables" - - [deps.Transducers.weakdeps] - Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" - 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.TriplotBase]] git-tree-sha1 = "4d4ed7f294cda19382ff7de4c137d24d16adc89b" uuid = "981d1d27-644d-49a2-9326-4793e63143c3" version = "0.1.0" -[[deps.TruncatedStacktraces]] -deps = ["InteractiveUtils", "MacroTools", "Preferences"] -git-tree-sha1 = "ea3e54c2bdde39062abf5a9758a23735558705e1" -uuid = "781d530d-4396-4725-bb49-402e4bee1e77" -version = "1.4.0" - [[deps.URIs]] -git-tree-sha1 = "bef26fb046d031353ef97a82e3fdb6afe7f21b1a" +git-tree-sha1 = "908fec9df6c5de98548ead82a468c95ccf6cd263" uuid = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4" -version = "1.6.1" +version = "1.7.0" [[deps.UUIDs]] deps = ["Random", "SHA"] @@ -3064,9 +3018,9 @@ weakdeps = ["ConstructionBase", "ForwardDiff", "InverseFunctions", "LaTeXStrings PrintfExt = "Printf" [[deps.UnsafeAtomics]] -git-tree-sha1 = "0f30765c32d66d58e41f4cb5624d4fc8a82ec13b" +git-tree-sha1 = "21b39bfb1fab6156b61fbcba4c86c57b6216d2c3" uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f" -version = "0.3.1" +version = "0.3.2" weakdeps = ["LLVM"] [deps.UnsafeAtomics.extensions] @@ -3079,15 +3033,15 @@ version = "0.2.0" [[deps.VectorizationBase]] deps = ["ArrayInterface", "CPUSummary", "HostCPUFeatures", "IfElse", "LayoutPointers", "Libdl", "LinearAlgebra", "SIMDTypes", "Static", "StaticArrayInterface"] -git-tree-sha1 = "d1d9a935a26c475ebffd54e9c7ad11627c43ea85" +git-tree-sha1 = "807a234dc5e6132dd6cf4c9317ca0917c4001ab3" uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f" -version = "0.21.72" +version = "0.21.74" [[deps.VectorizedRNG]] deps = ["Distributed", "Random", "SLEEFPirates", "UnPack", "VectorizationBase"] -git-tree-sha1 = "5ca83562ba95272d8709c6c91e31e23c3c4c9825" +git-tree-sha1 = "2f294b26d806de1273608d5871080ef5c6e897f5" uuid = "33b4df10-0173-11e9-2a0c-851a7edac40e" -version = "0.2.25" +version = "0.2.26" weakdeps = ["Requires", "StaticArraysCore"] [deps.VectorizedRNG.extensions] @@ -3113,9 +3067,9 @@ version = "0.1.3" [[deps.WeightInitializers]] deps = ["ConcreteStructs", "GPUArraysCore", "LinearAlgebra", "Random", "SpecialFunctions", "Statistics"] -git-tree-sha1 = "d79b71da9e7be904db615bdb99187d30753822a4" +git-tree-sha1 = "d60c958a009b5a25cf8ea4eb3f9a2cf7352a9c6f" uuid = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" -version = "1.3.1" +version = "1.3.4" [deps.WeightInitializers.extensions] AMDGPUExt = "AMDGPU" @@ -3193,9 +3147,9 @@ version = "6.0.2+0" [[deps.Xorg_libXi_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libXext_jll", "Xorg_libXfixes_jll"] -git-tree-sha1 = "a376af5c7ae60d29825164db40787f15c80c7c54" +git-tree-sha1 = "dcb316b3ce0941f195537dda56bea4517fcd3ff5" uuid = "a51aa0fd-4e3c-5386-b890-e753decda492" -version = "1.8.3+0" +version = "1.8.4+0" [[deps.Xorg_libXinerama_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libXext_jll"] @@ -3217,9 +3171,9 @@ version = "0.9.12+0" [[deps.Xorg_libpciaccess_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"] -git-tree-sha1 = "4909eb8f1cbf6bd4b1c30dd18b2ead9019ef2fad" +git-tree-sha1 = "58972370b81423fc546c56a60ed1a009450177c3" uuid = "a65dc6b1-eb27-53a1-bb3e-dea574b5389e" -version = "0.18.1+0" +version = "0.19.0+0" [[deps.Xorg_libxcb_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libXau_jll", "Xorg_libXdmcp_jll"] @@ -3277,9 +3231,9 @@ version = "1.4.7+0" [[deps.Xorg_xkeyboard_config_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_xkbcomp_jll"] -git-tree-sha1 = "00af7ebdc563c9217ecc67776d1bbf037dbcebf4" +git-tree-sha1 = "2e59214e017a55cb87474a00fa76035c82ac0e17" uuid = "33bec58e-1273-512f-9401-5d533626f822" -version = "2.44.0+0" +version = "2.47.0+2" [[deps.Xorg_xtrans_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -3300,27 +3254,29 @@ version = "1.5.7+1" [[deps.Zygote]] deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "GPUArraysCore", "IRTools", "InteractiveUtils", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "PrecompileTools", "Random", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"] -git-tree-sha1 = "a29cbf3968d36022198bcc6f23fdfd70f7caf737" +git-tree-sha1 = "b6a713b80d2b8fd515bf46e9c499dd9cb96bae3d" uuid = "e88e6eb3-aa80-5325-afca-941959d7151f" -version = "0.7.10" +version = "0.7.12" [deps.Zygote.extensions] ZygoteAtomExt = "Atom" + ZygoteCUDAExt = "CUDA" ZygoteColorsExt = "Colors" ZygoteDistancesExt = "Distances" ZygoteTrackerExt = "Tracker" [deps.Zygote.weakdeps] Atom = "c52e3926-4ff0-5f6e-af25-54175e0327b1" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" [[deps.ZygoteRules]] deps = ["ChainRulesCore", "MacroTools"] -git-tree-sha1 = "434b3de333c75fc446aa0d19fc394edafd07ab08" +git-tree-sha1 = "c6a86c133861234450ab260dee01b42abd604095" uuid = "700de1a5-db45-46bc-99cf-38207098b444" -version = "0.2.7" +version = "0.2.8" [[deps.cuDNN]] deps = ["CEnum", "CUDA", "CUDA_Runtime_Discovery", "CUDNN_jll"] @@ -3354,9 +3310,9 @@ version = "0.2.3+0" [[deps.libaom_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "371cc681c00a3ccc3fbc5c0fb91f58ba9bec1ecf" +git-tree-sha1 = "ef17c47d22224aaecc76e597ab21a072e025cf7b" uuid = "a4ae2306-e953-59d6-aa16-d00cac43593b" -version = "3.13.1+0" +version = "3.14.1+0" [[deps.libass_jll]] deps = ["Artifacts", "Bzip2_jll", "FreeType2_jll", "FriBidi_jll", "HarfBuzz_jll", "JLLWrappers", "Libdl", "Zlib_jll"] @@ -3377,9 +3333,9 @@ version = "0.2.2+0" [[deps.libdrm_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libpciaccess_jll"] -git-tree-sha1 = "63aac0bcb0b582e11bad965cef4a689905456c03" +git-tree-sha1 = "28e57478e8a160d346a19c28b3fffb9273bcc9c2" uuid = "8e53e030-5e6c-5a89-a30b-be5b7263a166" -version = "2.4.125+1" +version = "2.4.134+0" [[deps.libevdev_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -3401,9 +3357,9 @@ version = "1.28.1+0" [[deps.libpng_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"] -git-tree-sha1 = "e2a7072fc0cdd7949528c1455a3e5da4122e1153" +git-tree-sha1 = "e51150d5ab85cee6fc36726850f0e627ad2e4aba" uuid = "b53b4c65-9356-5827-b1ea-8c7a1a84506f" -version = "1.6.56+0" +version = "1.6.58+0" [[deps.libsixel_jll]] deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "libpng_jll"] diff --git a/docs/src/tutorials/basic_cpu.md b/docs/src/tutorials/basic_cpu.md index ce97015..e04f93e 100644 --- a/docs/src/tutorials/basic_cpu.md +++ b/docs/src/tutorials/basic_cpu.md @@ -132,23 +132,26 @@ Here, we use synthetic data generated by the package. rng = StableRNG(111) n_site_test = 60 (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( - rng, DoubleMM.DoubleMMCase(); n_site_test, scenario=Val((:omit_r0,))) + rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,))) n_site = size(y_o,2) - n_site_test -i_test = n_site .+ (1:n_site_test) i_train = 1:n_site -test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], - y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site = i_test) +i_test = n_site .+ (1:n_site_test) train_data = (; xM = xM[:, i_train], xP = xP[:, i_train], - y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site = i_train) + y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site_train = 1:length(i_train)) +test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], + y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site_train = zeros(length(i_test))) ``` +The 5th entry, \`, denotes the index of sites in the training +dataset. For the test dataset, we indicate the new site by using zero. + Lets look at them. ``` julia map(size,train_data) ``` - (xM = (5, 800), xP = (16, 800), y_o = (8, 800), y_unc = (8, 800), i_site = (800,)) + (xM = (5, 740), xP = (16, 740), y_o = (8, 740), y_unc = (8, 740), i_site_train = (740,)) All of them have 800 columns, corresponding to 800 sites. There are 5 site-covaritas, 16 values of model drivers, and 8 observations per site. @@ -179,6 +182,29 @@ train_dataloader = MLUtils.DataLoader( CA.getdata.(values(train_data)), batchsize=n_batch, partial=false) ``` +## Setting up the Approximation strategy and the initial parameters + +Here, we are using the `MeanScalingHVIApproximation` approximation of the posterior +density, where the ML model predicts the means of the model at unconstrained scale +and a multiplier, i.e offset at log scale, of the main diagonal of the +covariance matrix for each site. + +We need to specify the structure of scaling blocks (here we use on block for +all parameters), the magnitude of the variance (that will be multiplied by the +site factor) at log-scale. + +Given this information, the initial values to be optimized and some information +in the approximator can be initialized. + +``` julia +block_ends = [length(θM)] # one scaling factor of all parameters +σ = FT(0.1) .* θM[block_ends] # standard deviation of 10% of values of template +logσ2 = FT(2) .* log.(σ) # transform to log_var scale +approx = MeanScalingHVIApproximation(block_ends, logσ2) + +(;ϕqc, approx) = init_hybrid_ϕq(approx, θP, θM, transP; n_site, transM) +``` + ## The Machine-Learning model The machine-learning (ML) part predicts parameters of the posterior of site-specific @@ -190,7 +216,8 @@ The ML model predicts the components of θM and an additional uncertainty factor per site. ``` julia -n_out = length(θM) + 1 # number of individuals to predict, and uncertainty factor +#n_out = length(θM) + 1 # number of individuals to predict, and uncertainty factor +n_out = get_numberof_MLinputs(approx, θM) n_input = n_covar = size(xM,1) g_chain = SimpleChain( @@ -241,38 +268,10 @@ transforms logistic predctions around 0.5 to the range of high prior probability of the parameters, and transforms ML predictions near 0 or 1 towards the outer lower probability ranges. -## Setting up the Approximation strategy and the initial parameters - -Here, we are using the `MeanScalingHVIApproximation` approximation of the posterior -density, where the ML model predicts the means of the model at unconstrained scale -and a multiplier, i.e offset at log scale, of the main diagonal of the -covariance matrix for each site. - -We need to specify the structure of scaling blocks (here we use on block for -all parameters), the magnitude of the variance (that will be multiplied by the -site factor) at log-scale. - -Given this information, the initial values to be optimized and some information -in the approximator can be initialized. - -``` julia -block_ends = [length(θM)] # one scaling factor of all parameters -σ = FT(0.1) .* θM[block_ends] # standard deviation of 10% of values of template -logσ2 = FT(2) .* log.(σ) # transform to log_var scale -approx = MeanScalingHVIApproximation(block_ends, logσ2) - -(;ϕqc, approx) = init_hybrid_ϕq(approx, θP, θM, transP; n_site, transM) -``` - ## Assembling the information All the specifications above are stored in a [`HybridProblem`](@ref) structure. -Before, a [`PBMSiteApplicator`](@ref) is constructed that efficiently -translates an invocation given -a vector of global parameters, and a matrix of site parameters to -invocation of the process based model (PBM), defined at the beginning. - ``` julia f_batch = PBMSiteApplicator(f_doubleMM; θP, θM, θFix, xPvec=xP[:,1]) diff --git a/docs/src/tutorials/basic_cpu.qmd b/docs/src/tutorials/basic_cpu.qmd index 3915a4e..88af48c 100644 --- a/docs/src/tutorials/basic_cpu.qmd +++ b/docs/src/tutorials/basic_cpu.qmd @@ -139,16 +139,20 @@ Here, we use synthetic data generated by the package. rng = StableRNG(111) n_site_test = 60 (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( - rng, DoubleMM.DoubleMMCase(); n_site_test, scenario=Val((:omit_r0,))) + rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,))) n_site = size(y_o,2) - n_site_test -i_test = n_site .+ (1:n_site_test) i_train = 1:n_site -test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], - y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site = i_test) +i_test = n_site .+ (1:n_site_test) train_data = (; xM = xM[:, i_train], xP = xP[:, i_train], - y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site = i_train) + y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site_train = 1:length(i_train)) +test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], + y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site_train = zeros(length(i_test))) + ``` +The 5th entry, `, denotes the index of sites in the training +dataset. For the test dataset, we indicate the new site by using zero. + ```{julia} #| echo: false #| eval: false @@ -327,6 +331,7 @@ ENV["MLDATADEVICES_SILENCE_WARN_NO_GPU"] = 1 ```{julia} using OptimizationOptimisers import Zygote +import CommonSolve: solve solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) diff --git a/docs/src/tutorials/basic_cpu_mean.md b/docs/src/tutorials/basic_cpu_mean.md index 8c6892a..c5bc79b 100644 --- a/docs/src/tutorials/basic_cpu_mean.md +++ b/docs/src/tutorials/basic_cpu_mean.md @@ -136,15 +136,17 @@ Here, we use synthetic data generated by the package. ``` julia rng = StableRNG(111) n_site_test = 60 +# (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( +# rng, DoubleMM.DoubleMMCase(); n_site_test, scenario=Val((:omit_r0,))) (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( - rng, DoubleMM.DoubleMMCase(); n_site_test, scenario=Val((:omit_r0,))) + rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,))) n_site = size(y_o,2) - n_site_test i_test = n_site .+ (1:n_site_test) i_train = 1:n_site test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], - y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site = i_test) + y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site_train = zeros(length(i_test))) train_data = (; xM = xM[:, i_train], xP = xP[:, i_train], - y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site = i_train) + y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site_train = 1:length(i_train)) ``` Lets look at them. @@ -153,7 +155,7 @@ Lets look at them. map(size,train_data) ``` - (xM = (5, 800), xP = (16, 800), y_o = (8, 800), y_unc = (8, 800), i_site = (800,)) + (xM = (5, 740), xP = (16, 740), y_o = (8, 740), y_unc = (8, 740), i_site_train = (740,)) All of them have 800 columns, corresponding to 800 sites. There are 5 site-covaritas, 16 values of model drivers, and 8 observations per site. @@ -346,7 +348,7 @@ however, an applicator for `n_site_pred` sites can be obtained by ``` julia (; probo) = solve(probo_sites, solver; rng, callback = callback_loss(100), # output during fitting - epochs = 20, + epochs = 40, #is_inferred = Val(true), # activate type-checks ); ``` diff --git a/docs/src/tutorials/basic_cpu_mean.qmd b/docs/src/tutorials/basic_cpu_mean.qmd index 7dc45cd..101d552 100644 --- a/docs/src/tutorials/basic_cpu_mean.qmd +++ b/docs/src/tutorials/basic_cpu_mean.qmd @@ -24,6 +24,12 @@ the uncertainty, but estimates uncertainty scaling with the parameter value at unconstrained scale. First load necessary packages. +```{julia} +#| eval: false +#| echo: false +Pkg.activate("docs/src/tutorials") +``` + ```{julia} using HybridVariationalInference using HybridVariationalInference: HybridVariationalInference as HVI @@ -143,24 +149,17 @@ Here, we use synthetic data generated by the package. ```{julia} rng = StableRNG(111) n_site_test = 60 +# (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( +# rng, DoubleMM.DoubleMMCase(); n_site_test, scenario=Val((:omit_r0,))) (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( - rng, DoubleMM.DoubleMMCase(); n_site_test, scenario=Val((:omit_r0,))) + rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,))) n_site = size(y_o,2) - n_site_test i_test = n_site .+ (1:n_site_test) i_train = 1:n_site test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], - y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site = i_test) + y_o = y_o[:, i_test], y_unc = y_unc[:, i_test], i_site_train = zeros(length(i_test))) train_data = (; xM = xM[:, i_train], xP = xP[:, i_train], - y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site = i_train) -``` - -```{julia} -#| echo: false -#| eval: false -() -> begin - (; xM, θP_true, θMs_true, xP, y_true, y_o, y_unc) = - gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,))) -end + y_o = y_o[:, i_train], y_unc = y_unc[:, i_train], i_site_train = 1:length(i_train)) ``` Lets look at them. @@ -384,7 +383,7 @@ however, an applicator for `n_site_pred` sites can be obtained by ```{julia} (; probo) = solve(probo_sites, solver; rng, callback = callback_loss(100), # output during fitting - epochs = 20, + epochs = 40, #is_inferred = Val(true), # activate type-checks ); ``` diff --git a/docs/src/tutorials/blocks_corr.md b/docs/src/tutorials/blocks_corr.md index 203076a..a976024 100644 --- a/docs/src/tutorials/blocks_corr.md +++ b/docs/src/tutorials/blocks_corr.md @@ -90,6 +90,7 @@ length(get_hybridproblem_ϕq(prob)), length(get_hybridproblem_ϕq(prob_ind)) ``` julia using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) diff --git a/docs/src/tutorials/blocks_corr.qmd b/docs/src/tutorials/blocks_corr.qmd index e18a200..9ca1b5d 100644 --- a/docs/src/tutorials/blocks_corr.qmd +++ b/docs/src/tutorials/blocks_corr.qmd @@ -94,6 +94,7 @@ length(get_hybridproblem_ϕq(prob)), length(get_hybridproblem_ϕq(prob_ind)) ```{julia} using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) diff --git a/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-10-output-1.png b/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-10-output-1.png index 8a647e7..3ff7b13 100644 Binary files a/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-10-output-1.png and b/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-10-output-1.png differ diff --git a/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-11-output-1.png b/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-11-output-1.png index 260c17a..6f81cb1 100644 Binary files a/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-11-output-1.png and b/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-11-output-1.png differ diff --git a/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-12-output-1.png b/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-12-output-1.png index 6dcd287..e25c736 100644 Binary files a/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-12-output-1.png and b/docs/src/tutorials/blocks_corr_files/figure-commonmark/cell-12-output-1.png differ diff --git a/docs/src/tutorials/corr_site_global.md b/docs/src/tutorials/corr_site_global.md index 223f1d7..70fe206 100644 --- a/docs/src/tutorials/corr_site_global.md +++ b/docs/src/tutorials/corr_site_global.md @@ -18,7 +18,8 @@ by predicting different (transformed) site parameters, $\zeta_{Ms}$ , given the sampled (transformed) global parameters, $\zeta_P$. $$ -p(\zeta_{Ms}, \zeta_P) = p(\zeta_{Ms} | \zeta_P) p(\zeta_P)$$ +p(\zeta_{Ms}, \zeta_P) = p(\zeta_{Ms} | \zeta_P) p(\zeta_P) +$$ This comes at the cost of running the ML model forward for each sampled global parameter, rather than just once in each sampling of the posterior. @@ -98,6 +99,7 @@ prob_cond = HybridProblem(probo_uncond; g=g_chain_scaled, ϕg=ϕg0, pbm_covars) ``` julia using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) @@ -135,7 +137,7 @@ it has only very weak correlations with the site parameters, $r_1$ and $K_1$. ``` julia i_out = 4 -fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y)",ylabel="sd(y)") +fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y_pred)",ylabel="sd(y_pred)") ymean_cond = [mean(y_cond[i_out,s,:]) for s in axes(y_cond, 2)] ysd_cond = [std(y_cond[i_out,s,:]) for s in axes(y_cond, 2)] scatter!(ax, ymean_cond, ysd_cond, label="conditional") diff --git a/docs/src/tutorials/corr_site_global.qmd b/docs/src/tutorials/corr_site_global.qmd index 19246cf..a2505d2 100644 --- a/docs/src/tutorials/corr_site_global.qmd +++ b/docs/src/tutorials/corr_site_global.qmd @@ -29,7 +29,8 @@ by predicting different (transformed) site parameters, $\zeta_{Ms}$ , given the sampled (transformed) global parameters, $\zeta_P$. $$ -p(\zeta_{Ms}, \zeta_P) = p(\zeta_{Ms} | \zeta_P) p(\zeta_P)$$ +p(\zeta_{Ms}, \zeta_P) = p(\zeta_{Ms} | \zeta_P) p(\zeta_P) +$$ This comes at the cost of running the ML model forward for each sampled global parameter, rather than just once in each sampling of the posterior. @@ -108,6 +109,7 @@ prob_cond = HybridProblem(probo_uncond; g=g_chain_scaled, ϕg=ϕg0, pbm_covars) ```{julia} using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) @@ -143,7 +145,7 @@ it has only very weak correlations with the site parameters, $r_1$ and $K_1$. ```{julia} #| output: true i_out = 4 -fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y)",ylabel="sd(y)") +fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y_pred)",ylabel="sd(y_pred)") ymean_cond = [mean(y_cond[i_out,s,:]) for s in axes(y_cond, 2)] ysd_cond = [std(y_cond[i_out,s,:]) for s in axes(y_cond, 2)] scatter!(ax, ymean_cond, ysd_cond, label="conditional") diff --git a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-10-output-1.png b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-10-output-1.png index d507178..ea7c741 100644 Binary files a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-10-output-1.png and b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-10-output-1.png differ diff --git a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-11-output-1.png b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-11-output-1.png index e12851f..e416569 100644 Binary files a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-11-output-1.png and b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-11-output-1.png differ diff --git a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-12-output-1.png b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-12-output-1.png index 41f9a07..2562a68 100644 Binary files a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-12-output-1.png and b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-12-output-1.png differ diff --git a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-9-output-1.png b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-9-output-1.png index 6fd62cc..b7e7428 100644 Binary files a/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-9-output-1.png and b/docs/src/tutorials/corr_site_global_files/figure-commonmark/cell-9-output-1.png differ diff --git a/docs/src/tutorials/inspect_results.md b/docs/src/tutorials/inspect_results.md index 2413deb..6881eb5 100644 --- a/docs/src/tutorials/inspect_results.md +++ b/docs/src/tutorials/inspect_results.md @@ -48,7 +48,7 @@ Lets look at the results. size(θsP), size(θsMs_tr) ``` - ((1, 400), (800, 2, 400)) + ((1, 400), (740, 2, 400)) The last dimension is the number of samples, the second-last dimension is the respective parameter. `θsMs_tr` has an additional dimension denoting @@ -139,7 +139,7 @@ sampling the posterior and predictive posterior and returns the additional size(y) ``` - (8, 800, 400) + (8, 740, 400) Again, the last dimension is the sample. The other dimensions correspond to the observations we provided for the fitting: diff --git a/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-10-output-1.png b/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-10-output-1.png index 92ef5f0..78f06d0 100644 Binary files a/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-10-output-1.png and b/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-10-output-1.png differ diff --git a/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-14-output-1.png b/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-14-output-1.png index 66844db..c37586f 100644 Binary files a/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-14-output-1.png and b/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-14-output-1.png differ diff --git a/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-8-output-1.png b/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-8-output-1.png index 163116d..706b324 100644 Binary files a/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-8-output-1.png and b/docs/src/tutorials/inspect_results_files/figure-commonmark/cell-8-output-1.png differ diff --git a/docs/src/tutorials/intermediate/basic_cpu_mean_results.jld2 b/docs/src/tutorials/intermediate/basic_cpu_mean_results.jld2 index 0db034a..f79d9ae 100644 Binary files a/docs/src/tutorials/intermediate/basic_cpu_mean_results.jld2 and b/docs/src/tutorials/intermediate/basic_cpu_mean_results.jld2 differ diff --git a/docs/src/tutorials/intermediate/basic_cpu_results.jld2 b/docs/src/tutorials/intermediate/basic_cpu_results.jld2 index 4a62651..2436c25 100644 Binary files a/docs/src/tutorials/intermediate/basic_cpu_results.jld2 and b/docs/src/tutorials/intermediate/basic_cpu_results.jld2 differ diff --git a/docs/src/tutorials/logden_user.md b/docs/src/tutorials/logden_user.md index 5f7db13..417c137 100644 --- a/docs/src/tutorials/logden_user.md +++ b/docs/src/tutorials/logden_user.md @@ -74,23 +74,32 @@ $$ Negating and dropping the constants $-\frac{1}{2} ln (2 \pi)$ and $-\frac{1}{2} ln \sigma^2$ $$ - ln x + \frac{1}{2} \left(\frac{ (\ln x-\mu)^2 }{\sigma^2} \right)$$ + ln x + \frac{1}{2} \left(\frac{ (\ln x-\mu)^2 }{\sigma^2} \right) + $$ ``` julia function neg_logden_lognormalep_lognormal(y_pred, y_obs::AbstractArray{ET}, y_unc; - σ2 = log(abs2(ET(0.02)) + ET(1))) where ET + σ2 = log(abs2(ET(0.01)/ET(0.5)) + ET(1))) where ET lnx = log.(CA.getdata(y_obs)) μ = log.(CA.getdata(y_pred)) - nlogL = sum(lnx .+ abs2.(lnx .- μ) ./ (ET(2) .* σ2)) - #nlogL = sum(lnx + (log(σ2) .+ abs2.(lnx .- μ) ./ σ2) ./ ET(2)) # nonconstant σ2 - return (nlogL) + nlogL_sites = vec(sum( + lnx .+ abs2.(lnx .- μ) ./ (ET(2) .* σ2) + #lnx + (log(σ2) .+ abs2.(lnx .- μ) ./ σ2) ./ ET(2) # nonconstant σ2 + ; dims=1)) + + #nlogL = sum(lnx + (log(σ2) .+ abs2.(lnx .- μ) ./ σ2) ./ ET(2)) # nonconstant σ2 + return (nlogL_sites) end ``` +The σ2 parameters used here is computed according to log(cv + 1) (Wutzler et al. 2020). If information on the different relative error by observation was available, we could pass that information using the DataLoader with `y_unc`, rather than assuming a constant relative error across observations. +Note, that we sum across observations (in rows), to return a vector of logdensities for +each site (in columns). + ## Update the problem and redo the inversion HybridProblem has keyword argument `py` to specify the function of negative Log-Likelihood. @@ -99,6 +108,7 @@ HybridProblem has keyword argument `py` to specify the function of negative Log- prob_lognormal = HybridProblem(prob; py = neg_logden_lognormalep_lognormal) using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) @@ -186,7 +196,7 @@ looks similar and shows correlations between site parameters, $r_1$ and $K_1$. ``` julia i_out = 4 -fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y)",ylabel="sd(y)") +fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y_pred)",ylabel="sd(y_pred)") ymean_normal = [mean(y_normal[i_out,s,:]) for s in axes(y_normal, 2)] ysd_normal = [std(y_normal[i_out,s,:]) for s in axes(y_normal, 2)] scatter!(ax, ymean_normal, ysd_normal, label="normal") diff --git a/docs/src/tutorials/logden_user.qmd b/docs/src/tutorials/logden_user.qmd index 23b86e2..5242a07 100644 --- a/docs/src/tutorials/logden_user.qmd +++ b/docs/src/tutorials/logden_user.qmd @@ -83,23 +83,32 @@ $$ Negating and dropping the constants $-\frac{1}{2} ln (2 \pi)$ and $-\frac{1}{2} ln \sigma^2$ $$ - ln x + \frac{1}{2} \left(\frac{ (\ln x-\mu)^2 }{\sigma^2} \right)$$ + ln x + \frac{1}{2} \left(\frac{ (\ln x-\mu)^2 }{\sigma^2} \right) + $$ ```{julia} function neg_logden_lognormalep_lognormal(y_pred, y_obs::AbstractArray{ET}, y_unc; - σ2 = log(abs2(ET(0.02)) + ET(1))) where ET + σ2 = log(abs2(ET(0.01)/ET(0.5)) + ET(1))) where ET lnx = log.(CA.getdata(y_obs)) μ = log.(CA.getdata(y_pred)) - nlogL = sum(lnx .+ abs2.(lnx .- μ) ./ (ET(2) .* σ2)) - #nlogL = sum(lnx + (log(σ2) .+ abs2.(lnx .- μ) ./ σ2) ./ ET(2)) # nonconstant σ2 - return (nlogL) + nlogL_sites = vec(sum( + lnx .+ abs2.(lnx .- μ) ./ (ET(2) .* σ2) + #lnx + (log(σ2) .+ abs2.(lnx .- μ) ./ σ2) ./ ET(2) # nonconstant σ2 + ; dims=1)) + + #nlogL = sum(lnx + (log(σ2) .+ abs2.(lnx .- μ) ./ σ2) ./ ET(2)) # nonconstant σ2 + return (nlogL_sites) end ``` +The σ2 parameters used here is computed according to log(cv + 1) (Wutzler et al. 2020). If information on the different relative error by observation was available, we could pass that information using the DataLoader with `y_unc`, rather than assuming a constant relative error across observations. +Note, that we sum across observations (in rows), to return a vector of logdensities for +each site (in columns). + ## Update the problem and redo the inversion HybridProblem has keyword argument `py` to specify the function of negative Log-Likelihood. @@ -108,6 +117,7 @@ HybridProblem has keyword argument `py` to specify the function of negative Log- prob_lognormal = HybridProblem(prob; py = neg_logden_lognormalep_lognormal) using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) @@ -192,7 +202,7 @@ looks similar and shows correlations between site parameters, $r_1$ and $K_1$. ```{julia} #| output: true i_out = 4 -fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y)",ylabel="sd(y)") +fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y_pred)",ylabel="sd(y_pred)") ymean_normal = [mean(y_normal[i_out,s,:]) for s in axes(y_normal, 2)] ysd_normal = [std(y_normal[i_out,s,:]) for s in axes(y_normal, 2)] scatter!(ax, ymean_normal, ysd_normal, label="normal") diff --git a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-10-output-1.png b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-10-output-1.png index 2a4c704..b53c5e1 100644 Binary files a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-10-output-1.png and b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-10-output-1.png differ diff --git a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-11-output-1.png b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-11-output-1.png index 8222378..e9323cc 100644 Binary files a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-11-output-1.png and b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-11-output-1.png differ diff --git a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-12-output-1.png b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-12-output-1.png index 8c0988b..6b0edbd 100644 Binary files a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-12-output-1.png and b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-12-output-1.png differ diff --git a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-8-output-1.png b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-8-output-1.png index bbf97cb..c9cade7 100644 Binary files a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-8-output-1.png and b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-8-output-1.png differ diff --git a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-9-output-1.png b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-9-output-1.png index 179bbef..38d52c6 100644 Binary files a/docs/src/tutorials/logden_user_files/figure-commonmark/cell-9-output-1.png and b/docs/src/tutorials/logden_user_files/figure-commonmark/cell-9-output-1.png differ diff --git a/docs/src/tutorials/lux_gpu.md b/docs/src/tutorials/lux_gpu.md index a1e858f..d803fb9 100644 --- a/docs/src/tutorials/lux_gpu.md +++ b/docs/src/tutorials/lux_gpu.md @@ -104,6 +104,7 @@ Currently, only CUDA is tested with this `HybridVariationalInference` package. import CUDA, cuDNN # so that gpu_device() returns a CUDADevice using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) @@ -126,7 +127,7 @@ n_sample_pred = 400 ``` If `gdev_P` is not an `AbstractGPUDevice` then all the results are on CPU. -This is the case, if running this tutorial on a machine without GPU/CUDA setup. +This is the case when running this tutorial on a machine without GPU/CUDA setup. If `gdev_P` is an `AbstractGPUDevice` then the results are GPUArrays and need to be transferred to CPU. @@ -134,7 +135,7 @@ and need to be transferred to CPU. typeof(θsMs_dev) ``` - ComponentArrays.ComponentArray{Float32, 3, Array{Float32, 3}, Tuple{ComponentArrays.Shaped1DAxis{(800,)}, ComponentArrays.Axis{(r1 = 1, K1 = 2)}, ComponentArrays.Shaped1DAxis{(400,)}}} + ComponentArrays.ComponentArray{Float32, 3, Array{Float32, 3}, Tuple{ComponentArrays.Shaped1DAxis{(740,)}, ComponentArrays.Axis{(r1 = 1, K1 = 2)}, ComponentArrays.Shaped1DAxis{(400,)}}} Handling of a `ComponentArrays` backed by GPUArrays can result in errors of scalar indexing. Therefore, use a semicolon diff --git a/docs/src/tutorials/lux_gpu.qmd b/docs/src/tutorials/lux_gpu.qmd index d3acd18..b1e445f 100644 --- a/docs/src/tutorials/lux_gpu.qmd +++ b/docs/src/tutorials/lux_gpu.qmd @@ -122,6 +122,7 @@ end import CUDA, cuDNN # so that gpu_device() returns a CUDADevice using OptimizationOptimisers +import CommonSolve: solve import Zygote solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) diff --git a/docs/src/tutorials/penalty.md b/docs/src/tutorials/penalty.md index bc4b17e..211f584 100644 --- a/docs/src/tutorials/penalty.md +++ b/docs/src/tutorials/penalty.md @@ -11,15 +11,13 @@ the solver to converge to global minimum. ## Motivation The basic cost in HVI is the negative log of the joint probability, i.e. -the likelihood of the observations given the parameters \* prior probability +the likelihood of the observations given the parameters times the prior probability of the parameters. -Sometimes there is additional knowledge not encoded in the prior, such as -one parameter must be larger than another, or entropy-weights of the -ML-parameters, and the solver accept a function to add additional loss terms. -The loglikelihood function assigns a cost to the mismatch between predictions and -observations. This often needs to be customized to the specific inversion. - +Sometimes one wants to specify additional knowledge not encoded in the prior, +such as one parameter must be larger than another, or entropy-weights of the +ML-parameters. +For such cases, the solver accept a function that computes additional loss terms. This guide walks through the specification of such additional penalties. First load necessary packages. @@ -39,7 +37,7 @@ function assuming observation error to be distributed independently normal. ``` julia fname = "intermediate/basic_cpu_results.jld2" print(abspath(fname)) -prob = probo_normal = load(fname, "probo"); +prob = load(fname, "probo"); ``` ## Write function to compute the penalty loss @@ -50,30 +48,40 @@ In this example we want to avoid local minima when parameter, `r1`, is larger th 70% of the maximum observation. ``` julia -# compute the maximum of observed rates at each site +# compute the maximum of observed rates at each training site y_obs = get_hybridproblem_train_dataloader(prob).data[3] const y_obs_max = map(col -> maximum(x -> isfinite(x) ? x : zero(x), col), eachcol(y_obs)) function compute_penalty_r1(y_pred::AbstractMatrix, addq_pred::AbstractMatrix, - θMs_tr::AbstractMatrix, θP::AbstractVector, i_sites, - ϕg, ϕq::AbstractVector) - # get the maximum of current batch from closure of this function - y_obs_max_sites = y_obs_max[i_sites] - # add a penalty if r1 is larger than 0.95 times the maximum - penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- 0.95 .* y_obs_max_sites) + θMs_tr::AbstractMatrix{T}, θP::AbstractVector, i_sites_train, + ϕq::AbstractVector) where T + penalty = if isempty(i_sites_train) || iszero(i_sites_train[1]) + # for unobserved sites (or test data), assign zero penalty + n_site = size(θMs_tr,1) + zeros(T, n_site) + else + # get the maximum of current batch from closure of this function + y_obs_max_sites = y_obs_max[i_sites_train] + # add a penalty if r1 is larger than 0.95 times the maximum + penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- 0.95 .* y_obs_max_sites) + end (; penalty) end ``` -The PenaltyComputer receives argument, `i_sites`, which can be used to index precomputed observation maxima. +The penalty is computed for each site in the batch separately. +Here, the `i_sites_train` argument is used to index into precomputed observation maxima. ## Update the problem and redo the inversion -HybridProblem has keyword argument `penalty_computer` to specify the Callable +HybridProblem provides the `penalty_computer` keyword argument to specify the Callable that computes the penalty. It defaults to `ZeroPenaltyComputer`, which returns zero penalty cost. -We can pass the function directly or alternatively construct a [`CustomPenaltyComputer`](@ref) and update the problem. +The argument accepts the function directly +(or alternatively construct a [`CustomPenaltyComputer`](@ref)). +Next, the updated problem is solved, using [`HybridPointSolver`](@ref) and +[`HybridPosteriorSolver`](@ref). ``` julia #prob_pen = HybridProblem(prob; penalty_computer = compute_penalty_r1) @@ -102,8 +110,10 @@ solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) ## Inspect the computed maxima -Function predic_hvi also evaluates the penalties. Internally, the penalty function is -called for each sample, but only the average is computed and returned. +The [`predict_hvi`](@ref) function also evaluates the penalties. +Internally, the penalty function is +called for each sample, but predict_hvi computes and returns the average +for each site. ``` julia rng = StableRNGs.StableRNG(112) @@ -112,8 +122,8 @@ n_sample_pred = 200 size(penalties) ``` -The penalties object is a ComponentMatrix, and we can look at a specific site -and a named component returned by +The penalties object is a ComponentMatrix, and we can look at a specific site (rows) +and a named component (column) returned by ``` julia i_site = 3 @@ -123,11 +133,12 @@ penalties[i_site, :penalty] ## Writing a customized PenaltyComputer In the above example, the maximum of the observations in the batch -are accesses by a global variable. +are accessed by a global variable, which can lead to type stability and +performance problems. This can be improved. The precomputed maxima can be stored in a struct implementing type `AbstractPenaltyComputer` -and function `compute_penalty`. +with associated function `compute_penalty`. ``` julia struct R1PenaltyComputer{T} <: AbstractPenaltyComputer where T @@ -138,14 +149,21 @@ function R1PenaltyComputer(ys::AbstractMatrix) R1PenaltyComputer(r_max) end function HybridVariationalInference.compute_penalty( - pc::R1PenaltyComputer, + pc::R1PenaltyComputer{T}, y_pred::AbstractMatrix, addq_pred::AbstractMatrix, θMs_tr::AbstractMatrix, θP::AbstractVector, - i_sites::AbstractVector, - ϕg, ϕq::AbstractVector - ) - # @assert pc.r_max[i_sites] == 0.95 .* map(col -> maximum(x -> isfinite(x) ? x : zero(x), col), eachcol(y_obs)) - # add a penalty if r1 is larger r_max - penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- pc.r_max[i_sites]) + i_sites_train::AbstractVector, + ϕq::AbstractVector + ) where T + penalty = if isempty(i_sites_train) || iszero(i_sites_train[1]) + # for unobserved sites (or test data), assign zero penalty + n_site = size(θMs_tr,1) + zeros(T, n_site) + else + # get the maximum of current batch from struct + r_max_sites = pc.r_max[i_sites_train] + # add a penalty if r1 is larger than the precomputed threshold + penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- r_max_sites) + end (;penalty) end diff --git a/docs/src/tutorials/penalty.qmd b/docs/src/tutorials/penalty.qmd index 81ff9df..6112836 100644 --- a/docs/src/tutorials/penalty.qmd +++ b/docs/src/tutorials/penalty.qmd @@ -21,15 +21,13 @@ the solver to converge to global minimum. ## Motivation The basic cost in HVI is the negative log of the joint probability, i.e. -the likelihood of the observations given the parameters * prior probability +the likelihood of the observations given the parameters times the prior probability of the parameters. -Sometimes there is additional knowledge not encoded in the prior, such as -one parameter must be larger than another, or entropy-weights of the -ML-parameters, and the solver accept a function to add additional loss terms. -The loglikelihood function assigns a cost to the mismatch between predictions and -observations. This often needs to be customized to the specific inversion. - +Sometimes one wants to specify additional knowledge not encoded in the prior, +such as one parameter must be larger than another, or entropy-weights of the +ML-parameters. +For such cases, the solver accept a function that computes additional loss terms. This guide walks through the specification of such additional penalties. First load necessary packages. @@ -48,7 +46,7 @@ function assuming observation error to be distributed independently normal. ```{julia} fname = "intermediate/basic_cpu_results.jld2" print(abspath(fname)) -prob = probo_normal = load(fname, "probo"); +prob = load(fname, "probo"); ``` ## Write function to compute the penalty loss @@ -59,30 +57,40 @@ In this example we want to avoid local minima when parameter, `r1`, is larger th 70% of the maximum observation. ```{julia} -# compute the maximum of observed rates at each site +# compute the maximum of observed rates at each training site y_obs = get_hybridproblem_train_dataloader(prob).data[3] const y_obs_max = map(col -> maximum(x -> isfinite(x) ? x : zero(x), col), eachcol(y_obs)) function compute_penalty_r1(y_pred::AbstractMatrix, addq_pred::AbstractMatrix, - θMs_tr::AbstractMatrix, θP::AbstractVector, i_sites, - ϕg, ϕq::AbstractVector) - # get the maximum of current batch from closure of this function - y_obs_max_sites = y_obs_max[i_sites] - # add a penalty if r1 is larger than 0.95 times the maximum - penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- 0.95 .* y_obs_max_sites) + θMs_tr::AbstractMatrix{T}, θP::AbstractVector, itrain_sites, + ϕq::AbstractVector) where T + penalty = if isempty(itrain_sites) || iszero(itrain_sites[1]) + # for unobserved sites (or test data), assign zero penalty + n_site = size(θMs_tr,1) + zeros(T, n_site) + else + # get the maximum of current batch from closure of this function + y_obs_max_sites = y_obs_max[itrain_sites] + # add a penalty if r1 is larger than 0.95 times the maximum + penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- 0.95 .* y_obs_max_sites) + end (; penalty) end ``` -The PenaltyComputer receives argument, `i_sites`, which can be used to index precomputed observation maxima. +The penalty is computed for each site in the batch separately. +Here, the `itrain_sites` argument is used to index into precomputed observation maxima. ## Update the problem and redo the inversion -HybridProblem has keyword argument `penalty_computer` to specify the Callable +HybridProblem provides the `penalty_computer` keyword argument to specify the Callable that computes the penalty. It defaults to `ZeroPenaltyComputer`, which returns zero penalty cost. -We can pass the function directly or alternatively construct a [`CustomPenaltyComputer`](@ref) and update the problem. +The argument accepts the function directly +(or alternatively construct a [`CustomPenaltyComputer`](@ref)). +Next, the updated problem is solved, using [`HybridPointSolver`](@ref) and +[`HybridPosteriorSolver`](@ref). ```{julia} #prob_pen = HybridProblem(prob; penalty_computer = compute_penalty_r1) @@ -110,8 +118,11 @@ solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3) ``` ## Inspect the computed maxima -Function predic_hvi also evaluates the penalties. Internally, the penalty function is -called for each sample, but only the average is computed and returned. +The [`predict_hvi`](@ref) function also evaluates the penalties. +Internally, the penalty function is +called for each sample, but predict_hvi computes and returns the average +for each site. + ```{julia} rng = StableRNGs.StableRNG(112) n_sample_pred = 200 @@ -119,8 +130,8 @@ n_sample_pred = 200 size(penalties) ``` -The penalties object is a ComponentMatrix, and we can look at a specific site -and a named component returned by +The penalties object is a ComponentMatrix, and we can look at a specific site (rows) +and a named component (column) returned by ```{julia} i_site = 3 @@ -130,11 +141,12 @@ penalties[i_site, :penalty] ## Writing a customized PenaltyComputer In the above example, the maximum of the observations in the batch -are accesses by a global variable. +are accessed by a global variable, which can lead to type stability and +performance problems. This can be improved. The precomputed maxima can be stored in a struct implementing type `AbstractPenaltyComputer` -and function `compute_penalty`. +with associated function `compute_penalty`. ```{julia} struct R1PenaltyComputer{T} <: AbstractPenaltyComputer where T @@ -145,14 +157,21 @@ function R1PenaltyComputer(ys::AbstractMatrix) R1PenaltyComputer(r_max) end function HybridVariationalInference.compute_penalty( - pc::R1PenaltyComputer, + pc::R1PenaltyComputer{T}, y_pred::AbstractMatrix, addq_pred::AbstractMatrix, θMs_tr::AbstractMatrix, θP::AbstractVector, - i_sites::AbstractVector, - ϕg, ϕq::AbstractVector - ) - # @assert pc.r_max[i_sites] == 0.95 .* map(col -> maximum(x -> isfinite(x) ? x : zero(x), col), eachcol(y_obs)) - # add a penalty if r1 is larger r_max - penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- pc.r_max[i_sites]) + itrain_sites::AbstractVector, + ϕq::AbstractVector + ) where T + penalty = if isempty(itrain_sites) || iszero(itrain_sites[1]) + # for unobserved sites (or test data), assign zero penalty + n_site = size(θMs_tr,1) + zeros(T, n_site) + else + # get the maximum of current batch from struct + r_max_sites = pc.r_max[itrain_sites] + # add a penalty if r1 is larger than the precomputed threshold + penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- r_max_sites) + end (;penalty) end diff --git a/intermediate/basic_cpu_results.jld2 b/intermediate/basic_cpu_results.jld2 deleted file mode 100644 index 9ff7979..0000000 Binary files a/intermediate/basic_cpu_results.jld2 and /dev/null differ diff --git a/projects/doubleMM.jl b/projects/doubleMM.jl index 9faa1c6..9e4988e 100644 --- a/projects/doubleMM.jl +++ b/projects/doubleMM.jl @@ -36,7 +36,7 @@ prob0_ = HybridProblem(DoubleMM.DoubleMMCase(); scenario); ) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario); n_site, n_batch = get_hybridproblem_n_site_and_batch(prob0_; scenario) ζP_true, ζMs_true = log.(θP_true), log.(θMs_true) -i_sites = 1:n_site +itrain_sites = 1:n_site n_site, n_batch = get_hybridproblem_n_site_and_batch(prob0_; scenario) train_dataloader = MLUtils.DataLoader( (xM, xP, y_o, y_unc, 1:n_site); @@ -138,10 +138,10 @@ end int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector( ϕg = 1:length(prob0.ϕg), θP = prob0.θP)) loss_gf = get_loss_gf(prob0.g, prob0.transM, prob0.transP, prob0.f, Float32[], py, int_ϕθP) - loss_gf(vcat(prob3.ϕg, prob3.θP), xM, xP, y_o, y_unc, i_sites)[1] - loss_gf(vcat(prob3o.ϕg, prob3o.θP), xM, xP, y_o, y_unc, i_sites)[1] + loss_gf(vcat(prob3.ϕg, prob3.θP), xM, xP, y_o, y_unc, itrain_sites)[1] + loss_gf(vcat(prob3o.ϕg, prob3o.θP), xM, xP, y_o, y_unc, itrain_sites)[1] # - loss_gf(vcat(prob2o.ϕg, prob2o.θP), xM, xP, y_o, y_unc, i_sites)[1] + loss_gf(vcat(prob2o.ϕg, prob2o.θP), xM, xP, y_o, y_unc, itrain_sites)[1] end end @@ -662,10 +662,10 @@ lineplot!(plt, 0, 1) ζsMs_hmc = log.(θsMs_hmc) # int_pms = interpreters.PMs # par_pos = int_pms(1:length(int_pms)) - #i_sites = 1:10 - i_sites = 1:5 - #i_sites = 6:10 - #i_sites = 11:15 + #itrain_sites = 1:10 + itrain_sites = 1:5 + #itrain_sites = 6:10 + #itrain_sites = 11:15 scen = vcat( fill(:hvi,size(ζsP_hvi,2)), fill(:hmc,size(ζsP_hmc,2)), @@ -679,11 +679,11 @@ lineplot!(plt, 0, 1) ζsP_hvi[i_par, :], ζsP_hmc[i_par,:], ζsP_hvi_indep[i_par,:], ζsP_hvi_neglect_cor[i_par,:]), variable = lower_lastdigits.(keys(θP_true)[i_par]), - site = "site $(i_sites[1])", + site = "site $(itrain_sites[1])", Method = scen ) end - dfMs = mapreduce(vcat, i_sites) do i_site + dfMs = mapreduce(vcat, itrain_sites) do i_site mapreduce(vcat, axes(θM,1)) do i_par #pos = par_pos.Ms[i_par, i_site] DataFrame( @@ -704,11 +704,11 @@ lineplot!(plt, 0, 1) DataFrame( value = ζP_true[i_par], variable = lower_lastdigits.(keys(θP)[i_par]), - site = "site $(i_sites[1])", + site = "site $(itrain_sites[1])", Method = :true ) end, - mapreduce(vcat, i_sites) do i_site + mapreduce(vcat, itrain_sites) do i_site mapreduce(vcat, axes(θM,1)) do i_par DataFrame( value = ζMs_true[i_par, i_site], @@ -733,7 +733,7 @@ lineplot!(plt, 0, 1) axis=(xlabelvisible=false,yticklabelsvisible=false), scales(Color = (; palette = color_methods)), ); - legend!(fig[length(i_sites),1], ffig, ; tellwidth=false, halign=:left, valign=:bottom , margin=(10, 10, 10, 10)) + legend!(fig[length(itrain_sites),1], ffig, ; tellwidth=false, halign=:left, valign=:bottom , margin=(10, 10, 10, 10)) fig end () -> begin @@ -758,7 +758,7 @@ lineplot!(plt, 0, 1) i_obss = [1,4,8] #i_obss = 1:8 dfy = mapreduce(vcat, i_obss) do i_obs - mapreduce(vcat, i_sites) do i_site + mapreduce(vcat, itrain_sites) do i_site vcat( DataFrame( value = y_hmc[i_obs,i_site,:], @@ -780,7 +780,7 @@ lineplot!(plt, 0, 1) end end dfyt = mapreduce(vcat, i_obss) do i_obs - mapreduce(vcat, i_sites) do i_site + mapreduce(vcat, itrain_sites) do i_site vcat( DataFrame( value = y_true[i_obs,i_site], @@ -839,7 +839,7 @@ Compute standard deviation and correlation for predicted parameters on unconstra _ζsP: n_P x n_pred matrix of draws of predicted cross-sites parameters _ζsMs: n_site x n_M x n_pred of draws of predicted physical parameters -returns sdP (n_P), sdMs (n_site x n_M), cor_PMs n_P + (n_M * length(i_sites)) square matrix +returns sdP (n_P), sdMs (n_site x n_M), cor_PMs n_P + (n_M * length(itrain_sites)) square matrix """ function compute_sd_cor_PMs(_ζsP, _ζsMs; i_sites_inspect = [1,2,3]) mP = mean(_ζsP; dims=2) diff --git a/src/AbstractHybridProblem.jl b/src/AbstractHybridProblem.jl index 2614891..2c7cec5 100644 --- a/src/AbstractHybridProblem.jl +++ b/src/AbstractHybridProblem.jl @@ -45,8 +45,8 @@ end compute_penalty(::PenaltyComputerOrFunction, y_pred::AbstractMatrix, addq_pred::AbstractMatrix, θMs::AbstractMatrix, θP::AbstractVector, - i_sites::AbstractVector{<:Int}, - ϕg, ϕq::AbstractVector) + itrain_sites::AbstractVector{<:Int}, + ϕq::AbstractVector) Add additional loss terms during the HVI fit. The basic cost in HVI is the negative log of the joint probability, i.e. @@ -62,8 +62,9 @@ Arguments - addq_pred::AbstractMatrix: Additional quantities computed by the PBM - θMs_tr::AbstractMatrix: site parameters (with sites in rows and parameters in columns) - θP::AbstractVector: global parameters -- i_sites: indices of sites in the minibatch, useful for using precoputed quantities -- ϕg: ML-model parameters, +- itrain_sites: indices of sites in the minibatch, useful for using precoputed quantities + Note, that this refers to index of the site within the training dataset + rather than some global identifier. - ϕq::AbstractVector, additional parameters of the posterior Returns a NamedTuple @@ -259,11 +260,11 @@ Return a DataLoader that provides a tuple of - `xP`: Iterator of process-model drivers, with one element per site - `y_o`: matrix of observations with added noise, with one column per site - `y_unc`: matrix `sizeof(y_o)` of uncertainty information -- `i_sites`: Vector of indices of sites in the minibatch +- `itrain_sites`: Vector of indices of sites in the minibatch """ function get_hybridproblem_train_dataloader end function get_hybridproblem_test_data end - +function get_i_sites_test end # to support construct_dataloader_from_synthetic """ construct_dataloader_from_synthetic(rng::AbstractRNG, prob::AbstractHybridProblem; @@ -272,25 +273,27 @@ function get_hybridproblem_test_data end Construct a dataloader based on `gen_hybridproblem_synthetic`. """ function construct_dataloader_from_synthetic(rng::AbstractRNG, prob::AbstractHybridProblem; - scenario = (), n_batch, n_site_test = 0, + scenario = (), n_batch, #gdev = :use_gpu ∈ scenario ? gpu_device() : identity, ) - (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, prob; n_site_test, scenario) - n_site = size(xM,2) - n_site_test + (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, prob; scenario) + i_sites_test = get_i_sites_test(prob; scenario) + n_site = size(xM, 2) @assert size(xP,2) == n_site @assert size(y_o,2) == n_site @assert size(y_unc,2) == n_site - i_sites = 1:n_site - train_loader = MLUtils.DataLoader((CA.getdata(xM), CA.getdata(xP), y_o, y_unc, i_sites); + i_sites = setdiff(1:n_site, i_sites_test) + itrain_sites = 1:length(i_sites) # index into training data + train_loader = MLUtils.DataLoader(( + CA.getdata(xM)[:,i_sites], CA.getdata(xP)[:,i_sites], + y_o[:,i_sites], y_unc[:,i_sites], itrain_sites); batchsize = n_batch, partial = false) return (train_loader) end """ - gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader; gdev_M, gdev_P, - batchsize = dataloader.batchsize, - partial = dataloader.partial + gdev_hybridproblem_dataloader(dataloader; gdev_M, gdev_P, ) Put relevant parts of the DataLoader to gpu, depending on scenario. @@ -302,10 +305,10 @@ function gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader; gdevs = n # gdev = gpu_device(), # gdev_M = :use_gpu ∈ _val_value(scenario) ? gdev : identity, # gdev_P = :f_on_gpu ∈ _val_value(scenario) ? gdev : identity, - batchsize = dataloader.batchsize, - partial = dataloader.partial ) - # xM, xP, y_o, y_unc, i_sites = dataloader.data + batchsize = dataloader.batchsize + partial = dataloader.partial + # xM, xP, y_o, y_unc, itrain_sites = dataloader.data # xM_dev = gdev_M(xM) # xP_dev, y_o_dev, y_unc_dev = (gdev_P(xP), gdev_P(y_o), gdev_P(y_unc)) data_dev = gdev_hybridproblem_data(dataloader.data; gdev_M, gdev_P) @@ -313,6 +316,11 @@ function gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader; gdevs = n return(train_loader_dev) end +function gdev_hybridproblem_data(view::WeightedObsView; kwargs...) + data = gdev_hybridproblem_data(view.data; kwargs...) + WeightedObsView(data, view.weights) +end + function gdev_hybridproblem_data(data::Tuple; gdevs = nothing, gdev_M = gdevs.gdev_M, gdev_P = gdevs.gdev_P, @@ -321,10 +329,10 @@ function gdev_hybridproblem_data(data::Tuple; gdevs = nothing, # gdev_M = :use_gpu ∈ _val_value(scenario) ? gdev : identity, # gdev_P = :f_on_gpu ∈ _val_value(scenario) ? gdev : identity, ) - xM, xP, y_o, y_unc, i_sites = data + xM, xP, y_o, y_unc, itrain_sites = data xM_dev = gdev_M(xM) xP_dev, y_o_dev, y_unc_dev = (gdev_P(xP), gdev_P(y_o), gdev_P(y_unc)) - (xM_dev, xP_dev, y_o_dev, y_unc_dev, i_sites) + (xM_dev, xP_dev, y_o_dev, y_unc_dev, itrain_sites) end """ @@ -402,6 +410,16 @@ end Return a AbstractHVIApproximation that should be used with this problem """ function get_hybridproblem_HVIApproximation end +""" + get_hybridproblem_ranef(::AbstractHybridProblem; scenario) + +Return a AbstractRandomEffects that should be used with this problem. +The default is `NullRandomEffects()`, which does not add any random effects. +""" +function get_hybridproblem_ranef(prob::AbstractHybridProblem; scenario = ()) + NullRandomEffects() +end + diff --git a/src/DoubleMM/f_doubleMM.jl b/src/DoubleMM/f_doubleMM.jl index 82c58f4..a885aa8 100644 --- a/src/DoubleMM/f_doubleMM.jl +++ b/src/DoubleMM/f_doubleMM.jl @@ -131,12 +131,15 @@ function f_doubleMM_sites(θc_tr::CA.ComponentMatrix, xPc::CA.ComponentMatrix) # #(rep_fac .* p1') # move to computation below to save allocation # end # + # inline function to fuse all into one dot-expression + + #@inline θ(par::Val) = is_valid .* CA.getdata(view(θc_tr,:, par))' + + # need to multiply already each parameter by is_valid to prevent nan_inf_gradients r0 = is_valid .* CA.getdata(θc_tr[:, Val(:r0)])' r1 = is_valid .* CA.getdata(θc_tr[:, Val(:r1)])' K1 = is_valid .* CA.getdata(θc_tr[:, Val(:K1)])' K2 = is_valid .* CA.getdata(θc_tr[:, Val(:K2)])' - # - #, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par # each variable is a matrix (n_obs x n_site) #r0 .+ r1 .* S1 ./ (K1 .+ S1) .* S2 ./ (K2 .+ S2) @@ -352,13 +355,13 @@ function HVI.get_hybridproblem_train_dataloader(prob::DoubleMMCase; scenario::Va # HVI.get_hybridproblem_n_covar, which be default relies on the train_dataloader dl = construct_dataloader_from_synthetic(rng, prob; scenario, n_batch, kwargs...) if (:driverNAN ∈ scen) - (xM, xP, y_o, y_unc, i_sites) = dl.data + (xM, xP, y_o, y_unc, itrain_sites) = dl.data # set the last two entries of the S1 drivers and observations of the second site NaN is_obs = 7:8 i_site = 2 xP[is_obs,i_site] .= NaN y_o[is_obs,i_site] .= NaN - train_loader = MLUtils.DataLoader((CA.getdata(xM), CA.getdata(xP), y_o, y_unc, i_sites); + train_loader = MLUtils.DataLoader((CA.getdata(xM), CA.getdata(xP), y_o, y_unc, itrain_sites); batchsize = n_batch, partial = false) else dl @@ -366,38 +369,42 @@ function HVI.get_hybridproblem_train_dataloader(prob::DoubleMMCase; scenario::Va end function HVI.get_hybridproblem_test_data(prob::DoubleMMCase; scenario::Val{scen}, - rng::AbstractRNG = StableRNG(211), kwargs... + rng::AbstractRNG = StableRNG(111), kwargs... ) where {scen} - n_site_test = 60 - (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic( - rng, prob; scenario, n_site_test) - n_site_all = size(xM,2) - i_test = (n_site_all - n_site_test + 1):n_site_all - (; xM = xM[:, i_test], xP = xP[:, i_test], y_o = y_o[:, i_test], - y_unc = y_unc[:, i_test], i_sites = i_test) + (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, prob; scenario) + i_sites_test = HVI.get_i_sites_test(prob; scenario) + itrain_sites = zeros(length(i_sites_test)) # index into training data + (; xM = xM[:, i_sites_test], xP = xP[:, i_sites_test], y_o = y_o[:, i_sites_test], + y_unc = y_unc[:, i_sites_test], itrain_sites) +end + +function HVI.get_i_sites_test(prob::DoubleMMCase; scenario::Val{scen}) where {scen} + rng = StableRNG(222) + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + n_site_test = n_site ÷ 10 + i_sites_test = sample(rng, 1:n_site, n_site_test, replace=false) end function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase; - scenario::Val{scen}, n_site_test = 0) where {scen} + scenario::Val{scen}) where {scen} n_covar_pc = 2 n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) pt = get_hybridproblem_par_templates(prob; scenario) - n_siteall = n_site + n_site_test n_covar = get_hybridproblem_n_covar(prob; scenario) (; transP, transM) = get_hybridproblem_transforms(prob; scenario) n_θM = length(pt.θM) FloatType = get_hybridproblem_float_type(prob; scenario) - (; xM, ζMs_true) = gen_cov_pred(rng, FloatType, n_covar_pc, n_covar, n_siteall, + (; xM, ζMs_true) = gen_cov_pred(rng, FloatType, n_covar_pc, n_covar, n_site, inverse(transM)(pt.θM); scenario, rhodec = 8, is_using_dropout = false) - int_θMs_sites = ComponentArrayInterpreter(pt.θM, (n_siteall,)) - transM_sites = StackedArray(transM, n_siteall) + int_θMs_sites = ComponentArrayInterpreter(pt.θM, (n_site,)) + transM_sites = StackedArray(transM, n_site) θMs_true = int_θMs_sites(transM_sites(ζMs_true)) f_batch = get_hybridproblem_PBmodel(prob; scenario) - f = create_nsite_applicator(f_batch, n_siteall) + f = create_nsite_applicator(f_batch, n_site) #xP = fill((; S1 = xP_S1, S2 = xP_S2), n_siteall) - int_xP_sites = ComponentArrayInterpreter(int_xP1, (n_siteall,)) - xP = int_xP_sites(vcat(repeat(xP_S1, 1, n_siteall), repeat(xP_S2, 1, n_siteall))) + int_xP_sites = ComponentArrayInterpreter(int_xP1, (n_site,)) + xP = int_xP_sites(vcat(repeat(xP_S1, 1, n_site), repeat(xP_S2, 1, n_site))) #xP[:S1,:] #θP = get_θP(prob) # for DoubleMMCase par_templates gives correct θP θP = get_hybridproblem_θP(prob; scenario) @@ -405,7 +412,6 @@ function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase; σ_o = FloatType(0.01) #σ_o = FloatType(0.002) logσ2_o = FloatType(2) .* log.(σ_o) - #σ_o = 0.002 y_o = y_true .+ randn(rng, FloatType, size(y_true)) .* σ_o (; xM, @@ -431,6 +437,7 @@ function HVI.get_hybridproblem_cor_ends(prob::DoubleMMCase; scenario::Val{scen}) end end +# TODO: separate parameters and computing tools from specification of the Problem function HVI.get_hybridproblem_ϕq(prob::DoubleMMCase; scenario::Val{scen}) where {scen} approx = get_hybridproblem_HVIApproximation(prob; scenario) FT = get_hybridproblem_float_type(prob; scenario) @@ -438,7 +445,13 @@ function HVI.get_hybridproblem_ϕq(prob::DoubleMMCase; scenario::Val{scen}) wher (;θP, θM) = get_hybridproblem_par_templates(prob; scenario) n_site, _ = get_hybridproblem_n_site_and_batch(prob; scenario) (;transP, transM) = get_hybridproblem_transforms(prob; scenario) - (;ϕqc, approx) = tmp = init_hybrid_ϕunc(approx, cor_ends, zero(FT); θM, transM, n_site) + (;ϕqc, approx) = init_hybrid_ϕunc(approx, cor_ends, zero(FT); θM, transM, n_site) + pt = get_hybridproblem_par_templates(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + #ranef = get_ranef_computer(ranef_spec, keys(pt.θM), n_site, one(eltype(ϕq0))) + ranef = get_ranef_computer(ranef_spec, keys(pt.θM), n_site, one(FT)) + ϕq_ranef = setup_ϕq_ranef(ranef) + ϕqc = CA.ComponentVector(ϕqc, ranef = ϕq_ranef) # for DoubleMMCase templates gives the correct values ϕqP = HVI.update_μP_by_θP(ϕqc, θP, transP) end @@ -453,10 +466,19 @@ function HVI.get_hybridproblem_HVIApproximation(prob::DoubleMMCase; scenario::Va (;θP, θM) = get_hybridproblem_par_templates(prob; scenario) FT = eltype(θM) block_ends = [length(θM)] - MeanScalingHVIApproximation(block_ends,FT(2) .* log.(FT(0.1) .* θM[block_ends])) + MeanScalingHVIApproximation(block_ends, FT(2) .* log.(FT(0.1) .* θM[block_ends])) elseif (:sepvar ∈ scen) MeanVarSepHVIApproximation() else MeanHVIApproximationMat() end -end \ No newline at end of file +end + +function HVI.get_hybridproblem_ranef(prob::DoubleMMCase; scenario::Val{scen}) where {scen} + ranef_spec = if any((:ranef) .∈ Ref(scen)) + par_ranef = (:r1, :K1) + RandomEffects(par_ranef) + else + NullRandomEffects() + end +end diff --git a/src/HybridProblem.jl b/src/HybridProblem.jl index 8d2686c..9b17de2 100644 --- a/src/HybridProblem.jl +++ b/src/HybridProblem.jl @@ -11,8 +11,8 @@ Fields: - `py`: Likelihood function - `transM::Stacked`, `transP::Stacked`: bijectors transforming from unconstrained to constrained scale for site-specific and global parameters respectively. -- `train_dataloader::MLUtils.DataLoader`: providing Tuple of matrices - `(xM, xP, y_o, y_unc, i_sites)`: covariates, model drivers, observations, +- `train_dataloader`: e.g. `MLUtils.DataLoader` providing Tuple of matrices + `(xM, xP, y_o, y_unc, itrain_sites)`: covariates, model drivers, observations, observation uncertainties and index of provided sites. - `test_data::Tuple of the same form as with `train_dataloader` for testset data. - `n_covar::Int`, `n_site::Int`, `n_batch::Int`: number covariates, @@ -38,13 +38,14 @@ struct HybridProblem <: AbstractHybridProblem transM::Stacked transP::Stacked cor_ends::@NamedTuple{P::Vector{Int}, M::Vector{Int}} # = (P=(1,),M=(1,)) - train_dataloader::MLUtils.DataLoader + train_dataloader::Any # MLUtils.DataLoader, but may be different such as WeightedDataLoader test_data::NamedTuple n_site::Int n_batch::Int pbm_covars::NTuple{_N, Symbol} where _N approx::AbstractHVIApproximation penalty_computer::PenaltyComputerOrFunction + ranef::AbstractRandomEffects #penalty_computer:: #inner constructor to constrain the types function HybridProblem( @@ -57,7 +58,7 @@ struct HybridProblem <: AbstractHybridProblem py, transM::Stacked, transP::Stacked, - train_dataloader::MLUtils.DataLoader, + train_dataloader, test_data::NamedTuple, n_site::Int, n_batch::Int; @@ -65,11 +66,12 @@ struct HybridProblem <: AbstractHybridProblem pbm_covars::NTuple{N,Symbol} = (), approx::AbstractHVIApproximation = MeanHVIApproximationMat(), penalty_computer::PenaltyComputerOrFunction = ZeroPenaltyComputer(), + ranef::AbstractRandomEffects = NullRandomEffects(), ) where N new( θM, f_batch, g, ϕg, ϕq, priors, py, transM, transP, cor_ends, train_dataloader, test_data, n_site, n_batch, pbm_covars, - approx, penalty_computer) + approx, penalty_computer, ranef) end end @@ -85,10 +87,13 @@ function init_hybrid_ϕq( transP::Stacked, cor_ends::NamedTuple = (P = [length(θP)], M = [length(θM)]); n_site::Integer, + # need dummy, so that ϕq.ranef is a ComponentVector rather than plain Array + ϕq_ranef::CA.ComponentVector = CA.ComponentVector{eltype(θM)}(dummy=θM[1:0]), kwargs..., ) FT = promote_type(eltype(θP), eltype(θM)) (;ϕqc, approx) = init_hybrid_ϕunc(approx, cor_ends, zero(FT); θM, n_site, kwargs...) + ϕqc = CA.ComponentVector(ϕqc; ranef = ϕq_ranef) ϕqP = update_μP_by_θP(ϕqc, θP, transP) (;ϕqc = ϕqP, approx) end @@ -133,6 +138,7 @@ function update_hybridProblem(prob::AbstractHybridProblem; scenario, ϕunc = nothing, approx::AbstractHVIApproximation = get_hybridproblem_HVIApproximation(prob; scenario), penalty_computer::PenaltyComputerOrFunction = get_hybridproblem_penalty_computer(prob; scenario), + ranef = get_hybridproblem_ranef(prob; scenario), ) n_batch_before = get_hybridproblem_n_site_and_batch(prob; scenario)[2] cor_ends_new = if !isnothing(cor_ends) @@ -149,15 +155,21 @@ function update_hybridProblem(prob::AbstractHybridProblem; scenario, if !isnothing(ϕunc) ϕq = CA.ComponentVector(ϕq; ϕunc...) end + if ranef != get_hybridproblem_ranef(prob; scenario) + # when ranef has been updated, create new default parameter vector + ranefc = get_ranef_computer(ranef, keys(θM), n_site, one(eltype(ϕq))) + ϕq_ranef = setup_ϕq_ranef(ranefc) + ϕq = CA.ComponentVector((;ϕq..., ranef = ϕq_ranef)) + end if n_batch != n_batch_before # if updating n_btach, then need to adjust f_batch and train_dataloader train_dataloader = MLUtils.DataLoader( - train_dataloader.data, batchsize=n_batch, partial=train_dataloader.partial, shuffle=train_dataloader.shuffle) + train_dataloader.data, batchsize=n_batch) f_batch = create_nsite_applicator(f_batch, n_batch) end HybridProblem(θM, ϕq, g, ϕg, f_batch, priors, py, transM, transP, train_dataloader, test_data, n_site, n_batch; cor_ends = cor_ends_new, pbm_covars, - approx, penalty_computer) + approx, penalty_computer, ranef) end function HybridProblem(prob::HybridProblem; kwargs... ) @@ -307,4 +319,8 @@ end function get_hybridproblem_HVIApproximation(prob::HybridProblem; scenario = ()) prob.approx -end \ No newline at end of file +end + +function get_hybridproblem_ranef(prob::HybridProblem; scenario = ()) + prob.ranef +end diff --git a/src/HybridSolver.jl b/src/HybridSolver.jl index 816c992..f03c1d9 100644 --- a/src/HybridSolver.jl +++ b/src/HybridSolver.jl @@ -1,5 +1,16 @@ abstract type AbstractHybridSolver end +""" +A Solver that solves for the mode of hybrid variational inference problem +([`AbstractHybridProblem`](@ref)). + + +It does not perform the full posterior inference, +but only finds the optimal parameters of the variational distribution. +This is useful for initializing the full posterior inference. + +- `alg`: The optimization algorithm to use for solving the problem, e.g. `Adam()` +""" struct HybridPointSolver{A} <: AbstractHybridSolver alg::A end @@ -12,6 +23,8 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve is_inferred::Val{is_infer} = Val(false), ad_backend_loss = AutoZygote(), epochs, + epochs_callback = max(1, floor(Int, epochs/10)), + callback = (state, loss_val) -> false, is_omitting_NaNbatches = false, is_omit_priors::Val{omit_priors} = Val(false), clusters::AbstractVector{<:Integer} = @@ -22,35 +35,44 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve gdevs = isnothing(gdevs) ? get_gdev_MP(scenario) : gdevs pt = get_hybridproblem_par_templates(prob; scenario) g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario) - (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + # TODO: separate parameters from problem description - right now optimized ϕq0 = get_hybridproblem_ϕq(prob; scenario) + (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(ϕq0))) + ϕq_ranef = setup_ϕq_ranef(ranef) + if :ranef ∉ keys(ϕq0) + ϕq0 = CA.ComponentVector(ϕq0, ranef = ϕq_ranef) + else + @assert size(ϕq0[Val(:ranef)]) == size(ϕq_ranef) + end ϕP0 = ϕq0[Val(:μP)] - intϕ = ComponentArrayInterpreter(CA.ComponentVector(ϕg=1:length(ϕg0), ϕP=ϕP0)) + #intϕ = ComponentArrayInterpreter(CA.ComponentVector(ϕg=1:length(ϕg0), ϕq=ϕq0)) #ϕ0_cpu = vcat(ϕg0, pt.θP .* FT(0.9)) # slightly disturb θP_true - ϕ0_cpu = vcat(ϕg0, ϕP0) - n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + ϕ0_cpu = CA.ComponentVector(ϕg=ϕg0, ϕq=ϕq0) + intϕ = ComponentArrayInterpreter(ϕ0_cpu) n_sites_cluster = [count(==(element),clusters) for element in 1:maximum(clusters)] frac_cluster_all = (1 / cluster_rep) ./ n_sites_cluster[clusters] train_loader = get_hybridproblem_train_dataloader(prob; scenario) - #TODO provide different test data - # TODO use 1/10 of the training data - # currently HybridProblem returns only applciators of size n_batch and n_site + test_data = get_hybridproblem_test_data(prob; scenario) # i_test = rand(1:n_site, Integer(floor(n_site/10))) # test_data = map(train_loader.data) do data_comp # ndims(data_comp) == 2 ? data_comp[:, i_test] : data_comp[i_test] # end - test_data = train_loader.data gdev = gdevs.gdev_M if gdev isa MLDataDevices.AbstractGPUDevice ϕ0_dev = gdev(ϕ0_cpu) g_dev = gdev(g) train_loader_dev = gdev_hybridproblem_dataloader(train_loader; gdevs) - test_data_dev = gdev_hybridproblem_data(test_data; gdevs) + test_data_dev = gdev_hybridproblem_data( + test_data[keys(test_data)[1:5]]; gdevs) else ϕ0_dev = ϕ0_cpu g_dev = g train_loader_dev = train_loader - test_data_dev = test_data + test_data_dev = test_data[keys(test_data)[1:5]] end f = get_hybridproblem_PBmodel(prob; scenario) ftest = create_nsite_applicator(f, size(test_data[1],2)) @@ -68,22 +90,24 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve priorsP = Tuple(priors[k] for k in keys(pt.θP)) priorsM = Tuple(priors[k] for k in keys(pt.θM)) #intP = ComponentArrayInterpreter(pt.θP) - intθP = ComponentArrayInterpreter(pt.θP) - intθMs = ComponentArrayInterpreter((n_batch,), pt.θM) + # intθP = ComponentArrayInterpreter(pt.θP) + # intθM = ComponentArrayInterpreter(pt.θM) penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) loss_gf = get_loss_gf(g_dev, transM, transP, f_dev, py, intϕ; n_site_batch=n_batch, par_templates = pt, cdev=infer_cdev(gdevs), pbm_covars, priorsP, priorsM, is_omit_priors, penalty_computer, - intθMs, intθP, frac_cluster_all, + #intθM, intθP, + frac_cluster_all, ranef, ) loss_gf_test = get_loss_gf(g_dev, transM, transP, ftest_dev, py, intϕ; n_site_batch=n_site_test, par_templates = pt, cdev=infer_cdev(gdevs), pbm_covars, priorsP, priorsM, is_omit_priors, penalty_computer, - intθMs, intθP, frac_cluster_all, + #intθM, intθP, + frac_cluster_all, ranef, ) # call loss function once l1 = is_infer ? @@ -91,65 +115,98 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve # using ShareAdd; @usingany Cthulhu # @descend_code_warntype loss_gf(ϕ0_dev, first(train_loader_dev)...) loss_gf(ϕ0_dev, first(train_loader_dev)...; is_testmode=true)[1] + l1t = loss_gf_test(ϕ0_dev, test_data_dev...; is_testmode=true)[1] # and gradient # xMg, xP, y_o, y_unc = first(train_loader_dev) # gr1 = Zygote.gradient( # p -> loss_gf(p, xMg, xP, y_o, y_unc)[1], # ϕ0_dev) # Zygote.gradient(ϕ0_dev -> loss_gf(ϕ0_dev, data1...)[1], ϕ0_dev) - if is_omitting_NaNbatches - # implement training loop by hand to skip minibatches with NaN gradients - ps = CA.getdata(ϕ0_dev) - opt_st_new = Optimisers.setup(solver.alg, ps) - n_skips = 0 - # prepare DI.gradient, need to access and update outside cope data_batch - # because cannot redefine fopt_loss_gf - data_batch = first(train_loader_dev) - is_testmode = false - function fopt_loss_gf(ϕ) - #@show first(data_batch[5], 2) - loss_gf(ϕ, data_batch...; is_testmode)[1] + # if is_omitting_NaNbatches + # # implement training loop by hand to skip minibatches with NaN gradients + # ps = CA.getdata(ϕ0_dev) + # opt_st_new = Optimisers.setup(solver.alg, ps) + # n_skips = 0 + # # prepare DI.gradient, need to access and update outside cope data_batch + # # because cannot redefine fopt_loss_gf + # data_batch = first(train_loader_dev) + # is_testmode = false + # function fopt_loss_gf(ϕ) + # #@show first(data_batch[5], 2) + # loss_gf(ϕ, data_batch...; is_testmode)[1] + # end + # ad_prep = DI.prepare_gradient(fopt_loss_gf, ad_backend_loss, zero(ps)) + # grad = similar(ps) + # stime = time() + # for epoch in 1:epochs + # is_testmode = false + # #i,data_batch = first(enumerate(loader)) + # for (i, data_batch_) in enumerate(train_loader_dev) + # data_batch = data_batch_ # propagate outside for to scope of fopt_loss_gf + # DI.gradient!(fopt_loss_gf, grad, ad_prep, ad_backend_loss, ps) + # if any(isnan.(grad)) + # n_skips += 1 + # #println("Skipped NaN : Batch $i") + # print(",$i") + # else + # Optimisers.update!(opt_st_new, ps, grad) + # end + # end + # ttime = time() - stime + # # compute loss for test data + # l = loss_gf_test(ps, test_data_dev...; is_testmode = true) + # println() + # @show round(ttime, digits=1), epoch, l.nLy, l.neg_log_prior, l.loss_penalty + # # TODO log + # end + # res = nothing + # ϕ = intϕ(ps) + # else + loss_test = let test_data_dev = test_data_dev + (state) -> loss_gf_test(state.u, test_data_dev...; is_testmode=true) end - ad_prep = DI.prepare_gradient(fopt_loss_gf, ad_backend_loss, zero(ps)) - grad = similar(ps) - stime = time() - for epoch in 1:epochs - is_testmode = false - #i,data_batch = first(enumerate(loader)) - for (i, data_batch_) in enumerate(train_loader_dev) - data_batch = data_batch_ # propagate outside for to scope of fopt_loss_gf - DI.gradient!(fopt_loss_gf, grad, ad_prep, ad_backend_loss, ps) - if any(isnan.(grad)) - n_skips += 1 - #println("Skipped NaN : Batch $i") - print(",$i") - else - Optimisers.update!(opt_st_new, ps, grad) - end - end - ttime = time() - stime - # compute loss for test data - l = loss_gf_test(ps, test_data_dev...; is_testmode = true) - println() - @show round(ttime, digits=1), epoch, l.nLy, l.neg_log_prior, l.loss_penalty - # TODO log - end - res = nothing - ϕ = intϕ(ps) - else + callback_epochs = get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...; is_testmode=false)[1], ad_backend_loss) # use CA.getdata(ϕ0_dev), i.e. the plain vector to avoid recompiling for specific CA # loss_gf re-attaches the axes optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader_dev) - res = Optimization.solve(optprob, solver.alg; epochs, kwargs...) + res = Optimization.solve(optprob, solver.alg; + epochs, callback = callback_epochs, kwargs...) ϕ = intϕ(res.u) - end - θP = !isempty(ϕ.ϕP) ? cpu_ca(apply_preserve_axes(transP, cpu_ca(ϕ).ϕP)) : CA.ComponentVector{eltype(ϕ)}() - probo = HybridProblem(prob; ϕg=cpu_ca(ϕ).ϕg, θP) + # end + #θP = !isempty(ϕ.ϕP) ? cpu_ca(apply_preserve_axes(transP, cpu_ca(ϕ).ϕq.μP)) : CA.ComponentVector{eltype(ϕ)}() + # TODO check which components live on gpu and which on cpu + ϕq_opt = ComponentArrayInterpreter(ϕq0)(gdevs.gdev_P(ϕ.ϕq)) + probo = HybridProblem(prob; ϕg=cpu_ca(ϕ).ϕg, ϕq=ϕq_opt) (; ϕ, resopt=res, probo) end +function get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) + n_per_epoch = n_site ÷ n_batch + if epochs_callback == 0 + callback_epochs = callback + else + callback_epochs = function(state, l) + if (state.iter == 1) || (state.iter % (n_per_epoch * epochs_callback) == 0) + l_test = loss_test(state)[1] + println("epoch = $((state.iter) ÷ n_per_epoch), iter = $(state.iter), loss_train=$l, loss_test=$l_test") + end + return callback(state, l) + end + end +end + +""" +A Solver that solves for the full posterior of a variational inference problem +([`AbstractHybridProblem`](@ref)). + +Parameters +- `alg`: The optimization algorithm to use for solving the problem, e.g. `Adam()` +- `n_MC`: The number of Monte Carlo samples to use for estimating the gradients of the loss function. +""" struct HybridPosteriorSolver{A} <: AbstractHybridSolver alg::A n_MC::Int @@ -166,13 +223,16 @@ function HybridPosteriorSolver(solver::HybridPosteriorSolver; end """ - solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolver; ...) + solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolver; epochs, ...) Perform the inversion of HVI Problem. -Optional keyword arguments +Arguments - `prob`: The AbstractHybridProblem to solve. - `scenario`: Scenario to query prob, defaults to `Val(())`. +- `epochs`: number of epochs to train, i.e. number of passes through the whole dataset. + +Optional keyword arguments - `rng`: Random generator, defaults to `Random.default_rng()`. - `gdevs`: `NamedTuple` `(;gdev_M, gdev_P)` functions to move computation and data of ML model on and PBM respectively @@ -180,6 +240,10 @@ Optional keyword arguments defaults to [`get_gdev_MP`](@ref)`(scenario)` - `θmean_quant` default to `0.0`: deprecated - `is_inferred`: set to `Val(true)` to activate type stability checks +- `is_omit_priors`: set to `Val(true)` to omit priors in the loss computation, which can be useful for debugging or if priors are not implemented for a specific scenario (e.g. on gpu) +- `clusters`: vector of cluster assignments for each site, defaults to each site being its own cluster. Clusters are used to compute the loss in a way that accounts for clustering of sites, which can be useful if there are many sites and the number of Monte Carlo samples is limited. +- `cluster_rep`: number of times to repeat each cluster in the loss computation, defaults to 1. Repeating clusters can be useful to effectively increase the number of Monte Carlo samples when the number of clusters is small. +- `epochs_callback`: number of epochs between progress output on evaluating testdata Returns a `NamedTuple` of - `probo`: A copy of the HybridProblem, with updated optimized parameters @@ -193,6 +257,7 @@ Returns a `NamedTuple` of more information on convergence. """ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolver; + epochs, scenario::Val{scen}=Val(()), rng=Random.default_rng(), gdevs = get_gdev_MP(scenario), θmean_quant=0.0, @@ -202,6 +267,8 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS clusters::AbstractVector{<:Integer} = 1:first(get_hybridproblem_n_site_and_batch(prob; scenario)), cluster_rep = 1, + epochs_callback = max(1, floor(Int, epochs/10)), + callback = (state, loss_val) -> false, kwargs... ) where {scen, is_infer, omit_priors} pt = get_hybridproblem_par_templates(prob; scenario) @@ -212,17 +279,26 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) n_sites_cluster = [count(==(element),clusters) for element in 1:maximum(clusters)] frac_cluster_all = (1 / cluster_rep) ./ n_sites_cluster[clusters] - ϕq = get_hybridproblem_ϕq(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(pt.θM))) + ϕq_ranef = setup_ϕq_ranef(ranef) + ϕq = CA.ComponentVector(get_hybridproblem_ϕq(prob; scenario), ranef = ϕq_ranef) (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq) int_ϕq = interpreters.ϕq int_ϕg_ϕq = interpreters.ϕg_ϕq - transMs = StackedArray(transM, n_batch) priors = get_hybridproblem_priors(prob; scenario) priorsP = Tuple(priors[k] for k in keys(pt.θP)) priorsM = Tuple(priors[k] for k in keys(pt.θM)) zero_prior_logdensity = omit_priors ? 0f0 : get_zero_prior_logdensity( priorsP, priorsM, pt.θP, pt.θM) train_loader = get_hybridproblem_train_dataloader(prob; scenario) + test_data = get_hybridproblem_test_data(prob; scenario) + # i_test = rand(1:n_site, Integer(floor(n_site/10))) + # test_data = map(train_loader.data) do data_comp + # ndims(data_comp) == 2 ? data_comp[:, i_test] : data_comp[i_test] + # end + n_batch_test = size(test_data[1],2) if first(train_loader)[1] isa CA.ComponentArray @warn("ML model covariates (1) were provided as ComponentArray. " * "Consider providing them as a plain array.") @@ -235,52 +311,76 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS ϕ0_dev = gdevs.gdev_M(ϕ) g_dev = gdevs.gdev_M(g) # zygote fails if gdev is a CPUDevice, although should be non-op train_loader_dev = gdev_hybridproblem_dataloader(train_loader; gdevs) + test_data_dev = gdev_hybridproblem_data(test_data[keys(test_data)[1:5]]; gdevs) else ϕ0_dev = ϕ g_dev = g train_loader_dev = train_loader + test_data_dev = test_data[keys(test_data)[1:5]] end f = get_hybridproblem_PBmodel(prob; scenario) + f_test = create_nsite_applicator(f, n_batch_test) if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice f_dev = gdevs.gdev_P(f) #fmap(gdevs.gdev_P, f) + f_test_dev = gdevs.gdev_P(f_test) else f_dev = f + f_test_dev = f_test end py = get_hybridproblem_neg_logden_obs(prob; scenario) penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) - # intθP = ComponentArrayInterpreter(pt.θP, (solver.n_MC,)) - # intθMs = ComponentArrayInterpreter((n_batch,), pt.θM, (solver.n_MC,)) - intθP = ComponentArrayInterpreter(pt.θP) - intθMs = ComponentArrayInterpreter((n_batch,), pt.θM) - - priors_θP_mean, priors_θMs_mean = construct_priors_θ_mean( - prob, ϕ0_dev.ϕg, keys(pt.θM), pt.θP, θmean_quant, g_dev, transM, transP; - scenario, gdevs, pbm_covars) + # intθP = ComponentArrayInterpreter(pt.θP) + # intθMs = ComponentArrayInterpreter((n_batch,), pt.θM) loss_elbo = get_loss_elbo( - g_dev, transP, transMs, f_dev, py; - solver.n_MC, solver.n_MC_cap, cor_ends, priors_θP_mean, priors_θMs_mean, - cdev=infer_cdev(gdevs), pbm_covars, pt.θP, int_ϕq, int_ϕg_ϕq, priorsP, priorsM, - is_omit_priors, zero_prior_logdensity, approx, penalty_computer, intθMs, intθP, + g_dev, transP, transM, f_dev, py, n_batch; + n_MC = solver.n_MC, n_MC_cap = solver.n_MC_cap, cor_ends, + cdev=infer_cdev(gdevs), pbm_covars, + par_templates = pt, + #pt.θP, + int_ϕq, int_ϕg_ϕq, priorsP, priorsM, + is_omit_priors, zero_prior_logdensity, approx, penalty_computer, + ranef, + #intθMs, intθP, + frac_cluster_all, + ) + loss_elbo_test = get_loss_elbo( + g_dev, transP, transM, f_test_dev, py, n_batch_test; + solver.n_MC, solver.n_MC_cap, cor_ends, + cdev=infer_cdev(gdevs), pbm_covars, + par_templates = pt, + #pt.θP, + int_ϕq, int_ϕg_ϕq, priorsP, priorsM, + is_omit_priors, zero_prior_logdensity, approx, penalty_computer, + ranef, + #intθMs, intθP, frac_cluster_all, ) + # test loss function once # tmp = first(train_loader_dev) # using ShareAdd # @usingany Cthulhu # @descend_code_warntype loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...) # omit for type stability in AD + @assert length(first(train_loader_dev)) == 5 l0 = #is_infer ? # (Test.@inferred loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...; is_testmode=true)) : loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...; is_testmode=false) + l0t = loss_elbo_test(ϕ0_dev, rng, test_data_dev...; is_testmode=true) + loss_test = let test_data_dev = test_data_dev + (state) -> loss_elbo_test(state.u, rng, test_data_dev...; is_testmode=true) + end + callback_epochs = get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) optf = Optimization.OptimizationFunction( (ϕ, data) -> first(loss_elbo(ϕ, rng, data...; is_testmode=false)), Optimization.AutoZygote()) optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader_dev) - res = Optimization.solve(optprob, solver.alg; kwargs...) + res = Optimization.solve(optprob, solver.alg; callback = callback_epochs, epochs, kwargs...) ϕc = interpreters.ϕg_ϕq(cpu_device()(res.u)) ϕq = ϕc[Val(:ϕq)]; ϕg = ϕc[Val(:ϕg)]; @@ -313,39 +413,49 @@ The loss function takes in addition to ϕ, data that changes with minibatch - `xP`: drivers for the processmodel: Iterator of size n_site - `y_o`, `y_unc`: matrix of observations and uncertainties, sites in columns """ -function get_loss_elbo(g, transP, transMs, f, py; +function get_loss_elbo(g, transP, transM, f, py, n_batch; n_MC, n_MC_mean = max(n_MC,20), n_MC_cap=n_MC, - cor_ends, priors_θP_mean, priors_θMs_mean, cdev, pbm_covars, θP::AbstractVector{T}, + cor_ends, cdev, pbm_covars, + par_templates, + #θP::AbstractVector{T}, int_ϕq, int_ϕg_ϕq, priorsP, priorsM, penalty_computer = ZeroPenaltyComputer(), is_omit_priors, zero_prior_logdensity, approx, - intθMs, intθP, + ranef::AbstractRandomEffectsComputer, + #intθMs, intθP, frac_cluster_all, -) where T +) + T = eltype(par_templates.θP) + intθP = ComponentArrayInterpreter(par_templates.θP) + intθMs = ComponentArrayInterpreter((n_batch,), par_templates.θM) + transMs = StackedArray(transM, n_batch) + let g = g, transP = transP, transMs = transMs, f = f, py = py, n_MC = n_MC, n_MC_cap = n_MC_cap, n_MC_mean = n_MC_mean, cor_ends = cor_ends, int_ϕq = get_concrete(int_ϕq), int_ϕg_ϕq = get_concrete(int_ϕg_ϕq), - priors_θP_mean = priors_θP_mean, priors_θMs_mean = priors_θMs_mean, cdev = cdev, - pbm_covar_indices = get_pbm_covar_indices(θP, pbm_covars), + cdev = cdev, + pbm_covar_indices = get_pbm_covar_indices(par_templates.θP, pbm_covars), trans_mP=StackedArray(transP, n_MC_mean), trans_mMs=StackedArray(transMs.stacked, n_MC_mean), priorsP=priorsP, priorsM=priorsM, penalty_computer=penalty_computer, is_omit_priors = is_omit_priors, zero_prior_logdensity = zero_prior_logdensity, approx = approx, intθMs = get_concrete(intθMs), intθP = get_concrete(intθP), + ranef = ranef frac_cluster_all = convert.(T, frac_cluster_all) - function loss_elbo(ϕ, rng, xM, xP, y_o, y_unc, i_sites; is_testmode) + + function loss_elbo(ϕ, rng::Random.AbstractRNG, xM, xP, y_o, y_unc, itrain_sites; is_testmode) #ϕc = int_ϕg_ϕq(ϕ) neg_elbo_gtf( - rng, ϕ, g, f, py, xM, xP, y_o, y_unc, i_sites; + rng, ϕ, g, f, py, xM, xP, y_o, y_unc, itrain_sites; int_ϕq, int_ϕg_ϕq, - n_MC, n_MC_cap, n_MC_mean, cor_ends, priors_θP_mean, priors_θMs_mean, + n_MC, n_MC_cap, n_MC_mean, cor_ends, cdev, pbm_covar_indices, transP, transMs, trans_mP, trans_mMs, priorsP, priorsM, penalty_computer, #ϕg = ϕc.ϕg, ϕq = ϕc.ϕq, is_testmode, is_omit_priors, zero_prior_logdensity, approx, - intθMs, intθP, frac_cluster_all, + intθMs, intθP, ranef, frac_cluster_all, ) end end @@ -390,12 +500,12 @@ function compute_elbo_components( g_dev = g data_dev = data end - (xM, xP, y_o, y_unc, i_sites) = data_dev + (xM, xP, y_o, y_unc, itrain_sites) = data_dev n_site_pred = size(xP,2) @assert size(xM, 2) == n_site_pred @assert size(y_o, 2) == n_site_pred @assert size(y_unc, 2) == n_site_pred - @assert length(i_sites) == n_site_pred + @assert length(itrain_sites) == n_site_pred f_batch = get_hybridproblem_PBmodel(prob; scenario) f = (n_site_pred == n_batch) ? f : create_nsite_applicator(f_batch, n_site_pred) py = get_hybridproblem_neg_logden_obs(prob; scenario) @@ -404,53 +514,7 @@ function compute_elbo_components( # prob, ϕ0_dev.ϕg, keys(θM), θP, θmean_quant, g_dev, transM; # scenario, gdev, cdev, pbm_covars) neg_elbo_gtf_components( - rng, ϕ0_dev, g_dev, transPMs_batch, f, py, xM, xP, y_o, y_unc, i_sites, interpreters; + rng, ϕ0_dev, g_dev, transPMs_batch, f, py, xM, xP, y_o, y_unc, itrain_sites, interpreters; solver.n_MC, solver.n_MC_cap, cor_ends, priors_θ_mean) end -""" -In order to let mean of θ stay close to initial point parameter estimates -construct a prior on mean θ to a Normal around initial prediction. -""" -function construct_priors_θ_mean(prob, ϕg, keysθM, θP, θmean_quant, g_dev, transM, transP; - scenario::Val{scen}, gdevs, pbm_covars, - ) where {scen} - iszero(θmean_quant) ? ([],[]) : - begin - n_θM = length(keysθM) - gdev=gdevs.gdev_M - #cdev=infer_cdev(gdevs) - n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) - # all_loader = MLUtils.DataLoader( - # get_hybridproblem_train_dataloader(prob; scenario).data, batchsize = n_site) - # xM_all = first(all_loader)[1] - is_gpu = :use_gpu ∈ scen - xM_all_cpu = get_hybridproblem_train_dataloader(prob; scenario).data[1] - xM_all = is_gpu ? gdev(xM_all_cpu) : xM_all_cpu - ζP = apply_preserve_axes(inverse(transP), θP) - pbm_covar_indices = get_pbm_covar_indices(θP, pbm_covars) - xMP_all = _append_each_covars(xM_all, CA.getdata(ζP), pbm_covar_indices) - transMs = StackedArray(transM, n_site) - # ζMs = g_dev(xMP_all, CA.getdata(ϕg))' # transpose to par-last for StackedArray - # ζMs_cpu = cdev(ζMs) - # θMs_tr = transMs(ζMs_cpu) - θMs_tr = gtrans( - g_dev, transMs, xMP_all, CA.getdata(ϕg), n_θM; cdev=cpu_device(), is_testmode = true) - priors_dict = get_hybridproblem_priors(prob; scenario) - priorsP = [priors_dict[k] for k in keys(θP)] - priors_θP_mean = map(priorsP, θP) do priorsP, θPi - fit_narrow_normal(θPi, priorsP, θmean_quant) - end - priorsM = Tuple(priors_dict[k] for k in keysθM) - i_par = 1 - i_site = 1 - priors_θMs_mean = map(Iterators.product(axes(θMs_tr)...)) do (i_site, i_par) - #@show i_par, i_site - fit_narrow_normal(θMs_tr[i_site, i_par], priorsM[i_par], θmean_quant) - end - # # concatenate to a flat vector - # int_n_site = get_ca_int_PMs(n_site) - # int_n_site(vcat(priors_θP_mean, vec(priors_θMs_mean))) - priors_θP_mean, priors_θMs_mean - end -end diff --git a/src/HybridVariationalInference.jl b/src/HybridVariationalInference.jl index 44cb719..2e699c2 100644 --- a/src/HybridVariationalInference.jl +++ b/src/HybridVariationalInference.jl @@ -28,16 +28,22 @@ using KernelAbstractions import NaNMath # ignore missing observations in logDensity using DifferentiationInterface: DifferentiationInterface as DI import Zygote +import ForwardDiff, PreallocationTools using IterTools: IterTools using PDMats using Distances, Clustering #using OptimizationOptimisers +using RecursiveArrayTools: VectorOfArray, AbstractVectorOfArray export DoubleMM include("OneBasedVectorWithZero.jl") +export cat_namedtuple_lastdim, index_at_dim include("util.jl") +export WeightedDataLoader +include("WeightedDataLoader.jl") + export extend_stacked_nrow, StackedArray #public Exp #julia 1.10 public: https://github.com/JuliaLang/julia/pull/55097 @@ -45,6 +51,11 @@ VERSION >= v"1.11.0-DEV.469" && eval(Meta.parse("public Exp")) VERSION >= v"1.11.0-DEV.469" && eval(Meta.parse("public Logistic")) include("bijectors_utils.jl") +export AbstractRandomEffects, RandomEffects, NullRandomEffects, NullRandomEffectsComputer +export AbstractCovariancePrior, CVPrior_LKJ_Cauchy +export compute_nLranef, add_ranef, setup_ϕq_ranef, get_ranef_computer +include("RandomEffects.jl") + export AbstractHVIApproximation, AbstractMeanHVIApproximation export get_numberof_MLinputs export MeanHVIApproximation, MeanHVIApproximationMat @@ -74,19 +85,20 @@ include("PBMApplicator.jl") export AbstractHybridProblem, AbstractPenaltyComputer, CustomPenaltyComputer, compute_penalty, - get_hybridproblem_MLapplicator, get_hybridproblem_PBmodel, - get_hybridproblem_penalty_computer, - get_hybridproblem_ϕq, get_hybridproblem_θP, - get_hybridproblem_float_type, gen_hybridproblem_synthetic, - get_hybridproblem_par_templates, get_hybridproblem_transforms, - get_hybridproblem_train_dataloader, - get_hybridproblem_test_data, - get_hybridproblem_neg_logden_obs, - get_hybridproblem_n_covar, # default - get_hybridproblem_n_site_and_batch, - get_hybridproblem_cor_ends, - get_hybridproblem_priors, - get_hybridproblem_pbmpar_covars, + get_hybridproblem_MLapplicator, get_hybridproblem_PBmodel, + get_hybridproblem_penalty_computer, + get_hybridproblem_ϕq, get_hybridproblem_θP, + get_hybridproblem_float_type, gen_hybridproblem_synthetic, + get_hybridproblem_par_templates, get_hybridproblem_transforms, + get_hybridproblem_train_dataloader, + get_hybridproblem_test_data, + get_hybridproblem_neg_logden_obs, + get_hybridproblem_n_covar, # default + get_hybridproblem_n_site_and_batch, + get_hybridproblem_cor_ends, + get_hybridproblem_priors, + get_hybridproblem_pbmpar_covars, + get_hybridproblem_ranef, gen_cov_pred, construct_dataloader_from_synthetic, gdev_hybridproblem_dataloader, gdev_hybridproblem_data, @@ -123,6 +135,7 @@ export get_ca_starts, get_ca_ends, get_cor_count include("cholesky.jl") export neg_elbo_gtf, sample_posterior, predict_hvi, ZeroPenaltyComputer +export predict_hvi_and_compute_elbo_components export get_hybridproblem_correlation_Ms, get_hybridproblem_cholesky_correlation_Ms export get_marginal_std include("elbo_dev.jl") @@ -131,6 +144,9 @@ include("elbo_scaling.jl") include("elbo.jl") include("elbo2.jl") +#export get_loss_ran_tr_f +include("fit_ranef_site.jl") + export init_hybrid_params, init_hybrid_ϕunc include("init_hybrid_params.jl") diff --git a/src/OneBasedVectorWithZero.jl b/src/OneBasedVectorWithZero.jl index 04cff25..320a0df 100644 --- a/src/OneBasedVectorWithZero.jl +++ b/src/OneBasedVectorWithZero.jl @@ -3,7 +3,9 @@ A thin wrapper over an `AbstractVector` that exposes a linear 1-based indexing API mapping `v[i]` to `data[axes(data, 1)[i]]` on the underlying storage -and provides a value at index 0 (defaulting to zero) that is not stored in the underlying vector. +and provides a value at index 0 (defaulting to zero) that is not stored in +the underlying vector. +The zero is not counted in the length of the vector, so `length(v)` is the same as `length(data)`. Example usage: ```jldoctest; output=false diff --git a/src/RandomEffects.jl b/src/RandomEffects.jl new file mode 100644 index 0000000..084f259 --- /dev/null +++ b/src/RandomEffects.jl @@ -0,0 +1,236 @@ +abstract type AbstractRandomEffects end +abstract type AbstractRandomEffectsComputer{T} end + +abstract type AbstractCovariancePrior{N,T} end + +struct NullRandomEffectsComputer{T} <: AbstractRandomEffectsComputer{T} + n_par::Int + n_site::Int +end +compute_nLranef(re::NullRandomEffectsComputer{T}, rm) where T = zero(T) +add_ranef(re::NullRandomEffectsComputer, μ, ϕq_ranef, itrain_sites) = μ +function setup_ϕq_ranef(re::NullRandomEffectsComputer{T}) where T + res = CA.ComponentVector(β=reshape( T[], 0, re.n_site)) +end +function sample_ranef(re::NullRandomEffectsComputer{T}, ϕq_ranef, n_site, n_sample) where T + res = fill( zero(T), re.n_par, n_site, n_sample) +end +function get_choleskyΣ_ranef(re::NullRandomEffectsComputer{T}, ϕq_ranef) where T + zeros(T, 0, 0) # empty subset of parameters that have a randon effect +end +function get_choleskyΣ_par(re::NullRandomEffectsComputer{T}, ϕq_ranef) where T + zeros(T, re.n_par, re.n_par) +end + +struct NullRandomEffects <: AbstractRandomEffects; end +function get_ranef_computer( + rn::NullRandomEffects, + θM_keys::NTuple{NM,Symbol}, n_site::Integer, float_template::T=1.0, + ) where {NM,T} + NullRandomEffectsComputer{T}(length(θM_keys), Int(n_site)) +end + + +""" + RandomEffects(parameters; γ=0.5), η=5.0, γv = SA.SVector(γ,γ,...)) + +Each site-predicted parameter-sub-vector can have a site-specific random effect +θ_i = θ_i_ml + β_i +The random effects are estimated with the global parameters. +They are assumed to be drawn from Centered Normal distribution +with covariance Σ, whose coefficients of the upper cholesky-factor are also +estimated as global parameters. +In order to encourage low variance and low correlation, a prior of +zero-centered Cauchy distribution (default scale 0.5) is applied to main diagonal, +and a LKJ prior (default scale 5) is applied to the correlation matrix. +Scale of the Cauchy distribution defines the MAD. +""" +struct RandomEffects{N,T} <: AbstractRandomEffects + parameters::NTuple{N, Symbol} + prior_Σ::AbstractCovariancePrior{N,T} +end +function RandomEffects(parameters, float_template::T=1.0; + γ::T=T(0.5), η::T=T(5.0), γv = SA.SVector(ntuple(_ -> γ, length(parameters)))) where T + N = length(parameters) + prior_Σ = CVPrior_LKJ_Cauchy(γv, η) + RandomEffects{N,T}(parameters, prior_Σ) +end + +""" + +""" +struct RandomEffectsComputer{N,T,NM,NNM} <: AbstractRandomEffectsComputer{T} + parameters::NTuple{N, Symbol} + prior_Σ::AbstractCovariancePrior{N,T} # changed type to float_template + ncomp_U::Int8 # number of components to describe correlation + P_col::SA.SMatrix{NM, N, Bool, NNM} # projection matrix of subset of random to all + n_site::Int # number of sites to setup parameter vector +end +function get_ranef_computer( + rn::RandomEffects{N}, θM_keys::NTuple{NM,Symbol}, n_site::Integer, float_template::T=1.0 + ) where {N,NM,T} + prior_Σ = convert_prior(rn.prior_Σ, float_template) + # γ = convert.(T, rn.γ) + # η = convert(T,rn.η) + # γ = convert.(T, rn.γ) + # η = convert(T,rn.η) + # prior_Σ=CVPrior_LKJ_Cauchy(γ, η) + ncomp_U=Int8(sumn(N)) + pos = SA.SVector{N,Int8}(Tuple(findfirst(==(s), θM_keys) for s in rn.parameters)) + P_col = SA.SMatrix{NM,N}(j == k for j in 1:NM, k in pos) + RandomEffectsComputer(rn.parameters, prior_Σ, ncomp_U, P_col, convert(Int, n_site)) +end + + +function compute_nLranef(re::RandomEffectsComputer{N,T}, ϕq_ranef) where {N,T} + # get cholesky factor of covariance matrix from optimized parameters + U = get_choleskyΣ_ranef(re, ϕq_ranef) + # compute the logdensity of the random effects given the covariance + β = ϕq_ranef[Val(:β)] + # βi = first(eachrow(β)) + # TODO compute for entire array (solve required only once) + logden_rm = sum(eachrow(β)) do βi + #log_density_mvn_cholesky(UpperTriangular(U), βi) does not work with Zygote + log_density_mvn_cholesky(U, βi) + end + # compute the prior of the estimated covariance matrix + logden_Σ = logpdf(re.prior_Σ, U) + cost = -(logden_rm + logden_Σ) + #@show ϕq_ranef.σ, cost + cost +end + +""" + get_choleskyΣ_ranef(re::AbstractRandomEffectsComputer, ϕq_ranef::ComponentVector) + +Return the upper cholesky factor of Covariance of the Random effects. +""" +function get_choleskyΣ_ranef(re::RandomEffectsComputer{N,T}, ϕq_ranef) where {N,T} + coef_Ucorr = ϕq_ranef[Val(:coef_U)] # parameterization of cholesky of correlation + σ = max.(T(1e-10), ϕq_ranef[Val(:σ)]) # main diagonal of cholesky factor of covariance + U = transformU_cholesky1(coef_Ucorr; n=N) * diagm(σ) +end +function get_choleskyΣ_par(re::RandomEffectsComputer{N,T}, ϕq_ranef) where {N,T} + U = get_choleskyΣ_ranef(re, ϕq_ranef) + U_full = re.P_col * U * re.P_col' + U_full +end + +""" +Assume μ of size n_par x n_site +""" +function add_ranef(re::RandomEffectsComputer, μ, ϕq_ranef, itrain_sites) + ranef = CA.getdata(ϕq_ranef[Val(:β)][itrain_sites,:]) + # moved construction of projection matrix to RandomEffectsComputer + P_col = re.P_col + #ranef_full = ranef * P_col' + ranef_full = P_col * ranef' # par-first format + μadd = μ .+ ranef_full + μadd +end + +""" +Provide an initial estimate of the correct structure for the parameters to +estimate for given random effects computer. +""" +function setup_ϕq_ranef(re::RandomEffectsComputer{N,T}) where {N,T} + coef_Ucorr = uutri2vec(cholesky(I(N).*one(T)).U) + # U = vec2utri(coef_U; n = N); U' * U + CA.ComponentVector( + β = zeros(T, re.n_site, N), + coef_U = coef_Ucorr, + σ = fill(T(0.01), N) + ) +end + +""" +provide a sample n_par x n_site x n_sample, where +n_par is the dimension of parameters including those columns that have no +random effect. +""" +function sample_ranef(re::RandomEffectsComputer{N,T}, ϕq_ranef, n_site, n_sample) where {N,T} + U = get_choleskyΣ_ranef(re, ϕq_ranef) + rn = randn(T, N, n_site * n_sample) + r0 = similar(rn) + mul!(r0, U', rn) + out = similar(rn, size(re.P_col, 1), n_site * n_sample) + mul!(out, re.P_col, r0) + res = reshape(out, size(re.P_col, 1), n_site, n_sample) +end +function sample_ranef_nonmut(re::RandomEffectsComputer{N,T}, ϕq_ranef, n_site, n_sample) where {N,T} + U = get_choleskyΣ_ranef(re, ϕq_ranef) + rn = randn(N, n_site * n_sample ) + r0 = U' * rn + r = reshape(r0, N, n_site, n_sample) + #r1 = first(eachslice(r, dims=3)) + # resv = map(eachslice(r, dims=3)) do r1 + # r1' * re.P_col' + # end + # res = stack(resv) + resr = reshape(re.P_col * r0, size(re.P_col, 1), n_site, n_sample) + #res2 = permutedims(resr, (2,1,3)) + res, res2 +end + + + + + +""" +Assigns a Cauchy-prior_Σ with scale γ to the main diagonal of the covariance +and an LKJ prior_Σ with scale η on the correlation matrix. + +The default parameterization follows the +[STAN recommendation](https://mc-stan.org/docs/2_19/stan-users-guide/multivariate-hierarchical-priors-section.html) +assigning a Cauchy scale of 0.5 and an LKJ scale of 5.0 +This encourages to decrease variations correlations. +""" +struct CVPrior_LKJ_Cauchy{N,T} <: AbstractCovariancePrior{N,T} + dCauchy::SA.SVector{N, Cauchy{T}} + dLKJ::LKJ{T,Int8} +end + +function CVPrior_LKJ_Cauchy(γ::SA.SVector{N,T}, η::T) where {N,T} + dCauchy = Cauchy.(zero(T), γ)::SA.SVector{N, Cauchy{T}} + dLKJ = LKJ(Int8(N), η)::LKJ{T,Int8} + CVPrior_LKJ_Cauchy{N, T}(dCauchy, dLKJ) +end + +function CVPrior_LKJ_Cauchy(n::Integer, float_template::T=1.0; + γ::T=T(0.5), η::T=T(5.0), γv = SA.SVector(ntuple(_ -> γ, n))) where T + CVPrior_LKJ_Cauchy(γv, η, ) +end + +convert_prior(prior::CVPrior_LKJ_Cauchy{N,T}, float_template::T) where {N,T} = prior +function convert_prior(prior::CVPrior_LKJ_Cauchy{N,T}, float_template::TN) where {N,T,TN} + #γ = SA.SVector(Tuple(T(params(d)[2]) for d in prior.dCauchy)) + γ = map(d -> TN(params(d)[2]), prior.dCauchy) # already returns SVector + #η = convert.(TN,params(prior.dLKJ)) + η = TN(params(prior.dLKJ)[2]) + + CVPrior_LKJ_Cauchy(γ, η) +end + +# function rand_prior(d::CVPrior_LKJ_Cauchy{N,T}, n_sample) +# end + + +function Distributions.logpdf(d::CVPrior_LKJ_Cauchy{N,T}, cm) where {N,T} + cm_s = cm + I*T(1e-8) # to deal with initial zero covariance + τ = sqrt.(diag(cm_s)) + logpdf_cauchy = sum(logpdf_.(d.dCauchy, τ)) + corrm = cm_s ./ (τ * τ') + logpdf_lkj = logpdf_(d.dLKJ, corrm) + return logpdf_cauchy + logpdf_lkj +end + +logpdf_(d::Distribution, x) = logpdf(d,x) +function ChainRulesCore.rrule(::typeof(logpdf_), d::Distribution, x) + # avoid propagating to Distribution d but only to x, otherwise Δy is passed through + function logpdf_pullback(Δy) + (NoTangent(), NoTangent(), Δy) + end + return logpdf_(d,x), logpdf_pullback +end + + diff --git a/src/WeightedDataLoader.jl b/src/WeightedDataLoader.jl new file mode 100644 index 0000000..58f8929 --- /dev/null +++ b/src/WeightedDataLoader.jl @@ -0,0 +1,78 @@ +# Prompt: +# Implement a WeightedDataLoader in Julia that wraps MLUtils.DataLoader and performs weighted resampling. The implementation should satisfy the following requirements: +# Define a WeightedObsView{D} struct with fields data::D and weights::Vector{Float64}. +# Implement MLUtils.numobs and MLUtils.getobs for WeightedObsView. In getobs, use StatsBase.sample with Weights to draw length(indices) weighted indices with replacement, then call getobs on the underlying data. Shield the sampling and data fetching from Zygote using ChainRulesCore.@ignore_derivatives. Return data as-is without any type conversion. +# Define a WeightedDataLoader constructor that takes data, weights, and keyword arguments batchsize=32 and any extra kwargs forwarded to MLUtils.DataLoader. It should construct a WeightedObsView and wrap it in a MLUtils.DataLoader, returning a native MLUtils.DataLoader so it is directly recognized by Optimization.jl. +# Write a @testset using SimpleChains, Optimisers, Optimization, OptimizationOptimisers, ChainRulesCore, MLUtils, and StatsBase that tests: +# The loader is a native MLUtils.DataLoader +# Correct batch shapes and types are preserved for Float32, Float64, Int arrays, NTuple, and NamedTuple datasets +# Minority class oversampling works correctly +# End-to-end training with Optimization.solve and OptimizationOptimisers.Adam reduces the loss, where the loss function receives a batch directly as p (not the loader) +# Add methods getindex and length WeightedObsView to work on the data field + +""" + WeightedObsView{D} + +A view over a dataset that supports weighted resampling via `MLUtils.getobs`. +""" +struct WeightedObsView{D} + data::D + weights::Vector{Float64} +end + +""" + MLUtils.numobs(wov::WeightedObsView) + +Return the number of observations in the underlying dataset. +""" +MLUtils.numobs(wov::WeightedObsView) = MLUtils.numobs(wov.data) + +""" + MLUtils.getobs(wov::WeightedObsView, indices) + +Draw `length(indices)` weighted samples (with replacement) using `StatsBase.sample` +with `Weights`, then fetch those observations from the underlying dataset. +Both sampling and data fetching are shielded from Zygote via +`ChainRulesCore.@ignore_derivatives`. +""" +function MLUtils.getobs(wov::WeightedObsView, indices) + n = length(indices) + obs = ChainRulesCore.@ignore_derivatives begin + w = Weights(wov.weights) + sampled_indices = StatsBase.sample(1:MLUtils.numobs(wov.data), w, n; replace=true) + MLUtils.getobs(wov.data, sampled_indices) + end + return obs +end + +# ============================================================ +# WeightedDataLoader constructor +# ============================================================ + +""" + WeightedDataLoader(data, weights; batchsize=32, kwargs...) + +Construct a weighted data loader that performs weighted resampling. + +Wraps `data` in a `WeightedObsView` and then in a `MLUtils.DataLoader`, +returning a native `MLUtils.DataLoader` so it is directly recognised by +Optimization.jl. + +# Arguments +- `data`: Any dataset compatible with `MLUtils.numobs` / `MLUtils.getobs`. +- `weights`: A `Vector{Float64}` of per-observation sampling weights. +- `batchsize`: Number of observations per batch (default: 32). +- `kwargs...`: Additional keyword arguments forwarded to `MLUtils.DataLoader`. +""" +function WeightedDataLoader(data, weights::AbstractVector; batchsize::Int=32, kwargs...) + wov = WeightedObsView(data, Vector{Float64}(weights)) + return MLUtils.DataLoader(wov; batchsize=batchsize, kwargs...) +end + + +# ============================================================ +# Base interface +# ============================================================ + +Base.length(wov::WeightedObsView) = length(wov.data) +Base.getindex(wov::WeightedObsView, i) = wov.data[i] diff --git a/src/bijectors_utils.jl b/src/bijectors_utils.jl index 029edc2..4061150 100644 --- a/src/bijectors_utils.jl +++ b/src/bijectors_utils.jl @@ -169,9 +169,10 @@ Return a Jacobian for each row in x. function with_logabsdet_jacobians(sb::StackedArray, x::AbstractArray) (y, logjacs_vec) = with_logabsdet_jacobians(sb.stacked, vec(x)) ym = reshape(y, size(x)) - # move sum to elbo, here return all components + # move (weighted) sum to elbo, here return all components #logjacs = sum(reshape(logjacs_vec, sb.nrow, :); dims = 2)[:,1] - logjacs = reshape(logjacs_vec, sb.nrow, :) + logjacs = reshape(logjacs_vec, size(ym)) + #Main.@infiltrate_main return (ym, logjacs) end diff --git a/src/cholesky.jl b/src/cholesky.jl index 47b4ce6..03de5c2 100644 --- a/src/cholesky.jl +++ b/src/cholesky.jl @@ -144,7 +144,7 @@ and transforms it to an UpperTriangular that satisfies diag(U' * U) = 1. This can be used to fit parameters that yield an upper Cholesky-Factor -of a Covariance matrix. +of a correlation matrix. It uses the upper triangular matrix rather than the lower because it involves a sum across columns, whereas the alternative of a lower triangular diff --git a/src/clustering.jl b/src/clustering.jl index 179cd8e..e7c3662 100644 --- a/src/clustering.jl +++ b/src/clustering.jl @@ -53,8 +53,8 @@ function refit_clusters(rng, probo, solver, xM ; while (n_aggsplits_i > 0) && (length(cnts_clusters_totest) > 0) #global cnts, clusters, i_splits i_cluster = argmax(cnts_clusters_totest) - i_sites = findall(isequal.(clusters, i_cluster)) - σM = vec(median(σMs[:,i_sites]; dims=2)) + itrain_sites = findall(isequal.(clusters, i_cluster)) + σM = vec(median(σMs[:,itrain_sites]; dims=2)) # (; is_overdispersed, clusters, clusters_sub) = split_cluster( clusters, i_cluster, X, Ucor, σMs; n_cluster_sub) @@ -93,16 +93,22 @@ function extract_MLpred(probo, xM; scenario = Val(())) (; X, σMs) end -function cluster_records(X_matrix::AbstractMatrix, Ucor::AbstractMatrix{T}, σM::AbstractVector; n_cluster_sub=4, cluster_ids = 1:n_cluster_sub, ) where T +function cluster_records(X_matrix::AbstractMatrix, Ucor::AbstractMatrix, σM::AbstractVector; kwargs... ) + UΣ = Ucor * diagm(σM) + cluster_records(X_matrix::AbstractMatrix, UΣ; kwargs...) +end + +function cluster_records(X_matrix::AbstractMatrix, UΣ::AbstractMatrix{T}; n_cluster_sub=4, cluster_ids = 1:n_cluster_sub, ) where T # x_i -x_j are distributed N(0, 2Σ) #Σ = 2 * HybridVariationalInference.compute_cov(Ucor, σM) @assert n_cluster_sub == length(cluster_ids) "Length of cluster_ids must match n_cluster" - invΣ = HybridVariationalInference.compute_invcov(Ucor, σM) / T(2) + #invΣ = HybridVariationalInference.compute_invcov(Ucor, σM) / T(2) + invΣ = inv(UΣ) * inv(UΣ') # Precompute differences: (X_i - X_j) for all i, j # Use broadcasting to compute all pairwise differences #diffs = X_matrix' .- X_matrix # Shape: (n_vars, n_rows, n_rows) n_rows, n_vars = size(X_matrix) # e.g., 2000×10 - n_rows < n_cluster_sub && error("Cannot cluster $nrows records into $n_cluster_sub clusters.") + n_rows < n_cluster_sub && error("Cannot cluster $n_rows records into $n_cluster_sub clusters.") if n_rows == n_cluster_sub # assign each record to its own cluster clusters0 = 1:n_rows @@ -221,20 +227,20 @@ function check_overdispersion( end function split_cluster(clusters, i_cluster, X, Ucor, σMs; n_cluster_sub=4) - i_sites = findall(isequal.(clusters, i_cluster)) - n_sites_cl = length(i_sites) + itrain_sites = findall(isequal.(clusters, i_cluster)) + n_sites_cl = length(itrain_sites) if n_sites_cl == 1 return (; is_overdispersed = false, clusters, clusters_sub = eltype(clusters)[]) end - X_cluster = X[i_sites,:] - σM = vec(median(σMs[:,i_sites]; dims=2)) + X_cluster = X[itrain_sites,:] + σM = vec(median(σMs[:,itrain_sites]; dims=2)) # vec(std(X_cluster; dims = 1)) is_overdispersed = check_overdispersion(X_cluster, Ucor, σM)[1] if is_overdispersed n_cluster_sub_i = min(n_cluster_sub, n_sites_cl) cluster_ids = vcat(i_cluster, maximum(clusters) .+ (1:(n_cluster_sub_i-1))) clusters_sub = cluster_records(X_cluster, Ucor, σM; cluster_ids, n_cluster_sub = n_cluster_sub_i); - clusters[i_sites] = clusters_sub + clusters[itrain_sites] = clusters_sub else clusters_sub = eltype(clusters)[] end @@ -245,7 +251,7 @@ function split_cluster(clusters, i_cluster, X, Ucor, σMs; n_cluster_sub=4) # Main.@infiltrate_main # clusters_sub = vcat(i_cluster, maximum(clusters) .+ (1:(n_sites_cl-1))) - # clusters[i_sites] = cluster_sub + # clusters[itrain_sites] = cluster_sub # (; is_overdispersed = false, clusters, clusters_sub) # end end @@ -253,7 +259,7 @@ end function map_by_cluster(f, clusters) cluster_ids = unique(clusters) map(cluster_ids) do cluster_id - i_sites = findall(isequal.(clusters, cluster_id)) - f(i_sites) + itrain_sites = findall(isequal.(clusters, cluster_id)) + f(itrain_sites) end end diff --git a/src/elbo.jl b/src/elbo.jl index 199332e..cc823e8 100644 --- a/src/elbo.jl +++ b/src/elbo.jl @@ -13,12 +13,12 @@ expected value of the likelihood of observations. - `f`: mechanistic model - `py`: negative log-likelihood of observations given predictions: `function(y_ob, y_pred, y_unc)` -- `xM`, `xP`, `y_ob`, `y_unc`, `i_sites`: information of the sites in the current minibatch +- `xM`, `xP`, `y_ob`, `y_unc`, `itrain_sites`: information of the sites in the current minibatch - `xM`: matrix of covariates (n_cov x n_site_batch) - `xP`: model drivers, iterable of (n_site_batch) - `y_ob`: matrix of observations (n_obs x n_site_batch) - `y_unc`: observation uncertainty provided to py (same size as y_ob) - - `i_sites`: indices of sites for current minibatch + - `itrain_sites`: indices of sites incurrent minibatch within training dataset - `interpreters`: NamedTuple as generated by `gen_hybridproblem_synthetic` with entries: - `ϕg_ϕq`: extract components of parameter of 1) means of global PBM, 2) ML-weights, and 3) additional parameters of approximation q @@ -29,12 +29,12 @@ expected value of the likelihood of observations. """ function neg_elbo_gtf(args...; kwargs...) # TODO prior and penalty loss - (;nLjoint, entropy_ζ, loss_penalty, - nLy, nLprior_P, nLprior_M, neg_log_jac, + (;nLjoint, entropy_ζ, loss_penalty, nLRanef, + # nLy, nLprior_P, nLprior_Ms, neglogjac_P, neglogjac_Ms, nLsites, #nLmean_θ ) = neg_elbo_gtf_components(args...; kwargs...) # negative of log_joint - need to subtract entropy_ζ - nL = nLjoint + loss_penalty - entropy_ζ #+ nLmean_θ + nL = nLjoint - entropy_ζ + loss_penalty + nLRanef #+ nLmean_θ # if !isfinite(nL) # @show nL # @show nLjoint, entropy_ζ, loss_penalty, nLy, @@ -45,14 +45,11 @@ function neg_elbo_gtf(args...; kwargs...) end function neg_elbo_gtf_components(rng, ϕ::AbstractVector{FT}, g, f, py, - xM::AbstractMatrix, xP, y_ob, y_unc, i_sites::AbstractVector{<:Number}; + xM::AbstractMatrix, xP, y_ob, y_unc, itrain_sites::AbstractVector{<:Number}; int_ϕg_ϕq::AbstractComponentArrayInterpreter, int_ϕq::AbstractComponentArrayInterpreter, n_MC=12, n_MC_mean=n_MC, n_MC_cap=n_MC, cdev=cpu_device(), - priors_θP_mean=[], - priors_θMs_mean=[], - #priors_θ_mean=[], cor_ends, # =(P=(1,),M=(1,)) pbm_covar_indices, transP, transMs, @@ -65,20 +62,21 @@ function neg_elbo_gtf_components(rng, ϕ::AbstractVector{FT}, g, f, py, zero_prior_logdensity, approx::AbstractHVIApproximation, intθP, intθMs, + ranef::AbstractRandomEffectsComputer, frac_cluster_all, ) where {FT} ϕc = int_ϕg_ϕq(ϕ) VT= typeof(@view(ϕ[1:1])) - ϕg = CA.getdata(ϕc.ϕg)::VT - ϕq = CA.getdata(ϕc.ϕq)::VT + ϕg = CA.getdata(ϕc[Val(:ϕq)]) + ϕqc = ϕc[Val(:ϕq)] + #ϕq = CA.getdata(ϕqc)::VT if(!all(isfinite.(ϕ))) - @show ϕq + @show ϕqc @show ϕg error("encountered non-finite optimized parameters") end - n_MCr = isempty(priors_θP_mean) ? n_MC : max(n_MC, n_MC_mean) - ζsP, ζsMs_tr, σ = generate_ζ(approx, rng, g, ϕ, xM; n_MC=n_MCr, cor_ends, pbm_covar_indices, - int_ϕq, int_ϕg_ϕq, is_testmode, i_sites) + ζsP, ζsMs_tr, σ = generate_ζ(approx, rng, g, ϕ, xM; n_MC, cor_ends, pbm_covar_indices, + int_ϕq, int_ϕg_ϕq, is_testmode, itrain_sites, ranef) ζsP_cpu = cdev(ζsP) # fetch to CPU, because for <1000 sites (n_batch) this is faster ζsMs_tr_cpu = cdev(ζsMs_tr) # fetch to CPU, because for <1000 sites (n_batch) this is faster # @@ -86,29 +84,8 @@ function neg_elbo_gtf_components(rng, ϕ::AbstractVector{FT}, g, f, py, loss_comps = neg_elbo_ζtf( ζsP_cpu[:,1:n_MC], ζsMs_tr_cpu[:,:,1:n_MC], σ, f, py, xP, y_ob, y_unc; n_MC_cap, transP, transMs, priorsP, priorsM, - penalty_computer, ϕg, ϕq, is_omit_priors, zero_prior_logdensity, - i_sites, intθMs, intθP, frac_cluster_all) - # - # maybe: provide trans_mP and trans_mMs with creating cost function - # not used any more and merging named tuples takes long - # nLmean_θ = _compute_negloglik_meanθ(ζsP_cpu, ζsMs_cpu; - # trans_mP, trans_mMs, priors_θP_mean, priors_θMs_mean, i_sites, ) - # (;loss_comps..., nLmean_θ) -end - -function _compute_negloglik_meanθ(ζsP::AbstractMatrix{FT}, ζsMs_tr; - priors_θP_mean, priors_θMs_mean, i_sites, trans_mP, trans_mMs, -) where FT - if isempty(priors_θP_mean) - return zero(FT) - end - θsP, θsMs_tr = transform_ζs(ζsP, ζsMs_tr; trans_mP, trans_mMs) - mean_θP = mean(CA.getdata(θsP); dims=(2))[:, 1] - nLmean_θP = map((d, θi) -> -logpdf(d, θi), priors_θP_mean, mean_θP) - mean_θMs = mean(θsMs_tr; dims=(3))[:, :, 1] - nLmean_θMs = map((d, θi) -> -logpdf(d, θi), priors_θMs_mean[i_sites], mean_θMs) - nLmean_θ = sum(nLmean_θP) + sum(nLmean_θMs) - convert(FT,nLmean_θ)::FT + penalty_computer, ϕg, ϕqc, is_omit_priors, zero_prior_logdensity, + itrain_sites, intθMs, intθP, ranef, frac_cluster_all) end """ @@ -143,22 +120,33 @@ function neg_elbo_ζtf(ζsP::AbstractArray{T}, ζsMs_tr, σ, f, py, xP, y_ob, y_ transP, transMs=StackedArray(transM, size(ζsMs_tr, 2)), priorsP, priorsM, - penalty_computer, ϕg, ϕq, + penalty_computer, ϕg, ϕqc, is_omit_priors::Val, zero_prior_logdensity, - i_sites, intθP, intθMs, + itrain_sites, intθP, intθMs, + ranef::AbstractRandomEffectsComputer, frac_cluster_all, ) where T n_MC = size(ζsP,2) #@show ζsMs_tr[1,4,:] # fourth component goes to NaN at some time - if !all(isfinite.(ζsMs_tr)) - return (; - nLjoint=T(1e9), entropy_ζ=zero(T), loss_penalty=zero(T), nLy=zero(T), - neg_log_prior=T(1e9), neg_log_jac=zero(T)) + n_site_pred = size(ζsMs_tr,1) + @assert length(itrain_sites) == n_site_pred + frac_cluster = if iszero(itrain_sites[1]) + # unknown site indices, e.g. test deata + zeros(eltype(frac_cluster_all),length(itrain_sites)) + else + frac_cluster_all[itrain_sites] end - frac_cluster = frac_cluster_all[i_sites] + _entropy_ζ = compute_entropyζ_from_σ(σ; + frac_cluster, + n_θP = size(ζsP,1), + n_θM = size(ζsMs_tr,2), + n_site = n_site_pred, + ) + !all(isfinite.(ζsMs_tr)) && return(get_null_elbo_components_sites(ζsMs_tr)) f_sample = (ζP, ζMs_tr) -> begin - θP, θMs_tr, logjac_P, logjac_Ms = transform_and_logjac_ζ(ζP, ζMs_tr; transP, transMs) + θP, θMs_tr, logjac_P, logjac_Ms = transform_and_logjac_ζ( + ζP, ζMs_tr; transP, transMs) if !all(isfinite.(θMs_tr)) i_row = findfirst(θM -> !all(isfinite.(θM)), eachrow(θMs_tr)) @show i_row #info "encountered non-finite θMs_tr at $(i_row)th site" @@ -173,26 +161,34 @@ function neg_elbo_ζtf(ζsP::AbstractArray{T}, ζsMs_tr, σ, f, py, xP, y_ob, y_ # using ShareAdd # @usingany Cthulhu # @descend_code_warntype f(θP, θMs, xP) - nLy_i = py(y_ob, y_pred_i, y_unc) - # MAYBE avoid convert by making sure penalty_computer returns proper type - # Test.@inferred compute_penalty(penalty_computer, y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), i_sites, ϕg, ϕq)[1] - # loss_penalty_i = convert.(typeof(nLy_i),first(compute_penalty(penalty_computer, - # y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), i_sites, ϕg, ϕq))) - loss_penalty_i = compute_penalty(penalty_computer, - y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), i_sites, ϕg, ϕq)[1] - nLprior_P_i, nLprior_M_is = compute_priors_logdensity(priorsP, priorsM, θP, θMs_tr, - is_omit_priors, zero_prior_logdensity) - # make sure names to not match outer, otherwise Box type instability - # scale Likelihood and penalties to estimate all-site case from batch case - # scale nLy, priorsM, log_jac (sum for Exp), loss_penalty, and entropy - # essentially all, except prior_θP - # penalty should also be scaled, but then it does not select good parameters - #loss_penalty_if = sum(loss_penalty_i .* frac_cluster) - loss_penalty_if = sum(loss_penalty_i) - nLprior_M_if = sum(nLprior_M_is .* frac_cluster) - neg_log_jac_if = -logjac_P -sum(logjac_Ms .* frac_cluster) - (nLy_i, nLprior_P_i, nLprior_M_if, neg_log_jac_if, loss_penalty_if) - #(nLy_i, 0.0, 0.0, 0.0) + # sum across sites + f_sample_pre(logjac_P, logjac_Ms, θP, θMs_tr, y_pred_i, addq_pred_i; + py, y_ob, y_unc, + penalty_computer, itrain_sites, ϕqc, + intθMs, intθP, + priorsP, priorsM, is_omit_priors, zero_prior_logdensity, + frac_cluster, + ) + # nLy_i = py(y_ob, y_pred_i, y_unc) + # # MAYBE avoid convert by making sure penalty_computer returns proper type + # # Test.@inferred compute_penalty(penalty_computer, y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), itrain_sites, ϕq)[1] + # # loss_penalty_i = convert.(typeof(nLy_i),first(compute_penalty(penalty_computer, + # # y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), itrain_sites, ϕq))) + # loss_penalty_i = compute_penalty(penalty_computer, + # y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), itrain_sites, ϕqc)[1] + # nLprior_P_i, nLprior_M_is = compute_priors_logdensity(priorsP, priorsM, θP, θMs_tr, + # is_omit_priors, zero_prior_logdensity) + # # make sure names to not match outer, otherwise Box type instability + # # scale Likelihood and penalties to estimate all-site case from batch case + # # scale nLy, priorsM, log_jac (sum for Exp), loss_penalty, and entropy + # # essentially all, except prior_θP + # # penalty should also be scaled, but then it does not select good parameters + # #loss_penalty_if = sum(loss_penalty_i .* frac_cluster) + # loss_penalty_if = sum(loss_penalty_i) + # nLprior_M_if = sum(nLprior_M_is .* frac_cluster) + # neg_log_jac_if = -logjac_P -sum(logjac_Ms .* frac_cluster) + # (nLy_i, nLprior_P_i, nLprior_M_if, neg_log_jac_if, loss_penalty_if) + # #(nLy_i, 0.0, 0.0, 0.0) end # only Vector inferred, need to provide type hint # make that all components use the same Float type @@ -203,7 +199,125 @@ function neg_elbo_ζtf(ζsP::AbstractArray{T}, ζsMs_tr, σ, f, py, xP, y_ob, y_ #@descend_code_warntype f_sample(first(eachcol(ζsP)), first(eachslice(ζsMs_tr; dims=3))) #Test.@inferred map(f_sample, eachcol(ζsP), eachslice(ζsMs_tr; dims=3)) map_res = map(f_sample, eachcol(ζsP), eachslice(ζsMs_tr; dims=3)) - nLys, nLpriors_P, nLpriors_M, neglogjacs, loss_penalties = vectuptotupvec(map_res) + compile_elbo_components_sites(map_res; ranef, ϕqc, entropy_ζ = _entropy_ζ, n_MC_cap ) +end + +""" +Compute the elbo components for a single parameter sample. +Site-specific vectors: nL, loss_penalty +""" +function f_sample_pre(logjac_P, logjac_Ms, θP, θMs_tr, y_pred_i, addq_pred_i; + py, y_ob, y_unc, + penalty_computer, itrain_sites, ϕqc, + intθMs, intθP, + priorsP, priorsM, is_omit_priors, zero_prior_logdensity, + frac_cluster, + ) + nLy_i = py(y_ob, y_pred_i, y_unc) # n_site + # TODO return by site + loss_penalty_i = compute_penalty(penalty_computer, + y_pred_i, addq_pred_i, intθMs(θMs_tr), intθP(θP), itrain_sites, ϕqc)[1] + nLprior_P_i, nLprior_M_is = compute_priors_logdensity(priorsP, priorsM, θP, θMs_tr, + is_omit_priors, zero_prior_logdensity) + nLprior_M_if = nLprior_M_is .* frac_cluster + neg_log_jac_P = -logjac_P + neg_log_jac_Ms = -logjac_Ms .* frac_cluster + (;nLys=nLy_i, nLprior_P=nLprior_P_i, mLprior_Ms=nLprior_M_if, + neg_log_jac_P, neg_log_jac_Ms, loss_penalties=loss_penalty_i) +end + +function compile_elbo_components_sites(map_res; # result of map(f_sample_pre, ...) + ranef::AbstractRandomEffectsComputer, ϕqc::CA.ComponentVector, + entropy_ζ, + n_MC_cap, + ) + nLys_v, nLpriors_P, nLpriors_M_v, neglogjacs_P, neglogjacs_Ms_v, loss_penalties_v = vectuptotupvec(map_res) + ndims(nLpriors_P .+ neglogjacs_P) != 1 && error( + "adapt changed dimension of nLpriors_P, neglogjacs_P") + nLys = VectorOfArray(nLys_v) + loss_penalties = VectorOfArray(loss_penalties_v) + nLpriors_M = VectorOfArray(nLpriors_M_v) + neglogjacs_Ms = VectorOfArray(neglogjacs_Ms_v) + nLRanef = compute_nLranef(ranef, ϕqc[Val(:ranef)]) + res0_sites = average_elbo_components_across_samples(nLys, nLpriors_P, nLpriors_M, + neglogjacs_P, neglogjacs_Ms, loss_penalties; + n_MC_cap, nLRanef, entropy_ζ ) +end + +function average_elbo_components_across_samples(nLys, nLpriors_P, nLpriors_M, + neglogjacs_P, neglogjacs_Ms, loss_penalties; + n_MC_cap, + nLRanef, + entropy_ζ + ) + n_site_pred, n_MC = size(nLys) # size(nLpriors_M) + # For robustness may compute the expectation only on the n_smallest values + # because its very sensitive to few large outliers + #nLys_smallest = nsmallest(n_MC_cap, nLys) # does not work with Zygote + if n_MC_cap != n_MC + @warn "neg_elbo_ζtf: TPDP n_MC_cap: implement for for logjac, loss_penalty, and neg_log_prior not capped" + #nLys_smallest = partialsort(nLys, 1:n_MC_cap) + end + nLy = sum_voa(nLys; dim_ret=1) ./ n_MC + # distribute cross-site cost equally across sites + nLprior_P = sum(nLpriors_P) / n_MC + nLprior_Ms = sum_voa(nLpriors_M; dim_ret=1) ./ n_MC + neglogjac_P = sum(neglogjacs_P) / n_MC + neglogjac_Ms = sum_voa(neglogjacs_Ms; dim_ret=1) ./ n_MC + loss_penalty_sites = sum_voa(loss_penalties; dim_ret=1) ./ n_MC + #n_θM, n_site = size(ζsMs_tr)[1:2] + # why without underscore boxed and type instability? + # if itrain_sites[1] == 1 + # #Main.@infiltrate_main + # @show nLy, entropy_ζ, nLmean_θ, n_MC, n_MC_cap, itrain_sites[1:3] + # @show std(nLys), std(nLys)/abs(nLy) + # @show std(nLys_smallest), std(nLys_smallest)/abs(nLy) + # end + # sum across sites for overall loss + nLjoint = sum(nLy) + nLprior_P + sum(nLprior_Ms) + neglogjac_P + sum(neglogjac_Ms) + nLsites = nLy .+ nLprior_Ms .+ neglogjac_Ms .+ loss_penalty_sites + (;nLjoint, entropy_ζ, loss_penalty = sum(loss_penalty_sites), nLRanef, + nLy, nLprior_P, nLprior_Ms, neglogjac_P, neglogjac_Ms, loss_penalty_sites, nLsites) +end + +function get_null_elbo_components_sites(θsMs_tr::AbstractArray{T,3}) where T + n_site = size(θsMs_tr,1) + return (; + nLjoint=fill(T(1e9), n_site), + entropy_ζ=zero(T), + loss_penalty=fill(zero(T), n_site), + nLy=fill(zero(T), n_site), + neg_log_prior=T(1e9), + neg_log_jac=zero(T)) +end + + + +""" +sum x across subarray for each entry in dim_ret. This is opposite of the +dims argument in sum. +""" +function sum_voa(x::Union{AbstractArray{T},AbstractVectorOfArray{T}}; + dim_ret::Integer + ) where T + # sum_voa(x; dims=1) is mutating for VectorOfArray -> implement by indexing + # i = first(axes(x, dim_ret)) + map(axes(x, dim_ret)) do i + sum(fix_dim(x, dim_ret, i))::T + end +end + +function fix_dim(A::Union{AbstractArray, AbstractVectorOfArray}, dim::Int, i::Int) + indices = ntuple(d -> d == dim ? i : Colon(), ndims(A)) + return @view A[indices...] +end + +function compute_summed_elbo_components(nLys, nLpriors_P, nLpriors_M, neglogjacs, loss_penalties; + n_MC_cap, + ranef::AbstractRandomEffectsComputer, ϕqc::CA.ComponentVector, + entropy_ζ + ) + n_MC = length(nLys) # For robustness may compute the expectation only on the n_smallest values # because its very sensitive to few large outliers #nLys_smallest = nsmallest(n_MC_cap, nLys) # does not work with Zygote @@ -218,30 +332,76 @@ function neg_elbo_ζtf(ζsP::AbstractArray{T}, ζsMs_tr, σ, f, py, xP, y_ob, y_ nLys_smallest = partialsort(nLys, 1:n_MC_cap) nLy = sum(nLys_smallest) / n_MC_cap end - # sum_log_σ = sum(log.(σ)) - # logdet_jacT2 = -sum_log_σ # log Prod(1/σ_i) = -sum log σ_i - #logdetΣ = 2 * sum(log.(σ)) # det(Σ) = Prod(σ_i^2) - # also scale entropy (that depends on logdetΣ) for only a fraction of sites in btach - n_θP = size(ζsP,1) - n_θM, n_site = size(ζsMs_tr)[1:2] - n_θ = n_θP + n_θM * n_site - @assert length(σ) == n_θ - σP = σ[1:n_θP] - σMs = reshape(σ[(n_θP+1):end], :, n_site) - #logdetΣ = 2 * (sum(log.(σ))) # det(Σ) = Prod(σ_i^2) - logdetΣ = 2 * (sum(log.(σP)) + sum(frac_cluster .* log.(σMs))) # det(Σ) = Prod(σ_i^2) - entropy_ζ = entropy_MvNormal(n_θ, logdetΣ) # defined in logden_normal - # if i_sites[1] == 1 + #n_θM, n_site = size(ζsMs_tr)[1:2] + # why without underscore boxed and type instability? + # if itrain_sites[1] == 1 + # #Main.@infiltrate_main + # @show nLy, entropy_ζ, nLmean_θ, n_MC, n_MC_cap, itrain_sites[1:3] + # @show std(nLys), std(nLys)/abs(nLy) + # @show std(nLys_smallest), std(nLys_smallest)/abs(nLy) + # end + nLjoint = nLy + nLprior_P + nLprior_M + neg_log_jac + nLRanef = compute_nLranef(ranef, ϕqc[Val(:ranef)]) + (;nLjoint, entropy_ζ, loss_penalty, + nLy, nLprior_P, nLprior_M, neg_log_jac, nLRanef) +end + +function compute_elbo_components(nLys, nLpriors_P, nLpriors_M, neglogjacs, loss_penalties; + n_MC_cap, + ranef::AbstractRandomEffectsComputer, ϕqc::CA.ComponentVector, + entropy_ζ + ) + n_MC = length(nLys) + # For robustness may compute the expectation only on the n_smallest values + # because its very sensitive to few large outliers + #nLys_smallest = nsmallest(n_MC_cap, nLys) # does not work with Zygote + if n_MC_cap == n_MC + nLy = sum(nLys) / n_MC + nLprior_P = sum(nLpriors_P) / n_MC + nLprior_M = sum(nLpriors_M) / n_MC + neg_log_jac = sum(neglogjacs) / n_MC + loss_penalty = sum(loss_penalties) / n_MC + else + @warn "neg_elbo_ζtf: TPDP n_MC_cap: implement for for logjac, loss_penalty, and neg_log_prior not capped" + nLys_smallest = partialsort(nLys, 1:n_MC_cap) + nLy = sum(nLys_smallest) / n_MC_cap + end + #n_θM, n_site = size(ζsMs_tr)[1:2] + # why without underscore boxed and type instability? + # if itrain_sites[1] == 1 # #Main.@infiltrate_main - # @show nLy, entropy_ζ, nLmean_θ, n_MC, n_MC_cap, i_sites[1:3] + # @show nLy, entropy_ζ, nLmean_θ, n_MC, n_MC_cap, itrain_sites[1:3] # @show std(nLys), std(nLys)/abs(nLy) # @show std(nLys_smallest), std(nLys_smallest)/abs(nLy) # end nLjoint = nLy + nLprior_P + nLprior_M + neg_log_jac + nLRanef = compute_nLranef(ranef, ϕqc[Val(:ranef)]) (;nLjoint, entropy_ζ, loss_penalty, - nLy, nLprior_P, nLprior_M, neg_log_jac) + nLy, nLprior_P, nLprior_M, neg_log_jac, nLRanef) +end + + + +""" +construct log-determinant of covariance matrix. +which depends on the σ = sqrt(diag(cov)). +The part for for θMs is reduced by considering that a minibatch contains +only a fraction of the data. +""" +function compute_entropyζ_from_σ(σ::AbstractVector{T}; frac_cluster, n_θP, n_θM, n_site) where T + n_θ = n_θP + n_θM * n_site + @assert length(σ) == n_θ + σP = σ[1:n_θP] + # n_site is last dimension: blocks in cov. + σMs = reshape(σ[(n_θP+1):end], :, n_site) + #logdetΣ = 2 * (sum(log.(σ))) # det(Σ) = Prod(σ_i^2) + logdetΣ = (T(2) * (sum(log.(σP)) + sum(frac_cluster .* log.(σMs')))) # det(Σ) = Prod(σ_i^2) + entropy_MvNormal(n_θ, logdetΣ) # defined in logden_normal end + + + function compute_priors_logdensity(priorsP, priorsM, θP, θMs_tr, ::Val{omit_priors}, zero_prior_logdensity) where {omit_priors} if omit_priors @@ -262,16 +422,33 @@ function compute_priors_logdensity(priorsP, priorsM, θP, θMs_tr, zero_prior_lo # prior for each parameter across vector (therefore Base.Fix1) of site #nLprior_Ms_pars = map(i_par -> -logpdf(priorsM[i_par], θMs_tr[:,i_par])::Vector{typeof(nlP0)}, 1:length(priorsM)) nLprior_Ms_pars = map(i_par -> -map( + # asserting the type fails with ForwardDiff and empty priorsP + # but otherwise not typestable Base.Fix1(logpdf, priorsM[i_par]), θMs_tr[:,i_par])::Vector{typeof(nlP0)}, 1:length(priorsM)) # aggregate across vars, for each site nLprior_Ms = reduce(+, nLprior_Ms_pars) if !isfinite(nlP0) || !all(isfinite.(nLprior_Ms)) - @show nlP0, sum(nLprior_Ms) - @show priorsM - @show θMs_tr - @show nLprior_Ms - error("inspect non-finite priors") + #Main.@infiltrate_main + () -> begin + #using ShareAdd + #@usingany UnicodePlots + histogram(θMs_tr[:,4]) + i_par = 2 + logpdf(priorsM[i_par], θMs_tr[:,i_par]) + quantile.(Ref(priorsM[i_par]), [0.025, 0.5, 0.975]) + x = quantile.(Ref(priorsM[i_par]), 0.025:0.05:0.975) + y = pdf.(Ref(priorsM[i_par]), x) + scatterplot(x, y) + lineplot(x,y) + density + end + #error("inspect non-finite priors") + # @warn("inspect non-finite priors") + # @show nlP0, sum(nLprior_Ms) + # @show priorsM + # @show θMs_tr + # @show nLprior_Ms end (; nLprior_P = nlP0, nLprior_Ms) end @@ -280,8 +457,8 @@ struct ZeroPenaltyComputer <: AbstractPenaltyComputer end function compute_penalty( ::ZeroPenaltyComputer, y_pred::AbstractMatrix, addq_pred::AbstractMatrix, θMs_tr::AbstractMatrix, θP::AbstractVector, - i_sites::AbstractVector, - ϕg, ϕq::AbstractVector) + itrain_sites::AbstractVector, + ϕq::AbstractVector) return (; penalty = fill(zero(eltype(θMs_tr)), size(θMs_tr,1))) end @@ -297,10 +474,14 @@ Prediction function for hybrid variational inference parameter model. ## Keyword arguments - `scenario` - `n_sample_pred` +- `n_sample_ranef`: how many samples shall share the same random effect - `xM`: covariates for the machine-learning model (ML): Matrix (n_θM x n_site_pred). Possibility to override the default from `get_hybridproblem_train_dataloader`. - `xP`: model drivers for process based model (PBM): Matrix with (n_site_pred) rows. Possibility to override the default from `get_hybridproblem_train_dataloader`. +- `frac_cluster`: how many sites in a cluster does each site in xP represent? + Default one, does not regards clusters. If influences the value of returned `entropy_ζ`. + Hence, it should be adjuter, for recomputing elbo components of observed sites. Returns an NamedTuple `(; y, θsP, θsMs_tr, entropy_ζ)` with entries - `y`: Array `(n_obs, n_site, n_sample_pred)` of model predictions. @@ -317,37 +498,41 @@ Returns an NamedTuple `(; y, θsP, θsMs_tr, entropy_ζ)` with entries - `ζsMs_tr`: ComponentArray `(n_site, n_θM, n_sample_pred)` of PBM model parameters that vary by site at the unconstrained scale. - `penalties`: output of problems penalty computer average across samples as a ComponentVector - Each component is a vector of length n_site. +- `logjacs_P`: Vector `(n_sample_pred)` of logarithm of determinant + of the transformation of ζP + that are kept constant across sites at the unconstrained scale. +- `logjacs_Ms`: Matrix `(n_site, n_sample_pred)` of logarithm of determinant + of the transformation of ζMs Note that for some approximations, such as `MeanVarSepHVIApproximation`, `prob.ϕq` contains uncertainty parameters that are specific to sites. For a proper prediction for new sites, update `prob.ϕq` before appropriately. If xM and xP is unspecified, the problem's sites are used and `prob.ϕq` is consistent. -If predicting for a subset of sites, specify keyword argument `i_sites` . +If predicting for a subset of training sites, specify keyword argument `itrain_sites`, +Otherwise the default of empty vector of site-ids is used. """ -function predict_hvi(rng, prob::AbstractHybridProblem; scenario=Val(()), +function predict_hvi(rng::AbstractRNG, prob::AbstractHybridProblem; scenario=Val(()), gdevs = get_gdev_MP(scenario), xM = nothing, xP = nothing, is_testmode = true, - i_sites = nothing, + itrain_sites = Int[], n_sample_pred = 200, + n_sample_ranef = n_sample_pred ÷ 10, + frac_cluster = ones(eltype(prob.ϕg), 0), kwargs... ) - if isnothing(xM) || isnothing(xP) - dl = get_hybridproblem_train_dataloader(prob; scenario) - dl_dev = gdev_hybridproblem_dataloader(dl; gdevs) - xM_dl, xP_dl = dl_dev.data[1:2] - i_sites = isnothing(i_sites) ? (1:size(xM_dl, 2)) : i_sites - xP = isnothing(xP) ? xP_dl[:,i_sites] : xP - xM = isnothing(xM) ? xM_dl[:,i_sites] : xM - end + # update to HybridProblem, if empty defaults were used + xM, xP, itrain_sites, frac_cluster = initX_from_prob(prob; scenario, gdevs, + xM, xP, itrain_sites, frac_cluster) # sample_posterior required consistent prob.ϕq and xM - (; θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr) = sample_posterior( - rng, prob, xM; scenario, gdevs, is_testmode, i_sites, n_sample_pred,kwargs...) + (; θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr, logjacs_P, logjacs_Ms) = sample_posterior( + rng, prob, xM; + scenario, gdevs, is_testmode, itrain_sites, n_sample_pred, n_sample_ranef, + frac_cluster, + kwargs...) # n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) - n_site_pred = size(θsMs_tr,1) # determined by size(xM) - @assert size(xP, 2) == n_site_pred + n_site_pred = size(xP, 2) f_batch = get_hybridproblem_PBmodel(prob; scenario) f = n_site_pred == n_batch ? f_batch : create_nsite_applicator(f_batch, n_site_pred) if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice @@ -367,8 +552,8 @@ function predict_hvi(rng, prob::AbstractHybridProblem; scenario=Val(()), penalties_sum = mapreduce((x,y) -> x .+ y, axes(y,3)) do i_MC CA.ComponentVector( compute_penalty(penalty_computer, - y[:,:,i_MC], addq[:,:,i_MC], intθMs(θsMs_tr[:,:,i_MC]), intθP(θsP[:,i_MC]), i_sites, - prob.ϕg, prob.ϕq) + y[:,:,i_MC], addq[:,:,i_MC], intθMs(θsMs_tr[:,:,i_MC]), intθP(θsP[:,i_MC]), itrain_sites, + prob.ϕq) ) # TODO separate from prob end # reshape into ComponentMatrix with site rows @@ -377,9 +562,136 @@ function predict_hvi(rng, prob::AbstractHybridProblem; scenario=Val(()), # ) # penalties = intPen(CA.getdata(penalties_sum)) ./ n_sample_pred penalties = reshape_penalty_matrix(penalties_sum) ./ n_sample_pred - (; y, addq, θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr, penalties) + (; y, addq, θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr, penalties, logjacs_P, logjacs_Ms) +end + +function initX_from_prob(prob; scenario, gdevs, xM, xP, itrain_sites, frac_cluster) + if isnothing(xM) || isnothing(xP) + dl = get_hybridproblem_train_dataloader(prob; scenario) + dl_dev = gdev_hybridproblem_dataloader(dl; gdevs) + xM_dl, xP_dl = dl_dev.data[1:2] + itrain_sites = isempty(itrain_sites) ? (1:size(xM_dl, 2)) : itrain_sites + xP = isnothing(xP) ? xP_dl[:,itrain_sites] : xP + xM = isnothing(xM) ? xM_dl[:,itrain_sites] : xM + end + n_site_pred = size(xM,2) + if isempty(frac_cluster) + frac_cluster = ones(eltype(frac_cluster), n_site_pred) + end + xM, xP, itrain_sites, frac_cluster +end + + +""" +Compute the elbo components based on precomputed parameters and predictions, +as returned by predict_hvi. +""" +function predict_hvi_and_compute_elbo_components( +rng::AbstractRNG, prob::AbstractHybridProblem; scenario=Val(()), + gdevs = get_gdev_MP(scenario), + is_testmode = true, + itrain_sites = 1:get_hybridproblem_n_site_and_batch(prob; scenario)[1], + n_sample_pred = 200, + n_sample_ranef = n_sample_pred ÷ 10, + frac_cluster_all = + ones(eltype(prob.ϕg), get_hybridproblem_n_site_and_batch(prob; scenario)[1]), + is_omit_priors = Val(false), +) + pt = get_hybridproblem_par_templates(prob; scenario) + priors = get_hybridproblem_priors(prob; scenario) + priorsP = Tuple(priors[k] for k in keys(pt.θP)) + priorsM = Tuple(priors[k] for k in keys(pt.θM)) + penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + py = get_hybridproblem_neg_logden_obs(prob; scenario) + ϕqc = prob.ϕq + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + dl = get_hybridproblem_train_dataloader(prob; scenario) + dl_dev = gdev_hybridproblem_dataloader(dl; gdevs) + xM_dl, xP_dl = dl_dev.data[1:2] + itrain_sites = dl_dev.data[5] + ii_sites = map(x -> findfirst(==(x), itrain_sites), itrain_sites) + frac_cluster = frac_cluster_all[ii_sites] + xP = xP_dl[:,ii_sites] + xM = xM_dl[:,ii_sites] + y_ob = dl_dev.data[3][:,ii_sites] + y_unc = dl_dev.data[4][:,ii_sites] + res_predict = a = predict_hvi(rng, prob; + scenario, gdevs, xP, xM, is_testmode, itrain_sites, + n_sample_pred, n_sample_ranef, frac_cluster, + ); + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(pt.θM))) + res_elbo = compute_elbo_components_sites( + a.entropy_ζ, a.logjacs_P, a.logjacs_Ms, a.θsP, a.θsMs_tr, a.y, a.addq, + itrain_sites, y_ob, y_unc; + py, + penalty_computer, + ranef, + ϕqc, + frac_cluster, + is_omit_priors, priorsP, priorsM, + ) + (;res_predict, res_elbo) +end + + +function compute_elbo_components_sites( + entropy_ζ, + logjacs_P::AbstractVector{T}, logjacs_Ms::AbstractMatrix{T}, + θsP::AbstractMatrix{T}, θsMs_tr::AbstractArray{T,3}, + y, addq, + itrain_sites, y_ob, y_unc; + py, + priorsP, priorsM, + penalty_computer, + ranef::AbstractRandomEffectsComputer, ϕqc, + frac_cluster, + is_omit_priors::Val{omit_priors} = Val(false), + n_MC_cap = size(θsP,2), + intθP = ComponentArrayInterpreter(isempty(θsP) ? CA.ComponentVector{T}() : θsP[:,1]), + intθMs = ComponentArrayInterpreter(θsMs_tr[:,:,1]), +) where {T, omit_priors} + #n_MC = size(θsMs_tr,3) + n_site = size(θsMs_tr,1) + !all(isfinite.(θsMs_tr)) && return(get_null_elbo_components_sites(θsMs_tr)) + zero_prior_logdensity = omit_priors ? zero(T) : get_zero_prior_logdensity( + priorsP, priorsM, intθP(CA.getdata(θsP[:,1])), intθMs(CA.getdata(θsMs_tr[:,:,1]))) + # call f_sample_pre with all the keyword arguments that do not change across samples + f_sample_pre_wrapper = (args...) -> begin + f_sample_pre(args...; + py, y_ob, y_unc, + penalty_computer, itrain_sites, ϕqc, + intθMs, intθP, + priorsP, priorsM, is_omit_priors, zero_prior_logdensity, + frac_cluster, + ) + end + map_res = map(f_sample_pre_wrapper, + logjacs_P, eachcol(logjacs_Ms), + eachcol(θsP), eachslice(θsMs_tr; dims=3), + eachslice(y; dims=3), eachslice(addq; dims=3) + ) + compile_elbo_components_sites(map_res; ranef, ϕqc, entropy_ζ, n_MC_cap) + + # nLys_v, nLpriors_P, nLpriors_M, neglogjacs, loss_penalties_v = vectuptotupvec(map_res) + # # mean for each site across samples (ignore n_mc_cat for now) + # #Main.@infiltrate_main + # nLys = VectorOfArray(nLys_v) + # loss_penalties = VectorOfArray(loss_penalties_v) + # ndims(nLpriors_P .+ nLpriors_M .+ neglogjacs .+ entropy_ζ) != 1 && error( + # "returning site-based nLprior_P .+ nLprior_M .+ neglogjac: adjust computing nLjoint") + # #ii_site = first(axes(map_res,1)) + # nLjoint = map(axes(map_res,1)) do ii_site + # mean(nLys[ii_site,:] .+ nLpriors_P .+ nLpriors_M .+ neglogjacs) + # end + # #nLjoint = nLys .+ nLprior_P .+ nLprior_M .+ neglogjac + # nLRanef = compute_nLranef(ranef, ϕqc[Val(:ranef)]) + # res0_sites = average_elbo_components_across_samples(nLys, nLpriors_P, nLpriors_M, neglogjacs, loss_penalties; + # n_MC_cap, nLRanef, entropy_ζ) end + """ sample_posterior(rng, prob, [xM::AbstractMatrix]; scenario=Val(()), kwargs...) @@ -397,7 +709,7 @@ Optional keyword arguments - `n_sample_pred`: number of samples to draw, defaults to 200 - `gdevs`: `NamedTuple(gdev_M, gdev_P)`: GPU devices for machine learning model and parameter transformtation, default to [`get_gdev_MP`](@ref)`(scenario)`. -- `is_inferred`: set to `Val(true)` to activate type stabilicy check for transformation +- `is_inferred`: set to `Val(true)` to activate type stability check for transformation Returns an NamedTuple `(; θsP, θsMs_tr, entropy_ζ)` with entries - `θsP`: ComponentArray `(n_θP, n_sample_pred)` of PBM model parameters @@ -407,33 +719,36 @@ Returns an NamedTuple `(; θsP, θsMs_tr, entropy_ζ)` with entries - `entropy_ζ`: The entropy of the log-determinant of the transformation of the set of model parameters, which is involved in uncertainty quantification. """ -function sample_posterior(rng, prob::AbstractHybridProblem; i_sites=nothing, scenario=Val(()), +function sample_posterior(rng::AbstractRNG, prob::AbstractHybridProblem; + itrain_sites=Int[], scenario=Val(()), gdevs = get_gdev_MP(scenario), - kwargs...) + kwargs... + ) dl = get_hybridproblem_train_dataloader(prob; scenario) dl_dev = gdev_hybridproblem_dataloader(dl; gdevs) xM_all = dl_dev.data[1] - if !isnothing(i_sites) - xM = xM_all[:,i_sites] - else + if isnothing(itrain_sites) || isempty(itrain_sites) xM = xM_all - i_sites = 1:size(xM, 2) + itrain_sites = 1:size(xM, 2) + else + xM = xM_all[:,itrain_sites] end - sample_posterior(rng, prob, xM; scenario, gdevs, i_sites,kwargs...) + sample_posterior(rng, prob, xM; scenario, gdevs, itrain_sites,kwargs...) end -function sample_posterior(rng, prob::AbstractHybridProblem, xM::AbstractMatrix; +function sample_posterior(rng::AbstractRNG, prob::AbstractHybridProblem, xM::AbstractMatrix{T}; scenario=Val(()), - n_sample_pred=200, + n_sample_pred = 200, + n_sample_ranef = n_sample_pred ÷ 10, gdevs = get_gdev_MP(scenario), approx = nothing, + frac_cluster = ones(T, get_hybridproblem_n_site_and_batch(prob; scenario)[1]), kwargs... -) +) where T n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) is_predict_batch = (n_batch == size(xM,2)) - n_site_pred = is_predict_batch ? n_batch : n_site - @assert size(xM, 2) == n_site_pred + n_site_pred = size(xM, 2) par_templates = get_hybridproblem_par_templates(prob; scenario) cor_ends = get_hybridproblem_cor_ends(prob; scenario) g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario) @@ -446,49 +761,72 @@ function sample_posterior(rng, prob::AbstractHybridProblem, xM::AbstractMatrix; int_ϕq = interpreters.ϕq transMs = StackedArray(transM, n_batch) g_dev, ϕ_dev = gdevs.gdev_M(g), gdevs.gdev_M(ϕ) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer(ranef_spec, keys(par_templates.θM),n_site, one(eltype(ϕq))) if isnothing(approx) approx = prob.approx # assuming has field approx, e.g. if its a HybridProblem end - (; θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr) = sample_posterior(rng, g_dev, ϕ_dev, xM; + (; θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr, logjacs_P, logjacs_Ms) = sample_posterior(rng, g_dev, ϕ_dev, xM; int_ϕg_ϕq, int_ϕq, transP, transM, - n_sample_pred, cdev=infer_cdev(gdevs), cor_ends, pbm_covar_indices, approx, + n_sample_pred, n_sample_ranef, + cdev=infer_cdev(gdevs), cor_ends, pbm_covar_indices, approx, + ranef, + frac_cluster, kwargs...) # attach ComponentArray structure intP = ComponentArrayInterpreter(par_templates.θP, (n_sample_pred,)) θsPc = intP(θsP) ζsPc = intP(ζsP) - intMs = ComponentArrayInterpreter((n_site,), par_templates.θM, (n_sample_pred,)) + intMs = ComponentArrayInterpreter((n_site_pred,), par_templates.θM, (n_sample_pred,)) θsMsc_tr = intMs(θsMs_tr) ζsMsc_tr = intMs(ζsMs_tr) - (; θsP=θsPc, θsMs_tr=θsMsc_tr, entropy_ζ, ζsP=ζsPc, ζsMs_tr=ζsMsc_tr) + (; θsP=θsPc, θsMs_tr=θsMsc_tr, entropy_ζ, ζsP=ζsPc, ζsMs_tr=ζsMsc_tr, logjacs_P, logjacs_Ms) end -function sample_posterior(rng, g, ϕ::AbstractVector, xM::AbstractMatrix; +function sample_posterior(rng::AbstractRNG, g, ϕ::AbstractVector, xM::AbstractMatrix; int_ϕg_ϕq::AbstractComponentArrayInterpreter, int_ϕq::AbstractComponentArrayInterpreter, transP, transM, n_sample_pred, + n_sample_ranef, cdev, cor_ends, pbm_covar_indices, is_inferred::Val{is_infer} = Val(false), is_testmode, approx::AbstractHVIApproximation, - i_sites, + itrain_sites, + ranef::AbstractRandomEffectsComputer, + frac_cluster::AbstractVector, ) where is_infer - ζsP_gpu, ζsMs_tr_gpu, σ = generate_ζ(approx, rng, g, CA.getdata(ϕ), CA.getdata(xM); + ζsP_gpu, ζsMs_tr_gpu, σ = is_infer ? + Test.@inferred(generate_ζ(approx, rng, g, CA.getdata(ϕ), CA.getdata(xM); + int_ϕg_ϕq, int_ϕq, + n_MC=n_sample_pred, n_sample_ranef, + cor_ends, pbm_covar_indices, is_testmode, itrain_sites, ranef, + )) : + generate_ζ(approx, rng, g, CA.getdata(ϕ), CA.getdata(xM); int_ϕg_ϕq, int_ϕq, - n_MC=n_sample_pred, cor_ends, pbm_covar_indices, is_testmode, i_sites) + n_MC=n_sample_pred, n_sample_ranef, + cor_ends, pbm_covar_indices, is_testmode, itrain_sites, ranef, + ) ζsP = cdev(ζsP_gpu) ζsMs_tr = cdev(ζsMs_tr_gpu) - logdetΣ = 2 * sum(log.(σ)) - entropy_ζ = entropy_MvNormal(length(σ), logdetΣ) + # logdetΣ = 2 * sum(log.(σ)) + # entropy_ζ = entropy_MvNormal(length(σ), logdetΣ) + n_θP = size(ζsP,1) + n_θM = size(ζsMs_tr, 2) + n_site = size(ζsMs_tr, 1) + entropy_ζ = is_infer ? + Test.@inferred(compute_entropyζ_from_σ(σ; frac_cluster, n_θP, n_θM, n_site)) : + compute_entropyζ_from_σ(σ; frac_cluster, n_θP, n_θM, n_site) trans_mP = StackedArray(transP, size(ζsP, 2)) trans_mMs = StackedArray(transM, size(ζsMs_tr, 1) * size(ζsMs_tr, 3)) - θsP, θsMs_tr = is_infer ? - Test.@inferred(transform_ζs(ζsP, ζsMs_tr; trans_mP, trans_mMs)) : - transform_ζs(ζsP, ζsMs_tr; trans_mP, trans_mMs) - (; θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr) + θsP, θsMs_tr, logjacs_P, logjacs_Ms = is_infer ? + #transform_and_logjac_ζ + Test.@inferred(transform_and_logjacs_ζs(ζsP, ζsMs_tr; trans_mP, trans_mMs)) : + transform_and_logjacs_ζs(ζsP, ζsMs_tr; trans_mP, trans_mMs) + (; θsP, θsMs_tr, entropy_ζ, ζsP, ζsMs_tr, logjacs_P, logjacs_Ms) end @@ -510,10 +848,14 @@ function generate_ζ( g, ϕ::AbstractVector{FT}, xM::MT; int_ϕg_ϕq::AbstractComponentArrayInterpreter, int_ϕq::AbstractComponentArrayInterpreter, - n_MC=3, cor_ends, pbm_covar_indices, + ranef::AbstractRandomEffectsComputer, + n_MC=3, + n_sample_ranef = size(xM,2) ÷ 10, + cor_ends, pbm_covar_indices, is_testmode, - i_sites, # = 1:size(xM,2), + itrain_sites, # = 1:size(xM,2), ) where {FT,MT} + n_site = size(xM,2) # see documentation of neg_elbo_gtf ϕc = int_ϕg_ϕq(CA.getdata(ϕ)) #VT= typeof(@view(ϕ[1:1])) @@ -521,14 +863,33 @@ function generate_ζ( ϕq = CA.getdata(ϕc[Val(:ϕq)]) ϕqc = int_ϕq(ϕq) μ_ζP = CA.getdata(ϕqc[Val(:μP)]) - # first pass: append μ_ζP_to covars, need ML prediction for magnitude of ζMs + # first pass only : append μ_ζP_to covars, need ML prediction for magnitude of ζMs # TODO replace pbm_covar_indices by ComponentArray? dimensions to be type-inferred? + n_site_pred = size(xM, 2) xMP0 = _append_each_covars(xM, CA.getdata(μ_ζP), pbm_covar_indices) ϕm0 = g(xMP0, ϕg; is_testmode) ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(approx, rng, - i_sites, ϕm0, ϕq; n_MC, cor_ends, int_ϕq) + itrain_sites, ϕm0, ϕq; n_MC, cor_ends, int_ϕq) + # if itrain_sites is given add known random effect to μ_ζMs_g0 + # if random effect is unknown, sample for each group in ζMs_parfirst_resids + ϕq_ranef = ϕqc[Val(:ranef)] n_θm = size(ζMs_parfirst_resids, 1) - μ_ζMs0 = ϕm0[1:n_θm, :] + μ_ζMs0_tr = ϕm0[1:n_θm, :]' + # if random effect is known, i.e. itrain_sites is provided, add to ζMs_resids + # otherwise add the same sampled random effect for n_sample_ranef sites to ζMs_resids + ζMs_parfirst_ranef_resids = if !isempty(itrain_sites) && !iszero(itrain_sites[1]) + # random effect (same for all samples, in 3rd dimension) is recycled + # adding it to μ is not valid, because μ may be resampled given ζPi covariates + add_ranef(ranef, ζMs_parfirst_resids, ϕq_ranef, itrain_sites) + else + # TODO move n_distinct_ranef and i_col out of the loop + n_distinct_ranef = ceil(Int, n_site / n_sample_ranef) + i_col = generate_repeated_integers(n_site_pred, n_sample_ranef) + β_distinct = sample_ranef(ranef, ϕq_ranef, n_distinct_ranef, n_MC) + β = β_distinct[:,i_col,:] + @assert size(β) == size(ζMs_parfirst_resids) + ζMs_parfirst_resids .+ β + end # if !all(isfinite.(μ_ζMs0)) # @show μ_ζMs0 # is_infinite_ϕg = !all(isfinite.(ϕg)) @@ -539,25 +900,27 @@ function generate_ζ( if pbm_covar_indices isa SA.SVector{0} # do not need to predict again but just add the residuals to μ_ζP and μ_ζMs #ζsP = μ_ζP .+ ζP_resids # n_par x n_MC # .+ on empty view does not work - ζsMs_tr = permutedims(μ_ζMs0 .+ ζMs_parfirst_resids, (2, 1, 3)) # n_site x n_par x n_MC + ζsMs_tr = μ_ζMs0_tr .+ permutedims(ζMs_parfirst_ranef_resids, (2, 1, 3)) # n_site x n_par x n_MC # if any(ζsMs[:,2,:] .> 80.0) # @show ζsMs - # @show ζMs_parfirst_resids + # @show ζMs_parfirst_ranef_resids # @show ϕc.ϕq.coef_logσ2_ζMs # error("encountered scaled residual outside envisoned range. Debug") # end else - #rP, rMs = first(zip(eachcol(ζP_resids), eachslice(ζMs_parfirst_resids;dims=3))) - ζsMs_vec = map(eachcol(ζsP), eachslice(ζMs_parfirst_resids; dims=3)) do ζP, rMs + #ζP, rMs = first(zip(eachcol(ζsP), eachslice(ζMs_parfirst_ranef_resids;dims=3))) + ζsMs_tr_vec = map(eachcol(ζsP), eachslice(ζMs_parfirst_ranef_resids; dims=3)) do ζP, rMs # second pass: append ζP rather than μ_ζP to covars to xM xMP = _append_each_covars(xM, CA.getdata(ζP), pbm_covar_indices) - μ_ζMst = ϕm = g(xMP, ϕg; is_testmode) - ζMs = (μ_ζMst .+ rMs)' # already transform to par-last form - ζMs + ϕm = g(xMP, ϕg; is_testmode) + μ_ζMs_gt = ϕm[1:n_θm, :] + # randeom effect already in rMs + ζMs_tr = (μ_ζMs_gt' .+ rMs') # already transform to par-last form + ζMs_tr end # ζsP = stack(map(first, ζst); dims=1) # n_MC x n_par # ζsMs = stack(map(x -> x[2], ζst); dims=1) # n_MC x n_site x n_par - ζsMs_tr = stack(ζsMs_vec) # n_site x n_par x n_MC + ζsMs_tr = stack(ζsMs_tr_vec) # n_site x n_par x n_MC end ζsP, ζsMs_tr, σ end @@ -581,7 +944,8 @@ function _append_each_covars(xM, ζP::AbstractVector, pbm_covar_indices::Abstrac _append_each_covars(xM, ζP_covar) end function _append_each_covars(xM, ζP_covar::AbstractVector) - #@show ζP, typeof(ζP) + isempty(ζP_covar) && return(xM) + #@show ζP_covar, typeof(ζP_covar), typeof(xM) @assert eltype(xM) == eltype(ζP_covar) #Main.@infiltrate_main ζP_rep = reduce(hcat, fill(ζP_covar, size(xM, 2))) @@ -633,7 +997,7 @@ function sample_ζresid_norm( #approx::Union{AbstractMeanHVIApproximation,AbstractMeanVarSepHVIApproximation}, approx::AbstractHVIApproximation, rng::Random.AbstractRNG, - i_sites, + itrain_sites, ϕm::AbstractMatrix, ϕq::AbstractVector, args...; n_MC, cor_ends, int_ϕq) @@ -650,7 +1014,7 @@ function sample_ζresid_norm( #z = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP) zP = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP) zMs = _create_randn(rng, CA.getdata(ζP), n_MC, n_θMs) # ζP only for type inference - sample_ζresid_norm(approx, i_sites, zP, zMs, CA.getdata(ϕm), ϕq, args...; + sample_ζresid_norm(approx, itrain_sites, zP, zMs, CA.getdata(ϕm), ϕq, args...; cor_ends, int_ϕq=get_concrete(int_ϕq) ) @@ -692,7 +1056,7 @@ end function sample_ζresid_norm( approx::Union{AbstractMeanScalingHVIApproximation, MeanHVIApproximationMat}, - i_sites, + itrain_sites, zP::AbstractMatrix, zMs::AbstractMatrix, ϕm::TM, ϕq::AbstractVector{T}; int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)), @@ -946,6 +1310,35 @@ function transform_ζs(ζsP::AbstractMatrix, ζsMs_tr::AbstractArray; θsP, θsMs_tr end +function transform_and_logjacs_ζs(ζsP::AbstractMatrix, ζsMs_tr::AbstractArray; + trans_mP::StackedArray=StackedArray(transP, n_MC), + trans_mMs::StackedArray=StackedArray(transM, n_MC * n_site_batch) +) + # transform to parameter-last that can apply transformations effectively + # n_site x n_par x n_MC -> n_MC x n_site x n_par + ζsMs_pl = permutedims(ζsMs_tr, (3, 1, 2)) + #θsMs_tr0 = trans_mMs(ζsMs_pl) + θsMs_pl, logjacs_M_comps = with_logabsdet_jacobians(trans_mMs, ζsMs_pl) + logjacs_Ms_pl = sum(logjacs_M_comps; dims = 3)[:,:,1] + logjacs_Ms = logjacs_Ms_pl' # n_site x n_MC + # backtransform to n_mc last for efficient mapping + θsMs_tr = permutedims(θsMs_pl, (2, 3, 1)) + #θsPt0 = trans_mP(ζsP') # Bijectors use copy and copy(ζsP') errors if ζsP is an empty CuArray + θsPt, logjacs_P_comps = with_logabsdet_jacobians(trans_mP, ζsP') + logjacs_P = sum(logjacs_P_comps; dims = 2)[:,1] + θsP = θsPt' + # θsP = if isempty(ζsP) + # # trans_mP(ζsP') of empty array has problems with AD + # # copy(ζsP')' # copy of empty array does no harm but ensures type is Adjoint + # ζsP # leads to type instability + # #θsMs_tr[1, 1:0, :] # workaround: extract empty matrix from first mc_batch of θsMs_tr not the same type + # else + # θsPt = trans_mP(ζsP') + # θsP = θsPt' + # end + θsP, θsMs_tr, logjacs_P, logjacs_Ms +end + function flatten_hybrid_pars(xsP::AbstractMatrix{FT}, xsMs::AbstractArray{FT,3}) where FT n_site_pred, n_θM, n_MC = size(xsMs) @assert size(xsP,2) == n_MC diff --git a/src/elbo2.jl b/src/elbo2.jl index f6b477c..f3cfb73 100644 --- a/src/elbo2.jl +++ b/src/elbo2.jl @@ -16,7 +16,7 @@ # and forward runs for many sites (760mus cpu vs 439ms gpu unit!) function sample_ζresid_norm(app::MeanHVIApproximation, - i_sites, + itrain_sites, zP::AbstractMatrix, zMs::AbstractMatrix, ϕm::TM, ϕq::AbstractVector{T}; int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)), diff --git a/src/elbo_sepvec.jl b/src/elbo_sepvec.jl index 80ce903..30cfd72 100644 --- a/src/elbo_sepvec.jl +++ b/src/elbo_sepvec.jl @@ -1,7 +1,7 @@ # Similar to MeanHVIApproximationMat # but own variance parameter for each predicted parameter function sample_ζresid_norm(approx::MeanVarSepHVIApproximation, - i_sites, + itrain_sites, zP::AbstractMatrix, zMs::AbstractMatrix, ϕm::TM, ϕq::AbstractVector{T}; int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)), @@ -20,7 +20,7 @@ function sample_ζresid_norm(approx::MeanVarSepHVIApproximation, UM = transformU_block_cholesky1(ρsM, cor_ends.M) # # make that ϕuncc[:logσ2_ζMs] stores the uncertainty for each site currently predicted - logσ2_logMs = ϕuncc[Val(:logσ2_ζMs)][:,i_sites] + logσ2_logMs = ϕuncc[Val(:logσ2_ζMs)][:,itrain_sites] logσ2_ζP = vec(CA.getdata(ϕuncc[Val(:logσ2_ζP)])) # CUDA cannot multiply BlockDiagonal * Diagonal, construct already those blocks σMs = reshape(exp.(logσ2_logMs ./ 2), n_θM, :) diff --git a/src/fit_ranef_site.jl b/src/fit_ranef_site.jl new file mode 100644 index 0000000..5a0d61a --- /dev/null +++ b/src/fit_ranef_site.jl @@ -0,0 +1,107 @@ +""" +Loss function for random effects, given fixed ζ. +""" +function get_loss_ran_tr_f(rng, prob, itrain_site; scenario = Val((:default,)), + train_data = NamedTuple{(:xM, :xP, :y, :y_unc, :itrain_site)}( + get_hybridproblem_train_dataloader(prob; scenario).data[:]), + is_omit_priors::Val{omit_priors} = Val(false), + frac_cluster = 1.0 + ) where omit_priors + # provide xM explicitly + res_predict_point = predict_point_hvi(rng, prob; train_data.xM, train_data.xP, ) + n_site = size(res_predict_point.θMs_tr, 1) + (;transM, transP) = get_hybridproblem_transforms(prob; scenario) + transMs = StackedArray(transM, n_site) + ζMs = inverse(transMs)(res_predict_point.θMs_tr)' + θP = CA.getdata(res_predict_point.θP) + intP = ComponentArrayInterpreter(res_predict_point.θP) + intMs1_tr = ComponentArrayInterpreter(res_predict_point.θMs_tr[[1],:]) + #ζMsc = intMs(ζMs) + #ζP = inverse(transP)(res_predict_point.θP) + pt = get_hybridproblem_par_templates(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(ζMs))) + f_batch = get_hybridproblem_PBmodel(prob; scenario) + f = create_nsite_applicator(f_batch, 1) + xP_itrain = train_data.xP[:,[itrain_site]] + ϕqc = prob.ϕq + ϕqc_cache = PreallocationTools.DiffCache(copy(ϕqc)) # copy to not modify orig + ϕq_ranef = prob.ϕq.ranef + intϕq_ranef = ComponentArrayInterpreter(ϕq_ranef) + priors = get_hybridproblem_priors(prob; scenario) + priorsP = Tuple(priors[k] for k in keys(pt.θP)) + priorsM = Tuple(priors[k] for k in keys(pt.θM)) + zero_prior_logdensity = omit_priors ? zero(eltype(pt.θP)) : get_zero_prior_logdensity( + priorsP, priorsM, pt.θP, pt.θM) + penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) + + let f=f, + y_o = train_data.y[:,[itrain_site]], + y_unc = train_data.y_unc[:,[itrain_site]], + py = get_hybridproblem_neg_logden_obs(prob; scenario), + ranef = ranef, ζMs = ζMs, θP = θP, + transMs = transMs, + xP_itrain = xP_itrain, + intϕq_ranef = intϕq_ranef, + ϕqc_cache = ϕqc_cache, + is_omit_priors = is_omit_priors, + zero_prior_logdensity = zero_prior_logdensity, + priorsP = priorsP, priorsM = priorsM, + frac_cluster = eltype(ζMs)(frac_cluster), + penalty_computer = penalty_computer, + intMs_tr = intMs1_tr, intP = intP, + rng = rng + + function loss_ran_tr_f(ranef_itrain::AbstractVector{T}) where T # uses mutation -> use NelderMead for few parameters + ϕqc1 = PreallocationTools.get_tmp(ϕqc_cache, ranef_itrain) + ϕqc1.ranef.β[itrain_site,:] .= ranef_itrain # mutation + # β0 = ϕq_ranef.β + # β = [r == itrain_site ? ranef_itrain[c] : β0[r,c] for r in axes(β0,1), c in axes(β0,2)] + # ϕq_ranef1 = CA.ComponentVector(;ϕq_ranef..., β) + ζMs_tr_ranef = add_ranef(ranef, ζMs, ϕqc1.ranef, [itrain_site])' + θMs_tr_ranef = transMs(ζMs_tr_ranef)[[itrain_site],:] + # do not add sampled alleatoric error here + y_pred, addq_pred = f( + θP, θMs_tr_ranef, xP_itrain) + () -> begin + # using ShareAdd + # @usingany UnicodePlots + pl = scatterplot(log10.(xP_itrain[:,1]), y_o[:,1]; label="obs", title="site $(itrain_site)") + scatterplot!(pl, log10.(xP_itrain[:,1]), y_pred[:,1]; label="pred") + θMs_tr_ranef + end + nLy = if !all(isfinite.(y_pred[isfinite.(y_o)])) + #@warn "encountered non-finite y_pred" + # random effect so large, that infinite after transformation + T(1e6) + else + res_py = py(y_o, y_pred, y_unc)[1] + # if !isfinite(res_py) + # @warn "encountered non-finite res_py" + # Main.@infiltrate_main + # end + res_py + end + # + nLprior_P, nLprior_Ms = + # @descend_code_warntype ( + compute_priors_logdensity(priorsP, priorsM, θP, θMs_tr_ranef, + is_omit_priors, zero_prior_logdensity) + nLprior_M = nLprior_Ms[1] * frac_cluster + # + loss_penalties = first(compute_penalty(penalty_computer, + y_pred, addq_pred, intMs_tr(θMs_tr_ranef), intP(θP), + [itrain_site], ϕqc1)) + + loss_penalty = loss_penalties[1] # * frac_cluster + # + nLjoint_pen = nLy + nLprior_M + loss_penalty #+ nLRanef + if !isfinite(nLjoint_pen) + nLjoint_pen = typeof(nLjoint_pen)(1e5) + end + #@show nLjoint_pen, nLy + (;nLjoint_pen, nLy, nLprior_M, loss_penalty) + end + end +end diff --git a/src/gf.jl b/src/gf.jl index be20b37..eabd4ca 100644 --- a/src/gf.jl +++ b/src/gf.jl @@ -62,8 +62,13 @@ the sampling step but returns the prediction at the mean in unconstrained space. - `gdevs` - `xM`: covariates for the machine-learning model (ML): Matrix (n_θM x n_site_pred). Possibility to override the default from `get_hybridproblem_train_dataloader`. + If the default it used, assume random effects are added. - `xP`: model drivers for process based model (PBM): Matrix with (n_site_pred) rows. Possibility to override the default from `get_hybridproblem_train_dataloader`. +- `itrain_sites`: Integer vector of known sites to predict. If positive length and not zero + then random effects for those sites are added. If `xM` is nothing, this defaults to + index in training dataset. If xM is provided, the default `itrain_sites = Int[]` results + in not adding random effects. Returns an NamedTuple `(; y, θMs_tr, θP)` with entries - `y`: Matrix `(n_obs, n_site)` of model predictions. @@ -76,16 +81,21 @@ function predict_point_hvi(rng, prob::AbstractHybridProblem; scenario=Val(()), gdevs = get_gdev_MP(scenario), xM = nothing, xP = nothing, is_testmode = true, + itrain_sites = Int[], kwargs... ) if isnothing(xM) || isnothing(xP) + is_sites_known = isnothing(xM) # predict for known problem training sites dl = get_hybridproblem_train_dataloader(prob; scenario) dl_dev = gdev_hybridproblem_dataloader(dl; gdevs) xM_dl, xP_dl = dl_dev.data[1:2] xM = isnothing(xM) ? xM_dl : xM xP = isnothing(xP) ? xP_dl : xP + if is_sites_known + itrain_sites = 1:size(xM,2) + end end - y_pred, addq_pred, θMs_tr, θP = gf(prob, xM, xP; scenario, gdevs, is_testmode, kwargs...) + y_pred, addq_pred, θMs_tr, θP = gf(prob, xM, xP; scenario, gdevs, is_testmode, itrain_sites, kwargs...) pt = get_hybridproblem_par_templates(prob) θPc = ComponentArrayInterpreter(pt.θP)(θP) θMsc = ComponentArrayInterpreter((size(θMs_tr,1),), pt.θM)(θMs_tr) @@ -107,6 +117,7 @@ function gf(prob::AbstractHybridProblem, xM::AbstractMatrix, xP::AbstractMatrix; scenario = Val(()), gdevs = nothing, #get_gdev_MP(scenario), is_inferred::Val{is_infer} = Val(false), + itrain_sites, kwargs... ) where is_infer gdevs = isnothing(gdevs) ? get_gdev_MP(scenario) : gdevs @@ -134,12 +145,16 @@ function gf(prob::AbstractHybridProblem, xM::AbstractMatrix, xP::AbstractMatrix; g_dev, ϕg_dev, xM_dev, ζP_dev = gdev(g), gdev(ϕg), gdev(CA.getdata(xM)), gdev(CA.getdata(ζP)) # most of the properties of prob are not type-inferred # hence result is not type-inferred, but may test at this context + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(ϕq))) + ϕq_ranef = ϕq[Val(:ranef)] res = is_infer ? Test.@inferred( gf( g_dev, transMs, transP, f_dev, xM_dev, xP, ϕg_dev, n_θM, ζP_dev, pbm_covar_indices; - cdev, kwargs...)) : + cdev, ranef, ϕq_ranef, itrain_sites, kwargs...)) : gf(g_dev, transMs, transP, f_dev, xM_dev, xP, ϕg_dev, n_θM, ζP_dev, pbm_covar_indices; - cdev, kwargs...) + cdev, ranef, ϕq_ranef, itrain_sites, kwargs...) end function gf(g::AbstractModelApplicator, transMs, transP, f, xM, xP, ϕg, n_θM, ζP; @@ -153,6 +168,7 @@ end function gf(g::AbstractModelApplicator, transMs, transP, f, xM, xP, ϕg, n_θM, ζP, pbm_covar_indices::AbstractVector{<:Integer}; + ranef::AbstractRandomEffectsComputer, ϕq_ranef, itrain_sites, cdev, is_testmode) # @show first(xM,5) # @show first(ϕg,5) @@ -163,7 +179,7 @@ function gf(g::AbstractModelApplicator, transMs, transP, f, xM, xP, ϕg, n_θM, # end #xMP = _append_PBM_covars(xM, intP(ζP), pbm_covars) xMP = _append_each_covars(xM, CA.getdata(ζP), pbm_covar_indices) - θMs_tr = gtrans(g, transMs, xMP, ϕg, n_θM; cdev, is_testmode) + θMs_tr = gtrans(g, transMs, xMP, ϕg, n_θM; ranef, ϕq_ranef, itrain_sites, cdev, is_testmode) # transPM = RRuleMonitor("transP", ζP -> transP(ζP)) # θP = transPM(CA.getdata(ζP)) θP = transP(CA.getdata(ζP)) @@ -180,11 +196,20 @@ end composition transM ∘ g: transformation after machine learning parameter prediction Provide a `transMs = StackedArray(transM, n_batch)` """ -function gtrans(g, transMs, xMP, ϕg, n_θM; cdev, is_testmode) +function gtrans(g, transMs, xMP, ϕg, n_θM; + ranef::AbstractRandomEffectsComputer, ϕq_ranef, itrain_sites, + cdev, is_testmode + ) ϕg = g(xMP, ϕg; is_testmode) - ζMs_tr = ϕg[1:n_θM,:]' # ignore the uncertainty-related parameters - ζMs_tr_cpu = cdev(ζMs_tr) - θMs_tr = transMs(ζMs_tr_cpu) + ζMs = ϕg[1:n_θM,:] # ignore the uncertainty-related parameters + ζMs_cpu0 = cdev(ζMs) + ζMs_ranef_tr_cpu = if !isempty(itrain_sites) && !iszero(itrain_sites[1]) + add_ranef(ranef, ζMs_cpu0, ϕq_ranef, itrain_sites)' + else + ζMs_cpu0' + end + #@show ζMs_tr_cpu .- ζMs_tr_cpu0 + θMs_tr = transMs(ζMs_ranef_tr_cpu) if !all(isfinite.(θMs_tr)) @info "gtrans: encountered non-finite parameters" #@show θMs_tr, ζMs_cpu, transMs @@ -200,12 +225,11 @@ Create a loss function for given - transM: transformation of parameters at unconstrained space - f(θMs_tr, θP): mechanistic model - py: `function(y_pred, y_obs, y_unc)` to compute negative log-likelihood, i.e. cost -- intϕ: interpreter attaching axis with components ϕg and ϕP -- intP: interpreter attaching axis to ζP = ϕP with components used by f, +- intϕ: interpreter attaching axis with components ϕg and ϕq The default, uses `intϕ(ϕ)` as a template - kwargs: additional keyword arguments passed to `gf`, such as `gdev` or `pbm_covars` -The loss function `loss_gf(ϕ, xM, xP, y_o, y_unc, i_sites)` takes +The loss function `loss_gf(ϕ, xM, xP, y_o, y_unc, itrain_sites)` takes - parameter vector ϕ - xM: matrix of covariate, sites in the batch are in columns - xP: iteration of drivers for each site @@ -213,7 +237,7 @@ The loss function `loss_gf(ϕ, xM, xP, y_o, y_unc, i_sites)` takes - y_unc: vector of uncertainty information for each observation Currently, hardcoes squared error loss of `(y_pred .- y_o) ./ σ`, with `σ = exp.(y_unc ./ 2)`. -- i_sites: index of sites in the batch +- itrain_sites: index of sites in the batch and returns a NamedTuple of - `nLjoint`: the negative-log of the joint parameter probability (Likelihood * prior) @@ -224,39 +248,48 @@ and returns a NamedTuple of - `neg_log_prior`: negative log-prior of `θMs_tr` and `θP` """ function get_loss_gf(g, transM, transP, f, py, - intϕ::AbstractComponentArrayInterpreter, - intP::AbstractComponentArrayInterpreter = ComponentArrayInterpreter( - intϕ(1:length(intϕ)).ϕP); + intϕ::AbstractComponentArrayInterpreter; + # intP::AbstractComponentArrayInterpreter = nothing, #ComponentArrayInterpreter( + # intϕ(1:length(intϕ))[:Val(:ϕq)][:Val(:ϕP)]); cdev=cpu_device(), par_templates::NamedTuple, pbm_covars, n_site_batch, penalty_computer = ZeroPenaltyComputer(), priorsP, priorsM, is_omit_priors::Val{omit_priors} = Val(false), - intθP, intθMs, + #intθP, intθM, + ranef::AbstractRandomEffectsComputer, frac_cluster_all, kwargs...) where omit_priors pt = par_templates zero_prior_logdensity = omit_priors ? zero(eltype(pt.θP)) : get_zero_prior_logdensity( - priorsP, priorsM, pt.θP, pt.θM) + priorsP, priorsM, pt.θP, pt.θM) + intθMs = ComponentArrayInterpreter((n_site_batch,), pt.θM) + intθP = ComponentArrayInterpreter(pt.θP) let g = g, transM = transM, transP = transP, f = f, intϕ = get_concrete(intϕ), transMs = StackedArray(transM, n_site_batch), cdev = cdev, - pbm_covar_indices = CA.getdata(intP(1:length(intP))[pbm_covars]), + pbm_covar_indices = CA.getdata(intθP(1:length(intθP))[pbm_covars]), zero_prior_logdensity = zero_prior_logdensity, is_omit_priors = is_omit_priors, priorsP = priorsP, priorsM = priorsM, penalty_computer = penalty_computer, intθMs = get_concrete(intθMs), intθP = get_concrete(intθP), + ranef = ranef, frac_cluster_all = convert.(eltype(pt.θP),frac_cluster_all), n_θM = length(priorsM), cpu_dev = cpu_device() # real cpu, different form infer_cdev(gdevs) that maybe idenetity #, intP = get_concrete(intP) #inv_transP = inverse(transP), kwargs = kwargs - function loss_gf(ϕ::AbstractVector{T}, xM, xP, y_o, y_unc, i_sites; is_testmode) where T + + function loss_gf(ϕ::AbstractVector{T}, xM, xP, y_o, y_unc, itrain_sites; + is_testmode, ignore_ranef::Val{ignore_ranef_val} = Val(false), + ) where {T, ignore_ranef_val} ϕc = intϕ(ϕ) + ϕqc = ϕc[Val(:ϕq)] # looses structure + ϕq_ranef = ϕqc[Val(:ranef)] # GPUArraysCore.allowscalar(() -> if !all(isfinite.(ϕ)) # @show ϕc.ϕP # error("invokded loss function loss_gf with non-finite parameters") @@ -270,18 +303,26 @@ function get_loss_gf(g, transM, transP, f, py, # y_pred, _, _ = apply_f_trans(ζP_cpu, ζMs_cpu, f, xP; transM, transP) if !all(isfinite.(ϕ)) @info "loss_gf: encountered non-finite ϕ" - @show ϕc.ϕP + @show ϕqc.ϕP #Main.@infiltrate_main end + n_site = size(xM,2) + ranef1 = ignore_ranef_val ? NullRandomEffectsComputer{T}(n_θM, n_site) : ranef y_pred, addq_pred, θMs_tr_pred, θP_pred = gf( - g, transMs, transP, f, xM, xP, CA.getdata(ϕc.ϕg), n_θM, - CA.getdata(ϕc.ϕP), - pbm_covar_indices; cdev, is_testmode, kwargs...) - # TODO check computation - frac_cluster = frac_cluster_all[i_sites] + g, transMs, transP, f, xM, xP, CA.getdata(ϕc[Val(:ϕg)]), n_θM, + ϕqc[Val(:μP)], + pbm_covar_indices; cdev, is_testmode, + ranef = ranef1, ϕq_ranef, itrain_sites, + kwargs...) + frac_cluster = if isempty(itrain_sites) || iszero(itrain_sites[1]) + ones(T, n_site) + else + frac_cluster_all[itrain_sites] + end #σ = exp.(y_unc ./ 2) #nLy = sum(abs2, (y_pred .- y_o) ./ σ) - nLy = py(y_o, y_pred, y_unc) + #nLy = sum(py(y_o, y_pred, y_unc)) + nLy = sum(py(y_o, y_pred, y_unc)) # logpdf is not typestable for Distribution{Univariate, Continuous} # logpdf_t = (prior, θ) -> logpdf(prior, θ)::eltype(θP_pred) # logpdf_tv = (prior, θ::AbstractVector) -> begin @@ -294,19 +335,28 @@ function get_loss_gf(g, transM, transP, f, py, nLprior_M = sum(nLprior_Ms .* frac_cluster) if !isfinite(nLprior_P) || !isfinite(nLprior_M) @info "loss_gf: encountered non-finite prior density" - @show θP_pred, θMs_tr_pred, ϕc.ϕP + @show θP_pred, θMs_tr_pred, ϕc.ϕq.μP error("debug get_loss_gf") end - ϕq = eltype(θP_pred)[] # no uncertainty parameters optimized loss_penalties = first(compute_penalty(penalty_computer, y_pred, addq_pred, intθMs(θMs_tr_pred), intθP(θP_pred), - i_sites, ϕc.ϕg, ϕq)) + itrain_sites, ϕqc)) #loss_penalty = sum(loss_penalties .* frac_cluster) loss_penalty = sum(loss_penalties) #@show nLy, neg_log_prior, loss_penalty - nLjoint_pen = nLy + nLprior_P + nLprior_M + loss_penalty + nLRanef = compute_nLranef(ranef1, ϕqc[Val(:ranef)]) + ndims(nLprior_P) != 0 && error( + "adapt changed dimension of nLpriors_P, nLpriors_M, neglogjacs, loss_penalties") + # divide scalar cost equally across sites + # n_site = length(nLy) + nLjoint_pen = nLy + nLprior_M + nLprior_P + loss_penalty + nLRanef + if !isfinite(nLjoint_pen) + @show (;nLjoint_pen, nLy, nLprior_M, nLprior_P, loss_penalty, nLRanef) + error("encountered non-finite nLjoint_pen in loss_gf") + end + return (;nLjoint_pen, y_pred, θMs_tr_pred, θP_pred, nLy, nLprior_P, - nLprior_M, loss_penalty) + nLprior_M, loss_penalty, nLRanef) end end end diff --git a/src/logden_normal.jl b/src/logden_normal.jl index 1d2c0b3..f372e1a 100644 --- a/src/logden_normal.jl +++ b/src/logden_normal.jl @@ -5,7 +5,9 @@ Compute the negative Log-density of `obs` for multiple independent normal distri given estimated means `μ` and estimated log of variance parameters `logσ2s`. All the arguments should be vectors of the same length. -If `obs`, `μ` are given as a matrix of several column-vectors, their summed +If `obs`, `μ` are given as a matrix of several column-vectors, +the likelihood of each column is returned, assuming each column having the same `logσ2s`. +depr: If `obs`, `μ` are given as a matrix of several column-vectors, their summed Likelihood is computed, assuming each column having the same `logσ2s`. Keyword argument `σfac` can be increased to put more weight on achieving @@ -13,8 +15,9 @@ a low uncertainty estimate and means closer to the observations to help an initial fit. The obtained parameters then can be used as starting values for a the proper fit with `σfac=1.0`. """ -function neg_logden_indep_normal(obs::AbstractArray, μ::AbstractArray, logσ2::AbstractArray{ET}; - σfac=one(ET)) where ET +function neg_logden_indep_normal(obs::AbstractArray{TO,D}, μ::AbstractArray{TP,D}, logσ2::AbstractArray{TO}; + i_finobs = .! isnan.(obs), + σfac=one(TO)) where {TO, TP, D} # log of independent Normal distributions # estimate independent uncertainty of each θM, rather than full covariance #nlogL = sum(σfac .* log.(σs) .+ 1 / 2 .* abs2.((obs .- μ) ./ σs)) @@ -31,19 +34,42 @@ function neg_logden_indep_normal(obs::AbstractArray, μ::AbstractArray, logσ2:: # σfac .* logσ2 .+ abs2.(obs_data .- μ_data) .* exp.(.-logσ2)) / convert(eltype(μ),2) # return (nlogL) # - i_finobs = .! isnan.(obs) - obs_data = CA.getdata(obs)[i_finobs] - μ_data = CA.getdata(μ)[i_finobs] - logσ2_fin = logσ2[i_finobs] - nlogL = sum( # observations might by NaN for missing - σfac .* logσ2_fin .+ abs2.(obs_data .- μ_data) .* exp.(.-logσ2_fin)) / convert(eltype(μ),2) - #Main.@infiltrate_main + # obs_data = CA.getdata(obs)[i_finobs] + # μ_data = CA.getdata(μ)[i_finobs] + # logσ2_fin = logσ2[i_finobs] + # # nlogL = sum( # observations might by NaN for missing + # # σfac .* logσ2_fin .+ abs2.(obs_data .- μ_data) .* exp.(.-logσ2_fin)) / convert(eltype(μ),2) + # nlogL = + # (σfac .* logσ2_fin .+ abs2.(obs_data .- μ_data) .* exp.(.-logσ2_fin)) ./ eltype(μ)(2) + nlogL = + i_finobs .* (σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) ./ eltype(μ)(2) + nlogL_sites = if D == 1 + sum(nlogL) # provided a single-site vector + elseif D == 2 + vec(sum(nlogL;dims=1)) + else + error("need to provide a matrix with observations with one column per site, " * + "but got array of dimension $D") + end + #nlogL_sites = colsum_finite_obs(nlogL, obs) + # if !all(isfinite.(nlogL_sites)) + # @show nlogL_sites, μ + # @warn("encountered non-finite loglikelihood in neg_logden_indep_normal") + # Main.@infiltrate_main + # end + return (nlogL_sites) + # function fcol(logσ2, obs, μ) + # nlogL = sum(isfinite(obs[i]) ? X[i] : zero(eltype(X)) # observations might by NaN for missing + # σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) / convert(eltype(μ),2) + # end + # map(fcol, eachdol(obs), eachcol(μ), eachcol(logσ2)) + - return (nlogL) end function neg_logden_indep_normal(obs::AbstractGPUArray, μ::AbstractGPUArray, logσ2::AbstractGPUArray{ET}; σfac=one(ET)) where ET + error("update neg_logden_indep_normal for GPU to each solsum rather than total sum") #cannot use NaNMath.sum on gpu, allocate vectors of non-NAN i_finobs = .! isnan.(obs) obs_data = CA.getdata(obs)[i_finobs] diff --git a/src/util.jl b/src/util.jl index 1022de1..bff0f47 100644 --- a/src/util.jl +++ b/src/util.jl @@ -118,3 +118,272 @@ function ChainRulesCore.rrule(::typeof(insert_zeros), v::AbstractVector, positio end return y, pullback end + +""" + cat_namedtuple_lastdim(nt_agg, nt; along) + +Reducing function that takes two NamedTuple objects of the same type and +concatenates each component along specified dimension. + +Optionally, the dimension at which to concatenate can be specified in +the `along` NamedTuple argument. +""" +function cat_namedtuple_lastdim(nt_agg::NamedTuple, nt::NamedTuple; along = map(ndims, nt)::NamedTuple) + NamedTuple( map(nt, keys(nt)) do comp, key + key => cat(nt_agg[key], comp; dims = along[key]) + end) +end + +""" + index_at_dim(x::AbstractArray{T, N}, i::AbstractVector{Int}; dim::Int) where {T, N} + +Index into array `x` along dimension `dim` using indices `i`, while selecting +all elements along all other dimensions. + +# Arguments +- `x::AbstractArray{T, N}`: Input array of type `T` and `N` dimensions. +- `i::AbstractVector{Int}`: Vector of indices to select along dimension `dim`. +- `dim::Int`: The dimension along which to index. + +# Returns +- An array of the same type as `x` with the same number of dimensions, where + the size along `dim` is `length(i)` and all other dimensions are unchanged. + +# Examples +```julia +x = reshape(1:24, 3, 2, 4) + +# Index along dimension 1 +index_at_dim(x, [1]; dim=1) # 1×2×4 array + +# Index along dimension 2 +index_at_dim(x, [1, 3]; dim=1) # 2×2×4 array + +# Index along dimension 3 +index_at_dim(x, [2, 4]; dim=3) # 3×2×2 array +``` +""" +function index_at_dim(x::AbstractArray{T, N}, i::AbstractVector{Int}; dim::Int) where {T, N} + colons = ntuple(d -> d == dim ? i : Colon(), N) + return x[colons...] +end + + + +using LinearAlgebra + +""" + log_density_mvn_cholesky(U,x) + +Compute the log-density of a zero-mean multivariate normal distribution +with covariance matrix C = U' * U, where U is the upper Cholesky factor. + +Arguments: +- U: upper triangular Cholesky factor of the covariance matrix (n × n) +- x: vector of length n (the sample) + +Returns: +- log p(x) ∈ ℝ: log-density at x +""" +function log_density_mvn_cholesky(U::AbstractMatrix{T}, x::AbstractVector{T}) where T + n = length(x) + # Solve L * y = x for y (forward substitution) + y = U' \ x # Efficient triangular solve + # Compute ||y||^2 = y' * y + quad_form = dot(y, y) # or: sum(abs2, y) + # Compute sum of log(diagonals) of U → this is log(sqrt(det(C))) + if any(diag(U) .< 0) + @info("log_density_mvn_cholesky: encountered diag(U) components smaller than zero: $(diag(U))") + # ignore_derivatives() do + # Main.@infiltrate_main + # end + end + log_det_C_half = sum(log, diag(U)) # = 0.5 * log|C| + # Full log-density formula + log2π = T(1.8378770664093453) #log(2π) + log_density = -T(0.5) * (quad_form + T(2) * log_det_C_half + n * log2π) + return log_density +end + +# Prompt: Write a Julia function replace_values(x::Matrix, i_sites::Vector{Int}, pos::Vector{Int}, y::Matrix) that returns a new matrix where x[i_sites, pos] is replaced by y and all other values are unchanged. The function must be fully non-mutating and compatible with Zygote automatic differentiation. Use one-hot projection matrices P_row and P_col to scatter y into the full matrix space via P_row * y * P_col', and blend with x using a binary mask derived from the outer product of row and column indicator vectors. It should use matrix comprehensions to form P_row and P_col. +""" + replace_values_matrix(x::Matrix, i_sites::Vector{<:Integer}, pos::Vector{<:Integer}, y::Matrix) + +Return a new matrix where the submatrix at positions `x[i_sites, pos]` is replaced by `y`, +while all other values remain unchanged. + +This function performs a non-mutating replacement using one-hot projection matrices and +a binary mask derived from the outer product of indicator vectors. It is designed to be +compatible with automatic differentiation frameworks like Zygote. + +### Parameters +- `x`: The input matrix of size `(m, n)` to be modified. +- `i_sites`: A vector of row indices (1-based) specifying which rows to replace. +- `pos`: A vector of column indices (1-based) specifying which columns to replace. +- `y`: The replacement matrix of size `(length(i_sites), length(pos))`. + +### Returns +- A new matrix of the same size as `x`, where `x[i_sites, pos]` is replaced by `y`. + +### Details +- The function constructs one-hot projection matrices `P_row` (size `m × length(i_sites)`) + and `P_col` (size `n × length(pos)`) using matrix comprehensions. +- The replacement values are scattered into the full matrix space via `P_row * y * P_col'`. +- A binary mask is created using the outer product of indicator vectors for `i_sites` and `pos`. +- The result is computed as `x .* (1 - mask) + (P_row * y * P_col') .* mask`, blending the original + matrix with the scattered replacement values. + +### Example +```julia +x = [1 2 3; 4 5 6; 7 8 9] +i_sites = [1, 3] +pos = [2, 3] +y = [10 11; 12 13] + +result = replace_values(x, i_sites, pos, y) +# result = [1 10 11; 4 5 6; 7 12 13] +""" +function replace_values_matrix(x::AbstractMatrix{T}, i_sites::AbstractVector{<:Integer}, pos::AbstractVector{<:Integer}, y::AbstractMatrix{T}) where T + # Precompute projection matrices (one-hot style) + P_row = [i == k for i in 1:size(x, 1), k in i_sites] # n_rows × length(i_sites) + P_col = [j == k for j in 1:size(x, 2), k in pos] # n_cols × length(pos) + # Project y into full matrix space: P_row * y * P_col' + y_full = P_row * y * P_col' + # Compute mask as outer product of indicator vectors + # row_mask = [i in i_sites for i in 1:size(x, 1)] + # col_mask = [j in pos for j in 1:size(x, 2)] + # replace_mask0 = row_mask .* col_mask' + replace_mask = sum(P_row, dims=2) .* sum(P_col, dims=2)' + res = (1 .- replace_mask) .* x .+ replace_mask .* y_full + res +end + +""" + replace_columns_matrix(x::Matrix, col_indices::Vector{Int}, y::Matrix) + +Return a new matrix where the specified columns of `x` are replaced by the columns of `y`. + +This function performs a non-mutating column replacement using one-hot projection matrices +and is compatible with automatic differentiation frameworks like Zygote. + +### Parameters +- `x`: The input matrix of size `(m, n)` to be modified. +- `col_indices`: A vector of column indices (1-based) specifying which columns to replace. +- `y`: The replacement matrix of size `(m, length(col_indices))`. + +### Returns +- A new matrix of the same size as `x`, where the columns at positions `col_indices` are replaced by the corresponding columns of `y`. + +### Details +- The function constructs a one-hot projection matrix `P_col` (size `n × length(col_indices)`) + using matrix comprehensions, where each column corresponds to a target column index. +- The replacement values are scattered into the full matrix space via `x * (I - P_col * P_col') + y * P_col'`. +- The operation is differentiable with respect to all inputs. + +### Example +```julia +x = [1 2 3; 4 5 6; 7 8 9] +col_indices = [1, 3] +y = [10 11; 12 13; 14 15] + +result = replace_columns_matrix(x, col_indices, y) +#result = HVI.replace_columns_matrix(x, col_indices, y) +# result == [10 2 11; 12 5 13; 14 8 15] +``` + +### Notes +- All column indices must be valid (1-based). +- The function is fully non-mutating and compatible with Zygote for automatic differentiation. +- The operation is differentiable with respect to all inputs. +""" +function replace_columns_matrix(x::AbstractMatrix{T}, col_indices::AbstractVector{<:Integer}, y::AbstractMatrix{T}) where T + # Get dimensions + m, n = size(x) + p = length(col_indices) + + # Validate inputs + # @assert p == size(y, 2) "Number of columns in y must match length of col_indices" + # @assert m == size(y, 1) "Number of rows in y must match number of rows in x" + # @assert all(1 .<= col_indices .<= n) "col_indices must be valid column indices" + + # Create one-hot projection matrix using matrix comprehension + # P_col: n × p matrix where each column is a one-hot vector for col_indices + P_col = ChainRulesCore.@ignore_derivatives [ + j_it == j_col ? one(T) : zero(T) for j_it in 1:n, j_col in col_indices] + + # Replace columns: keep original columns (I - P_col * P_col') and replace with y * P_col' + # This is equivalent to: x * (I - P_col * P_col') + y * P_col' + result = x * (LinearAlgebra.I - P_col * P_col') .+ y * P_col' + return result +end + +""" + generate_repeated_integers(n_MC::Int, n_sample_ranef::Int) -> Vector{Int} + +Generate a vector of increasing integers where each integer is repeated `n_sample_ranef` +times, except possibly the last one, such that the total length of the vector is exactly +`n_MC`. + +# Arguments +- `n_MC::Int`: The total length of the output vector. Must be a positive integer. +- `n_sample_ranef::Int`: The number of times each integer is repeated. Must be a positive integer. + +# Returns +- `Vector{Int}`: A vector of length `n_MC` where each integer `i` appears `n_sample_ranef` + times, except for the last integer which appears `mod(n_MC, n_sample_ranef)` times if + `n_MC` is not a multiple of `n_sample_ranef`, and `n_sample_ranef` times otherwise. + +# Examples +```julia-repl +julia> generate_repeated_integers(8, 5) +8-element Vector{Int64}: + 1, 1, 1, 1, 1, 2, 2, 2 + ``` + """ +function generate_repeated_integers(n_MC::Integer, n_sample_ranef::Integer) + # Calculate how many complete groups we need + n_groups = ceil(Int, n_MC / n_sample_ranef) + # Generate the full repeated vector + full_vec = repeat(1:n_groups, inner = n_sample_ranef) + # Trim to exactly n_MC elements + return full_vec[1:n_MC] +end + + +""" + colsum_finite_obs(X::AbstractMatrix, obs::AbstractMatrix) + +Return the column sums of `X`, using `obs` as a finiteness mask. +Equivalently, entries of `X` are treated as zero whenever the corresponding +entry of `obs` is not finite. + +Important: this function does **not** check whether entries of `X` are finite. +If `X[i, j]` is `NaN`, `Inf`, or `-Inf` and `obs[i, j]` is finite, then that +value is included in the sum. + +`X` and `obs` must have the same axes. For ordinary Julia `Matrix` objects, +this means they must have the same size. +""" +function colsum_finite_obs(X::AbstractMatrix, obs::AbstractMatrix) + # The implementation avoids constructing a cleaned copy of `X` and does not use + # explicit mutation. + axes(X) == axes(obs) || + throw(DimensionMismatch("X and obs must have the same axes")) + + [ + sum( + isfinite(obs[i, j]) ? X[i, j] : zero(eltype(X)) + for i in axes(X, 1); + init = zero(eltype(X)) + ) + for j in axes(X, 2) + ] +end + +function colsum_finite_obs(X::AbstractVector, obs::AbstractVector) + sum( + isfinite(obs[i]) ? X[i] : zero(eltype(X)) + for i in axes(X, 1); + init = zero(eltype(X)) + ) +end diff --git a/test/Project.toml b/test/Project.toml index a9f2606..8049cbb 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,6 +3,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" diff --git a/test/runtests.jl b/test/runtests.jl index 2495c47..3be95fc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,12 @@ using Test, SafeTestsets +() -> begin + #@usingany ReferenceRevision + refmain = open_process(rev = "main") + refmain = open_process(rev = "main", instantiate = true) + refmain.eval(:(using CommonSolve)) + refmain = open_process(rev = "main", instantiate = true) + close(refmain) +end const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml @time begin @@ -15,6 +23,8 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml @time @safetestset "test_util_gpu" include("test_util_gpu.jl") #@safetestset "test" include("test/test_ComponentArrayInterpreter.jl") @time @safetestset "test_ComponentArrayInterpreter" include("test_ComponentArrayInterpreter.jl") + #@safetestset "test" include("test/test_RandomEffects.jl") + @time @safetestset "test_PBMApplicator" include("test_RandomEffects.jl") #@safetestset "test" include("test/test_PBMApplicator.jl") @time @safetestset "test_PBMApplicator" include("test_PBMApplicator.jl") #@safetestset "test" include("test/test_ModelApplicator.jl") diff --git a/test/test_HybridProblem.jl b/test/test_HybridProblem.jl index 637cea7..a77ecfc 100644 --- a/test/test_HybridProblem.jl +++ b/test/test_HybridProblem.jl @@ -6,6 +6,7 @@ using StableRNGs using Random using Statistics using ComponentArrays: ComponentArrays as CA +using StaticArrays: StaticArrays as SA using Bijectors using DistributionFits using StatsFuns: logistic @@ -15,6 +16,7 @@ using MLUtils import Zygote using OptimizationOptimisers +import CommonSolve: solve using MLDataDevices using Suppressor @@ -27,8 +29,11 @@ cdev = cpu_device() #scenario = Val((:default, )) #scenario = Val((:MeanHVIApproxMat,)) #scenario = Val((:covarK2,)) +#scenario = Val((:covarK2,:ranef)) #scen = CP._val_value(scenario) #scenario = Val((:clustered_sites, )) +#scenario = Val((:scalingall, )) +#scenario = Val((:scalingall, :ranef)) function construct_problem(; scenario::Val{scen}) where scen FT = Float32 @@ -53,12 +58,14 @@ function construct_problem(; scenario::Val{scen}) where scen n_site_test = 60 # dependency on DeoubleMMCase -> take care of changes in covariates (; xM, θP_true, θMs_true, xP, y_true, y_o, y_unc - ) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); n_site_test,scenario) - i_test = n_site .+ (1:n_site_test) - test_data = (; xM = xM[:, i_test], xP = xP[:, i_test], y_true = y_true[:, i_test], - y_o = y_o[:, i_test], y_unc = y_unc[:, i_test]) + ) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario) + i_sites_test = sample(rng, 1:n_site, n_site_test, replace=false) + test_data = (; xM = xM[:, i_sites_test], xP = xP[:, i_sites_test], + y_o = y_o[:, i_sites_test], y_unc = y_unc[:, i_sites_test], + itrain_sites = zeros(length(i_sites_test))) approx = if (:scalingall ∈ scen) - MeanHVIApproximationMat([length(θM)]) + block_ends = [length(θM)] + MeanScalingHVIApproximation(block_ends, FT(2) .* log.(FT(0.1) .* θM[block_ends])) elseif (:MeanHVIApproxBlocks ∈ scen) MeanHVIApproximationMat() else @@ -66,7 +73,7 @@ function construct_problem(; scenario::Val{scen}) where scen end n_covar = size(xM,1) n_input = (:covarK2 ∈ scen) ? n_covar +1 : n_covar - n_out = get_numberof_MLinputs(approx, θM) + n_out = get_numberof_MLinputs(approx, θM) g_chain = SimpleChain( static(n_input), # input dimension (optional) # dense layer with bias that maps to 8 outputs and applies `tanh` activation @@ -78,16 +85,17 @@ function construct_problem(; scenario::Val{scen}) where scen # g, ϕg = construct_SimpleChainsApplicator(g_chain) # py = neg_logden_indep_normal - i_sites = 1:n_site - # get_train_loader = let xM = xM, xP = xP, y_o = y_o, y_unc = y_unc, i_sites = i_sites + itrain_sites = 1:n_site + # get_train_loader = let xM = xM, xP = xP, y_o = y_o, y_unc = y_unc, itrain_sites = itrain_sites # function inner_get_train_loader(; n_batch, kwargs...) - # MLUtils.DataLoader((xM, xP, y_o, y_unc, i_sites), batchsize=n_batch, partial=false) + # MLUtils.DataLoader((xM, xP, y_o, y_unc, itrain_sites), batchsize=n_batch, partial=false) # end # end - i_train = 1:n_site + i_train = setdiff(1:n_site, i_sites_test) + @assert sort(vcat(i_sites_test, i_train)) == 1:n_site train_dataloader = MLUtils.DataLoader( (CA.getdata(xM[:,i_train]), CA.getdata(xP[:,i_train]), y_o[:,i_train], - y_unc[:,i_train], i_sites[i_train]), batchsize=n_batch, partial=false) + y_unc[:,i_train], itrain_sites[i_train]), batchsize=n_batch, partial=false) θall = vcat(θP, θM) priors_dict = Dict{Symbol, Distribution}( keys(θall) .=> fit.(LogNormal, θall, QuantilePoint.(θall .* 3, 0.95))) @@ -95,9 +103,9 @@ function construct_problem(; scenario::Val{scen}) where scen # scale (0,1) outputs MLmodel to normal distribution fitted to priors translated to ζ priorsM = Tuple(priors_dict[k] for k in keys(θM)) lowers, uppers = get_quantile_transformed(priorsM, transM) - + range_scaled = 1:length(lowers) # do only scale means, but not the uncertainty factor app, ϕg0 = construct_ChainsApplicator(rng, g_chain, FT) - g_chain_scaled = NormalScalingModelApplicator(app, lowers, uppers, FT) + g_chain_scaled = NormalScalingModelApplicator(app, lowers, uppers, FT; range_scaled) #g_chain_scaled = app pbm_covars = (:covarK2 ∈ scen) ? (:K2,) : () f_batch = PBMSiteApplicator( @@ -105,10 +113,15 @@ function construct_problem(; scenario::Val{scen}) where scen xPvec=xP[:,1]) (; ϕqc, approx) = init_hybrid_ϕunc(approx, cor_ends, zero(FT); θM, transM, n_site) ϕq = CP.update_μP_by_θP(ϕqc, θP, transP) + ranef = if (:ranef ∈ scen) + RandomEffects((:r1,)) + else + NullRandomEffects() + end HybridProblem(θM, ϕq, g_chain_scaled, ϕg0, f_batch, priors_dict, py, transM, transP, train_dataloader, test_data, n_site, n_batch; - cor_ends, pbm_covars, approx, + cor_ends, pbm_covars, approx, ranef = ranef #ϕunc0, ) end @@ -159,28 +172,36 @@ test_without_flux = (scenario) -> begin n_sites_cluster, clusters = CP.get_clusters(n_site; scenario) frac_cluster_all = 1 ./ n_sites_cluster[clusters] train_loader = get_hybridproblem_train_dataloader(prob; scenario) - (xM, xP, y_o, y_unc, i_sites) = first(train_loader) + (xM, xP, y_o, y_unc, itrain_sites) = first(train_loader) f = get_hybridproblem_PBmodel(prob; scenario) py = get_hybridproblem_neg_logden_obs(prob; scenario) - par_templates = get_hybridproblem_par_templates(prob; scenario) #f(par_templates.θP, hcat(par_templates.θM, par_templates.θM), xP[1:2]) (; transM, transP) = get_hybridproblem_transforms(prob; scenario) pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) - intϕ = ComponentArrayInterpreter(CA.ComponentVector( - ϕg=1:length(ϕg0), ϕP=par_templates.θP)) priors = get_hybridproblem_priors(prob; scenario) - priorsP = Tuple(priors[k] for k in keys(par_templates.θP)) - priorsM = Tuple(priors[k] for k in keys(par_templates.θM)) + priorsP = Tuple(priors[k] for k in keys(pt.θP)) + priorsM = Tuple(priors[k] for k in keys(pt.θM)) # slightly disturb θP_true - p = p0 = vcat(ϕg0, par_templates.θP .* convert(eltype(ϕg0), 0.8)) - intθP = ComponentArrayInterpreter(pt.θP) - intθMs = ComponentArrayInterpreter((n_batch,), pt.θM) + + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(pt.θM))) + ϕq_ranef = setup_ϕq_ranef(ranef) + pc = CA.ComponentVector( + ϕg = ϕg0, + ϕq = CA.ComponentVector( + μP = pt.θP .* convert(eltype(ϕg0), 0.8), + ranef = ϕq_ranef + ) + ) + p = p0 = CA.getdata(pc) + intϕ = ComponentArrayInterpreter(pc) # Pass the site-data for the batches as separate vectors wrapped in a tuple loss_gf = get_loss_gf(g, transM, transP, f, py, intϕ; par_templates = pt, pbm_covars, n_site_batch = n_batch, priorsP, priorsM, - intθMs, intθP, frac_cluster_all, + frac_cluster_all, ranef, ) (_xM, _xP, _y_o, _y_unc, _i_sites) = first(train_loader) #l1 = loss_gf(p0, _xM, _xP, _y_o, _y_unc, _i_sites; is_testmode = false) @@ -211,7 +232,7 @@ test_without_flux = (scenario) -> begin pbm_covars, n_site_batch = n_site) l1, y_pred, θMs_pred, θP, nLy, neg_log_prior = loss_gf_sites( res.u, train_loader.data...) - @test isapprox(par_templates.θP, transP(intϕ(res.u).ϕP), rtol=0.5) + @test isapprox(pt.θP, transP(intϕ(res.u).ϕP), rtol=0.5) end end end @@ -220,7 +241,9 @@ end #scenario=Val((:default,)) test_without_flux(Val((:default,))) test_without_flux(Val((:covarK2,))) +test_without_flux(Val((:ranef, :covarK2,))) test_without_flux(Val((:clustered_sites,))) +test_without_flux(Val((:ranef, :clustered_sites))) import CUDA, cuDNN using GPUArraysCore @@ -237,22 +260,33 @@ test_with_flux = (scenario) -> begin rng = StableRNG(111) solver = HybridPointSolver(; alg=Adam(0.02)) (; ϕ, resopt, probo) = solve(prob, solver; scenario, rng, - #callback = callback_loss(100), maxiters = 1200 - #maxiters = 1200 - #maxiters = 20 - #maxiters=200, + #callback = callback_loss(100), epochs = 2, + #epochs_callback = 1, # print every epoch + epochs_callback = 0, # do not evaluate test and do not print gdevs = (; gdev_M=identity, gdev_P=identity), #gpu_handler = NullGPUDataHandler is_inferred = Val(true), ) + () -> begin + ref_solver = refmain.HybridPointSolver(; alg=Adam(0.02)) + ref_ans_solve = refmain.solve(prob, ref_solver; scenario, rng, + #callback = callback_loss(100), + epochs = 2, + #epochs_callback = 1, # print every epoch + epochs_callback = 0, # do not evaluate test and do not print + gdevs = (; gdev_M=identity, gdev_P=identity), + #gpu_handler = NullGPUDataHandler + is_inferred = Val(true), + ) + end (; θP) = get_hybridproblem_par_templates(prob; scenario) θPo = (() -> begin (; θP) = get_hybridproblem_par_templates(probo; scenario); θP end)() @test θPo.r0 < 1.5 * θP.r0 - @test ϕ.ϕP.K2 < 1.5 * log(θP.K2) + @test ϕ.ϕq.μP.K2 < 1.5 * log(θP.K2) (;y_pred, θMs_tr, θP) = tmp = predict_point_hvi(rng, probo; scenario) _,_,y_obs,_ = get_hybridproblem_train_dataloader(prob; scenario).data @test size(y_pred) == size(y_obs) @@ -265,17 +299,34 @@ test_with_flux = (scenario) -> begin #callback = callback_loss(100), maxiters = 1200, #maxiters = 20 # too small so that it yields error #maxiters=37, # still complains "need to specify maxiters or epochs" - epochs = 1, + epochs = 2, + epochs_callback = 1, # print every epoch + #epochs_callback = 0, # not progress output θmean_quant = 0.01, # test constraining mean to initial prediction gdevs = (; gdev_M=identity, gdev_P=identity), is_inferred = Val(true), ) θPt = get_hybridproblem_par_templates(prob; scenario).θP + ϕ.ϕq.ranef @test θP.r0 < 1.5 * θPt.r0 @test exp(ϕ.ϕq.μP.K2) == θP.K2 < 1.5 * θP.K2 n_sample_pred = 12 - (; y, addq, θsP, θsMs_tr, entropy_ζ) = predict_hvi(rng, probo; scenario, n_sample_pred); - _,_,y_obs,_ = get_hybridproblem_train_dataloader(prob; scenario).data + (; y, addq, θsP, θsMs_tr, entropy_ζ) = predict_hvi(rng, probo; + scenario, n_sample_pred, + ); + _,_,y_obs,_ = get_hybridproblem_train_dataloader(probo; scenario).data + @test size(y) == (size(y_obs)..., n_sample_pred) + yc = cdev(y) + _ = map(eachslice(yc; dims = 3)) do ycs + @test all(isfinite.(ycs[isfinite.(y_obs)])) + end + # test prediction without known random effects + testdata = get_hybridproblem_test_data(probo; scenario) + (; y, addq, θsP, θsMs_tr, entropy_ζ) = predict_hvi(rng, probo; + xM = testdata.xM, xP = testdata.xP, + scenario, n_sample_pred, + ); + _,_,y_obs,_ = get_hybridproblem_test_data(prob; scenario) @test size(y) == (size(y_obs)..., n_sample_pred) yc = cdev(y) _ = map(eachslice(yc; dims = 3)) do ycs @@ -286,6 +337,9 @@ end # test_with flux test_with_flux(Val((:default,))) test_with_flux(Val((:covarK2,))) +test_with_flux(Val((:ranef, :covarK2,))) +test_without_flux(Val((:ranef, :clustered_sites))) + #scenario = Val((:default,:useSitePBM)) test_with_flux_gpu = (scenario) -> begin @@ -304,6 +358,7 @@ test_with_flux_gpu = (scenario) -> begin #maxiters = 37, # smallest value by trial and error #maxiters = 20 # too small so that it yields error epochs = 2, + epochs_callback = 0, # not progress output θmean_quant = 0.01, # test constraining mean to initial prediction is_inferred = Val(true), gdevs = (; gdev_M=gpu_device(), gdev_P=identity),); @@ -315,6 +370,7 @@ test_with_flux_gpu = (scenario) -> begin (; probo, ϕ, resopt) = solve(prob, solver; scenario = scenf, #maxiters = 37, epochs = 2, + epochs_callback = 0, # not progress output gdevs = (; gdev_M=gpu_device(), gdev_P=identity), is_inferred = Val(true), ); @@ -332,7 +388,8 @@ test_with_flux_gpu = (scenario) -> begin test_correlation = () -> begin n_epoch = 20 # requires (; ϕ, resopt, probo) = solve(prob, solver; scenario = scenf, - maxiters = n_batches_in_epoch * n_epoch, + epochs = n_epoch, + epochs_callback = 2, # not progress output gdevs = (; gdev_M=gpu_device(), gdev_P=identity), callback = callback_loss(n_batches_in_epoch*5) ); @@ -361,9 +418,9 @@ test_with_flux_gpu = (scenario) -> begin cr = cor(CA.getdata(residθs')) pos_P = get_positions(ComponentArrayInterpreter(θs[:P,1])) - i_sites = [1,2,3] + itrain_sites = [1,2,3] #ax = map(x -> axes(x,1), get_hybridproblem_par_templates(probo; scenario = scenf)) - is = vcat(pos.P, vec(pos.Ms[i_sites,:])) + is = vcat(pos.P, vec(pos.Ms[itrain_sites,:])) cr[is,is] end end; @@ -387,6 +444,7 @@ test_with_flux_gpu = (scenario) -> begin #maxiters = 37, # smallest value by trial and error #maxiters = 20, # too small so that it yields error epochs = 1, + epochs_callback = 0, # not progress output #θmean_quant = 0.01, # TODO make possible on gpu gdevs = (; gdev_M=gpu_device(), gdev_P=gpu_device()), is_inferred = Val(true), @@ -404,7 +462,8 @@ end # test_with flux #test_with_flux_gpu(Val((:MeanHVIApproxBlocks,))) # do not test any more, its slower #scenario = Val(()) -test_with_flux_gpu(Val((:default,))) -test_with_flux_gpu(Val((:covarK2,))) +#test_with_flux_gpu(Val((:default,))) +test_with_flux_gpu(Val((:scalingall))) +test_with_flux_gpu(Val((:covarK2))) test_with_flux_gpu(Val((:useSitePBM,))) diff --git a/test/test_RandomEffects.jl b/test/test_RandomEffects.jl new file mode 100644 index 0000000..af9e2c3 --- /dev/null +++ b/test/test_RandomEffects.jl @@ -0,0 +1,137 @@ +using Test +using HybridVariationalInference +using HybridVariationalInference: HybridVariationalInference as HVI +using HybridVariationalInference: HybridVariationalInference as CP +using LinearAlgebra +using StaticArrays: StaticArrays as SA +using ComponentArrays: ComponentArrays as CA +using Distributions + +@testset "CVPrior_LKJ_Cauchy" begin + prior = CVPrior_LKJ_Cauchy(3) + cm = Hermitian([1.0 0.5 0.3; 0.5 1.0 0.4; 0.3 0.4 1.0] .* 2.0) + logpdf_prior = logpdf(prior, cm) + @test isfinite(logpdf_prior) +end + +@testset "CVPrior_LKJ_Cauchy one column" begin + prior = CVPrior_LKJ_Cauchy(1) + cm = Hermitian(reshape([3.0],1,1)) + logpdf_prior = logpdf(prior, cm) + @test isfinite(logpdf_prior) +end + +@testset "CVPrior_LKJ_Cauchy Float32" begin + FT = Float32 + prior = CVPrior_LKJ_Cauchy(3, one(FT)) + @test partype(prior.dLKJ) == FT + cm = Hermitian(FT[1.0 0.5 0.3; 0.5 1.0 0.4; 0.3 0.4 1.0] .* FT(2.0)) + logpdf_prior = logpdf(prior, cm) + @test isfinite(logpdf_prior) + @test logpdf_prior isa FT +end + +@testset "CVPrior_LKJ_Cauchy convert_prior" begin + prior64 = CVPrior_LKJ_Cauchy(1) + prior32 = HVI.convert_prior(prior64, one(Float32)) + @test logpdf(prior32, diagm([1.0f0])) isa Float32 +end + +() -> begin + d = Cauchy(0.0, 0.5) + #using StatsPlots + plot(d, xlim=(0,12)) +end + +@testset "NullRandomEffects" begin + nre = NullRandomEffects() + n_site = 200 + ranef = HVI.get_ranef_computer(nre, (:_,), n_site) + ϕq_ranef = setup_ϕq_ranef(ranef) + @test size(ϕq_ranef.β) == (0, n_site) + @test eltype(ϕq_ranef) == Float64 + μ = randn(0,n_site) + μ2 = add_ranef(ranef, μ, ϕq_ranef, 1:n_site) + @test μ2 == μ + # + ranef = HVI.get_ranef_computer(nre, (:_,), n_site, one(Float32)) + ϕq_ranef = setup_ϕq_ranef(ranef) + @test ϕq_ranef.β isa AbstractMatrix{Float32} + μ = randn(Float32, 0,n_site) + μ2 = add_ranef(ranef, μ, ϕq_ranef, 1:n_site) + @test μ2 == μ + # + β = HVI.sample_ranef(ranef, ϕq_ranef, 4, 3) + @test eltype(β) == Float32 + @test β == zeros(1,4,3) # one parameter in construction +end + + +@testset "RandomEffects" begin + η = 3.0 + d = LKJCholesky(2, η) + rL = rand(d) + L = rL.L + corrm = L * L' + # Array(rL) already gives corrm + res = randn(40000,2) + tau = [1.2,2.2] # sqrt of main diagonal + cm = diagm(tau) * corrm * diagm(tau) + ranef = (diagm(tau) * L * res') + #ranef = res * Array(L)' * diagm(tau) + #cov(ranef') # should roughly match cm + #cor(ranef') + # + θM = CA.ComponentVector(a=1.0, b=2.0, c=3.0) + par_ranef = (:c, :a) # positions (3,1) + pos_ranef = [findfirst(==(s),keys(θM)) for s in par_ranef] + prior_Σ = CVPrior_LKJ_Cauchy(length(par_ranef), η=η) + re0 = RandomEffects(par_ranef; η=η) + @test re0.prior_Σ == prior_Σ + n_site = 200 + re = HVI.get_ranef_computer(re0, keys(θM), n_site) + ϕq_ranef = setup_ϕq_ranef(re) + @test size(ϕq_ranef[Val(:β)]) == (n_site, length(par_ranef)) + U = HVI.transformU_cholesky1(ϕq_ranef[Val(:coef_U)]); U' * U + l1 = compute_nLranef(re, ϕq_ranef) + ϕq_ranef2 = CA.ComponentVector(ϕq_ranef; σ = ϕq_ranef.σ .+ eltype(ϕq_ranef)(0.02)) + l2 = compute_nLranef(re, ϕq_ranef2) + @test l1 < l2 + () -> begin + # using Zygote + # @usingany FiniteDiff + Zygote.gradient(x -> compute_nLranef(re, x), ϕq_ranef2) + end + # + re32 = HVI.get_ranef_computer(re0, keys(θM), n_site, one(Float32)) + ϕq_ranef = setup_ϕq_ranef(re32) + @test eltype(ϕq_ranef) == Float32 + ϕq_ranef.coef_U .= 0.1f0 + ϕq_ranef.σ .= [0.2, 0.5] + U_exp = HVI.get_choleskyΣ_ranef(re32, ϕq_ranef) + Σ_exp = U_exp' * U_exp + P_col = [j == k for j in 1:3, k in pos_ranef] + Σ_full = P_col * Σ_exp * P_col' + β = HVI.sample_ranef(re32, ϕq_ranef, n_site, 60)[pos_ranef,:,:] + @test eltype(β) == Float32 + @test size(β) == (2, n_site, 60) + βstacked = reshape(β, size(β,1),:) + @test cov(βstacked') ≈ Σ_exp atol=0.05 + ϕq_ranef.β = β[:,:,1]' + # + @test eltype(ϕq_ranef) == Float32 + l1 = compute_nLranef(re32, ϕq_ranef) + @test l1 isa Float32 + itrain_sites = 1:20 + μ = randn(Float32, length(θM), length(itrain_sites)) + μ_updated = add_ranef(re32, μ, ϕq_ranef, itrain_sites) + @test eltype(μ_updated) == Float32 + @test all((μ_updated .- μ)[2,:] .== 0) + @test all((μ_updated .- μ)[[3,1],:] .≈ ϕq_ranef.β[itrain_sites,:]') + # + # test extending cholesky fator to all parameter + U_par = HVI.get_choleskyΣ_par(re32, ϕq_ranef) + Σ = U_par' * U_par + @test Σ ≈ Σ_full +end + diff --git a/test/test_WeightedDataLoader.jl b/test/test_WeightedDataLoader.jl new file mode 100644 index 0000000..4fc7fc8 --- /dev/null +++ b/test/test_WeightedDataLoader.jl @@ -0,0 +1,214 @@ +using Test +using HybridVariationalInference +using HybridVariationalInference: HybridVariationalInference as HVI +using HybridVariationalInference: WeightedObsView +using SimpleChains +#using Optimisers +#using StatsBase +using MLUtils +using Optimization +using OptimizationOptimisers # provides Optimization-compatible Adam +import Statistics + + + # ---------------------------------------------------------- + # Helper: uniform weights + # ---------------------------------------------------------- + uniform_weights(n) = ones(Float64, n) ./ n + + # ========================================================== + # 1. Returns a native MLUtils.DataLoader + # ========================================================== + @testset "native MLUtils.DataLoader" begin + X = rand(Float32, 4, 100) + w = uniform_weights(100) + loader = WeightedDataLoader(X, w; batchsize=16) + @test loader isa MLUtils.DataLoader + end + + # ========================================================== + # 2. Correct batch shapes and types + # ========================================================== + N = 200 + BS = 32 + + @testset "Float32 matrix" begin + X = rand(Float32, 8, N) + w = uniform_weights(N) + loader = WeightedDataLoader(X, w; batchsize=BS, shuffle=false) + batch = first(loader) + @test batch isa Array{Float32} + @test size(batch, 1) == 8 + @test size(batch, 2) == BS + end + + @testset "Float64 matrix" begin + X = rand(Float64, 5, N) + w = uniform_weights(N) + loader = WeightedDataLoader(X, w; batchsize=BS, shuffle=false) + batch = first(loader) + @test batch isa Array{Float64} + @test size(batch) == (5, BS) + end + + @testset "Int array" begin + X = rand(1:10, 3, N) + w = uniform_weights(N) + loader = WeightedDataLoader(X, w; batchsize=BS, shuffle=false) + batch = first(loader) + @test eltype(batch) <: Integer + @test size(batch) == (3, BS) + end + + @testset "NTuple dataset" begin + X = rand(Float32, 4, N) + Y = rand(Float32, 2, N) + data = (X, Y) + w = uniform_weights(N) + loader = WeightedDataLoader(data, w; batchsize=BS, shuffle=false) + batch = first(loader) + @test batch isa Tuple + @test length(batch) == 2 + @test batch[1] isa Array{Float32} + @test size(batch[1]) == (4, BS) + @test size(batch[2]) == (2, BS) + end + + @testset "NamedTuple dataset" begin + X = rand(Float32, 4, N) + Y = rand(Float32, 1, N) + data = (x=X, y=Y) + w = uniform_weights(N) + loader = WeightedDataLoader(data, w; batchsize=BS, shuffle=false) + batch = first(loader) + @test batch isa NamedTuple + @test haskey(batch, :x) && haskey(batch, :y) + @test size(batch.x) == (4, BS) + @test size(batch.y) == (1, BS) + end + + # ========================================================== + # 3. Minority class oversampling + # ========================================================== + @testset "minority class oversampling" begin + # Build an imbalanced binary dataset: + # class 0: 190 samples, class 1: 10 samples + n_maj = 190 + n_min = 10 + n_total = n_maj + n_min + + labels = vcat(zeros(Int, n_maj), ones(Int, n_min)) + + # Inverse-frequency weights so minority is over-sampled + weights = vcat( + fill(1.0 / n_maj, n_maj), + fill(1.0 / n_min, n_min) + ) + weights ./= sum(weights) # normalise + + loader = WeightedDataLoader(labels, weights; batchsize=1000, shuffle=false) + + # Collect several batches and measure minority fraction + minority_fracs = Float64[] + for (i, batch) in enumerate(loader) + i > 5 && break + push!(minority_fracs, Statistics.mean(batch .== 1)) + end + + avg_frac = Statistics.mean(minority_fracs) + # With balanced weights the minority should appear ~50% of the time + # (within reasonable Monte-Carlo tolerance) + @test avg_frac > 0.30 # much more than the original 5% + @test avg_frac < 0.70 + end + + # ========================================================== + # 4. End-to-end training with Optimization.solve + # ========================================================== + @testset "end-to-end training with Optimization.solve" begin + # Build a tiny regression dataset y = W*x with W known + n_features = 4 + n_out = 2 + n_samples = 256 + BS_train = 64 + + W_true = randn(Float32, n_out, n_features) + X_train = randn(Float32, n_features, n_samples) + Y_train = W_true * X_train + + w = uniform_weights(n_samples) + loader = WeightedDataLoader((X_train, Y_train), w; + batchsize=BS_train, shuffle=true) + + @test loader isa MLUtils.DataLoader + + # Simple linear model: θ is a flattened (n_out × n_features) weight matrix + θ_init = zeros(Float32, n_out * n_features) + + # Loss: mean squared error over a single batch + # `p` receives one batch (a Tuple) from the loader + function loss(p, batch) + Xb, Yb = batch # unpack the tuple + W = reshape(p, n_out, n_features) + Ŷ = W * Xb + return sum(abs2, Ŷ .- Yb) / size(Xb, 2) + end + + # Compute initial loss on the full training set as reference + W0 = reshape(θ_init, n_out, n_features) + loss_init = sum(abs2, W0 * X_train .- Y_train) / n_samples + + # Set up Optimization problem with the DataLoader + opt_func = OptimizationFunction(loss, Optimization.AutoZygote()) + prob = OptimizationProblem(opt_func, θ_init, loader) + + # Train for a few epochs + sol = Optimization.solve(prob, OptimizationOptimisers.Adam(0.05f0); + epochs=30) + + W_sol = reshape(sol.u, n_out, n_features) + loss_final = sum(abs2, W_sol * X_train .- Y_train) / n_samples + + @test loss_final < loss_init # loss must decrease + @test loss_final < 0.5f0 * loss_init # by at least 50 % + end + + # ========================================================== + # Core requirement: loader.data[i] returns the i-th raw array + # ========================================================== + @testset "getindex returns raw array" begin + N = 100 + X1 = rand(Float32, 8, N) + X2 = rand(Float32, 4, N) + X3 = rand(Float32, 2, N) + + loader = WeightedDataLoader((X1, X2, X3), uniform_weights(N); + batchsize=32, partial=false) + + @test loader.data[1] === X1 + @test loader.data[2] === X2 + @test loader.data[3] === X3 + + @test size(loader.data[1]) == size(X1) + @test size(loader.data[2]) == size(X2) + @test size(loader.data[3]) == size(X3) + end + + # ========================================================== + # length returns the number of arrays in the tuple + # ========================================================== + @testset "length returns tuple length" begin + N = 100 + X1 = rand(Float32, 8, N) + X2 = rand(Float32, 4, N) + X3 = rand(Float32, 2, N) + + loader = WeightedDataLoader((X1, X2, X3), uniform_weights(N); + batchsize=32, partial=false) + + @test length(loader.data) == 3 + + loader2 = WeightedDataLoader((X1, X2), uniform_weights(N); + batchsize=32, partial=false) + @test length(loader2.data) == 2 + end diff --git a/test/test_bijectors_utils.jl b/test/test_bijectors_utils.jl index bc5c8de..9982590 100644 --- a/test/test_bijectors_utils.jl +++ b/test/test_bijectors_utils.jl @@ -142,6 +142,14 @@ end Xt2 = @inferred bse(copy(X')') # works also for adjoint @test Xt2 == Xt @inferred bse(X) + # + # with_logabsdet_jacobians + Xt3, logjac = with_logabsdet_jacobian(bse, X) # single value + Xt3, logjacs = CP.with_logabsdet_jacobians(bse, X) + @test Xt3 == Xt + @test sum(logjacs) == logjac + @test size(logjacs) == size(Xt3) # logjac for all components + # if gdev isa MLDataDevices.AbstractGPUDevice Xd = gdev(X) bse(Xd) @@ -160,5 +168,9 @@ end end end +@testset "with_logabsdet_jacobians" begin + +end + diff --git a/test/test_doubleMM.jl b/test/test_doubleMM.jl index 6c147b1..fddd0f7 100644 --- a/test/test_doubleMM.jl +++ b/test/test_doubleMM.jl @@ -1,10 +1,12 @@ using Test using HybridVariationalInference using HybridVariationalInference: HybridVariationalInference as CP +#using HybridVariationalInference: HybridVariationalInference as HVI using StableRNGs using Random using Statistics using ComponentArrays: ComponentArrays as CA +using StaticArrays: StaticArrays as SA using Bijectors using SimpleChains @@ -19,6 +21,8 @@ using CUDA: CUDA using Flux using GPUArraysCore +using Suppressor + gdev = gpu_device() cdev = cpu_device() @@ -191,7 +195,7 @@ end θc = intθ(θ) y = CP.DoubleMM.f_doubleMM_sites(θc, xPM)[1] #y = CP.DoubleMM.f_doubleMM(θ, xPM, θpos) - res = fneglogden(y_o, y, y_unc) + res = sum(fneglogden(y_o, y, y_unc)) # sum across sites res end end @@ -275,80 +279,171 @@ end n_site, n_site_batch = get_hybridproblem_n_site_and_batch(prob; scenario) frac_cluster_all = fill(1, n_site) f = get_hybridproblem_PBmodel(prob; scenario) - f2 = create_nsite_applicator(f, n_site) py = get_hybridproblem_neg_logden_obs(prob; scenario) priors = get_hybridproblem_priors(prob; scenario) priorsP = Tuple(priors[k] for k in keys(par_templates.θP)) priorsM = Tuple(priors[k] for k in keys(par_templates.θM)) + par_ranef = (:r1, :K1) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + #ranef_spec = RandomEffects(par_ranef) + #ranef_spec = NullRandomEffects() + ranef = get_ranef_computer(ranef_spec, keys(pt.θM), n_site, one(eltype(pt.θM))) + ϕq_ranef = setup_ϕq_ranef(ranef) intϕ = ComponentArrayInterpreter(CA.ComponentVector( - ϕg = 1:length(ϕg0), ϕP = par_templates.θP)) + ϕg = 1:length(ϕg0), + ϕq = CA.ComponentVector( + μP=par_templates.θP, + ranef=ϕq_ranef + ) + )) p = p0 = vcat(ϕg0, CP.apply_preserve_axes(inverse(transP), par_templates.θP) .- - convert(eltype(ϕg0), 0.1)) # slightly disturb θP_true + convert(eltype(ϕg0), 0.1), # slightly disturb θP_true + ϕq_ranef, + ) #p = p0 = vcat(ϕg_opt1, par_templates.θP); # almost true # Pass the site-data for the batches as separate vectors wrapped in a tuple # train_loader = MLUtils.DataLoader( # (xM, xP, y_o, y_unc, i_sites), batchsize = n_site_batch) train_loader = get_hybridproblem_train_dataloader(prob; scenario) - @assert train_loader.data == (xM, xP, y_o, y_unc, i_sites) + itrain_sites = train_loader.data[5] + n_site_train = length(itrain_sites) + i_sites_train_all = setdiff(1:size(xM,2),CP.get_i_sites_test(prob; scenario)) + @assert train_loader.data[1:4] == (xM[:,i_sites_train_all], xP[:,i_sites_train_all], y_o[:,i_sites_train_all], y_unc[:,i_sites_train_all]) pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) - intθP = ComponentArrayInterpreter(pt.θP) - intθMs_batch = ComponentArrayInterpreter((n_batch,), pt.θM) - intθMs_site = ComponentArrayInterpreter((n_site,), pt.θM) + #intθP = ComponentArrayInterpreter(pt.θP) + #intθMs_batch = ComponentArrayInterpreter((n_batch,), pt.θM) + #intθMs_site = ComponentArrayInterpreter((n_site,), pt.θM) #loss_gf = get_loss_gf(g, transM, f, intϕ; gdev = identity) loss_gf = get_loss_gf(g, transM, transP, f, py, intϕ; pbm_covars, n_site_batch = n_batch, priorsP, priorsM, par_templates, - intθMs = intθMs_batch, intθP, frac_cluster_all) + ranef, frac_cluster_all, + is_omit_priors = Val(true), + ) + f2 = create_nsite_applicator(f, n_site_train) loss_gf_site = get_loss_gf(g, transM, transP, f2, py, intϕ; - pbm_covars, n_site_batch = n_site, priorsP, priorsM, par_templates, - intθMs = intθMs_site, intθP, frac_cluster_all) - nLjoint = @inferred first(loss_gf(p0, first(train_loader)...; is_testmode=true)) + pbm_covars, n_site_batch = n_site_train, priorsP, priorsM, par_templates, + ranef, frac_cluster_all, is_omit_priors = Val(true)) + nLcomponents = @inferred loss_gf(p0, first(train_loader)...; is_testmode=true) + () -> begin + ref_loss_gf = get_loss_gf(g, transM, transP, f, py, intϕ; + pbm_covars, n_site_batch = n_batch, priorsP, priorsM, par_templates, + ranef, frac_cluster_all, + is_omit_priors = Val(true), + ) + ref_nLcomponents = ref_loss_gf(p0, first(train_loader)...; is_testmode=true) + + end + nLjoint = @inferred first(nLcomponents) + # p01 = copy(p0); p01[end] = 5.0 + p01 = copy(p0); + if ranef_spec isa RandomEffects + p01[intϕ(1:length(p01)).ϕq.ranef.β[1:5]] .= 2.0 # one of the beta + end + nLjoint2 = @inferred first(loss_gf(p01, first(train_loader)...; is_testmode=true)) + @test nLjoint2 >= nLjoint + nLcomponents3 = loss_gf(p01, first(train_loader)...; is_testmode=true, ignore_ranef=Val(true)) + nLjoint3 = @inferred first(nLcomponents3) + @test nLjoint3 ≈ nLjoint - nLcomponents.nLRanef (xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch) = first(train_loader) # @usingany Cthulhu # @descend_code_warntype loss_gf(p0, xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch) - Zygote.gradient( + tmp = Zygote.gradient( p0 -> first(loss_gf( p0, xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch; is_testmode=false)), CA.getdata(p0)) + #intϕ(tmp[1]).ϕq.ranef.β[i_sites_batch,:] + tmp = Zygote.gradient( + p0 -> first(loss_gf( + p0, xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch; is_testmode=false)), CA.getdata(p01)) + # test that gradient of all random effects is zero with argument ignore_ranef=Val(true) + tmp3 = Zygote.gradient( + p0 -> first(loss_gf( + p0, xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch; + is_testmode=false,ignore_ranef=Val(true))), CA.getdata(p0)) + @test all(intϕ(tmp3[1]).ϕq.ranef .== zero(eltype(p0))) + intϕ(first(tmp)).ϕq + p0_norand = if !(ranef_spec isa NullRandomEffects) + # first optimized without random effects, to start from good parameter set + optf0 = Optimization.OptimizationFunction((ϕ, data) -> first(loss_gf( + ϕ, data...; is_testmode=false, ignore_ranef = Val(true))), + Optimization.AutoZygote()) + #Optimization.AutoFiniteDiff()) + optprob0 = Optimization.OptimizationProblem(optf0, CA.getdata(p0), train_loader) + res0 = Optimization.solve( + #optprob, Adam(0.02), callback = callback_loss(100), maxiters = 5000); + optprob0, Adam(0.02), + epochs = 30, + ) + intϕ(res0.u).ϕq.ranef + intϕ(p00).ϕq.ranef + res0.u + else + p0 + end + optf = Optimization.OptimizationFunction((ϕ, data) -> first(loss_gf(ϕ, data...; is_testmode=false)), Optimization.AutoZygote()) - optprob = OptimizationProblem(optf, CA.getdata(p0), train_loader) - - res = Optimization.solve( - #optprob, Adam(0.02), callback = callback_loss(100), maxiters = 5000); - optprob, Adam(0.02), maxiters = 2000) + #Optimization.AutoFiniteDiff()) + optprob = Optimization.OptimizationProblem(optf, CA.getdata(p0_norand), train_loader) + res = #@suppress begin + Optimization.solve(optprob, Adam(0.02), + #Optim.Options(checkfinite = false), + #callback = callback_loss(100), + epochs = 40, + #epochs = 90, + ) + #end (;nLjoint_pen, y_pred, θMs_tr_pred, θP_pred, nLy, nLprior_P, nLprior_M, loss_penalty) = loss_gf_site( res.u, train_loader.data...; is_testmode=true) #(nLjoint, y_pred, θMs_tr_pred, θP, nLy, nLprior_P, nLprior_M, loss_penalty) = loss_gf(p0, xM, xP, y_o, y_unc); + ϕq_opt = intϕ(res.u).ϕq + ϕq_opt.ranef + intϕ(p0).ϕq.ranef + #ϕq_opt.ranef.β[itrain_sites,:] + θMs_tr_pred = CA.ComponentArray(θMs_tr_pred, CA.getaxes(θMs_true')) - #TODO @test isapprox(par_templates.θP, intϕ(res.u).ϕP, rtol = 0.15) + #TODO @test isapprox(par_templates.θP, intϕ(res.u).μP, rtol = 0.15) #@test cor(vec(θMs_true), vec(θMs_tr_pred)) > 0.8 - @test cor(θMs_true'[:, 1], θMs_tr_pred[:, 1]) > 0.8 - @test cor(θMs_true'[:, 2], θMs_tr_pred[:, 2]) > 0.8 + @test cor(θMs_true'[i_sites_train_all, 1], θMs_tr_pred[:, 1]) > 0.8 + @test cor(θMs_true'[i_sites_train_all, 2], θMs_tr_pred[:, 2]) > 0.8 # started from low values -> increased but not too much above true values # logpdf.(priorsP, θP_pred) # logpdf.(priorsP, par_templates.θP) - @test all(transP(intϕ(p0).ϕP) .< θP_pred .< (1.2 .* θP_true)) + @test all(transP(intϕ(p0).ϕq.μP) .< θP_pred .< (1.2 .* θP_true)) @test all(0.8 .* θP_true .< θP_pred .< (1.2 .* θP_true)) () -> begin #@usingany UnicodePlots pdf(priorsP[1], θP_pred[1]) pdf(priorsP[1], θP_true[1]) - pdf(priorsP[1], transP(intϕ(p0).ϕP)[1]) - #pdf(priorsM[1], transP(intϕ(p0).ϕP)[1]) + pdf(priorsP[1], transP(intϕ(p0).μP)[1]) + #pdf(priorsM[1], transP(intϕ(p0).μP)[1]) quantile.(priorsM[2], [0.05, 0.5, 0.95]) - loss_gf(p0, xM, xP, y_o, y_unc, i_sites) + train_data = NamedTuple{(:xM, :xP, :y, :y_unc, :i_site)}(train_loader.data) + l0 = loss_gf_site(p0, train_data...; is_testmode=true) + lopt = loss_gf_site(res.u, train_data...; is_testmode=true) + _yt = f2(θP_true, θMs_true[:,i_sites_train_all]', train_data.xP)[1] + _yt == y_true[:,i_sites_train_all] + yp = f2(lopt.θP_pred, lopt.θMs_tr_pred, train_data.xP)[1] + yp == lopt.y_pred + scatterplot(vec(y_true[:,i_sites_train_all]), vec(train_data.y)) + scatterplot(vec(lopt.y_pred), vec(train_data.y)) + sum(py(train_data.y, lopt.y_pred, train_data.y_unc)) + sum(py(train_data.y, y_true[:,i_sites_train_all], train_data.y_unc)) + # better loglik with true ϕP2? no, wrong θMs cause wrong better θP + θP_test = [lopt.θP_pred[1], θP_true[2]] + ymod = f2(θP_test, lopt.θMs_tr_pred, train_data.xP)[1] + sum(py(train_data.y, ymod, train_data.y_unc)) # - scatterplot(θMs_true'[:,1], θMs_tr_pred[:,1]) - scatterplot(θMs_true'[:,2], θMs_tr_pred[:,2]) - scatterplot(log.(vec(θMs_true')), log.(vec(θMs_tr_pred))) - scatterplot(vec(y_pred), vec(y_o)) - hcat(par_templates.θP, intϕ(p0).ϕP, intϕ(res.u).ϕP, transP(intϕ(p0).ϕP), θP_pred) + scatterplot(θMs_true'[i_sites_train_all,1], θMs_tr_pred[:,1]) + scatterplot(θMs_true'[i_sites_train_all,2], θMs_tr_pred[:,2]) + scatterplot(vec(y_o[:,i_sites_train_all]), vec(y_pred)) + hcat(par_templates.θP, intϕ(p0).μP, intϕ(res.u).μP, transP(intϕ(p0).μP), θP_pred) end end diff --git a/test/test_elbo.jl b/test/test_elbo.jl index a2e3818..a75b0a8 100644 --- a/test/test_elbo.jl +++ b/test/test_elbo.jl @@ -38,18 +38,22 @@ end const prob = DoubleMM.DoubleMMCase() scenario = Val((:covarK2,)) -scenario = Val((:scalingall,)) scenario = Val((:sepvar,)) scenario = Val((:default,)) +scenario = Val((:ranef,)) scenario = Val((:clustered_sites,)) +scenario = Val((:clustered_sites,:ranef)) +#scenario = Val((:scalingall,)) # Also in clustered_sites One site uncertainty-scaling factor predicted by ML -pt = get_hybridproblem_par_templates(prob; scenario) -FT = eltype(pt.θM) +_pt = get_hybridproblem_par_templates(prob; scenario) +_FT = eltype(_pt.θM) #approx = MeanHVIApproximationMat() #approx = MeanVarSepHVIApproximation() -#approx = MeanScalingHVIApproximation([length(pt.θM)], FT(2) .* log.([FT(0.1) * pt.θM[end]])) +#approx = MeanScalingHVIApproximation([length(pt.θM)], _FT(2) .* log.([_FT(0.1) * pt.θM[end]])) test_scenario = (scenario) -> begin + scen = CP._val_value(scenario) + #@show scen #probc = HybridProblem(prob; scenario, approx); probc = HybridProblem(prob; scenario); # tmp = first(get_hybridproblem_train_dataloader(prob; scenario))[1] @@ -72,8 +76,10 @@ test_scenario = (scenario) -> begin # wrap inside function to not define(pollute) variables in level up _trainloader = get_hybridproblem_train_dataloader(probc; scenario) (_xM, _xP, _y_o, _y_unc, _i_sites) = _trainloader.data - @test _xM == xM - @test _y_o == y_o + i_sites_test_all = CP.get_i_sites_test(prob; scenario) + i_sites_train_all = setdiff(1:size(xM,2), i_sites_test_all) + @test _xM == xM[:,i_sites_train_all] + @test _y_o == y_o[:,i_sites_train_all] end; tmpf() # prediction by g(ϕg, XM) does not correspond to θMs_true, randomly initialized @@ -90,6 +96,15 @@ test_scenario = (scenario) -> begin priorsP = [priors[k] for k in keys(par_templates.θP)] priorsM = [priors[k] for k in keys(par_templates.θM)] + ranef_spec = if (:ranef ∈ scen) + RandomEffects((:r1, :K1)) + else + NullRandomEffects() + end + ranef = get_ranef_computer( + ranef_spec, keys(par_templates.θM), n_site, one(eltype(par_templates.θM))) + ϕq_ranef = setup_ϕq_ranef(ranef) + n_MC = 3 (; transP, transM) = get_hybridproblem_transforms(probc; scenario) cor_ends = get_hybridproblem_cor_ends(probc; scenario) @@ -97,8 +112,9 @@ test_scenario = (scenario) -> begin # transM = Stacked(elementwise(identity), elementwise(exp)) #transM = Stacked(elementwise(identity), elementwise(exp), elementwise(exp)) # test mismatch (;ϕqc, approx) = tmp = init_hybrid_ϕq( - probc.approx, par_templates.θP, par_templates.θM, transP, cor_ends; transM, n_site) - probc = HybridProblem(probc; approx) # update approx in probc + probc.approx, par_templates.θP, par_templates.θM, transP, cor_ends; + transM, n_site, ϕq_ranef) + probc = HybridProblem(probc; approx, ϕq = ϕqc) # update approx in probc # (ϕunc0, approx) = init_hybrid_ϕunc(cor_ends, zero(FT)) # ϕq0 = CP.update_μP_by_θP(ϕunc0, θP_true, transP) (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕqc) @@ -118,14 +134,14 @@ test_scenario = (scenario) -> begin g_gpu = ggdev(g_flux) end - i_sites = 1:n_batch + itrain_sites = 1:n_batch ζsP, ζsMs_tr, σ = @inferred ( # @usingany Cthulhu # @descend_code_warntype ( CP.generate_ζ( - probc.approx, rng, g, ϕ_ini, xM[:, i_sites]; + probc.approx, rng, g, ϕ_ini, xM[:, itrain_sites]; n_MC, cor_ends, pbm_covar_indices, - i_sites, + itrain_sites, ranef, int_ϕq=interpreters.ϕq, int_ϕg_ϕq=interpreters.ϕg_ϕq, is_testmode = false) ) @@ -141,8 +157,8 @@ test_scenario = (scenario) -> begin gr = Zygote.gradient( ϕ -> begin _ζsP, _ζsMs_tr, _σ = CP.generate_ζ( - probc.approx, rng, g, ϕ, xM[:, i_sites]; - i_sites, + probc.approx, rng, g, ϕ, xM[:, itrain_sites]; + itrain_sites, ranef, n_MC=8, cor_ends, pbm_covar_indices, int_ϕq=interpreters.ϕq, int_ϕg_ϕq=interpreters.ϕg_ϕq, is_testmode = true) @@ -151,6 +167,45 @@ test_scenario = (scenario) -> begin @test gr[1] isa Vector end + @testset "predict_hvi" begin + n_sample_pred = 200 # 10_000 #2_400 + # not type stable, because of probc - check type stabiliy inside + ans_predict = predict_hvi(rng, probc; + scenario, n_sample_pred, + is_inferred = Val(true), + ); + (; y, θsP, θsMs_tr, entropy_ζ, logjacs_P, logjacs_Ms) = ans_predict + n_site_pred = size(θsMs_tr,1) + @test size(logjacs_P) == (n_sample_pred,) + @test size(logjacs_Ms) == (n_site_pred, n_sample_pred) + # check jacobian against transformation of a single (backtransformed) par_vector + @test logjacs_P[1] ≈ with_logabsdet_jacobian(transP, inverse(transP)(θsP[:,1]))[2] + @test logjacs_Ms[1,1] ≈ with_logabsdet_jacobian( + transM, inverse(transM)(θsMs_tr[1,:,1]))[2] + #size(_ζsMs), size(θsMs) + #size(_ζsP), size(θsP) + # + # below test that generated sample matches specified distribution + # + # n_site_pred = size(xM,2) + # train_data = NamedTuple{(:xM, :xP, :y, :y_unc, :i_site)}(probc.train_dataloader.data) + (;res_predict, res_elbo) = predict_hvi_and_compute_elbo_components(rng, probc; + scenario, n_sample_pred, + ) + @test res_elbo.nLjoint isa Float32 + @test res_elbo.entropy_ζ isa Float32 + @test res_elbo.nLprior_P isa Float32 + @test res_elbo.loss_penalty isa Float32 + @test res_elbo.nLRanef isa Float32 + @test res_elbo.neglogjac_P isa Float32 + @test res_elbo.nLy isa AbstractVector{Float32} + @test res_elbo.nLprior_Ms isa AbstractVector{Float32} + @test res_elbo.loss_penalty_sites isa AbstractVector{Float32} + @test res_elbo.neglogjac_Ms isa AbstractVector{Float32} + @test res_elbo.nLsites isa AbstractVector{Float32} + @test length(res_elbo.nLsites) == n_site_pred + end; + if !(:covarK2 ∈ CP._val_value(scenario)) && (probc.approx isa MeanHVIApproximation) # can only test distribution if g is not repeated @testset "generate_ζ check sd residuals $(last(CP._val_value(scenario)))" begin @@ -183,13 +238,13 @@ test_scenario = (scenario) -> begin #hcat(ϕ_ini, ϕ, _ϕ)[1:4,:] #hcat(ϕ_ini, ϕ, _ϕ)[(end-20):end,:] n_predict = 10_000 #8_000 - i_sites = 1:n_batch - xM_batch = xM[:, i_sites] + itrain_sites = 1:n_batch + xM_batch = xM[:, itrain_sites] _ζsP, _ζsMs_tr, _σ = @inferred ( # @descend_code_warntype ( CP.generate_ζ( probc.approx, rng, g, _ϕ, xM_batch; - i_sites, + itrain_sites, n_MC = n_predict, cor_ends, pbm_covar_indices, int_ϕq=interpreters.ϕq, int_ϕg_ϕq=interpreters.ϕg_ϕq, is_testmode = true) @@ -271,12 +326,12 @@ test_scenario = (scenario) -> begin ϕ = ggdev(CA.getdata(ϕ_ini)) @test g_gpu.μ isa GPUArraysCore.AbstractGPUArray # @test g_gpu.app isa HybridVariationalInferenceFluxExt.FluxApplicator - xMg_batch = ggdev(xM[:, i_sites]) + xMg_batch = ggdev(xM[:, itrain_sites]) ζsP_d, ζsMs_tr_d, σ_d = @inferred ( # @descend_code_warntype ( CP.generate_ζ( probc.approx, rng, g_gpu, ϕ, xMg_batch; - i_sites, + itrain_sites, n_MC, cor_ends, pbm_covar_indices, int_ϕq=interpreters.ϕq, int_ϕg_ϕq=interpreters.ϕg_ϕq, is_testmode = true)) @@ -291,7 +346,7 @@ test_scenario = (scenario) -> begin ϕ -> begin _ζsP, _ζsMs_tr, _σ = CP.generate_ζ( probc.approx, rng, g_gpu, ϕ, xMg_batch; - i_sites, + itrain_sites, n_MC, cor_ends, pbm_covar_indices, int_ϕq=interpreters.ϕq, int_ϕg_ϕq=interpreters.ϕg_ϕq, is_testmode = false) @@ -374,33 +429,34 @@ test_scenario = (scenario) -> begin @testset "neg_elbo_gtf cpu $(last(CP._val_value(scenario)))" begin scen = CP._val_value(scenario) - i_sites = 1:n_batch + itrain_sites = 1:n_batch transMs = StackedArray(transM, size(ζsMs_tr, 1)) #intθMs = ComponentArrayInterpreter((n_batch,), int_M) intθMs = get_concrete(ComponentArrayInterpreter((n_batch,), int_M)) n_sites_cluster, clusters = CP.get_clusters(n_site; scenario) frac_cluster_all = convert.(eltype(ϕ_ini), 1 ./ n_sites_cluster[clusters]) cost = @inferred ( + #@usingany Cthulhu #@descend_code_warntype ( neg_elbo_gtf(rng, ϕ_ini, g, f, py, - xM[:, i_sites], xP[:, i_sites], y_o[:, i_sites], y_unc[:, i_sites], i_sites; + xM[:, itrain_sites], xP[:, itrain_sites], y_o[:, itrain_sites], y_unc[:, itrain_sites], itrain_sites; int_ϕq, int_ϕg_ϕq, cor_ends, pbm_covar_indices, transP, transMs, priorsP, priorsM, is_testmode = true, is_omit_priors = Val(false), zero_prior_logdensity=zero(eltype(ϕ_ini)), - probc.approx, intθMs, intθP = int_P, frac_cluster_all + probc.approx, intθMs, intθP = int_P, frac_cluster_all, ranef ) ) #@test cost isa Float64 @test cost isa promote_type(eltype(xM), eltype(y_o), eltype(ϕ_ini)) gr = Zygote.gradient( ϕ -> neg_elbo_gtf(rng, ϕ, g, f, py, - xM[:, i_sites], xP[:, i_sites], y_o[:, i_sites], y_unc[:, i_sites], i_sites; + xM[:, itrain_sites], xP[:, itrain_sites], y_o[:, itrain_sites], y_unc[:, itrain_sites], itrain_sites; int_ϕq, int_ϕg_ϕq, cor_ends, pbm_covar_indices, transP, transMs, priorsP, priorsM, is_testmode = false, is_omit_priors = Val(false), zero_prior_logdensity=zero(eltype(ϕ_ini)), - probc.approx, intθMs, intθP = int_P, frac_cluster_all + probc.approx, intθMs, intθP = int_P, frac_cluster_all, ranef ), CA.getdata(ϕ_ini)) @test gr[1] isa Vector @@ -408,31 +464,31 @@ test_scenario = (scenario) -> begin if ggdev isa MLDataDevices.AbstractGPUDevice @testset "neg_elbo_gtf gpu $(last(CP._val_value(scenario)))" begin - i_sites = 1:n_batch + itrain_sites = 1:n_batch transMs = StackedArray(transM, size(ζsMs_tr, 1)) ϕ = ggdev(CA.getdata(ϕ_ini)) - xMg_batch = ggdev(xM[:, i_sites]) - xP_batch = xP[:, i_sites] # used in f which runs on CPU + xMg_batch = ggdev(xM[:, itrain_sites]) + xP_batch = xP[:, itrain_sites] # used in f which runs on CPU cost = @inferred ( #@descend_code_warntype ( neg_elbo_gtf(rng, ϕ, g_gpu, f, py, - xMg_batch, xP_batch, y_o[:, i_sites], y_unc[:, i_sites], i_sites; + xMg_batch, xP_batch, y_o[:, itrain_sites], y_unc[:, itrain_sites], itrain_sites; int_ϕq, int_ϕg_ϕq, n_MC=3, cor_ends, pbm_covar_indices, transP, transMs, priorsP, priorsM, is_testmode = true, is_omit_priors = Val(false), zero_prior_logdensity=zero(eltype(ϕ_ini)), - probc.approx, + probc.approx, ranef, ) ) @test cost isa Float64 gr = Zygote.gradient( ϕ -> neg_elbo_gtf(rng, ϕ, g_gpu, f, py, - xMg_batch, xP_batch, y_o[:, i_sites], y_unc[:, i_sites], i_sites; + xMg_batch, xP_batch, y_o[:, itrain_sites], y_unc[:, itrain_sites], itrain_sites; int_ϕq, int_ϕg_ϕq, n_MC=3, cor_ends, pbm_covar_indices, transP, transMs, priorsP, priorsM, is_testmode = false, is_omit_priors = Val(false), zero_prior_logdensity=zero(eltype(ϕ_ini)), - probc.approx, + probc.approx, ranef, ), ϕ) @test gr[1] isa GPUArraysCore.AbstractGPUVector @@ -446,17 +502,21 @@ test_scenario = (scenario) -> begin # @test length(intm_PMs_gen) == 402 # @test trans_PMs_gen.length_in == 402 n_sample_pred = 30 + n_sample_ranef = 5 + n_site_pred = size(xM,2) + frac_cluster = ones(eltype(ϕ_ini), n_site_pred) (; θsP, θsMs_tr, entropy_ζ) = - #Cthulhu.@descend_code_warntype ( @inferred ( + #Cthulhu.@descend_code_warntype ( sample_posterior(rng, g, ϕ_ini, xM; - i_sites = 1:size(xM, 2), + itrain_sites = 1:size(xM, 2), + #itrain_sites = Int[], int_ϕg_ϕq, int_ϕq, transP, transM, cdev = identity, - n_sample_pred, cor_ends, pbm_covar_indices, + n_sample_pred, n_sample_ranef, cor_ends, pbm_covar_indices, is_testmode = true, - probc.approx, + probc.approx, ranef, frac_cluster, ) ) @test θsP isa AbstractMatrix @@ -475,6 +535,8 @@ test_scenario = (scenario) -> begin ϕ_ini_g = ggdev(CA.getdata(ϕ_ini)) xMg = ggdev(xM) n_sample_pred = 30 + n_site_pred = size(xM,2) + frac_cluster = ones(eltype(ϕ_ini), n_site_pred) (; θsP, θsMs_tr, entropy_ζ) = #Cthulhu.@descend_code_warntype ( @inferred ( @@ -485,7 +547,7 @@ test_scenario = (scenario) -> begin cdev = identity, # do not transfer to CPU n_sample_pred, cor_ends, pbm_covar_indices, is_testmode = true, - probc.approx, + probc.approx, ranef, frac_cluster, ) ) # this variant without the problem, does not attach axes @@ -531,6 +593,7 @@ end # test_scenario #test_scenario(Val((:scalingall,))) test_scenario(Val((:clustered_sites,))) +test_scenario(Val((:clustered_sites, :ranef,))) # with random effects in all parameters test_scenario(Val((:default,))) test_scenario(Val((:sepvar,))) diff --git a/test/test_logden_normal.jl b/test/test_logden_normal.jl index a27fdc0..d94c5ba 100644 --- a/test/test_logden_normal.jl +++ b/test/test_logden_normal.jl @@ -10,15 +10,15 @@ using LinearAlgebra loglik_norm(y, μ, σ) = -1 / 2 .* (Distributions.log2π .+ 2 .* log.(σ) .+ abs2.(y .- μ) ./ abs2.(σ)) loglik_norm_l(y, μ, logσ) = -1 / 2 .* (Distributions.log2π .+ 2 .* logσ .+ abs2.(y .- μ) ./ abs2.(exp.(logσ))) logden_norm_l(y, μ, logσ) = -1 / 2 .* (2 .* logσ .+ abs2.(y .- μ) ./ abs2.(exp.(logσ))) - neg_logden_norm_l2(y, μ, logσ2) = (logσ2 .+ abs2.(y .- μ) .* exp.(-logσ2)) ./ 2 + neg_logden_norm_l2(y, μ, logσ2) = (logσ2 .+ abs2.(y .- μ) .* exp.(.-logσ2)) ./ 2 # first test that neg_logden_norm_l2 returns values of logpdf(Normal) up to an additive μ = [1.0, 1.0] σ = [1.1, 2.0] logσ2 = log.(abs2.(σ)) y = [1.2, 1.1] - tmp_true = logpdf.(Normal.(μ, σ), y) - dlogpdf = tmp_true .- tmp_true[1] + logpdf_true = logpdf.(Normal.(μ, σ), y) + dlogpdf = logpdf_true .- logpdf_true[1] #loglik_norm(y, μ, σ) #loglik_norm_l(y, μ, log.(σ)) #tmp = logden_norm_l(y, μ, log.(σ)) @@ -26,8 +26,15 @@ using LinearAlgebra tmp = neg_logden_norm_l2(y, μ, logσ2) @test isapprox(tmp .- tmp[1], -dlogpdf) - # next test that the sum of neg_logden_norm_l2 corresponds to - @test neg_logden_indep_normal(y, μ, logσ2) ≈ sum(tmp) + # next test that the sum of neg_logden_norm_l2 corresponds to tested + res_vec = neg_logden_indep_normal(y, μ, logσ2) + @test res_vec ≈ sum(tmp) + + # test several columns + μ3 = hcat(μ,μ,μ) + y3 = hcat(y,y,y) + resm = neg_logden_indep_normal(y3, μ3, logσ2) + @test resm == [res_vec, res_vec, res_vec] end; @testset "entropy_MvNormal" begin diff --git a/test/test_missingdriver.jl b/test/test_missingdriver.jl index 94ebdce..a21209c 100644 --- a/test/test_missingdriver.jl +++ b/test/test_missingdriver.jl @@ -13,6 +13,7 @@ gdev = gpu_device() cdev = cpu_device() using OptimizationOptimisers +import CommonSolve: solve using Lux # in order to load extension @@ -93,6 +94,7 @@ function test_driverNaN(scenario::Val{scen}) where scen #callback = callback_loss(100), # output during fitting #callback = callback_loss(10), # output during fitting epochs = 2, + #epochs_callback = 0, # not progress output is_omit_priors = Val(:f_on_gpu ∈ scen), # prior computation does not work on gpu scenario, ); @@ -109,6 +111,7 @@ function test_driverNaN(scenario::Val{scen}) where scen (; probo, interpreters) = solve(prob, solver; rng, callback = callback_loss(10), # output during fitting epochs = 2, + epochs_callback = 0, # not progress output scenario, ); @test all(isfinite.(probo.θP)) diff --git a/test/test_no_globals.jl b/test/test_no_globals.jl index 8ac64df..8007de1 100644 --- a/test/test_no_globals.jl +++ b/test/test_no_globals.jl @@ -14,12 +14,15 @@ cdev = cpu_device() using OptimizationOptimisers using Lux # in order to load extension +import CommonSolve: solve # scenario = Val(()); scen=() +# scenario = Val((:ranef,)); scen=(:ranef,) function test_no_globals(scenario::Val{scen}) where scen scenario = Val((scen..., :no_globals)) prob = HybridProblem(DoubleMM.DoubleMMCase(); scenario); θP0, θM0 = get_hybridproblem_par_templates(prob) + #ranef_spec = get_hybridproblem_ranef(prob; scenario) @test isempty(θP0) solver_point = HybridPointSolver(; alg=Adam(0.02)) rng = StableRNG(111) @@ -27,6 +30,7 @@ function test_no_globals(scenario::Val{scen}) where scen #callback = callback_loss(100), # output during fitting #callback = callback_loss(10), # output during fitting epochs = 2, + epochs_callback = 0, # no progress output is_omit_priors = Val(:f_on_gpu ∈ scen), # prior computation does not work on gpu scenario, ); @@ -44,6 +48,7 @@ function test_no_globals(scenario::Val{scen}) where scen #callback = callback_loss(10), # output during fitting is_omit_priors = Val(:f_on_gpu ∈ scen), # prior computation does not work on gpu epochs = 2, + epochs_callback = 0, # no progress output scenario, ); @test all(isfinite.(CP.get_hybridproblem_θP(probo))) diff --git a/test/test_sample_zeta.jl b/test/test_sample_zeta.jl index 2c27982..275d08b 100644 --- a/test/test_sample_zeta.jl +++ b/test/test_sample_zeta.jl @@ -102,7 +102,7 @@ function test_with_scenario(scenario) int_ϕq = get_concrete(ComponentArrayInterpreter(ϕc.ϕq)) n_MC_pred = 300 # larger n_MC to test σ2 n_site_batch = size(ϕc.Ms,2) - i_sites = 1:n_site_batch + itrain_sites = 1:n_site_batch #rng = StableRNG(111) # @inferred gives any, while Cthulhu inferres concrete type # ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq; @@ -111,9 +111,9 @@ function test_with_scenario(scenario) # n_MC=n_MC_pred, cor_ends, int_ϕq)) # ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq; # n_MC=n_MC_pred, cor_ends, int_ϕq) - i_sites = 1:n_site_batch + itrain_sites = 1:n_site_batch ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm( - approx, rng, i_sites, ϕc.Ms, ϕc.ϕq; + approx, rng, itrain_sites, ϕc.Ms, ϕc.ϕq; n_MC=n_MC_pred, cor_ends, int_ϕq) #@code_warntype CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq; n_MC=n_MC_pred, cor_ends, int_ϕq) #@usingany Cthulhu @@ -125,7 +125,7 @@ function test_with_scenario(scenario) gr = Zygote.gradient(ϕc -> begin ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm( - approx, rng, i_sites, ϕc.Ms, ϕc.ϕq; + approx, rng, itrain_sites, ϕc.Ms, ϕc.ϕq; n_MC, cor_ends, int_ϕq) sum(ζP_resids) + sum(ζMs_parfirst_resids) end, ϕc)[1] @@ -164,9 +164,9 @@ function test_with_scenario(scenario) # ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm( # approx, rng, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq); # n_MC = n_MC_pred, cor_ends, int_ϕq) - i_sites = 1:n_site_batch + itrain_sites = 1:n_site_batch ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm( - approx, rng, i_sites, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq); + approx, rng, itrain_sites, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq); n_MC = n_MC_pred, cor_ends, int_ϕq) #@descend_code_warntype CP.sample_ζresid_norm(rng, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq); n_MC = n_MC_pred, cor_ends, int_ϕq) @test ζP_resids isa GPUArraysCore.AbstractGPUArray @@ -196,26 +196,26 @@ function test_with_scenario(scenario) () -> begin CP.sample_ζresid_norm( #@benchmark CP.sample_ζresid_norm( - approx, rng, i_sites, ϕc.Ms, ϕc.ϕq; + approx, rng, itrain_sites, ϕc.Ms, ϕc.ϕq; n_MC, cor_ends, int_ϕq) # CP.sample_ζresid_norm( #@benchmark CP.sample_ζresid_norm( - approx, rng, i_sites, ϕcd.Ms, ϕcd.ϕq; + approx, rng, itrain_sites, ϕcd.Ms, ϕcd.ϕq; n_MC, cor_ends, int_ϕq) # ϕc_few = CA.ComponentVector(ϕc; Ms = ϕc.Ms[:,1:n_site_few]); Zygote.gradient(ϕc -> begin #@benchmark Zygote.gradient(ϕc -> begin # many small allocs ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm( - approx, rng, i_sites, ϕc.Ms, ϕc.ϕq; + approx, rng, itrain_sites, ϕc.Ms, ϕc.ϕq; n_MC, cor_ends, int_ϕq) sum(ζP_resids) + sum(ζMs_parfirst_resids) end, ϕc_few)[1] Zygote.gradient(ϕc -> begin # many small allocs #@benchmark Zygote.gradient(ϕc -> begin # many small allocs ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm( - approx, rng, i_sites, CA.getdata(ϕc.Ms), CA.getdata(ϕc.ϕq); + approx, rng, itrain_sites, CA.getdata(ϕc.Ms), CA.getdata(ϕc.ϕq); n_MC, cor_ends, int_ϕq) sum(ζP_resids) + sum(ζMs_parfirst_resids) end, ϕcd_few)[1] diff --git a/test/test_util.jl b/test/test_util.jl index 913bec6..5ee3e0d 100644 --- a/test/test_util.jl +++ b/test/test_util.jl @@ -1,7 +1,11 @@ using Test using HybridVariationalInference: vectuptotupvec_allowmissing, vectuptotupvec, insert_zeros +using HybridVariationalInference: replace_columns_matrix using HybridVariationalInference: HybridVariationalInference as HVI using Zygote +using Distributions +using LinearAlgebra + @testset "OneBasedVectorWithZero" begin @@ -134,3 +138,78 @@ end; gr = Zygote.gradient(x -> sum(skipmissing(vectuptotupvec_allowmissing(x)[1])), vectupm) end; +@testset "log_density_mvn_cholesky" begin + # Test cases: different dimensions and covariance matrices + test_cases = [ + (2, 0.1), # 2D, small variance + #(3, 1.0), # 3D, moderate + #(5, 2.0), # 5D, larger + (10, 0.5), # 10D, small + ] + # n, scale = test_cases[3] + for (n, scale) in test_cases + # Generate a random positive definite covariance matrix + A = randn(n, n) + C = A * A' + scale * I # Ensure positive definite + U = cholesky(C).L' + x = randn(n) + log_p_custom = HVI.log_density_mvn_cholesky(U, x) + # Compute log-density using Distributions.jl + dist = MvNormal(zeros(n), C) + log_p_builtin = logpdf(dist, x) + @test isapprox(log_p_custom, log_p_builtin, atol=1e-10) + end + # test Zygote can handle this function + () -> begin + A = rand(3,3) + U = cholesky(A * A' + I).L' + Zygote.gradient(x -> HVI.log_density_mvn_cholesky(U, x), randn(3)) + end +end + +@testset "replace_columns_matrix" begin + # Test case 1: Basic replacement + x = [1 2 3; 4 5 6; 7 8 9] + col_indices = [1, 3] + y = [10 11; 12 13; 14 15] + result = replace_columns_matrix(x, col_indices, y) + expected = [10 2 11; 12 5 13; 14 8 15] + @test result ≈ expected + # TODO - do not know why error occurs, have no j + # tmp = Zygote.gradient(y -> sum(replace_columns_matrix(x, col_indices, y)), y) + # + # Test case 2: Replace all columns + x = [1 2; 3 4] + col_indices = [1, 2] + y = [5 6; 7 8] + # + result = replace_columns_matrix(x, col_indices, y) + expected = [5 6; 7 8] + @test result ≈ expected + # + # Test case 3: Replace single column + x = [1 2 3; 4 5 6] + col_indices = [2] + y = reshape([10; 20], :,1) + result = replace_columns_matrix(x, col_indices, y) + expected = [1 10 3; 4 20 6] + @test result ≈ expected + # + # Test case 4: Empty replacement + x = [1.0 2; 3 4] + col_indices = Int[] + y = Matrix{Float64}(undef, 2, 0) + result = replace_columns_matrix(x, col_indices, y) + expected = x + @test result ≈ expected + # + # Test case 5: Differentiation compatibility (Zygote) + x = [1.0 2.0; 3.0 4.0] + col_indices = [1] + y = reshape([5.0; 6.0], :, 1) + # Check that gradient can be computed + # TODO + # gradient((x,y) -> sum(replace_columns_matrix(x, col_indices, y)), x,y) +end + +