Make Enzyme the only compatible AD engine

.github/workflows/ci.yml

@@ -18,7 +18,7 @@ jobs:
         group:
           - Core 
         version:
-          - '1.9'
+          - '1.10'
           - '1'
         os:
           - ubuntu-latest

.gitignore

@@ -26,3 +26,4 @@ lib/ComradePigeons/Manifest.toml
 imaging_pol.md
 .CondaPkg
 test/CMAES.bson
+examples/intermediate/PolarizedImaging/m87polarized.uvfits

Project.toml

@@ -1,7 +1,7 @@
 name = "Comrade"
 uuid = "99d987ce-9a1e-4df8-bc0b-1ea019aa547b"
 authors = ["Paul Tiede <[email protected]>"]
-version = "0.10.5"
+version = "0.11.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -15,7 +15,7 @@ DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HypercubeTransform = "9ec9aee3-0fd3-44c2-8e61-a50acc66f3c8"
@@ -48,44 +48,44 @@ VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca"
 [weakdeps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 Dynesty = "eb527566-0f3e-4aab-bb5f-9d2e403dba70"
-NestedSamplers = "41ceaf6f-1696-4a54-9b49-2e7a9ec3782e"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Pigeons = "0eb8d820-af6a-4919-95ae-11206f830c31"
 Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468"
 
 [extensions]
 ComradeAdvancedHMCExt = "AdvancedHMC"
 ComradeDynestyExt = "Dynesty"
-ComradeNestedExt = "NestedSamplers"
+ComradeEnzymeExt = "Enzyme"
 ComradeOptimizationExt = "Optimization"
 ComradePigeonsExt = "Pigeons"
 ComradePyehtimExt = "Pyehtim"
 
 [compat]
-Accessors = "0.1"
 AbstractMCMC = "3, 4, 5"
+Accessors = "0.1"
 AdvancedHMC = "0.6"
 ArgCheck = "2"
 AstroTime = "0.6,0.7"
 ChainRulesCore = "1"
-ComradeBase = "0.7"
+ComradeBase = "0.8"
 DelimitedFiles = "1"
 DensityInterface = "0.4"
-DimensionalData = "0.26, 0.27"
-Distributions = "0.24,0.25"
-DocStringExtensions = "0.6,0.7,0.8, 0.9"
+DimensionalData = "0.27, 0.28"
+Distributions = "0.25"
+DocStringExtensions = "0.8, 0.9"
 Dynesty = "0.4"
-Enzyme = "0.11, 0.12"
-FillArrays = "0.12, 0.13, 1"
+Enzyme = "0.12"
+EnzymeCore = "0.7"
+FillArrays = "1"
 ForwardDiff = "0.9, 0.10"
 HypercubeTransform = "0.4"
 IntervalSets = "0.6, 0.7"
-LogDensityProblemsAD = "1"
 LogDensityProblems = "2"
+LogDensityProblemsAD = "1"
 Makie = "0.21"
 NamedTupleTools = "0.13,0.14"
-NestedSamplers = "0.8"
-Optimization = "3"
+Optimization = "4"
 PaddedViews = "0.5"
 ParameterHandling = "0.4, 0.5"
 Pigeons = "0.3, 0.4"
@@ -97,25 +97,25 @@ Reexport = "1"
 SpecialFunctions = "0.10, 1, 2"
 StaticArraysCore = "1"
 Statistics = "1.8"
-StatsBase = "0.31,0.32,0.33, 0.34"
-StructArrays = "0.3,0.4,0.5,0.6"
+StatsBase = "0.33,0.34"
+StructArrays = "0.5,0.6"
 Tables = "1"
 TransformVariables = "0.8"
 VLBIImagePriors = "0.8"
-VLBILikelihoods = "^0.2.1"
-VLBISkyModels = "^0.5.5"
-julia = "1.9"
+VLBILikelihoods = "^0.2.6"
+VLBISkyModels = "0.6"
+julia = "1.10"
 
 [extras]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"
 Dynesty = "eb527566-0f3e-4aab-bb5f-9d2e403dba70"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
-NestedSamplers = "41ceaf6f-1696-4a54-9b49-2e7a9ec3782e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Pigeons = "0eb8d820-af6a-4919-95ae-11206f830c31"
 Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "AdvancedHMC", "Dynesty", "Makie", "NestedSamplers", "Optimization", "Pigeons", "Pyehtim"]
