Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Multifrequency support for instrument modeling #383

Open
wants to merge 35 commits into
base: main
Choose a base branch
from
Open
Changes from 1 commit
Commits
Show all changes
35 commits
Select commit Hold shift + click to select a range
b78cc0f
Fix up site array
ptiede Dec 16, 2024
116d83d
Small adjustment to gradient
ptiede Dec 16, 2024
8baa5c0
Fix up site array
ptiede Dec 16, 2024
d71f3c2
Small adjustment to gradient
ptiede Dec 16, 2024
8b2d72c
Merge branch 'ptiede-frequencysitearr' of https://github.com/ptiede/C…
ptiede Dec 16, 2024
5d4d2b0
Add frequency segmentation and make SiteArray indexing work
ptiede Dec 16, 2024
7d5946b
Fix indexing again
ptiede Dec 16, 2024
c5b1085
Small bug fix to prevent SO
ptiede Dec 16, 2024
eab1d25
Fix forward_jones to work with frequency dependent instrument terms
ptiede Dec 17, 2024
2d08fe2
Update to Pyehtim 0.2
ptiede Dec 17, 2024
8ecc3ad
Fix Pyehtim extension
ptiede Dec 17, 2024
0e48002
Allow Comrade to work with new Pyehtim version
ptiede Dec 18, 2024
bb877ac
Make sitelookup aware of difference frequencies
ptiede Dec 18, 2024
d6d2744
Fix dependencies in examples
ptiede Dec 18, 2024
c2786b7
Test multifrequency
ptiede Dec 18, 2024
dedacd2
Start making CalTable frequency friendly
ptiede Dec 18, 2024
3fc98d0
Caltable now plots gains with frequency labels
ptiede Dec 19, 2024
4d52d6f
small bug fix in caltable and tests
ptiede Dec 19, 2024
2fa66a6
Fix CalTable tests
ptiede Dec 19, 2024
295ba32
Fix bug in caltable tests
ptiede Dec 19, 2024
1443391
Fix really dumb scan times issue
ptiede Dec 19, 2024
eca1e0c
Silly copy pasta mistake
ptiede Dec 19, 2024
b4688af
Fix this
ptiede Dec 19, 2024
39fa84f
finish up tests
ptiede Dec 21, 2024
c5f434e
tests passing locally
ptiede Dec 21, 2024
c35e479
rule was deleted
ptiede Dec 21, 2024
08b56d0
Update geometric model
ptiede Dec 21, 2024
27439fe
Small updates to tther examples
ptiede Dec 21, 2024
cc0f8c6
Remove the Pkg.develop in GeometricModeling since there is some stran…
ptiede Dec 22, 2024
0926991
Remove CondaPkg.toml
ptiede Dec 22, 2024
4a27ccd
Improve coverage
ptiede Dec 22, 2024
e0691b7
Try again
ptiede Dec 22, 2024
b624938
Make docs directly in literate
ptiede Dec 22, 2024
fe398fb
try tutorials again
ptiede Dec 22, 2024
7e37869
make analytic models take up less memory
ptiede Dec 22, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Make sitelookup aware of difference frequencies
  • Loading branch information
ptiede committed Dec 18, 2024
commit bb877ac9bd1ff43b14c5282d1136cd6520625da7
25 changes: 19 additions & 6 deletions src/instrument/instrument_transforms.jl
Original file line number Diff line number Diff line change
@@ -5,9 +5,12 @@ inner_transform(t::AbstractInstrumentTransform) = t.inner_transform
function TV.transform_with(flag::TV.LogJacFlag, m::AbstractInstrumentTransform, x, index)
y, ℓ, index = _instrument_transform_with(flag, m, x, index)
sm = m.site_map
return SiteArray(y, sm.times, sm.frequencies, sm.sites), ℓ, index
return SiteArray(y, sm), ℓ, index
end

EnzymeRules.inactive_type(::Type{AbstractArray{<:IntegrationTime}}) = true
EnzymeRules.inactive_type(::Type{AbstractArray{<:FrequencyChannel}}) = true

