Add frequency segmentation and make SiteArray indexing work

src/instrument/instrument.jl

@@ -10,12 +10,15 @@ struct IntegrationTime{T}
     dt::T
 end
 
-Base.in(t::Number, ts::IntegrationTime) = (ts.t0 - ts.dt/2) ≤ t < (ts.t0 + ts.dt)
+Base.in(t::Number, ts::IntegrationTime) = (ts.t0 - ts.dt/2) ≤ t < (ts.t0 + ts.dt/2)
 Base.isless(t::IntegrationTime, ts::IntegrationTime) = t.t0 < ts.t0
 Base.isless(s::Number, t::IntegrationTime) = s < (t.t0 - t.dt/2)
 Base.isless(t::IntegrationTime, s::Number) = (t.t0 + t.dt/2) < s
 mjd(ts::IntegrationTime) = ts.mjd
+Base.in(t::IntegrationTime, ::Base.Colon) = true
 
+_center(ts::IntegrationTime) = ts.t0
+_region(ts::IntegrationTime) = ts.dt
 
 struct FrequencyChannel{T, I<:Integer}
     central::T
@@ -24,6 +27,17 @@ struct FrequencyChannel{T, I<:Integer}
 end
 Base.in(f::Number, fs::FrequencyChannel) = (fs.central-fs.bandwidth/2) ≤ f < (fs.central+fs.bandwidth/2)
 channel(fs::FrequencyChannel) = fs.channel
+Base.isless(t::FrequencyChannel, ts::FrequencyChannel) = _center(t) < _center(ts)
+Base.isless(s::Number, t::FrequencyChannel) = s < (_center(t) - _region(t)/2)
+Base.isless(t::FrequencyChannel, s::Number) = (_center(t) + _region(t)/2) < s
+Base.in(x, t::FrequencyChannel) = in(t, x)
+Base.in(t::FrequencyChannel, ::Base.Colon) = true
+
+
+
+_center(fs::FrequencyChannel) = fs.central
+_region(fs::FrequencyChannel) = fs.bandwidth
+
 
 
 include("site_array.jl")

src/instrument/model.jl

@@ -195,8 +195,8 @@ function _construct_baselinemap(T, F, bl, x::SiteArray)
         t = T[i]
         f = F[i]
         s1, s2 = bl[i]
-        i1 = findall(x->(t∈x[1])&&(x[2]==s1), tsf)
-        i2 = findall(x->(t∈x[1])&&(x[2]==s2), tsf)
+        i1 = findall(x->(t∈x[1])&&(x[2]==s1)&&(f∈x[3]), tsf)
+        i2 = findall(x->(t∈x[1])&&(x[2]==s2)&&(f∈x[3]), tsf)
         length(i1) > 1 && throw(AssertionError("Multiple indices found for $t, $((s1)) in SiteArray"))
         length(i2) > 1 && throw(AssertionError("Multiple indices found for $t, $((s2)) in SiteArray"))
         (isnothing(i1) | isempty(i1)) && throw(AssertionError("$t, $f, $((s1)) not found in SiteArray"))

src/instrument/priors/array_priors.jl

@@ -71,40 +71,49 @@ function build_sitemap(d::ArrayPrior, array)
     # Ok to so this we are going to construct the schema first over sites.
     # At the end we may re-order depending on the schema ordering we want
     # to use.
-    tlists = map(keys(sites_prior)) do s
+    lists = map(keys(sites_prior)) do s
         seg = segmentation(sites_prior[s])
         # get all the indices where this site is present
         inds_s = findall(x->((x[1]==s)||x[2]==s), array[:sites])
         # Get all the unique times
         ts = unique(T[inds_s])
+        fs = unique(F[inds_s])
         # Now makes the acceptable time stamps given the segmentation
         tstamp = timestamps(seg, array)
+        fchan  = freqchannels(SpectralWindow(), array)
         # Now we find commonalities
         times = eltype(tstamp)[]
-        for t in tstamp
-            if any(x->x∈t, ts) && (!(t.t0 ∈ times))
+        freqs = eltype(fchan)[]
+        for t in tstamp, f in fchan
+            if any(x->x∈t, ts) && any(x->x∈f, fs) && ((!(t.t0 ∈ times)) || (!(f.central ∈ freqs)))
                 push!(times, t)
+                push!(freqs, f)
             end
         end
-        return times
+        return times, freqs
     end
+    tlists = first.(lists)
+    flists = last.(lists)
     # construct the schema
     slist = mapreduce((t,s)->fill(s, length(t)), vcat, tlists, keys(sites_prior))
     tlist = reduce(vcat, tlists)
