Skip to content

Commit

Permalink
work on SDI for contrast calculations
Browse files Browse the repository at this point in the history
  • Loading branch information
scexao6 committed Mar 13, 2024
1 parent 7e7b21c commit 9242686
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 13 deletions.
1 change: 1 addition & 0 deletions Project.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ authors = ["Miles Lucas <[email protected]>"]
version = "0.9.1"

[deps]
BiweightStats = "5bf8a1e9-d5f8-4697-9608-80edd97af0ad"
ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
CoordinateTransformations = "150eb455-5306-5404-9cee-2592286d6298"
Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
Expand Down
15 changes: 10 additions & 5 deletions src/metrics/contrast.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ using LinearAlgebra: dot
using Photometry
using ProgressLogging
using StaticKernels
using Setfield

"""
contrast_curve(alg, cube, angles, psf;
Expand Down Expand Up @@ -56,7 +57,7 @@ function contrast_curve(alg, cube, angles, psf, args...;
# measure the noise and throughput in consecutive resolution elements
# across azimuthal branches
@info "Calculating Throughput"
reduced_empty = alg(cube, angles, args...; kwargs...)
reduced_empty = alg(cube, angles, args...; angles=angles, kwargs...)

through, meta = throughput(alg, cube, angles, psf, args...;
fwhm=fwhm, inner_rad=inner_rad, fc_rad_sep=fc_rad_sep, theta=theta,
Expand Down Expand Up @@ -285,9 +286,9 @@ function _fix_range(radii, cube::MultiAnnulusView)
end


function throughput(alg, cube::AbstractArray{T,4}, angles, psf_model, scales;
function throughput(alg, cube::AbstractArray{T,4}, angles, psf_model::AbstractArray{V,3}, scales;
fwhm, nbranch=1, theta=0, inner_rad=1, fc_rad_sep=3,
snr=100, reduced_empty = nothing, flux_scale=ones(size(cube, 3)), kwargs...) where T
snr=100, reduced_empty = nothing, kwargs...) where {T,V}
maxfcsep = minimum(size(cube)[begin:begin + 1]) ÷ (2 * fwhm) - 1
# too large separation between companions in the radial patterns
3 ≤ fc_rad_sep ≤ maxfcsep || error("`fc_rad_sep` should lie ∈[3, $(maxfcsep)], got $fc_rad_sep")
Expand All @@ -312,6 +313,7 @@ function throughput(alg, cube::AbstractArray{T,4}, angles, psf_model, scales;
angle_per_branch = 360 / nbranch
output = similar(cube, length(radii), nbranch)

mean_psf = mean(psf_model, dims=3)[:, :, begin]
tmp_cube = similar(cube)
fake_comps_full = zero(reduced_empty)
tmp_frame = similar(reduced_empty)
Expand All @@ -332,9 +334,11 @@ function throughput(alg, cube::AbstractArray{T,4}, angles, psf_model, scales;

A = snr * noise[ann]

inject!(tmp_frame, psf_model; x, y, amp=mean(A .* flux_scale), fwhm)
inject!(tmp_frame, mean_psf; x, y, amp=A, fwhm)
for wl_idx in axes(tmp_cube, 3)
inject!(tmp_cube[:, :, wl_idx, :], psf_model, angles; x, y, amp=A * flux_scale[wl_idx], fwhm)
wave_slice = @view tmp_cube[:, :, wl_idx, :]
psf_slice = @view psf_model[:, :, wl_idx]
inject!(wave_slice, psf_slice, angles; x, y, amp=A, fwhm)
end

return CircularAperture(x, y, fwhm / 2)
Expand All @@ -354,6 +358,7 @@ function throughput(alg, cube::AbstractArray{T,4}, angles, psf_model, scales;
return output, (distance=radii, fake_comps=fake_comps_full, noise=noise)
end


"""
throughput(alg, cube, angles, psf, position;
fwhm, snr=100,
Expand Down
38 changes: 30 additions & 8 deletions src/sdi.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,13 +85,17 @@ end

DoubleSDI(alg) = DoubleSDI(alg, alg)

function process(sdi::DoubleSDI, spcube::AbstractArray{T,4}, angles, scales; method=:deweight, kwargs...) where T
function process(sdi::DoubleSDI, spcube::AbstractArray{T,4}, angles, scales; method=:deweight, flux_scale=ones(T, size(spcube, 3)),kwargs...) where T
nx, ny, nλ, n = size(spcube)
spec_resids = similar(spcube, nx, ny, n)
# do first pass in spectral domain
Threads.@threads for tidx in axes(spcube, 4)
cube = @view spcube[:, :, :, tidx]
scaled_cube = scale(cube, scales)
# apply scaling factor to better subtract speckles
for wl_idx in axes(scaled_cube, 3)
@. scaled_cube[:, :, wl_idx] *= flux_scale[wl_idx]
end
angs = Zeros(nλ)
if :ref in keys(kwargs)
cube_ref = view(kwargs[:ref], :, :, :, tidx)
Expand All @@ -101,6 +105,10 @@ function process(sdi::DoubleSDI, spcube::AbstractArray{T,4}, angles, scales; met
R = subtract(sdi.alg_spec, scaled_cube; angles=angs, kwargs...)
end
spec_resid = invscale(R, scales)
# renormalize by dividing scale factor
for wl_idx in axes(spec_resid, 3)
@. spec_resid[:, :, wl_idx] /= flux_scale[wl_idx]
end
spec_resids[:, :, tidx] .= collapse(spec_resid)
end
# do second pass in temporal domain
Expand All @@ -125,14 +133,14 @@ julia> data, angles, scales = # load data...
julia> res = SliceSDI(Classic(), GreeDS(15))(data, angles, scales)
```
"""
struct SliceSDI{ALG<:ADIAlgorithm, ALG2<:ADIAlgorithm} <: SDIAlgorithm
struct SliceSDI{ALG<:Union{ADIAlgorithm,Nothing}, ALG2<:ADIAlgorithm} <: SDIAlgorithm
alg_spec::ALG
alg_temp::ALG2
end

SliceSDI(alg) = SliceSDI(alg, alg)

function process(sdi::SliceSDI, spcube::AbstractArray{T,4}, angles, scales; kwargs...) where T
function process(sdi::SliceSDI, spcube::AbstractArray{T,4}, angles, scales; flux_scale=ones(T, size(spcube, 3)), kwargs...) where T
nλ = size(spcube, 3)
temp_resids = similar(spcube, Base.front(size(spcube)))
# do first pass in temporal domain
Expand All @@ -145,9 +153,23 @@ function process(sdi::SliceSDI, spcube::AbstractArray{T,4}, angles, scales; kwar
temp_resids[:, :, waveidx] .= sdi.alg_temp(cube, angles; kwargs...)
end
end
# do second pass in temporal domain
scaled_resid_cube = scale(temp_resids, scales)
angs = Zeros(nλ)
resid = sdi.alg_spec(scaled_resid_cube, angs; angles=angs, kwargs...)
return invscale(resid, maximum(scales))
# do second pass in spectral domain
# scale residuals by ~wavelength (resids have alread been derotated!)
if sdi.alg_spec !== nothing
scaled_resid_cube = scale(temp_resids, scales)
for wl_idx in axes(scaled_resid_cube, 3)
@. scaled_resid_cube[:, :, wl_idx] *= flux_scale[wl_idx]
end
resid_cube = subtract(sdi.alg_spec, scaled_resid_cube; angles=Zeros(nλ), kwargs...)
# invscale back to astrometric
rescaled_resid = invscale(resid_cube, scales)
# collapse
# apply flux scaling before collapsing
for wl_idx in axes(rescaled_resid, 3)
@. rescaled_resid[:, :, wl_idx] /= flux_scale[wl_idx]
end
return collapse(rescaled_resid)
else
return collapse(temp_resids)
end
end

0 comments on commit 9242686

Please sign in to comment.