Relhum function on GPU using CuPy

geocat_cupy.yml

@@ -0,0 +1,26 @@
+name: geocat
+channels:
+  - conda-forge
+  - ncar
+dependencies:
+  - geocat-comp
+  - matplotlib
+  - jupyter
+  - cartopy
+  - cupy
+  - cudatoolkit=11.6
+  - dask-jobqueue
+  - intake
+  - intake-esm
+  - jupyterlab
+  - nc-time-axis
+  - dask
+  - h5py
+  - jupyterlab-nvdashboard
+  - nodejs
+  - numba
+  - dask-labextension
+  - ipywidgets
+  - xarray
+  - distributed
+  - netcdf4

src/geocat/comp/comp_util.py

@@ -8,3 +8,13 @@ def _is_duck_array(value):
     return (hasattr(value, "ndim") and hasattr(value, "shape") and
             hasattr(value, "dtype") and hasattr(value, "__array_function__") and
             hasattr(value, "__array_ufunc__"))
+
+
+def _import_cupy():
+    """imports the cupy and checks if not installed."""
+    try:
+        import cupy as cp
+        return cp
+    except ImportError as e:
+        print(f"Cupy is not installed for GPU computation!")
+        pass  # module doesn't exist, deal with it.

src/geocat/comp/meteorology.py

@@ -3,6 +3,7 @@
 import typing
 import xarray as xr
 import warnings
