Merge pull request #414 from ptiede/ptiede-multisky

Add ability to have multiple grids for a sky model

Author: ptiede

docs/src/api.md

@@ -87,6 +87,7 @@ Comrade.Coherencies
 Comrade.AbstractSkyModel
 Comrade.SkyModel
 Comrade.FixedSkyModel
+Comrade.MultiSkyModel
 Comrade.idealvisibilities
 Comrade.skymodel(::Comrade.AbstractVLBIPosterior, ::Any)
 ```

src/Comrade.jl

@@ -65,7 +65,7 @@ import ComradeBase: flux, radialextent, intensitymap, intensitymap!,
     amplitudemap
 include("observations/observations.jl")
 include("instrument/instrument.jl")
-include("skymodels/models.jl")
+include("skymodels/abstract.jl")
 include("posterior/abstract.jl")
 include("inference/inference.jl")
 include("visualizations/visualizations.jl")

src/posterior/vlbiposterior.jl

@@ -1,4 +1,4 @@
-struct VLBIPosterior{D, T, P, MS <: ObservedSkyModel, MI <: AbstractInstrumentModel, ADMode <: Union{Nothing, EnzymeCore.Mode}} <: AbstractVLBIPosterior
+struct VLBIPosterior{D, T, P, MS <: AbstractSkyModel, MI <: AbstractInstrumentModel, ADMode <: Union{Nothing, EnzymeCore.Mode}} <: AbstractVLBIPosterior
     data::D
     lklhds::T
     prior::P

src/skymodels/abstract.jl

@@ -0,0 +1,75 @@
+"""
+    AbstractSkyModel
+
+The abstract type for Comrade Sky Models. For a concrete implementation see [`SkyModel`](@ref).
+
+Any subtype must implement the following methods
+ - `ObservedSkyModel(m::AbstractSkyModel, array::AbstractArrayConfiguration)`: Constructs an observed sky model
+    given the sky model `m` and the array configuration `array`. This method is used to compute the visibilities
+    and image of the sky model.
+
+The following methods have default implementations:
+ - `idealvisibilities(m::AbstractSkyModel, x)`: Computes the ideal visibilities of the sky model `m`
+    given the model parameters `x`.
+ - `skymodel(m::AbstractSkyModel, x)`: Returns the sky model image given the model parameters `x`.
+ - `domain(m::AbstractSkyModel)`: Returns the domain of the sky model `m`.
+ - `set_array(m::AbstractSkyModel, array::AbstractArrayConfiguration)`: Sets the array configuration
+    for the sky model `m` and returns the observed sky model and prior.
+ - `set_prior(m::AbstractSkyModel, array::AbstractArrayConfiguration)`: Sets the prior for the sky model
+    `m` given the array configuration `array`. This is used to set the priors for the model parameters.
+
+"""
+abstract type AbstractSkyModel end
+
+function Base.show(io::IO, mime::MIME"text/plain", m::AbstractSkyModel)
+    T = typeof(m)
+    ST = split(split(" $T", '{')[1], ".")[end]
+    printstyled(io, ST; bold = true, color = :blue)
+    println(io)
+    println(io, "  with map: $(skymodel(m))")
+    # GT = typeof(domain(m))
+    # SGT = split("$GT", '{')[1]
+    print(io, "   on grid: \n")
+    show(io, mime, domain(m))
+    return print(io, "\n   )\n")
+end
+
+skymodel(m::AbstractSkyModel) = getfield(m, :f)
+
+
+function set_array(m::AbstractSkyModel, array::AbstractArrayConfiguration)
+    return ObservedSkyModel(m, array), set_prior(m, array)
+end
+
+function domain(m::AbstractSkyModel; kwargs...)
+    return getfield(m, :grid)
+end
+
+"""
+    idealvisibilities(m::AbstractSkyModel, x)
+
+Computes the ideal non-corrupted visibilities of the sky model `m` given the model parameters `x`.
+"""
+function idealvisibilities(m::AbstractSkyModel, x)
+    skym = skymodel(m, x.sky)
+    return visibilitymap(skym, domain(m))
+end
+
+function skymodel(m::AbstractSkyModel, x)
+    return m.f(x, m.metadata)
+end
+
+function set_prior(m::AbstractSkyModel, array::AbstractArrayConfiguration)
+    return getfield(m, :prior)
+end
+
+struct ObservedSkyModel{F, G <: VLBISkyModels.AbstractDomain, M} <: AbstractSkyModel
+    f::F
+    grid::G
+    metadata::M
+end
+
+
+include("models.jl")
+include("fixed.jl")
+include("multi.jl")

src/skymodels/fixed.jl

@@ -0,0 +1,48 @@
+"""
+    FixedSkyModel(m::AbstractModel, grid::AbstractRectiGrid; algorithm = NFFTAlg())
+
+Construct a sky model that has no free parameters. This is useful for models where the
+image structure is known apriori but the instrument model is unknown.
+
+# Arguments
+
+ - `m` : The fixed sky model.
+- `grid` : The domain on which the model is defined. This defines the field of view and resolution
+   of the model. Note that if `f` produces a analytic model then this field of view isn't used
+   directly in the computation of the visibilities.
+
+# Optional Arguments
+- `algorithm` : The Fourier transform algorithm used to compute the visibilities. The default is
+   `NFFTAlg()` which uses a non-uniform fast Fourier transform. Other options can be found by using
+   `subtypes(VLBISkyModels.FourierTransform)`
+"""
+Base.@kwdef struct FixedSkyModel{M <: AbstractModel, G, A <: FourierTransform} <: AbstractSkyModel
+    model::M
+    grid::G
+    algorithm::A = NFFTAlg()
+end
+
+skymodel(m::FixedSkyModel) = m.model
+
+function FixedSkyModel(m::AbstractModel, grid::AbstractRectiGrid; algorithm = NFFTAlg())
+    return FixedSkyModel(m, grid, algorithm)
+end
+
+function ObservedSkyModel(m::FixedSkyModel, arr::AbstractArrayConfiguration)
+    gfour = FourierDualDomain(m.grid, arr, m.algorithm)
+    img = intensitymap(m.model, gfour)
+    vis = visibilitymap(m.model, gfour)
+    return ObservedSkyModel(m, gfour, (; img, vis))
+end
+
+function set_prior(::FixedSkyModel, ::AbstractArrayConfiguration)
+    return NamedTuple()
+end
+
+function idealvisibilities(m::ObservedSkyModel{<:FixedSkyModel}, x)
+    return m.metadata.vis
+end
+
+function skymodel(m::ObservedSkyModel{<:FixedSkyModel}, x)
+    return m.f.model
+end

src/skymodels/models.jl

@@ -1,25 +1,6 @@
 export SkyModel, FixedSkyModel
 
 
-"""
-    AbstractSkyModel
-
-The abstract type for Comrade Sky Models. For a concrete implementation see [`SkyModel`](@ref).
-
-Any subtype must implement the following methods
-
- - `set_array(m::AbstractSkyModel, array::AbstractArrayConfiguration)`: Sets the array configuration
-    for the sky model `m` and returns the observed sky model and prior.
-
-The following methods have default implementations:
- - `idealvisibilities(m::AbstractSkyModel, x)`: Computes the ideal visibilities of the sky model `m`
-    given the model parameters `x`.
- - `skymodel(m::AbstractSkyModel, x)`: Returns the sky model image given the model parameters `x`.
- - `domain(m::AbstractSkyModel)`: Returns the domain of the sky model `m`.
-"""
-abstract type AbstractSkyModel end
-
-
 struct SkyModel{F, P, G <: AbstractDomain, A <: FourierTransform, M} <: AbstractSkyModel
     f::F
     prior::P
@@ -28,16 +9,6 @@ struct SkyModel{F, P, G <: AbstractDomain, A <: FourierTransform, M} <: Abstract
     metadata::M
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", m::AbstractSkyModel)
-    T = typeof(m)
-    ST = split(split(" $T", '{')[1], ".")[end]
-    printstyled(io, ST; bold = true, color = :blue)
-    println(io)
-    println(io, "  with map: $(m.f)")
-    GT = typeof(m.grid)
-    SGT = split("$GT", '{')[1]
-    return print(io, "   on grid: $SGT")
-end
 
 """
     SkyModel(f, prior, grid::AbstractRectiGrid; algorithm = NFFTAlg(), metadata=nothing)
@@ -72,15 +43,6 @@ function VLBISkyModels.FourierDualDomain(grid::AbstractRectiGrid, array::Abstrac
     return FourierDualDomain(grid, domain(array; executor), alg)
 end
 
-struct ObservedSkyModel{F, G <: VLBISkyModels.AbstractDomain, M} <: AbstractSkyModel
-    f::F
-    grid::G
-    metadata::M
-end
-
-function domain(m::AbstractSkyModel; kwargs...)
-    return getfield(m, :grid)
-end
 
 # If we are using a analytic model then we don't need to plan the FT and we
 # can save some memory by not storing the plans.
@@ -110,69 +72,3 @@ function ObservedSkyModel(m::SkyModel, arr::AbstractArrayConfiguration)
     end
     return ObservedSkyModel(m.f, g, m.metadata)
 end
-
-
-function set_array(m::AbstractSkyModel, array::AbstractArrayConfiguration)
-    return ObservedSkyModel(m, array), m.prior
-end
-
-"""
-    idealvisibilities(m::AbstractSkyModel, x)
-
-Computes the ideal non-corrupted visibilities of the sky model `m` given the model parameters `x`.
-"""
-function idealvisibilities(m::AbstractSkyModel, x)
-    skym = skymodel(m, x.sky)
-    return visibilitymap(skym, domain(m))
-end
-
-function skymodel(m::AbstractSkyModel, x)
-    return m.f(x, m.metadata)
-end
-
-"""
-    FixedSkyModel(m::AbstractModel, grid::AbstractRectiGrid; algorithm = NFFTAlg())
-
-Construct a sky model that has no free parameters. This is useful for models where the
-image structure is known apriori but the instrument model is unknown.
-
-# Arguments
-
- - `m` : The fixed sky model.
-- `grid` : The domain on which the model is defined. This defines the field of view and resolution
-   of the model. Note that if `f` produces a analytic model then this field of view isn't used
-   directly in the computation of the visibilities.
-
-# Optional Arguments
-- `algorithm` : The Fourier transform algorithm used to compute the visibilities. The default is
-   `NFFTAlg()` which uses a non-uniform fast Fourier transform. Other options can be found by using
-   `subtypes(VLBISkyModels.FourierTransform)`
-"""
-Base.@kwdef struct FixedSkyModel{M <: AbstractModel, G, A <: FourierTransform} <: AbstractSkyModel
-    model::M
-    grid::G
-    algorithm::A = NFFTAlg()
-end
-
-function FixedSkyModel(m::AbstractModel, grid::AbstractRectiGrid; algorithm = NFFTAlg())
-    return FixedSkyModel(m, grid, algorithm)
-end
-
-function ObservedSkyModel(m::FixedSkyModel, arr::AbstractArrayConfiguration)
-    gfour = FourierDualDomain(m.grid, arr, m.algorithm)
-    img = intensitymap(m.model, gfour)
-    vis = visibilitymap(m.model, gfour)
-    return ObservedSkyModel(m, gfour, (; img, vis))
-end
-
-function set_array(m::FixedSkyModel, array::AbstractArrayConfiguration)
-    return ObservedSkyModel(m, array), NamedTuple()
-end
-
-function idealvisibilities(m::ObservedSkyModel{<:FixedSkyModel}, x)
-    return m.metadata.vis
-end
-
-function skymodel(m::ObservedSkyModel{<:FixedSkyModel}, x)
-    return m.f.model
-end

src/skymodels/multi.jl

@@ -0,0 +1,68 @@
+export MultiSkyModel
+
+"""
+    MultiSkyModel(skymodels::NamedTuple)
+
+Create a sky model that is a collection of multiple components. This is useful when
+you, e.g., want to decompose the sky into multiple grids, such as for wide-field imaging
+where the is a core component and a very far away component. 
+
+# Arguments
+ - `skymodels` : A named tuple of sky models, where each model is a [`SkyModel`](@ref).
+
+# Example
+```julia
+julia> s1 = SkyModel(...)
+julia> s2 = SkyModel(...)
+julia> stot = MultiSkyModel((core = m1, far = m2))
+julia> skymodel(stot, x)
+(core = ..., far = ...)
+```
+"""
+struct MultiSkyModel{N, T} <: AbstractSkyModel
+    skymodels::NamedTuple{N, T}
+end
+
+function ObservedSkyModel(m::MultiSkyModel, arr::AbstractArrayConfiguration)
+    skymodels = map(m.skymodels) do sm
+        ObservedSkyModel(sm, arr)
+    end
+    return MultiSkyModel(skymodels)
+end
+
+function set_prior(m::MultiSkyModel, array::AbstractArrayConfiguration)
+    prs = map(m.skymodels) do sm
+        set_prior(sm, array)
+    end
+    return prs
+end
+
+function idealvisibilities(m::MultiSkyModel{N}, x) where {N}
+    sm = m.skymodels
+    vis = map(N) do n
+        Base.@_inline_meta
+        @inline idealvisibilities(getproperty(sm, n), (; sky = getproperty(x.sky, n)))
+    end
+    return reduce(+, vis)
+end
+
+function domain(m::MultiSkyModel)
+    return map(m.skymodels) do sm
+        domain(sm)
+    end
+end
+
+function skymodel(m::MultiSkyModel{N}, x) where {N}
+    sm = m.skymodels
+    skyms = map(N) do n
+        Base.@_inline_meta
+        skymodel(getproperty(sm, n), getproperty(x, n))
+    end
+    return NamedTuple{N}(skyms)
+end
+
+function skymodel(m::MultiSkyModel)
+    return map(m.skymodels) do sm
+        skymodel(sm)
+    end
+end

test/Core/models.jl

@@ -13,6 +13,7 @@ using Tables
 using Plots
 import TransformVariables as TV
 using VLBIImagePriors
+using JET
 
 ntequal(x::NamedTuple{N}, y::NamedTuple{N}) where {N} = map(_ntequal, (x), (y))
 ntequal(x, y) = false
@@ -128,6 +129,32 @@ end
         @test Comrade.skymodel(oskym, x) == m
         @test Comrade.idealvisibilities(oskym, (; sky = x)) ≈ Comrade.idealvisibilities(oskyf, (; sky = x))
     end
+
+    @testset "MultiSkyModel" begin
+        _, vis, amp, lcamp, cphase = load_data()
+        skytot = MultiSkyModel((dynamic = skym, static = skyf))
+        oskytot, ptot = Comrade.set_array(skytot, arrayconfig(vis))
+
+        show(IOBuffer(), MIME"text/plain"(), skytot)
+
+
+        x = rand(Comrade.NamedDist(ptot))
+
+        oskym, = Comrade.set_array(skym, arrayconfig(vis))
+        oskyf, = Comrade.set_array(skyf, arrayconfig(vis))
+
+        tt = Comrade.skymodel(oskytot, x)
+        @test tt.dynamic == Comrade.skymodel(oskym, x.dynamic)
+        @test tt.static == Comrade.skymodel(oskyf, x.static)
+
+        vtot = Comrade.idealvisibilities(oskytot, (; sky = x))
+        vdyn = Comrade.idealvisibilities(oskym, (; sky = x.dynamic))
+        vstat = Comrade.idealvisibilities(oskyf, (; sky = x.static))
+
+        @test_opt Comrade.idealvisibilities(oskytot, (; sky = x))
+
+        @test vtot ≈ vdyn + vstat
+    end
 end
 
 @testset "GMRF" begin