Move ReactantPlan to extension

Author: ptiede

test/integration/Comrade/comimager.jl

@@ -35,7 +35,6 @@ using Distributions
 using VLBIImagePriors
 using LogExpFunctions
 import TransformVariables as TV
-include("reactant_nfft.jl")
 
 using Downloads
 using Distributions
@@ -49,38 +48,6 @@ using Test
 const dataf = joinpath("/home/ptiede/.julia/dev/Comrade/examples/Data/SR1_M87_2017_096_lo_hops_netcal_StokesI.uvfits")
 
 # TODO upstream to VLBISkyModels
-struct ReactantAlg <: VLBISkyModels.NUFT end
-
-function VLBISkyModels.plan_nuft_spatial(
-    ::ReactantAlg,
-    imgdomain::ComradeBase.AbstractRectiGrid,
-    visdomain::ComradeBase.UnstructuredDomain,
-)
-    visp = domainpoints(visdomain)
-    uv2 = similar(visp.U, (2, length(visdomain)))
-    dpx = pixelsizes(imgdomain)
-    dx = dpx.X
-    dy = dpx.Y
-    rm = ComradeBase.rotmat(imgdomain)'
-    # Here we flip the sign because the NFFT uses the -2pi convention
-    uv2[1, :] .= -VLBISkyModels._rotatex.(visp.U, visp.V, Ref(rm)) .* dx
-    uv2[2, :] .= -VLBISkyModels._rotatey.(visp.U, visp.V, Ref(rm)) .* dy
-    return ReactantNFFTPlan(uv2, size(imgdomain))
-end
-
-function VLBISkyModels.make_phases(
-    ::ReactantAlg,
-    imgdomain::ComradeBase.AbstractRectiGrid,
-    visdomain::ComradeBase.UnstructuredDomain,
-)
-    return VLBISkyModels.make_phases(NFFTAlg(), imgdomain, visdomain)
-end
-
-function VLBISkyModels._jlnuft!(out, A::ReactantNFFTPlan, inp::Reactant.AnyTracedRArray)
-    LinearAlgebra.mul!(out, A, inp)
-    return nothing
-end
-
 
 
 # TODO Make ReactantLogExpFunctionsExt.
@@ -89,102 +56,6 @@ function LogExpFunctions.logistic(@nospecialize x::Reactant.TracedRNumber)
 end
 LogExpFunctions.log1pexp(x::Reactant.TracedRNumber) = log(1 + exp(x))
 
-#!!!! TODO Everything in this block needs to be upstreamed to TransformVariables.jl
-# The major problem is that most of them require @allowscalar which is not very nice enforce in their
-# codebase since it is expensive for CPU code. The other temporary solution is to make a ReactantTransformVariablesExt.jl package.
-# function TV.transform_with(
-#     flag::TV.NoLogJac, t::TV.ScalarTransform, x::Reactant.AnyTracedRVector, index
-# )
-#     return transform(t, @allowscalar x[index]), flag, index + 1
-# end
-
-# TODO Upstream to TransformVariables.jl
-# function TV.transform_with(
-#     ::TV.LogJac, t::TV.ScalarTransform, x::Reactant.AnyTracedRVector, index
-# )
-#     return TV.transform_and_logjac(t, @allowscalar x[index])..., index + 1
-# end
-
-# TODO This is needed for TransformVariables but @allowscalar is rather annoying here.
-# function TV._transform_tuple(flag::TV.LogJacFlag, x::Reactant.AnyTracedRVector, index, ts)
-#     tfirst = first(ts)
-#     out = TV.transform_with(flag, tfirst, x, index)
-#     @allowscalar yfirst = out[1]
-#     @allowscalar ℓfirst = out[2]
-#     @allowscalar index′ = out[3]
-#     # yrest, ℓrest, index′′
-#     trest = Base.tail(ts)
-#     outrest = TV._transform_tuple(flag, x, index′, trest)
-#     @allowscalar yrest = outrest[1]
-#     @allowscalar ℓrest = outrest[2]
-#     @allowscalar index′′ = outrest[3]
-#     return (yfirst, yrest...), ℓfirst + ℓrest, index′′
-# end
-
-# TODO Upstream to TransformVariables.jl
-# function TV._transform_tuple(
-#     flag::TV.LogJacFlag, x::Reactant.AnyTracedRVector, index, ::Tuple{}
-# )
-#     return (), TV.logjac_zero(flag, eltype(x)), index
-# end
-
-# # TODO Upstream to TransformVariables.jl
-# TV.logjac_zero(::TV.LogJac, ::Type{T}) where {T<:Reactant.RNumber} = log(one(T))
-
-# TODO Upstream to TransformVariables.jl (essentially identical just need to loosen types)
-# function TV.transform_with(
-#     flag::TV.LogJacFlag,
-#     t::TV.ArrayTransformation{TV.Identity},
-#     x::Reactant.AnyTracedRVector,
-#     index,
-# )
-#     (; dims) = t
-#     index′ = index + dimension(t)
-#     y = reshape(x[index:(index′ - 1)], t.dims)
-#     return y, TV.logjac_zero(flag, eltype(x)), index′
-# end
-
-#!!!! End of TransformVariables.jl upstream block
-
-# TODO Upstream to VLBIImagePriors (allowscalar needed for Reactant)
-# function TV.transform_with(
-#     flag::TV.LogJacFlag, ::AngleTransform, y::Reactant.AnyTracedRVector, index
-# )
-#     T = eltype(y)
-#     ℓi = TV.logjac_zero(flag, T)
-#     x1 = @allowscalar y[index]
-#     x2 = @allowscalar y[index + 1]
-#     r = sqrt(x1^2 + x2^2)
-#     # Use log-normal with μ = 0, σ = 1/4
-#     σ = oftype(r, 1 / 4)
-#     if !(flag isa TV.NoLogJac)
-#         lr = log(r)
-#         ℓi = -lr^2 * inv(2 * σ^2) - lr
-#     end
-
-#     return atan(x1, x2), ℓi, index + 2
-# end
-
-# TODO to upstream to VLBIImagePriors
-# function TV.transform_with(
-#     flag::TV.LogJacFlag,
-#     t::TV.ArrayTransformation{<:AngleTransform},
-#     y::Reactant.AnyTracedRVector,
-#     index,
-# )
-#     (; inner_transformation, dims) = t
-#     T = eltype(y)
-#     ℓ = TV.logjac_zero(flag, T)
-#     out = similar(y, dims)
-#     @trace for i in eachindex(out)
-#         θ, ℓi, index2 = TV.transform_with(flag, inner_transformation, y, index)
-#         index = index2
-#         ℓ += ℓi
-#         @allowscalar out[i] = θ
-#     end
-#     return out, ℓ, index
-# end
-
 # TODO Make Distributions package that is compatible with Reactant
 Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a))
 function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber)