heat_index with list comprehension to address #596

docs/release-notes.rst

@@ -7,12 +7,13 @@ Release Notes
 
 v2024.04.0 (Unreleased)
 -----------------------
-Upcoming release
 
 Internal Changes
 ^^^^^^^^^^^^^^^^
 * Drop Python 3.9 Support by `Cora Schneck`_ in (:pr:`599`)
 * Reduce DeprecationWarnings in testing by `Cora Schneck`_ in (:pr:`582`)
+* Prevent `heat_index` from running unnecessary calculations that may generate warnings by `Anissa Zacharias`_ in (:pr:`596`)
+
 
 
 v2024.03.0 (March 29, 2024)

geocat/comp/meteorology.py

@@ -106,30 +106,33 @@ def _heat_index(temperature: np.ndarray,
                         (relative_humidity * 0.094)) + temperature) * 0.5
 
     # http://ehp.niehs.nih.gov/1206273/
-    heatindex = xr.where(temperature < 40, temperature, heatindex)
+    heatindex = np.where(temperature < 40, temperature, heatindex)
 
     # if all t values less than critical, return hi
     # otherwise perform calculation
     eqtype = 0
     if not all(temperature.ravel() < crit[0]):
         eqtype = 1
 
-        heatindex = xr.where(heatindex > crit[0],
+        heatindex = np.where(heatindex > crit[0],
                              _nws_eqn(coeffs, temperature, relative_humidity),
                              heatindex)
-
         # adjustments
-        heatindex = xr.where(
-            np.logical_and(relative_humidity < 13,
-                           np.logical_and(temperature > 80, temperature < 112)),
-            heatindex - ((13 - relative_humidity) / 4) * np.sqrt(
-                (17 - abs(temperature - 95)) / 17), heatindex)
-
-        heatindex = xr.where(
-            np.logical_and(relative_humidity > 85,
-                           np.logical_and(temperature > 80, temperature < 87)),
-            heatindex + ((relative_humidity - 85.0) / 10.0) *
-            ((87.0 - temperature) / 5.0), heatindex)
+        heatindex = [
+            hi - ((13 - rh) / 4) * np.sqrt(
+                (17 - abs(t - 95)) / 17) if rh < 13 and 80 < t < 112 else hi
+            for hi, rh, t in zip(heatindex.ravel(), relative_humidity.ravel(),
+                                 temperature.ravel())
+        ]
+
+        heatindex = [
+            hi + ((rh - 85.0) / 10.0) *
+            ((87.0 - t) / 5.0) if rh > 85 and 80 < t < 87 else hi
+            for hi, rh, t in zip(heatindex, relative_humidity.ravel(),
+                                 temperature.ravel())
+        ]
+
+        heatindex = np.asarray(heatindex).reshape(temperature.shape)
 
     return heatindex
 
@@ -466,17 +469,21 @@ def _xheat_index(temperature: xr.DataArray,
                              heatindex)
 
         # adjustments
-        heatindex = xr.where(
-            np.logical_and(relative_humidity < 13,
-                           np.logical_and(temperature > 80, temperature < 112)),
-            heatindex - ((13 - relative_humidity) / 4) * np.sqrt(
-                (17 - abs(temperature - 95)) / 17), heatindex)
-
-        heatindex = xr.where(
-            np.logical_and(relative_humidity > 85,
-                           np.logical_and(temperature > 80, temperature < 87)),
-            heatindex + ((relative_humidity - 85.0) / 10.0) *
-            ((87.0 - temperature) / 5.0), heatindex)
+        heatindex.data = da.reshape(
+            da.block([
+                hi - ((13 - rh) / 4) * np.sqrt(
+                    (17 - abs(t - 95)) / 17) if rh < 13 and 80 < t < 112 else hi
+                for hi, rh, t in zip(heatindex.data.ravel(
+                ), relative_humidity.data.ravel(), temperature.data.ravel())
+            ]), temperature.shape)
+
+        heatindex.data = da.reshape(
+            da.block([
+                hi + ((rh - 85.0) / 10.0) *
+                ((87.0 - t) / 5.0) if rh > 85 and 80 < t < 87 else hi
+                for hi, rh, t in zip(heatindex, relative_humidity.data.ravel(),
+                                     temperature.data.ravel())
+            ]), temperature.shape)
 
     return heatindex, eqtype
 

test/test_meteorology.py

@@ -91,6 +91,14 @@ def test_multi_dimensional_input(self) -> None:
                            np.asarray(self.ncl_gt_2).reshape(2, 5),
                            atol=0.005)
 
+    def test_multi_dimensional_xarray_input(self) -> None:
+        assert np.allclose(heat_index(xr.DataArray(self.t2.reshape(2, 5)),
+                                      xr.DataArray(self.rh2.reshape(2, 5)),
+                                      True),
+                           xr.DataArray(
+                               np.asarray(self.ncl_gt_2).reshape(2, 5)),
+                           atol=0.005)
+
     def test_alt_coef(self) -> None:
         assert np.allclose(heat_index(self.t2, self.rh2, True),
                            self.ncl_gt_2,