feature: experimental gribjump source (#689)

* wip: first rough draft of a GribJumpSource

* wip: experimental tests for easier development

* format changes using pre-commit hooks and fix small bug

* add prototype for the GribJumpSource based on SimpleFieldList

* tests: add a few more simple tests and add NO_GRIBJUMP flag for pytest

* tidy: small cleanup, improve variable naming and fix type hints

* use original type in request dictionaries to make .sel more intuitive

* assign grid index to each value in to_xarray

* tidy: add some more error handling and improve tests

* refactor: introduce ExtractionRequest wrapper that combines pygribjump.ExtractionRequest and original fdb request dict

* feat(test): modify (now failing) test to expect latitude and longitude information

* feat: wip: allow reference lat/lons to be loaded from an fdb reference field

* refactor: move hardcoded test fixtures into pytest fixtures

* test: add failing test showing bug with geography for gridded extracts

* docs: add notebook draft with example usage of gribjump source

* tidy: move validation that extract request share the same ranges

* docs: add documentation for gribjump source

* docs: small fixes of markdown syntax

* feat: wip experiment to verify gridspec of reference field

* fix: force flattened array in xarray dataset creation

* refactor: tidy up the metadata enrichment a bit

* refactor: create ExtractionRequestCollection

* refactor: use FDBRetriever to load reference metadata

* docs: add example for masks and indices to notebook

* feat: enforce that masks are 1D boolean arrays

* refactor: simplify by condensing request splitting utilities into one function

* feat: remove verifiation functionality for now, to be added later

* tidy: small renamings and docstrings

* tidy: comments

* fix: type hint and name

* feat: convert masks to ranges once for significant speedups

* test: add another test for mask_to_ranges

* tidy: small comment/docstring changes

* tidy: make warning about missing validation in docs more explicit

* docs: improve wording of warning

* fix: allow fdb and gribjump to be configured via FDB5_CONFIG and GRIBJUMP_HOME

* chore: update docstring

* feat: pass log context to gribjump

* refactor: simplify gribjump log context

* test: add t_gribjump.grib test data with expver xxxx

* add pygribjump as an optional dependency

* docs: clarify gribjump install instructions and dependency handling

* docs: move warning before parameters section

* docs: clarify parameter description and types

* add pyfdb as a gribjump group dependency and update docs

* last docs and typo fixes

* docs: change notebook to also set FDB_HOME

* docs: reference gribjump example notebook in missing locations

Author: andreas-grafberger

docs/examples/gribjump.ipynb

@@ -30,6 +30,7 @@ Data sources
     polytope_feature.ipynb
     s3.ipynb
     wekeo.ipynb
+    gribjump.ipynb
 
 GRIB
 ++++++

docs/examples/index.rst

@@ -66,6 +66,8 @@ We can get data from a given source by using :func:`from_source`:
       - retrieve data from `WEkEO`_ using the WEkEO grammar
     * - :ref:`data-sources-wekeocds`
       - retrieve `CDS <https://cds.climate.copernicus.eu/>`_ data stored on `WEkEO`_ using the `cdsapi`_ grammar
+    * - :ref:`data-sources-gribjump`
+      - retrieve data from the `FDB (Fields DataBase)`_ using the `gribjump`_ library
     * - :ref:`data-sources-zarr`
       - load data from a `Zarr <https://zarr.readthedocs.io/en/stable/>`_ store
 
@@ -1231,6 +1233,85 @@ wekeocds
       - :ref:`/examples/wekeo.ipynb`
 
 
+.. _data-sources-gribjump:
+
+gribjump
+--------
+
+.. py:function:: from_source("gribjump", request, *, ranges=None, mask=None, indices=None, fetch_coords_from_fdb=False, fdb_kwargs=None, **kwargs)
+  :noindex:
+
+  The ``gribjump`` source enables fast retrieval of GRIB message subsets from the `FDB (Fields DataBase)`_ using the `gribjump <https://github.com/ecmwf/gribjump/>`_ library.
+  Both `pygribjump <https://pypi.org/project/pygribjump/>`_ and `pyfdb`_ must be installed. The `pygribjump`_ package uses `findlibs <https://github.com/ecmwf/findlibs>`_ to locate an installation of the `gribjump`_ library.
+  If the library is not available on your system, you can install it via the `gribjumplib <https://pypi.org/project/gribjumplib/>`_ wheel from PyPI.
+  Installing `gribjumplib` from PyPI will also automatically install `fdb5lib <https://pypi.org/project/fdb5lib/>`_ and other dependencies, which may take priority over any existing installations on your system.
+
+  .. warning::
+    ⚠️ This source is **experimental** and may change in future versions without
+    warning. It performs **no validation** that the specified grid indices,
+    masks, or ranges correspond to the fields' actual underlying grids.
+    **Incorrect usage may silently return wrong data points.**
+    The provided ranges or masks might correspond to unexpected points on the
+    grid.  This source is also currently **not thread-safe**.
+
+  Exactly one of the parameters ``ranges``, ``mask`` or ``indices`` must be specified at a time.
+
+  :param request: the FDB request as a dictionary. GribJump requires strict value formatting
+      (e.g., hdates as "YYYYMMDD", not "YYYY-MM-DD"). Format errors may result in "DataNotFound" errors.
+  :type request: dict
+  :param ranges: a list of tuples specifying the ranges of 1D grid indices to retrieve in the form
+      [(start1, end1), (start2, end2), ...]. Ranges are exclusive, meaning that the end index is not included in the range.
+  :type ranges: list[tuple[int, int]], optional
+  :param mask: a 1D boolean mask specifying which grid points to retrieve
+  :type mask: numpy.array, optional
+  :param indices: a 1D array of grid indices to retrieve
+  :type indices: numpy.array, optional
+  :param fetch_coords_from_fdb: if ``True``, loads the first field's metadata from
+      the FDB to extract the coordinates at the specified indices. If ``False``, the
+      coordinates are not loaded and no separate FDB request is made.
+      Default is ``False``. Please note that no validation is performed to
+      ensure that all fields in the requests share the same grid.
+  :type fetch_coords_from_fdb: bool, optional
+  :param fdb_kwargs: only used when ``fetch_coords_from_fdb=True``. A dict of
+      keyword arguments passed to the `pyfdb.FDB` constructor. This allows to
+      specify the FDB configuration, user configuration, etc. If not provided,
+      the default configuration is used. These arguments are only passed to the
+      FDB when fetching coordinates and are not used by GribJump for the
+      extraction itself.
+  :type fdb_kwargs: dict, optional
+
+
+  The following example retrieves a subset from a GRIB message in the FDB using a boolean mask:
+
+  .. code-block:: python
+
+      import earthkit.data as ekd
+      import numpy as np
+
+      request = {
+          "class": "od",
+          "type": "fc",
+          "stream": "oper",
+          "expver": "0001",
+          "repres": "gg",
+          "levtype": "sfc",
+          "param": "2t",
+          "date": "20250703",
+          "time": 0,
+          "step": list(range(0, 24, 6)),
+          "domain": "g",
+      }
+
+      ranges = [(0, 10), (20, 30)]
+
+      source = ekd.from_source("gribjump", request, ranges=ranges)
+      ds = source.to_xarray()
+
+  Further examples:
+
+      - :ref:`/examples/gribjump.ipynb`
+
+
 
 .. _data-sources-zarr:
 

docs/guide/sources.rst

@@ -51,6 +51,7 @@ Alternatively, you can install the following components:
     - covjsonkit: provides access to CoverageJSON data served by the :ref:`data-sources-polytope` source
     - s3: provides access to non-public :ref:`s3 <data-sources-s3>` buckets (new in version *0.11.0*)
     - geotiff: adds GeoTIFF support (new in version *0.11.0*). Please note that this is not included in the ``[all]`` option and has to be invoked separately.
+    - gribjump: provides access to the :ref:`data-sources-gribjump` source
     - zarr: provides access to the :ref:`data-sources-zarr` source (new in version *0.15.0*). Please note that this is not included in the ``[all]`` option and has to be invoked separately.
 
 E.g. to add :ref:`data-sources-mars`  support you can use:
@@ -85,3 +86,9 @@ FDB
 +++++
 
 For FDB (Fields DataBase) access FDB5 must be installed on the system. See the `FDB documentation <https://fields-database.readthedocs.io/en/latest/>`_ for details.
+
+
+GribJump
+++++++++++++
+
+For FDB access with GribJump, both FDB5 and GribJump must be installed on the system. See the `GribJump project <https://github.com/ecmwf/gribjump>`_ for details.

docs/install.rst

@@ -47,7 +47,7 @@ dependencies = [
   "xarray>=0.19",
 ]
 optional-dependencies.all = [
-  "earthkit-data[cds,covjsonkit,ecmwf-opendata,fdb,geo,geopandas,mars,odb,polytope,projection,s3,wekeo]",
+  "earthkit-data[cds,covjsonkit,ecmwf-opendata,fdb,geo,geopandas,gribjump,mars,odb,polytope,projection,s3,wekeo]",
 ]
 optional-dependencies.cds = [ "cdsapi>=0.7.2" ]
 optional-dependencies.ci = [ "numpy" ]
@@ -70,6 +70,7 @@ optional-dependencies.fdb = [ "pyfdb>=0.1" ]
 optional-dependencies.geo = [ "earthkit-geo>=0.2" ]
 optional-dependencies.geopandas = [ "geopandas" ]
 optional-dependencies.geotiff = [ "pyproj", "rasterio", "rioxarray" ]
+optional-dependencies.gribjump = [ "pyfdb>=0.1", "pygribjump" ]
 optional-dependencies.mars = [ "ecmwf-api-client>=1.6.1" ]
 optional-dependencies.odb = [ "pyodc" ]
 optional-dependencies.polytope = [ "polytope-client>=0.7.6" ]