Use earthkit-utils (#672)

* Use earthkit-utils

Author: sandorkertesz

environment.yml

@@ -29,6 +29,7 @@ dependencies:
   - earthkit-meteo>=0.0.1
   - covjsonkit>=0.1.1
   - earthkit-geo>=0.2.0
+  - earthkit-utils>=0.0.1
 - tqdm>=4.63.0
 - lru-dict
 - markdown
@@ -56,5 +57,4 @@ dependencies:
 - nbsphinx
 - ipykernel
 - geopandas
-- array-api-compat
 - deprecation

pyproject.toml

@@ -26,11 +26,11 @@ classifiers = [
 ]
 dynamic = [ "version" ]
 dependencies = [
-  "array-api-compat",
   "cfgrib>=0.9.10.1",
   "dask",
   "deprecation",
   "earthkit-meteo>=0.0.1",
+  "earthkit-utils>=0.0.1",
   "eccodes>=1.7",
   "entrypoints",
   "filelock",

src/earthkit/data/__init__.py

@@ -15,6 +15,8 @@
     # Local copy or not installed with setuptools
     __version__ = "999"
 
+from earthkit.utils import array as array_api
+
 from earthkit.data.translators import transform
 from earthkit.data.wrappers import get_wrapper as from_object
 
@@ -32,7 +34,6 @@
 from .sources.array_list import ArrayField
 from .targets import create_target
 from .targets import to_target
-from .utils import array as array_api
 from .utils.examples import download_example_file
 from .utils.examples import remote_example_file
 

src/earthkit/data/core/fieldlist.py

@@ -13,17 +13,17 @@
 from functools import cached_property
 
 import deprecation
+from earthkit.utils.array import array_namespace
+from earthkit.utils.array import array_to_numpy
+from earthkit.utils.array import convert_array
+from earthkit.utils.array import get_backend
 
 from earthkit.data.core import Base
 from earthkit.data.core.index import Index
 from earthkit.data.core.index import MaskIndex
 from earthkit.data.core.index import MultiIndex
 from earthkit.data.decorators import cached_method
 from earthkit.data.decorators import detect_out_filename
-from earthkit.data.utils.array import array_namespace
-from earthkit.data.utils.array import array_to_numpy
-from earthkit.data.utils.array import convert_array
-from earthkit.data.utils.array import get_backend
 from earthkit.data.utils.metadata.args import metadata_argument
 
 

src/earthkit/data/sources/array_list.py

@@ -10,9 +10,10 @@
 import logging
 import math
 
+from earthkit.utils.array import array_namespace
+
 from earthkit.data.core.fieldlist import Field
 from earthkit.data.indexing.fieldlist import ClonedFieldCore
-from earthkit.data.utils.array import array_namespace
 
 LOG = logging.getLogger(__name__)
 

src/earthkit/data/testing.py

@@ -15,6 +15,8 @@
 from importlib import import_module
 from unittest.mock import patch
 
+from earthkit.utils.testing import get_array_backend
+
 from earthkit.data import from_object
 from earthkit.data import from_source
 from earthkit.data.readers.text import TextReader
@@ -121,17 +123,7 @@ def modules_installed(*modules):
 
 NO_POLYTOPE = not os.path.exists(os.path.expanduser("~/.polytopeapirc"))
 NO_COVJSONKIT = not modules_installed("covjsonkit")
-NO_PYTORCH = not modules_installed("torch")
 NO_RIOXARRAY = not modules_installed("rioxarray")
-NO_CUPY = not modules_installed("cupy")
-NO_JAX = not modules_installed("jax")
-if not NO_CUPY:
-    try:
-        import cupy as cp
-
-        a = cp.ones(2)
-    except Exception:
-        NO_CUPY = True
 
 NO_S3_AUTH = not modules_installed("aws_requests_auth")
 NO_GEO = not modules_installed("earthkit-data")
@@ -187,34 +179,8 @@ def load_nc_or_xr_source(path, mode):
         return from_object(xarray.open_dataset(path))
 
 
-def check_array_type(array, expected_backend, dtype=None):
-    from earthkit.data.utils.array import get_backend
-
-    b1 = get_backend(array)
-    b2 = get_backend(expected_backend)
-
-    assert b1 == b2, f"{b1=}, {b2=}"
-
-    expected_dtype = dtype
-    if expected_dtype is not None:
-        assert b2.match_dtype(array, expected_dtype), f"{array.dtype}, {expected_dtype=}"
-
-
-def get_array_namespace(backend):
-    if backend is None:
-        backend = "numpy"
-
-    from earthkit.data.utils.array import get_backend
-
-    return get_backend(backend).namespace
-
-
-ARRAY_BACKENDS = ["numpy"]
-if not NO_PYTORCH:
-    ARRAY_BACKENDS.append("pytorch")
-
-if not NO_CUPY:
-    ARRAY_BACKENDS.append("cupy")
+# Array backends
+ARRAY_BACKENDS = get_array_backend(["numpy", "torch", "cupy", "jax"], raise_on_missing=False)
 
 
 def make_tgz(target_dir, target_name, paths):