Fixed input args bandpass and affine2d to expect ASDF files, updated …

…documentation with examples.

docs/jwst/ami_analyze/description.rst

@@ -36,7 +36,7 @@ other options:
                     rotation search values. The default setting of '-3 3 1'
                     results in search values of [-3, -2, -1, 0, 1, 2, 3].
 
-:--bandpass: Synphot spectrum or suitable array to override filter/source 
+:--bandpass: ASDF file containing suitable array to override filter/source 
              (default=None)
 
 :--usebp: If True, exclude pixels marked DO_NOT_USE from fringe fitting 
@@ -47,10 +47,72 @@ other options:
 :--chooseholes: If not None, fit only certain fringes e.g. ['B4','B5','B6','C2']
                 (default=None)
 
-:--affine2d: User-defined Affine2d object (default=None)
+:--affine2d: ASDF file containing user-defined affine parameters (default=None)
 
 :--run_bpfix: Run Fourier bad pixel fix on cropped data (default=True)
-
+
+
+Creating ASDF files
+^^^^^^^^^^^^^^^^^^^
+The optional arguments `bandpass` and `affine2d` must be written to `ASDF <https://asdf-standard.readthedocs.io/>`_ 
+files to be used by the step. The step expects the contents to be stored with particular keys but the format is not currently
+enforced by a schema; incorrect ASDF file contents will cause the step to revert back to the defaults for each argument.
+
+Examples of how to create ASDF files containing the properly formatted information for each of the arguments follows.
+
+.. code-block:: python
+
+   # Create a F380M filter + A0V source bandpass ASDF file
+
+   import asdf
+   from jwst.ami import utils
+
+   filt='F380M'
+   src = 'A0V'
+   nspecbin=19
+
+   filt_spec = utils.get_filt_spec(filt)
+   src_spec = utils.get_src_spec(src)
+
+   bandpass = utils.combine_src_filt(filt_spec, 
+                                  src_spec, 
+                                  trim=0.01, 
+                                  nlambda=nspecbin)
+
+   # this bandpass has shape (19, 2); each row is [throughput, wavelength]
+
+   asdf_name = 'bandpass_f380m_a0v.asdf'
+
+   tree = {"bandpass": bandpass}
+
+   with open(asdf_name, 'wb') as fh:
+        af = asdf.AsdfFile(tree)
+        af.write_to(fh)
+   af.close()
+
+
+.. code-block:: python
+
+   # Create an affine transform ASDF file to use for the model
+
+   import asdf
+   tree = {
+   'mx': 1., # dimensionless x-magnification
+   'my': 1., # dimensionless y-magnification
+   'sx': 0., # dimensionless x shear
+   'sy': 0., # dimensionless y shear
+   'xo': 0., # x-offset in pupil space
+   'yo': 0., # y-offset in pupil space
+   'rotradccw': None }
+
+   affineasdf = 'affine.asdf'
+
+   with open(affineasdf, 'wb') as fh:
+        af = asdf.AsdfFile(tree)
+        af.write_to(fh)
+   af.close()
+
+
 
 Inputs
 ------

jwst/ami/ami_analyze_step.py

@@ -2,6 +2,11 @@
 
 from ..stpipe import Step
 from . import ami_analyze
+from . import utils
+
+import numpy as np
+import asdf
+import synphot
 
 __all__ = ["AmiAnalyzeStep"]
 
@@ -16,11 +21,11 @@ class AmiAnalyzeStep(Step):
         rotation = float(default=0.0)           # Rotation initial guess [deg]
         psf_offset = string(default='0.0 0.0') # PSF offset values to use to create the model array
         rotation_search = string(default='-3 3 1') # Rotation search parameters: start, stop, step
-        bandpass = any(default=None) # Synphot spectrum or array to override filter/source
+        bandpass = string(default=None) # ASDF file containing array or Synphot spectrum to override filter/source
         usebp = boolean(default=True) # If True, exclude pixels marked DO_NOT_USE from fringe fitting
         firstfew = integer(default=None) # If not None, process only the first few integrations
         chooseholes = string(default=None) # If not None, fit only certain fringes e.g. ['B4','B5','B6','C2']