+from .comp_util import _import_cupy
 
 
 def _dewtemp(
@@ -177,11 +178,10 @@ def _nws_eqn(coeffs, temp, rel_hum):
     return heatindex
 
 
-def _relhum(
-        t: typing.Union[np.ndarray, list, float],
-        w: typing.Union[np.ndarray, xr.DataArray, list,
-                        float], p: typing.Union[np.ndarray, xr.DataArray, list,
-                                                float]) -> np.ndarray:
+def _relhum(t: typing.Union[np.ndarray, list, float],
+            w: typing.Union[np.ndarray, xr.DataArray, list, float],
+            p: typing.Union[np.ndarray, xr.DataArray, list, float],
+            use_gpu: bool = False) -> np.ndarray:
     """Calculates relative humidity with respect to ice, given temperature,
     mixing ratio, and pressure.
 
@@ -222,7 +222,15 @@ def _relhum(
     `relhum_water <https://www.ncl.ucar.edu/Document/Functions/Built-in/relhum_water.shtml>`_
     """
 
-    table = np.asarray([
+    #check if use_gpu input is true import cupy and convert all python arrays to cupy
+    xp = np
+    if (use_gpu):
+        xp = _import_cupy()
+        t = xp.asarray(t)
+        w = xp.asarray(w)
+        p = xp.asarray(p)
+
+    table = xp.asarray([
         0.01403, 0.01719, 0.02101, 0.02561, 0.03117, 0.03784, 0.04584, 0.05542,
         0.06685, 0.08049, 0.09672, 0.1160, 0.1388, 0.1658, 0.1977, 0.2353,
         0.2796, 0.3316, 0.3925, 0.4638, 0.5472, 0.6444, 0.7577, 0.8894, 1.042,
@@ -254,7 +262,7 @@ def _relhum(
     mintemp = 173.16
 
     # replace values of t above and below max and min values for temperature
-    t = np.clip(t, mintemp, maxtemp)
+    t = xp.clip(t, mintemp, maxtemp)
 
     it = (t - mintemp).astype(int)
     t2 = mintemp + it
@@ -265,7 +273,7 @@ def _relhum(
     rh = (w * (p - 0.378 * es) / (0.622 * es)) * 100
 
     # if any value is below 0.0001, set to 0.0001
-    rh = np.clip(rh, 0.0001, None)
+    rh = xp.clip(rh, 0.0001, None)
 
     return rh
 
@@ -480,7 +488,10 @@ def _xheat_index(temperature: xr.DataArray,
     return heatindex, eqtype
 
 
-def _xrelhum(t: xr.DataArray, w: xr.DataArray, p: xr.DataArray) -> xr.DataArray:
+def _xrelhum(t: xr.DataArray,
+             w: xr.DataArray,
+             p: xr.DataArray,
+             use_gpu: bool = False) -> xr.DataArray:
     """Calculates relative humidity with respect to ice, given temperature,
     mixing ratio, and pressure.
 
@@ -520,47 +531,64 @@ def _xrelhum(t: xr.DataArray, w: xr.DataArray, p: xr.DataArray) -> xr.DataArray:
     `relhum_water <https://www.ncl.ucar.edu/Document/Functions/Built-in/relhum_water.shtml>`_
     """
 
-    table = da.from_array([
-        0.01403, 0.01719, 0.02101, 0.02561, 0.03117, 0.03784, 0.04584, 0.05542,
-        0.06685, 0.08049, 0.09672, 0.1160, 0.1388, 0.1658, 0.1977, 0.2353,
-        0.2796, 0.3316, 0.3925, 0.4638, 0.5472, 0.6444, 0.7577, 0.8894, 1.042,
-        1.220, 1.425, 1.662, 1.936, 2.252, 2.615, 3.032, 3.511, 4.060, 4.688,
-        5.406, 6.225, 7.159, 8.223, 9.432, 10.80, 12.36, 14.13, 16.12, 18.38,
-        20.92, 23.80, 27.03, 30.67, 34.76, 39.35, 44.49, 50.26, 56.71, 63.93,
-        71.98, 80.97, 90.98, 102.1, 114.5, 128.3, 143.6, 160.6, 179.4, 200.2,
-        223.3, 248.8, 276.9, 307.9, 342.1, 379.8, 421.3, 466.9, 517.0, 572.0,
-        632.3, 698.5, 770.9, 850.2, 937.0, 1032.0, 1146.6, 1272.0, 1408.1,
-        1556.7, 1716.9, 1890.3, 2077.6, 2279.6, 2496.7, 2729.8, 2980.0, 3247.8,
-        3534.1, 3839.8, 4164.8, 4510.5, 4876.9, 5265.1, 5675.2, 6107.8, 6566.2,
-        7054.7, 7575.3, 8129.4, 8719.2, 9346.50, 10013.0, 10722.0, 11474.0,
-        12272.0, 13119.0, 14017.0, 14969.0, 15977.0, 17044.0, 18173.0, 19367.0,
-        20630.0, 21964.0, 23373.0, 24861.0, 26430.0, 28086.0, 29831.0, 31671.0,
-        33608.0, 35649.0, 37796.0, 40055.0, 42430.0, 44927.0, 47551.0, 50307.0,
-        53200.0, 56236.0, 59422.0, 62762.0, 66264.0, 69934.0, 73777.0, 77802.0,
-        82015.0, 86423.0, 91034.0, 95855.0, 100890.0, 106160.0, 111660.0,
-        117400.0, 123400.0, 129650.0, 136170.0, 142980.0, 150070.0, 157460.0,
-        165160.0, 173180.0, 181530.0, 190220.0, 199260.0, 208670.0, 218450.0,
-        228610.0, 239180.0, 250160.0, 261560.0, 273400.0, 285700.0, 298450.0,
-        311690.0, 325420.0, 339650.0, 354410.0, 369710.0, 385560.0, 401980.0,
-        418980.0, 436590.0, 454810.0, 473670.0, 493170.0, 513350.0, 534220.0,
-        555800.0, 578090.0, 601130.0, 624940.0, 649530.0, 674920.0, 701130.0,
-        728190.0, 756110.0, 784920.0, 814630.0, 845280.0, 876880.0, 909450.0,
-        943020.0, 977610.0, 1013250.0, 1049940.0, 1087740.0, 1087740.
-    ])
+    #check if use_gpu input is true import cupy and convert all python arrays to cupy
+    xp = np
+    if (use_gpu):
+        xp = _import_cupy()
+        #convert the xarray dataarrays to work with CuPy
+        if "dask" in str(type(t.data)):
+            t = xr.DataArray(t.data.map_blocks(xp.asarray))
+            w = xr.DataArray(w.data.map_blocks(xp.asarray))
+            p = xr.DataArray(p.data.map_blocks(xp.asarray))
+        else:
+            t = xr.DataArray(xp.asarray(t.data))
+            w = xr.DataArray(xp.asarray(w.data))
+            p = xr.DataArray(xp.asarray(p.data))
+
+    table = da.from_array(
+        xp.asarray([
+            0.01403, 0.01719, 0.02101, 0.02561, 0.03117, 0.03784, 0.04584,
+            0.05542, 0.06685, 0.08049, 0.09672, 0.1160, 0.1388, 0.1658, 0.1977,
+            0.2353, 0.2796, 0.3316, 0.3925, 0.4638, 0.5472, 0.6444, 0.7577,
+            0.8894, 1.042, 1.220, 1.425, 1.662, 1.936, 2.252, 2.615, 3.032,
+            3.511, 4.060, 4.688, 5.406, 6.225, 7.159, 8.223, 9.432, 10.80,
+            12.36, 14.13, 16.12, 18.38, 20.92, 23.80, 27.03, 30.67, 34.76,
+            39.35, 44.49, 50.26, 56.71, 63.93, 71.98, 80.97, 90.98, 102.1,
+            114.5, 128.3, 143.6, 160.6, 179.4, 200.2, 223.3, 248.8, 276.9,
+            307.9, 342.1, 379.8, 421.3, 466.9, 517.0, 572.0, 632.3, 698.5,
+            770.9, 850.2, 937.0, 1032.0, 1146.6, 1272.0, 1408.1, 1556.7, 1716.9,
+            1890.3, 2077.6, 2279.6, 2496.7, 2729.8, 2980.0, 3247.8, 3534.1,
+            3839.8, 4164.8, 4510.5, 4876.9, 5265.1, 5675.2, 6107.8, 6566.2,
+            7054.7, 7575.3, 8129.4, 8719.2, 9346.50, 10013.0, 10722.0, 11474.0,
+            12272.0, 13119.0, 14017.0, 14969.0, 15977.0, 17044.0, 18173.0,
+            19367.0, 20630.0, 21964.0, 23373.0, 24861.0, 26430.0, 28086.0,
+            29831.0, 31671.0, 33608.0, 35649.0, 37796.0, 40055.0, 42430.0,
+            44927.0, 47551.0, 50307.0, 53200.0, 56236.0, 59422.0, 62762.0,
+            66264.0, 69934.0, 73777.0, 77802.0, 82015.0, 86423.0, 91034.0,
+            95855.0, 100890.0, 106160.0, 111660.0, 117400.0, 123400.0, 129650.0,
+            136170.0, 142980.0, 150070.0, 157460.0, 165160.0, 173180.0,
+            181530.0, 190220.0, 199260.0, 208670.0, 218450.0, 228610.0,
+            239180.0, 250160.0, 261560.0, 273400.0, 285700.0, 298450.0,
+            311690.0, 325420.0, 339650.0, 354410.0, 369710.0, 385560.0,
+            401980.0, 418980.0, 436590.0, 454810.0, 473670.0, 493170.0,
+            513350.0, 534220.0, 555800.0, 578090.0, 601130.0, 624940.0,
+            649530.0, 674920.0, 701130.0, 728190.0, 756110.0, 784920.0,
+            814630.0, 845280.0, 876880.0, 909450.0, 943020.0, 977610.0,
+            1013250.0, 1049940.0, 1087740.0, 1087740.
+        ]))
 
     maxtemp = 375.16
     mintemp = 173.16
 
-    # replace values of t above and below max and min values for temperature
     t = t.clip(mintemp, maxtemp)
+    it = (t - mintemp).astype(int)
 
-    it = (t - mintemp).astype(int).data
     t2 = mintemp + it
 
     it_shape = it.shape
 
-    es = (t2 + 1 - t) * table[it.ravel()].reshape(it_shape) + (
-        t - t2) * table[it.ravel() + 1].reshape(it_shape)
+    es = (t2 + 1 - t.data) * table[it.data.ravel()].reshape(it_shape) + (
+        t - t2) * table[it.data.ravel() + 1].reshape(it_shape)
     es = es * 0.1
 
     rh = (w * (p - 0.378 * es) / (0.622 * es)) * 100
@@ -764,11 +792,10 @@ def heat_index(
     return heatindex
 
 
-def relhum(
-    temperature: typing.Union[np.ndarray, xr.DataArray, list, float],
-    mixing_ratio: typing.Union[np.ndarray, xr.DataArray, list, float],
-    pressure: typing.Union[np.ndarray, xr.DataArray, list, float]
-) -> typing.Union[np.ndarray, xr.DataArray]:
+def relhum(temperature: typing.Union[np.ndarray, xr.DataArray, list, float],
+           mixing_ratio: typing.Union[np.ndarray, xr.DataArray, list, float],
+           pressure: typing.Union[np.ndarray, xr.DataArray, list, float],
+           use_gpu: bool = False) -> typing.Union[np.ndarray, xr.DataArray]:
     """This function calculates the relative humidity given temperature, mixing
     ratio, and pressure.
 
@@ -826,7 +853,8 @@ def relhum(
                 "relhum: if using xarray, all inputs must be xarray")
 
         # call internal computation function
-        relative_humidity = _xrelhum(temperature, mixing_ratio, pressure)
+        relative_humidity = _xrelhum(temperature, mixing_ratio, pressure,
+                                     use_gpu)
 
         # set xarray attributes
         relative_humidity.attrs['long_name'] = "relative humidity"
@@ -840,7 +868,8 @@ def relhum(
         pressure = np.asarray(pressure)
 
         # function call for non-dask/xarray
-        relative_humidity = _relhum(temperature, mixing_ratio, pressure)
+        relative_humidity = _relhum(temperature, mixing_ratio, pressure,
+                                    use_gpu)
 
     return relative_humidity
 

test/test_meteorology.py

@@ -218,47 +218,49 @@ def setUpClass(cls):
         # make dask client to reference in subsequent tests
         cls.client = dd.Client()
 
-    def test_float_input(self):
+    def test_float_input(self, use_gpu=False):
         p = 1000. * 100
         t = 18. + 273.15
         q = 6. / 1000.
 
-        assert np.allclose(relhum(t, q, p), self.rh_gt_1, atol=0.1)
+        assert np.allclose(relhum(t, q, p, use_gpu), self.rh_gt_1, atol=0.1)
 
-    def test_list_input(self):
+    def test_list_input(self, use_gpu=False):
 
-        assert np.allclose(relhum(self.t_def, self.q_def, self.p_def),
+        assert np.allclose(relhum(self.t_def, self.q_def, self.p_def, use_gpu),
                            self.rh_gt_2,
                            atol=0.1)
 
-    def test_numpy_input(self):
+    def test_numpy_input(self, use_gpu=False):
         p = np.asarray(self.p_def)
         t = np.asarray(self.t_def)
         q = np.asarray(self.q_def)
 
-        assert np.allclose(relhum(t, q, p), self.rh_gt_2, atol=0.1)
+        assert np.allclose(relhum(t, q, p, use_gpu), self.rh_gt_2, atol=0.1)
 
-    def test_dims_error(self):
+    def test_dims_error(self, use_gpu=False):
         self.assertRaises(ValueError, relhum, self.t_def[:10], self.q_def[:10],
-                          self.p_def[:9])
+                          self.p_def[:9], use_gpu)
 
-    def test_xarray_type_error(self):
+    def test_xarray_type_error(self, use_gpu=False):
         self.assertRaises(TypeError, relhum, self.t_def,
-                          xr.DataArray(self.q_def), self.p_def)
+                          xr.DataArray(self.q_def), self.p_def, use_gpu)
 
-    def test_dask_compute(self):
+    def test_dask_compute(self, use_gpu=False):
         p = xr.DataArray(self.p_def).chunk(10)
         t = xr.DataArray(self.t_def).chunk(10)
         q = xr.DataArray(self.q_def).chunk(10)
 
-        assert np.allclose(relhum(t, q, p), self.rh_gt_2, atol=0.1)
+        assert np.allclose(relhum(t, q, p, use_gpu).data,
+                           self.rh_gt_2,
+                           atol=0.1)
 
-    def test_dask_lazy(self):
+    def test_dask_lazy(self, use_gpu=False):
         p = xr.DataArray(self.p_def).chunk(10)
         t = xr.DataArray(self.t_def).chunk(10)
         q = xr.DataArray(self.q_def).chunk(10)
 
-        assert isinstance(relhum(t, q, p).data, dask.array.Array)
+        assert isinstance(relhum(t, q, p, use_gpu).data, dask.array.Array)
 
 
 class Test_relhum_water(unittest.TestCase):