From 572c81796d6391be099799d01670d1cc41055a7f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Agn=C3=A8s=20Fert=C3=A9?= <ferte@slac.stanford.edu>
Date: Mon, 26 Feb 2024 16:51:53 -0800
Subject: [PATCH] Write and read mask and variance plane and change check on
 formatv.

---
 .../obs/fiberspectrograph/data_manager.py     | 22 +++++++++++++++--
 python/lsst/obs/fiberspectrograph/spectrum.py | 24 ++++++++++++-------
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/python/lsst/obs/fiberspectrograph/data_manager.py b/python/lsst/obs/fiberspectrograph/data_manager.py
index 5787b6c..389fcf6 100644
--- a/python/lsst/obs/fiberspectrograph/data_manager.py
+++ b/python/lsst/obs/fiberspectrograph/data_manager.py
@@ -53,14 +53,15 @@ def make_hdulist(self):
         """
         hdu1 = self.make_primary_hdu()
         hdu2 = self.make_wavelength_hdu()
-        return astropy.io.fits.HDUList([hdu1, hdu2])
+        hdu3, hdu4 = self.make_maskvariance_hdu()
+        return astropy.io.fits.HDUList([hdu1, hdu2, hdu3, hdu4])
 
     def make_fits_header(self):
         """Return a FITS header built from a Spectrum"""
         hdr = astropy.io.fits.Header()
 
         hdr["FORMAT_V"] = self.FORMAT_VERSION
-        hdr.update(self.spectrum.md)
+        hdr.update(self.spectrum.metadata)
 
         # WCS headers - Use -TAB WCS definition
         wcs_cards = [
@@ -88,6 +89,23 @@ def make_primary_hdu(self):
         )
         return hdu
 
+    def make_maskvariance_hdu(self):
+        """Return the HDU for the mask and variance plane."""
+        hdr_mask = astropy.io.fits.Header()
+        hdr_mask["EXTTYPE"] = 'MASK    '
+        hdr_mask["EXTNAME"] = 'MASK    '
+        hdu_mask = astropy.io.fits.ImageHDU(
+            data=self.spectrum.mask, header=hdr_mask
+        )
+
+        hdr_variance = astropy.io.fits.Header()
+        hdr_variance["EXTTYPE"] = 'VARIANCE'
+        hdr_variance["EXTNAME"] = 'VARIANCE'
+        hdu_variance = astropy.io.fits.ImageHDU(
+            data=self.spectrum.variance, header=hdr_variance
+        )
+        return hdu_mask, hdu_variance
+
     def make_wavelength_hdu(self):
         """Return the wavelength HDU built from a Spectrum."""
 
diff --git a/python/lsst/obs/fiberspectrograph/spectrum.py b/python/lsst/obs/fiberspectrograph/spectrum.py
index f6fb299..71bbf28 100644
--- a/python/lsst/obs/fiberspectrograph/spectrum.py
+++ b/python/lsst/obs/fiberspectrograph/spectrum.py
@@ -45,7 +45,7 @@ class FiberSpectrum:
         Optional Detector ID for this data.
     """
 
-    def __init__(self, wavelength, flux, md=None, detectorId=0):
+    def __init__(self, wavelength, flux, md=None, detectorId=0, mask=None, variance=None):
         self.wavelength = wavelength
         self.flux = flux
         self.metadata = md
@@ -53,10 +53,10 @@ def __init__(self, wavelength, flux, md=None, detectorId=0):
         self.info = ObservationInfo(md)
         self.detector = FiberSpectrograph().getCamera()[detectorId]
 
-        mask = afwImage.MaskX(1, 1)
-        self.getPlaneBitMask = mask.getPlaneBitMask  # ughh, awful Mask API
-        self.mask = np.zeros(flux.shape, dtype=mask.array.dtype)
-        self.variance = np.zeros_like(flux)
+        mask_temp = afwImage.MaskX(1, 1)
+        self.getPlaneBitMask = mask_temp.getPlaneBitMask  # ughh, awful Mask API
+        self.mask = mask
+        self.variance = variance
 
     def getDetector(self):
         """Get fiber spectrograph detector."
@@ -100,16 +100,24 @@ def readFits(cls, path):
 
         fitsfile = astropy.io.fits.open(path)
         md = dict(fitsfile[0].header)
+        format_v = md["FORMAT_V"]
 
-        if md["FORMAT_V"] == 1:
+        if format_v == 1:
             flux = fitsfile[0].data
             wavelength = fitsfile[md["PS1_0"]].data[md["PS1_1"]].flatten()
 
             wavelength = u.Quantity(wavelength, u.Unit(md["CUNIT1"]), copy=False)
+
+            mask_temp = afwImage.MaskX(1, 1)
+            mask = np.zeros(flux.shape, dtype=mask_temp.array.dtype)
+            variance = np.zeros_like(flux)
+            if len(fitsfile) == 4:
+                mask = fitsfile[2].data
+                variance = fitsfile[3].data
         else:
-            raise ValueError(f"FORMAT_V has changed from 1 to {md["FORMAT_V"]}")
+            raise ValueError(f"FORMAT_V has changed from 1 to {format_v}")
 
-        return cls(wavelength, flux, md)
+        return cls(wavelength, flux, md=md, mask=mask, variance=variance)
 
     def writeFits(self, path):
         """Write a Spectrum to disk.