function TV.inverse_at!(x::AbstractArray, index, t::AbstractInstrumentTransform, y::SiteArray)
itrf = inner_transform(t)
return TV.inverse_at!(x, index, itrf, parent(y))
@@ -50,15 +53,25 @@ end
return yout, ℓ, index
end

EnzymeRules.inactive_type(::Type{<:SiteLookup}) = true

@inline function site_sum(y, site_map::SiteLookup)
yout = similar(y)
@inbounds for site in site_map.lookup
# yout = similar(y)
vals = values(lookup(site_map))
@inbounds for site in vals
# i0 = site[begin]
# yout[i0] = y[i0]
# acc = zero(eltype(y))
# for idx in site
# acc += y[idx]
# yout[idx] = acc
# end
ys = @view y[site]
# y should never alias so we should be fine here.
youts = @view yout[site]
cumsum!(youts, ys)
# youts = @view yout[site]
cumsum!(ys, ys)
end
return yout
return y
end

# function ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(_instrument_transform_with), flag, m::MarkovInstrumentTransform, x, index)
15 changes: 7 additions & 8 deletions src/instrument/priors/array_priors.jl
Original file line number Diff line number Diff line change
@@ -76,21 +76,20 @@ function build_sitemap(d::ArrayPrior, array)
# 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])
ts = T[inds_s]
fs = F[inds_s]
tfs = zip(ts, fs)
# Now makes the acceptable time stamps given the segmentation
tstamp = timestamps(seg, array)
fchan = freqchannels(SpectralWindow(), array)
# Now we find commonalities
times = eltype(tstamp)[]
freqs = eltype(fchan)[]
tf = Tuple{eltype(tstamp), 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)
if any(x->(x[1]∈t && x[2]∈f), tfs) && ((!((t.t0, f.central) ∈ tf)))
push!(tf, (t, f))
end
end
return times, freqs
return first.(tf), last.(tf)
end
tlists = first.(lists)
flists = last.(lists)
11 changes: 8 additions & 3 deletions src/instrument/priors/refant.jl
Original file line number Diff line number Diff line change
@@ -43,14 +43,19 @@ end
function reference_indices(array::AbstractArrayConfiguration, st::SiteLookup, r::SEFDReference)
tarr = array.tarr
t = unique(st.times)
f = unique(st.frequencies)
sefd = NamedTuple{Tuple(tarr.sites)}(Tuple(tarr.SEFD1 .+ tarr.SEFD2))
fixedinds = map(eachindex(t)) do i
inds = findall(==(t[i]), st.times)
fixedinds = Int[]
for i in eachindex(t), j in eachindex(f)
inds = findall(x->((st.times[x]==t[i])&&(st.frequencies[x]==f[j])), eachindex(st.times))
if isempty(inds)
continue
end
sites = Tuple(st.sites[inds])
@assert length(sites) <= length(sefd) "Error in reference site generation. Too many sites"
sp = select(sefd, sites)
_, ind = findmin(values(sp))
return inds[ind]
push!(fixedinds, inds[ind])
end
return fixedinds, Fill(r.value, length(fixedinds))
end
41 changes: 36 additions & 5 deletions src/instrument/site_array.jl
Original file line number Diff line number Diff line change
@@ -148,8 +148,18 @@ struct SiteLookup{L<:NamedTuple, N, Ti<:AbstractArray{<:IntegrationTime, N}, Fr<
sites::Sy
end

times(s::SiteLookup) = s.times
sites(s::SiteLookup) = s.sites
frequencies(s::SiteLookup) = s.frequencies
lookup(s::SiteLookup) = s.lookup

EnzymeRules.inactive(::typeof(times), ::SiteLookup) = nothing
EnzymeRules.inactive(::typeof(frequencies), ::SiteLookup) = nothing
EnzymeRules.inactive(::typeof(sites), ::SiteLookup) = nothing
EnzymeRules.inactive(::typeof(lookup), ::SiteLookup) = nothing

function sitemap!(f, out::AbstractArray, gains::AbstractArray, slook::SiteLookup)
map(slook.lookup) do site
map(lookup(slook)) do site
ysite = @view gains[site]
outsite = @view out[site]
outsite .= f.(ysite)
@@ -163,9 +173,10 @@ function sitemap(f, gains::AbstractArray{T}, slook::SiteLookup) where {T}
end

function sitemap!(::typeof(cumsum), out::AbstractArray, gains::AbstractArray, slook::SiteLookup)
map(slook.lookup) do site
map(lookup(slook)) do site
ys = @view gains[site]
cumsum!(ys, ys)
nothing
end
return out
end
@@ -180,8 +191,28 @@ function SiteLookup(s::SiteArray)
end

function SiteLookup(times::AbstractVector, frequencies::AbstractArray, sites::AbstractArray)
slist = Tuple(sort(unique(sites)))
return SiteLookup(NamedTuple{slist}(map(p->findall(==(p), sites), slist)), times, frequencies, sites)
slist = sort(unique(sites))
flist = sort(unique(frequencies))
if length(flist) == 1
return SiteLookup(NamedTuple{Tuple(slist)}(map(p->findall(==(p), sites), slist)), times, frequencies, sites)
else
# Find sites first
sflist = Symbol[]
inds = Vector{Int}[]
for s in slist
sinds = findall(==(s), sites)
for (i,f) in enumerate(flist)
finds = findall(==(f), @view(frequencies[sinds]))
if !isempty(finds)
ss = Symbol(s, i)
push!(sflist, ss)
push!(inds, finds)
end
end
end
lookup = NamedTuple{Tuple(sflist)}(Tuple(inds))
return SiteLookup(lookup, times, frequencies, sites)
end
end

"""
@@ -191,7 +222,7 @@ Construct a site array with the entries `arr` and the site ordering implied by
`sitelookup`.
"""
function SiteArray(a::AbstractArray, map::SiteLookup)
return SiteArray(a, map.times, map.frequencies, map.sites)
return SiteArray(a, times(map), frequencies(map), sites(map))
end

function SiteArray(data::SiteArray{T, N},
54 changes: 54 additions & 0 deletions test/Core/models.jl
Original file line number Diff line number Diff line change
@@ -408,6 +408,60 @@ end

end


@testset "Multifrequency" begin
dcoh2 = deepcopy(dcoh)
dcoh2.config[:Fr][200:end] .= 345e9
vis = CoherencyMatrix.(Comrade.measurement(dcoh2), Ref(CirBasis()))

G = JonesG() do x
gR = exp(x.lgR + 1im*x.gpR)
gL = gR*exp(x.lgrat + 1im*x.gprat)
return gR, gL
end

D = JonesD() do x
dR = complex(x.dRx, x.dRy)
dL = complex(x.dLx, x.dLy)
return dR, dL
end

R = JonesR(;add_fr=true)
J = JonesSandwich(*, G, D, R)
intprior = (
lgR = ArrayPrior(IIDSitePrior(ScanSeg(), Normal(0.0, 0.1))),
gpR = ArrayPrior(IIDSitePrior(ScanSeg(), Normal(0.0, inv(π ^2))); phase=true, refant=SEFDReference(0.0)),
lgrat= ArrayPrior(IIDSitePrior(ScanSeg(), Normal(0.0, 0.1)), phase=false),
gprat= ArrayPrior(IIDSitePrior(ScanSeg(), Normal(0.0, 0.1))),
dRx = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
dRy = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
dLx = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
dLy = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
)
intm = InstrumentModel(J, intprior)
ointm, printm = Comrade.set_array(intm, arrayconfig(dcoh))

@testset "ObservedArrayPrior" begin
@inferred logpdf(printm, rand(printm))
x = rand(printm)
@test logpdf(printm, x) ≈ logpdf(printm2, x)
@test asflat(printm) isa TV.AbstractTransform
p = rand(printm)
t = asflat(printm)
pout = TV.transform(t, TV.inverse(t, p))
dp = ntequal(p, pout)
@test dp.lgR
@test dp.lgrat
@test dp.gprat
@test dp.dRx
@test dp.dRy
@test dp.dLx
@test dp.dLy
end


end

@testset "Integration" begin
_,dvis, amp, lcamp, cphase, dcoh = load_data()
ts = Comrade.timestamps(ScanSeg(), arrayconfig(dvis))
Loading