Merge branch 'main' of https://github.com/ptiede/Comrade.jl

Author: ptiede

Project.toml

@@ -1,7 +1,7 @@
 name = "Comrade"
 uuid = "99d987ce-9a1e-4df8-bc0b-1ea019aa547b"
 authors = ["Paul Tiede <ptiede91@gmail.com>"]
-version = "0.11.17"
+version = "0.11.18"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -82,7 +82,7 @@ FillArrays = "1"
 HypercubeTransform = "^0.4.11"
 IntervalSets = "0.6, 0.7"
 LogDensityProblems = "2"
-Makie = "0.22"
+Makie = "0.24"
 Measurements = "2"
 NamedTupleTools = "0.13,0.14"
 Optimization = "4"
@@ -103,7 +103,7 @@ Tables = "1"
 TransformVariables = "0.8"
 VLBIImagePriors = "0.10"
 VLBILikelihoods = "^0.2.6"
-VLBISkyModels = "0.6"
+VLBISkyModels = "^0.6.15"
 julia = "1.10"
 
 [extras]

docs/make.jl

@@ -10,6 +10,7 @@ using Pyehtim, VLBISkyModels, InteractiveUtils
 using AbstractMCMC, Random, HypercubeTransform
 using CairoMakie
 
+const ComradeMakieExt = Base.get_extension(Comrade, :ComradeMakieExt)
 
 Documenter.DocMeta.setdocmeta!(Comrade, :DocTestSetup, :(using Comrade); recursive = true)
 

docs/src/api.md

@@ -201,7 +201,10 @@ A user must first load `Makie` or a `Makie` backend, e.g., `CairoMakie` to use t
 
 ```@docs
 Comrade.plotfields
+Comrade.plotfields!
 Comrade.axisfields
 Comrade.plotcaltable
 Comrade.baselineplot
+Comrade.baselineplot!
+ComradeMakieExt.BaselinePlot
 ```

examples/advanced/HybridImaging/Project.toml

@@ -13,7 +13,7 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029"
 
 [compat]
-CairoMakie = "0.13"
+CairoMakie = "0.15"
 Optimization = "4"
 Pyehtim = "0.2"
 StableRNGs = "1"

examples/advanced/HybridImaging/main.jl

@@ -38,7 +38,7 @@ using Pyehtim
 
 # For reproducibility we use a stable random number genreator
 using StableRNGs
-rng = StableRNG(11)
+rng = StableRNG(12)
 
 
 # To download the data visit https://doi.org/10.25739/g85n-f134
@@ -149,13 +149,13 @@ skym = SkyModel(sky, skyprior, g; metadata = skymetadata)
 # This is everything we need to specify our posterior distribution, which our is the main
 # object of interest in image reconstructions when using Bayesian inference.
 using Enzyme
-post = VLBIPosterior(skym, intmodel, dvis; admode = set_runtime_activity(Enzyme.Reverse))
+post = VLBIPosterior(skym, intmodel, dvis)
 
 # We can sample from the prior to see what the model looks like
 using DisplayAs #hide
 using CairoMakie
 xrand = prior_sample(rng, post)
-gpl = imagepixels(μas2rad(200.0), μas2rad(200.0), 128, 128)
+gpl = refinespatial(g, 3)
 fig = imageviz(intensitymap(skymodel(post, xrand), gpl));
 fig |> DisplayAs.PNG |> DisplayAs.Text #hide
 
@@ -168,9 +168,8 @@ fig |> DisplayAs.PNG |> DisplayAs.Text #hide
 # For optimization we will use the `Optimization.jl` package and the LBFGS optimizer.
 # To use this we use the [`comrade_opt`](@ref) function
 using Optimization
-using OptimizationOptimJL
 xopt, sol = comrade_opt(
-    post, LBFGS();
+    post, Optimization.LBFGS();
     initial_params = xrand, maxiters = 2000, g_tol = 1.0e0
 );
 
@@ -249,12 +248,10 @@ figd |> DisplayAs.PNG |> DisplayAs.Text #hide
 
 # Now let's check the residuals using draws from the posterior
 fig = Figure(; size = (600, 400))
-ax, = baselineplot(fig[1, 1], res[1], :uvdist, :res, label = "MAP", color = :blue, colorim = :red, marker = :circle)
+ax, = plotfields!(fig[1, 1], res[1], :uvdist, :res, scatter_kwargs = (; label = "MAP", color = :blue, colorim = :red, marker = :circle), legend = false)
 for s in sample(chain, 10)
     baselineplot!(ax, residuals(post, s)[1], :uvdist, :res, alpha = 0.2, label = "Draw")
 end