+    flist = reduce(vcat, flists)
+
 
     tlistre = similar(tlist)
     slistre = similar(slist)
+    flistre = similar(flist)
     # Now rearrange so we have time site ordering (sites are the fastest changing)
     tuni = sort(unique(getproperty.(tlist, :t0)))
     ind0 = 1
     for t in tuni
         ind = findall(x->x.t0==t, tlist)
         tlistre[ind0:ind0+length(ind)-1] .= tlist[ind]
         slistre[ind0:ind0+length(ind)-1] .= slist[ind]
+        flistre[ind0:ind0+length(ind)-1] .= flist[ind]
         ind0 += length(ind)
     end
-    freqs = Fill(F[1], length(tlistre))
-    return SiteLookup(tlistre, freqs, slistre)
+    return SiteLookup(tlistre, flistre, slistre)
 end
 
 function ObservedArrayPrior(d::ArrayPrior, array::EHTArrayConfiguration)

src/instrument/priors/segmentation.jl

@@ -11,21 +11,25 @@ the following functions:
 """
 abstract type Segmentation end
 
+abstract type TimeSegmentation <: Segmentation end
+
+abstract type FrequencySegmentation <: Segmentation end
+
 # Track is for quantities that remain static across an entire observation
 """
     $(TYPEDEF)
 
 Data segmentation such that the quantity is constant over a `track`, i.e., the observation "night".
 """
-struct TrackSeg <: Segmentation end
+struct TrackSeg <: TimeSegmentation end
 
 # Scan is for quantities that are constant across a scan
 """
     $(TYPEDEF)
 
 Data segmentation such that the quantity is constant over a `scan`.
 """
-struct ScanSeg <: Segmentation end
+struct ScanSeg <: TimeSegmentation end
 
 
 # Integration is for quantities that change every integration time
@@ -35,10 +39,10 @@ struct ScanSeg <: Segmentation end
 Data segmentation such that the quantity is constant over the time stamps in the data.
 If the data is scan-averaged before then `IntegSeg` will be identical to `ScanSeg`.
 """
-struct IntegSeg <: Segmentation end
+struct IntegSeg <: TimeSegmentation end
 
 """
-    timestamps(seg::Segmentation, array::AbstractArrayConfiguration)
+    timestamps(seg::TimeSegmentation, array::AbstractArrayConfiguration)
 
 Return the time stamps or really a vector of integration time regions for a
 given segmentation scheme `seg` and array configuration `array`.
@@ -48,10 +52,11 @@ function timestamps end
 function timestamps(::ScanSeg, array)
     st     = array.scans
     mjd    = array.mjd
-
     # Shift the central time to the middle of the scan
     dt = (st.stop .- st.start)
+    dt[end] = dt[end]+0.5
     t0 = st.start .+ dt./2
+
     return IntegrationTime.(mjd, t0, dt)
 end
 
@@ -60,7 +65,6 @@ end
 
 function timestamps(::IntegSeg, array)
     ts = unique(array[:Ti])
-    st     = array.scans
     mjd    = array.mjd
 
     # TODO build in the dt into the data format
@@ -74,14 +78,26 @@ function timestamps(::IntegSeg, array)
 end
 
 function timestamps(::TrackSeg, array)
-    st     = array.scans
     mjd    = array.mjd
 
     tstart, tend = extrema(array[:Ti])
+    tend = tend + 1
     dt = tend - tstart
     if iszero(dt)
         dt = 1/3600
     end
     # TODO build in the dt into the data format
     return (IntegrationTime(mjd, (tend-tstart)/2 + tstart, dt),)
 end
+
+
+struct SpectralWindow <: FrequencySegmentation end
+
+
+function freqchannels(::SpectralWindow, array)
+    Fr = unique(array[:Fr])
+    if length(Fr) > 1
+        Fr[end] = nextfloat(Fr[end])
+    end
+    return FrequencyChannel.(Fr, array.bandwidth, 1:length(Fr))
+end

src/instrument/site_array.jl

@@ -16,7 +16,7 @@ which will grab the first 10 time and frequency points for the ALMA site.
 Otherwise indexing into the array will return an element whose time, frequency, and site are
 the element of the `times`, `frequencies`, and `sites` arrays respectively.
 """
-struct SiteArray{T, N, A<:AbstractArray{T,N}, Ti<:AbstractArray{<:IntegrationTime, N}, Fr<:AbstractArray{<:Number, N}, Sy<:AbstractArray{<:Any, N}} <: AbstractArray{T, N}
+struct SiteArray{T, N, A<:AbstractArray{T,N}, Ti<:AbstractArray{<:IntegrationTime, N}, Fr<:AbstractArray{<:FrequencyChannel, N}, Sy<:AbstractArray{<:Any, N}} <: AbstractArray{T, N}
     data::A
     times::Ti
     frequencies::Fr
