Merge pull request #31 from lsst/tickets/DM-38745

DM-38745: Add option to run uvspec

python/lsst/atmospec/formatters.py

@@ -19,11 +19,14 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-__all__ = ['SpectractorSpectrumFormatter', 'SpectractorImageFormatter']
+__all__ = ['SpectractorSpectrumFormatter',
+           'SpectractorImageFormatter',
+           'SpectractorFitParametersFormatter']
 
 from lsst.daf.butler.formatters.file import FileFormatter
 from spectractor.extractor.spectrum import Spectrum
 from spectractor.extractor.images import Image
+from spectractor.fit.fitter import read_fitparameter_json, write_fitparameter_json
 
 
 class SpectractorSpectrumFormatter(FileFormatter):
@@ -46,3 +49,14 @@ def _readFile(self, path, pytype=None):
 
     def _writeFile(self, inMemoryDataset):
         inMemoryDataset.save_image(self.fileDescriptor.location.path)
+
+
+class SpectractorFitParametersFormatter(FileFormatter):
+    extension = '.json'
+    unsupportedParameters = None
+
+    def _readFile(self, path, pytype=None):
+        return read_fitparameter_json(path)
+
+    def _writeFile(self, inMemoryDataset):
+        write_fitparameter_json(self.fileDescriptor.location.path, inMemoryDataset)

python/lsst/atmospec/processStar.py

@@ -22,6 +22,7 @@
 __all__ = ['ProcessStarTask', 'ProcessStarTaskConfig']
 
 import os
+import shutil
 import numpy as np
 import matplotlib.pyplot as plt
 
@@ -30,6 +31,7 @@
 import lsst.geom as geom
 from lsst.ip.isr import IsrTask
 import lsst.pex.config as pexConfig
+from lsst.pex.config import FieldValidationError
 import lsst.pipe.base as pipeBase
 import lsst.pipe.base.connectionTypes as cT
 from lsst.pipe.base.task import TaskError
@@ -88,9 +90,35 @@ class ProcessStarTaskConnections(pipeBase.PipelineTaskConnections,
         storageClass="SpectractorImage",
         dimensions=("instrument", "visit", "detector"),
     )
+    spectrumForwardModelFitParameters = cT.Output(
+        name="spectrumForwardModelFitParameters",
+        doc="The full forward model fit parameters.",
+        storageClass="SpectractorFitParameters",
+        dimensions=("instrument", "visit", "detector"),
+    )
+    spectrumLibradtranFitParameters = cT.Output(
+        name="spectrumLibradtranFitParameters",
+        doc="The fitted Spectractor atmospheric parameters from fitting the atmosphere with libradtran"
+            " on the spectrum.",
+        storageClass="SpectractorFitParameters",
+        dimensions=("instrument", "visit", "detector"),
+    )
+    spectrogramLibradtranFitParameters = cT.Output(
+        name="spectrogramLibradtranFitParameters",
+        doc="The fitted Spectractor atmospheric parameters from fitting the atmosphere with libradtran"
+            " directly on the spectrogram.",
+        storageClass="SpectractorFitParameters",
+        dimensions=("instrument", "visit", "detector"),
+    )
 
     def __init__(self, *, config=None):
         super().__init__(config=config)
