Format and update

Author: ptiede

examples/advanced/FitPS/main.jl

@@ -128,7 +128,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; imgdata = (Comrade.CentroidData((0.0, 0.0), beamsize(dcphase)/10.0, grid), ))
+post = VLBIPosterior(skym, dlcamp, dcphase; imgdata = (Comrade.CentroidData((0.0, 0.0), beamsize(dcphase) / 10.0, grid),))
 
 # ## Reconstructing the Image
 
@@ -143,7 +143,7 @@ post = VLBIPosterior(skym, dlcamp, dcphase; imgdata = (Comrade.CentroidData((0.0
 # We also need to import Enzyme to allow for automatic differentiation.
 using Optimization, OptimizationOptimisers
 # tpost = asflat(post)
-xopt, sol = comrade_opt(post, Adam(); maxiters = 5000, initial_params=xopt)
+xopt, sol = comrade_opt(post, Adam(); maxiters = 5000, initial_params = xopt)
 
 using CairoMakie
 using DisplayAs #hide

examples/intermediate/ClosureImaging/main.jl

@@ -126,7 +126,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; imgdata = (Comrade.CentroidData((0.0, 0.0), μas2rad(0.1), grid), ))
+post = VLBIPosterior(skym, dlcamp, dcphase; imgdata = (Comrade.CentroidData((0.0, 0.0), μas2rad(0.1), grid),))
 
 # ## Reconstructing the Image
 

examples/intermediate/PolarizedImaging/main.jl

@@ -301,7 +301,7 @@ intmodel = InstrumentModel(J, intprior)
 # Putting it all together, we form our likelihood and posterior objects for optimization and
 # sampling, and specifying to use Enzyme.Reverse with runtime activity for AD.
 using Enzyme
-post = VLBIPosterior(skym, intmodel, dvis; imgdata = (Comrade.CentroidData((0.0, 0.0), μas2rad(0.1), grid), ))
+post = VLBIPosterior(skym, intmodel, dvis; imgdata = (Comrade.CentroidData((0.0, 0.0), μas2rad(0.1), grid),))
 
 # ## Reconstructing the Image and Instrument Effects
 
@@ -383,7 +383,7 @@ fig |> DisplayAs.PNG |> DisplayAs.Text
 # other imaging examples. For example
 # ```julia
 # using AdvancedHMC
-chain = sample(rng, post, NUTS(0.8), 2_000, n_adapts=1000, progress=true, initial_params=xopt)
+chain = sample(rng, post, NUTS(0.8), 2_000, n_adapts = 1000, progress = true, initial_params = xopt)
 # ```
 
 

examples/intermediate/StokesIImaging/main.jl

@@ -144,7 +144,7 @@ intmodel = InstrumentModel(G, intpr)
 # gradients of the posterior we also need to load `Enzyme.jl`. Under the hood, Comrade will use
 # Enzyme to compute the gradients of the posterior.
 using Enzyme
-post = VLBIPosterior(skym, intmodel, dvis; imgdata = (Comrade.CentroidData((0.0, 0.0), μas2rad(0.1), grid), ))
+post = VLBIPosterior(skym, intmodel, dvis; imgdata = (Comrade.CentroidData((0.0, 0.0), μas2rad(0.1), grid),))
 
 # ## Optimization and Sampling
 

src/posterior/abstract.jl

@@ -96,7 +96,7 @@ data products observed.
 end
 
 @inline function forward_model(d::AbstractVLBIPosterior, θ)
-    forward_model_map(datatype(typeof(d)), d, θ)
+    return forward_model_map(datatype(typeof(d)), d, θ)
 end
 
 """
@@ -160,14 +160,11 @@ end
     return zero(img)
 end
 
-@inline function logdensityofimg(lklhds::Tuple{}, img::AbstractArray{T}) where T
+@inline function logdensityofimg(lklhds::Tuple{}, img::AbstractArray{T}) where {T}
     return zero(T)
 end
 
 
-
-
-
 include("likelihood.jl")
 include("vlbiposterior.jl")
 include("transformed.jl")

src/posterior/likelihood.jl

@@ -114,12 +114,12 @@ struct CentroidData{M, N, G}
 end
 
 
-struct _Centroid{F,T}
+struct _Centroid{F, T}
     f::F
     Σ::T
 end
 
-struct NormalFast{T,S}
+struct NormalFast{T, S}
     μ::T
     Σ::S
 end
@@ -137,7 +137,7 @@ end
 function makelikelihood(data::CentroidData)
     Σ = data.noise .^ 2
     meas = data.measurement
-    f = SVector∘centroid
+    f = SVector ∘ centroid
     ℓ = ConditionedLikelihood(_Centroid(f, Σ), SVector(meas))
     return ℓ
 end

src/posterior/vlbiposterior.jl

@@ -17,7 +17,6 @@ admode(post::VLBIPosterior) = post.admode
 end
 
 
-
 """
     VLBIPosterior(skymodel::SkyModel, instumentmodel::InstrumentModel, 
                   dataproducts::EHTObservationTable...; 
@@ -96,7 +95,7 @@ VLBIPosterior(
     skymodel::AbstractSkyModel, dataproducts::EHTObservationTable...;
     admode = EnzymeCore.set_runtime_activity(EnzymeCore.Reverse), kwargs...
 ) =
-    VLBIPosterior(skymodel, IdealInstrumentModel(), dataproducts...;admode, kwargs..., )
+    VLBIPosterior(skymodel, IdealInstrumentModel(), dataproducts...; admode, kwargs...)
 
 function combine_prior(skyprior, instrumentmodelprior)
     return NamedDist((sky = skyprior, instrument = instrumentmodelprior))

src/skymodels/abstract.jl

@@ -59,10 +59,37 @@ end
 
 function idealmaps(::DualData, m::AbstractSkyModel, x)
     skym = skymodel(m, x.sky)
-    dmap = dualmap(skym, domain(m))
-    return ComradeBase.imgmap(dmap), ComradeBase.vismap(dmap)
+    # We include this special method to sometime optimize the dualmap computation
+    # for specific sky models
+    img, vis = fastdualmap(skym, domain(m))
+    return img, vis
 end
 
+function fastdualmap(skym::AbstractModel, grid::VLBISkyModels.AbstractFourierDualDomain)
+    dm = dualmap(skym, grid)
+    return ComradeBase.imgmap(dm), ComradeBase.vismap(dm)
+end
+
+function fastdualmap(skym::ContinuousImage, grid::VLBISkyModels.AbstractFourierDualDomain)
+    VLBISkyModels.checkgrid(axisdims(skym), imgdomain(grid)) ||
+        throw(DomainError("Image domain does not match skymodel image domain"))
+    img = VLBISkyModels.make_map(m)
+    return img, visibilitymap(skym, grid)
+end
+
+
+function fastdualmap(skym::VLBISkyModels.CompositeModel{<:ContinuousImage}, grid::VLBISkyModels.AbstractFourierDualDomain)
+    VLBISkyModels.checkgrid(axisdims(skym.m1), ComradeBase.imgdomain(grid)) || throw(DomainError("Image domain does not match skymodel image domain"))
+    img = VLBISkyModels.make_map(skym.m1)
+    img2 = intensitymap(skym.m2, grid)
+    img2 .+= img #copy into this one because img will alias otherwise
+    return img, visibilitymap(skym, grid)
+end
+
+fastdualmap(m::VLBISkyModels.CompositeModel{M1, <:ContinuousImage}, grid::VLBISkyModels.AbstractFourierDualDomain) where {M1} =
+    fastdualmap(swap(m), grid)
+
+
 function skymodel(m::AbstractSkyModel, x)
     return m.f(x, m.metadata)
 end

test/Core/models.jl

@@ -150,8 +150,8 @@ end
         @test tt.dynamic == Comrade.skymodel(oskym, x.dynamic)
         @test tt.static == Comrade.skymodel(oskyf, x.static)
 
-        vtot  = last(Comrade.idealmaps(oskytot, (; sky = x)))
-        vdyn  = last(Comrade.idealmaps(oskym, (; sky = x.dynamic)))
+        vtot = last(Comrade.idealmaps(oskytot, (; sky = x)))
+        vdyn = last(Comrade.idealmaps(oskym, (; sky = x.dynamic)))
         vstat = last(Comrade.idealmaps(oskyf, (; sky = x.static)))
 
         @test_opt Comrade.idealmaps(oskytot, (; sky = x))