update Forward Models for MIRI

Author: jemorrison

src/stdatamodels/jwst/transforms/models.py

@@ -1707,7 +1707,7 @@ def evaluate(self, x, y, wavelength, order):
         return x + dx, y + dy, x, y, order
 
 
-class _NIRISSForwardGrismDispersion(_ForwardGrismDispersionBase):
+class _NIRISSMIRIForwardGrismDispersion(_ForwardGrismDispersionBase):
     def __init__(self, dispaxis, *args, **kwargs):
         self.dispaxis = dispaxis
         super().__init__(*args, **kwargs)
@@ -1735,8 +1735,8 @@ def evaluate(self, x, y, x0, y0, order):
 
         Returns
         -------
-        x, y : float or np.ndarray
-            The x, y values in the direct image, same as x0, y0.
+        x0, y0 : float or np.ndarray
+            The x0, y9 values in the direct image
         lambda : float or np.ndarray
             Wavelength in angstroms
         order : int
@@ -1788,10 +1788,10 @@ def evaluate(self, x, y, x0, y0, order):
         wavelength = dxr | tab | lmodel
         model = mapping | Const1D(x00) & Const1D(y00) & wavelength & Const1D(order)
 
-        return model(x, y, x0, y0, order)
+        return model(x, y, x0, y0, order)  # returns x0, y0 , lambda, order
 
 
-class NIRISSForwardRowGrismDispersion(_NIRISSForwardGrismDispersion):
+class NIRISSForwardRowGrismDispersion(_NIRISSMIRIForwardGrismDispersion):
     """
     Calculate the wavelengths of vertically dispersed NIRISS grism data.
 
@@ -1851,7 +1851,7 @@ def __init__(
         )
 
 
-class NIRISSForwardColumnGrismDispersion(_NIRISSForwardGrismDispersion):
+class NIRISSForwardColumnGrismDispersion(_NIRISSMIRIForwardGrismDispersion):
     """
     Calculate the wavelengths for horizontally dispersed NIRISS grism data.
 
@@ -2038,13 +2038,13 @@ def __init__(self, orders, lmodels=None, xmodels=None, ymodels=None):
             name=name,
         )
 
-    def evaluate(self, x, y, wavelength, order):
+    def evaluate(self, x0, y0, wavelength, order):
         """
         Transform from the direct image plane to the dispersed plane.
 
         Parameters
         ----------
-        x, y : float or np.ndarray
+        x0, y0 : float or np.ndarray
             Input x, y location in the direct image
         wavelength : float or np.ndarray
             Wavelength in microns
@@ -2054,15 +2054,15 @@ def evaluate(self, x, y, wavelength, order):
         Returns
         -------
         x, y : float or np.ndarray
-            The x, y values in the dispersed plane.
+            The x (dx + x0_) , y(dy + y0) values in the dispersed plane.
         x0, y0 : float or np.ndarray
-            Source object x-center, y-center in the direct image. Same as input x,y
+            Source object x-center, y-center in the direct image.
         order : int
             Output spectral order, same as input
         """
         wavelength = np.atleast_1d(wavelength)
-        x = np.atleast_1d(x)
-        y = np.atleast_1d(y)
+        x0 = np.atleast_1d(x0)
+        y0 = np.atleast_1d(y0)
 
         if (wavelength < 0).any():
             raise ValueError("Wavelength should be greater than zero")
@@ -2076,21 +2076,34 @@ def evaluate(self, x, y, wavelength, order):
         xmodel = self.xmodels[iorder]
         ymodel = self.ymodels[iorder]
 
-        dx = _poly_with_spatial_dependence(t, x, y, xmodel)
-        dy = _poly_with_spatial_dependence(t, x, y, ymodel)
+        dx = _poly_with_spatial_dependence(t, x0, y0, xmodel)
+        dy = _poly_with_spatial_dependence(t, x0, y0, ymodel)
 