@@ -111,47 +111,36 @@ function site(arr::SiteArray, p)
 end
 
 
-function time(arr::SiteArray, a::Union{AbstractInterval, IntegrationTime})
+function time(arr::SiteArray, a::Union{AbstractInterval, Real})
     inds = findall(in(a), times(arr))
     nd = view(parent(arr), inds)
     return SiteArray(nd, view(times(arr), inds), view(frequencies(arr), inds), view(sites(arr), inds))
 end
 
-function frequency(arr::SiteArray, a::Union{AbstractInterval, FrequencyChannel})
+function frequency(arr::SiteArray, a::Union{AbstractInterval, Real})
     inds = findall(in(a), times(arr))
     nd = view(parent(arr), inds)
     return SiteArray(nd, view(times(arr), inds), view(frequencies(arr), inds), view(sites(arr), inds))
 end
 
+_equalorin(x::T, y::T) where {T} = x == y 
+_equalorin(x::Real, y) = x ∈ y
+_equalorin(x, y::Real) = y ∈ x
+_equalorin(x, y) = y ∈ x
+_equalorin(x, ::typeof(Base.Colon())) = true
+const Indexable = Union{Integer, AbstractArray{<:Integer}, BitArray}
 
-@inline function _maybe_all(arr, X)
-    if X isa Base.Colon
-        ext = extrema(arr)
-        return ClosedInterval(ext[1], ext[2])
-    else
-        return X
-    end
-end
-
-function Base.getindex(arr::SiteArray; F=Base.Colon(), S=Base.Colon(), T=Base.Colon())
-    T2 = _maybe_all(times(arr), T)
-    F2 = _maybe_all(frequencies(arr), F)
-    S2 = S isa Base.Colon ? unique(sites(arr)) : S
-    return select_region(arr, S2, T2, F2)
-end
-
-function select_region(arr::SiteArray, S::Symbol, T::Union{Real, AbstractInterval}, F::Union{Real, AbstractInterval})
-    select_region(arr, (S,), T, F)
-end
-
-
-function select_region(arr::SiteArray, site, Ti::Union{Real, AbstractInterval}, Fr::Union{Real, AbstractInterval})
-    inds = findall(i->((Comrade.sites(arr)[i] ∈ site)&&(Ti ∈ Comrade.times(arr)[i])&&(Fr ∈ Comrade.frequencies(arr)[i])), eachindex(arr))
+function Base.getindex(arr::SiteArray; Fr=Base.Colon(), S=Base.Colon(), Ti=Base.Colon())
+    Fr2 = isa(Fr, Indexable) ? unique(arr.frequencies)[Fr] : Fr
+    S2 = isa(S, Indexable) ? unique(arr.sites)[S] : S
+    Ti2 = isa(Ti, Indexable) ? unique(arr.times)[Ti] : Ti
+    inds = findall(i->(_equalorin(S2, Comrade.sites(arr)[i])&&_equalorin(Ti2, Comrade.times(arr)[i])&&_equalorin(Fr2, Comrade.frequencies(arr)[i])), eachindex(arr))
     nd = view(parent(arr), inds)
     return SiteArray(nd, view(times(arr), inds), view(frequencies(arr), inds), view(sites(arr), inds))
 end
 
-struct SiteLookup{L<:NamedTuple, N, Ti<:AbstractArray{<:IntegrationTime, N}, Fr<:AbstractArray{<:Number, N}, Sy<:AbstractArray{<:Any, N}}
+
+struct SiteLookup{L<:NamedTuple, N, Ti<:AbstractArray{<:IntegrationTime, N}, Fr<:AbstractArray{<:FrequencyChannel, N}, Sy<:AbstractArray{<:Any, N}}
     lookup::L
     times::Ti
     frequencies::Fr
@@ -204,7 +193,10 @@ function SiteArray(a::AbstractArray, map::SiteLookup)
     return SiteArray(a, map.times, map.frequencies, map.sites)
 end
 
-function SiteArray(data::SiteArray{T, N}, times::AbstractArray{<:IntegrationTime, N}, frequencies::AbstractArray{<:Number, N}, sites::AbstractArray{<:Number, N}) where {T, N}
+function SiteArray(data::SiteArray{T, N},
+                   times::AbstractArray{<:IntegrationTime, N}, 
+                   frequencies::AbstractArray{<:FrequencyChannel, N}, 
+                   sites::AbstractArray{<:Number, N}) where {T, N}
     return data
 end