Fix spectrum plots for different states (#403)

* fix _get_spectrum function to consider state

* add state option to get_range_spectrum

* remove operations reduntant to decorator

* fix range spctrogram segments

* remove extra line

* remove unused import

* Apply suggestions from code review

Co-authored-by: Evan Goetz <32753745+eagoetz@users.noreply.github.com>

---------

Co-authored-by: Iara Ota <iara.ota@ligo.org>
Co-authored-by: Evan Goetz <32753745+eagoetz@users.noreply.github.com>

gwsumm/data/range.py

@@ -117,21 +117,37 @@ def get_range_spectrogram(channel, segments, config=None, cache=None,
             add_spectrogram(outspec if 'energy' in rangekwargs else
                             outspec**(1/2.), key=key)
 
-    if return_:
-        return globalv.SPECTROGRAMS.get(key, SpectrogramList())
+    if not return_:
+        return
+
+    # return correct data, according to state segment
+    out = SpectrogramList()
+    for specgram in globalv.SPECTROGRAMS[key]:
+        for seg in segments:
+            if abs(seg) < specgram.dt.value:
+                continue
+            if specgram.span.intersects(seg):
+                common = specgram.span & type(seg)(seg[0],
+                                                   seg[1] + specgram.dt.value)
+                s = specgram.crop(*common)
+                if s.shape[0]:
+                    out.append(s)
+    return out.coalesce()
 
 
 @use_segmentlist
 def get_range_spectrum(channel, segments, config=None, cache=None, query=True,
                        nds=None, return_=True, nproc=1, datafind_error='raise',
                        frametype=None, stride=60, fftlength=None, overlap=None,
-                       method=None, which='all', **rangekwargs):
+                       method=None, which='all', state=None, **rangekwargs):
     """Compute percentile spectra of the range integrand from a set of
     spectrograms
     """
     name = str(channel)
-    cmin = '%s.min' % name
-    cmax = '%s.max' % name
+    if state:
+        name += f',{state}'
+    cmin = f'{name}.min'
+    cmax = f'{name}.max'
 
     if name not in globalv.SPECTRUM:
         speclist = get_range_spectrogram(
@@ -164,4 +180,4 @@ def get_range_spectrum(channel, segments, config=None, cache=None, query=True,
         return globalv.SPECTRUM[cmin]
     if which == 'max':
         return globalv.SPECTRUM[cmax]
-    raise ValueError("Unrecognised value for `which`: %r" % which)
+    raise ValueError(f"Unrecognised value for `which`: {which}")

gwsumm/data/spectral.py

@@ -34,7 +34,6 @@
 
 from astropy import units
 
-from gwpy.segments import DataQualityFlag
 from gwpy.frequencyseries import FrequencySeries
 from gwpy.spectrogram import SpectrogramList
 
@@ -328,17 +327,11 @@ def apply_transfer_function_series(specgram, tfunc):
 def get_spectrum(channel, segments, config=None, cache=None,
                  query=True, nds=None, format='power', return_=True,
                  frametype=None, nproc=1, datafind_error='raise',
-                 **fftparams):
-    """Retrieve the time-series and generate a spectrogram of the given
+                 state=None, **fftparams):
+    """Retrieve the time-series and generate a spectrum of the given
     channel
     """
     channel = get_channel(channel)
-    if isinstance(segments, DataQualityFlag):
-        name = ','.join([channel.ndsname, segments.name])
-        segments = segments.active
-    else:
-        name = channel.ndsname
-    name += ',%s' % format
 
     # read data for all sub-channels
     specs = []
@@ -351,6 +344,7 @@ def get_spectrum(channel, segments, config=None, cache=None,
                                    return_=return_, frametype=frametype,
                                    nproc=nproc,
                                    datafind_error=datafind_error,
+                                   state=state,
                                    **fftparams))
     if return_ and len(channels) == 1:
         return specs[0]
@@ -363,19 +357,17 @@ def get_spectrum(channel, segments, config=None, cache=None,
 
 def _get_spectrum(channel, segments, config=None, cache=None, query=True,
                   nds=None, format='power', return_=True, which='all',
-                  **fftparams):
-    """Retrieve the time-series and generate a spectrogram of the given
+                  state=None, **fftparams):
+    """Retrieve the time-series and generate a spectrum of the given
     channel
     """
     channel = get_channel(channel)
-    if isinstance(segments, DataQualityFlag):
-        name = ','.join([channel.ndsname, segments.name])
-        segments = segments.active
-    else:
-        name = channel.ndsname
-    name += ',%s' % format
-    cmin = '%s.min' % name
-    cmax = '%s.max' % name
+
+    name = f'{channel.ndsname},{format}'
+    if state:
+        name = f'{channel.ndsname},{state},{format}'
+    cmin = f'{name}.min'
+    cmax = f'{name}.max'
 
     if name not in globalv.SPECTRUM:
         if os.path.isfile(channel.ndsname):
@@ -424,4 +416,4 @@ def _get_spectrum(channel, segments, config=None, cache=None, query=True,
         return globalv.SPECTRUM[cmin]
     if which == 'max':
         return globalv.SPECTRUM[cmax]
-    raise ValueError("Unrecognised value for `which`: %r" % which)
+    raise ValueError(f"Unrecognised value for `which`: {which}")

gwsumm/plot/builtin.py

@@ -512,10 +512,11 @@ def _draw(self):
                 data = get_coherence_spectrum(
                     [str(channel), str(channel2)], valid, query=False)
             elif self.type == 'range-spectrum':
-                data = get_range_spectrum(str(channel), valid, query=False)
+                data = get_range_spectrum(str(channel), valid, query=False,
+                                          state=valid)
             elif self.type == 'cumulative-range-spectrum':
-                data = get_range_spectrum(
-                    str(channel), valid, query=False, which='mean')
+                data = get_range_spectrum(str(channel), valid, query=False,
+                                          which='mean', state=valid)
                 if str(data.unit) == 'Mpc':
                     data = (data**3).cumsum() ** (1/3.)
                 else:
@@ -527,7 +528,8 @@ def _draw(self):
             else:
                 try:
                     data = get_spectrum(str(channel), valid, query=False,
-                                        format=sdform, method=method)
+                                        format=sdform, method=method,
+                                        state=valid)
                 except ValueError as exc:
                     # math op failed beacuse one of the datasets is empty
                     if (