Add multifrequency capabilities for ContinuousImage

.gitignore

@@ -9,4 +9,3 @@
 examples/Manifest.toml
 .vscode/launch.json
 .vscode/settings.json
-.vscode/settings.json

Project.toml

@@ -8,6 +8,7 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Comrade = "99d987ce-9a1e-4df8-bc0b-1ea019aa547b"
 ComradeBase = "6d8c423b-a35f-4ef1-850c-862fe21f82c4"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
@@ -18,6 +19,7 @@ FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GridInterpolations = "bb4c363b-b914-514b-8517-4eb369bc008a"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 NamedTupleTools = "d9ec5142-1e00-5aa0-9d6a-321866360f50"
@@ -30,6 +32,7 @@ Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029"
 
 [weakdeps]
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
@@ -42,7 +45,7 @@ AbstractFFTs = "1"
 Accessors = "0.1"
 ArgCheck = "2"
 ChainRulesCore = "1"
-ComradeBase = "0.8.7"
+ComradeBase = ">=0.8.7"
 DelimitedFiles = "1"
 DimensionalData = "0.27, 0.28, 0.29.2"
 DocStringExtensions = "0.6,0.7,0.8,0.9"
@@ -52,7 +55,7 @@ FITSIO = "0.17"
 FillArrays = "1"
 ForwardDiff = "0.9, 0.10"
 LinearAlgebra = "1.8"
-Makie = "0.21"
+Makie = "0.22"
 NFFT = "0.10, 0.11, 0.12, 0.13"
 NamedTupleTools = "0.13,0.14"
 PaddedViews = "0.5"

src/VLBISkyModels.jl

@@ -69,6 +69,7 @@ include(joinpath("models", "models.jl"))
 include("utility.jl")
 # include("rules.jl")
 include(joinpath("visualizations", "vis.jl"))
+include(joinpath("models", "multifrequency.jl"))
 include("io.jl")
 
 end

src/models/multidomain/freqtaylor.jl

@@ -23,6 +23,13 @@ struct TaylorSpectral{N,P,T<:NTuple{N},F<:Real,P0} <: FrequencyParams{P}
         return new{N,typeof(param),typeof(index),typeof(freq0),typeof(p0)}(param, index,
                                                                            freq0, p0)
     end
+
+    function TaylorSpectral(index::NTuple{N}, freq0::Real) where {N}
+        param = Nothing
+        p0 = 0
+        return new{N,typeof(param),typeof(index),typeof(freq0),typeof(p0)}(param, index,
+                                                                           freq0, 0)
+    end
 end
 
 function TaylorSpectral(param, index::Real, freq0, p0=zero(param))

src/models/multifrequency.jl

