JP-3739: Update NRM geometry, affine distortion use

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.py

@@ -126,15 +126,17 @@
     # Throughput (combined filter and source spectrum) calculated here
     bandpass = utils.handle_bandpass(bandpass, throughput_model)
 
-    rotsearch_d = np.append(
+
+
+    if affine2d is None:
+        rotsearch_d = np.append(
         np.arange(
             rotsearch_parameters[0], rotsearch_parameters[1], rotsearch_parameters[2]
         ),
         rotsearch_parameters[1],
-    )
+        )
 
-    log.info(f"Initial values to use for rotation search: {rotsearch_d}")
-    if affine2d is None:
+        log.info(f"Initial values to use for rotation search: {rotsearch_d}")
         # affine2d object, can be overridden by user input affine.
         # do rotation search on uncropped median image (assuming rotation constant over exposure)
         # replace remaining NaNs in median image with median of surrounding 8 (non-NaN) pixels
@@ -173,6 +175,14 @@
             oversample,
             holeshape,
         )
+        log.info(f'Found rotation: {affine2d.rotradccw:.4f} rad ({np.rad2deg(affine2d.rotradccw):.4f} deg)')
+        # the affine2d returned here has only rotation...
+        # to use rotation and scaling/shear, do some matrix multiplication here??
+
+    log.info('Using affine transform with parameters:')
+    log.info(f'\tmx={affine2d.mx}\tmy={affine2d.my}')
+    log.info(f'\tsx={affine2d.sx}\tsy={affine2d.sy}')
+    log.info(f'\txo={affine2d.xo}\tyo={affine2d.yo}')
 
     niriss = instrument_data.NIRISS(filt,
                                     nrm_model,

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 @@
         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='commissioning') # ASDF file containing user-defined affine parameters OR 'commssioning'
         run_bpfix = boolean(default=True) # Run Fourier bad pixel fix on cropped data
     """
 
@@ -32,6 +37,66 @@
             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(f'Using affine transform from ASDF file {self.affine2d}')
+            # now self.affine2d updated from string to object
+            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 +148,23 @@
                 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 affine2d == 'commissioning':
+                    affine2d = utils.Affine2d(mx=9.92820e-01,
+                        my=9.98540e-01,
+                        sx=6.18605e-03,
+                        sy=-7.27008e-03,
+                        xo=0,
+                        yo=0,
+                        name='commissioning')
+                    self.log.info("Using affine parameters from commissioning.")
+                else:
+                    affine2d = self.override_affine2d()
+            # and if it is None, rotation search done in apply_LG_plus
+
             # 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/find_affine2d_parameters.py

@@ -98,7 +98,7 @@
     crosscorr_rots = []
 
     for (rot, aff) in zip(rotdegs, affine2d_list):
-        jw = lg_model.NrmModel(mask='jwst_g7s6c', holeshape=holeshape, over=over, affine2d=aff)
+        jw = lg_model.LgModel(mask='jwst_ami', holeshape=holeshape, over=over, affine2d=aff)
 
         jw.set_pixelscale(pixel)
         # psf_offset in data coords & pixels.  Does it get rotated?  Second order errors poss.

jwst/ami/instrument_data.py

@@ -6,7 +6,7 @@
 import logging
 import numpy as np
 
-from .mask_definitions import NRM_mask_definitions
+from .mask_definition_ami import NRM_definition
 from . import utils
 from . import bp_fix
 from stdatamodels.jwst.datamodels import dqflags
@@ -91,12 +91,12 @@
         # only one NRM on JWST:
         self.telname = "JWST"
         self.instrument = "NIRISS"
-        self.arrname = "jwst_g7s6c"
-        self.holeshape = "hex"
-        self.mask = NRM_mask_definitions(
+        self.arrname = "jwst_ami"
+        self.holeshape = 'hex'
+        self.mask = NRM_definition(
+            nrm_model,
             maskname=self.arrname,
-            chooseholes=self.chooseholes,
-            holeshape=self.holeshape,
+            chooseholes=self.chooseholes
         )
 
         # save affine deformation of pupil object or create a no-deformation object.
@@ -179,9 +179,18 @@
         pscale_deg = np.mean([pscaledegx, pscaledegy])
         self.pscale_rad = np.deg2rad(pscale_deg)
         self.pscale_mas = pscale_deg * (60 * 60 * 1000)
-        self.pav3 = input_model.meta.pointing.pa_v3
+        #####
+
+        # Until further notice this is fine 
+        #self.pav3 = input_model.meta.pointing.pa_v3
+        self.pav3 = input_model.meta.wcsinfo.roll_ref
+
         self.vparity = input_model.meta.wcsinfo.vparity
         self.v3iyang = input_model.meta.wcsinfo.v3yangle
+
+        self.wcsinfo = input_model.meta.wcsinfo
+
+
         self.parangh = input_model.meta.wcsinfo.roll_ref
         self.crpix1 = input_model.meta.wcsinfo.crpix1
         self.crpix2 = input_model.meta.wcsinfo.crpix2

jwst/ami/leastsqnrm.py

@@ -800,6 +800,43 @@
 
     return arr
 
+def t3_amplitudes(amps, n=7):
+    """
+    Populate the triple-product amplitude array
+    (NOT closure amplitudes)
+
+    Parameters
+    ----------
+    amps: 1D float array
+        fringe visibility between each pair of holes
+
+    n: integer
+        number of holes
+
+    Returns
+    -------
+    cpamps: 1D float array
+        triple product amplitude array
+    """
+
+    arr = populate_symmamparray(amps, n=n)
+
+    cpamps = np.zeros(int(comb(n, 3)))
+
+    nn = 0
+    for kk in range(n - 2):
+        for ii in range(n - kk - 2):
+            for jj in range(n - kk - ii - 2):
+                cpamps[nn + jj] = (
+                    arr[kk, ii + kk + 1]
+                    * arr[ii + kk + 1, jj + ii + kk + 2]
+                    * arr[jj + ii + kk + 2, kk]
+                )
+
+            nn += jj + 1
+
+    return cpamps
+
 
 def redundant_cps(deltaps, n=7):
     """
@@ -934,3 +971,38 @@
                 nn = nn + ll + 1
 
     return cas
+
+def q4_phases(deltaps, n=7):
+    """
+    Calculate phases for each set of 4 holes
+
+    Parameters
+    ----------
+    deltaps: 1D float array
+        pistons between each pair of holes
+
+    n: integer
+        number of holes
+
+    Returns
+    -------
+    quad_phases: 1D float array
+        quad phases
+    """
+    arr = populate_antisymmphasearray(deltaps, n=n)  # fringe phase array
+    nn = 0
+    quad_phases = np.zeros(int(comb(n, 4)))
+
+    for ii in range(n - 3):
+        for jj in range(n - ii - 3):
+            for kk in range(n - jj - ii - 3):
+                for ll in range(n - jj - ii - kk - 3):
+                    quad_phases[nn + ll] = (
+                        arr[ii, jj + ii + 1]
+                        + arr[ll + ii + jj + kk + 3, kk + jj + ii + 2]
+                        - arr[ii, kk + ii + jj + 2]
+                        - arr[jj + ii + 1, ll + ii + jj + kk + 3]
+                        )
+                nn = nn + ll + 1
+
+    return quad_phases