-        return x + dx, y + dy, x, y, order
+        return x0 + dx, y0 + dy, x0, y0, order
 
 
-class MIRIWFSSForwardDispersion(_ForwardGrismDispersionBase):
+class MIRIWFSSForwardDispersion(_NIRISSMIRIForwardGrismDispersion):
     """
-    Calculate the wavelengths of vertically dispersed MIRI WFSS data.
+    Calculate the wavelengths of the dispersed MIRI WFSS data.
 
     The dispersion polynomial is relative to the input x,y pixels
     in the direct image for a given wavelength.
+
+    The dispersion polynomial is relative to the input x,y pixels
+    in the direct image for a given wavelength. This transform uses
+    a generic method for both MIRI and NIRISS. For MIRI the theta
+    parameter = 0,
     """
 
-    def __init__(self, orders, lmodels=None, xmodels=None, ymodels=None):
+    def __init__(
+        self,
+        orders,
+        lmodels=None,
+        xmodels=None,
+        ymodels=None,
+        theta=None,
+        sampling=10,
+    ):
         """
         Initialize the model.
 
@@ -2117,67 +2130,24 @@ def __init__(self, orders, lmodels=None, xmodels=None, ymodels=None):
             C0(x0, y0) + C1(x0, y0) * t + C2(x0, y0) * t^2 + C3(x0,y0) * t^3.
             The inner list corresponds to the 2-D polynomials (C0, C1, C2, C3).
             The outer list corresponds to the different spectral orders.
+        theta : float
+            Set = 0 for MIRI.
+        sampling : int, optional
+            Number of sampling points in t to use; these will be linearly interpolated.
         """
+        self.theta = 0.0
+        dispaxis = "column"
         name = "miri_wfss_forward_dispersion"
         super().__init__(
+            dispaxis,
             orders,
             lmodels=lmodels,
             xmodels=xmodels,
             ymodels=ymodels,
             name=name,
+            sampling=sampling,
         )
 
-    def evaluate(self, x, y, x0, y0, order):
-        """
-        Transform from the dispersed plane into the direct image plane.
-
-        Only the xmodel and lmodels are used.
-
-        Parameters
-        ----------
-        x, y :  float or np.ndarray
-            Input x, y location in the dispersed image.
-        x0, y0 : float or np.ndarray
-            Source object x-center, y-center in the direct image.
-        order : int
-            Input spectral order
-
-        Returns
-        -------
-        x, y : float or np.ndarray
-            The x, y values in the direct image, same as x0, y0.
-        lambda : float or np.ndarray
-            Wavelength in microns
-        order : int
-            Output spectral order, same as input
-        """
-        try:
-            iorder = self._order_mapping[int(order.flatten()[0])]
-        except KeyError as err:
-            raise ValueError("Specified order is not available") from err
-
-        # The next two lines are to get around the fact that
-        # modeling.standard_broadcasting=False does not work.
-        x00 = x0.flatten()[0]
-        y00 = y0.flatten()[0]
-
-        t = np.linspace(0, 1, self.sampling)  # sample t
-        xmodel = self.xmodels[iorder]
-        lmodel = self.lmodels[iorder]
-
-        dx = _poly_with_spatial_dependence(t, x00, y00, xmodel)
-
-        so = np.argsort(dx)
-        tab = Tabular1D(dx[so], t[so], bounds_error=False, fill_value=None)
-        # wavelength model takes in x, x0.
-        # it then subtracts them to get dx; that's what SubtractUfunc does
-        # next it finds the t value for that dx from the lookup table, interpolating linearly
-        # finally it applies the lmodel of t to get the wavelength
-        dxr = astmath.SubtractUfunc()
-        wavelength = dxr | tab | lmodel
-        model = Mapping((2, 3, 0, 2, 4)) | Const1D(x00) & Const1D(y00) & wavelength & Const1D(order)
-        return model(x, y, x0, y0, order)
-
 
 class Rotation3DToGWA(Model):
     """