Refactor allocation with MetaGraph (#807)

Fixes #785.

---------

Co-authored-by: Martijn Visser <[email protected]>

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0-rc1"
 manifest_format = "2.0"
-project_hash = "f38abc22b48e63729fc0df2250576035135df64d"
+project_hash = "05a69b11bfa835080edd1b88f4f3d29dd068d662"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "5d2e21d7b0d8c22f67483ef95ebdc39c0e6b6003"
@@ -498,6 +498,12 @@ git-tree-sha1 = "b12f05108e405dadcc2aff0008db7f831374e051"
 uuid = "29a986be-02c6-4525-aec4-84b980013641"
 version = "2.0.0"
 
+[[deps.FileIO]]
+deps = ["Pkg", "Requires", "UUIDs"]
+git-tree-sha1 = "299dc33549f68299137e51e6d49a13b5b1da9673"
+uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+version = "1.16.1"
+
 [[deps.FileWatching]]
 uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
 
@@ -669,6 +675,12 @@ git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
 uuid = "82899510-4779-5014-852e-03e436cf321d"
 version = "1.0.0"
 
+[[deps.JLD2]]
+deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Printf", "Reexport", "Requires", "TranscodingStreams", "UUIDs"]
+git-tree-sha1 = "9bbb5130d3b4fa52846546bca4791ecbdfb52730"
+uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+version = "0.4.38"
+
 [[deps.JLLWrappers]]
 deps = ["Artifacts", "Preferences"]
 git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca"
@@ -899,6 +911,12 @@ deps = ["Artifacts", "Libdl"]
 uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
 version = "2.28.2+1"
 
+[[deps.MetaGraphsNext]]
+deps = ["Graphs", "JLD2", "SimpleTraits"]
+git-tree-sha1 = "8dd4f3f8a643d53e61ff9115749f522c35a38f3f"
+uuid = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+version = "0.6.0"
+
 [[deps.Missings]]
 deps = ["DataAPI"]
 git-tree-sha1 = "f66bdc5de519e8f8ae43bdc598782d35a25b1272"
@@ -1171,7 +1189,7 @@ uuid = "ae029012-a4dd-5104-9daa-d747884805df"
 version = "1.3.0"
 
 [[deps.Ribasim]]
-deps = ["Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimeZones", "TimerOutputs", "TranscodingStreams"]
+deps = ["Accessors", "Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "EnumX", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimeZones", "TimerOutputs", "TranscodingStreams"]
 path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "0.4.0"

core/Project.toml

@@ -5,6 +5,7 @@ manifest = "../Manifest.toml"
 version = "0.4.0"
 
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
 CodecLz4 = "5ba52731-8f18-5e0d-9241-30f10d1ec561"
@@ -17,6 +18,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -26,6 +28,7 @@ JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Legolas = "741b9549-f6ed-4911-9fbf-4a1c0c97f0cd"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9"
@@ -39,6 +42,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 TranscodingStreams = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
 
 [compat]
+Accessors = "0.1"
 Aqua = "0.8"
 Arrow = "2.3"
 BasicModelInterface = "0.1"
@@ -55,16 +59,18 @@ Dates = "<0.0.1,1"
 Dictionaries = "0.3.25"
 DiffEqCallbacks = "2.29.1"
 Documenter = "0.27,1"
+EnumX = "1.0"
 FiniteDiff = "2.21"
 ForwardDiff = "0.10"
-Graphs = "1.6"
+Graphs = "1.9"
 HiGHS = "1.7"
 IOCapture = "0.2"
 IterTools = "1.4"
 JuMP = "1.15"
 Legolas = "0.5"
 Logging = "<0.0.1,1"
 LoggingExtras = "1"
+MetaGraphsNext = "0.6"
 OrdinaryDiffEq = "6.7"
 PreallocationTools = "0.4"
 ReTestItems = "1.20"

core/src/Ribasim.jl

@@ -20,6 +20,7 @@ import HiGHS
 import JuMP
 import TranscodingStreams
 
+using Accessors: @set
 using Arrow: Arrow, Table
 using CodecLz4: LZ4FrameCompressor
 using CodecZstd: ZstdCompressor
@@ -29,8 +30,9 @@ using DataInterpolations: LinearInterpolation, derivative
 using Dates
 using DBInterface: execute, prepare
 using Dictionaries: Indices, Dictionary, gettoken, dictionary
-using ForwardDiff: pickchunksize
 using DiffEqCallbacks
+using EnumX
+using ForwardDiff: pickchunksize
 using Graphs:
     add_edge!,
     adjacency_matrix,
@@ -46,6 +48,8 @@ using Graphs:
 using Legolas: Legolas, @schema, @version, validate, SchemaVersion, declared
 using Logging: current_logger, min_enabled_level, with_logger
 using LoggingExtras: EarlyFilteredLogger, LevelOverrideLogger
+using MetaGraphsNext:
+    MetaGraphsNext, MetaGraph, label_for, labels, outneighbor_labels, inneighbor_labels
 using OrdinaryDiffEq
 using PreallocationTools: DiffCache, FixedSizeDiffCache, get_tmp
 using SciMLBase