+        if not config.doFullForwardModelDeconvolution:
+            self.outputs.remove("spectrumForwardModelFitParameters")
+        if not config.doFitAtmosphere:
+            self.outputs.remove("spectrumLibradtranFitParameters")
+        if not config.doFitAtmosphereOnSpectrogram:
+            self.outputs.remove("spectrogramLibradtranFitParameters")
 
 
 class ProcessStarTaskConfig(pipeBase.PipelineTaskConfig,
@@ -434,6 +462,23 @@ class ProcessStarTaskConfig(pipeBase.PipelineTaskConfig,
         doc="Which filter in the reference catalog to match to?",
         default="phot_g_mean"
     )
+    # This is a post-processing function in Spectractor and therefore isn't
+    # controlled by its top-level function, and thus doesn't map to a
+    # spectractor.parameters ALL_CAPS config option
+    doFitAtmosphere = pexConfig.Field(
+        dtype=bool,
+        doc="Use uvspec to fit the atmosphere? Requires the binary to be available.",
+        default=False
+    )
+    # This is a post-processing function in Spectractor and therefore isn't
+    # controlled by its top-level function, and thus doesn't map to a
+    # spectractor.parameters ALL_CAPS config option
+    doFitAtmosphereOnSpectrogram = pexConfig.Field(
+        dtype=bool,
+        doc="Experimental option to use uvspec to fit the atmosphere directly on the spectrogram?"
+            " Requires the binary to be available.",
+        default=False
+    )
 
     def setDefaults(self):
         self.isr.doWrite = False
@@ -449,6 +494,16 @@ def setDefaults(self):
         self.charImage.detection.includeThresholdMultiplier = 3
         self.isr.overscan.fitType = 'MEDIAN_PER_ROW'
 
+    def validate(self):
+        super().validate()
+        uvspecPath = shutil.which('uvspec')
+        if uvspecPath is None and self.doFitAtmosphere is True:
+            raise FieldValidationError(self.__class__.doFitAtmosphere, self, "uvspec is not in the path,"
+                                       " but doFitAtmosphere is True.")
+        if uvspecPath is None and self.doFitAtmosphereOnSpectrogram is True:
+            raise FieldValidationError(self.__class__.doFitAtmosphereOnSpectrogram, self, "uvspec is not in"
+                                       " the path, but doFitAtmosphere is True.")
+
 
 class ProcessStarTask(pipeBase.PipelineTask):
     """Task for the spectral extraction of single-star dispersed images.
@@ -865,13 +920,19 @@ def run(self, *, inputExp, inputCentroid, dataIdDict):
         else:  # it's a raw tuple
             centroid = inputCentroid  # TODO: put this support in the docstring
 
-        spectraction = spectractor.run(inputExp, *centroid, target)
+        spectraction = spectractor.run(inputExp, *centroid, target,
+                                       self.config.doFitAtmosphere,
+                                       self.config.doFitAtmosphereOnSpectrogram)
 
         self.log.info("Finished processing %s" % (dataIdDict))
 
-        return pipeBase.Struct(spectractorSpectrum=spectraction.spectrum,
-                               spectractorImage=spectraction.image,
-                               spectraction=spectraction)
+        return pipeBase.Struct(
+            spectractorSpectrum=spectraction.spectrum,
+            spectractorImage=spectraction.image,
+            spectrumForwardModelFitParameters=spectraction.spectrumForwardModelFitParameters,
+            spectrumLibradtranFitParameters=spectraction.spectrumLibradtranFitParameters,
+            spectrogramLibradtranFitParameters=spectraction.spectrogramLibradtranFitParameters
+        )
 
     def runAstrometry(self, butler, exp, icSrc):
         refObjLoaderConfig = ReferenceObjectLoader.ConfigClass()

python/lsst/atmospec/spectraction.py

@@ -38,6 +38,9 @@
                                              dumpParameters,
                                              run_spectrogram_deconvolution_psf2d)
 from spectractor.extractor.spectrum import Spectrum, calibrate_spectrum  # noqa: E402
+from spectractor.fit.fit_spectrum import SpectrumFitWorkspace, run_spectrum_minimisation  # noqa: E402
+from spectractor.fit.fit_spectrogram import (SpectrogramFitWorkspace,  # noqa: E402
+                                             run_spectrogram_minimisation)
 
 from lsst.daf.base import DateTime  # noqa: E402
 from .utils import getFilterAndDisperserFromExp  # noqa: E402
@@ -358,8 +361,13 @@ def setAdrParameters(self, spectrum, exp):
 
         airmass = vi.getBoresightAirmass()
         spectrum.adr_params = [dec, hourAngle, temperature, pressure, humidity, airmass]
+        spectrum.pressure = pressure
+        spectrum.humidity = humidity
+        spectrum.airmass = airmass
+        spectrum.temperature = temperature
 
-    def run(self, exp, xpos, ypos, target, outputRoot=None, plotting=True):
+    def run(self, exp, xpos, ypos, target, doFitAtmosphere, doFitAtmosphereOnSpectrogram,
+            outputRoot=None, plotting=True):
         # run option kwargs in the original code, seems to ~always be True
         atmospheric_lines = True
 
@@ -453,11 +461,31 @@ def run(self, exp, xpos, ypos, target, outputRoot=None, plotting=True):
 
         # Full forward model extraction:
         # adds transverse ADR and order 2 subtraction
-        w = None
+        ffmWorkspace = None
         if parameters.SPECTRACTOR_DECONVOLUTION_FFM:
-            w = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE, plot=True, live_fit=False,
-                                             amplitude_priors_method="spectrum")
-            spectrum = run_ffm_minimisation(w, method="newton", niter=2)
+            ffmWorkspace = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE,
+                                                        plot=True,
+                                                        live_fit=False,
+                                                        amplitude_priors_method="spectrum")
+            spectrum = run_ffm_minimisation(ffmWorkspace, method="newton", niter=2)
+
+        # Fit the atmosphere on the spectrum using uvspec binary
+        spectrumAtmosphereWorkspace = None
+        if doFitAtmosphere:
+            spectrumAtmosphereWorkspace = SpectrumFitWorkspace(spectrum,
+                                                               fit_angstrom_exponent=True,
+                                                               verbose=parameters.VERBOSE,
+                                                               plot=True)
+            run_spectrum_minimisation(spectrumAtmosphereWorkspace, method="newton")
+
+        # Fit the atmosphere directly on the spectrogram using uvspec binary
+        spectrogramAtmosphereWorkspace = None
+        if doFitAtmosphereOnSpectrogram:
+            spectrogramAtmosphereWorkspace = SpectrogramFitWorkspace(spectrum,
+                                                                     fit_angstrom_exponent=True,
+                                                                     verbose=parameters.VERBOSE,
+                                                                     plot=True)
+            run_spectrogram_minimisation(spectrogramAtmosphereWorkspace, method="newton")
 
         # Save the spectrum
         self._ensureFitsHeader(spectrum)  # SIMPLE is missing by default
@@ -470,10 +498,12 @@ def run(self, exp, xpos, ypos, target, outputRoot=None, plotting=True):
         result = Spectraction()
         result.spectrum = spectrum
         result.image = image
-        result.w = w
+        result.spectrumForwardModelFitParameters = ffmWorkspace.params if ffmWorkspace is not None else None
+        result.spectrumLibradtranFitParameters = (spectrumAtmosphereWorkspace.params if
+                                                  spectrumAtmosphereWorkspace is not None else None)
+        result.spectrogramLibradtranFitParameters = (spectrogramAtmosphereWorkspace.params if
+                                                     spectrogramAtmosphereWorkspace is not None else None)
 
-        # XXX technically this should be a pipeBase.Struct I think
-        # change it if it matters
         return result