From 629bbce957cb4711672b55d5b4d24401e5ac72bb Mon Sep 17 00:00:00 2001
From: Stephen Bailey <stephen.code@gmail.com>
Date: Fri, 11 Aug 2023 11:35:59 -0700
Subject: [PATCH] read_spectra(filename, rows=...) option

---
 py/desispec/io/spectra.py        | 11 +++++++++--
 py/desispec/test/test_spectra.py | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/py/desispec/io/spectra.py b/py/desispec/io/spectra.py
index fdbf66a75..a6b6443f7 100644
--- a/py/desispec/io/spectra.py
+++ b/py/desispec/io/spectra.py
@@ -201,6 +201,7 @@ def read_spectra(
     infile,
     single=False,
     targetids=None,
+    rows=None,
     skip_hdus=None,
     select_columns={
         "FIBERMAP": None,
@@ -219,6 +220,7 @@ def read_spectra(
         infile (str): path to read
         single (bool): if True, keep spectra as single precision in memory.
         targetids (list): Optional, list of targetids to read from file, if present.
+        rows (list): Optional, list of rows to read from file
         skip_hdus (list): Optional, list/set/tuple of HDUs to skip
         select_columns (dict): Optional, dictionary to select column names to be read. Default, all columns are read.
 
@@ -230,6 +232,8 @@ def read_spectra(
     Note that WAVE, FLUX, and IVAR are always required.
 
     If a table HDU is not listed in `select_columns`, all of its columns will be read
+
+    User can optionally specify targetids OR rows, but not both
     """
     log = get_logger()
     infile = checkgzip(infile)
@@ -245,14 +249,17 @@ def read_spectra(
     hdus = fitsio.FITS(infile, mode="r")
     nhdu = len(hdus)
 
+    if targetids is not None and rows is not None:
+        raise ValueError('Set rows or targetids but not both')
+
     if targetids is not None:
         targetids = np.atleast_1d(targetids)
         file_targetids = hdus["FIBERMAP"].read(columns="TARGETID")
         rows = np.where(np.isin(file_targetids, targetids))[0]
         if len(rows) == 0:
             return Spectra()
-    else:
-        rows = None
+    elif rows is not None:
+        rows = np.asarray(rows)
 
     if skip_hdus is None:
         skip_hdus = set()  #- empty set, include everything
diff --git a/py/desispec/test/test_spectra.py b/py/desispec/test/test_spectra.py
index 1cad15664..b20701074 100644
--- a/py/desispec/test/test_spectra.py
+++ b/py/desispec/test/test_spectra.py
@@ -227,6 +227,25 @@ def test_read_targetids(self):
         # targetid 10 doesn't appear because it wasn't in the input file, ok
         self.assertTrue(np.all(comp_subset.fibermap['TARGETID'] == np.array([2,2,4,4,4,0,0,0,0])))
 
+    def test_read_rows(self):
+        """Test reading specific rows"""
+
+        # manually create the spectra and write
+        spec = Spectra(bands=self.bands, wave=self.wave, flux=self.flux,
+            ivar=self.ivar, mask=self.mask, resolution_data=self.res,
+            fibermap=self.fmap1, meta=self.meta, extra=self.extra)
+
+        write_spectra(self.fileio, spec)
+
+        rows = [1,3]
+        subset = read_spectra(self.fileio, rows=rows)
+        self.assertTrue(np.all(spec.fibermap[rows] == subset.fibermap))
+
+        with self.assertRaises(ValueError):
+            subset = read_spectra(self.fileio, rows=rows, targetids=[1,2])
+
+
+
     def test_read_columns(self):
         """test reading while subselecting columns"""
         # manually create the spectra and write