Implement GTI Reading and Handling

Project.toml

@@ -14,6 +14,8 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Metadata = "4fb893c1-3164-4f58-823a-cb4c64eabb4f"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb"
+Intervals = "d8418881-c3e1-53bb-8760-2df7ec849ed5"
 
 [compat]
 julia = "1.7"
@@ -26,7 +28,8 @@ DataFrames = "1.3"
 Metadata = "0.3"
 HDF5 = "0.16"
 NaNMath = "0.3"
-
+FITSIO = "0.16"
+Intervals = "1.8"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/Stingray.jl

@@ -1,6 +1,7 @@
 module Stingray
 
-using ResumableFunctions, StatsBase, Statistics, DataFrames, FFTW, Metadata, NaNMath
+using ResumableFunctions, StatsBase, Statistics, DataFrames
+using FFTW, Metadata, NaNMath, FITSIO, Intervals
 using ProgressBars: tqdm as show_progress 
 
 include("fourier.jl")
@@ -20,6 +21,12 @@ export avg_pds_from_events
 export avg_cs_from_events
 
 include("gti.jl")
+export load_gtis
+export get_total_gti_length
+export create_gti_mask
+export create_gti_from_condition
+export operations_on_gtis
+export get_btis
 export time_intervals_from_gtis
 export bin_intervals_from_gtis
 

src/gti.jl

