Add config options for new fit types and save all fit parameters.

Moves libradtran fiting to run from the correct file. Renames its output dataset from spectractorFitParameters to spectrumLibradtranFitParameters.
lsst · May 22, 2023 · 9bd9d1e · 9bd9d1e
1 parent f6413a8
commit 9bd9d1e
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 29 deletions.
diff --git a/python/lsst/atmospec/processStar.py b/python/lsst/atmospec/processStar.py
@@ -84,23 +84,41 @@ class ProcessStarTaskConnections(pipeBase.PipelineTaskConnections,
         storageClass="SpectractorSpectrum",
         dimensions=("instrument", "visit", "detector"),
     )
-    spectractorFitParameters = cT.Output(
-        name="spectractorFitParameters",
-        doc="The fitted Spectractor atmospheric parameters.",
-        storageClass="SpectractorFitParameters",
-        dimensions=("instrument", "visit", "detector"),
-    )
     spectractorImage = cT.Output(
         name="spectractorImage",
         doc="The Spectractor output image.",
         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("spectractorFitParameters")
+            self.outputs.remove("spectrumLibradtranFitParameters")
+        if not config.doFitAtmosphereOnSpectrogram:
+            self.outputs.remove("spectrogramLibradtranFitParameters")
 
 
 class ProcessStarTaskConfig(pipeBase.PipelineTaskConfig,
@@ -444,11 +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
@@ -470,6 +500,9 @@ def validate(self):
         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):
@@ -887,23 +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)
-
-        w = None
-        if self.config.doFitAtmosphere:
-            from spectractor.fit.fit_spectrum import SpectrumFitWorkspace, run_spectrum_minimisation
-            w = SpectrumFitWorkspace(spectraction.spectrum,
-                                     fit_angstrom_exponent=True,
-                                     verbose=True,
-                                     plot=True)
-            run_spectrum_minimisation(w, method="newton")
+        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,
-                               spectractorFitParameters=w.params if w is not None else None,
-                               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()

diff --git a/python/lsst/atmospec/spectraction.py b/python/lsst/atmospec/spectraction.py
@@ -38,6 +38,8 @@
                                              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, run_spectrogram_minimisation  # noqa: E402
 
 from lsst.daf.base import DateTime  # noqa: E402
 from .utils import getFilterAndDisperserFromExp  # noqa: E402
@@ -363,7 +365,8 @@ def setAdrParameters(self, spectrum, exp):
         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
 
@@ -457,11 +460,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
@@ -474,10 +497,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