diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..0e4598e
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,5 @@
+data/*
+plots/*
+weights/*
+slurm-*
+*.png
diff --git a/CIG/Manifest.toml b/CIG/Manifest.toml
new file mode 100644
index 0000000..81b3334
--- /dev/null
+++ b/CIG/Manifest.toml
@@ -0,0 +1,1275 @@
+# This file is machine-generated - editing it directly is not advised
+
+julia_version = "1.8.5"
+manifest_format = "2.0"
+project_hash = "4134fada2dd3cd44b649de79e47ca941bf914660"
+
+[[deps.AbstractFFTs]]
+deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
+git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef"
+uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+version = "1.5.0"
+
+[[deps.AbstractNFFTs]]
+deps = ["LinearAlgebra", "Printf"]
+git-tree-sha1 = "292e21e99dedb8621c15f185b8fdb4260bb3c429"
+uuid = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
+version = "0.8.2"
+
+[[deps.Accessors]]
+deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown", "Requires", "Test"]
+git-tree-sha1 = "b392ede862e506d451fc1616e79aa6f4c673dab8"
+uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+version = "0.1.38"
+
+[[deps.Adapt]]
+deps = ["LinearAlgebra", "Requires"]
+git-tree-sha1 = "50c3c56a52972d78e8be9fd135bfb91c9574c140"
+uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+version = "4.1.1"
+
+[[deps.ArgCheck]]
+git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4"
+uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+version = "2.3.0"
+
+[[deps.ArgParse]]
+deps = ["Logging", "TextWrap"]
+git-tree-sha1 = "22cf435ac22956a7b45b0168abbc871176e7eecc"
+uuid = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+version = "1.2.0"
+
+[[deps.ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+version = "1.1.1"
+
+[[deps.Artifacts]]
+uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+
+[[deps.Atomix]]
+deps = ["UnsafeAtomics"]
+git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be"
+uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
+version = "0.1.0"
+
+[[deps.Augmentor]]
+deps = ["ComputationalResources", "CoordinateTransformations", "FileIO", "IdentityRanges", "ImageCore", "ImageFiltering", "ImageTransformations", "InteractiveUtils", "Interpolations", "LinearAlgebra", "MLDataPattern", "OffsetArrays", "Rotations", "StaticArrays", "Statistics"]
+git-tree-sha1 = "47cd3b6a76a132c899ed798bb72cc1f1fe60f8d5"
+uuid = "02898b10-1f73-11ea-317c-6393d7073e15"
+version = "0.6.6"
+
+[[deps.AxisAlgorithms]]
+deps = ["LinearAlgebra", "Random", "SparseArrays", "WoodburyMatrices"]
+git-tree-sha1 = "66771c8d21c8ff5e3a93379480a2307ac36863f7"
+uuid = "13072b0f-2c55-5437-9ae7-d433b7a33950"
+version = "1.0.1"
+
+[[deps.AxisArrays]]
+deps = ["Dates", "IntervalSets", "IterTools", "RangeArrays"]
+git-tree-sha1 = "16351be62963a67ac4083f748fdb3cca58bfd52f"
+uuid = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+version = "0.4.7"
+
+[[deps.BFloat16s]]
+deps = ["LinearAlgebra", "Printf", "Random", "Test"]
+git-tree-sha1 = "2c7cc21e8678eff479978a0a2ef5ce2f51b63dff"
+uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
+version = "0.5.0"
+
+[[deps.BangBang]]
+deps = ["Accessors", "ConstructionBase", "InitialValues", "LinearAlgebra", "Requires"]
+git-tree-sha1 = "e2144b631226d9eeab2d746ca8880b7ccff504ae"
+uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+version = "0.4.3"
+
+[[deps.Base64]]
+uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+
+[[deps.Baselet]]
+git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e"
+uuid = "9718e550-a3fa-408a-8086-8db961cd8217"
+version = "0.1.1"
+
+[[deps.BasicInterpolators]]
+deps = ["LinearAlgebra", "Memoize", "Random"]
+git-tree-sha1 = "3f7be532673fc4a22825e7884e9e0e876236b12a"
+uuid = "26cce99e-4866-4b6d-ab74-862489e035e0"
+version = "0.7.1"
+
+[[deps.CEnum]]
+git-tree-sha1 = "eb4cb44a499229b3b8426dcfb5dd85333951ff90"
+uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+version = "0.4.2"
+
+[[deps.CUDA]]
+deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"]
+git-tree-sha1 = "fdd9dfb67dfefd548f51000cc400bb51003de247"
+uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
+version = "5.4.3"
+
+[[deps.CUDA_Driver_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "325058b426c2b421e3d2df3d5fa646d72d2e3e7e"
+uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc"
+version = "0.9.2+0"
+
+[[deps.CUDA_Runtime_Discovery]]
+deps = ["Libdl"]
+git-tree-sha1 = "33576c7c1b2500f8e7e6baa082e04563203b3a45"
+uuid = "1af6417a-86b4-443c-805f-a4643ffb695f"
+version = "0.3.5"
+
+[[deps.CUDA_Runtime_jll]]
+deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "afea94249b821dc754a8ca6695d3daed851e1f5a"
+uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
+version = "0.14.1+0"
+
+[[deps.CatIndices]]
+deps = ["CustomUnitRanges", "OffsetArrays"]
+git-tree-sha1 = "a0f80a09780eed9b1d106a1bf62041c2efc995bc"
+uuid = "aafaddc9-749c-510e-ac4f-586e18779b91"
+version = "0.2.2"
+
+[[deps.ChainRulesCore]]
+deps = ["Compat", "LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "3e4b134270b372f2ed4d4d0e936aabaefc1802bc"
+uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+version = "1.25.0"
+
+[[deps.ChangesOfVariables]]
+deps = ["InverseFunctions", "LinearAlgebra", "Test"]
+git-tree-sha1 = "799b25ca3a8a24936ae7b5c52ad194685fc3e6ef"
+uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
+version = "0.1.9"
+
+[[deps.ColorSchemes]]
+deps = ["ColorTypes", "ColorVectorSpace", "Colors", "FixedPointNumbers", "PrecompileTools", "Random"]
+git-tree-sha1 = "b5278586822443594ff615963b0c09755771b3e0"
+uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
+version = "3.26.0"
+
+[[deps.ColorTypes]]
+deps = ["FixedPointNumbers", "Random"]
+git-tree-sha1 = "b10d0b65641d57b8b4d5e234446582de5047050d"
+uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
+version = "0.11.5"
+
+[[deps.ColorVectorSpace]]
+deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "SpecialFunctions", "Statistics", "TensorCore"]
+git-tree-sha1 = "600cc5508d66b78aae350f7accdb58763ac18589"
+uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4"
+version = "0.9.10"
+
+[[deps.Colors]]
+deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
+git-tree-sha1 = "362a287c3aa50601b0bc359053d5c2468f0e7ce0"
+uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
+version = "0.12.11"
+
+[[deps.Compat]]
+deps = ["Dates", "LinearAlgebra", "TOML", "UUIDs"]
+git-tree-sha1 = "8ae8d32e09f0dcf42a36b90d4e17f5dd2e4c4215"
+uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
+version = "4.16.0"
+
+[[deps.CompilerSupportLibraries_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
+version = "1.0.1+0"
+
+[[deps.CompositionsBase]]
+git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad"
+uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b"
+version = "0.1.2"
+
+[[deps.ComputationalResources]]
+git-tree-sha1 = "52cb3ec90e8a8bea0e62e275ba577ad0f74821f7"
+uuid = "ed09eef8-17a6-5b46-8889-db040fac31e3"
+version = "0.3.2"
+
+[[deps.Conda]]
+deps = ["Downloads", "JSON", "VersionParsing"]
+git-tree-sha1 = "b19db3927f0db4151cb86d073689f2428e524576"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.10.2"
+
+[[deps.ConstructionBase]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "76219f1ed5771adbb096743bff43fb5fdd4c1157"
+uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
+version = "1.5.8"
+
+[[deps.ContextVariablesX]]
+deps = ["Compat", "Logging", "UUIDs"]
+git-tree-sha1 = "25cc3803f1030ab855e383129dcd3dc294e322cc"
+uuid = "6add18c4-b38d-439d-96f6-d6bc489c04c5"
+version = "0.1.3"
+
+[[deps.CoordinateTransformations]]
+deps = ["LinearAlgebra", "StaticArrays"]
+git-tree-sha1 = "f9d7112bfff8a19a3a4ea4e03a8e6a91fe8456bf"
+uuid = "150eb455-5306-5404-9cee-2592286d6298"
+version = "0.6.3"
+
+[[deps.Crayons]]
+git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
+uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
+version = "4.1.1"
+
+[[deps.CustomUnitRanges]]
+git-tree-sha1 = "1a3f97f907e6dd8983b744d2642651bb162a3f7a"
+uuid = "dc8bdbbb-1ca9-579f-8c36-e416f6a65cce"
+version = "1.0.2"
+
+[[deps.DSP]]
+deps = ["Compat", "FFTW", "IterTools", "LinearAlgebra", "Polynomials", "Random", "Reexport", "SpecialFunctions", "Statistics"]
+git-tree-sha1 = "0df00546373af8eee1598fb4b2ba480b1ebe895c"
+uuid = "717857b8-e6f2-59f4-9121-6e50c889abd2"
+version = "0.7.10"
+
+[[deps.DataAPI]]
+git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe"
+uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
+version = "1.16.0"
+
+[[deps.DataFrames]]
+deps = ["Compat", "DataAPI", "DataStructures", "Future", "InlineStrings", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrecompileTools", "PrettyTables", "Printf", "Random", "Reexport", "SentinelArrays", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"]
+git-tree-sha1 = "fb61b4812c49343d7ef0b533ba982c46021938a6"
+uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+version = "1.7.0"
+
+[[deps.DataStructures]]
+deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
+git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82"
+uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+version = "0.18.20"
+
+[[deps.DataValueInterfaces]]
+git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
+uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
+version = "1.0.0"
+
+[[deps.Dates]]
+deps = ["Printf"]
+uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
+
+[[deps.DefineSingletons]]
+git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c"
+uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52"
+version = "0.1.2"
+
+[[deps.Dierckx]]
+deps = ["Dierckx_jll"]
+git-tree-sha1 = "d1ea9f433781bb6ff504f7d3cb70c4782c504a3a"
+uuid = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
+version = "0.5.3"
+
+[[deps.Dierckx_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "6596b96fe1caff3db36415eeb6e9d3b50bfe40ee"
+uuid = "cd4c43a9-7502-52ba-aa6d-59fb2a88580b"
+version = "0.1.0+0"
+
+[[deps.Distances]]
+deps = ["LinearAlgebra", "SparseArrays", "Statistics", "StatsAPI"]
+git-tree-sha1 = "c7e3a542b999843086e2f29dac96a618c105be1d"
+uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+version = "0.10.12"
+
+[[deps.Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[[deps.DistributedArrays]]
+deps = ["Distributed", "LinearAlgebra", "Primes", "Random", "Serialization", "SparseArrays", "Statistics"]
+git-tree-sha1 = "8a2950a2a7fcc87d9d9f07433a96bfa70394a3e6"
+uuid = "aaf54ef3-cdf8-58ed-94cc-d582ad619b94"
+version = "0.6.7"
+
+[[deps.DocStringExtensions]]
+deps = ["LibGit2"]
+git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d"
+uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+version = "0.9.3"
+
+[[deps.Downloads]]
+deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+version = "1.6.0"
+
+[[deps.DrWatson]]
+deps = ["Dates", "FileIO", "JLD2", "LibGit2", "MacroTools", "Pkg", "Random", "Requires", "Scratch", "UnPack"]
+git-tree-sha1 = "17c9db646d6cb3f68a90053bf46faf5312d13b36"
+uuid = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
+version = "2.18.0"
+
+[[deps.EnzymeCore]]
+deps = ["Adapt"]
+git-tree-sha1 = "8f205a601760f4798a10f138c3940f0451d95188"
+uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
+version = "0.7.8"
+
+[[deps.ExprTools]]
+git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec"
+uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+version = "0.1.10"
+
+[[deps.FFTViews]]
+deps = ["CustomUnitRanges", "FFTW"]
+git-tree-sha1 = "cbdf14d1e8c7c8aacbe8b19862e0179fd08321c2"
+uuid = "4f61f5a4-77b1-5117-aa51-3ab5ef4ef0cd"
+version = "0.3.2"
+
+[[deps.FFTW]]
+deps = ["AbstractFFTs", "FFTW_jll", "LinearAlgebra", "MKL_jll", "Preferences", "Reexport"]
+git-tree-sha1 = "4820348781ae578893311153d69049a93d05f39d"
+uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+version = "1.8.0"
+
+[[deps.FFTW_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "4d81ed14783ec49ce9f2e168208a12ce1815aa25"
+uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a"
+version = "3.3.10+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.FileIO]]
+deps = ["Pkg", "Requires", "UUIDs"]
+git-tree-sha1 = "62ca0547a14c57e98154423419d8a342dca75ca9"
+uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+version = "1.16.4"
+
+[[deps.FileWatching]]
+uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
+
+[[deps.FixedPointNumbers]]
+deps = ["Statistics"]
+git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172"
+uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
+version = "0.8.5"
+
+[[deps.Future]]
+deps = ["Random"]
+uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
+
+[[deps.GPUArrays]]
+deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"]
+git-tree-sha1 = "04661708f5301394a1f1be86a07a89e835900db6"
+uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+version = "10.2.3"
+
+[[deps.GPUArraysCore]]
+deps = ["Adapt"]
+git-tree-sha1 = "ec632f177c0d990e64d955ccc1b8c04c485a0950"
+uuid = "46192b85-c4d5-4398-a991-12ede77f4527"
+version = "0.1.6"
+
+[[deps.GPUCompiler]]
+deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Preferences", "Scratch", "Serialization", "TOML", "TimerOutputs", "UUIDs"]
+git-tree-sha1 = "ab29216184312f99ff957b32cd63c2fe9c928b91"
+uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
+version = "0.26.7"
+
+[[deps.Ghostscript_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "43ba3d3c82c18d88471cfd2924931658838c9d8f"
+uuid = "61579ee1-b43e-5ca0-a5da-69d92c66a64b"
+version = "9.55.0+4"
+
+[[deps.Graphics]]
+deps = ["Colors", "LinearAlgebra", "NaNMath"]
+git-tree-sha1 = "d61890399bc535850c4bf08e4e0d3a7ad0f21cbd"
+uuid = "a2bd30eb-e257-5431-a919-1863eab51364"
+version = "1.1.2"
+
+[[deps.IdentityRanges]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "be8fcd695c4da16a1d6d0cd213cb88090a150e3b"
+uuid = "bbac6d45-d8f3-5730-bfe4-7a449cd117ca"
+version = "0.3.1"
+
+[[deps.ImageAxes]]
+deps = ["AxisArrays", "ImageCore", "Reexport", "SimpleTraits"]
+git-tree-sha1 = "794ad1d922c432082bc1aaa9fa8ffbd1fe74e621"
+uuid = "2803e5a7-5153-5ecf-9a86-9b4c37f5f5ac"
+version = "0.6.9"
+
+[[deps.ImageContrastAdjustment]]
+deps = ["ColorVectorSpace", "ImageCore", "ImageTransformations", "Parameters"]
+git-tree-sha1 = "2e6084db6cccab11fe0bc3e4130bd3d117092ed9"
+uuid = "f332f351-ec65-5f6a-b3d1-319c6670881a"
+version = "0.3.7"
+
+[[deps.ImageCore]]
+deps = ["AbstractFFTs", "Colors", "FixedPointNumbers", "Graphics", "MappedArrays", "MosaicViews", "OffsetArrays", "PaddedViews", "Reexport"]
+git-tree-sha1 = "db645f20b59f060d8cfae696bc9538d13fd86416"
+uuid = "a09fc81d-aa75-5fe9-8630-4744c3626534"
+version = "0.8.22"
+
+[[deps.ImageDistances]]
+deps = ["ColorVectorSpace", "Distances", "ImageCore", "ImageMorphology", "LinearAlgebra", "Statistics"]
+git-tree-sha1 = "6378c34a3c3a216235210d19b9f495ecfff2f85f"
+uuid = "51556ac3-7006-55f5-8cb3-34580c88182d"
+version = "0.2.13"
+
+[[deps.ImageFiltering]]
+deps = ["CatIndices", "ColorVectorSpace", "ComputationalResources", "DataStructures", "FFTViews", "FFTW", "ImageCore", "LinearAlgebra", "OffsetArrays", "Requires", "SparseArrays", "StaticArrays", "Statistics", "TiledIteration"]
+git-tree-sha1 = "bf96839133212d3eff4a1c3a80c57abc7cfbf0ce"
+uuid = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
+version = "0.6.21"
+
+[[deps.ImageGather]]
+deps = ["JUDI", "LinearAlgebra"]
+git-tree-sha1 = "e7f8a3e4bafabc4154558c1f3cd6f3e5915a534f"
+uuid = "355d8124-6b2e-49b5-aab5-cdbc0a5fccbe"
+version = "0.3.0"
+
+[[deps.ImageIO]]
+deps = ["FileIO", "Netpbm", "PNGFiles"]
+git-tree-sha1 = "0d6d09c28d67611c68e25af0c2df7269c82b73c7"
+uuid = "82e4d734-157c-48bb-816b-45c225c6df19"
+version = "0.4.1"
+
+[[deps.ImageMagick]]
+deps = ["FileIO", "ImageCore", "ImageMagick_jll", "InteractiveUtils", "Libdl", "Pkg", "Random"]
+git-tree-sha1 = "5bc1cb62e0c5f1005868358db0692c994c3a13c6"
+uuid = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
+version = "1.2.1"
+
+[[deps.ImageMagick_jll]]
+deps = ["Artifacts", "Ghostscript_jll", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "OpenJpeg_jll", "Zlib_jll", "libpng_jll"]
+git-tree-sha1 = "d65554bad8b16d9562050c67e7223abf91eaba2f"
+uuid = "c73af94c-d91f-53ed-93a7-00f77d67a9d7"
+version = "6.9.13+0"
+
+[[deps.ImageMetadata]]
+deps = ["AxisArrays", "ColorVectorSpace", "ImageAxes", "ImageCore", "IndirectArrays"]
+git-tree-sha1 = "ae76038347dc4edcdb06b541595268fca65b6a42"
+uuid = "bc367c6b-8a6b-528e-b4bd-a4b897500b49"
+version = "0.9.5"
+
+[[deps.ImageMorphology]]
+deps = ["ColorVectorSpace", "ImageCore", "LinearAlgebra", "TiledIteration"]
+git-tree-sha1 = "68e7cbcd7dfaa3c2f74b0a8ab3066f5de8f2b71d"
+uuid = "787d08f9-d448-5407-9aad-5290dd7ab264"
+version = "0.2.11"
+
+[[deps.ImageQualityIndexes]]
+deps = ["ColorVectorSpace", "ImageCore", "ImageDistances", "ImageFiltering", "OffsetArrays", "Statistics"]
+git-tree-sha1 = "1198f85fa2481a3bb94bf937495ba1916f12b533"
+uuid = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
+version = "0.2.2"
+
+[[deps.ImageShow]]
+deps = ["Base64", "FileIO", "ImageCore", "OffsetArrays", "Requires", "StackViews"]
+git-tree-sha1 = "832abfd709fa436a562db47fd8e81377f72b01f9"
+uuid = "4e3cecfd-b093-5904-9786-8bbb286a6a31"
+version = "0.3.1"
+
+[[deps.ImageTransformations]]
+deps = ["AxisAlgorithms", "ColorVectorSpace", "CoordinateTransformations", "IdentityRanges", "ImageCore", "Interpolations", "OffsetArrays", "Rotations", "StaticArrays"]
+git-tree-sha1 = "e4cc551e4295a5c96545bb3083058c24b78d4cf0"
+uuid = "02fcd773-0e25-5acc-982a-7f6622650795"
+version = "0.8.13"
+
+[[deps.Images]]
+deps = ["AxisArrays", "Base64", "ColorVectorSpace", "FileIO", "Graphics", "ImageAxes", "ImageContrastAdjustment", "ImageCore", "ImageDistances", "ImageFiltering", "ImageIO", "ImageMagick", "ImageMetadata", "ImageMorphology", "ImageQualityIndexes", "ImageShow", "ImageTransformations", "IndirectArrays", "OffsetArrays", "Random", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "StatsBase", "TiledIteration"]
+git-tree-sha1 = "8b714d5e11c91a0d945717430ec20f9251af4bd2"
+uuid = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+version = "0.24.1"
+
+[[deps.IndirectArrays]]
+git-tree-sha1 = "c2a145a145dc03a7620af1444e0264ef907bd44f"
+uuid = "9b13fd28-a010-5f03-acff-a1bbcff69959"
+version = "0.5.1"
+
+[[deps.InitialValues]]
+git-tree-sha1 = "4da0f88e9a39111c2fa3add390ab15f3a44f3ca3"
+uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
+version = "0.3.1"
+
+[[deps.InlineStrings]]
+deps = ["Parsers"]
+git-tree-sha1 = "45521d31238e87ee9f9732561bfee12d4eebd52d"
+uuid = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
+version = "1.4.2"
+
+[[deps.InplaceOps]]
+deps = ["LinearAlgebra", "Test"]
+git-tree-sha1 = "50b41d59e7164ab6fda65e71049fee9d890731ff"
+uuid = "505f98c9-085e-5b2c-8e89-488be7bf1f34"
+version = "0.3.0"
+
+[[deps.IntegerMathUtils]]
+git-tree-sha1 = "b8ffb903da9f7b8cf695a8bead8e01814aa24b30"
+uuid = "18e54dd8-cb9d-406c-a71d-865a43cbb235"
+version = "0.1.2"
+
+[[deps.IntelOpenMP_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"]
+git-tree-sha1 = "10bd689145d2c3b2a9844005d01087cc1194e79e"
+uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0"
+version = "2024.2.1+0"
+
+[[deps.InteractiveUtils]]
+deps = ["Markdown"]
+uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+
+[[deps.Interpolations]]
+deps = ["AxisAlgorithms", "ChainRulesCore", "LinearAlgebra", "OffsetArrays", "Random", "Ratios", "Requires", "SharedArrays", "SparseArrays", "StaticArrays", "WoodburyMatrices"]
+git-tree-sha1 = "b7bc05649af456efc75d178846f47006c2c4c3c7"
+uuid = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+version = "0.13.6"
+
+[[deps.IntervalSets]]
+deps = ["Random", "RecipesBase", "Statistics"]
+git-tree-sha1 = "dba9ddf07f77f60450fe5d2e2beb9854d9a49bd0"
+uuid = "8197267c-284f-5f27-9208-e0e47529a953"
+version = "0.7.10"
+
+[[deps.InverseFunctions]]
+deps = ["Dates", "Test"]
+git-tree-sha1 = "a779299d77cd080bf77b97535acecd73e1c5e5cb"
+uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
+version = "0.1.17"
+
+[[deps.InvertedIndices]]
+git-tree-sha1 = "0dc7b50b8d436461be01300fd8cd45aa0274b038"
+uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
+version = "1.3.0"
+
+[[deps.IrrationalConstants]]
+git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2"
+uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
+version = "0.2.2"
+
+[[deps.IterTools]]
+git-tree-sha1 = "42d5f897009e7ff2cf88db414a389e5ed1bdd023"
+uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+version = "1.10.0"
+
+[[deps.IterativeSolvers]]
+deps = ["LinearAlgebra", "Printf", "Random", "RecipesBase", "SparseArrays"]
+git-tree-sha1 = "59545b0a2b27208b0650df0a46b8e3019f85055b"
+uuid = "42fd0dbc-a981-5370-80f2-aaf504508153"
+version = "0.9.4"
+
+[[deps.IteratorInterfaceExtensions]]
+git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
+uuid = "82899510-4779-5014-852e-03e436cf321d"
+version = "1.0.0"
+
+[[deps.JLD2]]
+deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"]
+git-tree-sha1 = "783c1be5213a09609b23237a0c9e5dfd258ae6f2"
+uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+version = "0.5.7"
+
+[[deps.JLLWrappers]]
+deps = ["Artifacts", "Preferences"]
+git-tree-sha1 = "be3dc50a92e5a386872a493a10050136d4703f9b"
+uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
+version = "1.6.1"
+
+[[deps.JOLI]]
+deps = ["ChainRulesCore", "Distributed", "DistributedArrays", "FFTW", "InplaceOps", "InteractiveUtils", "IterativeSolvers", "Libdl", "LinearAlgebra", "NFFT", "Nullables", "Printf", "PyCall", "Random", "SharedArrays", "SparseArrays", "SpecialFunctions", "Wavelets"]
+git-tree-sha1 = "97beb81d38c0928b58280f5daee8673546732ea4"
+uuid = "bb331ad6-a1cf-11e9-23da-9bcb53c69f6f"
+version = "0.8.5"
+
+[[deps.JSON]]
+deps = ["Dates", "Mmap", "Parsers", "Unicode"]
+git-tree-sha1 = "31e996f0a15c7b280ba9f76636b3ff9e2ae58c9a"
+uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+version = "0.21.4"
+
+[[deps.JUDI]]
+deps = ["ChainRulesCore", "DSP", "Dierckx", "Distributed", "FFTW", "JOLI", "LinearAlgebra", "OrderedCollections", "Printf", "PyCall", "Random", "Requires", "SegyIO"]
+git-tree-sha1 = "21f9b9ae041be1caba073bb496235bceaa1d2ff0"
+uuid = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
+version = "3.4.7"
+
+[[deps.JpegTurbo_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "25ee0be4d43d0269027024d75a24c24d6c6e590c"
+uuid = "aacddb02-875f-59d6-b918-886e6ef4fbf8"
+version = "3.0.4+0"
+
+[[deps.JuliaNVTXCallbacks_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "af433a10f3942e882d3c671aacb203e006a5808f"
+uuid = "9c1d0b0a-7046-5b2e-a33f-ea22f176ac7e"
+version = "0.2.1+0"
+
+[[deps.JuliaVariables]]
+deps = ["MLStyle", "NameResolution"]
+git-tree-sha1 = "49fb3cb53362ddadb4415e9b73926d6b40709e70"
+uuid = "b14d175d-62b4-44ba-8fb7-3064adc8c3ec"
+version = "0.2.4"
+
+[[deps.KernelAbstractions]]
+deps = ["Adapt", "Atomix", "EnzymeCore", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
+git-tree-sha1 = "e73a077abc7fe798fe940deabe30ef6c66bdde52"
+uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+version = "0.9.29"
+
+[[deps.LERC_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "36bdbc52f13a7d1dcb0f3cd694e01677a515655b"
+uuid = "88015f11-f218-50d7-93a8-a6af411a945d"
+version = "4.0.0+0"
+
+[[deps.LLVM]]
+deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"]
+git-tree-sha1 = "020abd49586480c1be84f57da0017b5d3db73f7c"
+uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
+version = "8.0.0"
+
+[[deps.LLVMExtra_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "c2636c264861edc6d305e6b4d528f09566d24c5e"
+uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
+version = "0.0.30+0"
+
+[[deps.LLVMLoopInfo]]
+git-tree-sha1 = "2e5c102cfc41f48ae4740c7eca7743cc7e7b75ea"
+uuid = "8b046642-f1f6-4319-8d3c-209ddc03c586"
+version = "1.0.0"
+
+[[deps.LaTeXStrings]]
+git-tree-sha1 = "dda21b8cbd6a6c40d9d02a73230f9d70fed6918c"
+uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+version = "1.4.0"
+
+[[deps.LazyArtifacts]]
+deps = ["Artifacts", "Pkg"]
+uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+
+[[deps.LearnBase]]
+git-tree-sha1 = "a0d90569edd490b82fdc4dc078ea54a5a800d30a"
+uuid = "7f8f8fb0-2700-5f03-b4bd-41f8cfc144b6"
+version = "0.4.1"
+
+[[deps.LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+version = "0.6.3"
+
+[[deps.LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
+version = "7.84.0+0"
+
+[[deps.LibGit2]]
+deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
+uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
+
+[[deps.LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+version = "1.10.2+0"
+
+[[deps.Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+
+[[deps.Libtiff_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "LERC_jll", "Libdl", "XZ_jll", "Zlib_jll", "Zstd_jll"]
+git-tree-sha1 = "b404131d06f7886402758c9ce2214b636eb4d54a"
+uuid = "89763e89-9b03-5906-acba-b20f662cd828"
+version = "4.7.0+0"
+
+[[deps.LinearAlgebra]]
+deps = ["Libdl", "libblastrampoline_jll"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[[deps.LittleCMS_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll"]
+git-tree-sha1 = "fa7fd067dca76cadd880f1ca937b4f387975a9f5"
+uuid = "d3a379c0-f9a3-5b72-a4c0-6bf4d2e8af0f"
+version = "2.16.0+0"
+
+[[deps.LogExpFunctions]]
+deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
+git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea"
+uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+version = "0.3.28"
+
+[[deps.Logging]]
+uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+
+[[deps.MKL_jll]]
+deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"]
+git-tree-sha1 = "f046ccd0c6db2832a9f639e2c669c6fe867e5f4f"
+uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
+version = "2024.2.0+0"
+
+[[deps.MLDataPattern]]
+deps = ["LearnBase", "MLLabelUtils", "Random", "SparseArrays", "StatsBase"]
+git-tree-sha1 = "2741aabe5ec93c63cc0a5277bb1c2cb819b01bb1"
+uuid = "9920b226-0b2a-5f5f-9153-9aa70a013f8b"
+version = "0.5.5"
+
+[[deps.MLLabelUtils]]
+deps = ["LearnBase", "MappedArrays", "StatsBase"]
+git-tree-sha1 = "fd75d4b0c4016e047bbb6263eecf7ae3891af522"
+uuid = "66a33bbf-0c2b-5fc8-a008-9da813334f0a"
+version = "0.5.7"
+
+[[deps.MLStyle]]
+git-tree-sha1 = "bc38dff0548128765760c79eb7388a4b37fae2c8"
+uuid = "d8e11817-5142-5d16-987a-aa16d5891078"
+version = "0.4.17"
+
+[[deps.MacroTools]]
+deps = ["Markdown", "Random"]
+git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df"
+uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.5.13"
+
+[[deps.MappedArrays]]
+git-tree-sha1 = "2dab0221fe2b0f2cb6754eaa743cc266339f527e"
+uuid = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
+version = "0.4.2"
+
+[[deps.Markdown]]
+deps = ["Base64"]
+uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
+
+[[deps.MbedTLS_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
+version = "2.28.0+0"
+
+[[deps.Memoize]]
+deps = ["MacroTools"]
+git-tree-sha1 = "2b1dfcba103de714d31c033b5dacc2e4a12c7caa"
+uuid = "c03570c3-d221-55d1-a50c-7939bbd78826"
+version = "0.4.4"
+
+[[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"
+uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
+version = "1.2.0"
+
+[[deps.Mmap]]
+uuid = "a63ad114-7e13-5084-954f-fe012c677804"
+
+[[deps.MosaicViews]]
+deps = ["MappedArrays", "OffsetArrays", "PaddedViews", "StackViews"]
+git-tree-sha1 = "7b86a5d4d70a9f5cdf2dacb3cbe6d251d1a61dbe"
+uuid = "e94cdb99-869f-56ef-bcf0-1ae2bcbe0389"
+version = "0.3.4"
+
+[[deps.MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+version = "2022.2.1"
+
+[[deps.NFFT]]
+deps = ["AbstractNFFTs", "BasicInterpolators", "Distributed", "FFTW", "FLoops", "LinearAlgebra", "Printf", "Random", "Reexport", "SnoopPrecompile", "SparseArrays", "SpecialFunctions"]
+git-tree-sha1 = "93a5f32dd6cf09456b0b81afcb8fc29f06535ffd"
+uuid = "efe261a4-0d2b-5849-be55-fc731d526b0d"
+version = "0.13.3"
+
+[[deps.NVTX]]
+deps = ["Colors", "JuliaNVTXCallbacks_jll", "Libdl", "NVTX_jll"]
+git-tree-sha1 = "53046f0483375e3ed78e49190f1154fa0a4083a1"
+uuid = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
+version = "0.3.4"
+
+[[deps.NVTX_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "ce3269ed42816bf18d500c9f63418d4b0d9f5a3b"
+uuid = "e98f9f5b-d649-5603-91fd-7774390e6439"
+version = "3.1.0+2"
+
+[[deps.NaNMath]]
+deps = ["OpenLibm_jll"]
+git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4"
+uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+version = "1.0.2"
+
+[[deps.NameResolution]]
+deps = ["PrettyPrint"]
+git-tree-sha1 = "1a0fa0e9613f46c9b8c11eee38ebb4f590013c5e"
+uuid = "71a1bf82-56d0-4bbc-8a3c-48b961074391"
+version = "0.1.5"
+
+[[deps.Netpbm]]
+deps = ["ColorVectorSpace", "FileIO", "ImageCore"]
+git-tree-sha1 = "09589171688f0039f13ebe0fdcc7288f50228b52"
+uuid = "f09324ee-3d7c-5217-9330-fc30815ba969"
+version = "1.0.1"
+
+[[deps.NetworkOptions]]
+uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
+version = "1.2.0"
+
+[[deps.Nullables]]
+git-tree-sha1 = "8f87854cc8f3685a60689d8edecaa29d2251979b"
+uuid = "4d1e1d77-625e-5b40-9113-a560ec7a8ecd"
+version = "1.0.0"
+
+[[deps.OffsetArrays]]
+deps = ["Adapt"]
+git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e"
+uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+version = "1.14.1"
+
+[[deps.OpenBLAS_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
+uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
+version = "0.3.20+0"
+
+[[deps.OpenJpeg_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "libpng_jll"]
+git-tree-sha1 = "f4cb457ffac5f5cf695699f82c537073958a6a6c"
+uuid = "643b3616-a352-519d-856d-80112ee9badc"
+version = "2.5.2+0"
+
+[[deps.OpenLibm_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
+version = "0.8.1+0"
+
+[[deps.OpenSpecFun_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
+uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
+version = "0.5.5+0"
+
+[[deps.OrderedCollections]]
+git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5"
+uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+version = "1.6.3"
+
+[[deps.PNGFiles]]
+deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"]
+git-tree-sha1 = "f809158b27eba0c18c269cf2a2be6ed751d3e81d"
+uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883"
+version = "0.3.17"
+
+[[deps.PaddedViews]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "0fac6313486baae819364c52b4f483450a9d793f"
+uuid = "5432bcbf-9aad-5242-b902-cca2824c8663"
+version = "0.5.12"
+
+[[deps.Parameters]]
+deps = ["OrderedCollections", "UnPack"]
+git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
+uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+version = "0.12.3"
+
+[[deps.Parsers]]
+deps = ["Dates", "PrecompileTools", "UUIDs"]
+git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821"
+uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
+version = "2.8.1"
+
+[[deps.Pkg]]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+version = "1.8.0"
+
+[[deps.Polynomials]]
+deps = ["LinearAlgebra", "RecipesBase", "Requires", "Setfield", "SparseArrays"]
+git-tree-sha1 = "1a9cfb2dc2c2f1bd63f1906d72af39a79b49b736"
+uuid = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
+version = "4.0.11"
+
+[[deps.PooledArrays]]
+deps = ["DataAPI", "Future"]
+git-tree-sha1 = "36d8b4b899628fb92c2749eb488d884a926614d3"
+uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
+version = "1.4.3"
+
+[[deps.PrecompileTools]]
+deps = ["Preferences"]
+git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f"
+uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+version = "1.2.1"
+
+[[deps.Preferences]]
+deps = ["TOML"]
+git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6"
+uuid = "21216c6a-2e73-6563-6e65-726566657250"
+version = "1.4.3"
+
+[[deps.PrettyPrint]]
+git-tree-sha1 = "632eb4abab3449ab30c5e1afaa874f0b98b586e4"
+uuid = "8162dcfd-2161-5ef2-ae6c-7681170c5f98"
+version = "0.2.0"
+
+[[deps.PrettyTables]]
+deps = ["Crayons", "LaTeXStrings", "Markdown", "PrecompileTools", "Printf", "Reexport", "StringManipulation", "Tables"]
+git-tree-sha1 = "66b20dd35966a748321d3b2537c4584cf40387c7"
+uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
+version = "2.3.2"
+
+[[deps.Primes]]
+deps = ["IntegerMathUtils"]
+git-tree-sha1 = "cb420f77dc474d23ee47ca8d14c90810cafe69e7"
+uuid = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
+version = "0.5.6"
+
+[[deps.Printf]]
+deps = ["Unicode"]
+uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[[deps.PyCall]]
+deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
+git-tree-sha1 = "9816a3826b0ebf49ab4926e2b18842ad8b5c8f04"
+uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+version = "1.96.4"
+
+[[deps.PyPlot]]
+deps = ["Colors", "LaTeXStrings", "PyCall", "Sockets", "Test", "VersionParsing"]
+git-tree-sha1 = "0371ca706e3f295481cbf94c8c36692b072285c2"
+uuid = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+version = "2.11.5"
+
+[[deps.Quaternions]]
+deps = ["LinearAlgebra", "Random", "RealDot"]
+git-tree-sha1 = "994cc27cdacca10e68feb291673ec3a76aa2fae9"
+uuid = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
+version = "0.7.6"
+
+[[deps.REPL]]
+deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
+uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+
+[[deps.Random]]
+deps = ["SHA", "Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[deps.Random123]]
+deps = ["Random", "RandomNumbers"]
+git-tree-sha1 = "4743b43e5a9c4a2ede372de7061eed81795b12e7"
+uuid = "74087812-796a-5b5d-8853-05524746bad3"
+version = "1.7.0"
+
+[[deps.RandomNumbers]]
+deps = ["Random"]
+git-tree-sha1 = "c6ec94d2aaba1ab2ff983052cf6a606ca5985902"
+uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143"
+version = "1.6.0"
+
+[[deps.RangeArrays]]
+git-tree-sha1 = "b9039e93773ddcfc828f12aadf7115b4b4d225f5"
+uuid = "b3c3ace0-ae52-54e7-9d0b-2c1406fd6b9d"
+version = "0.3.2"
+
+[[deps.Ratios]]
+deps = ["Requires"]
+git-tree-sha1 = "1342a47bf3260ee108163042310d26f2be5ec90b"
+uuid = "c84ed2f1-dad5-54f0-aa8e-dbefe2724439"
+version = "0.4.5"
+
+[[deps.RealDot]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "9f0a1b71baaf7650f4fa8a1d168c7fb6ee41f0c9"
+uuid = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9"
+version = "0.1.0"
+
+[[deps.RecipesBase]]
+deps = ["PrecompileTools"]
+git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff"
+uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+version = "1.3.4"
+
+[[deps.Reexport]]
+git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
+uuid = "189a3867-3050-52da-a836-e630ba90ab69"
+version = "1.2.2"
+
+[[deps.Requires]]
+deps = ["UUIDs"]
+git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
+uuid = "ae029012-a4dd-5104-9daa-d747884805df"
+version = "1.3.0"
+
+[[deps.Rotations]]
+deps = ["LinearAlgebra", "Quaternions", "Random", "RecipesBase", "StaticArrays"]
+git-tree-sha1 = "5680a9276685d392c87407df00d57c9924d9f11e"
+uuid = "6038ab10-8711-5258-84ad-4b1120ba62dc"
+version = "1.7.1"
+
+[[deps.SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+version = "0.7.0"
+
+[[deps.Scratch]]
+deps = ["Dates"]
+git-tree-sha1 = "3bac05bc7e74a75fd9cba4295cde4045d9fe2386"
+uuid = "6c6a2e73-6563-6170-7368-637461726353"
+version = "1.2.1"
+
+[[deps.SegyIO]]
+deps = ["Distributed", "Printf", "Test"]
+git-tree-sha1 = "0fc24db28695a80aa59c179372b11165d371188a"
+uuid = "157a0f19-4d44-4de5-a0d0-07e2f0ac4dfa"
+version = "0.8.5"
+
+[[deps.SentinelArrays]]
+deps = ["Dates", "Random"]
+git-tree-sha1 = "305becf8af67eae1dbc912ee9097f00aeeabb8d5"
+uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
+version = "1.4.6"
+
+[[deps.Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[deps.Setfield]]
+deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"]
+git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac"
+uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+version = "1.1.1"
+
+[[deps.SharedArrays]]
+deps = ["Distributed", "Mmap", "Random", "Serialization"]
+uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
+
+[[deps.SimpleTraits]]
+deps = ["InteractiveUtils", "MacroTools"]
+git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231"
+uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
+version = "0.9.4"
+
+[[deps.SlimPlotting]]
+deps = ["ColorSchemes", "Requires", "Statistics"]
+git-tree-sha1 = "e9b0658e186ce6bc020ca6bf24305fb835716470"
+uuid = "f6d04670-764e-495b-a720-91c3c9a588ff"
+version = "1.0.0"
+
+[[deps.SnoopPrecompile]]
+deps = ["Preferences"]
+git-tree-sha1 = "e760a70afdcd461cf01a575947738d359234665c"
+uuid = "66db9d55-30c0-4569-8b51-7e840670fc0c"
+version = "1.0.3"
+
+[[deps.Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+
+[[deps.SortingAlgorithms]]
+deps = ["DataStructures"]
+git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085"
+uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
+version = "1.2.1"
+
+[[deps.SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[deps.SpecialFunctions]]
+deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
+git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14"
+uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
+version = "2.4.0"
+
+[[deps.SplittablesBase]]
+deps = ["Setfield", "Test"]
+git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5"
+uuid = "171d559e-b47b-412a-8079-5efa626c420e"
+version = "0.1.15"
+
+[[deps.StackViews]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "46e589465204cd0c08b4bd97385e4fa79a0c770c"
+uuid = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
+version = "0.1.1"
+
+[[deps.StaticArrays]]
+deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore", "Statistics"]
+git-tree-sha1 = "777657803913ffc7e8cc20f0fd04b634f871af8f"
+uuid = "90137ffa-7385-5640-81b9-e52037218182"
+version = "1.9.8"
+
+[[deps.StaticArraysCore]]
+git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682"
+uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+version = "1.4.3"
+
+[[deps.Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[deps.StatsAPI]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed"
+uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
+version = "1.7.0"
+
+[[deps.StatsBase]]
+deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
+git-tree-sha1 = "d1bf48bfcc554a3761a133fe3a9bb01488e06916"
+uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+version = "0.33.21"
+
+[[deps.StringManipulation]]
+deps = ["PrecompileTools"]
+git-tree-sha1 = "a04cabe79c5f01f4d723cc6704070ada0b9d46d5"
+uuid = "892a3eda-7b42-436c-8928-eab12a02cf0e"
+version = "0.3.4"
+
+[[deps.TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+version = "1.0.0"
+
+[[deps.TableTraits]]
+deps = ["IteratorInterfaceExtensions"]
+git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
+uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
+version = "1.0.1"
+
+[[deps.Tables]]
+deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"]
+git-tree-sha1 = "598cd7c1f68d1e205689b1c2fe65a9f85846f297"
+uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+version = "1.12.0"
+
+[[deps.Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+version = "1.10.1"
+
+[[deps.TensorCore]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6"
+uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50"
+version = "0.1.1"
+
+[[deps.Test]]
+deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
+uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[[deps.TextWrap]]
+git-tree-sha1 = "43044b737fa70bc12f6105061d3da38f881a3e3c"
+uuid = "b718987f-49a8-5099-9789-dcd902bef87d"
+version = "1.0.2"
+
+[[deps.TiledIteration]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "5683455224ba92ef59db72d10690690f4a8dc297"
+uuid = "06e1c1a7-607b-532d-9fad-de7d9aa2abac"
+version = "0.3.1"
+
+[[deps.TimerOutputs]]
+deps = ["ExprTools", "Printf"]
+git-tree-sha1 = "3a6f063d690135f5c1ba351412c82bae4d1402bf"
+uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+version = "0.5.25"
+
+[[deps.TranscodingStreams]]
+git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
+uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+version = "0.11.3"
+
+[[deps.Transducers]]
+deps = ["Accessors", "Adapt", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "ConstructionBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "SplittablesBase", "Tables"]
+git-tree-sha1 = "7deeab4ff96b85c5f72c824cae53a1398da3d1cb"
+uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999"
+version = "0.4.84"
+
+[[deps.UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
+[[deps.UnPack]]
+git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
+uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+version = "1.0.2"
+
+[[deps.Unicode]]
+uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[deps.UnsafeAtomics]]
+git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278"
+uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f"
+version = "0.2.1"
+
+[[deps.UnsafeAtomicsLLVM]]
+deps = ["LLVM", "UnsafeAtomics"]
+git-tree-sha1 = "bf2c553f25e954a9b38c9c0593a59bb13113f9e5"
+uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
+version = "0.1.5"
+
+[[deps.VersionParsing]]
+git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.3.0"
+
+[[deps.Wavelets]]
+deps = ["DSP", "FFTW", "LinearAlgebra", "Reexport", "SpecialFunctions", "Statistics"]
+git-tree-sha1 = "58f7491c21ab2b1d69368c7f7e8a6a93cbf8b7bf"
+uuid = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
+version = "0.9.5"
+
+[[deps.WoodburyMatrices]]
+deps = ["LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "5f24e158cf4cee437052371455fe361f526da062"
+uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6"
+version = "0.5.6"
+
+[[deps.XZ_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "15e637a697345f6743674f1322beefbc5dcd5cfc"
+uuid = "ffd25f8a-64ca-5728-b0f7-c24cf3aae800"
+version = "5.6.3+0"
+
+[[deps.Zlib_jll]]
+deps = ["Libdl"]
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+version = "1.2.12+3"
+
+[[deps.Zstd_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "555d1076590a6cc2fdee2ef1469451f872d8b41b"
+uuid = "3161d3a3-bdf6-5164-811a-617609db77b4"
+version = "1.5.6+1"
+
+[[deps.libblastrampoline_jll]]
+deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
+uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"
+version = "5.1.1+0"
+
+[[deps.libpng_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"]
+git-tree-sha1 = "b70c870239dc3d7bc094eb2d6be9b73d27bef280"
+uuid = "b53b4c65-9356-5827-b1ea-8c7a1a84506f"
+version = "1.6.44+0"
+
+[[deps.nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+version = "1.48.0+0"
+
+[[deps.oneTBB_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "7d0ea0f4895ef2f5cb83645fa689e52cb55cf493"
+uuid = "1317d2d5-d96f-522e-a858-c73665f53c3e"
+version = "2021.12.0+0"
+
+[[deps.p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
+version = "17.4.0+0"
diff --git a/CIG/Project.toml b/CIG/Project.toml
new file mode 100644
index 0000000..6b77f6e
--- /dev/null
+++ b/CIG/Project.toml
@@ -0,0 +1,14 @@
+[deps]
+ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+Augmentor = "02898b10-1f73-11ea-317c-6393d7073e15"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+DrWatson = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+ImageGather = "355d8124-6b2e-49b5-aab5-cdbc0a5fccbe"
+Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
+PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+SegyIO = "157a0f19-4d44-4de5-a0d0-07e2f0ac4dfa"
+SlimPlotting = "f6d04670-764e-495b-a720-91c3c9a588ff"
diff --git a/CIG/collect_cigs.jl b/CIG/collect_cigs.jl
new file mode 100644
index 0000000..fce791b
--- /dev/null
+++ b/CIG/collect_cigs.jl
@@ -0,0 +1,25 @@
+using JLD2
+
+include("../config.jl")
+
+params = Config.get_parameters()
+
+j, nsamples = parse.(Int, ARGS[1:2])
+n_offsets = params["n_offsets"]
+down_rate = params["down_rate"]
+nx = params["nx"] ÷ down_rate
+nz = params["nz"] ÷ down_rate
+
+CIGs = zeros(Float32, n_offsets, nx, nz, 1, nsamples)
+
+for i in 1:nsamples
+    @load "data/$j/sample_$i.jld2" rtm
+    CIGs[:, :, :, 1, i] = rtm
+end
+
+JLD2.@save "data/cigs_iteration_j=$j.jld2" CIGs
+
+CIG0 = CIGs[:, :, :, 1:1, 1:1]
+
+@load "data/posteriors_iteration_j=$(j-1).jld2" x0
+JLD2.@save "data/posteriors_iteration_j=$(j-1).jld2" x0 CIG0
diff --git a/CIG/simulate.jl b/CIG/simulate.jl
new file mode 100644
index 0000000..64946f7
--- /dev/null
+++ b/CIG/simulate.jl
@@ -0,0 +1,169 @@
+
+# import Pkg; Pkg.instantiate()
+using DrWatson
+using JLD2, JUDI, SegyIO, ImageGather
+using ArgParse
+using Statistics, Images
+using FFTW
+using LinearAlgebra
+using Augmentor
+try
+    @eval using PyPlot  # Try to load PyPlot the first time
+catch e
+    @warn "Failed to load PyPlot on the first attempt: $e"
+    @info "Retrying to load PyPlot..."
+    @eval using PyPlot  # Retry loading PyPlot
+end
+using PyPlot, SlimPlotting
+using Random
+
+seed = 1
+Random.seed!(seed)
+
+include("../config.jl")
+
+j, n_samples, rank, comm_size = parse.(Int, ARGS[1:4])
+
+# Calculate the number of samples each process should handle
+samples_per_process = div(n_samples, comm_size)
+remainder = mod(n_samples, comm_size)
+
+# Determine the start and end indices for each rank
+if rank < remainder
+    n_start = rank * (samples_per_process + 1)
+    n_end = n_start + samples_per_process
+else
+    n_start = rank * samples_per_process + remainder
+    n_end = n_start + samples_per_process - 1
+end
+
+n_start = n_start + 1
+n_end = n_end + 1
+
+println("Process $rank handling range: $n_start to $n_end")
+
+PyPlot.rcdefaults()
+
+params = Config.get_parameters()
+n_offsets = params["n_offsets"]
+offset_start = params["offset_start"]
+offset_end = params["offset_end"]
+f0 = params["f0"]
+timeD = params["timeD"]
+timeR = params["timeR"]
+TD = params["TD"]
+dtD = params["dtD"]
+dtS = params["dtS"]
+nbl = params["nbl"]
+down_rate = params["down_rate"]
+
+plot_path = joinpath("plots", "CIG")
+
+function ContJitter(l::Number, num::Int)
+    #l = length, num = number of samples
+    interval_width = l/num
+    interval_center = range(interval_width/2, stop = l-interval_width/2, length=num)
+    randomshift = interval_width .* rand(Float32, num) .- interval_width/2
+    return interval_center .+ randomshift
+end
+
+offsetrange = range(offset_start, stop=offset_end, length=n_offsets)
+wavelet = ricker_wavelet(TD, dtS, f0)
+wavelet = filter_data(wavelet, dtS; fmin=3f0, fmax=Inf)
+d = (12.5f0, 12.5f0)
+o = (0f0, 0f0)
+
+# use no salt (after smoothing) as the background
+@load "data/posteriors_iteration_j=$(j-1).jld2" x0
+x_no_salt = x0
+
+# use salt (after smoothing) as the target
+@load "data/initial.jld2" x
+x_salt = x
+
+# Down Sampled velocity models in scripts/prepare_training.jl
+n = (size(x_salt)[1], size(x_salt)[2])
+f0 = f0 / down_rate
+d = d .* down_rate
+
+# Setup model structure
+nsrc = 16    # number of sources
+model = Model(n, d, o, (1f0./imresize(x_salt[:,:,1,1], n)).^2f0; nb=nbl)
+nxrec = n[1]
+xrec = range(0f0, stop=(n[1]-1)*d[1], length=nxrec)
+yrec = 0f0 # WE have to set the y coordiante to zero (or any number) for 2D modeling
+zrec = range(d[1], stop=d[1], length=nxrec)
+# Set up receiver structure
+recGeometry = Geometry(xrec, yrec, zrec; dt=dtD, t=timeD, nsrc=nsrc)
+wb = 16
+ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
+zsrc = convertToCell(range((wb-1)*d[1], stop=(wb-1)*d[1], length=nsrc))
+snr = 12f0
+
+# Setup operators
+gaussian = 20
+
+for i in n_start:n_end
+    filename = "data/$j/sample_$i.jld2"
+    if isfile(filename)
+        @info "Skipping sample $i as file $filename already exists."
+        continue  # Skip this iteration and move to the next one
+    end
+
+    Base.flush(Base.stdout)
+    @info "sample $i out of $(size(x_salt)[end]) samples"
+    # call function that generate background
+    x_back = x_no_salt[:,:,1,i];
+    x_salt_i= x_salt[:,:,1,i]
+
+    # add water bottom
+    nwb = 15
+    x_back[:, 1:nwb] .= 1.48
+    x_salt_i[:, 1:nwb] .= 1.48
+
+    # Set up source structure
+    xsrc = convertToCell(ContJitter((n[1]-1)*d[1], nsrc))
+    srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtD, t=timeD)
+    q = judiVector(srcGeometry, wavelet)
+    opt = Options(isic=true)
+    F = judiModeling(model, srcGeometry, recGeometry, options=opt)
+    @time d_obs = F(1f0./x_salt_i.^2f0) * q
+    J = judiExtendedJacobian(F(1f0./x_back.^2f0), q, offsetrange)
+    d_obs0 = F(1f0./x_back.^2f0) * q
+    noise_ = deepcopy(d_obs)
+    for l = 1:nsrc
+        noise_.data[l] = randn(Float32, size(d_obs.data[l]))
+        noise_.data[l] = real.(ifft(fft(noise_.data[l]).*fft(q.data[1])))
+    end
+    noise_ = noise_/norm(noise_) * norm(d_obs) * 10f0^(-snr/20f0)
+    @time rtm = J' * (d_obs0 - (d_obs + noise_))
+
+    # save_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o q d_obs i rtm snr offset_start offset_end n_offsets
+    # @tagsave(
+    #         joinpath(joinpath(plot_path, "cig"), savename(save_dict, "jld2"; digits=6)),
+    #         save_dict;
+    #         safe=true
+    #     );
+
+    rtm[:, :, 1:20] .= 0;
+    for z = 1:n[2]
+        rtm[:,:,z] .*= z * d[2]
+    end
+
+    JLD2.@save filename rtm
+
+    if false
+        plot_velocity(x_back', d, perc=95, vmax=4.8, aspect=true, name="idx=$i"); 
+        savefig(joinpath(plot_path, "x_back_idx=$(i)"), bbox_inches="tight", dpi=300);
+        plot_velocity(x_salt_i', d, perc=95, vmax=4.8, aspect=true, name="idx=$i"); 
+        savefig(joinpath(plot_path, "x_target_idx=$(i)"), bbox_inches="tight", dpi=300);
+        rtm[:, :, 1:20] .= 0;
+        for z = 1:n[2]
+            rtm[:,:,z] .*= z * d[2]
+        end
+
+        cig_img_suffix = "_gaussian=$(gaussian)"
+        cig_img_fname = "cig_idx=$(i)" * cig_img_suffix
+        plot_cig(rtm, plot_path, cig_img_fname)
+    end
+end
diff --git a/CIG/test.jl b/CIG/test.jl
new file mode 100644
index 0000000..11d9c9d
--- /dev/null
+++ b/CIG/test.jl
@@ -0,0 +1,48 @@
+using DrWatson
+using Random
+using JLD2
+
+j, n_samples, n_procs = parse.(Int, ARGS[1:3])
+temp_dir = "data/temp"
+
+random_number = rand(1:10000)
+data = @strdict random_number
+file_path = joinpath(temp_dir, savename(data, "jld2"; digits=6))
+@JLD2.save file_path random_number
+
+println("Process with random number $random_number saved as $file_path")
+
+while length(readdir(temp_dir)) < n_procs
+    sleep(0.1)
+end
+
+numbers = []
+for file in readdir(temp_dir)
+    @load joinpath(temp_dir, file) random_number
+    push!(numbers, random_number)
+end
+
+sorted_numbers = sort(numbers)
+rank = findfirst(x -> x == random_number, sorted_numbers) - 1 
+
+# Calculate the number of samples each process should handle
+samples_per_process = div(n_samples, n_procs)
+remainder = mod(n_samples, n_procs)
+
+# Determine the start and end indices for this process
+function get_sample_range(rank, samples_per_process, remainder)
+    if rank < remainder
+        n_start = rank * (samples_per_process + 1)
+        n_end = n_start + samples_per_process
+    else
+        n_start = rank * samples_per_process + remainder
+        n_end = n_start + samples_per_process - 1
+    end
+    return n_start + 1, n_end + 1
+end
+
+n_start, n_end = get_sample_range(rank, samples_per_process, remainder)
+println("Process rank $rank handling range: $n_start to $n_end")
+
+# # Clean up temporary directory (optional, if needed)
+# rm(temp_dir, force=true)
diff --git a/CNF/Manifest.toml b/CNF/Manifest.toml
new file mode 100644
index 0000000..2b9eb5b
--- /dev/null
+++ b/CNF/Manifest.toml
@@ -0,0 +1,1360 @@
+# This file is machine-generated - editing it directly is not advised
+
+julia_version = "1.8.5"
+manifest_format = "2.0"
+project_hash = "7e8b35736817d441122f9bb3f8208e3c8a7e5051"
+
+[[deps.AbstractFFTs]]
+deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
+git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef"
+uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+version = "1.5.0"
+
+[[deps.AbstractNFFTs]]
+deps = ["LinearAlgebra", "Printf"]
+git-tree-sha1 = "292e21e99dedb8621c15f185b8fdb4260bb3c429"
+uuid = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
+version = "0.8.2"
+
+[[deps.AbstractTrees]]
+git-tree-sha1 = "03e0550477d86222521d254b741d470ba17ea0b5"
+uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+version = "0.3.4"
+
+[[deps.Accessors]]
+deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown", "Requires", "Test"]
+git-tree-sha1 = "b392ede862e506d451fc1616e79aa6f4c673dab8"
+uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+version = "0.1.38"
+
+[[deps.Adapt]]
+deps = ["LinearAlgebra", "Requires"]
+git-tree-sha1 = "cde29ddf7e5726c9fb511f340244ea3481267608"
+uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+version = "3.7.2"
+
+[[deps.AliasTables]]
+deps = ["PtrArrays", "Random"]
+git-tree-sha1 = "9876e1e164b144ca45e9e3198d0b689cadfed9ff"
+uuid = "66dad0bd-aa9a-41b7-9441-69ab47430ed8"
+version = "1.1.3"
+
+[[deps.ArgCheck]]
+git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4"
+uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+version = "2.3.0"
+
+[[deps.ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+version = "1.1.1"
+
+[[deps.ArrayInterface]]
+deps = ["Compat", "IfElse", "LinearAlgebra", "Requires", "SparseArrays", "Static"]
+git-tree-sha1 = "81f0cb60dc994ca17f68d9fb7c942a5ae70d9ee4"
+uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+version = "5.0.8"
+
+[[deps.Artifacts]]
+uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+
+[[deps.Atomix]]
+deps = ["UnsafeAtomics"]
+git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be"
+uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
+version = "0.1.0"
+
+[[deps.Augmentor]]
+deps = ["ComputationalResources", "CoordinateTransformations", "FileIO", "IdentityRanges", "ImageCore", "ImageFiltering", "ImageTransformations", "InteractiveUtils", "Interpolations", "LinearAlgebra", "MLDataPattern", "OffsetArrays", "Rotations", "StaticArrays", "Statistics"]
+git-tree-sha1 = "47cd3b6a76a132c899ed798bb72cc1f1fe60f8d5"
+uuid = "02898b10-1f73-11ea-317c-6393d7073e15"
+version = "0.6.6"
+
+[[deps.AxisAlgorithms]]
+deps = ["LinearAlgebra", "Random", "SparseArrays", "WoodburyMatrices"]
+git-tree-sha1 = "66771c8d21c8ff5e3a93379480a2307ac36863f7"
+uuid = "13072b0f-2c55-5437-9ae7-d433b7a33950"
+version = "1.0.1"
+
+[[deps.AxisArrays]]
+deps = ["Dates", "IntervalSets", "IterTools", "RangeArrays"]
+git-tree-sha1 = "16351be62963a67ac4083f748fdb3cca58bfd52f"
+uuid = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+version = "0.4.7"
+
+[[deps.BFloat16s]]
+deps = ["LinearAlgebra", "Printf", "Random", "Test"]
+git-tree-sha1 = "a598ecb0d717092b5539dbbe890c98bac842b072"
+uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
+version = "0.2.0"
+
+[[deps.BSON]]
+git-tree-sha1 = "4c3e506685c527ac6a54ccc0c8c76fd6f91b42fb"
+uuid = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+version = "0.3.9"
+
+[[deps.BangBang]]
+deps = ["Accessors", "ConstructionBase", "InitialValues", "LinearAlgebra", "Requires"]
+git-tree-sha1 = "e2144b631226d9eeab2d746ca8880b7ccff504ae"
+uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+version = "0.4.3"
+
+[[deps.Base64]]
+uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+
+[[deps.Baselet]]
+git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e"
+uuid = "9718e550-a3fa-408a-8086-8db961cd8217"
+version = "0.1.1"
+
+[[deps.BasicInterpolators]]
+deps = ["LinearAlgebra", "Memoize", "Random"]
+git-tree-sha1 = "3f7be532673fc4a22825e7884e9e0e876236b12a"
+uuid = "26cce99e-4866-4b6d-ab74-862489e035e0"
+version = "0.7.1"
+
+[[deps.CEnum]]
+git-tree-sha1 = "eb4cb44a499229b3b8426dcfb5dd85333951ff90"
+uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+version = "0.4.2"
+
+[[deps.CUDA]]
+deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CompilerSupportLibraries_jll", "ExprTools", "GPUArrays", "GPUCompiler", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "TimerOutputs"]
+git-tree-sha1 = "6717cb9a3425ebb7b31ca4f832823615d175f64a"
+uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
+version = "3.13.1"
+
+[[deps.CatIndices]]
+deps = ["CustomUnitRanges", "OffsetArrays"]
+git-tree-sha1 = "a0f80a09780eed9b1d106a1bf62041c2efc995bc"
+uuid = "aafaddc9-749c-510e-ac4f-586e18779b91"
+version = "0.2.2"
+
+[[deps.ChainRules]]
+deps = ["Adapt", "ChainRulesCore", "Compat", "Distributed", "GPUArraysCore", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "SparseInverseSubset", "Statistics", "StructArrays", "SuiteSparse"]
+git-tree-sha1 = "be227d253d132a6d57f9ccf5f67c0fb6488afd87"
+uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+version = "1.71.0"
+
+[[deps.ChainRulesCore]]
+deps = ["Compat", "LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "3e4b134270b372f2ed4d4d0e936aabaefc1802bc"
+uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+version = "1.25.0"
+
+[[deps.ChangesOfVariables]]
+deps = ["InverseFunctions", "LinearAlgebra", "Test"]
+git-tree-sha1 = "799b25ca3a8a24936ae7b5c52ad194685fc3e6ef"
+uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
+version = "0.1.9"
+
+[[deps.CodecZlib]]
+deps = ["TranscodingStreams", "Zlib_jll"]
+git-tree-sha1 = "bce6804e5e6044c6daab27bb533d1295e4a2e759"
+uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
+version = "0.7.6"
+
+[[deps.ColorSchemes]]
+deps = ["ColorTypes", "ColorVectorSpace", "Colors", "FixedPointNumbers", "PrecompileTools", "Random"]
+git-tree-sha1 = "b5278586822443594ff615963b0c09755771b3e0"
+uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
+version = "3.26.0"
+
+[[deps.ColorTypes]]
+deps = ["FixedPointNumbers", "Random"]
+git-tree-sha1 = "b10d0b65641d57b8b4d5e234446582de5047050d"
+uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
+version = "0.11.5"
+
+[[deps.ColorVectorSpace]]
+deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "SpecialFunctions", "Statistics", "TensorCore"]
+git-tree-sha1 = "600cc5508d66b78aae350f7accdb58763ac18589"
+uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4"
+version = "0.9.10"
+
+[[deps.Colors]]
+deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
+git-tree-sha1 = "362a287c3aa50601b0bc359053d5c2468f0e7ce0"
+uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
+version = "0.12.11"
+
+[[deps.CommonSubexpressions]]
+deps = ["MacroTools"]
+git-tree-sha1 = "cda2cfaebb4be89c9084adaca7dd7333369715c5"
+uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
+version = "0.3.1"
+
+[[deps.Compat]]
+deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
+git-tree-sha1 = "d476eaeddfcdf0de15a67a948331c69a585495fa"
+uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
+version = "3.47.0"
+
+[[deps.CompilerSupportLibraries_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
+version = "1.0.1+0"
+
+[[deps.CompositionsBase]]
+git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad"
+uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b"
+version = "0.1.2"
+
+[[deps.ComputationalResources]]
+git-tree-sha1 = "52cb3ec90e8a8bea0e62e275ba577ad0f74821f7"
+uuid = "ed09eef8-17a6-5b46-8889-db040fac31e3"
+version = "0.3.2"
+
+[[deps.Conda]]
+deps = ["Downloads", "JSON", "VersionParsing"]
+git-tree-sha1 = "b19db3927f0db4151cb86d073689f2428e524576"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.10.2"
+
+[[deps.ConstructionBase]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "76219f1ed5771adbb096743bff43fb5fdd4c1157"
+uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
+version = "1.5.8"
+
+[[deps.ContextVariablesX]]
+deps = ["Compat", "Logging", "UUIDs"]
+git-tree-sha1 = "25cc3803f1030ab855e383129dcd3dc294e322cc"
+uuid = "6add18c4-b38d-439d-96f6-d6bc489c04c5"
+version = "0.1.3"
+
+[[deps.CoordinateTransformations]]
+deps = ["LinearAlgebra", "StaticArrays"]
+git-tree-sha1 = "f9d7112bfff8a19a3a4ea4e03a8e6a91fe8456bf"
+uuid = "150eb455-5306-5404-9cee-2592286d6298"
+version = "0.6.3"
+
+[[deps.CustomUnitRanges]]
+git-tree-sha1 = "1a3f97f907e6dd8983b744d2642651bb162a3f7a"
+uuid = "dc8bdbbb-1ca9-579f-8c36-e416f6a65cce"
+version = "1.0.2"
+
+[[deps.DSP]]
+deps = ["Compat", "FFTW", "IterTools", "LinearAlgebra", "Polynomials", "Random", "Reexport", "SpecialFunctions", "Statistics"]
+git-tree-sha1 = "0df00546373af8eee1598fb4b2ba480b1ebe895c"
+uuid = "717857b8-e6f2-59f4-9121-6e50c889abd2"
+version = "0.7.10"
+
+[[deps.DataAPI]]
+git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe"
+uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
+version = "1.16.0"
+
+[[deps.DataStructures]]
+deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
+git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82"
+uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+version = "0.18.20"
+
+[[deps.DataValueInterfaces]]
+git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
+uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
+version = "1.0.0"
+
+[[deps.Dates]]
+deps = ["Printf"]
+uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
+
+[[deps.DefineSingletons]]
+git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c"
+uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52"
+version = "0.1.2"
+
+[[deps.DelimitedFiles]]
+deps = ["Mmap"]
+uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+
+[[deps.DensityInterface]]
+deps = ["InverseFunctions", "Test"]
+git-tree-sha1 = "80c3e8639e3353e5d2912fb3a1916b8455e2494b"
+uuid = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
+version = "0.4.0"
+
+[[deps.DiffResults]]
+deps = ["StaticArraysCore"]
+git-tree-sha1 = "782dd5f4561f5d267313f23853baaaa4c52ea621"
+uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+version = "1.1.0"
+
+[[deps.DiffRules]]
+deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"]
+git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272"
+uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
+version = "1.15.1"
+
+[[deps.Distances]]
+deps = ["LinearAlgebra", "SparseArrays", "Statistics", "StatsAPI"]
+git-tree-sha1 = "c7e3a542b999843086e2f29dac96a618c105be1d"
+uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+version = "0.10.12"
+
+[[deps.Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[[deps.DistributedArrays]]
+deps = ["Distributed", "LinearAlgebra", "Primes", "Random", "Serialization", "SparseArrays", "Statistics"]
+git-tree-sha1 = "8a2950a2a7fcc87d9d9f07433a96bfa70394a3e6"
+uuid = "aaf54ef3-cdf8-58ed-94cc-d582ad619b94"
+version = "0.6.7"
+
+[[deps.Distributions]]
+deps = ["AliasTables", "ChainRulesCore", "DensityInterface", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns", "Test"]
+git-tree-sha1 = "d7477ecdafb813ddee2ae727afa94e9dcb5f3fb0"
+uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
+version = "0.25.112"
+
+[[deps.DocStringExtensions]]
+deps = ["LibGit2"]
+git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d"
+uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+version = "0.9.3"
+
+[[deps.Downloads]]
+deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+version = "1.6.0"
+
+[[deps.EnzymeCore]]
+deps = ["Adapt"]
+git-tree-sha1 = "8f205a601760f4798a10f138c3940f0451d95188"
+uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
+version = "0.7.8"
+
+[[deps.ExprTools]]
+git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec"
+uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+version = "0.1.10"
+
+[[deps.FFTViews]]
+deps = ["CustomUnitRanges", "FFTW"]
+git-tree-sha1 = "cbdf14d1e8c7c8aacbe8b19862e0179fd08321c2"
+uuid = "4f61f5a4-77b1-5117-aa51-3ab5ef4ef0cd"
+version = "0.3.2"
+
+[[deps.FFTW]]
+deps = ["AbstractFFTs", "FFTW_jll", "LinearAlgebra", "MKL_jll", "Preferences", "Reexport"]
+git-tree-sha1 = "4820348781ae578893311153d69049a93d05f39d"
+uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+version = "1.8.0"
+
+[[deps.FFTW_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "4d81ed14783ec49ce9f2e168208a12ce1815aa25"
+uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a"
+version = "3.3.10+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.FileIO]]
+deps = ["Pkg", "Requires", "UUIDs"]
+git-tree-sha1 = "62ca0547a14c57e98154423419d8a342dca75ca9"
+uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+version = "1.16.4"
+
+[[deps.FileWatching]]
+uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
+
+[[deps.FillArrays]]
+deps = ["LinearAlgebra", "PDMats", "SparseArrays", "Statistics"]
+git-tree-sha1 = "6a70198746448456524cb442b8af316927ff3e1a"
+uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
+version = "1.13.0"
+
+[[deps.FixedPointNumbers]]
+deps = ["Statistics"]
+git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172"
+uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
+version = "0.8.5"
+
+[[deps.Flux]]
+deps = ["AbstractTrees", "Adapt", "ArrayInterface", "CUDA", "CodecZlib", "Colors", "DelimitedFiles", "Functors", "Juno", "LinearAlgebra", "MacroTools", "NNlib", "NNlibCUDA", "Pkg", "Printf", "Random", "Reexport", "SHA", "SparseArrays", "Statistics", "StatsBase", "Test", "ZipFile", "Zygote"]
+git-tree-sha1 = "511b7c48eebb602a8f63e7d6c63e25633468dc16"
+uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+version = "0.12.10"
+
+[[deps.ForwardDiff]]
+deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
+git-tree-sha1 = "cf0fe81336da9fb90944683b8c41984b08793dad"
+uuid = "f6369f11-7733-5829-9624-2563aa707210"
+version = "0.10.36"
+
+[[deps.Functors]]
+git-tree-sha1 = "223fffa49ca0ff9ce4f875be001ffe173b2b7de4"
+uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+version = "0.2.8"
+
+[[deps.Future]]
+deps = ["Random"]
+uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
+
+[[deps.GPUArrays]]
+deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"]
+git-tree-sha1 = "2e57b4a4f9cc15e85a24d603256fe08e527f48d1"
+uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+version = "8.8.1"
+
+[[deps.GPUArraysCore]]
+deps = ["Adapt"]
+git-tree-sha1 = "2d6ca471a6c7b536127afccfa7564b5b39227fe0"
+uuid = "46192b85-c4d5-4398-a991-12ede77f4527"
+version = "0.1.5"
+
+[[deps.GPUCompiler]]
+deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "TimerOutputs", "UUIDs"]
+git-tree-sha1 = "19d693666a304e8c371798f4900f7435558c7cde"
+uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
+version = "0.17.3"
+
+[[deps.Ghostscript_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "43ba3d3c82c18d88471cfd2924931658838c9d8f"
+uuid = "61579ee1-b43e-5ca0-a5da-69d92c66a64b"
+version = "9.55.0+4"
+
+[[deps.Graphics]]
+deps = ["Colors", "LinearAlgebra", "NaNMath"]
+git-tree-sha1 = "d61890399bc535850c4bf08e4e0d3a7ad0f21cbd"
+uuid = "a2bd30eb-e257-5431-a919-1863eab51364"
+version = "1.1.2"
+
+[[deps.HypergeometricFunctions]]
+deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"]
+git-tree-sha1 = "7c4195be1649ae622304031ed46a2f4df989f1eb"
+uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a"
+version = "0.3.24"
+
+[[deps.IRTools]]
+deps = ["InteractiveUtils", "MacroTools"]
+git-tree-sha1 = "950c3717af761bc3ff906c2e8e52bd83390b6ec2"
+uuid = "7869d1d1-7146-5819-86e3-90919afe41df"
+version = "0.4.14"
+
+[[deps.IdentityRanges]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "be8fcd695c4da16a1d6d0cd213cb88090a150e3b"
+uuid = "bbac6d45-d8f3-5730-bfe4-7a449cd117ca"
+version = "0.3.1"
+
+[[deps.IfElse]]
+git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1"
+uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
+version = "0.1.1"
+
+[[deps.ImageAxes]]
+deps = ["AxisArrays", "ImageCore", "Reexport", "SimpleTraits"]
+git-tree-sha1 = "794ad1d922c432082bc1aaa9fa8ffbd1fe74e621"
+uuid = "2803e5a7-5153-5ecf-9a86-9b4c37f5f5ac"
+version = "0.6.9"
+
+[[deps.ImageContrastAdjustment]]
+deps = ["ColorVectorSpace", "ImageCore", "ImageTransformations", "Parameters"]
+git-tree-sha1 = "2e6084db6cccab11fe0bc3e4130bd3d117092ed9"
+uuid = "f332f351-ec65-5f6a-b3d1-319c6670881a"
+version = "0.3.7"
+
+[[deps.ImageCore]]
+deps = ["AbstractFFTs", "Colors", "FixedPointNumbers", "Graphics", "MappedArrays", "MosaicViews", "OffsetArrays", "PaddedViews", "Reexport"]
+git-tree-sha1 = "db645f20b59f060d8cfae696bc9538d13fd86416"
+uuid = "a09fc81d-aa75-5fe9-8630-4744c3626534"
+version = "0.8.22"
+
+[[deps.ImageDistances]]
+deps = ["ColorVectorSpace", "Distances", "ImageCore", "ImageMorphology", "LinearAlgebra", "Statistics"]
+git-tree-sha1 = "6378c34a3c3a216235210d19b9f495ecfff2f85f"
+uuid = "51556ac3-7006-55f5-8cb3-34580c88182d"
+version = "0.2.13"
+
+[[deps.ImageFiltering]]
+deps = ["CatIndices", "ColorVectorSpace", "ComputationalResources", "DataStructures", "FFTViews", "FFTW", "ImageCore", "LinearAlgebra", "OffsetArrays", "Requires", "SparseArrays", "StaticArrays", "Statistics", "TiledIteration"]
+git-tree-sha1 = "bf96839133212d3eff4a1c3a80c57abc7cfbf0ce"
+uuid = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
+version = "0.6.21"
+
+[[deps.ImageIO]]
+deps = ["FileIO", "Netpbm", "PNGFiles"]
+git-tree-sha1 = "0d6d09c28d67611c68e25af0c2df7269c82b73c7"
+uuid = "82e4d734-157c-48bb-816b-45c225c6df19"
+version = "0.4.1"
+
+[[deps.ImageMagick]]
+deps = ["FileIO", "ImageCore", "ImageMagick_jll", "InteractiveUtils", "Libdl", "Pkg", "Random"]
+git-tree-sha1 = "5bc1cb62e0c5f1005868358db0692c994c3a13c6"
+uuid = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
+version = "1.2.1"
+
+[[deps.ImageMagick_jll]]
+deps = ["Artifacts", "Ghostscript_jll", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "OpenJpeg_jll", "Zlib_jll", "libpng_jll"]
+git-tree-sha1 = "d65554bad8b16d9562050c67e7223abf91eaba2f"
+uuid = "c73af94c-d91f-53ed-93a7-00f77d67a9d7"
+version = "6.9.13+0"
+
+[[deps.ImageMetadata]]
+deps = ["AxisArrays", "ColorVectorSpace", "ImageAxes", "ImageCore", "IndirectArrays"]
+git-tree-sha1 = "ae76038347dc4edcdb06b541595268fca65b6a42"
+uuid = "bc367c6b-8a6b-528e-b4bd-a4b897500b49"
+version = "0.9.5"
+
+[[deps.ImageMorphology]]
+deps = ["ColorVectorSpace", "ImageCore", "LinearAlgebra", "TiledIteration"]
+git-tree-sha1 = "68e7cbcd7dfaa3c2f74b0a8ab3066f5de8f2b71d"
+uuid = "787d08f9-d448-5407-9aad-5290dd7ab264"
+version = "0.2.11"
+
+[[deps.ImageQualityIndexes]]
+deps = ["ColorVectorSpace", "ImageCore", "ImageDistances", "ImageFiltering", "OffsetArrays", "Statistics"]
+git-tree-sha1 = "1198f85fa2481a3bb94bf937495ba1916f12b533"
+uuid = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
+version = "0.2.2"
+
+[[deps.ImageShow]]
+deps = ["Base64", "FileIO", "ImageCore", "OffsetArrays", "Requires", "StackViews"]
+git-tree-sha1 = "832abfd709fa436a562db47fd8e81377f72b01f9"
+uuid = "4e3cecfd-b093-5904-9786-8bbb286a6a31"
+version = "0.3.1"
+
+[[deps.ImageTransformations]]
+deps = ["AxisAlgorithms", "ColorVectorSpace", "CoordinateTransformations", "IdentityRanges", "ImageCore", "Interpolations", "OffsetArrays", "Rotations", "StaticArrays"]
+git-tree-sha1 = "e4cc551e4295a5c96545bb3083058c24b78d4cf0"
+uuid = "02fcd773-0e25-5acc-982a-7f6622650795"
+version = "0.8.13"
+
+[[deps.Images]]
+deps = ["AxisArrays", "Base64", "ColorVectorSpace", "FileIO", "Graphics", "ImageAxes", "ImageContrastAdjustment", "ImageCore", "ImageDistances", "ImageFiltering", "ImageIO", "ImageMagick", "ImageMetadata", "ImageMorphology", "ImageQualityIndexes", "ImageShow", "ImageTransformations", "IndirectArrays", "OffsetArrays", "Random", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "StatsBase", "TiledIteration"]
+git-tree-sha1 = "8b714d5e11c91a0d945717430ec20f9251af4bd2"
+uuid = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+version = "0.24.1"
+
+[[deps.IndirectArrays]]
+git-tree-sha1 = "c2a145a145dc03a7620af1444e0264ef907bd44f"
+uuid = "9b13fd28-a010-5f03-acff-a1bbcff69959"
+version = "0.5.1"
+
+[[deps.InitialValues]]
+git-tree-sha1 = "4da0f88e9a39111c2fa3add390ab15f3a44f3ca3"
+uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
+version = "0.3.1"
+
+[[deps.InplaceOps]]
+deps = ["LinearAlgebra", "Test"]
+git-tree-sha1 = "50b41d59e7164ab6fda65e71049fee9d890731ff"
+uuid = "505f98c9-085e-5b2c-8e89-488be7bf1f34"
+version = "0.3.0"
+
+[[deps.IntegerMathUtils]]
+git-tree-sha1 = "b8ffb903da9f7b8cf695a8bead8e01814aa24b30"
+uuid = "18e54dd8-cb9d-406c-a71d-865a43cbb235"
+version = "0.1.2"
+
+[[deps.IntelOpenMP_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"]
+git-tree-sha1 = "10bd689145d2c3b2a9844005d01087cc1194e79e"
+uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0"
+version = "2024.2.1+0"
+
+[[deps.InteractiveUtils]]
+deps = ["Markdown"]
+uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+
+[[deps.Interpolations]]
+deps = ["AxisAlgorithms", "ChainRulesCore", "LinearAlgebra", "OffsetArrays", "Random", "Ratios", "Requires", "SharedArrays", "SparseArrays", "StaticArrays", "WoodburyMatrices"]
+git-tree-sha1 = "b7bc05649af456efc75d178846f47006c2c4c3c7"
+uuid = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+version = "0.13.6"
+
+[[deps.IntervalSets]]
+deps = ["Random", "RecipesBase", "Statistics"]
+git-tree-sha1 = "dba9ddf07f77f60450fe5d2e2beb9854d9a49bd0"
+uuid = "8197267c-284f-5f27-9208-e0e47529a953"
+version = "0.7.10"
+
+[[deps.InverseFunctions]]
+deps = ["Dates", "Test"]
+git-tree-sha1 = "a779299d77cd080bf77b97535acecd73e1c5e5cb"
+uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
+version = "0.1.17"
+
+[[deps.InvertibleNetworks]]
+deps = ["CUDA", "ChainRulesCore", "Flux", "JOLI", "LinearAlgebra", "NNlib", "Random", "Statistics", "Wavelets"]
+git-tree-sha1 = "b8bd9465a5c05318152db2343fdef102f369fa06"
+repo-rev = "master"
+repo-url = "https://github.com/ziyiyin97/InvertibleNetworks.jl.git"
+uuid = "b7115f24-5f92-4794-81e8-23b0ddb121d3"
+version = "2.2.6"
+
+[[deps.IrrationalConstants]]
+git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2"
+uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
+version = "0.2.2"
+
+[[deps.IterTools]]
+git-tree-sha1 = "42d5f897009e7ff2cf88db414a389e5ed1bdd023"
+uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+version = "1.10.0"
+
+[[deps.IterativeSolvers]]
+deps = ["LinearAlgebra", "Printf", "Random", "RecipesBase", "SparseArrays"]
+git-tree-sha1 = "59545b0a2b27208b0650df0a46b8e3019f85055b"
+uuid = "42fd0dbc-a981-5370-80f2-aaf504508153"
+version = "0.9.4"
+
+[[deps.IteratorInterfaceExtensions]]
+git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
+uuid = "82899510-4779-5014-852e-03e436cf321d"
+version = "1.0.0"
+
+[[deps.JLD2]]
+deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"]
+git-tree-sha1 = "783c1be5213a09609b23237a0c9e5dfd258ae6f2"
+uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+version = "0.5.7"
+
+[[deps.JLLWrappers]]
+deps = ["Artifacts", "Preferences"]
+git-tree-sha1 = "be3dc50a92e5a386872a493a10050136d4703f9b"
+uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
+version = "1.6.1"
+
+[[deps.JOLI]]
+deps = ["ChainRulesCore", "Distributed", "DistributedArrays", "FFTW", "InplaceOps", "InteractiveUtils", "IterativeSolvers", "Libdl", "LinearAlgebra", "NFFT", "Nullables", "Printf", "PyCall", "Random", "SharedArrays", "SparseArrays", "SpecialFunctions", "Wavelets"]
+git-tree-sha1 = "97beb81d38c0928b58280f5daee8673546732ea4"
+uuid = "bb331ad6-a1cf-11e9-23da-9bcb53c69f6f"
+version = "0.8.5"
+
+[[deps.JSON]]
+deps = ["Dates", "Mmap", "Parsers", "Unicode"]
+git-tree-sha1 = "31e996f0a15c7b280ba9f76636b3ff9e2ae58c9a"
+uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+version = "0.21.4"
+
+[[deps.JpegTurbo_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "25ee0be4d43d0269027024d75a24c24d6c6e590c"
+uuid = "aacddb02-875f-59d6-b918-886e6ef4fbf8"
+version = "3.0.4+0"
+
+[[deps.JuliaVariables]]
+deps = ["MLStyle", "NameResolution"]
+git-tree-sha1 = "49fb3cb53362ddadb4415e9b73926d6b40709e70"
+uuid = "b14d175d-62b4-44ba-8fb7-3064adc8c3ec"
+version = "0.2.4"
+
+[[deps.Juno]]
+deps = ["Base64", "Logging", "Media", "Profile"]
+git-tree-sha1 = "07cb43290a840908a771552911a6274bc6c072c7"
+uuid = "e5e0dc1b-0480-54bc-9374-aad01c23163d"
+version = "0.8.4"
+
+[[deps.KernelAbstractions]]
+deps = ["Adapt", "Atomix", "EnzymeCore", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
+git-tree-sha1 = "e73a077abc7fe798fe940deabe30ef6c66bdde52"
+uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+version = "0.9.29"
+
+[[deps.LERC_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "36bdbc52f13a7d1dcb0f3cd694e01677a515655b"
+uuid = "88015f11-f218-50d7-93a8-a6af411a945d"
+version = "4.0.0+0"
+
+[[deps.LLVM]]
+deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
+git-tree-sha1 = "f044a2796a9e18e0531b9b3072b0019a61f264bc"
+uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
+version = "4.17.1"
+
+[[deps.LLVMExtra_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "070e4b5b65827f82c16ae0916376cb47377aa1b5"
+uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
+version = "0.0.18+0"
+
+[[deps.LaTeXStrings]]
+git-tree-sha1 = "dda21b8cbd6a6c40d9d02a73230f9d70fed6918c"
+uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+version = "1.4.0"
+
+[[deps.LazyArtifacts]]
+deps = ["Artifacts", "Pkg"]
+uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+
+[[deps.LearnBase]]
+git-tree-sha1 = "a0d90569edd490b82fdc4dc078ea54a5a800d30a"
+uuid = "7f8f8fb0-2700-5f03-b4bd-41f8cfc144b6"
+version = "0.4.1"
+
+[[deps.LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+version = "0.6.3"
+
+[[deps.LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
+version = "7.84.0+0"
+
+[[deps.LibGit2]]
+deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
+uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
+
+[[deps.LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+version = "1.10.2+0"
+
+[[deps.Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+
+[[deps.Libtiff_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "LERC_jll", "Libdl", "XZ_jll", "Zlib_jll", "Zstd_jll"]
+git-tree-sha1 = "b404131d06f7886402758c9ce2214b636eb4d54a"
+uuid = "89763e89-9b03-5906-acba-b20f662cd828"
+version = "4.7.0+0"
+
+[[deps.LinearAlgebra]]
+deps = ["Libdl", "libblastrampoline_jll"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[[deps.LittleCMS_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll"]
+git-tree-sha1 = "fa7fd067dca76cadd880f1ca937b4f387975a9f5"
+uuid = "d3a379c0-f9a3-5b72-a4c0-6bf4d2e8af0f"
+version = "2.16.0+0"
+
+[[deps.LogExpFunctions]]
+deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
+git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea"
+uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+version = "0.3.28"
+
+[[deps.Logging]]
+uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+
+[[deps.MKL_jll]]
+deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"]
+git-tree-sha1 = "f046ccd0c6db2832a9f639e2c669c6fe867e5f4f"
+uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
+version = "2024.2.0+0"
+
+[[deps.MLDataPattern]]
+deps = ["LearnBase", "MLLabelUtils", "Random", "SparseArrays", "StatsBase"]
+git-tree-sha1 = "2741aabe5ec93c63cc0a5277bb1c2cb819b01bb1"
+uuid = "9920b226-0b2a-5f5f-9153-9aa70a013f8b"
+version = "0.5.5"
+
+[[deps.MLLabelUtils]]
+deps = ["LearnBase", "MappedArrays", "StatsBase"]
+git-tree-sha1 = "fd75d4b0c4016e047bbb6263eecf7ae3891af522"
+uuid = "66a33bbf-0c2b-5fc8-a008-9da813334f0a"
+version = "0.5.7"
+
+[[deps.MLStyle]]
+git-tree-sha1 = "bc38dff0548128765760c79eb7388a4b37fae2c8"
+uuid = "d8e11817-5142-5d16-987a-aa16d5891078"
+version = "0.4.17"
+
+[[deps.MacroTools]]
+deps = ["Markdown", "Random"]
+git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df"
+uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.5.13"
+
+[[deps.MappedArrays]]
+git-tree-sha1 = "2dab0221fe2b0f2cb6754eaa743cc266339f527e"
+uuid = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
+version = "0.4.2"
+
+[[deps.Markdown]]
+deps = ["Base64"]
+uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
+
+[[deps.MbedTLS_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
+version = "2.28.0+0"
+
+[[deps.Media]]
+deps = ["MacroTools", "Test"]
+git-tree-sha1 = "75a54abd10709c01f1b86b84ec225d26e840ed58"
+uuid = "e89f7d12-3494-54d1-8411-f7d8b9ae1f27"
+version = "0.5.0"
+
+[[deps.Memoize]]
+deps = ["MacroTools"]
+git-tree-sha1 = "2b1dfcba103de714d31c033b5dacc2e4a12c7caa"
+uuid = "c03570c3-d221-55d1-a50c-7939bbd78826"
+version = "0.4.4"
+
+[[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"
+uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
+version = "1.2.0"
+
+[[deps.Mmap]]
+uuid = "a63ad114-7e13-5084-954f-fe012c677804"
+
+[[deps.MosaicViews]]
+deps = ["MappedArrays", "OffsetArrays", "PaddedViews", "StackViews"]
+git-tree-sha1 = "7b86a5d4d70a9f5cdf2dacb3cbe6d251d1a61dbe"
+uuid = "e94cdb99-869f-56ef-bcf0-1ae2bcbe0389"
+version = "0.3.4"
+
+[[deps.MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+version = "2022.2.1"
+
+[[deps.NFFT]]
+deps = ["AbstractNFFTs", "BasicInterpolators", "Distributed", "FFTW", "FLoops", "LinearAlgebra", "Printf", "Random", "Reexport", "SnoopPrecompile", "SparseArrays", "SpecialFunctions"]
+git-tree-sha1 = "93a5f32dd6cf09456b0b81afcb8fc29f06535ffd"
+uuid = "efe261a4-0d2b-5849-be55-fc731d526b0d"
+version = "0.13.3"
+
+[[deps.NNlib]]
+deps = ["Adapt", "Atomix", "ChainRulesCore", "GPUArraysCore", "KernelAbstractions", "LinearAlgebra", "Pkg", "Random", "Requires", "Statistics"]
+git-tree-sha1 = "72240e3f5ca031937bd536182cb2c031da5f46dd"
+uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+version = "0.8.21"
+
+[[deps.NNlibCUDA]]
+deps = ["Adapt", "CUDA", "LinearAlgebra", "NNlib", "Random", "Statistics"]
+git-tree-sha1 = "b05a082b08a3af0e5c576883bc6dfb6513e7e478"
+uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
+version = "0.2.6"
+
+[[deps.NaNMath]]
+deps = ["OpenLibm_jll"]
+git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4"
+uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+version = "1.0.2"
+
+[[deps.NameResolution]]
+deps = ["PrettyPrint"]
+git-tree-sha1 = "1a0fa0e9613f46c9b8c11eee38ebb4f590013c5e"
+uuid = "71a1bf82-56d0-4bbc-8a3c-48b961074391"
+version = "0.1.5"
+
+[[deps.Netpbm]]
+deps = ["ColorVectorSpace", "FileIO", "ImageCore"]
+git-tree-sha1 = "09589171688f0039f13ebe0fdcc7288f50228b52"
+uuid = "f09324ee-3d7c-5217-9330-fc30815ba969"
+version = "1.0.1"
+
+[[deps.NetworkOptions]]
+uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
+version = "1.2.0"
+
+[[deps.Nullables]]
+git-tree-sha1 = "8f87854cc8f3685a60689d8edecaa29d2251979b"
+uuid = "4d1e1d77-625e-5b40-9113-a560ec7a8ecd"
+version = "1.0.0"
+
+[[deps.OffsetArrays]]
+deps = ["Adapt"]
+git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e"
+uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+version = "1.14.1"
+
+[[deps.OpenBLAS_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
+uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
+version = "0.3.20+0"
+
+[[deps.OpenJpeg_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "libpng_jll"]
+git-tree-sha1 = "f4cb457ffac5f5cf695699f82c537073958a6a6c"
+uuid = "643b3616-a352-519d-856d-80112ee9badc"
+version = "2.5.2+0"
+
+[[deps.OpenLibm_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
+version = "0.8.1+0"
+
+[[deps.OpenSpecFun_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
+uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
+version = "0.5.5+0"
+
+[[deps.OrderedCollections]]
+git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5"
+uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+version = "1.6.3"
+
+[[deps.PDMats]]
+deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
+git-tree-sha1 = "949347156c25054de2db3b166c52ac4728cbad65"
+uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
+version = "0.11.31"
+
+[[deps.PNGFiles]]
+deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"]
+git-tree-sha1 = "f809158b27eba0c18c269cf2a2be6ed751d3e81d"
+uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883"
+version = "0.3.17"
+
+[[deps.PaddedViews]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "0fac6313486baae819364c52b4f483450a9d793f"
+uuid = "5432bcbf-9aad-5242-b902-cca2824c8663"
+version = "0.5.12"
+
+[[deps.Parameters]]
+deps = ["OrderedCollections", "UnPack"]
+git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
+uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+version = "0.12.3"
+
+[[deps.Parsers]]
+deps = ["Dates", "PrecompileTools", "UUIDs"]
+git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821"
+uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
+version = "2.8.1"
+
+[[deps.Pkg]]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+version = "1.8.0"
+
+[[deps.Polynomials]]
+deps = ["LinearAlgebra", "RecipesBase", "Requires", "Setfield", "SparseArrays"]
+git-tree-sha1 = "1a9cfb2dc2c2f1bd63f1906d72af39a79b49b736"
+uuid = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
+version = "4.0.11"
+
+[[deps.PrecompileTools]]
+deps = ["Preferences"]
+git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f"
+uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+version = "1.2.1"
+
+[[deps.Preferences]]
+deps = ["TOML"]
+git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6"
+uuid = "21216c6a-2e73-6563-6e65-726566657250"
+version = "1.4.3"
+
+[[deps.PrettyPrint]]
+git-tree-sha1 = "632eb4abab3449ab30c5e1afaa874f0b98b586e4"
+uuid = "8162dcfd-2161-5ef2-ae6c-7681170c5f98"
+version = "0.2.0"
+
+[[deps.Primes]]
+deps = ["IntegerMathUtils"]
+git-tree-sha1 = "cb420f77dc474d23ee47ca8d14c90810cafe69e7"
+uuid = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
+version = "0.5.6"
+
+[[deps.Printf]]
+deps = ["Unicode"]
+uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[[deps.Profile]]
+deps = ["Printf"]
+uuid = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
+
+[[deps.PtrArrays]]
+git-tree-sha1 = "77a42d78b6a92df47ab37e177b2deac405e1c88f"
+uuid = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d"
+version = "1.2.1"
+
+[[deps.PyCall]]
+deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
+git-tree-sha1 = "9816a3826b0ebf49ab4926e2b18842ad8b5c8f04"
+uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+version = "1.96.4"
+
+[[deps.PyPlot]]
+deps = ["Colors", "LaTeXStrings", "PyCall", "Sockets", "Test", "VersionParsing"]
+git-tree-sha1 = "0371ca706e3f295481cbf94c8c36692b072285c2"
+uuid = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+version = "2.11.5"
+
+[[deps.QuadGK]]
+deps = ["DataStructures", "LinearAlgebra"]
+git-tree-sha1 = "cda3b045cf9ef07a08ad46731f5a3165e56cf3da"
+uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+version = "2.11.1"
+
+[[deps.Quaternions]]
+deps = ["LinearAlgebra", "Random", "RealDot"]
+git-tree-sha1 = "994cc27cdacca10e68feb291673ec3a76aa2fae9"
+uuid = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
+version = "0.7.6"
+
+[[deps.REPL]]
+deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
+uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+
+[[deps.Random]]
+deps = ["SHA", "Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[deps.Random123]]
+deps = ["Random", "RandomNumbers"]
+git-tree-sha1 = "4743b43e5a9c4a2ede372de7061eed81795b12e7"
+uuid = "74087812-796a-5b5d-8853-05524746bad3"
+version = "1.7.0"
+
+[[deps.RandomNumbers]]
+deps = ["Random"]
+git-tree-sha1 = "c6ec94d2aaba1ab2ff983052cf6a606ca5985902"
+uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143"
+version = "1.6.0"
+
+[[deps.RangeArrays]]
+git-tree-sha1 = "b9039e93773ddcfc828f12aadf7115b4b4d225f5"
+uuid = "b3c3ace0-ae52-54e7-9d0b-2c1406fd6b9d"
+version = "0.3.2"
+
+[[deps.Ratios]]
+deps = ["Requires"]
+git-tree-sha1 = "1342a47bf3260ee108163042310d26f2be5ec90b"
+uuid = "c84ed2f1-dad5-54f0-aa8e-dbefe2724439"
+version = "0.4.5"
+
+[[deps.RealDot]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "9f0a1b71baaf7650f4fa8a1d168c7fb6ee41f0c9"
+uuid = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9"
+version = "0.1.0"
+
+[[deps.RecipesBase]]
+deps = ["PrecompileTools"]
+git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff"
+uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+version = "1.3.4"
+
+[[deps.Reexport]]
+git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
+uuid = "189a3867-3050-52da-a836-e630ba90ab69"
+version = "1.2.2"
+
+[[deps.Requires]]
+deps = ["UUIDs"]
+git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
+uuid = "ae029012-a4dd-5104-9daa-d747884805df"
+version = "1.3.0"
+
+[[deps.Rmath]]
+deps = ["Random", "Rmath_jll"]
+git-tree-sha1 = "852bd0f55565a9e973fcfee83a84413270224dc4"
+uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa"
+version = "0.8.0"
+
+[[deps.Rmath_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "58cdd8fb2201a6267e1db87ff148dd6c1dbd8ad8"
+uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
+version = "0.5.1+0"
+
+[[deps.Rotations]]
+deps = ["LinearAlgebra", "Quaternions", "Random", "RecipesBase", "StaticArrays"]
+git-tree-sha1 = "5680a9276685d392c87407df00d57c9924d9f11e"
+uuid = "6038ab10-8711-5258-84ad-4b1120ba62dc"
+version = "1.7.1"
+
+[[deps.SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+version = "0.7.0"
+
+[[deps.Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[deps.Setfield]]
+deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"]
+git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac"
+uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+version = "1.1.1"
+
+[[deps.SharedArrays]]
+deps = ["Distributed", "Mmap", "Random", "Serialization"]
+uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
+
+[[deps.SimpleTraits]]
+deps = ["InteractiveUtils", "MacroTools"]
+git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231"
+uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
+version = "0.9.4"
+
+[[deps.SlimPlotting]]
+deps = ["ColorSchemes", "Requires", "Statistics"]
+git-tree-sha1 = "e9b0658e186ce6bc020ca6bf24305fb835716470"
+uuid = "f6d04670-764e-495b-a720-91c3c9a588ff"
+version = "1.0.0"
+
+[[deps.SnoopPrecompile]]
+deps = ["Preferences"]
+git-tree-sha1 = "e760a70afdcd461cf01a575947738d359234665c"
+uuid = "66db9d55-30c0-4569-8b51-7e840670fc0c"
+version = "1.0.3"
+
+[[deps.Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+
+[[deps.SortingAlgorithms]]
+deps = ["DataStructures"]
+git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085"
+uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
+version = "1.2.1"
+
+[[deps.SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[deps.SparseInverseSubset]]
+deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
+git-tree-sha1 = "52962839426b75b3021296f7df242e40ecfc0852"
+uuid = "dc90abb0-5640-4711-901d-7e5b23a2fada"
+version = "0.1.2"
+
+[[deps.SpecialFunctions]]
+deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
+git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14"
+uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
+version = "2.4.0"
+
+[[deps.SplittablesBase]]
+deps = ["Setfield", "Test"]
+git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5"
+uuid = "171d559e-b47b-412a-8079-5efa626c420e"
+version = "0.1.15"
+
+[[deps.StackViews]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "46e589465204cd0c08b4bd97385e4fa79a0c770c"
+uuid = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
+version = "0.1.1"
+
+[[deps.Static]]
+deps = ["IfElse"]
+git-tree-sha1 = "5d2c08cef80c7a3a8ba9ca023031a85c263012c5"
+uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
+version = "0.6.6"
+
+[[deps.StaticArrays]]
+deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore", "Statistics"]
+git-tree-sha1 = "777657803913ffc7e8cc20f0fd04b634f871af8f"
+uuid = "90137ffa-7385-5640-81b9-e52037218182"
+version = "1.9.8"
+
+[[deps.StaticArraysCore]]
+git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682"
+uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+version = "1.4.3"
+
+[[deps.Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[deps.StatsAPI]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed"
+uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
+version = "1.7.0"
+
+[[deps.StatsBase]]
+deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
+git-tree-sha1 = "d1bf48bfcc554a3761a133fe3a9bb01488e06916"
+uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+version = "0.33.21"
+
+[[deps.StatsFuns]]
+deps = ["ChainRulesCore", "HypergeometricFunctions", "InverseFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"]
+git-tree-sha1 = "b423576adc27097764a90e163157bcfc9acf0f46"
+uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+version = "1.3.2"
+
+[[deps.StructArrays]]
+deps = ["Adapt", "ConstructionBase", "DataAPI", "GPUArraysCore", "SparseArrays", "StaticArrays", "Tables"]
+git-tree-sha1 = "f4dc295e983502292c4c3f951dbb4e985e35b3be"
+uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+version = "0.6.18"
+
+[[deps.SuiteSparse]]
+deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
+uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
+
+[[deps.TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+version = "1.0.0"
+
+[[deps.TableTraits]]
+deps = ["IteratorInterfaceExtensions"]
+git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
+uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
+version = "1.0.1"
+
+[[deps.Tables]]
+deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"]
+git-tree-sha1 = "598cd7c1f68d1e205689b1c2fe65a9f85846f297"
+uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+version = "1.12.0"
+
+[[deps.Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+version = "1.10.1"
+
+[[deps.TensorCore]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6"
+uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50"
+version = "0.1.1"
+
+[[deps.Test]]
+deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
+uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[[deps.TiledIteration]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "5683455224ba92ef59db72d10690690f4a8dc297"
+uuid = "06e1c1a7-607b-532d-9fad-de7d9aa2abac"
+version = "0.3.1"
+
+[[deps.TimerOutputs]]
+deps = ["ExprTools", "Printf"]
+git-tree-sha1 = "3a6f063d690135f5c1ba351412c82bae4d1402bf"
+uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+version = "0.5.25"
+
+[[deps.TranscodingStreams]]
+git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
+uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+version = "0.11.3"
+
+[[deps.Transducers]]
+deps = ["Accessors", "Adapt", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "ConstructionBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "SplittablesBase", "Tables"]
+git-tree-sha1 = "7deeab4ff96b85c5f72c824cae53a1398da3d1cb"
+uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999"
+version = "0.4.84"
+
+[[deps.UNet]]
+deps = ["Distributions", "Flux", "Random", "StatsBase", "Test"]
+path = "../../UNet.jl"
+uuid = "0d73aaa9-994a-4556-95d0-da67cb772a03"
+version = "0.2.1"
+
+[[deps.UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
+[[deps.UnPack]]
+git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
+uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+version = "1.0.2"
+
+[[deps.Unicode]]
+uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[deps.UnsafeAtomics]]
+git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278"
+uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f"
+version = "0.2.1"
+
+[[deps.UnsafeAtomicsLLVM]]
+deps = ["LLVM", "UnsafeAtomics"]
+git-tree-sha1 = "ead6292c02aab389cb29fe64cc9375765ab1e219"
+uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
+version = "0.1.1"
+
+[[deps.VersionParsing]]
+git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.3.0"
+
+[[deps.Wavelets]]
+deps = ["DSP", "FFTW", "LinearAlgebra", "Reexport", "SpecialFunctions", "Statistics"]
+git-tree-sha1 = "58f7491c21ab2b1d69368c7f7e8a6a93cbf8b7bf"
+uuid = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
+version = "0.9.5"
+
+[[deps.WoodburyMatrices]]
+deps = ["LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "5f24e158cf4cee437052371455fe361f526da062"
+uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6"
+version = "0.5.6"
+
+[[deps.XZ_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "15e637a697345f6743674f1322beefbc5dcd5cfc"
+uuid = "ffd25f8a-64ca-5728-b0f7-c24cf3aae800"
+version = "5.6.3+0"
+
+[[deps.ZipFile]]
+deps = ["Libdl", "Printf", "Zlib_jll"]
+git-tree-sha1 = "3593e69e469d2111389a9bd06bac1f3d730ac6de"
+uuid = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
+version = "0.9.4"
+
+[[deps.Zlib_jll]]
+deps = ["Libdl"]
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+version = "1.2.12+3"
+
+[[deps.Zstd_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "555d1076590a6cc2fdee2ef1469451f872d8b41b"
+uuid = "3161d3a3-bdf6-5164-811a-617609db77b4"
+version = "1.5.6+1"
+
+[[deps.Zygote]]
+deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "GPUArrays", "GPUArraysCore", "IRTools", "InteractiveUtils", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "PrecompileTools", "Random", "Requires", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"]
+git-tree-sha1 = "f816633be6dc5c0ed9ffedda157ecfda0b3b6a69"
+uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
+version = "0.6.72"
+
+[[deps.ZygoteRules]]
+deps = ["ChainRulesCore", "MacroTools"]
+git-tree-sha1 = "27798139afc0a2afa7b1824c206d5e87ea587a00"
+uuid = "700de1a5-db45-46bc-99cf-38207098b444"
+version = "0.2.5"
+
+[[deps.libblastrampoline_jll]]
+deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
+uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"
+version = "5.1.1+0"
+
+[[deps.libpng_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"]
+git-tree-sha1 = "b70c870239dc3d7bc094eb2d6be9b73d27bef280"
+uuid = "b53b4c65-9356-5827-b1ea-8c7a1a84506f"
+version = "1.6.44+0"
+
+[[deps.nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+version = "1.48.0+0"
+
+[[deps.oneTBB_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "7d0ea0f4895ef2f5cb83645fa689e52cb55cf493"
+uuid = "1317d2d5-d96f-522e-a858-c73665f53c3e"
+version = "2021.12.0+0"
+
+[[deps.p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
+version = "17.4.0+0"
diff --git a/CNF/Project.toml b/CNF/Project.toml
new file mode 100644
index 0000000..e26c51a
--- /dev/null
+++ b/CNF/Project.toml
@@ -0,0 +1,13 @@
+[deps]
+Augmentor = "02898b10-1f73-11ea-317c-6393d7073e15"
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+ImageQualityIndexes = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
+Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+InvertibleNetworks = "b7115f24-5f92-4794-81e8-23b0ddb121d3"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+SlimPlotting = "f6d04670-764e-495b-a720-91c3c9a588ff"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+UNet = "0d73aaa9-994a-4556-95d0-da67cb772a03"
diff --git a/CNF/train.jl b/CNF/train.jl
new file mode 100644
index 0000000..6fce7e4
--- /dev/null
+++ b/CNF/train.jl
@@ -0,0 +1,302 @@
+import Pkg; Pkg.instantiate()
+using LinearAlgebra, Random, Statistics
+using ImageQualityIndexes
+using Images, BSON, JLD2
+using FFTW
+using Augmentor
+using Random
+using DrWatson
+using InvertibleNetworks, Flux, UNet
+using PyPlot, SlimPlotting
+
+Random.seed!(2023)
+
+include("utils.jl")
+
+# Training hyperparameters
+T = Float32
+device = gpu
+plot_path = "plots"
+weights_path = "weights"
+
+PyPlot.rc("font", family="serif");
+
+lr           = 8f-4
+clipnorm_val = 3f0
+noise_lev_x  = 0.1f0
+noise_lev_init = deepcopy(noise_lev_x)
+noise_lev_y  = 0.0 
+
+offset_start = 1
+keep_offset_num = 21
+d = (12.5f0, 12.5f0)
+
+save_every   = 5
+plot_every   = 5
+n_condmean   = 12
+num_post_samples = 64
+
+batch_size, n_epochs, j, ntrain, ntest = 8, 2, 1, 16, 16
+batch_size, n_epochs, j, ntrain, ntest = parse.(Int, ARGS[1:5])
+
+@load "data/initial.jld2" x
+@load "data/cigs_iteration_j=$j.jld2" CIGs
+
+weights_path = joinpath(weights_path, "$j", "CNF")
+plot_path = joinpath(plot_path, "$j", "CNF")
+
+isdir(weights_path) || mkpath(weights_path)
+isdir(plot_path) || mkpath(plot_path)
+
+target_train = x[:, :, :, 1:ntrain]
+target_test = x[:, :, :, ntrain+1:ntrain+ntest]
+Y_train = CIGs[:, :, :, 1, 1:ntrain]
+Y_test = CIGs[:, :, :, 1, ntrain+1:ntrain+ntest]
+
+Y_train = permutedims(Y_train, [2, 3, 1, 4])
+Y_test = permutedims(Y_test, [2, 3, 1, 4])
+ 
+n = (size(target_train)[1], size(target_train)[2])
+
+# Depth Scaling Already Done in Simulate
+# for z = 1:n[2]
+# 	Y_train[:,z,:,:] .*= z * d[2]
+# 	Y_test[:,z,:,:] .*= z * d[2]
+# end
+
+# normalize CIGs
+max_y = quantile(abs.(vec(Y_train)),0.9999);
+Y_train ./= max_y;
+Y_test ./= max_y;
+
+n_x, n_y, chan_target, n_train = size(target_train)
+n_train = size(target_train)[end]
+println("n_train: ", n_train)
+N = n_x*n_y*chan_target
+chan_obs   = size(Y_train)[end-1]
+chan_cond = 1
+
+X_train  = target_train
+X_test   = target_test 
+
+vmax_v = maximum(X_train)
+vmin_v = minimum(X_train)
+
+n_batches    = cld(n_train, batch_size)-1
+n_train_safe = batch_size*n_batches
+
+# Summary network parametrs
+unet_lev = 4
+unet = Chain(Unet(chan_obs, chan_cond, unet_lev)|> device);
+
+trainmode!(unet, true); 
+unet = FluxBlock(unet);
+
+# Create conditional network
+L = 3
+K = 9
+n_hidden = 64
+low = 0.5f0
+
+cond_net = NetworkConditionalGlow(chan_target, chan_cond, n_hidden,  L, K;  split_scales=true, activation=SigmoidLayer(low=low,high=1.0f0)) |> device;
+G        = SummarizedNet(cond_net, unet)
+
+# Optimizer
+opt = Flux.Optimiser(ClipNorm(clipnorm_val), ADAM(lr))
+
+# Training logs 
+loss      = []; logdet_train = []; ssim      = []; l2_cm      = []; std_train = [];
+loss_test = []; logdet_test  = []; ssim_test = []; l2_cm_test = []; std_test = [];
+
+noise_lev_x_min =  1f-3
+noise_decay_per_epochs = div(n_epochs-50, Int(floor(log(noise_lev_x_min/noise_lev_init)/log(1f0/1.2f0))+1))
+
+println("start training")
+
+for e=1:n_epochs # epoch loop
+	idx_e = reshape(randperm(n_train)[1:n_train_safe], batch_size, n_batches) 
+
+	if (e >= 30) && (e <= n_epochs-20) && (mod(e,noise_decay_per_epochs) == 0)
+        global noise_lev_x /= 1.2f0
+        global noise_lev_x = max(noise_lev_x, noise_lev_x_min)
+    end
+
+	for b = 1:n_batches # batch loop
+		X = X_train[:, :, :, idx_e[:,b]];
+		Y = Y_train[:, :, :, idx_e[:,b]];
+
+		for i in 1:batch_size # quick data augmentation to prevent overfitting
+			if rand() > 0.5
+				X[:,:,:,i:i] = X[end:-1:1,:,:,i:i]
+				Y[:,:,:,i:i] = Y[end:-1:1,:,:,i:i]
+			end
+		end
+
+		X .+= noise_lev_x*randn(Float32, size(X)); # noises not related to inverse problem 
+		Y .+= noise_lev_y*randn(Float32, size(Y))
+		Y = Y |> device;
+		
+		Zx, Zy, lgdet = G.forward(X |> device, Y)
+
+		# Loss function is l2 norm 
+		append!(loss, norm(Zx)^2 / (N*batch_size))  # normalize by image size and batch size
+		append!(logdet_train, -lgdet / N) # logdet is internally normalized by batch size
+
+		# Set gradients of flow and summary network
+		dx, x, dy = G.backward(Zx / batch_size, Zx, Zy; Y_save = Y)
+
+		for p in get_params(G) 
+			Flux.update!(opt,p.data,p.grad)
+		end; clear_grad!(G)
+
+		print("Iter: epoch=", e, "/", n_epochs, ", batch=", b, "/", n_batches, 
+			"; f l2 = ",  loss[end], 
+			"; lgdet = ", logdet_train[end], "; f = ", loss[end] + logdet_train[end], "\n")
+		Base.flush(Base.stdout)
+	end
+	
+	if(mod(e,plot_every)==0) 
+		#get loss of training objective on test set corresponds to mutual information between summary statistic and x
+		@time l2_test_val, lgdet_test_val  = get_loss(G, X_test, Y_test; device=device, batch_size=batch_size)
+		append!(logdet_test, -lgdet_test_val)
+		append!(loss_test, l2_test_val)
+
+		# get conditional mean metrics over training batch
+		@time cm_l2_train, cm_ssim_train, cm_std_train = get_cm_l2_ssim(G, X_train[:,:,:,1:n_condmean], Y_train[:,:,:,1:n_condmean]; device=device, num_samples=num_post_samples)
+		append!(ssim, cm_ssim_train)
+		append!(l2_cm, cm_l2_train)
+		append!(std_train, cm_std_train) 
+
+		# get conditional mean metrics over testing batch
+		@time cm_l2_test, cm_ssim_test, cm_std_test = get_cm_l2_ssim(G, X_test[:,:,:,1:n_condmean], Y_test[:,:,:,1:n_condmean]; device=device, num_samples=num_post_samples)
+		append!(ssim_test, cm_ssim_test)
+		append!(l2_cm_test, cm_l2_test)
+		append!(std_test, cm_std_test)
+
+		for (test_x, test_y, file_str) in [[X_train, Y_train, "train"], [X_test, Y_test, "test"]]
+			num_cols = 7
+			plots_len = 2
+			all_sampls = size(test_x)[end]-1
+			fig = figure(figsize=(25, 5)); 
+			for (i,ind) in enumerate((1:div(all_sampls,3):all_sampls)[1:plots_len])
+				x = test_x[:,:,:,ind:ind] 
+				y = test_y[:,:,:,ind:ind]
+				y .+= noise_lev_y*randn(Float32, size(y));
+
+				# make samples from posterior for train sample 
+				X_post = posterior_sampler(G,  y, x; device=device, num_samples=num_post_samples,batch_size=batch_size)|> cpu
+				X_post_mean = mean(X_post,dims=4)
+				X_post_std  = std(X_post, dims=4)
+
+				x_hat =  X_post_mean[:,:,1,1]
+				x_gt =   x[:,:,1,1]
+				error_mean = abs.(x_hat-x_gt)
+
+				ssim_i = round(assess_ssim(x_hat, x_gt),digits=2)
+				rmse_i = round(sqrt(mean(error_mean.^2)),digits=4)
+				stdtotal_i = round(sqrt(mean(X_post_std.^2)),digits=4)
+
+				# It's wired but if not train_rtm then comment out these three lines otherwise keep_offset_num not defined below 
+				# if train_rtm
+				#     keep_offset_num = 1
+				# end
+
+				y_plot = y[:,:,div(keep_offset_num, 2)+1, 1]
+				# make RTM looks better
+				y_plot[:, 1:37] = zeros(size(y_plot)[1], 37)
+				a = quantile(abs.(vec(y_plot)), 98/100)
+
+				# if synthoseis
+				#     global y_plot = y_plot'
+				#     # global X_post = X_post'
+				#     global x_gt = x_gt'
+				#     global x_hat = x_hat'
+				#     global error_mean = error_mean'
+				#     # global X_post_std = X_post_std'
+				# end
+
+				subplot(plots_len,num_cols,(i-1)*num_cols+1); imshow(y_plot', vmin=-a,vmax=a,interpolation="none", cmap="gray")
+				axis("off"); title("Migration");#colorbar(fraction=0.046, pad=0.04);
+
+				subplot(plots_len,num_cols,(i-1)*num_cols+2); imshow(X_post[:,:,1,1]', vmin=vmin_v,vmax=vmax_v, interpolation="none", cmap="cet_rainbow4")
+				axis("off"); title("Posterior sample") #colorbar(fraction=0.046, pad=0.04);
+				
+				subplot(plots_len,num_cols,(i-1)*num_cols+3); imshow(X_post[:,:,1,2]', vmin=vmin_v,vmax=vmax_v, interpolation="none", cmap="cet_rainbow4")
+				axis("off");title("Posterior sample")  #colorbar(fraction=0.046, pad=0.04);title("Posterior sample")
+
+				subplot(plots_len,num_cols,(i-1)*num_cols+4); imshow(x_gt',  vmin=vmin_v,vmax=vmax_v,  interpolation="none", cmap="cet_rainbow4")
+				axis("off"); title(L"Reference $\mathbf{x^{*}}$") ; #colorbar(fraction=0.046, pad=0.04)
+
+				subplot(plots_len,num_cols,(i-1)*num_cols+5); imshow(x_hat' ,  vmin=vmin_v,vmax=vmax_v,  interpolation="none", cmap="cet_rainbow4")
+				axis("off"); title("Posterior mean | SSIM="*string(ssim_i)) ; #colorbar(fraction=0.046, pad=0.04)
+
+				subplot(plots_len,num_cols,(i-1)*num_cols+6); imshow(error_mean' , vmin=0,vmax=0.42, interpolation="none", cmap="magma")
+				axis("off");title("Error | RMSE="*string(rmse_i)) ;# cb = colorbar(fraction=0.046, pad=0.04)
+
+				subplot(plots_len,num_cols,(i-1)*num_cols+7); imshow(X_post_std[:,:,1,1]', vmin=0,vmax=0.42,interpolation="none", cmap="magma")
+				axis("off"); title("Posterior variance | RMS point-wise std="*string(stdtotal_i)) ;#cb =colorbar(fraction=0.046, pad=0.04)
+			end
+			tight_layout()
+			fig_name = @strdict chan_obs noise_lev_x noise_lev_init n_train e offset_start # offset_end n_offsets keep_offset_num
+			safesave(joinpath(plot_path, savename(fig_name; digits=6)*"_"*file_str*".png"), fig); close(fig)
+		end
+			
+		############# Training metric logs
+		if e != plot_every
+			sum_train = loss + logdet_train
+			sum_test = loss_test + logdet_test
+
+			fig = figure("training logs ", figsize=(10,12))
+			subplot(6,1,1); title("L2 Term: train="*string(loss[end])*" test="*string(loss_test[end]))
+			plot(range(0f0, 1f0, length=length(loss)), loss, label="train");
+			plot(range(0f0, 1f0, length=length(loss_test)),loss_test, label="test"); 
+			axhline(y=1,color="red",linestyle="--",label="Normal Noise")
+			ylim(bottom=0.,top=1.5)
+			xlabel("Parameter Update"); legend();
+
+			subplot(6,1,2); title("Logdet Term: train="*string(logdet_train[end])*" test="*string(logdet_test[end]))
+			plot(range(0f0, 1f0, length=length(logdet_train)),logdet_train);
+			plot(range(0f0, 1f0, length=length(logdet_test)),logdet_test);
+			xlabel("Parameter Update") ;
+
+			subplot(6,1,3); title("Total Objective: train="*string(sum_train[end])*" test="*string(sum_test[end]))
+			plot(range(0f0, 1f0, length=length(sum_train)),sum_train); 
+			plot(range(0f0, 1f0, length=length(sum_test)),sum_test); 
+			xlabel("Parameter Update") ;
+
+			subplot(6,1,4); title("SSIM train=$(ssim[end]) test=$(ssim_test[end])")
+			plot(range(0f0, 1f0, length=length(ssim)),ssim); 
+			plot(range(0f0, 1f0, length=length(ssim_test)),ssim_test); 
+			xlabel("Parameter Update") 
+
+			subplot(6,1,5); title("RMSE train=$(l2_cm[end]) test=$(l2_cm_test[end])")
+			plot(range(0f0, 1f0, length=length(l2_cm)),l2_cm); 
+			plot(range(0f0, 1f0, length=length(l2_cm_test)),l2_cm_test); 
+			xlabel("Parameter Update") 
+
+			subplot(6,1,6); title("RMS pointwise STD train=$(std_train[end]) test=$(std_test[end])")
+			plot(range(0f0, 1f0, length=length(std_train)),std_train); 
+			plot(range(0f0, 1f0, length=length(std_test)),std_test); 
+			xlabel("Parameter Update")
+
+			tight_layout()
+			fig_name = @strdict chan_obs noise_lev_x noise_lev_init n_train e offset_start # offset_end n_offsets keep_offset_num
+			safesave(joinpath(plot_path, savename(fig_name; digits=6)*"_log.png"), fig); close(fig)
+		end
+
+	end
+
+	if(mod(e,save_every)==0) 
+		unet_model = G.sum_net.model|> cpu;
+		G_save = deepcopy(G);
+		reset!(G_save.sum_net); # clear params to not save twice
+		Params = get_params(G_save) |> cpu;
+		# save_dict = @strdict chan_obs unet_lev unet_model n_train e noise_lev_x noise_lev_init lr n_hidden L K Params loss logdet_train l2_cm ssim loss_test logdet_test l2_cm_test ssim_test batch_size offset_start # offset_end n_offsets keep_offset_num; 
+		save_dict = @strdict n_train e unet_lev n_hidden L K Params unet_model
+		@tagsave(
+			joinpath(weights_path, savename(save_dict, "bson"; digits=6)),
+			save_dict;
+			safe=true
+		);
+	end
+end
diff --git a/CNF/update_fiducials.jl b/CNF/update_fiducials.jl
new file mode 100644
index 0000000..643b989
--- /dev/null
+++ b/CNF/update_fiducials.jl
@@ -0,0 +1,57 @@
+import Pkg; Pkg.instantiate()
+using LinearAlgebra, Random, Statistics
+using ImageQualityIndexes
+using Images, BSON, JLD2
+using FFTW
+using Augmentor
+using Random
+using DrWatson
+using InvertibleNetworks, Flux, UNet
+using PyPlot, SlimPlotting
+
+Random.seed!(2023)
+
+include("utils.jl")
+include("../config.jl")
+
+using .Config
+
+T = Float32
+device = gpu
+num_post_samples = 64
+batch_size = 8
+
+params = Config.get_parameters()
+
+offsets = params["n_offsets"]
+down_rate = params["down_rate"]
+nx = params["nx"] ÷ 2
+nz = params["nz"] ÷ 2
+
+j, nsamples, ntrain, epochs = parse.(Int, ARGS[1:4])
+
+net_path = "weights/$j/CNF/K=9_L=3_e=$(epochs)_n_hidden=64_n_train=$(ntrain)_unet_lev=4.bson"
+G = load_trained_network(net_path, offsets);
+
+x0 = zeros(Float32, nx, nz, 1, nsamples);
+@load "data/cigs_iteration_j=$j.jld2" CIGs;
+
+Y = CIGs[:, :, :, 1, :];
+Y = permutedims(Y, [2, 3, 1, 4]);
+
+for idx in 1:nsamples
+    y_hat = Y[:, :, :, idx:idx];
+    x_temp = zeros(nx, nz, 1, 1);
+    X_post = posterior_sampler(G, y_hat, x_temp; device=device, num_samples=num_post_samples, batch_size=batch_size) |> cpu
+    
+    x0[:, :, 1:1, idx:idx] = mean(X_post, dims=4)
+end
+
+# plot_velocity_model(x0[:, :, 1, 1]', "velocity.png", params);
+
+# TODO: Streamline plotting
+# imshow(x0[:, :, 1, 1]', vmin=minimum(x0), vmax=maximum(x0),  interpolation="none", cmap="cet_rainbow4")
+# savefig("plots/velocity.png", bbox_inches="tight", dpi=300)
+# close()
+
+JLD2.@save "data/posteriors_iteration_j=$j.jld2" x0
diff --git a/CNF/utils.jl b/CNF/utils.jl
new file mode 100644
index 0000000..d5897bc
--- /dev/null
+++ b/CNF/utils.jl
@@ -0,0 +1,420 @@
+function posterior_sampler(G, y, x; device=gpu, num_samples=1, batch_size=16)
+  size_x = size(x)
+    # make samples from posterior for train sample
+  X_forward = randn(Float32, size_x[1:end-1]...,batch_size) |> device
+    Y_train_latent_repeat = repeat(y |>cpu, 1, 1, 1, batch_size) |> device
+    _, Zy_fixed_train, _ = G.forward(X_forward, Y_train_latent_repeat); #needs to set the proper sizes here
+
+    X_post_train = zeros(Float32, size_x[1:end-1]...,num_samples)
+    for i in 1:div(num_samples, batch_size)
+      ZX_noise_i = randn(Float32, size_x[1:end-1]...,batch_size)|> device
+
+      X_post_train[:,:,:, (i-1)*batch_size+1 : i*batch_size] = G.inverse(
+        ZX_noise_i,
+        Zy_fixed_train
+      ) |> cpu;
+  end
+  X_post_train
+end
+
+function get_cm_l2_ssim(G, X_batch, Y_batch; device=gpu, num_samples=1)
+        # needs to be towards target so that it generalizes accross iteration
+        num_test = size(Y_batch)[end]
+        l2_total = 0
+        ssim_total = 0
+        std_total = 0
+        #get cm for each element in batch
+        for i in 1:num_test
+            y   = Y_batch[:,:,:,i:i]
+            x   = X_batch[:,:,:,i:i]
+
+            X_post = posterior_sampler(G, y, x; device=device, num_samples=num_samples, batch_size=batch_size)
+            x_hat =   mean(X_post; dims=4)[:,:,1,1]
+            x_gt =  (x[:,:,1,1]) |> cpu
+            X_post_std = std(X_post, dims=4)
+            ssim_total += assess_ssim(x_hat, x_gt)
+            l2_total   += sqrt(mean((x_hat - x_gt).^2))
+            std_total  += sqrt(mean(X_post_std.^2))
+        end
+    return l2_total / num_test, ssim_total / num_test, std_total / num_test
+end
+
+function get_loss(G, X_batch, Y_batch; device=gpu, batch_size=16)
+    l2_total = 0
+    logdet_total = 0
+    num_batches = div(size(Y_batch)[end], batch_size)
+    for i in 1:num_batches
+        x_i = X_batch[:,:,:,(i-1)*batch_size+1 : i*batch_size]
+        y_i = Y_batch[:,:,:,(i-1)*batch_size+1 : i*batch_size]
+
+        x_i .+= noise_lev_x*randn(Float32, size(x_i)); 
+        y_i .+= noise_lev_y*randn(Float32, size(y_i)); 
+        Zx, Zy, lgdet = G.forward(x_i|> device, y_i|> device) |> cpu;
+        l2_total     += norm(Zx)^2 / (N*batch_size)
+        logdet_total += lgdet / N
+    end
+
+    return l2_total / (num_batches), logdet_total / (num_batches)
+end
+
+
+
+# # load synthoseis data for training
+# using NPZ
+
+# # Function to read all .npy files in a directory and store in one array
+# function load_npy(directory)
+#     # Get a list of .npy files in the directory
+#     npy_files = filter(f -> endswith(f, ".npy"), readdir(directory; join=true))
+
+#     data_array = zeros(256, 256, 1, size(npy_files)[1])
+#     for (i, file) in enumerate(npy_files)
+#         # println("Reading file: $file")
+#         data = npzread(file)
+#         data_array[:, :, 1, i] = data
+#     end
+
+#     return data_array
+# end
+
+
+function load_compass()
+    repeat_train = Int(num_train/800)
+    
+    data_path = "/slimdata/rafaeldata/fwiuq_eod/rtms_oed.jld2" 
+    v_train = JLD2.jldopen(data_path, "r")["m_train"][:, :, :, 1:800];
+    v_train = cat([v_train for _ in 1:repeat_train]..., dims=4)
+    v_test = JLD2.jldopen(data_path, "r")["m_train"][:, :, :, 851:1000]
+    v_train = cat(v_train, v_test, dims=4)
+    return v_train
+end
+
+function load_compass_all()
+    data_path = "/slimdata/rafaeldata/fwiuq_eod/rtms_oed.jld2"
+    v_train = JLD2.jldopen(data_path, "r")["m_train"][:, :, :, 1:1000];
+    return v_train
+end
+
+function load_synthoseis_data()
+    directory = "/slimdata/tunadata/ambient_synthetic_test/synthoseis/synthoseis_data/3209_synthetic_velocities_256x256"
+    v_train = load_npy(directory);
+    v_train ./= 1000   # The resulting velocity is closer to compass model and the noise level takes best effect
+    # transpose velocity to make them consistent with compass model
+    v_train = permutedims(v_train, (2, 1, 3, 4)) 
+
+    return v_train
+end
+
+function load_synthoseis_w_wb_data(; nwb=20)
+    m_train = load_synthoseis_data()
+    tmp = zeros(size(m_train)[1], size(m_train)[2] + nwb, size(m_train)[3], size(m_train)[4])
+    tmp[:, nwb + 1:end, :, :] = m_train
+    tmp[:, 1:nwb, :, :] .= 1.48
+    m_train = deepcopy(tmp)
+
+    return m_train
+end
+
+function load_synthoseis_all_bkg(gaussian; nwb=20)
+    m_train = load_synthoseis_data()
+    m_back_all = zeros(size(m_train)[1], size(m_train)[2] + nwb, size(m_train)[3], size(m_train)[4])
+    m_back_all[:, 1:nwb, :, :] .= 1.48
+    for i in 1:size(m_train)[end]
+        x = m_train[:, :, 1, i]
+        # pert = 40
+        # pl = ElasticDistortion(pert, pert);
+        # xnew = augment(x, pl);
+        # no perturbation
+        xnew = x
+        m_back = 1f0./Float32.(imfilter(1f0./xnew, Kernel.gaussian(gaussian)));
+        
+        m_back_all[:, nwb+1:end, 1, i] = m_back
+    end
+
+    return m_back_all
+end
+
+function load_rand_prior_bkg()
+    # random walk. Draw from the fixed region
+    prior_vel = gen_prior_vel(vel_all)
+    m_back = 1f0./Float32.(imfilter(1f0./prior_vel, Kernel.gaussian(10)))
+    plot_vel_random(m_back, i, plot_path)
+end
+
+function background_1d_average()
+    m_mean = mean(m_train, dims=4)[:,:,1,1]
+    wb = maximum(find_water_bottom(m_mean.-minimum(m_mean)))
+    m0 = deepcopy(m_mean)
+    m0[:,wb+1:end] .= 1f0./Float32.(imfilter(1f0./m_mean[:,wb+1:end], Kernel.gaussian(10)))
+    return m0
+end
+
+function background_1d_gradient()
+    m_1d_gradient = reshape(repeat(range(minimum(m_train), stop=maximum(m_train), length=n[2]), inner=n[1]), n)
+    m0 = 1f0./Float32.(imfilter(1f0./m_1d_gradient, Kernel.gaussian(10)))
+    return m0
+end
+
+function load_y()
+    n_tot_sample = num_train + num_test
+    m0_train = deepcopy(m_train)
+    grad_train = zeros(Float32, size(m_train, 1), size(m_train, 2), keep_offset_num, n_tot_sample)
+    keep_offset_idx = div(n_offsets,2)+1-div(keep_offset_num, 2):div(n_offsets,2)+1+div(keep_offset_num, 2)
+
+    for i = 1:num_train
+        # misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+        if rtm_type == "ext-rtm"
+            if i <= 800
+                misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+                grad_train[:,:,1:keep_offset_num,i] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+                # grad_train[:,:,keep_offset_num+1:end,i] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-4"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+
+                # gen_ext_rtm.jl saves rtm with rtm.data, which has only 2 dimensions. size = (512, 256).
+                # rtm = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+                # rtm = reshape(rtm, (size(rtm)[1], size(rtm)[2], 1, 1))
+                # grad_train[:,:,:,i] = rtm
+            elseif 801 <= i <= 1600
+                i -= 800
+                misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+                grad_train[:,:,:,i+800] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+                # rtm = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname * "-2"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+                # rtm = reshape(rtm, (size(rtm)[1], size(rtm)[2], 1, 1))
+                # i += 800
+                # grad_train[:,:,:,i] = rtm
+            elseif 1601 <= i <= 2400
+                i -= 1600
+                misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+                grad_train[:,:,:,i+1600] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+                # rtm = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+                # rtm = reshape(rtm, (size(rtm)[1], size(rtm)[2], 1, 1))
+                # i += 1600
+                # grad_train[:,:,:,i] = rtm
+            elseif 2401 <= i <= 3200
+                i -= 2400
+                misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+                grad_train[:,:,:,i+2400] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-4"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+                # rtm = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+                # rtm = reshape(rtm, (size(rtm)[1], size(rtm)[2], 1, 1))
+                # i += 2400
+                # grad_train[:,:,:,i] = rtm
+            end
+        else
+            load_grad = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+            for k = 1:size(grad_train, 3)
+                grad_train[:,:,k,i] = load_grad
+            end
+        end
+        if background_type == "1d-average" || background_type == "1d-gradient"
+            m0_train[:,:,1,i] .= m0
+        end
+    end
+
+    # append test samples to the end of `grad_train`
+    for (idx, i) in enumerate(851:1000)
+        misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+        if rtm_type == "ext-rtm"
+            grad_train[:,:,1:keep_offset_num,num_train+idx] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+            # grad_train[:,:,keep_offset_num+1:end,num_train+idx] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-4"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+    
+            # grad_train[:,:,keep_offset_num+1:end,num_train+idx] = permutedims(JLD2.jldopen(joinpath("/slimdata/zyin62/francis_data/plots/FWIUQ/gen_ext-rtm-1d-average", savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+            # rtm = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+            # rtm = reshape(rtm, (size(rtm)[1], size(rtm)[2], 1, 1))
+            # grad_train[:,:,:,num_train+idx] = rtm
+        end
+    end
+
+    # for synthoseis cig data, depth scaling has been done in gen_cig.jl
+    # for z = 1:n[2]
+    #     grad_train[:,z,:,:] .*= z * d[2]
+    #     if z <= wb
+    #         # grad_train[:,z,:,:] .= 0 # no mute
+    #     end
+    # end
+
+    println("finish loading training and testing data")
+    return grad_train
+end
+
+
+function load_y_test(cig_fname)
+    grad_train = zeros(Float32, size(m_train, 1), size(m_train, 2), keep_offset_num, num_test)
+    keep_offset_idx = div(n_offsets,2)+1-div(keep_offset_num, 2):div(n_offsets,2)+1+div(keep_offset_num, 2)
+    
+    # append test samples to the end of `grad_train`
+    for (idx, i) in enumerate(start_test_idx:start_test_idx+149)
+        misc_dict = @strdict f0 dtD dtS nbl timeD timeR nsrc nxrec n d o i snr offset_start offset_end n_offsets
+        if rtm_type == "ext-rtm"
+            grad_train[:,:,:,idx] = permutedims(JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", cig_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"][keep_offset_idx, :, :], [2, 3, 1])
+    
+            # rtm = JLD2.jldopen(joinpath(joinpath("/slimdata/yunlindata/plots/FWIUQ", rtm_fname * "-1"), savename(misc_dict, "jld2"; digits=6)), "r")["rtm"]
+            # rtm = reshape(rtm, (size(rtm)[1], size(rtm)[2], 1, 1))
+            # grad_train[:,:,1,num_train+idx] = rtm
+        end
+    end
+
+    println("finish loading grad_train")
+    
+    if !synthoseis
+        for z = 1:n[2]
+            grad_train[:,z,:,:] .*= z * d[2]
+            if z <= wb
+                #grad_train[:,z,:,:] .= 0
+            end
+        end
+    end
+    
+    # normalize rtms
+    grad_train ./= quantile(abs.(vec(grad_train[:,:,:,end-149:end])),0.9999);
+     
+    return grad_train
+end 
+
+
+
+function load_trained_network(net_path, chan_obs; device=gpu)
+    # n_x, n_y, chan_target, n_test = size(X_test)
+    # N = n_x*n_y*chan_target
+    chan_cond = 1
+
+    # Summary network parametrs
+    unet_lev = 4
+    unet = Chain(Unet(chan_obs, chan_cond, unet_lev)|> device);
+    trainmode!(unet, true); 
+    unet = FluxBlock(unet);
+    
+    # Create conditional network
+    # L = 3
+    # K = 9
+    # n_hidden = 64
+    # low = 0.5f0
+    
+    Random.seed!(123);
+    
+    unet_lev = BSON.load(net_path)["unet_lev"];
+    n_hidden = BSON.load(net_path)["n_hidden"];
+    L = BSON.load(net_path)["L"];
+    K = BSON.load(net_path)["K"];
+    
+    unet = Unet(chan_obs,1,unet_lev);
+    trainmode!(unet, true);
+    unet = FluxBlock(Chain(unet)) |> device;
+    
+    cond_net = NetworkConditionalGlow(1, 1, n_hidden,  L, K;  freeze_conv=true, split_scales=true, activation=SigmoidLayer(low=0.5f0,high=1.0f0)) |> device;
+    G        = SummarizedNet(cond_net, unet)
+    
+    Params = BSON.load(net_path)["Params"]; 
+    set_params!(G,Params)
+    
+    # Load in unet summary net
+    G.sum_net.model = BSON.load(net_path)["unet_model"]; 
+    G = G |> device;
+
+    return G
+end
+
+function plot_cig(cig, plot_path, cig_img_fname, params)
+    println("start plotting CIG")
+    cig_fs = 40
+    down_rate = params["down_rate"]
+    d = params["d"] .* down_rate
+
+    nx = params["nx"] ÷ down_rate
+    nz = params["nz"] ÷ down_rate
+
+    n = (nx, nz)
+    n_offsets = params["n_offsets"]
+    offset_start = params["offset_start"]
+    offset_end = params["offset_end"]
+
+    # start plotting
+    y = cig
+
+    PyPlot.rc("figure", titlesize=cig_fs)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=cig_fs); PyPlot.rc("ytick", labelsize=cig_fs)
+    PyPlot.rc("axes", labelsize=cig_fs)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=cig_fs)     # Default fontsize for titles
+
+    ### X, Z position in km
+    xpos = 3.4f3 ÷ 2
+    zpos = 2.4f3 ÷ 2
+    xgrid = Int(round(xpos / d[1]))
+    zgrid = Int(round(zpos / d[2]))
+
+    # Create a figure and a 2x2 grid of subplots
+    fig, axs = subplots(2, 2, figsize=(20,12), gridspec_kw = Dict("width_ratios" => [3, 1], "height_ratios" => [1, 3]))
+    # Adjust the spacing between the plots
+    subplots_adjust(hspace=0.0, wspace=0.0)
+
+    vmin1, vmax1 = (-1, 1) .* quantile(abs.(vec(y[:,zgrid,:,1])), 0.99)
+    vmin2, vmax2 = (-1, 1) .* quantile(abs.(vec(y[:,:,div(n_offsets,2)+1,1])), 0.88)
+    vmin3, vmax3 = (-1, 1) .* quantile(abs.(vec(y[xgrid,:,:,1])), 0.999)
+    sca(axs[1, 1])
+    # Top left subplot
+    axs[1, 1].imshow(y[:,zgrid,:,1]', aspect="auto", cmap="gray", interpolation="none",vmin=vmin1, vmax=vmax1,
+    	extent=(0f0, (n[1]-1)*d[1], offset_start, offset_end))
+    axs[1, 1].set_ylabel("Offset [m]", fontsize=cig_fs)
+    axs[1, 1].set_xticklabels([])
+    axs[1, 1].set_xlabel("")
+    hlines(y=0, colors=:b, xmin=0, xmax=(n[1]-1)*d[1], linewidth=3)
+    vlines(x=xpos, colors=:b, ymin=offset_start, ymax=offset_end, linewidth=3)
+    # Bottom left subplot
+    sca(axs[2, 1])
+    axs[2, 1].imshow(y[:,:,div(n_offsets,2)+1,1]', aspect="auto", cmap="gray", interpolation="none",vmin=vmin2, vmax=vmax2,
+    extent=(0f0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0f0))
+    axs[2, 1].set_xlabel("X [m]", fontsize=cig_fs)
+    axs[2, 1].set_ylabel("Z [m]", fontsize=cig_fs)
+    axs[2, 1].set_xticks([0, 1000, 2000, 3000, 4000, 5000])
+    axs[2, 1].set_xticklabels(["0", "1000", "2000", "3000", "4000", "5000"])
+    axs[2, 1].set_yticks([1000, 2000, 3000])
+    axs[2, 1].set_yticklabels(["1000", "2000", "3000"])
+    # axs[2, 2].get_shared_x_axes().join(axs[1, 1], axs[2, 1])
+    vlines(x=xpos, colors=:b, ymin=0, ymax=(n[2]-1)*d[2], linewidth=3)
+    hlines(y=zpos, colors=:b, xmin=0, xmax=(n[1]-1)*d[1], linewidth=3)
+    # Top right subplot
+    axs[1, 2].set_visible(false)
+    # Bottom right subplot
+    sca(axs[2, 2])
+    axs[2, 2].imshow(y[xgrid,:,:,1], aspect="auto", cmap="gray", interpolation="none",vmin=vmin3, vmax=vmax3,
+    extent=(offset_start, offset_end, (n[2]-1)*d[2], 0f0))
+    axs[2, 2].set_xlabel("Offset [m]", fontsize=cig_fs)
+    # Share y-axis with bottom left
+    # axs[2, 2].get_shared_y_axes().join(axs[2, 2], axs[2, 1])
+    axs[2, 2].set_yticklabels([])
+    axs[2, 2].set_ylabel("")
+    vlines(x=0, colors=:b, ymin=0, ymax=(n[2]-1)*d[2], linewidth=3)
+    hlines(y=zpos, colors=:b, xmin=offset_end, xmax=offset_start, linewidth=3)
+    # Remove the space between subplots and hide the spines
+    for ax in reshape(axs, :)
+        for spine in ["top", "right", "bottom", "left"]
+            ax.spines[spine].set_visible(false)
+        end
+    end
+
+    PyPlot.savefig(joinpath(plot_path, cig_img_fname), bbox_inches="tight", dpi=300);
+    close(fig)
+    
+    # Zoom in offset image
+    PyPlot.rcdefaults()
+    matplotlib.pyplot.close()
+end
+
+function plot_velocity_model(x, filename, params)
+    vmin = quantile(vec(x), 0.05)  # 5th percentile
+    vmax = quantile(vec(x), 0.95)  # 95th percentile
+
+    d = params["d"]
+    down_rate = params["down_rate"]
+
+    nx = params["nx"] ÷ down_rate
+    nz = params["nz"] ÷ down_rate
+
+    n = (nx, nz)
+
+    fig, ax = subplots(figsize=(20,12)) 
+    extentfull = (0f0, (n[1]-1)*d[1], (n[end]-1)*d[end], 0f0)
+    cax = ax.imshow(x', vmin=vmin, vmax=vmax, extent=extentfull, aspect=0.45*(extentfull[2]-extentfull[1])/(extentfull[3]-extentfull[4]))
+    ax.set_xlabel("X [m]", fontsize=40)
+    ax.set_ylabel("Z [m]", fontsize=40)
+    savefig(filename, bbox_inches="tight", dpi=300)
+    close(fig)
+end
diff --git a/FNO/Manifest.toml b/FNO/Manifest.toml
new file mode 100644
index 0000000..acfb153
--- /dev/null
+++ b/FNO/Manifest.toml
@@ -0,0 +1,2016 @@
+# This file is machine-generated - editing it directly is not advised
+
+julia_version = "1.8.5"
+manifest_format = "2.0"
+project_hash = "4e3bbe64b6aee6c66c14b4fdb4ad81dda34ad3fb"
+
+[[deps.AbstractFFTs]]
+deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
+git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef"
+uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+version = "1.5.0"
+
+[[deps.AbstractTrees]]
+git-tree-sha1 = "2d9c9a55f9c93e8887ad391fbae72f8ef55e1177"
+uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+version = "0.4.5"
+
+[[deps.Accessors]]
+deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown", "Requires", "Test"]
+git-tree-sha1 = "b392ede862e506d451fc1616e79aa6f4c673dab8"
+uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+version = "0.1.38"
+
+[[deps.Adapt]]
+deps = ["LinearAlgebra", "Requires"]
+git-tree-sha1 = "cde29ddf7e5726c9fb511f340244ea3481267608"
+uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+version = "3.7.2"
+
+[[deps.ArgCheck]]
+git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4"
+uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+version = "2.3.0"
+
+[[deps.ArgParse]]
+deps = ["Logging", "TextWrap"]
+git-tree-sha1 = "22cf435ac22956a7b45b0168abbc871176e7eecc"
+uuid = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+version = "1.2.0"
+
+[[deps.ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+version = "1.1.1"
+
+[[deps.ArnoldiMethod]]
+deps = ["LinearAlgebra", "Random", "StaticArrays"]
+git-tree-sha1 = "d57bd3762d308bded22c3b82d033bff85f6195c6"
+uuid = "ec485272-7323-5ecc-a04f-4719b315124d"
+version = "0.4.0"
+
+[[deps.ArrayInterface]]
+deps = ["Adapt", "LinearAlgebra", "Requires", "SparseArrays", "SuiteSparse"]
+git-tree-sha1 = "16267cf279190ca7c1b30d020758ced95db89cd0"
+uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+version = "7.5.1"
+
+[[deps.Artifacts]]
+uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+
+[[deps.AssetRegistry]]
+deps = ["Distributed", "JSON", "Pidfile", "SHA", "Test"]
+git-tree-sha1 = "b25e88db7944f98789130d7b503276bc34bc098e"
+uuid = "bf4720bc-e11a-5d0c-854e-bdca1663c893"
+version = "0.1.0"
+
+[[deps.Atomix]]
+deps = ["UnsafeAtomics"]
+git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be"
+uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
+version = "0.1.0"
+
+[[deps.AxisAlgorithms]]
+deps = ["LinearAlgebra", "Random", "SparseArrays", "WoodburyMatrices"]
+git-tree-sha1 = "01b8ccb13d68535d73d2b0c23e39bd23155fb712"
+uuid = "13072b0f-2c55-5437-9ae7-d433b7a33950"
+version = "1.1.0"
+
+[[deps.AxisArrays]]
+deps = ["Dates", "IntervalSets", "IterTools", "RangeArrays"]
+git-tree-sha1 = "16351be62963a67ac4083f748fdb3cca58bfd52f"
+uuid = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+version = "0.4.7"
+
+[[deps.BFloat16s]]
+deps = ["LinearAlgebra", "Printf", "Random", "Test"]
+git-tree-sha1 = "dbf84058d0a8cbbadee18d25cf606934b22d7c66"
+uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
+version = "0.4.2"
+
+[[deps.BangBang]]
+deps = ["Accessors", "ConstructionBase", "InitialValues", "LinearAlgebra", "Requires"]
+git-tree-sha1 = "e2144b631226d9eeab2d746ca8880b7ccff504ae"
+uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+version = "0.4.3"
+
+[[deps.Base64]]
+uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+
+[[deps.Baselet]]
+git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e"
+uuid = "9718e550-a3fa-408a-8086-8db961cd8217"
+version = "0.1.1"
+
+[[deps.BatchedRoutines]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "441db9f0399bcfb4eeb8b891a6b03f7acc5dc731"
+uuid = "a9ab73d0-e05c-5df1-8fde-d6a4645b8d8e"
+version = "0.2.2"
+
+[[deps.BitBasis]]
+deps = ["LinearAlgebra", "StaticArrays"]
+git-tree-sha1 = "fe126b4ca000cdd8de9bf2b94ec40ec27b4b1f7f"
+uuid = "50ba71b6-fa0f-514d-ae9a-0916efc90dcf"
+version = "0.9.6"
+
+[[deps.BitFlags]]
+git-tree-sha1 = "0691e34b3bb8be9307330f88d1a3c3f25466c24d"
+uuid = "d1d4a3ce-64b1-5f1a-9ba4-7e7e69966f35"
+version = "0.1.9"
+
+[[deps.BitTwiddlingConvenienceFunctions]]
+deps = ["Static"]
+git-tree-sha1 = "f21cfd4950cb9f0587d5067e69405ad2acd27b87"
+uuid = "62783981-4cbd-42fc-bca8-16325de8dc4b"
+version = "0.1.6"
+
+[[deps.Blink]]
+deps = ["Base64", "Distributed", "HTTP", "JSExpr", "JSON", "Lazy", "Logging", "MacroTools", "Mustache", "Mux", "Pkg", "Reexport", "Sockets", "WebIO"]
+git-tree-sha1 = "bc93511973d1f949d45b0ea17878e6cb0ad484a1"
+uuid = "ad839575-38b3-5650-b840-f874b8c74a25"
+version = "0.12.9"
+
+[[deps.BufferedStreams]]
+git-tree-sha1 = "6863c5b7fc997eadcabdbaf6c5f201dc30032643"
+uuid = "e1450e63-4bb3-523b-b2a4-4ffa8c0fd77d"
+version = "1.2.2"
+
+[[deps.Bzip2_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "8873e196c2eb87962a2048b3b8e08946535864a1"
+uuid = "6e34b625-4abd-537c-b88f-471c36dfa7a0"
+version = "1.0.8+2"
+
+[[deps.CEnum]]
+git-tree-sha1 = "eb4cb44a499229b3b8426dcfb5dd85333951ff90"
+uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+version = "0.4.2"
+
+[[deps.CPUSummary]]
+deps = ["CpuId", "IfElse", "PrecompileTools", "Static"]
+git-tree-sha1 = "5a97e67919535d6841172016c9530fd69494e5ec"
+uuid = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"
+version = "0.2.6"
+
+[[deps.CUDA]]
+deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Preferences", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "UnsafeAtomicsLLVM"]
+git-tree-sha1 = "b2e58316a391bddc8b9234ca4aa484eba6878ef2"
+uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
+version = "4.4.2"
+
+[[deps.CUDA_Driver_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"]
+git-tree-sha1 = "498f45593f6ddc0adff64a9310bb6710e851781b"
+uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc"
+version = "0.5.0+1"
+
+[[deps.CUDA_Runtime_Discovery]]
+deps = ["Libdl"]
+git-tree-sha1 = "38f830504358e9972d2a0c3e5d51cb865e0733df"
+uuid = "1af6417a-86b4-443c-805f-a4643ffb695f"
+version = "0.2.4"
+
+[[deps.CUDA_Runtime_jll]]
+deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "5248d9c45712e51e27ba9b30eebec65658c6ce29"
+uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
+version = "0.6.0+0"
+
+[[deps.CUDNN_jll]]
+deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "75923dce4275ead3799b238e10178a68c07dbd3b"
+uuid = "62b44479-cb7b-5706-934f-f13b2eb2e645"
+version = "8.9.4+0"
+
+[[deps.Cairo_jll]]
+deps = ["Artifacts", "Bzip2_jll", "CompilerSupportLibraries_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "JLLWrappers", "LZO_jll", "Libdl", "Pixman_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Zlib_jll", "libpng_jll"]
+git-tree-sha1 = "009060c9a6168704143100f36ab08f06c2af4642"
+uuid = "83423d85-b0ee-5818-9007-b63ccbeb887a"
+version = "1.18.2+1"
+
+[[deps.CatIndices]]
+deps = ["CustomUnitRanges", "OffsetArrays"]
+git-tree-sha1 = "a0f80a09780eed9b1d106a1bf62041c2efc995bc"
+uuid = "aafaddc9-749c-510e-ac4f-586e18779b91"
+version = "0.2.2"
+
+[[deps.ChainRules]]
+deps = ["Adapt", "ChainRulesCore", "Compat", "Distributed", "GPUArraysCore", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "SparseInverseSubset", "Statistics", "StructArrays", "SuiteSparse"]
+git-tree-sha1 = "be227d253d132a6d57f9ccf5f67c0fb6488afd87"
+uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+version = "1.71.0"
+
+[[deps.ChainRulesCore]]
+deps = ["Compat", "LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "3e4b134270b372f2ed4d4d0e936aabaefc1802bc"
+uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+version = "1.25.0"
+
+[[deps.ChangesOfVariables]]
+deps = ["InverseFunctions", "LinearAlgebra", "Test"]
+git-tree-sha1 = "799b25ca3a8a24936ae7b5c52ad194685fc3e6ef"
+uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
+version = "0.1.9"
+
+[[deps.CloseOpenIntervals]]
+deps = ["Static", "StaticArrayInterface"]
+git-tree-sha1 = "05ba0d07cd4fd8b7a39541e31a7b0254704ea581"
+uuid = "fb6a15b2-703c-40df-9091-08a04967cfa9"
+version = "0.1.13"
+
+[[deps.Clustering]]
+deps = ["Distances", "LinearAlgebra", "NearestNeighbors", "Printf", "Random", "SparseArrays", "Statistics", "StatsBase"]
+git-tree-sha1 = "9ebb045901e9bbf58767a9f34ff89831ed711aae"
+uuid = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5"
+version = "0.15.7"
+
+[[deps.CodeTracking]]
+deps = ["InteractiveUtils", "UUIDs"]
+git-tree-sha1 = "7eee164f122511d3e4e1ebadb7956939ea7e1c77"
+uuid = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
+version = "1.3.6"
+
+[[deps.CodecZlib]]
+deps = ["TranscodingStreams", "Zlib_jll"]
+git-tree-sha1 = "bce6804e5e6044c6daab27bb533d1295e4a2e759"
+uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
+version = "0.7.6"
+
+[[deps.ColorSchemes]]
+deps = ["ColorTypes", "ColorVectorSpace", "Colors", "FixedPointNumbers", "PrecompileTools", "Random"]
+git-tree-sha1 = "13951eb68769ad1cd460cdb2e64e5e95f1bf123d"
+uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
+version = "3.27.0"
+
+[[deps.ColorTypes]]
+deps = ["FixedPointNumbers", "Random"]
+git-tree-sha1 = "b10d0b65641d57b8b4d5e234446582de5047050d"
+uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
+version = "0.11.5"
+
+[[deps.ColorVectorSpace]]
+deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "Requires", "Statistics", "TensorCore"]
+git-tree-sha1 = "a1f44953f2382ebb937d60dafbe2deea4bd23249"
+uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4"
+version = "0.10.0"
+
+[[deps.Colors]]
+deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
+git-tree-sha1 = "362a287c3aa50601b0bc359053d5c2468f0e7ce0"
+uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
+version = "0.12.11"
+
+[[deps.Combinatorics]]
+git-tree-sha1 = "08c8b6831dc00bfea825826be0bc8336fc369860"
+uuid = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+version = "1.0.2"
+
+[[deps.CommonSubexpressions]]
+deps = ["MacroTools"]
+git-tree-sha1 = "cda2cfaebb4be89c9084adaca7dd7333369715c5"
+uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
+version = "0.3.1"
+
+[[deps.Compat]]
+deps = ["Dates", "LinearAlgebra", "TOML", "UUIDs"]
+git-tree-sha1 = "8ae8d32e09f0dcf42a36b90d4e17f5dd2e4c4215"
+uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
+version = "4.16.0"
+
+[[deps.CompilerSupportLibraries_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
+version = "1.0.1+0"
+
+[[deps.CompositionsBase]]
+git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad"
+uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b"
+version = "0.1.2"
+
+[[deps.ComputationalResources]]
+git-tree-sha1 = "52cb3ec90e8a8bea0e62e275ba577ad0f74821f7"
+uuid = "ed09eef8-17a6-5b46-8889-db040fac31e3"
+version = "0.3.2"
+
+[[deps.ConcurrentUtilities]]
+deps = ["Serialization", "Sockets"]
+git-tree-sha1 = "ea32b83ca4fefa1768dc84e504cc0a94fb1ab8d1"
+uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb"
+version = "2.4.2"
+
+[[deps.Conda]]
+deps = ["Downloads", "JSON", "VersionParsing"]
+git-tree-sha1 = "b19db3927f0db4151cb86d073689f2428e524576"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.10.2"
+
+[[deps.ConstructionBase]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "76219f1ed5771adbb096743bff43fb5fdd4c1157"
+uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
+version = "1.5.8"
+
+[[deps.ContextVariablesX]]
+deps = ["Compat", "Logging", "UUIDs"]
+git-tree-sha1 = "25cc3803f1030ab855e383129dcd3dc294e322cc"
+uuid = "6add18c4-b38d-439d-96f6-d6bc489c04c5"
+version = "0.1.3"
+
+[[deps.CoordinateTransformations]]
+deps = ["LinearAlgebra", "StaticArrays"]
+git-tree-sha1 = "f9d7112bfff8a19a3a4ea4e03a8e6a91fe8456bf"
+uuid = "150eb455-5306-5404-9cee-2592286d6298"
+version = "0.6.3"
+
+[[deps.CpuId]]
+deps = ["Markdown"]
+git-tree-sha1 = "fcbb72b032692610bfbdb15018ac16a36cf2e406"
+uuid = "adafc99b-e345-5852-983c-f28acb93d879"
+version = "0.3.1"
+
+[[deps.CustomUnitRanges]]
+git-tree-sha1 = "1a3f97f907e6dd8983b744d2642651bb162a3f7a"
+uuid = "dc8bdbbb-1ca9-579f-8c36-e416f6a65cce"
+version = "1.0.2"
+
+[[deps.DataAPI]]
+git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe"
+uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
+version = "1.16.0"
+
+[[deps.DataStructures]]
+deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
+git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82"
+uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+version = "0.18.20"
+
+[[deps.DataValueInterfaces]]
+git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
+uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
+version = "1.0.0"
+
+[[deps.Dates]]
+deps = ["Printf"]
+uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
+
+[[deps.DefineSingletons]]
+git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c"
+uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52"
+version = "0.1.2"
+
+[[deps.DelimitedFiles]]
+deps = ["Mmap"]
+uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+
+[[deps.DiffResults]]
+deps = ["StaticArraysCore"]
+git-tree-sha1 = "782dd5f4561f5d267313f23853baaaa4c52ea621"
+uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+version = "1.1.0"
+
+[[deps.DiffRules]]
+deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"]
+git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272"
+uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
+version = "1.15.1"
+
+[[deps.Distances]]
+deps = ["LinearAlgebra", "SparseArrays", "Statistics", "StatsAPI"]
+git-tree-sha1 = "c7e3a542b999843086e2f29dac96a618c105be1d"
+uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+version = "0.10.12"
+
+[[deps.Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[[deps.DocStringExtensions]]
+deps = ["LibGit2"]
+git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d"
+uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+version = "0.9.3"
+
+[[deps.Downloads]]
+deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+version = "1.6.0"
+
+[[deps.DrWatson]]
+deps = ["Dates", "FileIO", "JLD2", "LibGit2", "MacroTools", "Pkg", "Random", "Requires", "Scratch", "UnPack"]
+git-tree-sha1 = "17c9db646d6cb3f68a90053bf46faf5312d13b36"
+uuid = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
+version = "2.18.0"
+
+[[deps.EnzymeCore]]
+deps = ["Adapt"]
+git-tree-sha1 = "8f205a601760f4798a10f138c3940f0451d95188"
+uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
+version = "0.7.8"
+
+[[deps.ExceptionUnwrapping]]
+deps = ["Test"]
+git-tree-sha1 = "dcb08a0d93ec0b1cdc4af184b26b591e9695423a"
+uuid = "460bff9d-24e4-43bc-9d9f-a8973cb893f4"
+version = "0.1.10"
+
+[[deps.Expat_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "1c6317308b9dc757616f0b5cb379db10494443a7"
+uuid = "2e619515-83b5-522b-bb60-26c02a35a201"
+version = "2.6.2+0"
+
+[[deps.ExprTools]]
+git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec"
+uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+version = "0.1.10"
+
+[[deps.FFTViews]]
+deps = ["CustomUnitRanges", "FFTW"]
+git-tree-sha1 = "cbdf14d1e8c7c8aacbe8b19862e0179fd08321c2"
+uuid = "4f61f5a4-77b1-5117-aa51-3ab5ef4ef0cd"
+version = "0.3.2"
+
+[[deps.FFTW]]
+deps = ["AbstractFFTs", "FFTW_jll", "LinearAlgebra", "MKL_jll", "Preferences", "Reexport"]
+git-tree-sha1 = "4820348781ae578893311153d69049a93d05f39d"
+uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+version = "1.8.0"
+
+[[deps.FFTW_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "4d81ed14783ec49ce9f2e168208a12ce1815aa25"
+uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a"
+version = "3.3.10+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.FileIO]]
+deps = ["Pkg", "Requires", "UUIDs"]
+git-tree-sha1 = "62ca0547a14c57e98154423419d8a342dca75ca9"
+uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+version = "1.16.4"
+
+[[deps.FileWatching]]
+uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
+
+[[deps.FillArrays]]
+deps = ["LinearAlgebra", "PDMats", "SparseArrays", "Statistics"]
+git-tree-sha1 = "6a70198746448456524cb442b8af316927ff3e1a"
+uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
+version = "1.13.0"
+
+[[deps.FixedPointNumbers]]
+deps = ["Statistics"]
+git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172"
+uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
+version = "0.8.5"
+
+[[deps.Flux]]
+deps = ["Adapt", "CUDA", "ChainRulesCore", "Functors", "LinearAlgebra", "MLUtils", "MacroTools", "NNlib", "NNlibCUDA", "OneHotArrays", "Optimisers", "Preferences", "ProgressLogging", "Random", "Reexport", "SparseArrays", "SpecialFunctions", "Statistics", "Zygote", "cuDNN"]
+git-tree-sha1 = "3e2c3704c2173ab4b1935362384ca878b53d4c34"
+uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+version = "0.13.17"
+
+[[deps.Fontconfig_jll]]
+deps = ["Artifacts", "Bzip2_jll", "Expat_jll", "FreeType2_jll", "JLLWrappers", "Libdl", "Libuuid_jll", "Zlib_jll"]
+git-tree-sha1 = "db16beca600632c95fc8aca29890d83788dd8b23"
+uuid = "a3f928ae-7b40-5064-980b-68af3947d34b"
+version = "2.13.96+0"
+
+[[deps.ForwardDiff]]
+deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
+git-tree-sha1 = "cf0fe81336da9fb90944683b8c41984b08793dad"
+uuid = "f6369f11-7733-5829-9624-2563aa707210"
+version = "0.10.36"
+
+[[deps.FreeType2_jll]]
+deps = ["Artifacts", "Bzip2_jll", "JLLWrappers", "Libdl", "Zlib_jll"]
+git-tree-sha1 = "5c1d8ae0efc6c2e7b1fc502cbe25def8f661b7bc"
+uuid = "d7e528f0-a631-5988-bf34-fe36492bcfd7"
+version = "2.13.2+0"
+
+[[deps.FriBidi_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "1ed150b39aebcc805c26b93a8d0122c940f64ce2"
+uuid = "559328eb-81f9-559d-9380-de523a88c83c"
+version = "1.0.14+0"
+
+[[deps.FunctionalCollections]]
+deps = ["Test"]
+git-tree-sha1 = "04cb9cfaa6ba5311973994fe3496ddec19b6292a"
+uuid = "de31a74c-ac4f-5751-b3fd-e18cd04993ca"
+version = "0.5.0"
+
+[[deps.Functors]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "64d8e93700c7a3f28f717d265382d52fac9fa1c1"
+uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+version = "0.4.12"
+
+[[deps.Future]]
+deps = ["Random"]
+uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
+
+[[deps.GPUArrays]]
+deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"]
+git-tree-sha1 = "2e57b4a4f9cc15e85a24d603256fe08e527f48d1"
+uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+version = "8.8.1"
+
+[[deps.GPUArraysCore]]
+deps = ["Adapt"]
+git-tree-sha1 = "2d6ca471a6c7b536127afccfa7564b5b39227fe0"
+uuid = "46192b85-c4d5-4398-a991-12ede77f4527"
+version = "0.1.5"
+
+[[deps.GPUCompiler]]
+deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Scratch", "TimerOutputs", "UUIDs"]
+git-tree-sha1 = "72b2e3c2ba583d1a7aa35129e56cf92e07c083e3"
+uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
+version = "0.21.4"
+
+[[deps.Gettext_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Libiconv_jll", "Pkg", "XML2_jll"]
+git-tree-sha1 = "9b02998aba7bf074d14de89f9d37ca24a1a0b046"
+uuid = "78b55507-aeef-58d4-861c-77aaff3498b1"
+version = "0.21.0+0"
+
+[[deps.Ghostscript_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "43ba3d3c82c18d88471cfd2924931658838c9d8f"
+uuid = "61579ee1-b43e-5ca0-a5da-69d92c66a64b"
+version = "9.55.0+4"
+
+[[deps.Glib_jll]]
+deps = ["Artifacts", "Gettext_jll", "JLLWrappers", "Libdl", "Libffi_jll", "Libiconv_jll", "Libmount_jll", "PCRE2_jll", "Zlib_jll"]
+git-tree-sha1 = "674ff0db93fffcd11a3573986e550d66cd4fd71f"
+uuid = "7746bdde-850d-59dc-9ae8-88ece973131d"
+version = "2.80.5+0"
+
+[[deps.GraphViz]]
+deps = ["FileIO", "Graphviz_jll", "Requires"]
+git-tree-sha1 = "da5580b236c5b6bb634e96ca35dcd7eaeeac1bd1"
+uuid = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
+version = "0.2.0"
+
+[[deps.Graphics]]
+deps = ["Colors", "LinearAlgebra", "NaNMath"]
+git-tree-sha1 = "d61890399bc535850c4bf08e4e0d3a7ad0f21cbd"
+uuid = "a2bd30eb-e257-5431-a919-1863eab51364"
+version = "1.1.2"
+
+[[deps.Graphite2_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "344bf40dcab1073aca04aa0df4fb092f920e4011"
+uuid = "3b182d85-2403-5c21-9c21-1e1f0cc25472"
+version = "1.3.14+0"
+
+[[deps.Graphs]]
+deps = ["ArnoldiMethod", "Compat", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"]
+git-tree-sha1 = "1dc470db8b1131cfc7fb4c115de89fe391b9e780"
+uuid = "86223c79-3864-5bf0-83f7-82e725a168b6"
+version = "1.12.0"
+
+[[deps.Graphviz_jll]]
+deps = ["Artifacts", "Cairo_jll", "Expat_jll", "JLLWrappers", "Libdl", "Pango_jll", "Pkg"]
+git-tree-sha1 = "a5d45833dda71048117e8a9828bef75c03b18b1c"
+uuid = "3c863552-8265-54e4-a6dc-903eb78fde85"
+version = "2.50.0+1"
+
+[[deps.HDF5]]
+deps = ["Compat", "HDF5_jll", "Libdl", "MPIPreferences", "Mmap", "Preferences", "Printf", "Random", "Requires", "UUIDs"]
+git-tree-sha1 = "e856eef26cf5bf2b0f95f8f4fc37553c72c8641c"
+uuid = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+version = "0.17.2"
+
+[[deps.HDF5_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LLVMOpenMP_jll", "LazyArtifacts", "LibCURL_jll", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "OpenSSL_jll", "TOML", "Zlib_jll", "libaec_jll"]
+git-tree-sha1 = "38c8874692d48d5440d5752d6c74b0c6b0b60739"
+uuid = "0234f1f7-429e-5d53-9886-15a909be8d59"
+version = "1.14.2+1"
+
+[[deps.HTTP]]
+deps = ["Base64", "CodecZlib", "ConcurrentUtilities", "Dates", "ExceptionUnwrapping", "Logging", "LoggingExtras", "MbedTLS", "NetworkOptions", "OpenSSL", "Random", "SimpleBufferStream", "Sockets", "URIs", "UUIDs"]
+git-tree-sha1 = "bc3f416a965ae61968c20d0ad867556367f2817d"
+uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+version = "1.10.9"
+
+[[deps.HarfBuzz_jll]]
+deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "Graphite2_jll", "JLLWrappers", "Libdl", "Libffi_jll"]
+git-tree-sha1 = "401e4f3f30f43af2c8478fc008da50096ea5240f"
+uuid = "2e76f6c2-a576-52d4-95c1-20adfe4de566"
+version = "8.3.1+0"
+
+[[deps.Hiccup]]
+deps = ["MacroTools", "Test"]
+git-tree-sha1 = "6187bb2d5fcbb2007c39e7ac53308b0d371124bd"
+uuid = "9fb69e20-1954-56bb-a84f-559cc56a8ff7"
+version = "0.2.2"
+
+[[deps.HistogramThresholding]]
+deps = ["ImageBase", "LinearAlgebra", "MappedArrays"]
+git-tree-sha1 = "7194dfbb2f8d945abdaf68fa9480a965d6661e69"
+uuid = "2c695a8d-9458-5d45-9878-1b8a99cf7853"
+version = "0.3.1"
+
+[[deps.HostCPUFeatures]]
+deps = ["BitTwiddlingConvenienceFunctions", "IfElse", "Libdl", "Static"]
+git-tree-sha1 = "8e070b599339d622e9a081d17230d74a5c473293"
+uuid = "3e5b6fbb-0976-4d2c-9146-d79de83f2fb0"
+version = "0.1.17"
+
+[[deps.Hwloc_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "dd3b49277ec2bb2c6b94eb1604d4d0616016f7a6"
+uuid = "e33a78d0-f292-5ffc-b300-72abe9b543c8"
+version = "2.11.2+0"
+
+[[deps.IJulia]]
+deps = ["Base64", "Conda", "Dates", "InteractiveUtils", "JSON", "Libdl", "Logging", "Markdown", "MbedTLS", "Pkg", "Printf", "REPL", "Random", "SoftGlobalScope", "Test", "UUIDs", "ZMQ"]
+git-tree-sha1 = "1702f79fa30f56b68d5b2fd6fb3a9a14ff6f9130"
+uuid = "7073ff75-c697-5162-941a-fcdaad2a7d2a"
+version = "1.25.0"
+
+[[deps.IRTools]]
+deps = ["InteractiveUtils", "MacroTools"]
+git-tree-sha1 = "950c3717af761bc3ff906c2e8e52bd83390b6ec2"
+uuid = "7869d1d1-7146-5819-86e3-90919afe41df"
+version = "0.4.14"
+
+[[deps.IfElse]]
+git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1"
+uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
+version = "0.1.1"
+
+[[deps.ImageAxes]]
+deps = ["AxisArrays", "ImageBase", "ImageCore", "Reexport", "SimpleTraits"]
+git-tree-sha1 = "2e4520d67b0cef90865b3ef727594d2a58e0e1f8"
+uuid = "2803e5a7-5153-5ecf-9a86-9b4c37f5f5ac"
+version = "0.6.11"
+
+[[deps.ImageBase]]
+deps = ["ImageCore", "Reexport"]
+git-tree-sha1 = "eb49b82c172811fd2c86759fa0553a2221feb909"
+uuid = "c817782e-172a-44cc-b673-b171935fbb9e"
+version = "0.1.7"
+
+[[deps.ImageBinarization]]
+deps = ["HistogramThresholding", "ImageCore", "LinearAlgebra", "Polynomials", "Reexport", "Statistics"]
+git-tree-sha1 = "33485b4e40d1df46c806498c73ea32dc17475c59"
+uuid = "cbc4b850-ae4b-5111-9e64-df94c024a13d"
+version = "0.3.1"
+
+[[deps.ImageContrastAdjustment]]
+deps = ["ImageBase", "ImageCore", "ImageTransformations", "Parameters"]
+git-tree-sha1 = "eb3d4365a10e3f3ecb3b115e9d12db131d28a386"
+uuid = "f332f351-ec65-5f6a-b3d1-319c6670881a"
+version = "0.3.12"
+
+[[deps.ImageCore]]
+deps = ["ColorVectorSpace", "Colors", "FixedPointNumbers", "MappedArrays", "MosaicViews", "OffsetArrays", "PaddedViews", "PrecompileTools", "Reexport"]
+git-tree-sha1 = "b2a7eaa169c13f5bcae8131a83bc30eff8f71be0"
+uuid = "a09fc81d-aa75-5fe9-8630-4744c3626534"
+version = "0.10.2"
+
+[[deps.ImageCorners]]
+deps = ["ImageCore", "ImageFiltering", "PrecompileTools", "StaticArrays", "StatsBase"]
+git-tree-sha1 = "24c52de051293745a9bad7d73497708954562b79"
+uuid = "89d5987c-236e-4e32-acd0-25bd6bd87b70"
+version = "0.1.3"
+
+[[deps.ImageDistances]]
+deps = ["Distances", "ImageCore", "ImageMorphology", "LinearAlgebra", "Statistics"]
+git-tree-sha1 = "08b0e6354b21ef5dd5e49026028e41831401aca8"
+uuid = "51556ac3-7006-55f5-8cb3-34580c88182d"
+version = "0.2.17"
+
+[[deps.ImageFiltering]]
+deps = ["CatIndices", "ComputationalResources", "DataStructures", "FFTViews", "FFTW", "ImageBase", "ImageCore", "LinearAlgebra", "OffsetArrays", "PrecompileTools", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "TiledIteration"]
+git-tree-sha1 = "432ae2b430a18c58eb7eca9ef8d0f2db90bc749c"
+uuid = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
+version = "0.7.8"
+
+[[deps.ImageIO]]
+deps = ["FileIO", "IndirectArrays", "JpegTurbo", "LazyModules", "Netpbm", "OpenEXR", "PNGFiles", "QOI", "Sixel", "TiffImages", "UUIDs"]
+git-tree-sha1 = "437abb322a41d527c197fa800455f79d414f0a3c"
+uuid = "82e4d734-157c-48bb-816b-45c225c6df19"
+version = "0.6.8"
+
+[[deps.ImageMagick]]
+deps = ["FileIO", "ImageCore", "ImageMagick_jll", "InteractiveUtils"]
+git-tree-sha1 = "c5c5478ae8d944c63d6de961b19e6d3324812c35"
+uuid = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
+version = "1.4.0"
+
+[[deps.ImageMagick_jll]]
+deps = ["Artifacts", "Ghostscript_jll", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "OpenJpeg_jll", "Zlib_jll", "libpng_jll"]
+git-tree-sha1 = "fa01c98985be12e5d75301c4527fff2c46fa3e0e"
+uuid = "c73af94c-d91f-53ed-93a7-00f77d67a9d7"
+version = "7.1.1+1"
+
+[[deps.ImageMetadata]]
+deps = ["AxisArrays", "ImageAxes", "ImageBase", "ImageCore"]
+git-tree-sha1 = "355e2b974f2e3212a75dfb60519de21361ad3cb7"
+uuid = "bc367c6b-8a6b-528e-b4bd-a4b897500b49"
+version = "0.9.9"
+
+[[deps.ImageMorphology]]
+deps = ["DataStructures", "ImageCore", "LinearAlgebra", "LoopVectorization", "OffsetArrays", "Requires", "TiledIteration"]
+git-tree-sha1 = "6f0a801136cb9c229aebea0df296cdcd471dbcd1"
+uuid = "787d08f9-d448-5407-9aad-5290dd7ab264"
+version = "0.4.5"
+
+[[deps.ImageQualityIndexes]]
+deps = ["ImageContrastAdjustment", "ImageCore", "ImageDistances", "ImageFiltering", "LazyModules", "OffsetArrays", "PrecompileTools", "Statistics"]
+git-tree-sha1 = "783b70725ed326340adf225be4889906c96b8fd1"
+uuid = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
+version = "0.3.7"
+
+[[deps.ImageSegmentation]]
+deps = ["Clustering", "DataStructures", "Distances", "Graphs", "ImageCore", "ImageFiltering", "ImageMorphology", "LinearAlgebra", "MetaGraphs", "RegionTrees", "SimpleWeightedGraphs", "StaticArrays", "Statistics"]
+git-tree-sha1 = "3ff0ca203501c3eedde3c6fa7fd76b703c336b5f"
+uuid = "80713f31-8817-5129-9cf8-209ff8fb23e1"
+version = "1.8.2"
+
+[[deps.ImageShow]]
+deps = ["Base64", "ColorSchemes", "FileIO", "ImageBase", "ImageCore", "OffsetArrays", "StackViews"]
+git-tree-sha1 = "3b5344bcdbdc11ad58f3b1956709b5b9345355de"
+uuid = "4e3cecfd-b093-5904-9786-8bbb286a6a31"
+version = "0.3.8"
+
+[[deps.ImageTransformations]]
+deps = ["AxisAlgorithms", "CoordinateTransformations", "ImageBase", "ImageCore", "Interpolations", "OffsetArrays", "Rotations", "StaticArrays"]
+git-tree-sha1 = "e0884bdf01bbbb111aea77c348368a86fb4b5ab6"
+uuid = "02fcd773-0e25-5acc-982a-7f6622650795"
+version = "0.10.1"
+
+[[deps.Images]]
+deps = ["Base64", "FileIO", "Graphics", "ImageAxes", "ImageBase", "ImageBinarization", "ImageContrastAdjustment", "ImageCore", "ImageCorners", "ImageDistances", "ImageFiltering", "ImageIO", "ImageMagick", "ImageMetadata", "ImageMorphology", "ImageQualityIndexes", "ImageSegmentation", "ImageShow", "ImageTransformations", "IndirectArrays", "IntegralArrays", "Random", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "StatsBase", "TiledIteration"]
+git-tree-sha1 = "12fdd617c7fe25dc4a6cc804d657cc4b2230302b"
+uuid = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+version = "0.26.1"
+
+[[deps.Imath_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "0936ba688c6d201805a83da835b55c61a180db52"
+uuid = "905a6f67-0a94-5f89-b386-d35d92009cd1"
+version = "3.1.11+0"
+
+[[deps.IndirectArrays]]
+git-tree-sha1 = "012e604e1c7458645cb8b436f8fba789a51b257f"
+uuid = "9b13fd28-a010-5f03-acff-a1bbcff69959"
+version = "1.0.0"
+
+[[deps.Inflate]]
+git-tree-sha1 = "d1b1b796e47d94588b3757fe84fbf65a5ec4a80d"
+uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9"
+version = "0.1.5"
+
+[[deps.InitialValues]]
+git-tree-sha1 = "4da0f88e9a39111c2fa3add390ab15f3a44f3ca3"
+uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
+version = "0.3.1"
+
+[[deps.IntegralArrays]]
+deps = ["ColorTypes", "FixedPointNumbers", "IntervalSets"]
+git-tree-sha1 = "be8e690c3973443bec584db3346ddc904d4884eb"
+uuid = "1d092043-8f09-5a30-832f-7509e371ab51"
+version = "0.1.5"
+
+[[deps.IntelOpenMP_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"]
+git-tree-sha1 = "10bd689145d2c3b2a9844005d01087cc1194e79e"
+uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0"
+version = "2024.2.1+0"
+
+[[deps.InteractiveUtils]]
+deps = ["Markdown"]
+uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+
+[[deps.Interpolations]]
+deps = ["Adapt", "AxisAlgorithms", "ChainRulesCore", "LinearAlgebra", "OffsetArrays", "Random", "Ratios", "Requires", "SharedArrays", "SparseArrays", "StaticArrays", "WoodburyMatrices"]
+git-tree-sha1 = "88a101217d7cb38a7b481ccd50d21876e1d1b0e0"
+uuid = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+version = "0.15.1"
+
+[[deps.IntervalSets]]
+deps = ["Random", "RecipesBase", "Statistics"]
+git-tree-sha1 = "dba9ddf07f77f60450fe5d2e2beb9854d9a49bd0"
+uuid = "8197267c-284f-5f27-9208-e0e47529a953"
+version = "0.7.10"
+
+[[deps.InverseFunctions]]
+deps = ["Dates", "Test"]
+git-tree-sha1 = "a779299d77cd080bf77b97535acecd73e1c5e5cb"
+uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
+version = "0.1.17"
+
+[[deps.IrrationalConstants]]
+git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2"
+uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
+version = "0.2.2"
+
+[[deps.IterTools]]
+git-tree-sha1 = "42d5f897009e7ff2cf88db414a389e5ed1bdd023"
+uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+version = "1.10.0"
+
+[[deps.IteratorInterfaceExtensions]]
+git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
+uuid = "82899510-4779-5014-852e-03e436cf321d"
+version = "1.0.0"
+
+[[deps.JLD2]]
+deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"]
+git-tree-sha1 = "a0746c21bdc986d0dc293efa6b1faee112c37c28"
+uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+version = "0.4.53"
+
+[[deps.JLLWrappers]]
+deps = ["Artifacts", "Preferences"]
+git-tree-sha1 = "be3dc50a92e5a386872a493a10050136d4703f9b"
+uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
+version = "1.6.1"
+
+[[deps.JSExpr]]
+deps = ["JSON", "MacroTools", "Observables", "WebIO"]
+git-tree-sha1 = "b413a73785b98474d8af24fd4c8a975e31df3658"
+uuid = "97c1335a-c9c5-57fe-bc5d-ec35cebe8660"
+version = "0.5.4"
+
+[[deps.JSON]]
+deps = ["Dates", "Mmap", "Parsers", "Unicode"]
+git-tree-sha1 = "31e996f0a15c7b280ba9f76636b3ff9e2ae58c9a"
+uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+version = "0.21.4"
+
+[[deps.JpegTurbo]]
+deps = ["CEnum", "FileIO", "ImageCore", "JpegTurbo_jll", "TOML"]
+git-tree-sha1 = "fa6d0bcff8583bac20f1ffa708c3913ca605c611"
+uuid = "b835a17e-a41a-41e7-81f0-2f016b05efe0"
+version = "0.1.5"
+
+[[deps.JpegTurbo_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "25ee0be4d43d0269027024d75a24c24d6c6e590c"
+uuid = "aacddb02-875f-59d6-b918-886e6ef4fbf8"
+version = "3.0.4+0"
+
+[[deps.JuliaInterpreter]]
+deps = ["CodeTracking", "InteractiveUtils", "Random", "UUIDs"]
+git-tree-sha1 = "2984284a8abcfcc4784d95a9e2ea4e352dd8ede7"
+uuid = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
+version = "0.9.36"
+
+[[deps.JuliaVariables]]
+deps = ["MLStyle", "NameResolution"]
+git-tree-sha1 = "49fb3cb53362ddadb4415e9b73926d6b40709e70"
+uuid = "b14d175d-62b4-44ba-8fb7-3064adc8c3ec"
+version = "0.2.4"
+
+[[deps.Kaleido_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "43032da5832754f58d14a91ffbe86d5f176acda9"
+uuid = "f7e6163d-2fa5-5f23-b69c-1db539e41963"
+version = "0.2.1+0"
+
+[[deps.KernelAbstractions]]
+deps = ["Adapt", "Atomix", "EnzymeCore", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
+git-tree-sha1 = "e73a077abc7fe798fe940deabe30ef6c66bdde52"
+uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+version = "0.9.29"
+
+[[deps.LERC_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "bf36f528eec6634efc60d7ec062008f171071434"
+uuid = "88015f11-f218-50d7-93a8-a6af411a945d"
+version = "3.0.0+1"
+
+[[deps.LLVM]]
+deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"]
+git-tree-sha1 = "839c82932db86740ae729779e610f07a1640be9a"
+uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
+version = "6.6.3"
+
+[[deps.LLVMExtra_jll]]
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "88b916503aac4fb7f701bb625cd84ca5dd1677bc"
+uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
+version = "0.0.29+0"
+
+[[deps.LLVMOpenMP_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "78211fb6cbc872f77cad3fc0b6cf647d923f4929"
+uuid = "1d63c593-3942-5779-bab2-d838dc0a180e"
+version = "18.1.7+0"
+
+[[deps.LRUCache]]
+deps = ["Serialization"]
+git-tree-sha1 = "b3cc6698599b10e652832c2f23db3cab99d51b59"
+uuid = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
+version = "1.6.1"
+
+[[deps.LZO_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "854a9c268c43b77b0a27f22d7fab8d33cdb3a731"
+uuid = "dd4b983a-f0e5-5f8d-a1b7-129d4a5fb1ac"
+version = "2.10.2+1"
+
+[[deps.LaTeXStrings]]
+git-tree-sha1 = "dda21b8cbd6a6c40d9d02a73230f9d70fed6918c"
+uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+version = "1.4.0"
+
+[[deps.LayoutPointers]]
+deps = ["ArrayInterface", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface"]
+git-tree-sha1 = "a9eaadb366f5493a5654e843864c13d8b107548c"
+uuid = "10f19ff3-798f-405d-979b-55457f8fc047"
+version = "0.1.17"
+
+[[deps.Lazy]]
+deps = ["MacroTools"]
+git-tree-sha1 = "1370f8202dac30758f3c345f9909b97f53d87d3f"
+uuid = "50d2b5c4-7a5e-59d5-8109-a42b560f39c0"
+version = "0.15.1"
+
+[[deps.LazyArtifacts]]
+deps = ["Artifacts", "Pkg"]
+uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+
+[[deps.LazyModules]]
+git-tree-sha1 = "a560dd966b386ac9ae60bdd3a3d3a326062d3c3e"
+uuid = "8cdb02fc-e678-4876-92c5-9defec4f444e"
+version = "0.3.1"
+
+[[deps.LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+version = "0.6.3"
+
+[[deps.LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
+version = "7.84.0+0"
+
+[[deps.LibGit2]]
+deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
+uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
+
+[[deps.LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+version = "1.10.2+0"
+
+[[deps.Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+
+[[deps.Libffi_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "0b4a5d71f3e5200a7dff793393e09dfc2d874290"
+uuid = "e9f186c6-92d2-5b65-8a66-fee21dc1b490"
+version = "3.2.2+1"
+
+[[deps.Libgcrypt_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libgpg_error_jll"]
+git-tree-sha1 = "8be878062e0ffa2c3f67bb58a595375eda5de80b"
+uuid = "d4300ac3-e22c-5743-9152-c294e39db1e4"
+version = "1.11.0+0"
+
+[[deps.Libgpg_error_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "c6ce1e19f3aec9b59186bdf06cdf3c4fc5f5f3e6"
+uuid = "7add5ba3-2f88-524e-9cd5-f83b8a55f7b8"
+version = "1.50.0+0"
+
+[[deps.Libiconv_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "61dfdba58e585066d8bce214c5a51eaa0539f269"
+uuid = "94ce4f54-9a6c-5748-9c1c-f9c7231a4531"
+version = "1.17.0+1"
+
+[[deps.Libmount_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "0c4f9c4f1a50d8f35048fa0532dabbadf702f81e"
+uuid = "4b2f31a3-9ecc-558c-b454-b3730dcb73e9"
+version = "2.40.1+0"
+
+[[deps.Libtiff_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "LERC_jll", "Libdl", "Pkg", "Zlib_jll", "Zstd_jll"]
+git-tree-sha1 = "3eb79b0ca5764d4799c06699573fd8f533259713"
+uuid = "89763e89-9b03-5906-acba-b20f662cd828"
+version = "4.4.0+0"
+
+[[deps.Libuuid_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "5ee6203157c120d79034c748a2acba45b82b8807"
+uuid = "38a345b3-de98-5d2b-a5d3-14cd9215e700"
+version = "2.40.1+0"
+
+[[deps.LinearAlgebra]]
+deps = ["Libdl", "libblastrampoline_jll"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[[deps.LittleCMS_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "Pkg"]
+git-tree-sha1 = "110897e7db2d6836be22c18bffd9422218ee6284"
+uuid = "d3a379c0-f9a3-5b72-a4c0-6bf4d2e8af0f"
+version = "2.12.0+0"
+
+[[deps.LogExpFunctions]]
+deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
+git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea"
+uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+version = "0.3.28"
+
+[[deps.Logging]]
+uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+
+[[deps.LoggingExtras]]
+deps = ["Dates", "Logging"]
+git-tree-sha1 = "f02b56007b064fbfddb4c9cd60161b6dd0f40df3"
+uuid = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+version = "1.1.0"
+
+[[deps.LoopVectorization]]
+deps = ["ArrayInterface", "CPUSummary", "ChainRulesCore", "CloseOpenIntervals", "DocStringExtensions", "ForwardDiff", "HostCPUFeatures", "IfElse", "LayoutPointers", "LinearAlgebra", "OffsetArrays", "PolyesterWeave", "PrecompileTools", "SIMDTypes", "SLEEFPirates", "SpecialFunctions", "Static", "StaticArrayInterface", "ThreadingUtilities", "UnPack", "VectorizationBase"]
+git-tree-sha1 = "8084c25a250e00ae427a379a5b607e7aed96a2dd"
+uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
+version = "0.12.171"
+
+[[deps.LoweredCodeUtils]]
+deps = ["JuliaInterpreter"]
+git-tree-sha1 = "260dc274c1bc2cb839e758588c63d9c8b5e639d1"
+uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
+version = "3.0.5"
+
+[[deps.MAT]]
+deps = ["BufferedStreams", "CodecZlib", "HDF5", "SparseArrays"]
+git-tree-sha1 = "1d2dd9b186742b0f317f2530ddcbf00eebb18e96"
+uuid = "23992714-dd62-5051-b70f-ba57cb901cac"
+version = "0.10.7"
+
+[[deps.MKL_jll]]
+deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"]
+git-tree-sha1 = "f046ccd0c6db2832a9f639e2c669c6fe867e5f4f"
+uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
+version = "2024.2.0+0"
+
+[[deps.MLStyle]]
+git-tree-sha1 = "bc38dff0548128765760c79eb7388a4b37fae2c8"
+uuid = "d8e11817-5142-5d16-987a-aa16d5891078"
+version = "0.4.17"
+
+[[deps.MLUtils]]
+deps = ["ChainRulesCore", "Compat", "DataAPI", "DelimitedFiles", "FLoops", "NNlib", "Random", "ShowCases", "SimpleTraits", "Statistics", "StatsBase", "Tables", "Transducers"]
+git-tree-sha1 = "b45738c2e3d0d402dffa32b2c1654759a2ac35a4"
+uuid = "f1d291b0-491e-4a28-83b9-f70985020b54"
+version = "0.4.4"
+
+[[deps.MPI]]
+deps = ["Distributed", "DocStringExtensions", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "PkgVersion", "PrecompileTools", "Requires", "Serialization", "Sockets"]
+git-tree-sha1 = "892676019c58f34e38743bc989b0eca5bce5edc5"
+uuid = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+version = "0.20.22"
+
+[[deps.MPICH_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"]
+git-tree-sha1 = "7715e65c47ba3941c502bffb7f266a41a7f54423"
+uuid = "7cb0a576-ebde-5e09-9194-50597f1243b4"
+version = "4.2.3+0"
+
+[[deps.MPIPreferences]]
+deps = ["Libdl", "Preferences"]
+git-tree-sha1 = "c105fe467859e7f6e9a852cb15cb4301126fac07"
+uuid = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
+version = "0.1.11"
+
+[[deps.MPItrampoline_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"]
+git-tree-sha1 = "70e830dab5d0775183c99fc75e4c24c614ed7142"
+uuid = "f1f71cc9-e9ae-5b93-9b94-4fe0e1ad3748"
+version = "5.5.1+0"
+
+[[deps.MacroTools]]
+deps = ["Markdown", "Random"]
+git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df"
+uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.5.13"
+
+[[deps.ManualMemory]]
+git-tree-sha1 = "bcaef4fc7a0cfe2cba636d84cda54b5e4e4ca3cd"
+uuid = "d125e4d3-2237-4719-b19c-fa641b8a4667"
+version = "0.1.8"
+
+[[deps.MappedArrays]]
+git-tree-sha1 = "2dab0221fe2b0f2cb6754eaa743cc266339f527e"
+uuid = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
+version = "0.4.2"
+
+[[deps.Markdown]]
+deps = ["Base64"]
+uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
+
+[[deps.Match]]
+deps = ["MacroTools", "OrderedCollections"]
+git-tree-sha1 = "5ac5e5267e17ccbd717bc7caaa57c5a20b80261d"
+uuid = "7eb4fadd-790c-5f42-8a69-bfa0b872bfbf"
+version = "2.1.0"
+
+[[deps.MbedTLS]]
+deps = ["Dates", "MbedTLS_jll", "MozillaCACerts_jll", "NetworkOptions", "Random", "Sockets"]
+git-tree-sha1 = "c067a280ddc25f196b5e7df3877c6b226d390aaf"
+uuid = "739be429-bea8-5141-9913-cc70e7f3736d"
+version = "1.1.9"
+
+[[deps.MbedTLS_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
+version = "2.28.0+0"
+
+[[deps.MetaGraphs]]
+deps = ["Graphs", "JLD2", "Random"]
+git-tree-sha1 = "1130dbe1d5276cb656f6e1094ce97466ed700e5a"
+uuid = "626554b9-1ddb-594c-aa3c-2596fe9399a5"
+version = "0.7.2"
+
+[[deps.MicroCollections]]
+deps = ["Accessors", "BangBang", "InitialValues"]
+git-tree-sha1 = "44d32db644e84c75dab479f1bc15ee76a1a3618f"
+uuid = "128add7d-3638-4c79-886c-908ea0c25c34"
+version = "0.2.0"
+
+[[deps.MicrosoftMPI_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "f12a29c4400ba812841c6ace3f4efbb6dbb3ba01"
+uuid = "9237b28f-5490-5468-be7b-bb81f5f5e6cf"
+version = "10.1.4+2"
+
+[[deps.Missings]]
+deps = ["DataAPI"]
+git-tree-sha1 = "ec4f7fbeab05d7747bdf98eb74d130a2a2ed298d"
+uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
+version = "1.2.0"
+
+[[deps.Mmap]]
+uuid = "a63ad114-7e13-5084-954f-fe012c677804"
+
+[[deps.MosaicViews]]
+deps = ["MappedArrays", "OffsetArrays", "PaddedViews", "StackViews"]
+git-tree-sha1 = "7b86a5d4d70a9f5cdf2dacb3cbe6d251d1a61dbe"
+uuid = "e94cdb99-869f-56ef-bcf0-1ae2bcbe0389"
+version = "0.3.4"
+
+[[deps.MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+version = "2022.2.1"
+
+[[deps.Mustache]]
+deps = ["Printf", "Tables"]
+git-tree-sha1 = "3b2db451a872b20519ebb0cec759d3d81a1c6bcb"
+uuid = "ffc61752-8dc7-55ee-8c37-f3e9cdd09e70"
+version = "1.0.20"
+
+[[deps.Mux]]
+deps = ["AssetRegistry", "Base64", "HTTP", "Hiccup", "MbedTLS", "Pkg", "Sockets"]
+git-tree-sha1 = "7295d849103ac4fcbe3b2e439f229c5cc77b9b69"
+uuid = "a975b10e-0019-58db-a62f-e48ff68538c9"
+version = "1.0.2"
+
+[[deps.NNlib]]
+deps = ["Adapt", "Atomix", "ChainRulesCore", "GPUArraysCore", "KernelAbstractions", "LinearAlgebra", "Pkg", "Random", "Requires", "Statistics"]
+git-tree-sha1 = "72240e3f5ca031937bd536182cb2c031da5f46dd"
+uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+version = "0.8.21"
+
+[[deps.NNlibCUDA]]
+deps = ["Adapt", "CUDA", "LinearAlgebra", "NNlib", "Random", "Statistics", "cuDNN"]
+git-tree-sha1 = "f94a9684394ff0d325cc12b06da7032d8be01aaf"
+uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
+version = "0.2.7"
+
+[[deps.NaNMath]]
+deps = ["OpenLibm_jll"]
+git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4"
+uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+version = "1.0.2"
+
+[[deps.NameResolution]]
+deps = ["PrettyPrint"]
+git-tree-sha1 = "1a0fa0e9613f46c9b8c11eee38ebb4f590013c5e"
+uuid = "71a1bf82-56d0-4bbc-8a3c-48b961074391"
+version = "0.1.5"
+
+[[deps.NearestNeighbors]]
+deps = ["Distances", "StaticArrays"]
+git-tree-sha1 = "3cebfc94a0754cc329ebc3bab1e6c89621e791ad"
+uuid = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
+version = "0.4.20"
+
+[[deps.Netpbm]]
+deps = ["FileIO", "ImageCore", "ImageMetadata"]
+git-tree-sha1 = "d92b107dbb887293622df7697a2223f9f8176fcd"
+uuid = "f09324ee-3d7c-5217-9330-fc30815ba969"
+version = "1.1.1"
+
+[[deps.NetworkOptions]]
+uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
+version = "1.2.0"
+
+[[deps.OMEinsum]]
+deps = ["AbstractTrees", "BatchedRoutines", "ChainRulesCore", "Combinatorics", "LinearAlgebra", "MacroTools", "OMEinsumContractionOrders", "Test", "TupleTools"]
+git-tree-sha1 = "3198068d99e1e28f358d9c77210d581fe13c572d"
+uuid = "ebe7aa44-baf0-506c-a96f-8464559b3922"
+version = "0.8.4"
+
+[[deps.OMEinsumContractionOrders]]
+deps = ["AbstractTrees", "JSON", "SparseArrays", "StatsBase", "Suppressor", "TreeWidthSolver"]
+git-tree-sha1 = "4b9c4220a66cc583ab8048660f5802406ecedb9c"
+uuid = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
+version = "0.9.3"
+
+[[deps.Observables]]
+git-tree-sha1 = "7438a59546cf62428fc9d1bc94729146d37a7225"
+uuid = "510215fc-4207-5dde-b226-833fc4488ee2"
+version = "0.5.5"
+
+[[deps.OffsetArrays]]
+deps = ["Adapt"]
+git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e"
+uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+version = "1.14.1"
+
+[[deps.OneHotArrays]]
+deps = ["Adapt", "ChainRulesCore", "Compat", "GPUArraysCore", "LinearAlgebra", "NNlib"]
+git-tree-sha1 = "963a3f28a2e65bb87a68033ea4a616002406037d"
+uuid = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
+version = "0.2.5"
+
+[[deps.OpenBLAS_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
+uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
+version = "0.3.20+0"
+
+[[deps.OpenEXR]]
+deps = ["Colors", "FileIO", "OpenEXR_jll"]
+git-tree-sha1 = "327f53360fdb54df7ecd01e96ef1983536d1e633"
+uuid = "52e1d378-f018-4a11-a4be-720524705ac7"
+version = "0.3.2"
+
+[[deps.OpenEXR_jll]]
+deps = ["Artifacts", "Imath_jll", "JLLWrappers", "Libdl", "Zlib_jll"]
+git-tree-sha1 = "8292dd5c8a38257111ada2174000a33745b06d4e"
+uuid = "18a262bb-aa17-5467-a713-aee519bc75cb"
+version = "3.2.4+0"
+
+[[deps.OpenJpeg_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "Pkg", "libpng_jll"]
+git-tree-sha1 = "76374b6e7f632c130e78100b166e5a48464256f8"
+uuid = "643b3616-a352-519d-856d-80112ee9badc"
+version = "2.4.0+0"
+
+[[deps.OpenLibm_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
+version = "0.8.1+0"
+
+[[deps.OpenMPI_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"]
+git-tree-sha1 = "bfce6d523861a6c562721b262c0d1aaeead2647f"
+uuid = "fe0851c0-eecd-5654-98d4-656369965a5c"
+version = "5.0.5+0"
+
+[[deps.OpenSSL]]
+deps = ["BitFlags", "Dates", "MozillaCACerts_jll", "OpenSSL_jll", "Sockets"]
+git-tree-sha1 = "38cb508d080d21dc1128f7fb04f20387ed4c0af4"
+uuid = "4d8831e6-92b7-49fb-bdf8-b643e874388c"
+version = "1.4.3"
+
+[[deps.OpenSSL_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "7493f61f55a6cce7325f197443aa80d32554ba10"
+uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95"
+version = "3.0.15+1"
+
+[[deps.OpenSpecFun_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
+uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
+version = "0.5.5+0"
+
+[[deps.Optimisers]]
+deps = ["ChainRulesCore", "Functors", "LinearAlgebra", "Random", "Statistics"]
+git-tree-sha1 = "c1fc26bab5df929a5172f296f25d7d08688fd25b"
+uuid = "3bd65402-5787-11e9-1adc-39752487f4e2"
+version = "0.2.20"
+
+[[deps.OrderedCollections]]
+git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5"
+uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+version = "1.6.3"
+
+[[deps.PCRE2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "efcefdf7-47ab-520b-bdef-62a2eaa19f15"
+version = "10.40.0+0"
+
+[[deps.PDMats]]
+deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
+git-tree-sha1 = "949347156c25054de2db3b166c52ac4728cbad65"
+uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
+version = "0.11.31"
+
+[[deps.PNGFiles]]
+deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"]
+git-tree-sha1 = "67186a2bc9a90f9f85ff3cc8277868961fb57cbd"
+uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883"
+version = "0.4.3"
+
+[[deps.PaddedViews]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "0fac6313486baae819364c52b4f483450a9d793f"
+uuid = "5432bcbf-9aad-5242-b902-cca2824c8663"
+version = "0.5.12"
+
+[[deps.Pango_jll]]
+deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "FriBidi_jll", "Glib_jll", "HarfBuzz_jll", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "e127b609fb9ecba6f201ba7ab753d5a605d53801"
+uuid = "36c8627f-9965-5494-a995-c6b170f724f3"
+version = "1.54.1+0"
+
+[[deps.Parameters]]
+deps = ["OrderedCollections", "UnPack"]
+git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
+uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+version = "0.12.3"
+
+[[deps.ParametricDFNOs]]
+deps = ["CUDA", "DrWatson", "FileIO", "Flux", "HDF5", "IJulia", "JLD2", "LaTeXStrings", "MAT", "MPI", "Parameters", "ParametricOperators", "PlotlyJS", "Printf", "Profile", "ProgressMeter", "PyCall", "PyPlot", "Random", "Revise", "Test", "Zygote"]
+path = "/pscratch/sd/r/richardr/FNO_CIG/ParametricDFNOs.jl/"
+uuid = "3b406863-5cae-494d-8b19-779e643ae09e"
+version = "2.0.2"
+
+[[deps.ParametricOperators]]
+deps = ["CUDA", "ChainRulesCore", "Combinatorics", "DataStructures", "FFTW", "Flux", "GraphViz", "JSON", "LRUCache", "LaTeXStrings", "MPI", "Match", "OMEinsum", "Random", "UUIDs"]
+path = "../../FNO_CIG/ParametricOperators.jl"
+uuid = "db9e0614-c73c-4112-a40c-114e5b366d0d"
+version = "1.1.3"
+
+[[deps.Parsers]]
+deps = ["Dates", "PrecompileTools", "UUIDs"]
+git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821"
+uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
+version = "2.8.1"
+
+[[deps.Pidfile]]
+deps = ["FileWatching", "Test"]
+git-tree-sha1 = "2d8aaf8ee10df53d0dfb9b8ee44ae7c04ced2b03"
+uuid = "fa939f87-e72e-5be4-a000-7fc836dbe307"
+version = "1.3.0"
+
+[[deps.Pixman_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LLVMOpenMP_jll", "Libdl"]
+git-tree-sha1 = "35621f10a7531bc8fa58f74610b1bfb70a3cfc6b"
+uuid = "30392449-352a-5448-841d-b1acce4e97dc"
+version = "0.43.4+0"
+
+[[deps.Pkg]]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+version = "1.8.0"
+
+[[deps.PkgVersion]]
+deps = ["Pkg"]
+git-tree-sha1 = "f9501cc0430a26bc3d156ae1b5b0c1b47af4d6da"
+uuid = "eebad327-c553-4316-9ea0-9fa01ccd7688"
+version = "0.3.3"
+
+[[deps.PlotlyBase]]
+deps = ["ColorSchemes", "Dates", "DelimitedFiles", "DocStringExtensions", "JSON", "LaTeXStrings", "Logging", "Parameters", "Pkg", "REPL", "Requires", "Statistics", "UUIDs"]
+git-tree-sha1 = "56baf69781fc5e61607c3e46227ab17f7040ffa2"
+uuid = "a03496cd-edff-5a9b-9e67-9cda94a718b5"
+version = "0.8.19"
+
+[[deps.PlotlyJS]]
+deps = ["Base64", "Blink", "DelimitedFiles", "JSExpr", "JSON", "Kaleido_jll", "Markdown", "Pkg", "PlotlyBase", "PlotlyKaleido", "REPL", "Reexport", "Requires", "WebIO"]
+git-tree-sha1 = "f198c8a80c08987a2915156e6e6131e5d73b97f4"
+uuid = "f0f68f2c-4968-5e81-91da-67840de0976a"
+version = "0.18.14"
+
+[[deps.PlotlyKaleido]]
+deps = ["Base64", "JSON", "Kaleido_jll"]
+git-tree-sha1 = "3210de4d88af7ca5de9e26305758a59aabc48aac"
+uuid = "f2990250-8cf9-495f-b13a-cce12b45703c"
+version = "2.2.5"
+
+[[deps.PolyesterWeave]]
+deps = ["BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "Static", "ThreadingUtilities"]
+git-tree-sha1 = "645bed98cd47f72f67316fd42fc47dee771aefcd"
+uuid = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad"
+version = "0.2.2"
+
+[[deps.Polynomials]]
+deps = ["LinearAlgebra", "RecipesBase", "Requires", "Setfield", "SparseArrays"]
+git-tree-sha1 = "1a9cfb2dc2c2f1bd63f1906d72af39a79b49b736"
+uuid = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
+version = "4.0.11"
+
+[[deps.PrecompileTools]]
+deps = ["Preferences"]
+git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f"
+uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+version = "1.2.1"
+
+[[deps.Preferences]]
+deps = ["TOML"]
+git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6"
+uuid = "21216c6a-2e73-6563-6e65-726566657250"
+version = "1.4.3"
+
+[[deps.PrettyPrint]]
+git-tree-sha1 = "632eb4abab3449ab30c5e1afaa874f0b98b586e4"
+uuid = "8162dcfd-2161-5ef2-ae6c-7681170c5f98"
+version = "0.2.0"
+
+[[deps.Printf]]
+deps = ["Unicode"]
+uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[[deps.Profile]]
+deps = ["Printf"]
+uuid = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
+
+[[deps.ProgressLogging]]
+deps = ["Logging", "SHA", "UUIDs"]
+git-tree-sha1 = "80d919dee55b9c50e8d9e2da5eeafff3fe58b539"
+uuid = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
+version = "0.1.4"
+
+[[deps.ProgressMeter]]
+deps = ["Distributed", "Printf"]
+git-tree-sha1 = "8f6bc219586aef8baf0ff9a5fe16ee9c70cb65e4"
+uuid = "92933f4c-e287-5a05-a399-4b506db050ca"
+version = "1.10.2"
+
+[[deps.PyCall]]
+deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
+git-tree-sha1 = "9816a3826b0ebf49ab4926e2b18842ad8b5c8f04"
+uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+version = "1.96.4"
+
+[[deps.PyPlot]]
+deps = ["Colors", "LaTeXStrings", "PyCall", "Sockets", "Test", "VersionParsing"]
+git-tree-sha1 = "0371ca706e3f295481cbf94c8c36692b072285c2"
+uuid = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+version = "2.11.5"
+
+[[deps.QOI]]
+deps = ["ColorTypes", "FileIO", "FixedPointNumbers"]
+git-tree-sha1 = "18e8f4d1426e965c7b532ddd260599e1510d26ce"
+uuid = "4b34888f-f399-49d4-9bb3-47ed5cae4e65"
+version = "1.0.0"
+
+[[deps.Quaternions]]
+deps = ["LinearAlgebra", "Random", "RealDot"]
+git-tree-sha1 = "994cc27cdacca10e68feb291673ec3a76aa2fae9"
+uuid = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
+version = "0.7.6"
+
+[[deps.REPL]]
+deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
+uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+
+[[deps.Random]]
+deps = ["SHA", "Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[deps.Random123]]
+deps = ["Random", "RandomNumbers"]
+git-tree-sha1 = "4743b43e5a9c4a2ede372de7061eed81795b12e7"
+uuid = "74087812-796a-5b5d-8853-05524746bad3"
+version = "1.7.0"
+
+[[deps.RandomNumbers]]
+deps = ["Random"]
+git-tree-sha1 = "c6ec94d2aaba1ab2ff983052cf6a606ca5985902"
+uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143"
+version = "1.6.0"
+
+[[deps.RangeArrays]]
+git-tree-sha1 = "b9039e93773ddcfc828f12aadf7115b4b4d225f5"
+uuid = "b3c3ace0-ae52-54e7-9d0b-2c1406fd6b9d"
+version = "0.3.2"
+
+[[deps.Ratios]]
+deps = ["Requires"]
+git-tree-sha1 = "1342a47bf3260ee108163042310d26f2be5ec90b"
+uuid = "c84ed2f1-dad5-54f0-aa8e-dbefe2724439"
+version = "0.4.5"
+
+[[deps.RealDot]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "9f0a1b71baaf7650f4fa8a1d168c7fb6ee41f0c9"
+uuid = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9"
+version = "0.1.0"
+
+[[deps.RecipesBase]]
+deps = ["PrecompileTools"]
+git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff"
+uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+version = "1.3.4"
+
+[[deps.Reexport]]
+git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
+uuid = "189a3867-3050-52da-a836-e630ba90ab69"
+version = "1.2.2"
+
+[[deps.RegionTrees]]
+deps = ["IterTools", "LinearAlgebra", "StaticArrays"]
+git-tree-sha1 = "4618ed0da7a251c7f92e869ae1a19c74a7d2a7f9"
+uuid = "dee08c22-ab7f-5625-9660-a9af2021b33f"
+version = "0.3.2"
+
+[[deps.Requires]]
+deps = ["UUIDs"]
+git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
+uuid = "ae029012-a4dd-5104-9daa-d747884805df"
+version = "1.3.0"
+
+[[deps.Revise]]
+deps = ["CodeTracking", "Distributed", "FileWatching", "JuliaInterpreter", "LibGit2", "LoweredCodeUtils", "OrderedCollections", "REPL", "Requires", "UUIDs", "Unicode"]
+git-tree-sha1 = "7f4228017b83c66bd6aa4fddeb170ce487e53bc7"
+uuid = "295af30f-e4ad-537b-8983-00126c2a3abe"
+version = "3.6.2"
+
+[[deps.Rotations]]
+deps = ["LinearAlgebra", "Quaternions", "Random", "RecipesBase", "StaticArrays"]
+git-tree-sha1 = "5680a9276685d392c87407df00d57c9924d9f11e"
+uuid = "6038ab10-8711-5258-84ad-4b1120ba62dc"
+version = "1.7.1"
+
+[[deps.SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+version = "0.7.0"
+
+[[deps.SIMD]]
+deps = ["PrecompileTools"]
+git-tree-sha1 = "98ca7c29edd6fc79cd74c61accb7010a4e7aee33"
+uuid = "fdea26ae-647d-5447-a871-4b548cad5224"
+version = "3.6.0"
+
+[[deps.SIMDTypes]]
+git-tree-sha1 = "330289636fb8107c5f32088d2741e9fd7a061a5c"
+uuid = "94e857df-77ce-4151-89e5-788b33177be4"
+version = "0.1.0"
+
+[[deps.SLEEFPirates]]
+deps = ["IfElse", "Static", "VectorizationBase"]
+git-tree-sha1 = "456f610ca2fbd1c14f5fcf31c6bfadc55e7d66e0"
+uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa"
+version = "0.6.43"
+
+[[deps.Scratch]]
+deps = ["Dates"]
+git-tree-sha1 = "3bac05bc7e74a75fd9cba4295cde4045d9fe2386"
+uuid = "6c6a2e73-6563-6170-7368-637461726353"
+version = "1.2.1"
+
+[[deps.Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[deps.Setfield]]
+deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"]
+git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac"
+uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+version = "1.1.1"
+
+[[deps.SharedArrays]]
+deps = ["Distributed", "Mmap", "Random", "Serialization"]
+uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
+
+[[deps.ShowCases]]
+git-tree-sha1 = "7f534ad62ab2bd48591bdeac81994ea8c445e4a5"
+uuid = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
+version = "0.1.0"
+
+[[deps.SimpleBufferStream]]
+git-tree-sha1 = "f305871d2f381d21527c770d4788c06c097c9bc1"
+uuid = "777ac1f9-54b0-4bf8-805c-2214025038e7"
+version = "1.2.0"
+
+[[deps.SimpleTraits]]
+deps = ["InteractiveUtils", "MacroTools"]
+git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231"
+uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
+version = "0.9.4"
+
+[[deps.SimpleWeightedGraphs]]
+deps = ["Graphs", "LinearAlgebra", "Markdown", "SparseArrays"]
+git-tree-sha1 = "4b33e0e081a825dbfaf314decf58fa47e53d6acb"
+uuid = "47aef6b3-ad0c-573a-a1e2-d07658019622"
+version = "1.4.0"
+
+[[deps.Sixel]]
+deps = ["Dates", "FileIO", "ImageCore", "IndirectArrays", "OffsetArrays", "REPL", "libsixel_jll"]
+git-tree-sha1 = "2da10356e31327c7096832eb9cd86307a50b1eb6"
+uuid = "45858cf5-a6b0-47a3-bbea-62219f50df47"
+version = "0.1.3"
+
+[[deps.Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+
+[[deps.SoftGlobalScope]]
+deps = ["REPL"]
+git-tree-sha1 = "986ec2b6162ccb95de5892ed17832f95badf770c"
+uuid = "b85f4697-e234-5449-a836-ec8e2f98b302"
+version = "1.1.0"
+
+[[deps.SortingAlgorithms]]
+deps = ["DataStructures"]
+git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085"
+uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
+version = "1.2.1"
+
+[[deps.SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[deps.SparseInverseSubset]]
+deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
+git-tree-sha1 = "52962839426b75b3021296f7df242e40ecfc0852"
+uuid = "dc90abb0-5640-4711-901d-7e5b23a2fada"
+version = "0.1.2"
+
+[[deps.SpecialFunctions]]
+deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
+git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14"
+uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
+version = "2.4.0"
+
+[[deps.SplittablesBase]]
+deps = ["Setfield", "Test"]
+git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5"
+uuid = "171d559e-b47b-412a-8079-5efa626c420e"
+version = "0.1.15"
+
+[[deps.StackViews]]
+deps = ["OffsetArrays"]
+git-tree-sha1 = "46e589465204cd0c08b4bd97385e4fa79a0c770c"
+uuid = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
+version = "0.1.1"
+
+[[deps.Static]]
+deps = ["IfElse"]
+git-tree-sha1 = "b366eb1eb68075745777d80861c6706c33f588ae"
+uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
+version = "0.8.9"
+
+[[deps.StaticArrayInterface]]
+deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "PrecompileTools", "Requires", "Static"]
+git-tree-sha1 = "c3668ff1a3e4ddf374fc4f8c25539ce7194dcc39"
+uuid = "0d7ed370-da01-4f52-bd93-41d350b8b718"
+version = "1.6.0"
+
+[[deps.StaticArrays]]
+deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore", "Statistics"]
+git-tree-sha1 = "777657803913ffc7e8cc20f0fd04b634f871af8f"
+uuid = "90137ffa-7385-5640-81b9-e52037218182"
+version = "1.9.8"
+
+[[deps.StaticArraysCore]]
+git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682"
+uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+version = "1.4.3"
+
+[[deps.Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[deps.StatsAPI]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed"
+uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
+version = "1.7.0"
+
+[[deps.StatsBase]]
+deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
+git-tree-sha1 = "5cf7606d6cef84b543b483848d4ae08ad9832b21"
+uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+version = "0.34.3"
+
+[[deps.StructArrays]]
+deps = ["Adapt", "ConstructionBase", "DataAPI", "GPUArraysCore", "SparseArrays", "StaticArrays", "Tables"]
+git-tree-sha1 = "f4dc295e983502292c4c3f951dbb4e985e35b3be"
+uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+version = "0.6.18"
+
+[[deps.SuiteSparse]]
+deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
+uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
+
+[[deps.Suppressor]]
+deps = ["Logging"]
+git-tree-sha1 = "6dbb5b635c5437c68c28c2ac9e39b87138f37c0a"
+uuid = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+version = "0.2.8"
+
+[[deps.TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+version = "1.0.0"
+
+[[deps.TableTraits]]
+deps = ["IteratorInterfaceExtensions"]
+git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
+uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
+version = "1.0.1"
+
+[[deps.Tables]]
+deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"]
+git-tree-sha1 = "598cd7c1f68d1e205689b1c2fe65a9f85846f297"
+uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+version = "1.12.0"
+
+[[deps.Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+version = "1.10.1"
+
+[[deps.TensorCore]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6"
+uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50"
+version = "0.1.1"
+
+[[deps.Test]]
+deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
+uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[[deps.TextWrap]]
+git-tree-sha1 = "43044b737fa70bc12f6105061d3da38f881a3e3c"
+uuid = "b718987f-49a8-5099-9789-dcd902bef87d"
+version = "1.0.2"
+
+[[deps.ThreadingUtilities]]
+deps = ["ManualMemory"]
+git-tree-sha1 = "eda08f7e9818eb53661b3deb74e3159460dfbc27"
+uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5"
+version = "0.5.2"
+
+[[deps.TiffImages]]
+deps = ["ColorTypes", "DataStructures", "DocStringExtensions", "FileIO", "FixedPointNumbers", "IndirectArrays", "Inflate", "Mmap", "OffsetArrays", "PkgVersion", "ProgressMeter", "SIMD", "UUIDs"]
+git-tree-sha1 = "38f139cc4abf345dd4f22286ec000728d5e8e097"
+uuid = "731e570b-9d59-4bfa-96dc-6df516fadf69"
+version = "0.10.2"
+
+[[deps.TiledIteration]]
+deps = ["OffsetArrays", "StaticArrayInterface"]
+git-tree-sha1 = "1176cc31e867217b06928e2f140c90bd1bc88283"
+uuid = "06e1c1a7-607b-532d-9fad-de7d9aa2abac"
+version = "0.5.0"
+
+[[deps.TimerOutputs]]
+deps = ["ExprTools", "Printf"]
+git-tree-sha1 = "3a6f063d690135f5c1ba351412c82bae4d1402bf"
+uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+version = "0.5.25"
+
+[[deps.TranscodingStreams]]
+git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
+uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+version = "0.11.3"
+
+[[deps.Transducers]]
+deps = ["Accessors", "Adapt", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "ConstructionBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "SplittablesBase", "Tables"]
+git-tree-sha1 = "7deeab4ff96b85c5f72c824cae53a1398da3d1cb"
+uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999"
+version = "0.4.84"
+
+[[deps.TreeWidthSolver]]
+deps = ["AbstractTrees", "BitBasis", "Combinatorics", "Graphs", "SparseArrays"]
+git-tree-sha1 = "11626397b7835f5ebe2b79dd12bb3876a420ef9b"
+uuid = "7d267fc5-9ace-409f-a54c-cd2374872a55"
+version = "0.3.3"
+
+[[deps.TupleTools]]
+git-tree-sha1 = "41e43b9dc950775eac654b9f845c839cd2f1821e"
+uuid = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
+version = "1.6.0"
+
+[[deps.URIs]]
+git-tree-sha1 = "67db6cc7b3821e19ebe75791a9dd19c9b1188f2b"
+uuid = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
+version = "1.5.1"
+
+[[deps.UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
+[[deps.UnPack]]
+git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
+uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+version = "1.0.2"
+
+[[deps.Unicode]]
+uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[deps.UnsafeAtomics]]
+git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278"
+uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f"
+version = "0.2.1"
+
+[[deps.UnsafeAtomicsLLVM]]
+deps = ["LLVM", "UnsafeAtomics"]
+git-tree-sha1 = "bf2c553f25e954a9b38c9c0593a59bb13113f9e5"
+uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
+version = "0.1.5"
+
+[[deps.VectorizationBase]]
+deps = ["ArrayInterface", "CPUSummary", "HostCPUFeatures", "IfElse", "LayoutPointers", "Libdl", "LinearAlgebra", "SIMDTypes", "Static", "StaticArrayInterface"]
+git-tree-sha1 = "e7f5b81c65eb858bed630fe006837b935518aca5"
+uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
+version = "0.21.70"
+
+[[deps.VersionParsing]]
+git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.3.0"
+
+[[deps.WebIO]]
+deps = ["AssetRegistry", "Base64", "Distributed", "FunctionalCollections", "JSON", "Logging", "Observables", "Pkg", "Random", "Requires", "Sockets", "UUIDs", "WebSockets", "Widgets"]
+git-tree-sha1 = "0eef0765186f7452e52236fa42ca8c9b3c11c6e3"
+uuid = "0f1e0344-ec1d-5b48-a673-e5cf874b6c29"
+version = "0.8.21"
+
+[[deps.WebSockets]]
+deps = ["Base64", "Dates", "HTTP", "Logging", "Sockets"]
+git-tree-sha1 = "4162e95e05e79922e44b9952ccbc262832e4ad07"
+uuid = "104b5d7c-a370-577a-8038-80a2059c5097"
+version = "1.6.0"
+
+[[deps.Widgets]]
+deps = ["Colors", "Dates", "Observables", "OrderedCollections"]
+git-tree-sha1 = "fcdae142c1cfc7d89de2d11e08721d0f2f86c98a"
+uuid = "cc8bc4a8-27d6-5769-a93b-9d913e69aa62"
+version = "0.6.6"
+
+[[deps.WoodburyMatrices]]
+deps = ["LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "c1a7aa6219628fcd757dede0ca95e245c5cd9511"
+uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6"
+version = "1.0.0"
+
+[[deps.XML2_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"]
+git-tree-sha1 = "6a451c6f33a176150f315726eba8b92fbfdb9ae7"
+uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a"
+version = "2.13.4+0"
+
+[[deps.XSLT_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libgcrypt_jll", "Libgpg_error_jll", "Libiconv_jll", "XML2_jll", "Zlib_jll"]
+git-tree-sha1 = "a54ee957f4c86b526460a720dbc882fa5edcbefc"
+uuid = "aed1982a-8fda-507f-9586-7b0439959a61"
+version = "1.1.41+0"
+
+[[deps.Xorg_libX11_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libxcb_jll", "Xorg_xtrans_jll"]
+git-tree-sha1 = "afead5aba5aa507ad5a3bf01f58f82c8d1403495"
+uuid = "4f6342f7-b3d2-589e-9d20-edeb45f2b2bc"
+version = "1.8.6+0"
+
+[[deps.Xorg_libXau_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "6035850dcc70518ca32f012e46015b9beeda49d8"
+uuid = "0c0b7dd1-d40b-584c-a123-a41640f87eec"
+version = "1.0.11+0"
+
+[[deps.Xorg_libXdmcp_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "34d526d318358a859d7de23da945578e8e8727b7"
+uuid = "a3789734-cfe1-5b06-b2d0-1dd0d9d62d05"
+version = "1.1.4+0"
+
+[[deps.Xorg_libXext_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libX11_jll"]
+git-tree-sha1 = "d2d1a5c49fae4ba39983f63de6afcbea47194e85"
+uuid = "1082639a-0dae-5f34-9b06-72781eeb8cb3"
+version = "1.3.6+0"
+
+[[deps.Xorg_libXrender_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Xorg_libX11_jll"]
+git-tree-sha1 = "47e45cd78224c53109495b3e324df0c37bb61fbe"
+uuid = "ea2f1a96-1ddc-540d-b46f-429655e07cfa"
+version = "0.9.11+0"
+
+[[deps.Xorg_libpthread_stubs_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "8fdda4c692503d44d04a0603d9ac0982054635f9"
+uuid = "14d82f49-176c-5ed1-bb49-ad3f5cbd8c74"
+version = "0.1.1+0"
+
+[[deps.Xorg_libxcb_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "XSLT_jll", "Xorg_libXau_jll", "Xorg_libXdmcp_jll", "Xorg_libpthread_stubs_jll"]
+git-tree-sha1 = "bcd466676fef0878338c61e655629fa7bbc69d8e"
+uuid = "c7cfdc94-dc32-55de-ac96-5a1b8d977c5b"
+version = "1.17.0+0"
+
+[[deps.Xorg_xtrans_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "e92a1a012a10506618f10b7047e478403a046c77"
+uuid = "c5fb5394-a638-5e4d-96e5-b29de1b5cf10"
+version = "1.5.0+0"
+
+[[deps.ZMQ]]
+deps = ["FileWatching", "PrecompileTools", "Sockets", "ZeroMQ_jll"]
+git-tree-sha1 = "18cfd00df3cbbebf8ea4ec7ea6bbceb3af716bd0"
+uuid = "c2297ded-f4af-51ae-bb23-16f91089e4e1"
+version = "1.3.0"
+
+[[deps.ZeroMQ_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "libsodium_jll"]
+git-tree-sha1 = "f02ce8f0fda1ed40f4d0d59a2ad05e35e8ac9b0e"
+uuid = "8f1865be-045e-5c20-9c9f-bfbfb0764568"
+version = "4.3.5+1"
+
+[[deps.Zlib_jll]]
+deps = ["Libdl"]
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+version = "1.2.12+3"
+
+[[deps.Zstd_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "555d1076590a6cc2fdee2ef1469451f872d8b41b"
+uuid = "3161d3a3-bdf6-5164-811a-617609db77b4"
+version = "1.5.6+1"
+
+[[deps.Zygote]]
+deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "GPUArrays", "GPUArraysCore", "IRTools", "InteractiveUtils", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "PrecompileTools", "Random", "Requires", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"]
+git-tree-sha1 = "f816633be6dc5c0ed9ffedda157ecfda0b3b6a69"
+uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
+version = "0.6.72"
+
+[[deps.ZygoteRules]]
+deps = ["ChainRulesCore", "MacroTools"]
+git-tree-sha1 = "27798139afc0a2afa7b1824c206d5e87ea587a00"
+uuid = "700de1a5-db45-46bc-99cf-38207098b444"
+version = "0.2.5"
+
+[[deps.cuDNN]]
+deps = ["CEnum", "CUDA", "CUDNN_jll"]
+git-tree-sha1 = "5a1ba43303c62f4a09b0d6751422de03424ab0cd"
+uuid = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
+version = "1.1.1"
+
+[[deps.libaec_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "46bf7be2917b59b761247be3f317ddf75e50e997"
+uuid = "477f73a3-ac25-53e9-8cc3-50b2fa2566f0"
+version = "1.1.2+0"
+
+[[deps.libblastrampoline_jll]]
+deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
+uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"
+version = "5.1.1+0"
+
+[[deps.libpng_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"]
+git-tree-sha1 = "b70c870239dc3d7bc094eb2d6be9b73d27bef280"
+uuid = "b53b4c65-9356-5827-b1ea-8c7a1a84506f"
+version = "1.6.44+0"
+
+[[deps.libsixel_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Pkg", "libpng_jll"]
+git-tree-sha1 = "7dfa0fd9c783d3d0cc43ea1af53d69ba45c447df"
+uuid = "075b6546-f08a-558a-be8f-8157d0f608a5"
+version = "1.10.3+1"
+
+[[deps.libsodium_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "f76d682d87eefadd3f165d8d9fda436464213142"
+uuid = "a9144af2-ca23-56d9-984f-0d03f7b5ccf8"
+version = "1.0.20+1"
+
+[[deps.nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+version = "1.48.0+0"
+
+[[deps.oneTBB_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "7d0ea0f4895ef2f5cb83645fa689e52cb55cf493"
+uuid = "1317d2d5-d96f-522e-a858-c73665f53c3e"
+version = "2021.12.0+0"
+
+[[deps.p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
+version = "17.4.0+0"
diff --git a/FNO/Project.toml b/FNO/Project.toml
new file mode 100644
index 0000000..833ed84
--- /dev/null
+++ b/FNO/Project.toml
@@ -0,0 +1,13 @@
+[deps]
+ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DrWatson = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+ParametricDFNOs = "3b406863-5cae-494d-8b19-779e643ae09e"
+PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
diff --git a/FNO/hello.jl b/FNO/hello.jl
new file mode 100644
index 0000000..2ee21bf
--- /dev/null
+++ b/FNO/hello.jl
@@ -0,0 +1,6 @@
+using MPI
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+println("Hello world, I am $(MPI.Comm_rank(comm)) of $(MPI.Comm_size(comm))")
+MPI.Barrier(comm)
diff --git a/FNO/train.jl b/FNO/train.jl
new file mode 100644
index 0000000..9239f81
--- /dev/null
+++ b/FNO/train.jl
@@ -0,0 +1,87 @@
+import Pkg; Pkg.instantiate()
+using ParametricDFNOs.DFNO_3D
+using DrWatson
+using MPI
+using CUDA
+using Statistics
+
+include("utils.jl")
+include("../config.jl")
+
+using .Utils
+using .Config
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+pe_count = MPI.Comm_size(comm)
+
+CUDA.device!(rank % 4)
+partition = [1,pe_count]
+
+nc_lift = 48
+nblocks = 4
+mx, mz, mh = 72, 72, 1
+
+nbatch, epochs, j, ntrain, nvalid = 1, 50, 1, 1, 0
+nbatch, epochs, j, ntrain, nvalid = parse.(Int, ARGS[1:5])
+
+params = Config.get_parameters()
+
+offsets = params["n_offsets"]
+down_rate = params["down_rate"]
+nx = params["nx"] ÷ 2
+nz = params["nz"] ÷ 2
+
+use_nz = false
+labels = @strdict use_nz offsets nbatch epochs ntrain nvalid j
+
+nc_in = use_nz ? 7 : offsets + 1 + 1 + 4 # offsets + 2 velocity models + indices
+nc_out = use_nz ? 1 : offsets
+
+nh = use_nz ? offsets : 1
+mh = use_nz ? mh : 1
+
+@info "Initializing model..."
+
+@assert MPI.Comm_size(comm) == prod(partition)
+modelConfig = DFNO_3D.ModelConfig(nc_in=nc_in, nc_lift=nc_lift, nc_out=nc_out, nx=nx, ny=nz, nz=nh, nt=1, mx=mx, my=mz, mz=mh, mt=1, nblocks=nblocks, partition=partition, dtype=Float32, relu01=false)
+
+x_path = "data/posteriors_iteration_j=$(j-1).jld2"
+y_path = "data/cigs_iteration_j=$j.jld2"
+
+x_train, y_train, x_valid, y_valid = read_velocity_cigs_offsets_as_nc(x_path, y_path, modelConfig, ntrain=ntrain, nvalid=nvalid)
+
+# x_train = reshape(x_train, nc_in, nx, nz, :)
+# y_train = reshape(y_train, nc_out, nx, nz, :)
+# x_valid = reshape(x_valid, nc_in, nx, nz, :)
+# y_valid = reshape(y_valid, nc_out, nx, nz, :)
+
+@info "Loaded data..."
+
+model = DFNO_3D.Model(modelConfig)
+θ = DFNO_3D.initModel(model)
+
+# # To train from a checkpoint
+# filename = "mt=25_mx=10_my=10_mz=10_nblocks=20_nc_in=5_nc_lift=20_nc_mid=128_nc_out=1_nd=20_nt=51_nx=20_ny=20_nz=20_p=8.jld2"
+# DFNO_3D.loadWeights!(θ, filename, "θ_save", partition)
+
+# Normalize CIGs
+max_y = quantile(abs.(vec(y_train)), 0.9999);
+y_train ./= max_y;
+y_valid ./= max_y;
+
+trainConfig = DFNO_3D.TrainConfig(
+    epochs=epochs,
+    x_train=x_train,
+    y_train=y_train,
+    x_valid=nvalid == 0 ? x_train : x_valid,
+    y_valid=nvalid == 0 ? y_train : y_valid,
+    plot_every=5,
+    nbatch=nbatch
+)
+
+DFNO_3D.train!(trainConfig, model, θ, plotEval=Utils.plot_cig_eval_wrapper(use_nz, labels=labels))
+
+MPI.Finalize()
diff --git a/FNO/train_control.jl b/FNO/train_control.jl
new file mode 100644
index 0000000..5fbeff9
--- /dev/null
+++ b/FNO/train_control.jl
@@ -0,0 +1,218 @@
+using ParametricDFNOs.DFNO_3D
+using DrWatson
+using MPI
+using CUDA
+using JLD2
+using HDF5
+
+include("FNO/utils.jl")
+include("config.jl")
+
+using .Utils
+using .Config
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+pe_count = MPI.Comm_size(comm)
+
+CUDA.device!(rank % 4)
+nvalid, offset_is_nz = parse.(Int, ARGS[1:2])
+
+filename = "/pscratch/sd/r/richardr/FNO_CIG/FNO-CIG/weights/DFNO_3D/mt=1_mx=36_my=36_mz=1_nblocks=4_nc_in=27_nc_lift=32_nc_mid=128_nc_out=21_nd=256_nt=1_ntrain=500_nvalid=10_nx=256_ny=128_nz=1_p=1.jld2"
+dataset_path = "/pscratch/sd/r/richardr/FNO_CIG/FNO-CIG/results/concatenated_data_99_quantile_scaled.jld2"
+
+file = filename # projectdir("weights", "DFNO_3D", filename)
+file = load(file)
+
+nblocks = file["nblocks"]
+mx = file["mx"] 
+mz = file["my"]
+mh = file["mz"] 
+nd = file["nd"] 
+nc_lift = file["nc_lift"]
+ntrain = file["ntrain"]
+# nvalid = file["nvalid"]
+
+partition = [1,pe_count]
+params = Config.get_parameters()
+
+down_rate = params["down_rate"]
+nx = params["nx"] ÷ down_rate
+nz = params["nz"] ÷ down_rate
+offsets = params["n_offsets"]
+
+use_nz = Bool(offset_is_nz)
+labels = @strdict use_nz offsets
+
+nc_in = use_nz ? 7 : offsets + 1 + 1 + 4 # offsets + 2 velocity models + indices
+nc_out = use_nz ? 1 : offsets
+
+nh = use_nz ? offsets : 1
+mh = use_nz ? mh : 1 
+
+@info "Initializing model..."
+@assert MPI.Comm_size(comm) == prod(partition)
+
+function read_velocity_cigs_offsets_as_nc_old(path::String, modelConfig::DFNO_3D.ModelConfig; ntrain::Int, nvalid::Int)
+
+    params = Config.get_parameters()
+    offset_start = params["read_offset_start"]
+    offsets = params["n_offsets"]
+    total = params["n_total"]
+
+    # Assumption that x is (nx, nz, 1, n). x0 is (nx, nz). CIG0 is (nh, nx, nz). CIG is (nh, nx, nz, 1, n)
+    function read_x_tensor_helper(file_name, key, indices; flip=false, init_index=1)
+        data = nothing
+        h5open(file_name, "r") do file
+            x_data = file[key[1]]
+            cigs_data = file[key[2]]
+
+            # Read proper indices of x and x0. NOTE: Disclude 3 because no z = 1 for background
+            x = x_data[indices[1], indices[2], 1, indices[4]]
+            x0 = x_data[indices[1], indices[2], 1, init_index]
+            cig0 = cigs_data[offset_start:offset_start+offsets-1, indices[1], indices[2], 1, init_index]
+
+            # Reshape to prepare for augmentation
+            x = reshape(x, :, map(range -> length(range), indices[1:4])...)
+            x0 = reshape(x0, :, map(range -> length(range), indices[1:3])..., 1)
+            cig0 = reshape(cig0, :, map(range -> length(range), indices[1:3])..., 1)
+
+            x0 = repeat(x0, outer=[1, 1, 1, 1, length(indices[4])])
+            cig0 = repeat(cig0, outer=[1, 1, 1, 1, length(indices[4])])
+
+            # Concat along dimension 1
+            data = cat(x, x0, cig0, dims=1)
+            flip && (data = reverse(data, dims=1))
+        end
+
+        # data_channels * nx * ny * nz * nt * n = data_channels * nx * nz * nh * 1 * n
+        return data
+    end
+
+    function read_x_tensor(file_name, key, indices)
+        requested = length(indices[4])
+        @assert (requested - nvalid) % (total - nvalid) == 0
+
+        nrounds = (requested - nvalid) ÷ (total - nvalid)
+        augmented_data = []
+
+        samples_per_round = (requested - nvalid) ÷ nrounds
+
+        for round in 1:nrounds
+            Random.seed!(round % 2)
+            
+            init_index = rand(1:samples_per_round)
+            new_indices = [indices[1:3]..., 1:samples_per_round]
+    
+            data = read_x_tensor_helper(file_name, key, new_indices, flip=(round <= nrounds ÷ 2), init_index=init_index)
+
+            if round == 1
+                augmented_data = data
+            else
+                augmented_data = cat(augmented_data, data, dims=ndims(augmented_data))
+            end
+        end
+
+        Random.seed!(1 % 2) # Simulate round 1
+        init_index = rand(1:samples_per_round+nvalid)
+
+        new_indices = [indices[1:3]..., 1:samples_per_round+nvalid]
+        validation = read_x_tensor_helper(file_name, key, new_indices, flip=false, init_index=init_index)
+
+        augmented_data = cat(augmented_data, validation[:, :, :, :, end-nvalid+1:end], dims=ndims(augmented_data))
+        return augmented_data
+    end
+    
+    function read_y_tensor_helper(file_name, key, indices; flip=false)
+        data = nothing
+        h5open(file_name, "r") do file
+            cigs_data = file[key]
+            data = cigs_data[offset_start:offset_start+offsets-1, indices[1], indices[2], 1, indices[5]] # first dim is offsets as channel, dim 4 which is t = 1:1
+        end
+
+        # channels * nx * ny * nz * nt * n = channels * nx * nz * nh * 1 * n
+        data = reshape(data, :, map(range -> length(range), indices[1:5])...)
+        flip && (data = reverse(data, dims=2))
+        return data
+    end
+
+    function read_y_tensor(file_name, key, indices)
+        requested = length(indices[5])
+        @assert (requested - nvalid) % (total - nvalid) == 0
+
+        nrounds = (requested - nvalid) ÷ (total - nvalid)
+        augmented_data = []
+
+        samples_per_round = (requested - nvalid) ÷ nrounds
+
+        for round in 1:nrounds
+            Random.seed!(round % 2)
+            
+            init_index = rand(1:samples_per_round)
+
+            new_indices = [indices[1:4]..., 1:samples_per_round]
+            data = read_y_tensor_helper(file_name, key, new_indices, flip=(round <= nrounds ÷ 2))
+
+            if round == 1
+                augmented_data = data
+            else
+                augmented_data = cat(augmented_data, data, dims=ndims(augmented_data))
+            end
+        end
+
+        Random.seed!(1 % 2) # Simulate round 1
+        init_index = rand(1:samples_per_round+nvalid)
+
+        new_indices = [indices[1:4]..., 1:samples_per_round+nvalid]
+        validation = read_y_tensor_helper(file_name, key, new_indices, flip=false)
+
+        augmented_data = cat(augmented_data, validation[:, :, :, :, :, end-nvalid+1:end], dims=ndims(augmented_data))
+        return augmented_data
+    end
+
+    dataConfig = DFNO_3D.DataConfig(modelConfig=modelConfig, 
+                                    ntrain=ntrain, 
+                                    nvalid=nvalid, 
+                                    x_file=path,
+                                    y_file=path,
+                                    x_key=["xs", "cigs"],
+                                    y_key="cigs")
+
+    return DFNO_3D.loadDistData(dataConfig, dist_read_x_tensor=read_x_tensor_helper, dist_read_y_tensor=read_y_tensor_helper)
+end
+
+ntrain = 1
+nvalid = 0
+nbatch = 1
+epochs = 500
+
+modelConfig = DFNO_3D.ModelConfig(nc_in=nc_in, nc_lift=nc_lift, nc_out=nc_out, nx=nx, ny=nz, nz=nh, nt=1, mx=mx, my=mz, mz=mh, mt=1, nblocks=nblocks, partition=partition, dtype=Float32, relu01=false)
+x_train, y_train, x_valid, y_valid = read_velocity_cigs_offsets_as_nc_old(dataset_path, modelConfig, ntrain=ntrain, nvalid=nvalid)
+
+@info "Loaded data..."
+
+model = DFNO_3D.Model(modelConfig)
+θ = DFNO_3D.initModel(model)
+
+# DFNO_3D.loadWeights!(θ, filename, "θ_save", partition)
+# labels = @strdict nx nz nblocks mx mz mh ntrain nvalid nc_lift offset_is_nz
+
+# y_predict = DFNO_3D.forward(model, θ, x_valid) |> cpu
+# plot_cig_eval(modelConfig, x_valid, y_valid, y_predict, use_nz=use_nz, additional=labels)
+
+trainConfig = DFNO_3D.TrainConfig(
+    epochs=epochs,
+    x_train=x_train,
+    y_train=y_train,
+    x_valid=nvalid == 0 ? x_train : x_valid,
+    y_valid=nvalid == 0 ? y_train : y_valid,
+    plot_every=epochs÷5,
+    nbatch=nbatch
+)
+
+DFNO_3D.train!(trainConfig, model, θ, plotEval=plot_cig_eval_wrapper(use_nz, labels=labels))
+
+
+MPI.Finalize()
diff --git a/FNO/utils.jl b/FNO/utils.jl
new file mode 100644
index 0000000..b63fccc
--- /dev/null
+++ b/FNO/utils.jl
@@ -0,0 +1,832 @@
+module Utils
+
+include("../config.jl")
+
+using HDF5
+using DrWatson
+using PyPlot
+using Images
+using ArgParse
+using Statistics
+using LinearAlgebra
+using ParametricDFNOs.DFNO_3D
+using FFTW
+using Random
+using .Config
+
+export create_wavelet, create_geometry, generate_noise, parse_commandline, ContJitter, plot_velocity_model, plot_cig, read_velocity_cigs_offsets_as_nc, read_velocity_cigs_offsets_as_nz, plot_cig_eval, plot_cig_eval_wrapper, plotLoss, plot_validation_cig, get_energy_cig, plot_optimize_eval, plot_optimize_loss, plot_control_histogram, plot_velocity_cigs, plot_rankings, plot_cig_diffs
+
+function plot_rankings(true_ranking, test_ranking)
+    fig = figure(figsize=(8, 12))
+
+    PyPlot.rc("figure", titlesize=8)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=8); PyPlot.rc("ytick", labelsize=8)
+    PyPlot.rc("axes", labelsize=8)
+    PyPlot.rc("axes", titlesize=8)
+
+    subplot(1,1,1)
+    plot(true_ranking, test_ranking, "o-")
+    xlabel("true ranking")
+    ylabel("test ranking")
+    title("Comparison of rankings")
+    tight_layout();
+
+    savefig(joinpath("plots", "DFNO_3D", "ranking_comparison.png"), bbox_inches="tight", dpi=300)
+    close(fig);
+end
+
+function plot_control_histogram(hist1, hist2; additional=Dict{String,Any}())
+
+    PyPlot.rc("figure", titlesize=10)
+    PyPlot.rc("font", family="serif"); 
+    PyPlot.rc("xtick", labelsize=10); 
+    PyPlot.rc("ytick", labelsize=10)
+    PyPlot.rc("axes", labelsize=10)
+    PyPlot.rc("axes", titlesize=10)
+
+    fig, axs = plt.subplots(2, 2)
+
+    # Plot histogram 1
+    axs[1, 1].hist(hist1, bins=20, color="blue", alpha=0.7)
+    axs[1, 1].set_title("Histogram of True CIGs")
+    axs[1, 1].set_xlabel("Value")
+    axs[1, 1].set_ylabel("Frequency")
+
+    # Plot histogram 2
+    axs[1, 2].hist(hist2, bins=20, color="green", alpha=0.7)
+    axs[1, 2].set_title("Histogram of Predicted CIGs")
+    axs[1, 2].set_xlabel("Value")
+    axs[1, 2].set_ylabel("Frequency")
+
+    # Plot combined histograms
+    axs[2, 1].hist([hist1, hist2], bins=20, color=["blue", "green"], alpha=0.7, label=["hist1", "hist2"])
+    axs[2, 1].set_title("Combined Histogram")
+    axs[2, 1].set_xlabel("Value")
+    axs[2, 1].set_ylabel("Frequency")
+    axs[2, 1].legend()
+
+    axs[2, 2].hist([hist1, hist2], bins=20, color=["blue", "green"], alpha=0.7, stacked=true)
+    axs[2, 2].set_title("Stacked Histogram")
+    axs[2, 2].set_xlabel("Value")
+    axs[2, 2].set_ylabel("Frequency")
+
+    tight_layout()
+
+    savefig(joinpath("plots", "DFNO_3D", savename(additional; digits=6) * "_DFNO_CIG_histogram.png"), bbox_inches="tight", dpi=300)
+    close(fig)
+end
+
+function plot_cig_diffs(modelConfig, y, ŷ; trainConfig=nothing, use_nz=false, additional=Dict{String,Any}())
+    params = Config.get_parameters()
+
+    offset_start = params["offset_start"]
+    offset_end = params["offset_end"]
+    d = params["d"]
+    n = [modelConfig.nx, modelConfig.ny]
+    n_offsets = params["n_offsets"]
+    down_rate = params["down_rate"]
+
+    # Downsample
+    d = d.*down_rate
+
+    y_plot = reshape(y, (modelConfig.nc_out, modelConfig.nt, modelConfig.nx, modelConfig.ny, modelConfig.nz, :))
+    y_predict = reshape(ŷ, (modelConfig.nc_out, modelConfig.nt, modelConfig.nx, modelConfig.ny, modelConfig.nz, :))
+
+    num_samples = size(y_plot, 6)
+    
+    PyPlot.rc("figure", titlesize=40)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=40); PyPlot.rc("ytick", labelsize=40)
+    PyPlot.rc("axes", labelsize=40)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=40)     # Default fontsize for titles
+
+    fig, axs = subplots(num_samples*2, 6, figsize=(50, num_samples*10), gridspec_kw = Dict("width_ratios" => [4, 1, 4, 1, 4, 1], "height_ratios" => vcat([[1, 3] for i in 1:num_samples]...)))
+
+    for i in 1:num_samples
+        output_CIG = use_nz ? y_predict[1, 1, :, :, :, i] : permutedims(y_predict[:, 1, :, :, 1, i], [2, 3, 1])
+        true_CIG = use_nz ? y_plot[1, 1, :, :, :, i] : permutedims(y_plot[:, 1, :, :, 1, i], [2, 3, 1])
+
+        plot_cig_helper(true_CIG, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 1:2])
+        plot_cig_helper(output_CIG, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 3:4])
+        plot_cig_helper(5f0 .* abs.(true_CIG - output_CIG), n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 5:6])
+
+        if i == 1
+            axs[1, 1].set_title("True CIG")
+            axs[1, 3].set_title("Predicted CIG")
+            axs[1, 5].set_title("5X Diff")
+        end
+    end
+
+    suptitle("Model and CIG Comparison", fontsize=20)
+    
+    figname = _getFigname(trainConfig, additional)
+    tight_layout()
+
+    savefig(joinpath("plots", "DFNO_3D", savename(figname; digits=6) * "_DFNO_CIG_fitting.png"), bbox_inches="tight", dpi=300)
+    close(fig)
+end
+
+function plot_velocity_cigs(modelConfig, x, y; additional=Dict{String,Any}())
+
+end
+
+function plot_cig_eval_wrapper(use_nz; labels=Dict{String,Any}())
+    function _wrapper(modelConfig, x_plot, y_plot, y_predict; trainConfig, additional=labels)
+        return plot_cig_eval(modelConfig, x_plot, y_plot, y_predict, use_nz=use_nz, trainConfig=trainConfig, additional=additional)
+    end
+    return _wrapper
+end
+
+function plot_cig_eval(modelConfig, x_plot, y_plot, y_predict; use_nz=false, trainConfig=nothing, additional=Dict{String,Any}())
+    params = Config.get_parameters()
+
+    offset_start = params["offset_start"]
+    offset_end = params["offset_end"]
+    d = params["d"]
+    n = [modelConfig.nx, modelConfig.ny]
+    n_offsets = params["n_offsets"]
+    down_rate = params["down_rate"]
+
+    # Downsample
+    d = d.*down_rate
+
+    # Reshape the data to fit the model configuration
+    x_plot = reshape(x_plot, (modelConfig.nc_in, modelConfig.nt, modelConfig.nx, modelConfig.ny, modelConfig.nz, :))
+    y_plot = reshape(y_plot, (modelConfig.nc_out, modelConfig.nt, modelConfig.nx, modelConfig.ny, modelConfig.nz, :))
+    y_predict = reshape(y_predict, (modelConfig.nc_out, modelConfig.nt, modelConfig.nx, modelConfig.ny, modelConfig.nz, :))
+
+    num_samples = size(x_plot, 6)
+    
+    PyPlot.rc("figure", titlesize=40)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=40); PyPlot.rc("ytick", labelsize=40)
+    PyPlot.rc("axes", labelsize=40)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=40)     # Default fontsize for titles
+
+    fig, axs = subplots(num_samples*2, 13, figsize=(110, num_samples*10), gridspec_kw = Dict("width_ratios" => [4, 4, 1, 4, 4, 1, 4, 1, 4, 1, 4, 1, 4], "height_ratios" => vcat([[1, 3] for i in 1:num_samples]...)))
+
+    for i in 1:num_samples
+        perturbed_model = x_plot[1, 1, :, :, 1, i]
+        init_background_model = x_plot[2, 1, :, :, 1, i]
+        input_CIG = use_nz ? x_plot[3, 1, :, :, :, i] : permutedims(x_plot[3:(3+n_offsets-1), 1, :, :, 1, i], [2, 3, 1])
+        output_CIG = use_nz ? y_predict[1, 1, :, :, :, i] : permutedims(y_predict[:, 1, :, :, 1, i], [2, 3, 1])
+        true_CIG = use_nz ? y_plot[1, 1, :, :, :, i] : permutedims(y_plot[:, 1, :, :, 1, i], [2, 3, 1])
+
+        axs[i*2-1, 1].set_visible(false)
+        axs[i*2-1, 4].set_visible(false)
+        axs[i*2-1, 13].set_visible(false)
+
+        plot_velocity_model_helper(init_background_model, n, d, axs[i*2, 1])
+        plot_velocity_model_helper(perturbed_model, n, d, axs[i*2, 4])
+
+        plot_cig_helper(input_CIG, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 2:3])
+        plot_cig_helper(true_CIG, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 5:6])
+        plot_cig_helper(output_CIG, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 7:8])
+        plot_cig_helper(5f0 .* abs.(true_CIG - output_CIG), n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 9:10])
+        plot_cig_helper(abs.(fftshift(fft(true_CIG - output_CIG))), n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 11:12])
+
+        plot_energy_vs_offset_helper(true_CIG, output_CIG, n, d, offset_start, offset_end, n_offsets, axs[i*2, 13])
+
+        # Label columns for the first sample
+        if i == 1
+            axs[1, 1].set_title("Init velocity model")
+            axs[1, 2].set_title("Init CIG")
+            axs[1, 4].set_title("Smoothed velocity model")
+            axs[1, 5].set_title("True CIG")
+            axs[1, 7].set_title("Predicted CIG")
+            axs[1, 9].set_title("5X Diff")
+            axs[1, 11].set_title("FFT of 1X difference")
+            axs[1, 13].set_title("% Energy Distribution")
+        end
+    end
+
+    suptitle("Model and CIG Comparison", fontsize=20)
+    
+    figname = _getFigname(trainConfig, additional)
+    tight_layout()
+
+    savefig(joinpath("plots", "DFNO_3D", savename(figname; digits=6) * "_DFNO_CIG_fitting.png"), bbox_inches="tight", dpi=300)
+    close(fig)
+end
+
+# NOTE: Make sure input has offsets as first dimension. Choice due to channel being the first output of DFNO
+function get_energy_cig(n, n_offsets, cig; dist=false)
+    cig = reshape(cig, n_offsets, :)
+    center_offset = n_offsets ÷ 2 + 1
+
+    @assert n_offsets % 2 == 1
+    range = n_offsets - center_offset
+
+    energies = [sum(abs2, fft(permutedims(cig[center_offset-distance:center_offset+distance, :], [2, 1]), [2])) for distance in 0:range]
+
+    if dist
+        reduce = ParReduce(eltype(energies))
+        energies = reduce(energies)
+    end
+
+    return (energies ./ energies[end]) # Last contains the total energy
+end
+
+# NOTE: All plot functions will receive offsets as last dimension. It is their responsibility to permute
+function plot_energy_vs_offset_helper(true_CIG, output_CIG, n, d, offset_start, offset_end, n_offsets, ax)
+
+    true_CIG = permutedims(true_CIG, [3, 1, 2])
+    output_CIG = permutedims(output_CIG, [3, 1, 2])
+
+    true_energy = get_energy_cig(n, n_offsets, true_CIG)
+    output_energy = get_energy_cig(n, n_offsets, output_CIG)
+
+    x_values = LinRange(0, offset_end, length(true_energy))
+    x_ticks = [x_values[1], x_values[Int(ceil(length(x_values) / 2))], x_values[end]]
+
+    ax.plot(x_values, true_energy, label="True Energy", color="blue")
+    ax.plot(x_values, output_energy, label="Output Energy", color="red")
+    ax.set_xlabel("Distance from 0 offset [m]")
+    ax.set_ylabel("% of Energy")
+    ax.set_xticks(x_ticks)
+    ax.legend()
+end
+
+function plot_velocity_model_helper(x, n, d, ax)
+    # Assume that vmin and vmax are computed similarly to the plot_cig_helper function
+    vmin = quantile(vec(x), 0.05)  # 5th percentile
+    vmax = quantile(vec(x), 0.95)  # 95th percentile
+    extentfull = (0f0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0f0)
+    cax = ax.imshow(x', vmin=vmin, vmax=vmax, extent=extentfull, aspect="auto")
+    ax.set_xlabel("X [m]")
+    ax.set_ylabel("Z [m]")
+end
+
+function plot_cig_helper(cig, n, d, offset_start, offset_end, n_offsets, axs)
+    y = reshape(cig, n[1], n[2], n_offsets)
+    
+    ### X, Z position in km
+    xpos = 3.6f3
+    zpos = 2.7f3
+    xgrid = Int(round(xpos / d[1]))
+    zgrid = Int(round(zpos / d[2]))
+    
+    # Adjust the spacing between the plots
+    subplots_adjust(hspace=0.0, wspace=0.0)
+    
+    vmin1, vmax1 = (-1, 1) .* quantile(abs.(vec(y[:,zgrid,:,1])), 0.99)
+    vmin2, vmax2 = (-1, 1) .* quantile(abs.(vec(y[:,:,div(n_offsets,2)+1,1])), 0.88)
+    vmin3, vmax3 = (-1, 1) .* quantile(abs.(vec(y[xgrid,:,:,1])), 0.999)
+    sca(axs[1, 1])
+    
+    # Top left subplot
+    axs[1, 1].imshow(y[:,zgrid,:,1]', aspect="auto", cmap="gray", interpolation="none",vmin=vmin1, vmax=vmax1,
+        extent=(0f0, (n[1]-1)*d[1], offset_start, offset_end))
+    axs[1, 1].set_ylabel("Offset [m]", fontsize=40)
+    axs[1, 1].set_xticklabels([])
+    axs[1, 1].set_xlabel("")
+    hlines(y=0, colors=:b, xmin=0, xmax=(n[1]-1)*d[1], linewidth=3)
+    vlines(x=xpos, colors=:b, ymin=offset_start, ymax=offset_end, linewidth=3)
+    
+    # Bottom left subplot
+    sca(axs[2, 1])
+    axs[2, 1].imshow(y[:,:,div(n_offsets,2)+1,1]', aspect="auto", cmap="gray", interpolation="none",vmin=vmin2, vmax=vmax2,
+    extent=(0f0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0f0))
+    axs[2, 1].set_xlabel("X [m]", fontsize=40)
+    axs[2, 1].set_ylabel("Z [m]", fontsize=40)
+    axs[2, 1].set_xticks([0, 2000, 4000, 6000])
+    axs[2, 1].set_xticklabels(["0", "2000", "4000", "6000"])
+    axs[2, 1].set_yticks([1000, 2000, 3000])
+    axs[2, 1].set_yticklabels(["1000", "2000", "3000"])
+    
+    # axs[2, 2].get_shared_x_axes().join(axs[1, 1], axs[2, 1])
+    vlines(x=xpos, colors=:b, ymin=0, ymax=(n[2]-1)*d[2], linewidth=3)
+    hlines(y=zpos, colors=:b, xmin=0, xmax=(n[1]-1)*d[1], linewidth=3)
+    
+    # Top right subplot
+    axs[1, 2].set_visible(false)
+    
+    # Bottom right subplot
+    sca(axs[2, 2])
+    axs[2, 2].imshow(y[xgrid,:,:,1], aspect="auto", cmap="gray", interpolation="none",vmin=vmin3, vmax=vmax3,
+    extent=(offset_start, offset_end, (n[2]-1)*d[2], 0f0))
+    axs[2, 2].set_xlabel("Offset [m]", fontsize=40)
+    # Share y-axis with bottom left
+    # axs[2, 2].get_shared_y_axes().join(axs[2, 2], axs[2, 1])
+    axs[2, 2].set_yticklabels([])
+    axs[2, 2].set_ylabel("")
+    vlines(x=0, colors=:b, ymin=0, ymax=(n[2]-1)*d[2], linewidth=3)
+    hlines(y=zpos, colors=:b, xmin=offset_end, xmax=offset_start, linewidth=3)
+end
+
+function _getFigname(config, additional::Dict)
+    isnothing(config) && return additional
+
+    nbatch = config.nbatch
+    epochs = config.epochs
+    ntrain = size(config.x_train, 3)
+    nvalid = size(config.x_valid, 3)
+    
+    figname = @strdict nbatch epochs ntrain nvalid
+    return merge(additional, figname)
+end
+
+function read_velocity_cigs_offsets_as_nc(x_path::String, y_path::String, modelConfig::DFNO_3D.ModelConfig; ntrain::Int, nvalid::Int)
+
+    params = Config.get_parameters()
+    offset_start = params["read_offset_start"]
+    offsets = params["n_offsets"]
+    total = params["n_total"]
+
+    # Assumption that x is (nx, nz, 1, n). x0 is (nx, nz). CIG0 is (nh, nx, nz). CIG is (nh, nx, nz, 1, n)
+    function read_x_tensor_helper(file_name, key, indices; flip=false, init_index=1)
+        data = nothing
+        h5open(file_name, "r") do file
+            x_data = file[key[1]]
+            cigs_data = file[key[2]]
+
+            # Read proper indices of x and x0. NOTE: Disclude 3 because no z = 1 for background
+            x = x_data[indices[1], indices[2], 1, indices[4]]
+            x0 = x_data[indices[1], indices[2], 1, init_index]
+            cig0 = cigs_data[offset_start:offset_start+offsets-1, indices[1], indices[2], 1, init_index]
+
+            # Reshape to prepare for augmentation
+            x = reshape(x, :, map(range -> length(range), indices[1:4])...)
+            x0 = reshape(x0, :, map(range -> length(range), indices[1:3])..., 1)
+            cig0 = reshape(cig0, :, map(range -> length(range), indices[1:3])..., 1)
+
+            x0 = repeat(x0, outer=[1, 1, 1, 1, length(indices[4])])
+            cig0 = repeat(cig0, outer=[1, 1, 1, 1, length(indices[4])])
+
+            # Concat along dimension 1
+            data = cat(x, x0, cig0, dims=1)
+            flip && (data = reverse(data, dims=1))
+        end
+
+        # data_channels * nx * ny * nz * nt * n = data_channels * nx * nz * nh * 1 * n
+        return data
+    end
+
+    function read_x_tensor(file_name, key, indices)
+        requested = length(indices[4])
+        @assert (requested - nvalid) % (total - nvalid) == 0
+
+        nrounds = (requested - nvalid) ÷ (total - nvalid)
+        augmented_data = []
+
+        samples_per_round = (requested - nvalid) ÷ nrounds
+
+        for round in 1:nrounds
+            Random.seed!(round % 2)
+            
+            init_index = rand(1:samples_per_round)
+            new_indices = [indices[1:3]..., 1:samples_per_round]
+    
+            data = read_x_tensor_helper(file_name, key, new_indices, flip=(round <= nrounds ÷ 2), init_index=init_index)
+
+            if round == 1
+                augmented_data = data
+            else
+                augmented_data = cat(augmented_data, data, dims=ndims(augmented_data))
+            end
+        end
+
+        Random.seed!(1 % 2) # Simulate round 1
+        init_index = rand(1:samples_per_round+nvalid)
+
+        new_indices = [indices[1:3]..., 1:samples_per_round+nvalid]
+        validation = read_x_tensor_helper(file_name, key, new_indices, flip=false, init_index=init_index)
+
+        augmented_data = cat(augmented_data, validation[:, :, :, :, end-nvalid+1:end], dims=ndims(augmented_data))
+        return augmented_data
+    end
+    
+    function read_y_tensor_helper(file_name, key, indices; flip=false)
+        data = nothing
+        h5open(file_name, "r") do file
+            cigs_data = file[key]
+            data = cigs_data[offset_start:offset_start+offsets-1, indices[1], indices[2], 1, indices[5]] # first dim is offsets as channel, dim 4 which is t = 1:1
+        end
+
+        # channels * nx * ny * nz * nt * n = channels * nx * nz * nh * 1 * n
+        data = reshape(data, :, map(range -> length(range), indices[1:5])...)
+        flip && (data = reverse(data, dims=2))
+        return data
+    end
+
+    function read_y_tensor(file_name, key, indices)
+        requested = length(indices[5])
+        @assert (requested - nvalid) % (total - nvalid) == 0
+
+        nrounds = (requested - nvalid) ÷ (total - nvalid)
+        augmented_data = []
+
+        samples_per_round = (requested - nvalid) ÷ nrounds
+
+        for round in 1:nrounds
+            Random.seed!(round % 2)
+            
+            init_index = rand(1:samples_per_round)
+
+            new_indices = [indices[1:4]..., 1:samples_per_round]
+            data = read_y_tensor_helper(file_name, key, new_indices, flip=(round <= nrounds ÷ 2))
+
+            if round == 1
+                augmented_data = data
+            else
+                augmented_data = cat(augmented_data, data, dims=ndims(augmented_data))
+            end
+        end
+
+        Random.seed!(1 % 2) # Simulate round 1
+        init_index = rand(1:samples_per_round+nvalid)
+
+        new_indices = [indices[1:4]..., 1:samples_per_round+nvalid]
+        validation = read_y_tensor_helper(file_name, key, new_indices, flip=false)
+
+        augmented_data = cat(augmented_data, validation[:, :, :, :, :, end-nvalid+1:end], dims=ndims(augmented_data))
+        return augmented_data
+    end
+
+    dataConfig = DFNO_3D.DataConfig(modelConfig=modelConfig, 
+                                    ntrain=ntrain, 
+                                    nvalid=nvalid, 
+                                    x_file=x_path,
+                                    y_file=y_path,
+                                    x_key=["x0", "CIG0"],
+                                    y_key="CIGs")
+
+    return DFNO_3D.loadDistData(dataConfig, dist_read_x_tensor=read_x_tensor_helper, dist_read_y_tensor=read_y_tensor_helper)
+end
+
+function read_velocity_cigs_offsets_as_nz(path::String, modelConfig::DFNO_3D.ModelConfig; ntrain::Int, nvalid::Int)
+
+    params = Config.get_parameters()
+    offset_start = params["read_offset_start"]
+
+    # Assumption that x is (nx, nz, 1, n). x0 is (nx, nz). CIG0 is (nh, nx, nz). CIG is (nh, nx, nz, 1, n)
+    function read_x_tensor(file_name, key, indices)
+        data = nothing
+        h5open(file_name, "r") do file
+            # Size to augment
+            target_zeros = zeros(modelConfig.dtype, 1, map(range -> length(range), indices[1:4])...)
+
+            x_data = file[key[1]]
+            cigs_data = file[key[2]]
+
+            # Read proper indices of x and x0. NOTE: Disclude 3 because no z = h = offset = 1 for background
+            x = x_data[indices[1], indices[2], 1, indices[4]]
+            x0 = x_data[indices[1], indices[2], 1, 1] # Use the first x for init model for now
+            cig0 = cigs_data[indices[3] .+ (offset_start - 1), indices[1], indices[2], 1, 1] # Use the first cig for init cig for now
+            cig0 = permutedims(cig0, [2, 3, 1])
+
+            # Reshape to prepare for augmentation
+            x = reshape(x, 1, length(indices[1]), length(indices[2]), 1, length(indices[4]))
+            x0 = reshape(x0, 1, length(indices[1]), length(indices[2]), 1, 1)
+            cig0 = reshape(cig0, 1, length(indices[1]), length(indices[2]), length(indices[3]), 1)
+
+            # Augment to full size
+            x = target_zeros .+ x
+            x0 = target_zeros .+ x0
+            cig0 = target_zeros .+ cig0
+
+            # Concat along dimension 1
+            data = cat(x, x0, cig0, dims=1)
+        end
+
+        # data_channels * nx * ny * nz * nt * n = data_channels * nx * nz * nh * 1 * n
+        return data
+    end
+    
+    function read_y_tensor(file_name, key, indices)
+        data = nothing
+        h5open(file_name, "r") do file
+            cigs_data = file[key]
+            cigs = cigs_data[indices[3] .+ (offset_start - 1), indices[1], indices[2], indices[4], indices[5]] # dim 4 which is t = 1:1
+            data = permutedims(cigs, [2, 3, 1, 4, 5])
+        end
+
+        # channels * nx * ny * nz * nt * n = channels * nx * nz * nh * 1 * n
+        return reshape(data, 1, size(data)...)
+    end
+
+    dataConfig = DFNO_3D.DataConfig(modelConfig=modelConfig, 
+                                    ntrain=ntrain, 
+                                    nvalid=nvalid, 
+                                    x_file=path,
+                                    y_file=path,
+                                    x_key=["xs", "cigs"],
+                                    y_key="cigs")
+
+    return DFNO_3D.loadDistData(dataConfig, dist_read_x_tensor=read_x_tensor, dist_read_y_tensor=read_y_tensor)
+end
+
+function create_wavelet(timeD, dtD, f0)
+    return ricker_wavelet(timeD, dtD, f0)
+end
+
+function create_geometry(n, d, nsrc, nxrec, dtD, timeD)
+    xrec = range(0f0, stop=(n[1]-1)*d[1], length=nxrec)
+    yrec = 0f0
+    zrec = range(d[1], stop=d[1], length=nxrec)
+    recGeometry = Geometry(xrec, yrec, zrec; dt=dtD, t=timeD, nsrc=nsrc)
+
+    ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
+    zsrc = convertToCell(range(d[1], stop=d[1], length=nsrc))
+    xsrc = convertToCell(ContJitter((n[1]-1)*d[1], nsrc))
+    srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtD, t=timeD)
+
+    return recGeometry, srcGeometry
+end
+
+function generate_noise(d_obs, nsrc, snr, q)
+    noise_ = deepcopy(d_obs)
+    for l = 1:nsrc
+        noise_.data[l] = randn(Float32, size(d_obs.data[l]))
+        noise_.data[l] = real.(ifft(fft(noise_.data[l]).*fft(q.data[1])))
+    end
+    noise_ = noise_/norm(noise_) * norm(d_obs) * 10f0^(-snr/20f0)
+    return noise_
+end
+
+function parse_commandline()
+    s = ArgParseSettings()
+    @add_arg_table s begin
+        "--startidx"
+            help = "Start index"
+            arg_type = Int
+            default = 1
+        "--endidx"
+            help = "End index"
+            arg_type = Int
+            default = 1000
+        "--n_offsets"
+            help = "num of offsets"
+            arg_type = Int
+            default = 51
+        "--offset_start"
+            help = "start of offset"
+            arg_type = Float32
+            default = -500f0
+        "--offset_end"
+            help = "end of offset"
+            arg_type = Float32
+            default = 500f0
+        "--keep_offset_num"
+            help = "keep how many offset during training"
+            arg_type = Int
+            default = 51
+    end
+    return parse_args(s)
+end
+
+# function de_z_shape_simple(G::NetworkGlow, X::AbstractArray{T, N}) where {T, N}
+#     G.split_scales && (Z_save = array_of_array(X, max(G.L-1,1)))
+
+#     logdet_ = 0
+#     for i=1:G.L
+#         (G.split_scales) && (X = G.squeezer.forward(X))
+#         if G.split_scales && (i < G.L || i == 1)    # don't split after last iteration
+#             X, Z = tensor_split(X)
+#             Z_save[i] = Z
+#             G.Z_dims[i] = collect(size(Z))
+#         end
+#     end
+#     G.split_scales && (X = cat_states(Z_save, X))
+
+#     return X
+# end
+
+# function z_shape_simple(G::NetworkGlow, ZX_test::AbstractArray{T, N}) where {T, N}
+#     Z_save, ZX = split_states(ZX_test[:], G.Z_dims)
+#     for i=G.L:-1:1
+#         if i < G.L
+#             ZX = tensor_cat(ZX, Z_save[i])
+#         end
+#         ZX = G.squeezer.inverse(ZX) 
+#     end
+#     ZX
+# end
+
+function ContJitter(l::Number, num::Int)
+    #l = length, num = number of samples
+    interval_width = l/num
+    interval_center = range(interval_width/2, stop = l-interval_width/2, length=num)
+    randomshift = interval_width .* rand(Float32, num) .- interval_width/2
+
+    return interval_center .+ randomshift
+end
+
+function plot_velocity_model(x, n, d, filename)
+    vmin = quantile(vec(x), 0.05)  # 5th percentile
+    vmax = quantile(vec(x), 0.95)  # 95th percentile
+
+    fig, ax = subplots(figsize=(20,12)) 
+    extentfull = (0f0, (n[1]-1)*d[1], (n[end]-1)*d[end], 0f0)
+    cax = ax.imshow(x', vmin=vmin, vmax=vmax, extent=extentfull, aspect=0.45*(extentfull[2]-extentfull[1])/(extentfull[3]-extentfull[4]))
+    ax.set_xlabel("X [m]", fontsize=40)
+    ax.set_ylabel("Z [m]", fontsize=40)
+    savefig(filename, bbox_inches="tight", dpi=300)
+    close(fig)
+end
+
+function plot_cig(cig, n, d, offset_start, offset_end, n_offsets, filename)
+    y = reshape(permutedims(cig, [2, 3, 1]), n[1], n[2], n_offsets, 1)
+
+    PyPlot.rc("figure", titlesize=40)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=40); PyPlot.rc("ytick", labelsize=40)
+    PyPlot.rc("axes", labelsize=40)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=40)     # Default fontsize for titles
+    
+    ### X, Z position in km
+    xpos = 3.6f3
+    zpos = 2.7f3
+    xgrid = Int(round(xpos / d[1]))
+    zgrid = Int(round(zpos / d[2]))
+    
+    # Create a figure and a 2x2 grid of subplots
+    fig, axs = subplots(2, 2, figsize=(20,12), gridspec_kw = Dict("width_ratios" => [3, 1], "height_ratios" => [1, 3]))
+    
+    # Adjust the spacing between the plots
+    subplots_adjust(hspace=0.0, wspace=0.0)
+    
+    vmin1, vmax1 = (-1, 1) .* quantile(abs.(vec(y[:,zgrid,:,1])), 0.99)
+    vmin2, vmax2 = (-1, 1) .* quantile(abs.(vec(y[:,:,div(n_offsets,2)+1,1])), 0.88)
+    vmin3, vmax3 = (-1, 1) .* quantile(abs.(vec(y[xgrid,:,:,1])), 0.999)
+    sca(axs[1, 1])
+    
+    # Top left subplot
+    axs[1, 1].imshow(y[:,zgrid,:,1]', aspect="auto", cmap="gray", interpolation="none",vmin=vmin1, vmax=vmax1,
+        extent=(0f0, (n[1]-1)*d[1], offset_start, offset_end))
+    axs[1, 1].set_ylabel("Offset [m]", fontsize=40)
+    axs[1, 1].set_xticklabels([])
+    axs[1, 1].set_xlabel("")
+    hlines(y=0, colors=:b, xmin=0, xmax=(n[1]-1)*d[1], linewidth=3)
+    vlines(x=xpos, colors=:b, ymin=offset_start, ymax=offset_end, linewidth=3)
+    
+    # Bottom left subplot
+    sca(axs[2, 1])
+    axs[2, 1].imshow(y[:,:,div(n_offsets,2)+1,1]', aspect="auto", cmap="gray", interpolation="none",vmin=vmin2, vmax=vmax2,
+    extent=(0f0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0f0))
+    axs[2, 1].set_xlabel("X [m]", fontsize=40)
+    axs[2, 1].set_ylabel("Z [m]", fontsize=40)
+    axs[2, 1].set_xticks([0, 1000, 2000, 3000, 4000, 5000])
+    axs[2, 1].set_xticklabels(["0", "1000", "2000", "3000", "4000", "5000"])
+    axs[2, 1].set_yticks([1000, 2000, 3000])
+    axs[2, 1].set_yticklabels(["1000", "2000", "3000"])
+    
+    axs[2, 2].get_shared_x_axes().join(axs[1, 1], axs[2, 1])
+    vlines(x=xpos, colors=:b, ymin=0, ymax=(n[2]-1)*d[2], linewidth=3)
+    hlines(y=zpos, colors=:b, xmin=0, xmax=(n[1]-1)*d[1], linewidth=3)
+    
+    # Top right subplot
+    axs[1, 2].set_visible(false)
+    
+    # Bottom right subplot
+    sca(axs[2, 2])
+    axs[2, 2].imshow(y[xgrid,:,:,1], aspect="auto", cmap="gray", interpolation="none",vmin=vmin3, vmax=vmax3,
+    extent=(offset_start, offset_end, (n[2]-1)*d[2], 0f0))
+    axs[2, 2].set_xlabel("Offset [m]", fontsize=40)
+    # Share y-axis with bottom left
+    axs[2, 2].get_shared_y_axes().join(axs[2, 2], axs[2, 1])
+    axs[2, 2].set_yticklabels([])
+    axs[2, 2].set_ylabel("")
+    vlines(x=0, colors=:b, ymin=0, ymax=(n[2]-1)*d[2], linewidth=3)
+    hlines(y=zpos, colors=:b, xmin=offset_end, xmax=offset_start, linewidth=3)
+    
+    # Remove the space between subplots and hide the spines
+    for ax in reshape(axs, :)
+        for spine in ["top", "right", "bottom", "left"]
+            ax.spines[spine].set_visible(false)
+        end
+    end
+    
+    savefig(filename, bbox_inches="tight", dpi=300);
+    close(fig)
+end
+
+function plotLoss(ep, Loss, Loss_valid, trainConfig::DFNO_3D.TrainConfig; additional=Dict())
+
+    ntrain = size(trainConfig.x_train, 3)
+    nbatches = Int(ntrain/trainConfig.nbatch)
+
+    loss_train = Loss[1:ep*nbatches]
+    loss_valid = Loss_valid[1:ep]
+    fig = figure(figsize=(20, 12))
+
+    PyPlot.rc("figure", titlesize=8)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=8); PyPlot.rc("ytick", labelsize=8)
+    PyPlot.rc("axes", labelsize=8)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=8)     # Default fontsize for titles
+
+    subplot(1,3,1)
+    plot(loss_train)
+    xlabel("batch iterations")
+    ylabel("loss")
+    title("training loss at epoch $ep")
+    subplot(1,3,2)
+    plot(1:nbatches:nbatches*ep, loss_valid);
+    xlabel("batch iterations")
+    ylabel("loss")
+    title("validation loss at epoch $ep")
+    subplot(1,3,3)
+    plot(loss_train);
+    plot(1:nbatches:nbatches*ep, loss_valid); 
+    xlabel("batch iterations")
+    ylabel("loss")
+    title("Objective function at epoch $ep")
+    legend(["training", "validation"])
+    tight_layout();
+
+    figname = _getFigname(trainConfig, additional)
+
+    savefig(joinpath("plots", "DFNO_3D", savename(figname; digits=6) * "_DFNO_CIG_loss.png"), bbox_inches="tight", dpi=300)
+    close(fig);
+end
+
+function plot_optimize_loss(Loss, labels)
+
+    fig = figure(figsize=(8, 12))
+
+    PyPlot.rc("figure", titlesize=8)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=8); PyPlot.rc("ytick", labelsize=8)
+    PyPlot.rc("axes", labelsize=8)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=8)     # Default fontsize for titles
+
+    subplot(1,1,1)
+    plot(Loss)
+    xlabel("iterations")
+    ylabel("loss")
+    title("loss after $(length(Loss)) iterations")
+    tight_layout();
+
+    savefig(joinpath("plots", "DFNO_3D", savename(labels; digits=6) * "_DFNO_CIG_OPT_loss.png"), bbox_inches="tight", dpi=300)
+    close(fig);
+end
+
+function plot_optimize_eval(modelConfig, start_velocity, start_cig, end_velocity, end_cig; trainConfig=nothing, additional=Dict{String,Any}())
+    params = Config.get_parameters()
+
+    offset_start = params["offset_start"]
+    offset_end = params["offset_end"]
+    d = params["d"]
+    n = [modelConfig.nx, modelConfig.ny]
+    n_offsets = params["n_offsets"]
+    down_rate = params["down_rate"]
+
+    # Downsample
+    d = d.*down_rate
+
+    # Reshape the data to fit the model configuration
+    start_velocity = reshape(start_velocity, (modelConfig.nx, modelConfig.ny, :))
+    end_velocity = reshape(end_velocity, (modelConfig.nx, modelConfig.ny, :))
+
+    start_cig = reshape(start_cig, (modelConfig.nc_out, modelConfig.nx, modelConfig.ny, :))
+    end_cig = reshape(end_cig, (modelConfig.nc_out, modelConfig.nx, modelConfig.ny, :))
+
+    num_samples = size(start_velocity, 3)
+    
+    PyPlot.rc("figure", titlesize=40)
+    PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=40); PyPlot.rc("ytick", labelsize=40)
+    PyPlot.rc("axes", labelsize=40)     # Default fontsize for x and y labels
+    PyPlot.rc("axes", titlesize=40)     # Default fontsize for titles
+
+    fig, axs = subplots(num_samples*2, 12, figsize=(105, num_samples*10), gridspec_kw = Dict("width_ratios" => [4, 4, 1, 4, 4, 1, 4, 4, 1, 4, 1, 4], "height_ratios" => vcat([[1, 3] for i in 1:num_samples]...)))
+
+    for i in 1:num_samples
+        start_velocity_sample = start_velocity[:, :, i]
+        end_velocity_sample = end_velocity[:, :, i]
+
+        start_cig_sample = permutedims(start_cig[:, :, :, i], [2, 3, 1])
+        end_cig_sample = permutedims(end_cig[:, :, :, i], [2, 3, 1])
+
+        axs[i*2-1, 1].set_visible(false)
+        axs[i*2-1, 4].set_visible(false)
+        axs[i*2-1, 7].set_visible(false)
+        axs[i*2-1, 12].set_visible(false)
+
+        plot_velocity_model_helper(start_velocity_sample, n, d, axs[i*2, 1])
+        plot_velocity_model_helper(end_velocity_sample, n, d, axs[i*2, 4])
+        plot_velocity_model_helper(abs.(start_velocity_sample - end_velocity_sample), n, d, axs[i*2, 7])
+
+        plot_cig_helper(start_cig_sample, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 2:3])
+        plot_cig_helper(end_cig_sample, n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 5:6])
+        plot_cig_helper(5f0 .* abs.(start_cig_sample - end_cig_sample), n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 8:9])
+        plot_cig_helper(abs.(fftshift(fft(start_cig_sample - end_cig_sample, [3]))), n, d, offset_start, offset_end, n_offsets, axs[i*2-1:i*2, 10:11])
+
+        plot_energy_vs_offset_helper(start_cig_sample, end_cig_sample, n, d, offset_start, offset_end, n_offsets, axs[i*2, 12])
+
+        # Label columns for the first sample
+        if i == 1
+            axs[2, 1].set_title("Init velocity model")
+            axs[1, 2].set_title("Target CIG")
+            axs[2, 4].set_title("Opt velocity model")
+            axs[1, 5].set_title("Opt CIG")
+            axs[2, 7].set_title("1X Diff velocity models")
+            axs[1, 8].set_title("5X Diff CIG")
+            axs[1, 10].set_title("FFT 1X CIG Diff")
+            axs[2, 12].set_title("% Energy vs Offsets")
+        end
+    end
+
+    suptitle("Optimization Comparison", fontsize=20)
+    tight_layout()
+
+    savefig(joinpath("plots", "DFNO_3D", savename(additional; digits=6) * "_DFNO_FIT_CIG_OPT_fitting.png"), bbox_inches="tight", dpi=300)
+    close(fig)
+end
+
+end
diff --git a/FNO/validation_control.jl b/FNO/validation_control.jl
new file mode 100644
index 0000000..a5ce92b
--- /dev/null
+++ b/FNO/validation_control.jl
@@ -0,0 +1,200 @@
+using ParametricDFNOs.DFNO_3D
+using DrWatson
+using MPI
+using CUDA
+using JLD2
+using HDF5
+
+include("FNO/utils.jl")
+include("config.jl")
+
+using .Utils
+using .Config
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+pe_count = MPI.Comm_size(comm)
+
+CUDA.device!(rank % 4)
+nvalid, offset_is_nz = parse.(Int, ARGS[1:2])
+
+filename = "/pscratch/sd/r/richardr/FNO_CIG/FNO-CIG/weights/DFNO_3D/mt=1_mx=36_my=36_mz=1_nblocks=4_nc_in=27_nc_lift=32_nc_mid=128_nc_out=21_nd=256_nt=1_ntrain=500_nvalid=10_nx=256_ny=128_nz=1_p=1.jld2"
+dataset_path = "/pscratch/sd/r/richardr/FNO_CIG/FNO-CIG/results/concatenated_data_99_quantile_scaled.jld2"
+
+file = filename # projectdir("weights", "DFNO_3D", filename)
+file = load(file)
+
+nblocks = file["nblocks"]
+mx = file["mx"] 
+mz = file["my"]
+mh = file["mz"] 
+nd = file["nd"] 
+nc_lift = file["nc_lift"]
+ntrain = file["ntrain"]
+# nvalid = file["nvalid"]
+
+partition = [1,pe_count]
+params = Config.get_parameters()
+
+down_rate = params["down_rate"]
+nx = params["nx"] ÷ down_rate
+nz = params["nz"] ÷ down_rate
+offsets = params["n_offsets"]
+
+use_nz = Bool(offset_is_nz)
+labels = @strdict use_nz offsets
+
+nc_in = use_nz ? 7 : offsets + 1 + 1 + 4 # offsets + 2 velocity models + indices
+nc_out = use_nz ? 1 : offsets
+
+nh = use_nz ? offsets : 1
+mh = use_nz ? mh : 1 
+
+@info "Initializing model..."
+@assert MPI.Comm_size(comm) == prod(partition)
+
+function read_velocity_cigs_offsets_as_nc_old(path::String, modelConfig::DFNO_3D.ModelConfig; ntrain::Int, nvalid::Int)
+
+    params = Config.get_parameters()
+    offset_start = params["read_offset_start"]
+    offsets = params["n_offsets"]
+    total = params["n_total"]
+
+    # Assumption that x is (nx, nz, 1, n). x0 is (nx, nz). CIG0 is (nh, nx, nz). CIG is (nh, nx, nz, 1, n)
+    function read_x_tensor_helper(file_name, key, indices; flip=false, init_index=1)
+        data = nothing
+        h5open(file_name, "r") do file
+            x_data = file[key[1]]
+            cigs_data = file[key[2]]
+
+            # Read proper indices of x and x0. NOTE: Disclude 3 because no z = 1 for background
+            x = x_data[indices[1], indices[2], 1, indices[4]]
+            x0 = x_data[indices[1], indices[2], 1, init_index]
+            cig0 = cigs_data[offset_start:offset_start+offsets-1, indices[1], indices[2], 1, init_index]
+
+            # Reshape to prepare for augmentation
+            x = reshape(x, :, map(range -> length(range), indices[1:4])...)
+            x0 = reshape(x0, :, map(range -> length(range), indices[1:3])..., 1)
+            cig0 = reshape(cig0, :, map(range -> length(range), indices[1:3])..., 1)
+
+            x0 = repeat(x0, outer=[1, 1, 1, 1, length(indices[4])])
+            cig0 = repeat(cig0, outer=[1, 1, 1, 1, length(indices[4])])
+
+            # Concat along dimension 1
+            data = cat(x, x0, cig0, dims=1)
+            flip && (data = reverse(data, dims=1))
+        end
+
+        # data_channels * nx * ny * nz * nt * n = data_channels * nx * nz * nh * 1 * n
+        return data
+    end
+
+    function read_x_tensor(file_name, key, indices)
+        requested = length(indices[4])
+        @assert (requested - nvalid) % (total - nvalid) == 0
+
+        nrounds = (requested - nvalid) ÷ (total - nvalid)
+        augmented_data = []
+
+        samples_per_round = (requested - nvalid) ÷ nrounds
+
+        for round in 1:nrounds
+            Random.seed!(round % 2)
+            
+            init_index = rand(1:samples_per_round)
+            new_indices = [indices[1:3]..., 1:samples_per_round]
+    
+            data = read_x_tensor_helper(file_name, key, new_indices, flip=(round <= nrounds ÷ 2), init_index=init_index)
+
+            if round == 1
+                augmented_data = data
+            else
+                augmented_data = cat(augmented_data, data, dims=ndims(augmented_data))
+            end
+        end
+
+        Random.seed!(1 % 2) # Simulate round 1
+        init_index = rand(1:samples_per_round+nvalid)
+
+        new_indices = [indices[1:3]..., 1:samples_per_round+nvalid]
+        validation = read_x_tensor_helper(file_name, key, new_indices, flip=false, init_index=init_index)
+
+        augmented_data = cat(augmented_data, validation[:, :, :, :, end-nvalid+1:end], dims=ndims(augmented_data))
+        return augmented_data
+    end
+    
+    function read_y_tensor_helper(file_name, key, indices; flip=false)
+        data = nothing
+        h5open(file_name, "r") do file
+            cigs_data = file[key]
+            data = cigs_data[offset_start:offset_start+offsets-1, indices[1], indices[2], 1, indices[5]] # first dim is offsets as channel, dim 4 which is t = 1:1
+        end
+
+        # channels * nx * ny * nz * nt * n = channels * nx * nz * nh * 1 * n
+        data = reshape(data, :, map(range -> length(range), indices[1:5])...)
+        flip && (data = reverse(data, dims=2))
+        return data
+    end
+
+    function read_y_tensor(file_name, key, indices)
+        requested = length(indices[5])
+        @assert (requested - nvalid) % (total - nvalid) == 0
+
+        nrounds = (requested - nvalid) ÷ (total - nvalid)
+        augmented_data = []
+
+        samples_per_round = (requested - nvalid) ÷ nrounds
+
+        for round in 1:nrounds
+            Random.seed!(round % 2)
+            
+            init_index = rand(1:samples_per_round)
+
+            new_indices = [indices[1:4]..., 1:samples_per_round]
+            data = read_y_tensor_helper(file_name, key, new_indices, flip=(round <= nrounds ÷ 2))
+
+            if round == 1
+                augmented_data = data
+            else
+                augmented_data = cat(augmented_data, data, dims=ndims(augmented_data))
+            end
+        end
+
+        Random.seed!(1 % 2) # Simulate round 1
+        init_index = rand(1:samples_per_round+nvalid)
+
+        new_indices = [indices[1:4]..., 1:samples_per_round+nvalid]
+        validation = read_y_tensor_helper(file_name, key, new_indices, flip=false)
+
+        augmented_data = cat(augmented_data, validation[:, :, :, :, :, end-nvalid+1:end], dims=ndims(augmented_data))
+        return augmented_data
+    end
+
+    dataConfig = DFNO_3D.DataConfig(modelConfig=modelConfig, 
+                                    ntrain=ntrain, 
+                                    nvalid=nvalid, 
+                                    x_file=path,
+                                    y_file=path,
+                                    x_key=["xs", "cigs"],
+                                    y_key="cigs")
+
+    return DFNO_3D.loadDistData(dataConfig, dist_read_x_tensor=read_x_tensor_helper, dist_read_y_tensor=read_y_tensor_helper)
+end
+
+modelConfig = DFNO_3D.ModelConfig(nc_in=nc_in, nc_lift=nc_lift, nc_out=nc_out, nx=nx, ny=nz, nz=nh, nt=1, mx=mx, my=mz, mz=mh, mt=1, nblocks=nblocks, partition=partition, dtype=Float32, relu01=false)
+_, _, x_valid, y_valid = read_velocity_cigs_offsets_as_nc_old(dataset_path, modelConfig, ntrain=ntrain, nvalid=nvalid)
+
+@info "Loaded data..."
+
+model = DFNO_3D.Model(modelConfig)
+θ = DFNO_3D.initModel(model)
+
+DFNO_3D.loadWeights!(θ, filename, "θ_save", partition)
+labels = @strdict nx nz nblocks mx mz mh ntrain nvalid nc_lift offset_is_nz
+
+y_predict = DFNO_3D.forward(model, θ, x_valid) |> cpu
+plot_cig_eval(modelConfig, x_valid, y_valid, y_predict, use_nz=use_nz, additional=labels)
+
+MPI.Finalize()
diff --git a/Manifest.toml b/Manifest.toml
new file mode 100644
index 0000000..c67db91
--- /dev/null
+++ b/Manifest.toml
@@ -0,0 +1,181 @@
+# This file is machine-generated - editing it directly is not advised
+
+julia_version = "1.8.5"
+manifest_format = "2.0"
+project_hash = "bcbf45ca7bd5851cab3a7af7a3ad54d9a84e7fbe"
+
+[[deps.ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+version = "1.1.1"
+
+[[deps.Artifacts]]
+uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+
+[[deps.Base64]]
+uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+
+[[deps.Dates]]
+deps = ["Printf"]
+uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
+
+[[deps.Downloads]]
+deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+version = "1.6.0"
+
+[[deps.FileIO]]
+deps = ["Pkg", "Requires", "UUIDs"]
+git-tree-sha1 = "62ca0547a14c57e98154423419d8a342dca75ca9"
+uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+version = "1.16.4"
+
+[[deps.FileWatching]]
+uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
+
+[[deps.InteractiveUtils]]
+deps = ["Markdown"]
+uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+
+[[deps.JLD2]]
+deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"]
+git-tree-sha1 = "783c1be5213a09609b23237a0c9e5dfd258ae6f2"
+uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+version = "0.5.7"
+
+[[deps.LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+version = "0.6.3"
+
+[[deps.LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
+version = "7.84.0+0"
+
+[[deps.LibGit2]]
+deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
+uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
+
+[[deps.LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+version = "1.10.2+0"
+
+[[deps.Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+
+[[deps.Logging]]
+uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+
+[[deps.MacroTools]]
+deps = ["Markdown", "Random"]
+git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df"
+uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.5.13"
+
+[[deps.Markdown]]
+deps = ["Base64"]
+uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
+
+[[deps.MbedTLS_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
+version = "2.28.0+0"
+
+[[deps.Mmap]]
+uuid = "a63ad114-7e13-5084-954f-fe012c677804"
+
+[[deps.MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+version = "2022.2.1"
+
+[[deps.NetworkOptions]]
+uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
+version = "1.2.0"
+
+[[deps.OrderedCollections]]
+git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5"
+uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+version = "1.6.3"
+
+[[deps.Pkg]]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+version = "1.8.0"
+
+[[deps.PrecompileTools]]
+deps = ["Preferences"]
+git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f"
+uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+version = "1.2.1"
+
+[[deps.Preferences]]
+deps = ["TOML"]
+git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6"
+uuid = "21216c6a-2e73-6563-6e65-726566657250"
+version = "1.4.3"
+
+[[deps.Printf]]
+deps = ["Unicode"]
+uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[[deps.REPL]]
+deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
+uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+
+[[deps.Random]]
+deps = ["SHA", "Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[deps.Requires]]
+deps = ["UUIDs"]
+git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
+uuid = "ae029012-a4dd-5104-9daa-d747884805df"
+version = "1.3.0"
+
+[[deps.SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+version = "0.7.0"
+
+[[deps.Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[deps.Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+
+[[deps.TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+version = "1.0.0"
+
+[[deps.Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+version = "1.10.1"
+
+[[deps.TranscodingStreams]]
+git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
+uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+version = "0.11.3"
+
+[[deps.UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
+[[deps.Unicode]]
+uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[deps.Zlib_jll]]
+deps = ["Libdl"]
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+version = "1.2.12+3"
+
+[[deps.nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+version = "1.48.0+0"
+
+[[deps.p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
+version = "17.4.0+0"
diff --git a/Project.toml b/Project.toml
new file mode 100644
index 0000000..40c0c8e
--- /dev/null
+++ b/Project.toml
@@ -0,0 +1,2 @@
+[deps]
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..ca90cb1
--- /dev/null
+++ b/README.md
@@ -0,0 +1,83 @@
+## TODO:
+
+1. Make Inference Loop
+2. Clean Plotting Code
+3. Run everything
+
+## Training Routine
+
+1. Make `initial.jld2` and `posteriors_iteration_j=0.jld2`
+```
+julia --project=CIG/ scripts/prepare_training.jl
+```
+
+For each j:
+
+2. Run generate CIGs for j -> `cigs_iteration_j=j.jld2`
+
+```
+sh scripts/launch_cigs.sh j nsamples nprocs
+```
+
+Collect CIGs when batch jobs finish
+```
+salloc --nodes=1 --constraint=gpu --gpus=1 --qos=interactive --time=04:00:00 --account=m3863_g --ntasks=1 --gpus-per-task=1 --gpu-bind=none
+srun julia --project=CIG/ CIG/collect_cigs.jl j nsamples
+```
+
+3. Train CNF and FNO parallely
+
+```
+sh scripts/train_fno.sh ntrain j epochs
+sh scripts/train_cnf.sh ntrain j epochs
+```
+
+Clean Directories # TODO: Fix inside `ParametricDFNOs.jl`
+
+```
+sh scripts/clean_dir.sh j
+```
+
+4. Run Posterior mean function -> `posteriors_iteration_j=j.jld2`
+
+```
+salloc --nodes=1 --constraint=gpu --gpus=1 --qos=interactive --time=04:00:00 --account=m3863_g --ntasks=1 --gpus-per-task=1 --gpu-bind=none
+srun julia --project=CNF/ CNF/update_fiducials.jl j nsamples ntrain epochs
+```
+
+## Inference Routine
+
+1. Make `initial.jld2` and `posteriors_iteration_j=0.jld2` (Inside test data directory)
+
+```
+sh scripts/preprare_for_testing.sh nsamples 1 (offset, number of samples to test)
+```
+
+2. Run Inference using Trained Networks
+
+```
+srun julia --project=. inference.jl 0 (1/0 dont-use/use FNO)
+```
+
+Internally, For each j:
+- Run generate CIGs for j -> `cigs_iteration_j=j.jld2` (Using Trained FNO / Using True Simulator)
+- Run Posterior mean function -> `posteriors_iteration_j=j.jld2`
+- Plot refined posterior and reference for each iteration
+
+
+## Example run:
+
+```
+julia --project=CIG/ scripts/prepare_training.jl
+sh scripts/launch_cigs.sh 1 850 34
+
+salloc --nodes=1 --constraint=gpu --gpus=1 --qos=interactive --time=04:00:00 --account=m3863_g --ntasks=1 --gpus-per-task=1 --gpu-bind=none
+srun julia --project=CIG/ CIG/collect_cigs.jl 1 850
+
+
+srun julia --project=CNF/ CNF/train.jl 2 2 0 2 2
+sh scripts/train_fno.sh 800 1 80 4
+sh scripts/train_cnf.sh 800 1 200
+
+julia --project=CNF/ CNF/update_fiducials.jl 1 850 800 165
+```
diff --git a/cnf.jl b/cnf.jl
deleted file mode 100644
index 3f84249..0000000
--- a/cnf.jl
+++ /dev/null
@@ -1 +0,0 @@
-# TODO: RICHARD
diff --git a/config.jl b/config.jl
new file mode 100644
index 0000000..1b239f7
--- /dev/null
+++ b/config.jl
@@ -0,0 +1,32 @@
+module Config
+
+export get_parameters
+
+function get_parameters()
+    return Dict(
+        "nx" => 512,
+        "nz" => 256,
+        "n_offsets" => 21,
+        "down_rate" => 2,
+        "read_offset_start" => 1,
+        "offset_start" => -500.0,
+        "offset_end" => 500.0,
+        "f0" => 0.015f0,
+        "timeD" => 3200f0,
+        "timeR" => 3200f0,
+        "TD" => 3200f0,
+        "dtD" => 4f0,
+        "dtS" => 4f0,
+        "nbl" => 120,
+        "d" => (12.5f0, 12.5f0),
+        "o" => (0f0, 0f0),
+        # "n" => (size(m_train, 1), size(m_train, 2)),
+        "nsrc" => 64,
+        "n_samples" => 64,
+        # "nxrec" => size(m_train, 1),
+        "snr" => 12f0,
+        "n_total" => 512
+    )
+end
+
+end
diff --git a/fno.jl b/fno.jl
deleted file mode 100644
index 3f84249..0000000
--- a/fno.jl
+++ /dev/null
@@ -1 +0,0 @@
-# TODO: RICHARD
diff --git a/inference.jl b/inference.jl
index d8da063..67da070 100644
--- a/inference.jl
+++ b/inference.jl
@@ -1,8 +1,77 @@
 # Inference ASPIRE
 
-# TODO: Init yobs and x0
+using JLD2
+using ParametricDFNOs.DFNO_3D
+using DrWatson
+using MPI
+using CUDA
+
+include("config.jl")
+
+using .Config
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+pe_count = MPI.Comm_size(comm)
+
+CUDA.device!(rank % 4)
+partition = [1, pe_count]
+
+nc_lift = 32
+nblocks = 4
+mx, mz, mh = 36, 36, 1
+
+params = Config.get_parameters()
+
+nh = 1
+nx = params["nx"]
+nz = params["nz"]
+offsets = params["n_offsets"]
+
+nc_in = offsets + 1 + 1 + 4 # offsets + 2 velocity models + indices
+nc_out = offsets
+
+offset, use_fno = parse.(Int, ARGS[1:2])
+
+# TODO: Make test Files, so its easy for FNO to read
+@load "data/test/initial.jld2" x y
+@load "data/test/cigs_iteration_j=0.jld2" CIGs
+@load "data/test/posteriors_iteration_j=0.jld2" x0 CIG0
+
+y_obs = y[offset + 1]
+x0 = x0[offset + 1]
+x = x[:, :, 1:1, ]
+
+function migrate(sim::TrueSimulator, x_path, y_path)
+
+end
+
+function migrate(sim::FNO_Simulator, x_path, y_path)
+    # TODO: MOve to FNO_Simulator
+    @assert MPI.Comm_size(comm) == prod(partition)
+    modelConfig = DFNO_3D.ModelConfig(nc_in=nc_in, nc_lift=nc_lift, nc_out=nc_out, nx=nx, ny=nz, nz=nh, nt=1, mx=mx, my=mz, mz=mh, mt=1, nblocks=nblocks, partition=partition, dtype=Float32, relu01=false)
+
+    model = DFNO_3D.Model(modelConfig)
+    θ = DFNO_3D.initModel(model)
+
+end
 
 for j = 1:J
+    x_path = "data/test/posteriors_iteration_j=$(j-1).jld2"
+    y_path = "data/test/cigs_iteration_j=$j.jld2"
+    _, _, x_valid, _ = read_velocity_cigs_offsets_as_nc(x_path, y_path, modelConfig, ntrain=offset, nvalid=nvalid)
+
+    filename = "mt=25_mx=10_my=10_mz=10_nblocks=20_nc_in=5_nc_lift=20_nc_mid=128_nc_out=1_nd=20_nt=51_nx=20_ny=20_nz=20_p=8.jld2"
+    DFNO_3D.loadWeights!(θ, filename, "θ_save", partition)
+
+    x_valid = reshape(x_valid, nc_in, nx, nz, :)
+    yhat = migrate(simulator, x_path, y_path)
+
+    post = CNF-1(z; yhat)
+
+
     # TODO: Calculate summary statistic yhat for yobs around fiducials x0
 
     # TODO: Update fiducials x0 by averaging posterior samples conditioned on yhat
diff --git a/plot.jl b/plot.jl
new file mode 100644
index 0000000..977e64c
--- /dev/null
+++ b/plot.jl
@@ -0,0 +1,193 @@
+function plot_metrics(x_gt, x_hat_fno_data, x_hat_true_data; n, d, title="")
+    """
+    Compare the posterior means of each iteration (both using FNO CIG and analytical CIG) with the ground truth velocity
+    Assume `x_hat_fno_data` and `x_hat_true_dat` have shape (nx, nz, n_iter)
+    """
+    n_iter = size(x_hat_fno_data, 3)
+    extentfull = (0f0, (n[1]-1)*d[1], (n[end]-1)*d[end], 0f0)
+
+    fig, axes = subplots(2, n_iter + 1, figsize=((n_iter + 1)*9, 12),)
+    
+    # plot ground truth in the last column
+    for i in 1:2
+        ax = axes[i, n_iter + 1]
+        vmin = quantile(vec(x_gt), 0.05)
+        vmax = quantile(vec(x_gt), 0.95)
+        cax = ax.imshow(x_gt', vmin=vmin, vmax=vmax, cmap="cet_rainbow4", extent=extentfull, aspect=0.45 * (extentfull[2] - extentfull[1]) / (extentfull[3] - extentfull[4]))
+        ax.set_xlabel("X [m]", fontsize=14)
+        ax.set_ylabel("Z [m]", fontsize=14)
+        ax.set_title("Ground Truth", fontsize=16)
+    end
+
+    # plot x_hat_fno_data in the first row
+    for j in 1:n_iter
+        ax = axes[1, j]
+        x = x_hat_fno_data[:, :, j]
+        vmin = quantile(vec(x), 0.05)
+        vmax = quantile(vec(x), 0.95)
+        cax = ax.imshow(x', vmin=vmin, vmax=vmax, cmap="cet_rainbow4", extent=extentfull, aspect=0.45 * (extentfull[2] - extentfull[1]) / (extentfull[3] - extentfull[4]))
+        ax.set_xlabel("X [m]", fontsize=14)
+        ax.set_ylabel("Z [m]", fontsize=14)
+        ax.set_title("Use FNO CIG Iter $j", fontsize=16)
+    end
+
+    # plot x_hat_true_data in the second row
+    for j in 1:n_iter
+        ax = axes[2, j]
+        x = x_hat_true_data[:, :, j]
+        vmin = quantile(vec(x), 0.05)
+        vmax = quantile(vec(x), 0.95)
+        cax = ax.imshow(x', vmin=vmin, vmax=vmax, cmap="cet_rainbow4", extent=extentfull, aspect=0.45 * (extentfull[2] - extentfull[1]) / (extentfull[3] - extentfull[4]))
+        ax.set_xlabel("X [m]", fontsize=14)
+        ax.set_ylabel("Z [m]", fontsize=14)
+        ax.set_title("Use analytical CIG Iter $j", fontsize=16)
+    end
+
+    # Add a main title for the figure
+    fig.suptitle(title, fontsize=20)
+    fig.tight_layout()
+    savefig("plots/result_inference.png", bbox_inches="tight", dpi=300)
+    close(fig)
+end
+
+
+function plot_cig_metric(cig_fno, cig_analytic, plot_path, params; fs=40)
+    """
+    A similar function as `plot_metrics`. This function compares FNO CIG and analytical CIG across multiple iterations.
+    The 3rd row is the 5X difference between FNO CIG and analytical CIG.
+    Assume `cig_fno`, `cig_analytic` both have shape (nx, nz, n_offset, n_iter)
+    """
+
+    n_iter = size(cig_fno, 4)  # Get the number of iterations
+
+    fig2, ax2 = subplots(3, n_iter, figsize=(15 * n_iter, 25))   # need to name it fig2, ax2. Otherwise fig, ax will mix with those defined in plot_cig
+
+    for i in 1:n_iter
+        cig_fno_i = cig_fno[:, :, :, i]
+        cig_analytic_i = cig_analytic[:, :, :, i]
+        cig_dif_i = abs.(cig_fno_i - cig_analytic_i) .* 5
+    
+        # plot FNO CIG
+        cig_img_fname = "cig_fno_iter_$i.png"
+        plot_cig(cig_fno_i, plot_path, cig_img_fname, params)
+        img1 = imread(joinpath(plot_path, cig_img_fname))
+        ax2[1, i].imshow(img1)
+        ax2[1, i].axis("off")  # Hide axes for a cleaner look
+        ax2[1, i].set_title("FNO predicted CIG - Iter $i", fontsize=fs)
+    
+        # plot analytical CIG
+        cig_img_fname = "cig_analytic_iter_$i.png"
+        plot_cig(cig_analytic_i, plot_path, cig_img_fname, params)
+        img2 = imread(joinpath(plot_path, cig_img_fname))
+        ax2[2, i].imshow(img2)
+        ax2[2, i].axis("off")
+        ax2[2, i].set_title("Analytical CIG - Iter $i", fontsize=fs)
+    
+        # plot (FNO CIG - analytical CIG) X 5
+        cig_img_fname = "cig_dif_iter_$i.png"
+        plot_cig(cig_dif_i, plot_path, cig_img_fname, params)
+        img3 = imread(joinpath(plot_path, cig_img_fname))
+        ax2[3, i].imshow(img3)
+        ax2[3, i].axis("off")
+        ax2[3, i].set_title("5X Difference - Iter $i", fontsize=fs)
+    end
+
+    fig2.tight_layout()
+
+    fig2.savefig(joinpath(plot_path, "cig_compare_all_iters.png"))
+end
+
+
+using PyCall
+function make_movie(X_post, plot_path; movie_suffix="X_post")
+    """
+    Make a movie of posterior samples
+    size(X_post) = (nx, nz, 1, n_post) 
+    """
+    down_rate = params["down_rate"]
+    d = params["d"] .* down_rate
+
+    nx = params["nx"] ÷ down_rate
+    nz = params["nz"] ÷ down_rate
+
+    n = (nx, nz)
+
+    # Update function to change the image per frame
+    function update_frame(frame)
+        frame_i = X_post[:, :, 1, frame]
+        img.set_data(frame_i')  # Update the image data
+        ax.set_title("Frame $frame")  # Update title with frame number
+        return img
+    end
+
+    fig, ax = plt.subplots()
+
+    # Initial plot (first image)
+    frame_1 = X_post[:, :, 1, 1]
+    # img = ax.imshow(frame_1', cmap="cet_rainbow4", vmin=minimum(m_train), vmax=maximum(m_train), extent=[0, size(X_post, 1) * d[1], size(X_post, 2) * d[2], 0])
+    img = ax.imshow(frame_1', cmap="cet_rainbow4", vmax=4.5, extent=[0, size(X_post, 1) * d[1], size(X_post, 2) * d[2], 0])
+    ax.set_title("Frame 1")
+    ax.set_xlabel("X [m]")
+    ax.set_ylabel("Z [m]")
+
+    # Create the animation
+    animation = pyimport("matplotlib.animation")
+    n_frames = size(X_post)[end]
+    anim = animation.FuncAnimation(fig, update_frame, frames=1:n_frames, interval=200)
+
+    # Save as a video or GIF
+    anim.save(joinpath(plot_path, "movie_$(movie_suffix).gif"), writer="ffmpeg")
+end
+
+
+function make_movie_X_post_w_CIG(X_post, plot_path, cig; movie_suffix="X_post")
+    """
+    Make a movie of posterior samples with a static image on the left.
+    size(X_post) = (nx, nz, 1, n_post) 
+    size(cig) = (nx, nz, n_offset)
+    """
+    down_rate = params["down_rate"]
+    d = params["d"] .* down_rate
+
+    nx = params["nx"] ÷ down_rate
+    nz = params["nz"] ÷ down_rate
+
+    n = (nx, nz)
+
+    # Update function to change the image per frame
+    function update_frame(frame)
+        frame_i = X_post[:, :, 1, frame]
+        img.set_data(frame_i')  # Update the image data on the right
+        ax[2].set_title("Frame $frame")  # Update title with frame number
+        return img
+    end
+
+    fig, ax = plt.subplots(1, 2, figsize=(12, 6))
+
+    # Load and plot CIG
+    cig_img_fname = "cig.png"
+    plot_cig(cig, plot_path, cig_img_fname, params)
+    img = imread(joinpath(plot_path, cig_img_fname))
+
+    ax[1].imshow(img)
+    ax[1].axis("off")  # Hide axes for a cleaner look
+    ax[1].set_title("input CIG")
+
+    # Plot the first frame of the movie
+    frame_1 = X_post[:, :, 1, 1]
+    img = ax[2].imshow(frame_1', cmap="cet_rainbow4", vmax=4.5, extent=[0, size(X_post, 1) * d[1], size(X_post, 2) * d[2], 0])
+    ax[2].set_title("Frame 1")
+    ax[2].set_xlabel("X [m]")
+    ax[2].set_ylabel("Z [m]")
+
+    # Create the animation
+    animation = pyimport("matplotlib.animation")
+    n_frames = size(X_post)[end]
+    anim = animation.FuncAnimation(fig, update_frame, frames=1:n_frames, interval=200)
+
+    fig.tight_layout()
+
+    # Save as a video or GIF
+    anim.save(joinpath(plot_path, "movie_$(movie_suffix).gif"), writer="ffmpeg")
+end
+
diff --git a/train.jl b/psedocode.jl
similarity index 100%
rename from train.jl
rename to psedocode.jl
diff --git a/scripts/clean_dir.sh b/scripts/clean_dir.sh
new file mode 100644
index 0000000..4818bfc
--- /dev/null
+++ b/scripts/clean_dir.sh
@@ -0,0 +1,10 @@
+# Move contents of weights/DFNO_3D to weights/$1, then recreate weights/DFNO_3D
+mv weights/DFNO_3D weights/"$1"
+mkdir weights/DFNO_3D
+
+# Move contents of plots/DFNO_3D to plots/$1, then recreate plots/DFNO_3D
+mv plots/DFNO_3D plots/"$1"
+mkdir plots/DFNO_3D
+
+# Move specific contents from FNO/plots/DFNO_3D to plots/$1/DFNO_3D, then recreate original folder
+mv FNO/plots/DFNO_3D/* plots/"$1"/DFNO_3D
diff --git a/scripts/gen_random.jl b/scripts/gen_random.jl
new file mode 100644
index 0000000..9657d73
--- /dev/null
+++ b/scripts/gen_random.jl
@@ -0,0 +1,45 @@
+using JLD2
+
+T = Float64
+J = 3
+
+nz = 128
+nx = 256
+N = 800
+
+offsets = 21
+
+x = rand(T, nx, nz, 1, N)
+y = rand(T, nx, nz, N)
+
+@JLD2.save "data/initial.jld2" x y
+
+for j in 0:J
+    x0 = rand(T, nx, nz, 1, N)
+    CIGs = rand(T, offsets, nx, nz, 1, N)
+    CIG0 = CIGs[:, :, :, 1:1, 1:1]
+
+    @JLD2.save "data/posteriors_iteration_j=$j.jld2" x0 CIG0
+    if j > 0
+        @JLD2.save "data/cigs_iteration_j=$j.jld2" CIGs
+    end
+end
+
+# BEFORE ANYTHING
+
+# posteriors_iteration_0.jld2
+x0 -> (nx, nz, N)
+
+# initial.jld2 # make one example save to initial.jld2. size (1, nx, nz)
+x  -> (nx, nz, N)
+y  -> (??????, N)
+
+# CIG SUMMARY STATISTIC SIMULATION
+
+# cigs_iteration_j=$j.jld2 # save in a separate file 
+CIGs -> (nx, nz, offsets, N)
+
+# CNF POST TRAINING POSTERIORS
+
+# posteriors_iteration_1.jld2
+x0 -> (nx, nz, N)
diff --git a/scripts/launch_cigs.sh b/scripts/launch_cigs.sh
new file mode 100644
index 0000000..34e1cad
--- /dev/null
+++ b/scripts/launch_cigs.sh
@@ -0,0 +1,48 @@
+#!/bin/bash
+
+# Check if three arguments (j, N and P) are provided
+if [ "$#" -ne 3 ]; then
+    echo "Usage: $0 <j> <N> <P>"
+    exit 1
+fi
+
+j=$1      # iteration index
+N=$2      # number of samples
+P=$3      # number of processors
+
+for (( i=0; i<P; i++ ))
+do
+    sbatch <<EOT
+#!/bin/bash
+
+#SBATCH --nodes=1
+#SBATCH --constraint=gpu
+#SBATCH --gpus=1
+#SBATCH --ntasks=1
+#SBATCH --gpus-per-task=1
+#SBATCH --gpu-bind=none
+#SBATCH --qos=regular
+#SBATCH --job-name=GEN_CIG_proc=${i}_iteration_j=${j}
+#SBATCH --mail-user=richardr2926@gmail.com
+#SBATCH --mail-type=ALL
+#SBATCH --time=01:10:00
+#SBATCH --account=m3863_g
+
+nvidia-smi
+export SLURM_CPU_BIND="cores"
+export PATH=\$PATH:\$HOME/.julia/bin
+export DFNO_3D_GPU=1
+export LD_PRELOAD=/opt/cray/pe/lib64/libmpi_gtl_cuda.so.0
+
+export JULIA_NUM_THREADS=$SLURM_CPUS_ON_NODE
+export DEVITO_LANGUAGE=openacc
+export DEVITO_ARCH=nvc
+export DEVITO_PLATFORM=nvidiaX
+
+module load cudnn/8.9.3_cuda12 julia/1.8.5
+
+srun julia --project=CIG/ CIG/simulate.jl ${j} ${N} ${i} ${P}
+
+exit 0
+EOT
+done
diff --git a/scripts/plot_prior.jl b/scripts/plot_prior.jl
new file mode 100644
index 0000000..1f75d81
--- /dev/null
+++ b/scripts/plot_prior.jl
@@ -0,0 +1,38 @@
+# To Plot a Prior
+using DrWatson
+using JLD2
+using PyPlot, SlimPlotting
+using Images
+
+include("CNF/utils.jl")
+include("config.jl")
+
+params = Config.get_parameters()
+down_rate = params["down_rate"]
+
+function plot_vel(x, filename; smooth=false, gaussian=20)
+    if smooth
+        x = x_smooth = 1f0./Float32.(imfilter(1f0./x[:, :, 1, 1], Kernel.gaussian(gaussian)))
+    end
+
+    imshow(x, vmin=1.5, vmax=4.8,  interpolation="none", cmap="cet_rainbow4")
+    savefig(filename, bbox_inches="tight", dpi=300)
+
+    close()
+end
+
+@load "data/posteriors_iteration_j=1.jld2" x0
+x = x0
+
+for i in 1:10
+    plot_vel(x0[:, :, 1, i]', "plots/posterior_1_$(i)_sample.png")
+end
+
+# @load "data/synthoseis_no_salt.jld2" data
+# x = data[1:down_rate:end, 1:down_rate:end, :, :]
+
+# # @load "data/initial.jld2" x
+
+# plot_vel(x[:, :, 1, 1]', "plots/no_salt.png")
+# plot_vel(x0[:, :, 1, 1]', "plots/no_salt_smoothed.png")
+# plot_vel(x[:, :, 1, 1]', "plots/no_salt_smoothed_true.png", smooth=true)
diff --git a/scripts/prepare_training.jl b/scripts/prepare_training.jl
new file mode 100644
index 0000000..ae085f1
--- /dev/null
+++ b/scripts/prepare_training.jl
@@ -0,0 +1,24 @@
+using JLD2
+using Images
+
+include("../config.jl")
+
+params = Config.get_parameters()
+down_rate = params["down_rate"]
+
+# First set of posteriors/fiducials/background-models
+@load "data/synthoseis_no_salt.jld2" data;
+x0 = data[1:down_rate:end, 1:down_rate:end, :, :]
+println(size(x0))
+
+for sample in 1:size(x0, 4)
+    x0[:, :, 1, sample] = 1f0./Float32.(imfilter(1f0./x0[:, :, 1, sample], Kernel.gaussian(20)))
+end
+
+JLD2.@save "data/posteriors_iteration_j=0.jld2" x0;
+
+# Ground truth
+@load "data/synthoseis_salt.jld2" data;
+x = data[1:down_rate:end, 1:down_rate:end, :, :]
+println(size(x))
+JLD2.@save "data/initial.jld2" x;
diff --git a/setup.sh b/scripts/setup.sh
similarity index 62%
rename from setup.sh
rename to scripts/setup.sh
index de427df..3576dfa 100644
--- a/setup.sh
+++ b/scripts/setup.sh
@@ -2,6 +2,13 @@ nvidia-smi
 export SLURM_CPU_BIND="cores"
 export PATH=$PATH:$HOME/.julia/bin
 export DFNO_3D_GPU=1
-# export LD_LIBRARY_PATH=
 export LD_PRELOAD=/opt/cray/pe/lib64/libmpi_gtl_cuda.so.0
+
+export JULIA_NUM_THREADS=$SLURM_CPUS_ON_NODE
+export DEVITO_LANGUAGE=openacc
+export DEVITO_ARCH=nvc
+export DEVITO_PLATFORM=nvidiaX
+
+# export LD_LIBRARY_PATH=
+
 module load cudnn/8.9.3_cuda12 julia/1.8.5
diff --git a/scripts/train_cnf.sh b/scripts/train_cnf.sh
new file mode 100644
index 0000000..2ddabc8
--- /dev/null
+++ b/scripts/train_cnf.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+sbatch <<EOT
+#!/bin/bash
+
+#SBATCH --nodes=1
+#SBATCH --constraint=gpu
+#SBATCH --gpus=1
+#SBATCH --ntasks=1
+#SBATCH --gpus-per-task=1
+#SBATCH --gpu-bind=none
+#SBATCH --qos=regular
+#SBATCH --job-name=CNF_ntrain=${1}_iteration_j=${2}_epochs=${3}
+#SBATCH --mail-user=richardr2926@gmail.com
+#SBATCH --mail-type=ALL
+#SBATCH --time=12:00:00
+#SBATCH --account=m3863_g
+
+nvidia-smi
+export SLURM_CPU_BIND="cores"
+export PATH=$PATH:$HOME/.julia/bin
+export DFNO_3D_GPU=1
+# export LD_LIBRARY_PATH=
+export LD_PRELOAD=/opt/cray/pe/lib64/libmpi_gtl_cuda.so.0
+module load cudnn/8.9.3_cuda12 julia/1.8.5
+
+srun julia --project=CNF/ CNF/train.jl 8 ${3} ${2} ${1} 50
+
+exit 0
+EOT
diff --git a/scripts/train_fno.sh b/scripts/train_fno.sh
new file mode 100644
index 0000000..304d6b8
--- /dev/null
+++ b/scripts/train_fno.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+sbatch <<EOT
+#!/bin/bash
+
+#SBATCH --nodes=1
+#SBATCH --constraint=gpu
+#SBATCH --gpus=${4}
+#SBATCH --ntasks=${4}
+#SBATCH --gpus-per-task=1
+#SBATCH --gpu-bind=none
+#SBATCH --qos=regular
+#SBATCH --job-name=FNO_ntrain=${1}_iteration_j=${2}_epochs=${3}_gpus=${4}
+#SBATCH --mail-user=richardr2926@gmail.com
+#SBATCH --mail-type=ALL
+#SBATCH --time=12:00:00
+#SBATCH --account=m3863_g
+
+nvidia-smi
+export SLURM_CPU_BIND="cores"
+export PATH=$PATH:$HOME/.julia/bin
+export DFNO_3D_GPU=1
+# export LD_LIBRARY_PATH=
+export LD_PRELOAD=/opt/cray/pe/lib64/libmpi_gtl_cuda.so.0
+module load cudnn/8.9.3_cuda12 julia/1.8.5
+
+/global/homes/r/richardr/.julia/bin/mpiexecjl julia --project=FNO/ FNO/train.jl 1 ${3} ${2} ${1} 2
+
+exit 0
+EOT
diff --git a/test.jl b/test.jl
deleted file mode 100644
index e69de29..0000000