+test = ["Test", "AdvancedHMC", "Dynesty", "Enzyme", "Makie", "Optimization", "Pigeons", "Pyehtim"]

benchmarks/benchmarks.jl

@@ -38,12 +38,9 @@ x0 = prior_sample(tpost)
 
 ℓ = logdensityof(tpost)
 @benchmark ℓ($x0)
-
-using Zygote
-using LogDensityProblemsAD
-@benchmark $(tpost)($x0)
-# 32 μs
-@benchmark Zygote.gradient($tpost, $x0)
-# 175 μs
+# 38.1 μs
+using Enzyme
+@benchmark Enzyme.gradient(Enzyme.Reverse, $(Const(tpost)), $x0)
+#107.3 μs
 
 # Now we do the eht-imaging benchmarks

docs/Project.toml

@@ -8,24 +8,20 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365"
 Dynesty = "eb527566-0f3e-4aab-bb5f-9d2e403dba70"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 HypercubeTransform = "9ec9aee3-0fd3-44c2-8e61-a50acc66f3c8"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
-MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0"
-MCMCDiagnostics = "6e857e4b-079a-58c4-aeab-bc2670384359"
-NestedSamplers = "41ceaf6f-1696-4a54-9b49-2e7a9ec3782e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
-OptimizationBBO = "3e6eede4-6085-4f62-9a71-46d9bc1eb92b"
-OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+Pigeons = "0eb8d820-af6a-4919-95ae-11206f830c31"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PolarizedTypes = "d3c5d4cd-a8ee-40d6-aac7-e34df5a20044"
 Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468"
 StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Documenter = "1"

docs/make.jl

@@ -4,7 +4,7 @@
 
 using Documenter, Pkg
 using DocumenterVitepress
-using Comrade, ComradeBase, AdvancedHMC, Dynesty, NestedSamplers, Optimization,
+using Comrade, ComradeBase, AdvancedHMC, Dynesty, Optimization,
       PolarizedTypes
 using Pyehtim, VLBISkyModels, InteractiveUtils
 using AbstractMCMC, Random, HypercubeTransform
