feat: bugfix mismatched device arrays

Author: Oisin-M

src/earthkit/meteo/extreme/array/cpf.py

@@ -14,13 +14,14 @@ def _cpf(clim, ens, epsilon=None, from_zero=False):
     xp = array_namespace(clim, ens)
     clim = xp.asarray(clim)
     ens = xp.asarray(ens)
+    device = xp.device(clim)
 
     nclim, npoints = clim.shape
     nens, _ = ens.shape
 
-    cpf = xp.zeros(npoints, dtype=xp.float32)
-    mask = xp.zeros(npoints, dtype=xp.bool)
-    prim = xp.zeros(npoints, dtype=xp.bool)
+    cpf = xp.zeros(npoints, dtype=xp.float32, device=device)
+    mask = xp.zeros(npoints, dtype=xp.bool, device=device)
+    prim = xp.zeros(npoints, dtype=xp.bool, device=device)
 
     # start scanning ensemble from iq_start
     iq_start = 0 if from_zero else nens // 2

src/earthkit/meteo/extreme/array/efi.py

@@ -34,6 +34,7 @@ def efi(clim, ens, eps=-0.1):
     xp = array_namespace(clim, ens)
     clim = xp.asarray(clim)
     ens = xp.asarray(ens)
+    device = xp.device(clim)
 
     # locate missing values
     missing_mask = xp.logical_or(xp.sum(xp.isnan(clim), axis=0), xp.sum(xp.isnan(ens), axis=0))
@@ -60,10 +61,10 @@ def efi(clim, ens, eps=-0.1):
     acoef = (1.0 - 2.0 * p[:-1]) * acosdiff + proddiff
 
     # compute EFI from coefficients
-    efi = xp.zeros(npoints)
+    efi = xp.zeros(npoints, device=device)
     ##################################
     if eps > 0:
-        efimax = xp.zeros(npoints)
+        efimax = xp.zeros(npoints, device=device)
         for icl in range(nclim - 1):
             mask = clim[icl + 1, :] > eps
             dEFI = xp.where(

src/earthkit/meteo/extreme/array/sot.py

@@ -39,17 +39,9 @@ def sot_func(qc_tail, qc, qf, eps=-1e-4, lower_bound=-10, upper_bound=10):
     qc = xp.asarray(qc)
     qf = xp.asarray(qf)
 
-    # TODO: check if this is necessary
-    # NOTE: work for numpy but not for other backends
-    # avoid divided by zero warning
-    err = xp.seterr(divide="ignore", invalid="ignore")
-
     min_den = xp.fmax(xp.asarray(eps), xp.asarray(0))
     sot = xp.where(xp.abs(qc_tail - qc) > min_den, (qf - qc_tail) / (qc_tail - qc), xp.nan)
 
-    # revert to original error state
-    xp.seterr(**err)
-
     mask_missing = xp.isnan(sot)
 
     # upper and lower bounds

src/earthkit/meteo/stats/array/quantiles.py

@@ -49,12 +49,13 @@ def iter_quantiles(
 
     xp = array_namespace(arr)
     arr = xp.asarray(arr)
+    device = xp.device(arr)
 
     if isinstance(which, int):
         n = which
-        qs = xp.linspace(0.0, 1.0, n + 1)
+        qs = xp.linspace(0.0, 1.0, n + 1, device=device)
     else:
-        qs = xp.asarray(which)
+        qs = xp.asarray(which, device=device)
 
     if method == "numpy_bulk":
         quantiles = xp.quantile(arr, qs, axis=axis)

src/earthkit/meteo/wind/array/wind.py

@@ -303,12 +303,13 @@ def windrose(speed, direction, sectors=16, speed_bins=None, percent=True):
     # TODO: atleast_1d is not part of the array API standard
     speed = xp.atleast_1d(speed)
     direction = xp.atleast_1d(direction)
+    device = xp.device(speed)
 
     dir_step = 360.0 / sectors
-    dir_bins = xp.asarray(
-        xp.linspace(int(-dir_step / 2), int(360 + dir_step / 2), int(360 / dir_step) + 2), dtype=speed.dtype
+    dir_bins = xp.linspace(
+        int(-dir_step / 2), int(360 + dir_step / 2), int(360 / dir_step) + 2, dtype=speed.dtype, device=device
     )
-    speed_bins = xp.asarray(speed_bins, dtype=speed.dtype)
+    speed_bins = xp.asarray(speed_bins, dtype=speed.dtype, device=device)
 
     # NOTE: np.histogram2d is only available in numpy. For other namespaces we use a fallback implementation
     # based on histogramdd. (See utils.compute.histogram2d). However, neither histogram2d nor
@@ -317,7 +318,6 @@ def windrose(speed, direction, sectors=16, speed_bins=None, percent=True):
         speed,
         direction,
         bins=[speed_bins, dir_bins],
-        density=False,
     )[0]
 
     # unify the north bins