From 28b12450b2e6bd80e570e5e248ddb9a27ed4736f Mon Sep 17 00:00:00 2001
From: William Jamieson <wjamieson@stsci.edu>
Date: Thu, 10 Oct 2024 13:30:54 -0400
Subject: [PATCH] Fix saturation flag issue for numpy 2+ in romancal (#305)

* Fix for roman numpy 2 issue

* Add tests

* Add changelog
---
 changes/305.bugfix.rst             |  1 +
 src/stcal/saturation/saturation.py |  8 ++++----
 tests/test_saturation.py           | 25 +++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 changes/305.bugfix.rst

diff --git a/changes/305.bugfix.rst b/changes/305.bugfix.rst
new file mode 100644
index 000000000..6f8471937
--- /dev/null
+++ b/changes/305.bugfix.rst
@@ -0,0 +1 @@
+Fix `IntEnum` saturation flag issue with numpy 2+ for romancal.
diff --git a/src/stcal/saturation/saturation.py b/src/stcal/saturation/saturation.py
index a2797470a..c607c64a5 100644
--- a/src/stcal/saturation/saturation.py
+++ b/src/stcal/saturation/saturation.py
@@ -63,10 +63,10 @@ def flag_saturated_pixels(
         updated pixel dq array
     """
     nints, ngroups, nrows, ncols = data.shape
-    dnu = dqflags["DO_NOT_USE"]
-    saturated = dqflags["SATURATED"]
-    ad_floor = dqflags["AD_FLOOR"]
-    no_sat_check = dqflags["NO_SAT_CHECK"]
+    dnu = int(dqflags["DO_NOT_USE"])
+    saturated = int(dqflags["SATURATED"])
+    ad_floor = int(dqflags["AD_FLOOR"])
+    no_sat_check = int(dqflags["NO_SAT_CHECK"])
 
     # Identify pixels flagged in reference file as NO_SAT_CHECK,
     no_sat_check_mask = np.bitwise_and(sat_dq, no_sat_check) == no_sat_check
diff --git a/tests/test_saturation.py b/tests/test_saturation.py
index fb207f776..cc1465a39 100644
--- a/tests/test_saturation.py
+++ b/tests/test_saturation.py
@@ -3,6 +3,7 @@
 Unit tests for saturation flagging
 
 """
+from enum import IntEnum
 
 import numpy as np
 
@@ -227,3 +228,27 @@ def test_zero_frame():
     # Check ZEROFRAME flagged elements are zeroed out.
     assert zframe[0, 0, 0] == 0.0
     assert zframe[1, 0, 1] == 0.0
+
+
+def test_intenum_flags():
+    """
+    In numpy 2.0 IntEnums induce a failure in bitwise_or. Romancal uses IntEnums
+    for clarity rather than raw dictionaries of integers
+    """
+
+    class DqFlags(IntEnum):
+        GOOD = 0
+        DO_NOT_USE = 1
+        SATURATED = 2
+        AD_FLOOR = 64
+        NO_SAT_CHECK = 2**21
+
+    # Create inputs, data, and saturation maps
+    data = np.zeros((1, 5, 20, 20), dtype=np.float32)
+    gdq = np.zeros((1, 5, 20, 20), dtype=np.uint32)
+    pdq = np.zeros((20, 20), dtype=np.uint32)
+    sat_thresh = np.ones((20, 20)) * 100000.0
+    sat_dq = np.zeros((20, 20), dtype=np.uint32)
+
+    # Simple test to check no errors are raised
+    _ = flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq, ATOD_LIMIT, DqFlags)