@@ -0,0 +1,162 @@
+export Multifrequency, TaylorSpectral, applyspectral, applyspectral!, generatemodel,
+       visibilitymap_numeric, build_imagecube, mfimagepixels
+
+"""Abstract type to hold all multifrequency spectral models"""
+abstract type AbstractSpectralModel end
+
+"""
+    $(TYPEDEF)
+Spectral Model object.
+
+# Fields
+$(FIELDS)
+"""
+struct Multifrequency{B<:ContinuousImage,F<:Number,S<:FrequencyParams} <:
+       ComradeBase.AbstractModel
+    """
+    Base image model (ContinuousImage only)
+    """
+    base::B
+    """
+    The base image reference frequency.
+    """
+    ν0::F
+    """
+    Multifrequency spectral model
+    """
+    spec::S
+end
+
+# defining the mandatory methods for a Comrade AbstractModel
+visanalytic(::Type{Multifrequency{B}}) where {B} = NotAnalytic()
+imanalytic(::Type{Multifrequency{B}}) where {B} = imanalytic(B)
+radialextent(::Type{Multifrequency{B}}) where {B} = radialextent(B)
+flux(::Type{Multifrequency{B}}) where {B} = flux(B)
+
+function intensity_point(M::Multifrequency, p)
+    I0 = parent(M.base)
+    I_img = applyspectral(I0, M.spec, p.Fr)
+    I_model = ContinuousImage(I_img, M.base.kernel)
+    return intensity_point(I_model, p)
+end
+
+#visibility_point(M::Multifrequency{B},p) where {B} = visibility_point(B)
+
+#"""
+#    $(TYPEDEF)
+#Taylor expansion spectral model of order n for multifrequency imaging. Applies to Continuous Image models only.
+#
+#This is the same multifrequency model implemented in ehtim (Chael et al., 2023).
+#
+# Fields
+#$(FIELDS)
+#"""
+#struct TaylorSpectral{C<:NTuple, ν0<:Real} <: AbstractSpectralModel
+#    """
+#    A tuple containing the Taylor expansion coefficients.
+#    c[1] is the spectral index α, c[2] is the spectral curvature β, etc.
+#    The coeffecients can be either constant values, or arrays with dimensions equal to that of the base image I0.
+#    """
+#    c::C
+#    """
+#    The expansion reference frequency.
+#    """
+#    ν0::C
+#end
+
+function order(::TaylorSpectral{N}) where {N}
+    return N
+end
+
+"""
+Applies Taylor Series spectral model to image data.
+
+Generates image data at frequency ν with respect to the reference frequency ν0.
+"""
+function applyspectral(I0::AbstractArray, spec::TaylorSpectral, ν::N) where {N<:Number}
+    data = copy(I0)
+    applyspectral!(data, spec, ν)
+    return data
+end
+
+function exparg(c::T, xlist::NTuple, i::Number) where {T<:Tuple{Vararg{<:AbstractArray}}}
+    return sum(getindex.(c, i) .* xlist)
+end
+
+function exparg(c::T, xlist::NTuple, i::Number) where {T<:Tuple{Vararg{<:Number}}}
+    return sum(c .* xlist)
+end
+
+function applyspectral!(I0::AbstractArray, spec::TaylorSpectral, ν::Number)
+    ν0 = spec.freq0
+    x = log(ν / ν0) # frequency to evaluate taylor expansion
+    c = spec.index
+
+    n = order(spec)
+    xlist = ntuple(i -> x^i, n) # creating a tuple to hold the frequency powers
+
+    for i in eachindex(I0) # doing expansion one pixel at a time
+        I0[i] = I0[i] * exp(exparg(c, xlist, i))
+    end
+    return I0
+end
+
+"""Given a multifrequency model (base image & spectral model), creates a new Continuous image model at frequency ν."""
+function generatemodel(MF::Multifrequency, ν::N) where {N<:Number}
+    image = parent(MF.base) # ContinuousImage -> SpatialIntensityMap
+    I0 = parent(image) # SpatialIntensityMap -> Array
+
+    data = applyspectral(I0, MF.spec, ν) # base image model, spectral model, frequency to generate new image
+
+    new_intensitymap = rebuild(image; data=data)
+    return ContinuousImage(new_intensitymap, MF.base.kernel)
+end
+
+function visibilitymap_numeric(m::Multifrequency{<:ContinuousImage},
+                               grid::AbstractFourierDualDomain)
+    checkspatialgrid(axisdims(m.base), grid.imgdomain) # compare base image dimensions to spatial dimensions of data cube
+    imgcube = build_imagecube(m, grid.imgdomain)
+    vis = applyft(forward_plan(grid), imgcube)
+    return applypulse!(vis, m.base.kernel, grid)
+end
+
+function checkspatialgrid(imgdims, grid)
+    return !(dims(imgdims) == dims(grid)[1:2]) &&
+           throw(ArgumentError("The image dimensions in `ContinuousImage`\n" *
+                               "and the spatial dimensions of the visibility grid passed to `visibilitymap`\n" *
+                               "do not match. This is not currently supported."))
+end
+
+"""
+Build a multifrequency image cube to hold images at all frequencies.
+"""
+function build_imagecube(m::Multifrequency, mfgrid::RectiGrid)
+    I0 = parent(m.base) # base image IntensityMap 
+    spec = m.spec
+    νlist = mfgrid.Fr
+
+    imgcube = allocate_imgmap(m.base, mfgrid) # build 3D cube of IntensityMap objects
+
+    for i in eachindex(νlist) # setting the image each frequency to first equal the base image and then apply spectral model
+        imgcube[:, :, i] .= I0
+        applyspectral!(@view(imgcube[:, :, i]), spec, νlist[i]) # @view modifies existing image cube inplace --- don't create new object
+    end
+
+    return imgcube
+end
+
+function mfimagepixels(fovx::Real, fovy::Real, nx::Integer, ny::Integer,
+                       νlist::AbstractVector,
+                       x0::Real=0, y0::Real=0; executor=Serial(),
+                       header=ComradeBase.NoHeader())
+    @assert (nx > 0) && (ny > 0) "Number of pixels must be positive"
+
+    psizex = fovx / nx
+    psizey = fovy / ny
+
+    xitr = X(LinRange(-fovx / 2 + psizex / 2 - x0, fovx / 2 - psizex / 2 - x0, nx))
+    yitr = Y(LinRange(-fovy / 2 + psizey / 2 - y0, fovy / 2 - psizey / 2 - y0, ny))
+    νlist = Fr(νlist)
+    grid = RectiGrid((X=xitr, Y=yitr, Fr=νlist); executor, header)
+    return grid
+end

test/multifrequency.jl