-ax.xlabel = "uv-distance (λ)"
-ax.ylabel = "Normalized Residuals"
 axislegend(ax, merge = true)
 fig |> DisplayAs.PNG |> DisplayAs.Text #hide
 

examples/beginner/GeometricModeling/Project.toml

@@ -14,7 +14,7 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029"
 
 [compat]
-CairoMakie = "0.12, 0.13"
+CairoMakie = "0.15"
 Distributions = "0.25"
 Optimization = "4"
 Pigeons = "0.4"

examples/beginner/GeometricModeling/main.jl

@@ -78,8 +78,8 @@ using Distributions, VLBIImagePriors
 prior = (
     radius = Uniform(μas2rad(10.0), μas2rad(30.0)),
     width = Uniform(μas2rad(1.0), μas2rad(10.0)),
-    ma = (Uniform(0.0, 0.5), Uniform(0.0, 0.5)),
-    mp = (Uniform(0, 2π), Uniform(0, 2π)),
+    ma = (Uniform(0.0, 0.5),),
+    mp = (Uniform(0, 2π),),
     τ = Uniform(0.0, 1.0),
     ξτ = Uniform(0.0, π),
     f = Uniform(0.0, 1.0),
@@ -119,8 +119,8 @@ logdensityof(
         sky = (
             radius = μas2rad(20.0),
             width = μas2rad(10.0),
-            ma = (0.3, 0.3),
-            mp = (π / 2, π),
+            ma = (0.3,),
+            mp = (π / 2,),
             τ = 0.1,
             ξτ = π / 2,
             f = 0.6,
@@ -218,16 +218,14 @@ DisplayAs.Text(DisplayAs.PNG(fig))
 # model is fitting the data we can plot the model and data products. As of Comrade 0.11.7 Makie
 # is the preferred plotting tool. For plotting data there are two classes of functions
 #  - `baselineplot` which gives complete control of plotting
-#  - `plotfields, axisfields` which are more automated and limited but will automatically add
+#  - `plotfields, plotfields!` which are more automated and limited but will automatically add
 #     labels, legends, titles etc.
-# We will demonstrate both below.
+# A reasonable workflow is to use `plotfields` to set up the initial figure and axis labels and then
+# then use `baselineplot!` to add additional plots to the axis. For example,
 lcsim, cpsim = simulate_observation(post, xopt; add_thermal_noise = false)
-fig = Figure(; size = (800, 300))
-ax1 = Axis(fig[1, 1], xlabel = "√Quadrangle Area", ylabel = "Log Closure Amplitude")
-baselineplot!(ax1, lcsim, uvdist, measwnoise, marker = :circle, label = "MAP", error = true)
-baselineplot!(ax1, dlcamp, uvdist, Comrade.measurement, marker = :+, color = :black, label = "Data")
-ax2 = Axis(fig[1, 2], xlabel = "√Triangle Area", ylabel = "Closure Phase")
-baselineplot!(ax2, cpsim, uvdist, mod2pi ∘ measwnoise, marker = :circle, label = "MAP", error = true)
+fig, ax1 = plotfields(lcsim, uvdist, measwnoise, scatter_kwargs = (; marker = :circle, label = "MAP"), figure_kwargs = (; size = (800, 300)), legend = false);
+baselineplot!(ax1, dlcamp, uvdist, measurement, marker = :+, color = :black, label = "Data")
+ax2, = plotfields!(fig[1, 2], cpsim, uvdist, mod2pi ∘ measwnoise, scatter_kwargs = (; marker = :circle, label = "MAP"), axis_kwargs = (ylabel = "Closure Phase (rad)",))
 baselineplot!(ax2, dcphase, uvdist, mod2pi ∘ measurement, marker = :+, color = :black, label = "Data")
 axislegend(ax1, framevisible = false)
 DisplayAs.Text(DisplayAs.PNG(fig))
@@ -237,7 +235,7 @@ DisplayAs.Text(DisplayAs.PNG(fig))
 # are marginalized over the posterior.
 fig = Figure(; size = (800, 300))
 ax1 = Axis(fig[1, 1], xlabel = "√Quadrangle Area", ylabel = "Log Closure Amplitude")
-ax2 = Axis(fig[1, 2], xlabel = "√Triangle Area", ylabel = "Closure Phase")
+ax2 = Axis(fig[1, 2], xlabel = "√Triangle Area", ylabel = "Closure Phase (rad)")
 for i in 1:10
     mobs = simulate_observation(post, sample(chain, 1)[1])
     mlca = mobs[1]
@@ -254,7 +252,6 @@ DisplayAs.Text(DisplayAs.PNG(fig))
 # The normalied residuals are the difference between the data
 # and the model, divided by the data's error:
 rd = residuals(post, chain[end])
-fig = Figure(; size = (800, 300))
-axisfields(fig[1, 1], rd[1], uvdist, :res)
-axisfields(fig[1, 2], rd[2], uvdist, :res)
+fig, ax = plotfields(rd[1], uvdist, :res, axis_kwargs = (; ylabel = "Norm. Res. LCA"))
+plotfields!(fig[2, 1], rd[2], uvdist, :res, axis_kwargs = (; ylabel = "Norm. Res. CP"))
 DisplayAs.Text(DisplayAs.PNG(fig))

examples/beginner/LoadingData/main.jl

@@ -58,24 +58,39 @@ coh = extract_table(obs, Coherencies())
 #-
 # We can also recover the array used in the observation using
 using DisplayAs
-plotfields(coh, :U, :V, axis_kwargs = (xreversed = true,)) |> DisplayAs.PNG |> DisplayAs.Text # Plot the baseline coverage
+plotfields(coh, U, V, axis_kwargs = (xreversed = true,)) |> DisplayAs.PNG |> DisplayAs.Text # Plot the baseline coverage
 
 # As of Comrade 0.11.7 Makie is the preferred plotting tool. For plotting data there are two
 # classes of functions:
 #  - `baselineplot` which gives complete control of plotting
 #  - `plotfields, axisfields` which are more automated and limited but will automatically add
 #     labels, legends, titles etc.
 fig = Figure(; size = (800, 600))
-axisfields(fig[1, 1], vis, :uvdist, :measurement)
-axisfields(fig[1, 2], amp, :uvdist, :measurement)
-axisfields(fig[2, 1], cphase, :uvdist, :measurement)
-axisfields(fig[2, 2], lcamp, :uvdist, :measurement)
+plotfields!(fig[1, 1], vis, uvdist, measurement)
+plotfields!(fig[1, 2], amp, uvdist, measurement)
+plotfields!(fig[2, 1], cphase, uvdist, measurement)
+plotfields!(fig[2, 2], lcamp, uvdist, measurement)
 fig |> DisplayAs.PNG |> DisplayAs.Text
 
-# And also the coherency matrices. Here we show how to use `baselineplot` to plot the data
+# And also the coherency matrices. Since the data products are a matrix we need to plot each one separately.
 fig = Figure(; size = (800, 600))
-baselineplot(fig[1, 1], coh, x -> uvdist(x) / 1.0e9, x -> abs(measwnoise(x)[1, 1]), error = true, axis = (ylabel = "RR", xlabel = "uv distance (Gλ)"))
-baselineplot(fig[2, 1], coh, x -> uvdist(x) / 1.0e9, x -> abs(measwnoise(x)[2, 1]), error = true, axis = (ylabel = "LR", xlabel = "uv distance (Gλ)"))
-baselineplot(fig[1, 2], coh, x -> uvdist(x) / 1.0e9, x -> abs(measwnoise(x)[1, 2]), error = true, axis = (ylabel = "RL", xlabel = "uv distance (Gλ)"))
-baselineplot(fig[2, 2], coh, x -> uvdist(x) / 1.0e9, x -> abs(measwnoise(x)[2, 2]), error = true, axis = (ylabel = "LL", xlabel = "uv distance (Gλ)"))
+plotfields!(fig[1, 1], coh, uvdist, x -> measwnoise(x)[1, 1], axis_kwargs = (ylabel = "RR", xlabel = "uv distance (Gλ)"))
+plotfields!(fig[2, 1], coh, uvdist, x -> measwnoise(x)[2, 1], axis_kwargs = (ylabel = "LR", xlabel = "uv distance (Gλ)"))
+plotfields!(fig[1, 2], coh, uvdist, x -> measwnoise(x)[1, 2], axis_kwargs = (ylabel = "RL", xlabel = "uv distance (Gλ)"))
+plotfields!(fig[2, 2], coh, uvdist, x -> measwnoise(x)[2, 2], axis_kwargs = (ylabel = "LL", xlabel = "uv distance (Gλ)"))
+fig
+
+# You can also plot a single baseline
+fig, ax = plotfields(coh, (:AA, :LM), Ti, x -> abs(measwnoise(x)[1, 1]), axis_kwargs = (; ylabel = "|RR|"))
+ax2 = plotfields!(fig[1, 2], coh, (:LM, :AZ), Ti, x -> abs(measwnoise(x)[1, 1]), axis_kwargs = (; ylabel = "|RR|"))
+fig
+
+# Finally, we provide a more low-level plotting function `baselineplot` which allows you to plot
+# any field against any other field. This is what `plotfields` calls under the hood. However, it
+# does not automatically add labels, legends, titles etc, but can add multiple baselines to the same plot.
+fig, ax = baselineplot(coh, (:AA, :LM), Ti, x -> abs(measwnoise(x)[1, 1]), label = "AA-LM")
+baselineplot!(ax, coh, (:LM, :AZ), Ti, x -> abs(measwnoise(x)[1, 1]), label = "LM-AZ")
+ax.ylabel = "|RR|"
+ax.xlabel = "Time (UTC)"
+axislegend(ax)
 fig

examples/intermediate/ClosureImaging/Project.toml

@@ -17,7 +17,7 @@ VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029"
 VLBILikelihoods = "90db92cd-0007-4c0a-8e51-dbf0782ce592"
 
 [compat]
-CairoMakie = "0.13"
+CairoMakie = "0.15"
 DisplayAs = "0.1"
 Distributions = "0.25"
 Optimization = "4"

examples/intermediate/ClosureImaging/main.jl

@@ -127,7 +127,7 @@ skym = SkyModel(sky, prior, grid; metadata = skymeta)
 # Since we are fitting closures we do not need to include an instrument model, since
 # the closure likelihood is approximately independent of gains in the high SNR limit.
 using Enzyme
-post = VLBIPosterior(skym, dlcamp, dcphase; admode = set_runtime_activity(Enzyme.Reverse))
+post = VLBIPosterior(skym, dlcamp, dcphase)
 
 # ## Reconstructing the Image
 
@@ -138,13 +138,12 @@ post = VLBIPosterior(skym, dlcamp, dcphase; admode = set_runtime_activity(Enzyme
 
 # To optimize our posterior `Comrade` provides the `comrade_opt` function. To use this
 # functionality a user first needs to import `Optimization.jl` and the optimizer of choice.
-# In this tutorial we will use Optim.jl's L-BFGS optimizer, which is defined in the sub-package
-# OptimizationOptimJL. We also need to import Enzyme to allow for automatic differentiation.
+# In this tutorial we will use Optiizations LBFGS optimizer.
+# We also need to import Enzyme to allow for automatic differentiation.
 using Optimization
-using OptimizationOptimJL
 xopt, sol = comrade_opt(
-    post, LBFGS();
-    maxiters = 1000, initial_params = prior_sample(rng, post)
+    post, Optimization.LBFGS();
+    maxiters = 2000, initial_params = prior_sample(rng, post)
 );
 
 
@@ -153,8 +152,8 @@ using CairoMakie
 using DisplayAs #hide
 res = residuals(post, xopt)
 fig = Figure(; size = (800, 300))
-axisfields(fig[1, 1], res[1], :uvdist, :res) |> DisplayAs.PNG |> DisplayAs.Text
-axisfields(fig[1, 2], res[2], :uvdist, :res) |> DisplayAs.PNG |> DisplayAs.Text
+plotfields!(fig[1, 1], res[1], :uvdist, :res);
+plotfields!(fig[1, 2], res[2], :uvdist, :res);
 fig |> DisplayAs.PNG |> DisplayAs.Text
 
 # Now let's plot the MAP estimate.
@@ -177,8 +176,8 @@ DisplayAs.Text(DisplayAs.PNG(fig)) #hide
 #     For our `metric` we use a diagonal matrix due to easier tuning.
 #-
 using AdvancedHMC
-chain = sample(rng, post, NUTS(0.8), 700; n_adapts = 500, progress = false, initial_params = xopt);
-
+out = sample(rng, post, NUTS(0.8), 700; n_adapts = 500, saveto = DiskStore(), initial_params = xopt);
+chain = load_samples(out)
 
 # !!! warning
 #     This should be run for longer!