From 9ad278ae4521f3513a0892b9d0777ab0683f9877 Mon Sep 17 00:00:00 2001
From: Paul Price <price@astro.princeton.edu>
Date: Tue, 3 May 2022 15:49:59 -0400
Subject: [PATCH 1/5] use double precision only for wavelength

Wavelength needs to be in double precision in order to preserve
the precision in the measurements, but other quantities (esp.
flux) don't need double precision, and datamodel I/O has been
reverted to single-precision for these to save disk space.
This commit adapts to these changes.
---
 include/pfs/drp/stella/FiberTrace.h   |  2 +-
 include/pfs/drp/stella/Spectrum.h     |  4 ++--
 include/pfs/drp/stella/SpectrumSet.h  |  4 ++--
 python/pfs/drp/stella/fiberProfile.py |  8 +++++---
 python/pfs/drp/stella/findLines.py    |  2 +-
 python/pfs/drp/stella/fitContinuum.py |  6 +++---
 python/pfs/drp/stella/fitLine.cc      |  5 +++--
 python/pfs/drp/stella/mergeArms.py    |  2 +-
 src/FiberTrace.cc                     |  2 +-
 tests/test_PfsArm.py                  | 14 +++++++-------
 tests/test_Spectrum.py                | 10 +++++-----
 tests/test_SpectrumSet.py             |  8 ++++----
 tests/test_extractSpectra.py          | 10 +++++-----
 tests/test_fitLine.py                 |  2 +-
 tests/test_profileNorm.py             |  2 +-
 15 files changed, 42 insertions(+), 39 deletions(-)

diff --git a/include/pfs/drp/stella/FiberTrace.h b/include/pfs/drp/stella/FiberTrace.h
index 23649696a..707e9a55d 100644
--- a/include/pfs/drp/stella/FiberTrace.h
+++ b/include/pfs/drp/stella/FiberTrace.h
@@ -127,7 +127,7 @@ class FiberTrace {
         ndarray::Array<double, 2, 1> const& profiles,
         ndarray::Array<bool, 2, 1> const& good,
         ndarray::Array<double, 1, 1> const& centers,
-        ndarray::Array<double, 1, 1> const& norm=ndarray::Array<double, 1, 1>()
+        ndarray::Array<Spectrum::ImageT, 1, 1> const& norm=ndarray::Array<Spectrum::ImageT, 1, 1>()
     );
 
   private:
diff --git a/include/pfs/drp/stella/Spectrum.h b/include/pfs/drp/stella/Spectrum.h
index e2083c21c..3e77ea7a4 100644
--- a/include/pfs/drp/stella/Spectrum.h
+++ b/include/pfs/drp/stella/Spectrum.h
@@ -13,8 +13,8 @@ namespace pfs { namespace drp { namespace stella {
  */
 class Spectrum {
   public:
-    using ImageT = double;
-    using VarianceT = double;
+    using ImageT = float;
+    using VarianceT = float;
     using Mask = lsst::afw::image::Mask<lsst::afw::image::MaskPixel>;
     using ImageArray = ndarray::Array<ImageT, 1, 1>;
     using ConstImageArray = ndarray::Array<const ImageT, 1, 1>;
diff --git a/include/pfs/drp/stella/SpectrumSet.h b/include/pfs/drp/stella/SpectrumSet.h
index a3905e64a..b58221f7d 100644
--- a/include/pfs/drp/stella/SpectrumSet.h
+++ b/include/pfs/drp/stella/SpectrumSet.h
@@ -20,8 +20,8 @@ class SpectrumSet {
     using SpectrumPtr = std::shared_ptr<Spectrum>;
     using ConstSpectrumPtr = std::shared_ptr<Spectrum const>;
     using Collection = std::vector<SpectrumPtr>;
-    using ImageArray = ndarray::Array<double, 2, 1>;
-    using CovarianceArray = ndarray::Array<double, 3, 1>;
+    using ImageArray = ndarray::Array<Spectrum::ImageT, 2, 1>;
+    using CovarianceArray = ndarray::Array<Spectrum::ImageT, 3, 1>;
     using MaskArray = ndarray::Array<lsst::afw::image::MaskPixel, 2, 1>;
     using WavelengthArray = ndarray::Array<double, 2, 1>;
     using iterator = Collection::iterator;
diff --git a/python/pfs/drp/stella/fiberProfile.py b/python/pfs/drp/stella/fiberProfile.py
index 9df34d8ef..4bf4f4d0b 100644
--- a/python/pfs/drp/stella/fiberProfile.py
+++ b/python/pfs/drp/stella/fiberProfile.py
@@ -27,7 +27,7 @@ class FiberProfile:
     profiles : array_like of `float`, shape ``(N, M)``
         Profiles for each swath, each of width
         ``M = 2*(radius + 1)*oversample + 1``.
-    norm : array_like of `float`, optional
+    norm : array_like of `np.float32`, optional
         Normalisation for each spectral pixel.
     """
     def __init__(self, radius, oversample, rows, profiles, norm=None):
@@ -220,7 +220,8 @@ def makeFiberTraceFromDetectorMap(self, detectorMap, fiberId):
         """
         return FiberTrace.fromProfile(fiberId, detectorMap.bbox.getDimensions(), self.radius,
                                       self.oversample, self.rows, self.profiles, ~self.profiles.mask,
-                                      detectorMap.getXCenter(fiberId), self.norm)
+                                      detectorMap.getXCenter(fiberId),
+                                      self.norm.astype(np.float32) if self.norm is not None else None)
 
     def makeFiberTrace(self, dimensions, centerFunc, fiberId):
         """Make a FiberTrace object
@@ -244,7 +245,8 @@ def makeFiberTrace(self, dimensions, centerFunc, fiberId):
         """
         rows = np.arange(dimensions.getY(), dtype=float)
         return FiberTrace.fromProfile(fiberId, dimensions, self.radius, self.oversample, self.rows,
-                                      self.profiles, ~self.profiles.mask, centerFunc(rows), self.norm)
+                                      self.profiles, ~self.profiles.mask, centerFunc(rows),
+                                      self.norm.astype(np.float32) if self.norm is not None else None)
 
     def extractSpectrum(self, maskedImage, detectorMap, fiberId):
         """Extract a single spectrum from an image
diff --git a/python/pfs/drp/stella/findLines.py b/python/pfs/drp/stella/findLines.py
index 9fb57cda1..1fa4bb377 100644
--- a/python/pfs/drp/stella/findLines.py
+++ b/python/pfs/drp/stella/findLines.py
@@ -125,7 +125,7 @@ def convolve(self, spectrum, continuum=None):
         size = 2*halfSize + 1
         xx = np.arange(size, dtype=float) - halfSize
         sigma = self.config.width
-        kernel = np.exp(-0.5*xx**2/sigma**2)/sigma/np.sqrt(2.0*np.pi)
+        kernel = (np.exp(-0.5*xx**2/sigma**2)/sigma/np.sqrt(2.0*np.pi)).astype(spectrum.flux.dtype)
 
         flux = np.convolve(spectrum.normFlux if continuum is None else spectrum.normFlux - continuum,
                            kernel, mode="same")
diff --git a/python/pfs/drp/stella/fitContinuum.py b/python/pfs/drp/stella/fitContinuum.py
index d8fa58565..bdf5e4948 100644
--- a/python/pfs/drp/stella/fitContinuum.py
+++ b/python/pfs/drp/stella/fitContinuum.py
@@ -141,13 +141,13 @@ def _fitContinuumImpl(self, values, good):
         fit : `numpy.ndarray`, floating-point
             Fit array.
         """
-        indices = np.arange(len(values), dtype=float)
+        indices = np.arange(len(values), dtype=values.dtype)
         knots, binned = binData(indices, values, good, self.config.numKnots)
         use = np.isfinite(knots) & np.isfinite(binned)
         if not np.any(use):
             raise FitContinuumError("No finite knots when fitting continuum")
         interp = makeInterpolate(knots[use], binned[use], self.fitType)
-        fit = np.array(interp.interpolate(indices))
+        fit = np.array(interp.interpolate(indices)).astype(values.dtype)
 
         if lsstDebug.Info(__name__).plot:
             import matplotlib.pyplot as plt
@@ -285,4 +285,4 @@ def binData(xx, yy, good, numBins):
         select = good[low:high]
         xBinned[ii] = np.median(xx[low:high][select]) if np.any(select) else np.nan
         yBinned[ii] = np.median(yy[low:high][select]) if np.any(select) else np.nan
-    return xBinned, yBinned
+    return xBinned.astype(xx.dtype), yBinned.astype(yy.dtype)
diff --git a/python/pfs/drp/stella/fitLine.cc b/python/pfs/drp/stella/fitLine.cc
index df08a6f13..bfdb36486 100644
--- a/python/pfs/drp/stella/fitLine.cc
+++ b/python/pfs/drp/stella/fitLine.cc
@@ -36,9 +36,10 @@ PYBIND11_PLUGIN(fitLine) {
     lsst::utils::python::addOutputOp(cls, "__repr__");
 
     mod.def("fitLine",
-            py::overload_cast<ndarray::Array<double const, 1, 1> const&,
+            py::overload_cast<ndarray::Array<Spectrum::ImageT const, 1, 1> const&,
                               ndarray::Array<lsst::afw::image::MaskPixel const, 1, 1> const&,
-                              float, float, lsst::afw::image::MaskPixel, std::size_t>(&fitLine<double>),
+                              float, float, lsst::afw::image::MaskPixel,
+                              std::size_t>(&fitLine<Spectrum::ImageT>),
             "flux"_a, "mask"_a, "peakPosition"_a, "rmsSize"_a, "badBitMask"_a,
             "fittingHalfSize"_a=0);
 
diff --git a/python/pfs/drp/stella/mergeArms.py b/python/pfs/drp/stella/mergeArms.py
index 26861dddd..67ab8ea5b 100644
--- a/python/pfs/drp/stella/mergeArms.py
+++ b/python/pfs/drp/stella/mergeArms.py
@@ -328,7 +328,7 @@ def normalizeSpectra(self, spectra):
         assert all(np.all(ss.fiberId == fiberId) for ss in spectra)  # Consistent fibers
 
         # Collect normalisations from each arm, interpolated to common wavelength sampling
-        norm = np.zeros((fiberId.size, wavelength.size))
+        norm = np.zeros((fiberId.size, wavelength.size), dtype=np.float32)
         for ss in spectra:
             badBitmask = ss.flags.get(*self.config.mask)
             for ii in range(len(ss)):
diff --git a/src/FiberTrace.cc b/src/FiberTrace.cc
index e1fd62c17..ba2fca83e 100644
--- a/src/FiberTrace.cc
+++ b/src/FiberTrace.cc
@@ -82,7 +82,7 @@ FiberTrace<ImageT, MaskT, VarianceT> FiberTrace<ImageT, MaskT, VarianceT>::fromP
     ndarray::Array<double, 2, 1> const& profiles,
     ndarray::Array<bool, 2, 1> const& good,
     ndarray::Array<double, 1, 1> const& centers,
-    ndarray::Array<double, 1, 1> const& norm
+    ndarray::Array<Spectrum::ImageT, 1, 1> const& norm
 ) {
     int const width = dims.getX();
     std::size_t const height = dims.getY();
diff --git a/tests/test_PfsArm.py b/tests/test_PfsArm.py
index a1bf0b586..3146423aa 100644
--- a/tests/test_PfsArm.py
+++ b/tests/test_PfsArm.py
@@ -30,11 +30,11 @@ def setUp(self):
         self.identity = Identity(self.visit, self.arm, self.spectrograph, self.pfsDesignId)
         self.fiberId = self.synthConfig.fiberId
         self.wavelength = np.vstack([np.linspace(600, 900, self.synthConfig.height)]*num)
-        self.flux = self.rng.uniform(size=shape)
-        self.mask = self.rng.uniform(high=2**31, size=shape).astype(np.int32)
-        self.sky = self.rng.uniform(size=shape)
-        self.norm = self.rng.uniform(size=shape)
-        self.covar = self.rng.uniform(size=(num, 3, self.synthConfig.height))
+        self.flux = self.rng.uniform(size=shape).astype(np.float32)
+        self.mask = self.rng.uniform(high=2**31, size=shape).astype(np.uint32)
+        self.sky = self.rng.uniform(size=shape).astype(np.float32)
+        self.norm = self.rng.uniform(size=shape).astype(np.float32)
+        self.covar = self.rng.uniform(size=(num, 3, self.synthConfig.height)).astype(np.float32)
         self.flags = MaskHelper()
         self.metadata = dict(FOO=12345, BAR=0.9876)  # FITS keywords get capitalised
 
@@ -48,9 +48,9 @@ def makeSpectra(self):
     def assertSpectra(self, spectra):
         """Check that the spectra match what's expected"""
         for name in ("visit", "arm", "spectrograph", "pfsDesignId"):
-            self.assertEqual(getattr(spectra.identity, name), getattr(self, name))
+            self.assertEqual(getattr(spectra.identity, name), getattr(self, name), msg=name)
         for name in ("fiberId", "wavelength", "flux", "mask", "sky", "norm", "covar"):
-            self.assertFloatsEqual(getattr(spectra, name), getattr(self, name))
+            self.assertFloatsEqual(getattr(spectra, name), getattr(self, name), msg=name)
         self.assertDictEqual(spectra.flags.flags, self.flags.flags)
         self.assertDictEqual({**self.metadata, **spectra.metadata}, spectra.metadata)
 
diff --git a/tests/test_Spectrum.py b/tests/test_Spectrum.py
index de7281b63..ed5417d26 100644
--- a/tests/test_Spectrum.py
+++ b/tests/test_Spectrum.py
@@ -23,12 +23,12 @@ def setUp(self):
         # Spectrum
         self.length = 123
         self.fiberId = 456
-        self.image = self.rng.uniform(size=self.length).astype(float)
+        self.image = self.rng.uniform(size=self.length).astype(np.float32)
         self.mask = lsst.afw.image.Mask(self.length, 1)
         self.mask.array[:] = self.rng.randint(0, 2**30, self.length).astype(np.int32)
-        self.background = self.rng.uniform(size=self.length).astype(float)
-        self.norm = self.rng.uniform(size=self.length).astype(float)
-        self.covariance = self.rng.uniform(size=(3, self.length)).astype(float)
+        self.background = self.rng.uniform(size=self.length).astype(np.float32)
+        self.norm = self.rng.uniform(size=self.length).astype(np.float32)
+        self.covariance = self.rng.uniform(size=(3, self.length)).astype(np.float32)
         self.wavelengthArray = np.arange(1, self.length + 1, dtype=float)
 
     def makeSpectrum(self):
@@ -145,7 +145,7 @@ def testBasics(self):
         # Set the variance and ensure it changes the covariance
         spectrum.covariance[:] = self.covariance
         self.assertFloatsEqual(spectrum.covariance, self.covariance)
-        variance = self.rng.uniform(size=self.length)
+        variance = self.rng.uniform(size=self.length).astype(np.float32)
         spectrum.variance = variance
         self.covariance[0] = variance
         self.assertFloatsEqual(spectrum.covariance, self.covariance)
diff --git a/tests/test_SpectrumSet.py b/tests/test_SpectrumSet.py
index 84ecbddbe..0c2154ec5 100644
--- a/tests/test_SpectrumSet.py
+++ b/tests/test_SpectrumSet.py
@@ -21,12 +21,12 @@ def setUp(self):
         # Spectrum
         self.length = 123
         self.fiberId = 456
-        self.image = self.rng.uniform(size=self.length).astype(float)
+        self.image = self.rng.uniform(size=self.length).astype(np.float32)
         self.mask = lsst.afw.image.Mask(self.length, 1)
         self.mask.array[:] = self.rng.randint(0, 2**30, self.length).astype(np.int32)
-        self.background = self.rng.uniform(size=self.length).astype(float)
-        self.norm = self.rng.uniform(size=self.length).astype(float)
-        self.covariance = self.rng.uniform(size=(3, self.length)).astype(float)
+        self.background = self.rng.uniform(size=self.length).astype(np.float32)
+        self.norm = self.rng.uniform(size=self.length).astype(np.float32)
+        self.covariance = self.rng.uniform(size=(3, self.length)).astype(np.float32)
         self.wavelengthArray = np.arange(1, self.length + 1, dtype=float)
 
     def makeSpectrum(self):
diff --git a/tests/test_extractSpectra.py b/tests/test_extractSpectra.py
index aaa10f46e..0ad1c4506 100644
--- a/tests/test_extractSpectra.py
+++ b/tests/test_extractSpectra.py
@@ -37,7 +37,7 @@ def setUp(self):
             lsst.afw.image.makeExposure(self.image), detectorMap=self.detMap
         ).profiles
         for fiberId in self.fiberProfiles:
-            self.fiberProfiles[fiberId].norm = np.full(self.synthConfig.height, self.flux, dtype=float)
+            self.fiberProfiles[fiberId].norm = np.full(self.synthConfig.height, self.flux, dtype=np.float32)
         self.fiberTraces = self.fiberProfiles.makeFiberTracesFromDetectorMap(self.detMap)
         self.assertEqual(len(self.fiberTraces), self.synthConfig.numFibers)
 
@@ -78,7 +78,7 @@ def assertSpectra(self, spectra, flux=None, mask=None):
                 expectMask = 0
 
             bbox = ft.trace.getBBox()
-            expectNorm = np.zeros(self.synthConfig.height, dtype=float)
+            expectNorm = np.zeros(self.synthConfig.height, dtype=np.float32)
             expectNorm[bbox.getMinY():bbox.getMaxY() + 1] = ft.trace.image.array.sum(axis=1)
 
             self.assertEqual(ss.fiberId, ft.fiberId)
@@ -127,7 +127,7 @@ def testUndersizedFiberTrace(self):
         newTrace = type(trace)(trace.trace.Factory(trace.trace, newBox), trace.fiberId)
         self.fiberTraces[index] = newTrace
         spectra = self.fiberTraces.extractSpectra(self.image)
-        expectFlux = np.zeros(self.synthConfig.height, dtype=float)
+        expectFlux = np.zeros(self.synthConfig.height, dtype=np.float32)
         expectFlux[middle:] = self.flux
         expectMask = np.zeros(self.synthConfig.height, dtype=np.int32)
         expectMask[:middle] = trace.trace.mask.getPlaneBitMask("NO_DATA")
@@ -138,7 +138,7 @@ def testUndersizedFiberTrace(self):
         newTrace = type(trace)(trace.trace.Factory(trace.trace, newBox), trace.fiberId)
         self.fiberTraces[index] = newTrace
         spectra = self.fiberTraces.extractSpectra(self.image)
-        expectFlux = np.zeros(self.synthConfig.height, dtype=float)
+        expectFlux = np.zeros(self.synthConfig.height, dtype=np.float32)
         expectFlux[:middle] = self.flux
         expectMask = np.zeros(self.synthConfig.height, dtype=np.int32)
         expectMask[middle:] = trace.trace.mask.getPlaneBitMask("NO_DATA")
@@ -155,7 +155,7 @@ def testSingular(self):
         index = self.synthConfig.numFibers//2
         row = self.synthConfig.height//2
         fiberId = self.synthConfig.fiberId[index]
-        expectFlux = np.full(self.synthConfig.height, self.flux, dtype=float)
+        expectFlux = np.full(self.synthConfig.height, self.flux, dtype=np.float32)
         expectMask = np.zeros(self.synthConfig.height, dtype=np.int32)
         expectFlux[row] = 0.0
         expectMask[row] = self.image.mask.getPlaneBitMask(["BAD_FIBERTRACE", "NO_DATA"])
diff --git a/tests/test_fitLine.py b/tests/test_fitLine.py
index 1c71f5878..8e8499aef 100644
--- a/tests/test_fitLine.py
+++ b/tests/test_fitLine.py
@@ -61,7 +61,7 @@ def makeSpectrum(length, center, amplitude, rmsSize, bgConst=0.0, bgSlope=0.0):
     return spectrum
 
 
-class FindLinesTestCase(lsst.utils.tests.TestCase):
+class FitLinesTestCase(lsst.utils.tests.TestCase):
     def setUp(self):
         self.length = 256
         self.center = 123.45
diff --git a/tests/test_profileNorm.py b/tests/test_profileNorm.py
index a3707c585..3d5ab2cba 100644
--- a/tests/test_profileNorm.py
+++ b/tests/test_profileNorm.py
@@ -59,7 +59,7 @@ def testProfileNorm(self):
         newSpectra = traces.extractSpectra(image)
         for new, old in zip(newSpectra, spectra):
             self.assertFloatsAlmostEqual(new.norm, old.normFlux, rtol=1.0e-7)
-            self.assertFloatsAlmostEqual(new.normFlux, 1.0, atol=1.0e-7)
+            self.assertFloatsAlmostEqual(new.normFlux, 1.0, atol=1.0e-6)
 
 
 class TestMemory(lsst.utils.tests.MemoryTestCase):

From ef89ac5687ddf6e5242b3359d18b9e5681227b62 Mon Sep 17 00:00:00 2001
From: Paul Price <price@astro.princeton.edu>
Date: Mon, 23 May 2022 14:46:17 -0400
Subject: [PATCH 2/5] symmetricTridiagonal: extend python bindings for
 workspace class

Helps to have access to the workspace in python to understand problems.
---
 python/pfs/drp/stella/symmetricTridiagonal.cc | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/pfs/drp/stella/symmetricTridiagonal.cc b/python/pfs/drp/stella/symmetricTridiagonal.cc
index 10a029853..1ceabc781 100644
--- a/python/pfs/drp/stella/symmetricTridiagonal.cc
+++ b/python/pfs/drp/stella/symmetricTridiagonal.cc
@@ -16,8 +16,13 @@ namespace {
 
 template <typename T>
 void declareSymmetricTridiagonal(py::module &mod, std::string const& suffix) {
-    py::class_<SymmetricTridiagonalWorkspace<T>> Workspace(
+    py::class_<SymmetricTridiagonalWorkspace<T>> cls(
         mod, ("SymmetricTridiagonalWorkspace" + suffix).c_str());
+    cls.def(py::init<>());
+    cls.def("reset", &SymmetricTridiagonalWorkspace<T>::reset, "num"_a);
+    cls.def_readonly("longArray1", &SymmetricTridiagonalWorkspace<T>::longArray1);
+    cls.def_readonly("longArray2", &SymmetricTridiagonalWorkspace<T>::longArray2);
+    cls.def_readonly("shortArray", &SymmetricTridiagonalWorkspace<T>::shortArray);
 
     mod.def(("solveSymmetricTridiagonal" + suffix).c_str(), &solveSymmetricTridiagonal<T>,
             "diagonal"_a, "offDiag"_a, "answer"_a, "workspace"_a=SymmetricTridiagonalWorkspace<T>());

From 89327652c57feac237843703a17a20b87d60aceb Mon Sep 17 00:00:00 2001
From: Paul Price <price@astro.princeton.edu>
Date: Mon, 23 May 2022 14:48:25 -0400
Subject: [PATCH 3/5] symmetricTridiagonal: replace assert failures with
 exceptions

It should not be possible to fail an assertion merely because a user
provided data with the wrong sizes.
---
 src/math/symmetricTridiagonal.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/math/symmetricTridiagonal.cc b/src/math/symmetricTridiagonal.cc
index fd600872b..e2bd9b93b 100644
--- a/src/math/symmetricTridiagonal.cc
+++ b/src/math/symmetricTridiagonal.cc
@@ -1,6 +1,7 @@
 #include "ndarray.h"
 #include "lsst/pex/exceptions.h"
 #include "pfs/drp/stella/math/symmetricTridiagonal.h"
+#include "pfs/drp/stella/utils/checkSize.h"
 
 namespace pfs {
 namespace drp {
@@ -36,12 +37,11 @@ ndarray::Array<T, 1, 1> solveSymmetricTridiagonal(
     SymmetricTridiagonalWorkspace<T> & workspace
 ) {
     std::size_t const num = diagonal.getNumElements();  // number of elements
+    utils::checkSize(offDiag.getNumElements(), num - 1, "diagonal vs offDiag");
+    utils::checkSize(answer.getNumElements(), num, "diagonal vs answer");
     std::size_t const last = num - 1;  // index of last element
     std::size_t const penultimate = num - 2;  // index of penultimate element
 
-    assert(offDiag.getNumElements() == num - 1);
-    assert(answer.getNumElements() == num);
-
     workspace.reset(num);
     auto & cPrime = workspace.shortArray;
     auto & dPrime = workspace.longArray1;

From 2c70493e17129f80adc779c0fe5c1af7e9acc5c5 Mon Sep 17 00:00:00 2001
From: Paul Price <price@astro.princeton.edu>
Date: Mon, 23 May 2022 15:00:48 -0400
Subject: [PATCH 4/5] symmetricTridiagonal: deal with singular matrix

The matrix for spectral extraction can be singular if two fibers rely on
the same single pixel.

This manifests in solveSymmetricTridiagonal as denominator=0, which leads
to all values in the solution being NaN, due to the back-propagation.
Instead, when we get denominator=0, treat that element like it doesn't have
any data (set diagonal=1, offDiag=0, answer=0) and set the solution for
that element to NAN.
---
 src/math/symmetricTridiagonal.cc   | 30 +++++++++++++---
 tests/test_symmetricTridiagonal.py | 57 ++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/src/math/symmetricTridiagonal.cc b/src/math/symmetricTridiagonal.cc
index e2bd9b93b..22603de4b 100644
--- a/src/math/symmetricTridiagonal.cc
+++ b/src/math/symmetricTridiagonal.cc
@@ -1,3 +1,4 @@
+#include "unordered_set"
 #include "ndarray.h"
 #include "lsst/pex/exceptions.h"
 #include "pfs/drp/stella/math/symmetricTridiagonal.h"
@@ -56,19 +57,40 @@ ndarray::Array<T, 1, 1> solveSymmetricTridiagonal(
         return solution;
     }
 
+    std::unordered_set<std::size_t> singular;
     cPrime[0] = offDiag[0]/diagonal[0];
     dPrime[0] = answer[0]/diagonal[0];
     for (std::size_t ii = 1, jj = 0; ii < last; ++ii, ++jj) {  // jj = ii - 1
-        T const denominator = diagonal[ii] - offDiag[jj]*cPrime[jj];
+        double const denominator = diagonal[ii] - offDiag[jj]*cPrime[jj];
+        if (denominator == 0.0) {
+            // Singularity detected. Excise this element so that it won't affect everything else.
+            // Pretend that diagonal[ii] = 1, offDiag[jj] = 0, answer[ii] = 0.
+            cPrime[ii] = 0.0;
+            dPrime[ii] = 0.0;
+            singular.insert(ii);
+            continue;
+        }
         cPrime[ii] = offDiag[ii]/denominator;
         dPrime[ii] = (answer[ii] - offDiag[jj]*dPrime[jj])/denominator;
     }
-    T const denominator = diagonal[last] - offDiag[penultimate]*cPrime[penultimate];
-    dPrime[last] = (answer[last] - offDiag[penultimate]*dPrime[penultimate])/denominator;
+    double const denominator = diagonal[last] - offDiag[penultimate]*cPrime[penultimate];
+    if (denominator == 0.0) {
+        // Singularity detected. Excise this element so that it won't affect everything else.
+        // Pretend that diagonal[last] = 1, offDiag[penultimate] = 0, answer[last] = 0.
+        dPrime[last] = 0.0;
+        singular.insert(last);
+    } else {
+        dPrime[last] = (answer[last] - offDiag[penultimate]*dPrime[penultimate])/denominator;
+    }
 
     solution[last] = dPrime[last];
     for (std::ptrdiff_t ii = penultimate, jj = last; ii >= 0; --ii, --jj) {  // jj = ii + 1
-        solution[ii] = (dPrime[ii] - cPrime[ii]*solution[jj]);
+        solution[ii] = dPrime[ii] - cPrime[ii]*solution[jj];
+    }
+
+    constexpr double nan = std::numeric_limits<T>::quiet_NaN();
+    for (auto const ii: singular) {
+        solution[ii] = nan;
     }
 
     return solution;
diff --git a/tests/test_symmetricTridiagonal.py b/tests/test_symmetricTridiagonal.py
index 73eeff89c..34ac0f7a8 100644
--- a/tests/test_symmetricTridiagonal.py
+++ b/tests/test_symmetricTridiagonal.py
@@ -91,6 +91,62 @@ def testSolve(self):
         for size in (10, 100, 1000):
             self.checkSolve(size, rng)
 
+    def testSingular(self):
+        """Test that we can solve singular matrices
+
+        These values are taken from fiber indices 510 through 519 of
+        visit=76003 at y=1025, where a cosmic ray has wiped out data and left
+        only a single pixel between fiber indices 515 and 516, the full value of
+        which is claimed by both.
+        """
+        diagonal = np.array(
+            [
+                2.04603214245271892e-01,
+                1.77478682035331681e-01,
+                1.86878709397182152e-01,
+                2.30714251670864207e-01,
+                1.93489959968508124e-01,
+                2.72308226631381165e-04,
+                2.74896376556077015e-04,
+                2.33667285475434305e-01,
+                2.31162677238954783e-01,
+                2.31444225252564806e-01,
+            ]
+        )
+        offDiag = np.array(
+            [
+                8.29542189990942528e-04,
+                1.50040111291657426e-03,
+                7.35440581909623209e-04,
+                5.02985261367372362e-04,
+                0.00000000000000000e00,
+                2.73599241240500741e-04,
+                0.00000000000000000e00,
+                3.89897626215648507e-04,
+                3.69925928467314269e-04,
+            ]
+        )
+        answer = np.array(
+            [
+                2.28096719322151216e02,
+                1.88301068255664319e02,
+                1.72575595095778453e02,
+                2.59294852978659037e02,
+                1.90605045408182889e02,
+                2.92346131950238675e-01,
+                2.93732146364626023e-01,
+                2.55098123020953864e02,
+                2.74973877207476960e02,
+                2.50663980370680378e02,
+            ]
+        )
+        self.assertEqual(diagonal[5]*diagonal[6] - offDiag[5]**2, 0.0)  # Singular
+        solution = solveSymmetricTridiagonal(diagonal, offDiag, answer)
+        good = np.ones_like(solution, dtype=bool)
+        good[6] = False
+        self.assertTrue(np.all(np.isfinite(solution[good])))
+        self.assertTrue(np.all(np.isnan(solution[~good])))
+
 
 class TestMemory(lsst.utils.tests.MemoryTestCase):
     pass
@@ -103,6 +159,7 @@ def setup_module(module):
 if __name__ == "__main__":
     setup_module(sys.modules["__main__"])
     from argparse import ArgumentParser
+
     parser = ArgumentParser(__file__)
     parser.add_argument("--display", help="Display backend")
     args, argv = parser.parse_known_args()

From 32670fa06d17ab9f54729eef9acea6952eac522b Mon Sep 17 00:00:00 2001
From: Paul Price <price@astro.princeton.edu>
Date: Mon, 23 May 2022 15:25:47 -0400
Subject: [PATCH 5/5] symmetricTridiagonal: add python typing

---
 .../pfs/drp/stella/symmetricTridiagonal.pyi   | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 python/pfs/drp/stella/symmetricTridiagonal.pyi

diff --git a/python/pfs/drp/stella/symmetricTridiagonal.pyi b/python/pfs/drp/stella/symmetricTridiagonal.pyi
new file mode 100644
index 000000000..887233295
--- /dev/null
+++ b/python/pfs/drp/stella/symmetricTridiagonal.pyi
@@ -0,0 +1,22 @@
+from typing import Optional
+
+import numpy as np
+
+class SymmetricTridiagonalWorkspace:
+    def __init__(self): ...
+    def reset(self, num: int): ...
+    longArray1: np.ndarray
+    longArray2: np.ndarray
+    shortArray: np.ndarray
+
+def solveSymmetricTridiagonal(
+    diagonal: np.ndarray,
+    offDiag: np.ndarray,
+    answer: np.ndarray,
+    workspace: Optional[SymmetricTridiagonalWorkspace] = None,
+) -> np.ndarray: ...
+def invertSymmetricTridiagonal(
+    diagonal: np.ndarray,
+    offDiag: np.ndarray,
+    workspace: Optional[SymmetricTridiagonalWorkspace] = None,
+) -> np.ndarray: ...