diff --git a/pycbc/strain/gate.py b/pycbc/strain/gate.py index 90f9f02309d..eceefba1861 100644 --- a/pycbc/strain/gate.py +++ b/pycbc/strain/gate.py @@ -17,6 +17,8 @@ """ from . import strain +from scipy import linalg + def _gates_from_cli(opts, gate_opt): """Parses the given `gate_opt` into something understandable by @@ -137,3 +139,41 @@ def add_gate_option_group(parser): "prior to FFT-ing the data for PSD " "estimation.") return gate_group + + +def gate_and_paint(data, lindex, rindex, invpsd, copy=True): + """Gates and in-paints data. + + Parameters + ---------- + data : TimeSeries + The data to gate. + lindex : int + The start index of the gate. + rindex : int + The end index of the gate. + invpsd : FrequencySeries + The inverse of the PSD. + copy : bool, optional + Copy the data before applying the gate. Otherwise, the gate will + be applied in-place. Default is True. + + Returns + ------- + TimeSeries : + The gated and in-painted time series. + """ + # Copy the data and zero inside the hole + if copy: + data = data.copy() + data[lindex:rindex] = 0 + + # get the over-whitened gated data + tdfilter = invpsd.astype('complex').to_timeseries() * invpsd.delta_t + owhgated_data = (data.to_frequencyseries() * invpsd).to_timeseries() + + # remove the projection into the null space + proj = linalg.solve_toeplitz(tdfilter[:(rindex - lindex)], + owhgated_data[lindex:rindex]) + data[lindex:rindex] -= proj + return data