bug fixes

bin/pycbc_inspiral

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/data/gravwav/bgadre/cond_envs/pycbc-precomputed-psd/bin/python3.9
 
 # Copyright (C) 2014 Alex Nitz
 #
@@ -25,6 +25,7 @@ import numpy
 import itertools
 import time
 from pycbc.pool import BroadcastPool as Pool
+from icecream import ic
 
 import pycbc
 from pycbc import vetoes, psd, waveform, strain, scheme, fft, DYN_RANGE_FAC, events
@@ -450,16 +451,18 @@ with ctx:
     for inc, seg in enumerate(segments):
         seg /= seg.psd
 
-        if hasattr(opt, 'psd_output') and opt.psd_output:
+
+        if hasattr(opt, 'associated_psd_output') and opt.associated_psd_output:
             key = str(inc)
-            start.append(int(seg.start))
-            end.append(int(seg.end))
+            # start_time = gwstrain.start_time + seg.analyze.start/gwstrain.sample_rate
+            # end_time = gwstrain.start_time + seg.analyze.stop/gwstrain.sample_rate
+            # ic(start_time, end_time)
+            start.append(int(seg.start_time))
+            end.append(int(seg.end_time))
             psd_group.create_dataset(key, data=seg.psd, compression='gzip',
                                  compression_opts=9, shuffle=True)
-            psd_group[key].attrs['epoch'] = int(seg.start)
+            psd_group[key].attrs['epoch'] = int(seg.start_time)
             psd_group[key].attrs['delta_f'] = seg.psd.delta_f
-            start.append(int(seg.start))
-            end.append(int(seg.end))
 
     if hasattr(opt, 'associated_psd_output') and opt.associated_psd_output:
         f[ifo + '/start_time'] = numpy.array(start, dtype=numpy.uint32)

pycbc/psd/read.py

@@ -225,8 +225,13 @@ def assosiate_psd_to_inspiral_segment(self, inp_seg, delta_f=None):
             err_msg += "PSDs are within range ({}, {})".format(
                 self.begin, self.end)
             raise ValueError(err_msg)
-        sidx = numpy.argpartition(
-            numpy.abs(self.start_times - inp_seg[0]), 2)[:2]
+
+        if len(self.start_times) > 2:
+            sidx = numpy.argpartition(
+                numpy.abs(self.start_times - inp_seg[0]), 2)[:2]
+        else:
+            sidx = np.array([0, 1])
+
         nearest = segments.segment(
             self.start_times[sidx[0]], self.end_times[sidx[0]])
         next_nearest = segments.segment(
@@ -252,7 +257,6 @@ def assosiate_psd_to_inspiral_segment(self, inp_seg, delta_f=None):
                                                   self.sample_rate),
                                               low_frequency_cutoff=self.f_low,
                                               trunc_method=self.invpsd_trunc_method)
-
         return best_psd
 
     def get_psd(self, index, delta_f=None):