@@ -48,8 +48,8 @@ makedocs(;
         "Extensions" => [
                         "ext/optimization.md",
                         "ext/ahmc.md",
-                        "ext/nested.md",
                         "ext/dynesty.md",
+                        "ext/pigeons.md"
                        ],
         "base_api.md",
         "api.md"

docs/src/api.md

@@ -114,6 +114,7 @@ Comrade.IdealInstrumentModel
 Comrade.InstrumentModel
 Comrade.SiteArray
 Comrade.SiteLookup
+Comrade.forward_jones
 ```
 
 
@@ -126,6 +127,7 @@ Comrade.loglikelihood
 Comrade.dataproducts
 Comrade.skymodel
 Comrade.instrumentmodel(::Comrade.AbstractVLBIPosterior)
+Comrade.instrumentmodel(::Comrade.AbstractVLBIPosterior, ::Any)
 Comrade.forward_model
 Comrade.prior_sample
 Comrade.likelihood
@@ -163,6 +165,7 @@ Comrade.dirty_image
 Comrade.dirty_beam
 Comrade.beamsize
 Comrade.apply_fluctuations
+Comrade.corr_image_prior
 Comrade.rmap
 ```
 

docs/src/base_api.md

@@ -88,8 +88,6 @@ ComradeBase.UnstructuredMap
 ComradeBase.baseimage
 ComradeBase.centroid
 ComradeBase.second_moment
-ComradeBase.load
-ComradeBase.save
 ComradeBase.stokes
 ```
 

docs/src/benchmarks.md

@@ -34,14 +34,14 @@ Our benchmark results are the following:
 
 | | Comrade (micro sec) | eht-imaging (micro sec) | Themis (micro sec)|
 |---|---|---|---|
-| posterior eval (min) | 31  | 445  | 55  |
-| posterior eval (mean) | 36  | 476  | 60  |
-| grad posterior eval (min) |  105 (ForwardDiff) | 1898  | 1809  |
-| grad posterior eval (mean) |  119 (ForwardDiff) | 1971 |  1866  |
+| posterior eval (min) | 31.1  | 445  | 55  |
+| posterior eval (mean) | 31.8  | 476  | 60  |
+| grad posterior eval (min) |  104 (Enzyme) | 1898  | 1809  |
+| grad posterior eval (mean) |  107 (Enzyme) | 1971 |  1866  |
 
 Therefore, for this test we found that `Comrade` was the fastest method in all tests. For the posterior evaluation we found that Comrade is > 10x faster than `eht-imaging`, and 2x faster then `Themis`. For gradient evaluations we have `Comrade` is > 15x faster than both `eht-imaging` and `Themis`.
 
-[^1]: Chael A, et al. *Inteferometric Imaging Directly with Closure Phases* 2018 ApJ 857 1 arXiv:1803/07088
+[^1]: Chael A, et al. *Interferometric Imaging Directly with Closure Phases* 2018 ApJ 857 1 arXiv:1803/07088
 
 ## Code
 
@@ -81,11 +81,12 @@ tpost = asflat(post)
 
 x0 = prior_sample(tpost)
 
-using Zygote
-@benchmark $(tpost)($x0)
-# 32 μs
-@benchmark Zygote.gradient($tpost, $x0)
-# 175 μs
+ℓ = logdensityof(tpost)
+@benchmark ℓ($x0)
+# 31.1 μs
+using Enzyme
+@benchmark Enzyme.gradient(Enzyme.Reverse, $(Const(ℓ)), $x0)
+# 104 μs
 ```
 
 ### eht-imaging Code
@@ -129,7 +130,7 @@ preh[1]["y0"] = {"prior_type": "flat", "min" : -eh.RADPERUAS*(40.0), "max" : eh.
 preh[1]["PA"] = {"prior_type": "flat", "min" : -np.pi, "max" : np.pi}
 
 # This is a hack to get the objective function and its gradient
-# we need to do this since the functions depend on some global variables
+# we need to do this since the functions depend on some global ehtim variables
 transform_param = eh.modeling.modeling_utils.transform_param
 def make_paraminit(param_map, meh, trial_model, model_prior):
     model_init = meh.copy()

docs/src/ext/ahmc.md

@@ -15,7 +15,7 @@ chain = sample(post, NUTS(0.8), 10_000; n_adapts=5_000)
 
 In addition our sample call has a few additional keyword arguments:
 
- - `adtype = Val(:Zygote)`: The autodiff package to use. Currently the default is `Zygote` and we recommend using this. Note that you must load Zygote before calling `sample`.
+ - `adtype = Val(:Enzyme)`: The autodiff package to use. Currently the only options is `Enzyme`. Note that you must load Enzyme before calling `sample`.
  - `saveto = MemoryStore()`: Specifies how to store the samples. The default is `MemoryStore` which stores the samples directly in RAM. For large models this is not a good idea. To save samples periodically to disk use [`DiskStore`](@ref), and then load the results with `load_samples`.
 
 Note that like most `AbstractMCMC` samplers the initial location can be specified with the `initial_params` argument.
@@ -26,7 +26,7 @@ Note that like most `AbstractMCMC` samplers the initial location can be specifie
 ```julia
 using Comrade
 using AdvancedHMC
-using Zygote
+using Enzyme
 
 # Some stuff to create a posterior object
 post # of type Comrade.Posterior