-        affine2d = any(default=None) # None or user-defined Affine2d object
+        affine2d = string(default=None) # ASDF file containing user-defined affine parameters
         run_bpfix = boolean(default=True) # Run Fourier bad pixel fix on cropped data
     """
 
@@ -32,6 +37,68 @@ def save_model(self, model, *args, **kwargs):
             kwargs["suffix"] = ["ami-oi", "amimulti-oi", "amilg"][kwargs.pop("idx")]
         return Step.save_model(self, model, *args, **kwargs)
 
+
+    def override_bandpass(self):
+        """
+        Read bandpass from asdf file. Expects an array of [effstims, wave_m] 
+        (i.e. np.array((effstims,wave_m)).T) stored as 'bandpass' in asdf file,
+        where effstims are normalized countrates (unitless) and wave_m is wavelengths 
+        across the filter at which to compute the model (meters).
+
+        """
+
+        try:
+            with asdf.open(self.bandpass, lazy_load=False) as af:
+                bandpass = np.array(af['bandpass'])
+
+            # assume it is an array of the correct shape
+            wavemin = np.min(bandpass[:,1])
+            wavemax = np.max(bandpass[:,1])
+            self.log.info('User-defined bandpass provided:') 
+            self.log.info('\tOVERWRITING ALL NIRISS-SPECIFIC FILTER/BANDPASS VARIABLES')
+            self.log.info(f'Using {bandpass.shape[0]} wavelengths for fit.')
+            self.log.info(f'Wavelength min: {wavemin:.3e} \t Wavelength max: {wavemax:.3e}')
+
+            # update attribute and return
+            self.bandpass = bandpass 
+            return bandpass
+
+        except:
+            message = (f'Could not read bandpass from {self.bandpass}. \
+                See documentation for info on creating a custom bandpass ASDF file.')
+            raise Exception(message)
+
+    def override_affine2d(self):
+        """
+        Read user-input affine transform from ASDF file. Make Affine2d object
+        (see utils.Affine2D class). Input should contain mx,my,sx,sy,xo,yo,rotradccw.
+        """
+        try:
+            with asdf.open(self.affine2d, lazy_load=False) as af:
+                affine2d = utils.Affine2d(
+                    mx = af['mx'],
+                    my = af['my'],
+                    sx = af['sx'],
+                    sy = af['sy'],
+                    xo = af['xo'],
+                    yo = af['yo'],
+                    rotradccw = af['rotradccw']
+                    )
+                self.log.info('Using input affine transform with parameters:')
+                self.log.info(f'\tmx={af['mx']}, my={af['my']}')
+                self.log.info(f'\tsx={af['sx']}, sy={af['sy']}')
+                self.log.info(f'\txo={af['xo']}, yo={af['yo']}')
+            self.affine2d = affine2d
+            return affine2d
+        except:
+            self.log.info(f'Could not read affine transfrom parameters from {self.affine2d}. \
+                See documentation for info on creating a custom affine2d ASDF file.')
+            self.log.info('**** DEFAULTING TO USE IDENTITY TRANSFORM ****')
+            affine2d = None
+
+        self.affine2d = affine2d
+        return affine2d
+
     def process(self, input):
         """
         Performs analysis of an AMI mode exposure by applying the LG algorithm.
@@ -83,6 +150,15 @@ def process(self, input):
                 raise RuntimeError("No THROUGHPUT reference file found. "
                                    "ami_analyze cannot continue.")
 
+            # If there's a user-defined bandpass or affine, handle it
+            if bandpass is not None:
+                bandpass = self.override_bandpass()
+
+            if affine2d is not None:
+                # if it is None, handled in apply_LG_plus
+                affine2d = self.override_affine2d()
+
+
             # Get the name of the NRM reference file to use
             nrm_reffile = self.get_reference_file(input_model, 'nrm')
             self.log.info(f'Using NRM reference file {nrm_reffile}')

jwst/ami/utils.py

@@ -1322,16 +1322,12 @@ def handle_bandpass(bandpass, throughput_model):
         Array of weights, wavelengths used to generate model
     """
     if bandpass is not None:
-        log.info(
-            "User-defined bandpass provided: OVERWRITING ALL NIRISS-SPECIFIC FILTER/BANDPASS VARIABLES"
-        )
         # bandpass can be user-defined synphot object or appropriate array
         if isinstance(bandpass, synphot.spectrum.SpectralElement):
             log.info("User-defined synphot spectrum provided")
             wl, wt = bandpass._get_arrays(bandpass.waveset)
             bandpass = np.array((wt, wl)).T
         else:
-            log.info("User-defined bandpass array provided")
             bandpass = np.array(bandpass)
 
     else: