update

Author: ptiede

examples/Project.toml

@@ -32,6 +32,7 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029"
 VLBILikelihoods = "90db92cd-0007-4c0a-8e51-dbf0782ce592"
+VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca"
 WGLMakie = "276b4fcb-3e11-5398-bf8b-a0c2d153d008"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

playground/enzyme_dft_vis.jl

@@ -38,14 +38,14 @@ dvis = extract_table(obs, ComplexVisibilities())
 
 
 function sky(θ, metadata)
-    c = θ
+    c = θ.c
     (;grid, cache) = metadata
     img = IntensityMap(c, grid)
     m = ContinuousImage(img, cache)
     return m
 end
 
-npix = 12
+npix = 48
 fovx = μas2rad(80.0)
 fovy = μas2rad(80.0)
 
@@ -56,24 +56,58 @@ cache = create_cache(NFFTAlg(dvis), buffer, DeltaPulse())
 
 
 using VLBIImagePriors
-prior = ImageUniform(npix, npix)
 
 skymetadata = (;grid, cache)
-# instrumentmetadata = (;gcache, gcachep)
-lklhd = RadioLikelihood(sky, dvis; skymeta=skymetadata)
+
+function instrument(θ, metadata)
+    (; lgamp, gphase) = θ
+    (; gcache, gcachep) = metadata
+    ## Now form our instrument model
+    gvis = exp.(lgamp)
+    gphase = exp.(1im.*gphase)
+    jgamp = jonesStokes(gvis, gcache)
+    jgphase = jonesStokes(gphase, gcachep)
+    return JonesModel(jgamp*jgphase)
+end
+
+gcache = jonescache(dvis, ScanSeg())
+gcachep = jonescache(dvis, ScanSeg(); autoref=SEFDReference((complex(1.0))))
+
+using VLBIImagePriors
+# Now we can form our metadata we need to fully define our model.
+metadata = (;gcache, gcachep)
+
+using Distributions
+using DistributionsAD
+distamp = station_tuple(dvis, Normal(0.0, 0.1); LM = Normal(1.0))
+
+distphase = station_tuple(dvis, DiagonalVonMises(0.0, inv(π^2)))
+
+
+prior = NamedDist(
+            c = ImageUniform(npix, npix),
+            lgamp = CalPrior(distamp, gcache),
+            gphase = CalPrior(distphase, gcachep),
+            )
+
+
+
+lklhd = RadioLikelihood(sky, instrument, dvis; skymeta=skymetadata, instrumentmeta=metadata)
 post = Posterior(lklhd, prior)
 
 tpost = asflat(post)
 ndim = dimension(tpost)
 
 using Enzyme
+using Zygote
 Enzyme.API.runtimeActivity!(true)
 # Enzyme.API.printall!(false)
 x0 = prior_sample(tpost)
 dx0 = zero(x0)
 lt=logdensityof(tpost)
 ℓ = logdensityof(tpost, x0)
-dtpost = deepcopy(tpost)
+gz,  = Zygote.gradient(lt, x0)
 using BenchmarkTools
-@time autodiff(Reverse, logdensityof, Const(tpost), Duplicated(x0, fill!(dx0, 0.0)))
-@btime autodiff(Reverse, logdensityof, $(Duplicated(tpost, dtpost)), Duplicated($x0, fill!($dx0, 0.0)))
+autodiff(Reverse, Const(lt), Active, Duplicated(x0, fill!(dx0, 0.0)))
+@benchmark autodiff(Reverse, logdensityof, $(Const(tpost)), Duplicated($x0, fill!($dx0, 0.0)))
+@benchmark Zygote.gradient($lt, $x0)

playground/enzyme_geom_vis.jl

@@ -21,7 +21,7 @@ using LinearAlgebra
 using StableRNGs
 rng = StableRNG(42)
 
-
+Enzyme.Compiler.bitcode_replacement!(false)
 
 # ## Load the Data
 
@@ -106,7 +106,7 @@ prior = NamedDist(
 
 # Putting it all together we form our likelihood and posterior objects for optimization and
 # sampling.
-lklhd = RadioLikelihood(sky, instrument, dvis; instrumentmeta = metadata)
+lklhd = RadioLikelihood(sky, instrument, dvis; instrumentmeta=metadata)
 post = Posterior(lklhd, prior)
 
 # ## Reconstructing the Image and Instrument Effects
@@ -122,10 +122,10 @@ ndim = dimension(tpost)
 # inference packages use this interface as well.
 using Zygote
 using Enzyme
-Enzyme.API.runtimeActivity!(true)
+# Enzyme.API.runtimeActivity!(true)
 
 x0 = randn(rng, ndim)
 ℓ = logdensityof(tpost)
 gz, = Zygote.gradient(ℓ, x0)
 dx0 = zero(x0)
-autodiff(Reverse, logdensityof, Duplicated(tpost, deepcopy(tpost)), Duplicated(x0, dx0))
+autodiff(Reverse, Const(ℓ), Active, Duplicated(x0, dx0))

src/calibration/jones.jl

@@ -1,5 +1,5 @@
 export JonesCache, TrackSeg, ScanSeg, FixedSeg, IntegSeg, jonesG, jonesD, jonesT,
-       ResponseCache, JonesModel, jonescache, station_tuple, jonesmap
+       ResponseCache, JonesModel, jonescache, station_tuple
 
 """
     $(TYPEDEF)