From 24011bf373391801fadc03bac26467148755ad5f Mon Sep 17 00:00:00 2001
From: Bhooshan Gadre <bhooshanudaygadre@gmail.com>
Date: Wed, 6 Nov 2024 14:20:21 +0100
Subject: [PATCH] bug fixes

---
 bin/pycbc_inspiral | 17 ++++++++++-------
 pycbc/psd/read.py  | 10 +++++++---
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/bin/pycbc_inspiral b/bin/pycbc_inspiral
index 993c4d471fe..3e191fec239 100644
--- a/bin/pycbc_inspiral
+++ b/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)
diff --git a/pycbc/psd/read.py b/pycbc/psd/read.py
index 2f0162050d7..0e7470cb124 100644
--- a/pycbc/psd/read.py
+++ b/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):