@@ -1,3 +1,194 @@
+function get_total_gti_length(gti::AbstractMatrix{<:Real}; minlen::Real=0.0)
+    lengths = diff(gti; dims =2)
+    return sum(x->x > minlen ? x : zero(x), lengths)
+end
+
+function load_gtis(fits_file::String, gtistring::String="GTI")
+    lchdulist = FITS(fits_file)
+    gtihdu = lchdulist[gtistring]
+    gti = get_gti_from_hdu(gtihdu)
+    close(lchdulist)
+    return gti
+end
+
+function get_gti_from_hdu(gtihdu::TableHDU)
+
+    if "START" in FITSIO.colnames(gtihdu)
+        startstr = "START"
+        stopstr = "STOP"
+    else
+        startstr = "Start"
+        stopstr = "Stop"
+    end
+
+    gtistart = read(gtihdu,startstr)
+    gtistop = read(gtihdu,stopstr)
+
+    return mapreduce(permutedims, vcat, 
+    [[a, b] for (a,b) in zip(gtistart, gtistop)])
+end
+
+function check_gtis(gti::AbstractMatrix)
+
+    if ndims(gti) != 2 || size(gti,2) != 2
+        throw(ArgumentError("Please check the formatting of the GTIs. 
+        They need to be provided as [[gti00 gti01]; [gti10 gti11]; ...]."))
+    end
+
+    gti_start = @view gti[:, 1]
+    gti_end = @view gti[:, 2]
+
+    if any(gti_end < gti_start)
+        throw(ArgumentError(
+            "The GTI end times must be larger than the GTI start times."
+        )) 
+    end
+
+    if any(@view(gti_start[2:end]) < @view(gti_end[1:end-1]))
+        throw(ArgumentError(
+            "This GTI has overlaps"
+        ))
+    end
+end
+
+function create_gti_mask(times::AbstractVector{<:Real},gtis::AbstractMatrix{<:Real};
+                         safe_interval::AbstractVector{<:Real}=[0,0], min_length::Real=0,
+                         dt::Real = -1, epsilon::Real = 0.001)
+
+    if length(times) == 0
+        throw(ArgumentError("Passing an empty time array to create_gti_mask"))
+    end
+
+    check_gtis(gtis)
+    mask = zeros(Bool,length(times))
+
+    if min_length>0
+        gtis = gtis[min_length .< @view(gtis[:,2]) - @view(gtis[:,1]),:]
+
+        if size(gtis,1) < 1
+            @warn "No GTIs longer than min_length $(min_length)"
+            return mask, gtis
+        end
+    end   
+
+    if dt < 0
+        dt = Statistics.median(diff(times))
+    end
+    epsilon_times_dt = epsilon * dt
+
+    new_gtis = [[0.0, 0.0] for _ in range(1,size(gtis,1))]
+    new_gti_mask = zeros(Bool, size(gtis,1))
+
+    gti_start = @view gtis[:, 1]
+    gti_end = @view gtis[:, 2]
+
+    for (ig,(limmin,limmax)) in enumerate(zip(gti_start,gti_end))
+        limmin += safe_interval[1]
+        limmax -= safe_interval[2]
+        if limmax - limmin >= min_length
+            new_gtis[ig][:] .= limmin, limmax
+            for (i,t) in enumerate(times) 
+                if t >= (limmin + dt / 2 - epsilon_times_dt) && t <= (limmax - dt / 2 + epsilon_times_dt)
+                mask[i] = true
+                end
+            end
+            new_gti_mask[ig] = true
+        end
+    end
+
+    return mask, mapreduce(permutedims, vcat, keepat!(new_gtis,new_gti_mask))
+end
+
+function create_gti_from_condition(time::AbstractVector{<:Real}, condition::AbstractVector{Bool};
+    safe_interval::AbstractVector{<:Real}=[0,0], dt::AbstractVector{<:Real}=Float64[])
+
+    if length(time) != length(condition)
+        throw(ArgumentError("The length of the condition and 
+        time arrays must be the same."))
+    end
+
+    idxs = contiguous_regions(condition)
+
+    if(isempty(dt))
+        dt = zero(time) .+ (time[2] .- time[1]) ./ 2
+    end
+
+    gtis = Vector{Float64}[]
+    for idx in eachrow(idxs)
+        startidx = idx[1]
+        stopidx = idx[2] - 1
+
+        t0 = time[startidx] - dt[startidx] + safe_interval[1]
+        t1 = time[stopidx] + dt[stopidx] - safe_interval[2]
+        if t1 - t0 < 0
+            continue
+        end
+        push!(gtis,[t0, t1])
+    end
+    return mapreduce(permutedims, vcat, gtis)
+end
+
+function operations_on_gtis(gti_list::AbstractVector{<:AbstractMatrix{T}}, 
+                            operation::Function) where {T<:Real}
+
+    required_interval = nothing
+
+    for gti in gti_list
+        check_gtis(gti)
+
+        combined_gti = Interval{T}[]
+        for ig in eachrow(gti)
+            push!(combined_gti,Interval{Closed,Open}(ig[1],ig[2]))
+        end
+        if isnothing(required_interval)
+            required_interval = IntervalSet(combined_gti)
+        else
+            required_interval = operation(required_interval, IntervalSet(combined_gti))
+        end
+    end
+
+    final_gti = Vector{T}[]
+
+    for interval in required_interval.items
+        push!(final_gti, [first(interval), last(interval)])
+    end
+
+    return mapreduce(permutedims, vcat, final_gti)
+end
+
+function get_btis(gtis::AbstractMatrix{<:Real})
+    if length(gtis) == 0
+        throw(ArgumentError("Empty GTI and no valid start_time and stop_time"))
+    end
+    return get_btis(gtis, gtis[1,1], gtis[end,2])
+end
+
+function get_btis(gtis::AbstractMatrix{T}, start_time, stop_time) where {T<:Real}
+    if length(gtis) == 0
+        return T[start_time stop_time]
+    end
+    check_gtis(gtis)
+
+    total_interval = Interval{T, Closed, Open}[Interval{T, Closed, Open}(start_time, stop_time)]
+    total_interval_set = IntervalSet(total_interval)
+
+    gti_interval = Interval{T, Closed, Open}[]
+    for ig in eachrow(gtis)
+        push!(gti_interval,Interval{T, Closed, Open}(ig[1],ig[2]))
+    end
+    gti_interval_set = IntervalSet(gti_interval)
+
+    bti_interval_set = setdiff(total_interval_set, gti_interval_set)
+
+    btis = Vector{T}[]
+
+    for interval in bti_interval_set.items
+        push!(btis, [first(interval), last(interval)])
+    end
+
+    return mapreduce(permutedims, vcat, btis)
+end
+
 function time_intervals_from_gtis(gtis::AbstractMatrix{<:Real}, segment_size::Real;
                                   fraction_step::Real=1, epsilon::Real=1e-5)  
     spectrum_start_times = Float64[]

src/utils.jl

@@ -4,3 +4,26 @@ function sum_if_not_none_or_initialize(A,B)
     end
     return A + B
 end
+
+function contiguous_regions(condition::AbstractVector{Bool})
+    # Find the indicies of changes in "condition"
+    d = diff(condition)
+    idx = findall(x->x != 0, d)
+
+    # We need to start things after the change in "condition". Therefore, 
+    # we'll shift the index by 1 to the right.
+    idx .+= 1
+
+    if condition[1]
+        # If the start of condition is True prepend a 0
+        idx = append!([1],idx)
+    end
+
+    if condition[end]
+        # If the end of condition is True, append the length of the array
+        idx = append!(idx,[condition.size+1]) # Edit
+    end
+
+    # Reshape the result into two columns
+    return reshape(idx,2,length(idx) ÷ 2)'
+end