From 706da936c20504b84277ede7930d503ab57f222c Mon Sep 17 00:00:00 2001
From: Brett <brettgraham@gmail.com>
Date: Tue, 30 Apr 2024 09:18:39 -0400
Subject: [PATCH 1/2] allow ami_analyze step to run on cal files

---
 jwst/ami/ami_analyze.py          |  4 ++++
 jwst/regtest/test_niriss_ami3.py | 25 +++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/jwst/ami/ami_analyze.py b/jwst/ami/ami_analyze.py
index 06d31eab03..dd7e24132c 100644
--- a/jwst/ami/ami_analyze.py
+++ b/jwst/ami/ami_analyze.py
@@ -3,6 +3,8 @@
 import numpy as np
 import copy
 
+from jwst.datamodels import CubeModel, ImageModel
+
 from .find_affine2d_parameters import find_rotation
 from . import instrument_data
 from . import nrm_core
@@ -73,6 +75,8 @@ def apply_LG_plus(
     # If the input image is 2D, expand all relevant extensions to be 3D
     # Incl. those not currently used?
     if len(input_model.data.shape) == 2:
+        if isinstance(input_copy, ImageModel):
+            input_copy = CubeModel(input_copy)
         input_copy.data = np.expand_dims(input_copy.data, axis=0)
         input_copy.dq = np.expand_dims(input_copy.dq, axis=0)
         # input_copy.err = np.expand_dims(input_copy.err, axis=0)
diff --git a/jwst/regtest/test_niriss_ami3.py b/jwst/regtest/test_niriss_ami3.py
index b7b8630e9c..0ec0f28e8e 100644
--- a/jwst/regtest/test_niriss_ami3.py
+++ b/jwst/regtest/test_niriss_ami3.py
@@ -17,6 +17,19 @@ def run_pipeline(rtdata_module):
     return rtdata
 
 
+@pytest.fixture(scope="module")
+def run_step_with_cal(rtdata_module):
+    rtdata = rtdata_module
+    # run step?
+    rtdata.get_data("niriss/ami/jw04478001001_03102_00001_nis_cal.fits")
+    args = [
+        'ami_analyze',
+        rtdata.input,
+    ]
+    Step.from_cmdline(args)
+    return rtdata
+
+
 @pytest.mark.bigdata
 @pytest.mark.parametrize("obs, suffix", [("012", "ami-oi"), ("015", "psf-ami-oi")])
 def test_niriss_ami3_exp(run_pipeline, obs, suffix, fitsdiff_default_kwargs):
@@ -45,3 +58,15 @@ def test_niriss_ami3_product(run_pipeline, fitsdiff_default_kwargs):
     diff = FITSDiff(rtdata.output, rtdata.truth, **fitsdiff_default_kwargs)
     assert diff.identical, diff.report()
 
+
+@pytest.mark.bigdata
+@pytest.mark.parametrize("suffix", ("ami-oi", "amimulti-oi", "amilg"))
+def test_niriss_ami3_cal(run_step_with_cal, suffix, fitsdiff_default_kwargs):
+    rtdata = run_step_with_cal
+
+    output = f"jw04478001001_03102_00001_nis_{suffix}.fits"
+    rtdata.output = output
+    rtdata.get_truth("truth/test_niriss_ami3/" + output)
+
+    diff = FITSDiff(rtdata.output, rtdata.truth, **fitsdiff_default_kwargs)
+    assert diff.identical, diff.report()

From 4fdc145ff96a8090ece546a3826a3d07d47e4c60 Mon Sep 17 00:00:00 2001
From: Brett <brettgraham@gmail.com>
Date: Tue, 30 Apr 2024 09:21:43 -0400
Subject: [PATCH 2/2] add changelog

---
 CHANGES.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGES.rst b/CHANGES.rst
index 3f7018eb61..213f7eefe6 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,6 +6,8 @@ ami
 
 - Replaced deprecated ``np.mat()`` with ``np.asmatrix()``. [#8415]
 
+- Allow ``ami_analyze`` to run on ``cal`` files. [#8451]
+
 assign_wcs
 ----------