@@ -0,0 +1,91 @@
+function test_methods(m::Multifrequency{B}) where {B}
+    @test visanalytic(m) == NotAnalytic()
+    @test imanalysic(m) == imanalytic(B)
+    @test radialextent(m) == radialextent(B)
+    @test flux(m) = flux(B)
+    GC.gc()
+    return nothing
+end
+
+#function test_ContinuousImageTaylorSpectral()
+#    return TaylorSpectral(index::NTuple{N}, freq0::Real)
+#end
+
+function gaussmodel(θ)
+    (; f, σG) = θ
+    g = f * stretched(Gaussian(), σG, σG)
+    return g
+end
+
+# unit test consisting of a constant spectral index gaussian
+
+@testset "Multifrequency Gaussian: Taylor Expansion" begin
+    ν0 = 8e9 # reference frequency (Hz)
+    ν = 12e9 # target frequency (Hz)
+    νlist = [ν0, ν]
+
+    # testing mfimagepixels & generating a multifrequency image grid
+    g = imagepixels(μas2rad(1000.0), μas2rad(1000.0), 256, 256)
+    mfgrid = mfimagepixels(μas2rad(1000.0), μas2rad(1000.0), 256, 256, νlist) # build multifrequency image grid
+
+    @test dims(g)[1].val ≈ dims(mfgrid)[1].val
+    @test dims(g)[2].val ≈ dims(mfgrid)[2].val
+    @test dims(mfgrid)[3].val == νlist
+
+    # 8 GHz Gaussian with total flux = 1.2 Jy
+    θ1 = (f=1.2, σG=μas2rad(100))
+    gauss1 = intensitymap(gaussmodel(θ1), g)
+    gaussmodel1 = ContinuousImage(gauss1, BSplinePulse{3}())
+
+    # 12 GHz Gaussian with total flux = 1.6 Jy
+    θ2 = (f=1.6, σG=μas2rad(100))
+    gauss2truth = intensitymap(gaussmodel(θ2), g)
+    gaussmodel2truth = ContinuousImage(gauss2truth, BSplinePulse{3}())
+
+    # testing spectral index map: constant spectral index and 0 spectral curvature
+    α0 = log(1.6 / 1.2) / log(12 / 8)
+    α = fill(α0, size(gauss1)) # spectral index map
+    β0 = 0.0
+    β = fill(β0, size(gauss1)) # spectral curvature map
+
+    spec1 = TaylorSpectral((α, β), ν0)
+    spec2 = TaylorSpectral((α0, β0), ν0)
+
+    @test VLBISkyModels.order(spec1) == 2
+    @test VLBISkyModels.order(spec2) == 2
+
+    # create a multifrequency object
+    mfgauss1 = Multifrequency(gaussmodel1, ν0, spec1)
+    mfgauss2 = Multifrequency(gaussmodel1, ν0, spec2)
+
+    # test intensity_point
+    p = (; X=0, Y=0, Fr=ν)
+    @test VLBISkyModels.intensity_point(mfgauss1, p) ==
+          VLBISkyModels.intensity_point(gaussmodel2truth, p)
+    @test VLBISkyModels.intensity_point(mfgauss2, p) ==
+          VLBISkyModels.intensity_point(gaussmodel2truth, p)
+
+    # test generatemodel
+    # generate model at new frequency & compare with ground truth
+    @test parent(generatemodel(mfgauss1, ν)) ≈ gauss2truth
+    @test parent(generatemodel(mfgauss2, ν)) ≈ gauss2truth
+
+    # test visibilitymap
+    # generate 100 visibilities: half at 8 GHz and half at 12 GHz
+    u = range(1, 10, 50) * 1e7 # random visibilities
+    v = range(1, 10, 50) * 1e7
+    f = similar(u) # each uv point has a frequency associated with it
+    f[1:25] .= νlist[1]
+    f[26:50] .= νlist[2]
+
+    fdd8 = FourierDualDomain(axisdims(gauss1), UnstructuredDomain((U=u[1:25], V=v[1:25])),
+                             NFFTAlg())
+    fdd12 = FourierDualDomain(axisdims(gauss2truth),
+                              UnstructuredDomain((U=u[26:50], V=v[26:50])), NFFTAlg())
+    mffdd = FourierDualDomain(RectiGrid((X=g.X, Y=g.Y, Fr=νlist)),
+                              UnstructuredDomain((U=u, V=v, Fr=f)), NFFTAlg())
+
+    # comparing multifrequency to single frequency results
+    @test visibilitymap(mfgauss1, mffdd)[1:25] ≈ visibilitymap(gaussmodel1, fdd8)
+    @test visibilitymap(mfgauss1, mffdd)[26:50] ≈ visibilitymap(gaussmodel2truth, fdd12)
+end

test/runtests.jl

@@ -193,4 +193,5 @@ end
     include("viz.jl")
     include("io.jl")
     include("stokesintensitymap.jl")
+    include("multifrequency.jl")
 end