From 4ad6f8c08d30b540e44c8e6f0b5cf528ca84970b Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 9 Sep 2025 11:04:32 +0100 Subject: [PATCH 1/6] Rebase to develop --- src/earthkit/data/indexing/tensor.py | 4 +- src/earthkit/data/utils/xarray/grib.py | 60 ++++++++++---- tests/xr_engine/test_xr_time.py | 22 +++++ tests/xr_engine/test_xr_write.py | 108 ++++++++----------------- 4 files changed, 103 insertions(+), 91 deletions(-) diff --git a/src/earthkit/data/indexing/tensor.py b/src/earthkit/data/indexing/tensor.py index 006e360eb..9a3b504c2 100644 --- a/src/earthkit/data/indexing/tensor.py +++ b/src/earthkit/data/indexing/tensor.py @@ -503,7 +503,9 @@ def make_valid_datetime(self, dims_map, dtype="datetime64[ns]"): dims_step = [keys[d] for d in dims] # use same dim order as in user_dims dims = [d for d in self.user_dims if d in dims_step] - assert len(dims) == len(dims_step), f"Duplicate dims in {dims}" + if len(dims) != len(dims_step): + continue + assert len(dims) == len(dims_step), f"{dims=} {dims_step=}" other_dims = [d for d in self.user_dims if d not in dims] if other_dims: diff --git a/src/earthkit/data/utils/xarray/grib.py b/src/earthkit/data/utils/xarray/grib.py index 4f81f74d1..f4d52eefd 100644 --- a/src/earthkit/data/utils/xarray/grib.py +++ b/src/earthkit/data/utils/xarray/grib.py @@ -16,15 +16,18 @@ from earthkit.data.utils.dates import step_to_grib from earthkit.data.utils.dates import time_to_grib from earthkit.data.utils.dates import to_datetime +from earthkit.data.utils.dates import to_timedelta LOG = logging.getLogger(__name__) -def update_metadata(metadata, compulsory): +def update_metadata(metadata, compulsory, step_len=0): if "valid_time" in metadata: dt = to_datetime(metadata.pop("valid_time")) metadata["date"] = dt.date() metadata["time"] = dt.time() + metadata["stepRange"] = step_to_grib(metadata.pop("stepRange", 0)) + metadata.pop("step", None) if "forecast_reference_time" in metadata: date, time = datetime_to_grib(to_datetime(metadata["forecast_reference_time"])) @@ -43,7 +46,19 @@ def update_metadata(metadata, compulsory): metadata["time"] = date.hour * 100 + date.minute if "step" in metadata: - metadata["step"] = step_to_grib(metadata["step"]) + if step_len is None: + metadata["step"] = step_to_grib(metadata["step"]) + elif step_len.total_seconds() == 0: + step = step_to_grib(metadata["step"]) + metadata["stepRange"] = step + metadata.pop("step", None) + elif step_len.total_seconds() > 0: + end = metadata["step"] + start = to_timedelta(end) - step_len + start = step_to_grib(start) + end = step_to_grib(end) + metadata["stepRange"] = f"{start}-{end}" + metadata.pop("step", None) if "stream" not in metadata: if "number" in metadata: @@ -104,7 +119,28 @@ def data_array_to_fields(da, metadata=None): else: coords[k] = coords[k].values - # print(f"{coords=}") + # extract metadata template from dataarray + if hasattr(da, "earthkit"): + template_metadata = da.earthkit.metadata + else: + raise ValueError("Earthkit attribute not found in DataArray. Required for conversion to FieldList!") + + # special treatment to set step related GRIB keys + compulsory_metadata = {} + step_len = datetime.timedelta(hours=0) + if "valid_time" in dims: + # when valid_time is a dimension we enforce the step to be 0 + compulsory_metadata["stepRange"] = 0 + else: + try: + step_range = template_metadata.get("stepRange", None) + if isinstance(step_range, str) and "-" in step_range: + step_len = to_timedelta(template_metadata.get("endStep", 0)) - to_timedelta( + template_metadata.get("startStep", 0) + ) + except TypeError as e: + print(f"Error calculating step length: {e}") + step_len = None for values in product(*[coords[dim] for dim in dims]): @@ -124,16 +160,12 @@ def data_array_to_fields(da, metadata=None): grib_metadata.update(dict(zip(components[k], grib_metadata[k][1:]))) # print(f"-> {grib_metadata=}") del grib_metadata[k] - update_metadata(grib_metadata, []) - - # extract metadata from object - if metadata is None: - if hasattr(da, "earthkit"): - metadata = da.earthkit.metadata - else: - raise ValueError( - "Earthkit attribute not found in DataArray. Required for conversion to FieldList!" - ) - metadata = metadata.override(grib_metadata) + for k in compulsory_metadata: + if k not in grib_metadata: + grib_metadata[k] = compulsory_metadata[k] + + update_metadata(grib_metadata, [], step_len=step_len) + + metadata = template_metadata.override(grib_metadata) yield ArrayField(xa_field.values, metadata) diff --git a/tests/xr_engine/test_xr_time.py b/tests/xr_engine/test_xr_time.py index 6c01efe7e..82baf8b68 100644 --- a/tests/xr_engine/test_xr_time.py +++ b/tests/xr_engine/test_xr_time.py @@ -574,3 +574,25 @@ def test_xr_time_step_range_2(kwargs, dims, step_units): assert ( ds[step_units[0]].attrs["units"] == step_units[1] ), f"step units mismatch {ds[step_units[0]].attrs['units']} != {step_units[1]}" + + +@pytest.mark.cache +def test_xr_time_forecast_per_month(): + ds_ek = from_source( + "url", earthkit_remote_test_data_file("test-data/xr_engine/date/2_months_6_hourly.grib") + ) + + ds = ds_ek.to_xarray(time_dim_mode="valid_time") + + ref = [] + start = np.datetime64("1979-01-01T06:00:00", "ns") + end = np.datetime64("1979-03-01T00:00:00", "ns") + while start <= end: + ref.append(np.datetime64(start)) + start += np.timedelta64(6, "h") + + dims = { + "valid_time": ref, + } + + compare_dims(ds, dims, order_ref_var="avg_dis") diff --git a/tests/xr_engine/test_xr_write.py b/tests/xr_engine/test_xr_write.py index 2cd4f1274..8f98d547e 100644 --- a/tests/xr_engine/test_xr_write.py +++ b/tests/xr_engine/test_xr_write.py @@ -16,6 +16,7 @@ from earthkit.data import to_target from earthkit.data.core.temporary import temp_file from earthkit.data.testing import earthkit_remote_test_data_file +from earthkit.data.utils.dates import datetime_to_grib @pytest.mark.cache @@ -256,6 +257,8 @@ def test_xr_write_bits_per_value(): assert r[0].metadata("bitsPerValue") == 8 + + @pytest.mark.cache @pytest.mark.parametrize("method", ["to_grib", "to_target_on_obj", "to_target_func"]) @pytest.mark.parametrize( @@ -334,6 +337,7 @@ def test_xr_write_to_grib_file_dataarray(method, kwargs): assert r.metadata("valid_datetime") == ref_valid + @pytest.mark.cache @pytest.mark.parametrize("method", ["to_grib", "to_target_on_obj", "to_target_func"]) @pytest.mark.parametrize( @@ -375,85 +379,37 @@ def test_xr_write_to_grib_file_dataset(method, kwargs): assert sorted(r.metadata("valid_datetime")) == sorted(ds_ek.metadata("valid_datetime")) -@pytest.mark.cache -@pytest.mark.parametrize("method", ["to_netcdf", "to_target_on_obj", "to_target_func"]) -@pytest.mark.parametrize( - "kwargs", - [ - {"profile": "mars", "time_dim_mode": "raw"}, - ], -) -def test_xr_write_to_netcdf_file_dataarray(method, kwargs): - ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/level/pl.grib")) - ds_ek = ds_ek.sel(param=["t"], level=[500, 850]) - - ref_t_vals = ds_ek.sel(param="t", step=6, level=500).to_numpy().flatten() - - import xarray as xr - - xr.set_options(keep_attrs=True) - - ds = ds_ek.to_xarray(**kwargs) - ds += 1 - - # data-array - with temp_file() as path: - if method == "to_netcdf": - ds["t"].earthkit.to_netcdf(path) - elif method == "to_target_on_obj": - ds["t"].earthkit.to_target("file", path, encoder="netcdf") - elif method == "to_target_func": - to_target("file", path, data=ds["t"], encoder="netcdf") - - r = xr.open_dataset(path, engine="netcdf4", decode_times=False) - - for name in ["t"]: - assert name in r.data_vars - - for name in ["latitude", "longitude"]: - assert name in r.coords - assert np.allclose(ref_t_vals + 1.0, r["t"].isel(step=1, level=0).to_numpy().flatten()) @pytest.mark.cache -@pytest.mark.parametrize("method", ["to_netcdf", "to_target_on_obj", "to_target_func"]) -@pytest.mark.parametrize( - "kwargs", - [ - {"profile": "mars", "time_dim_mode": "raw"}, - ], -) -def test_xr_write_to_netcdf_file_dataset(method, kwargs): - ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/level/pl.grib")) - ds_ek = ds_ek.sel(param=["t", "r"], level=[500, 850]) - - ref_t_vals = ds_ek.sel(param="t", step=6, level=500).to_numpy().flatten() - ref_r_vals = ds_ek.sel(param="r", step=6, level=500).to_numpy().flatten() - - import xarray as xr - - xr.set_options(keep_attrs=True) - - ds = ds_ek.to_xarray(**kwargs) - ds += 1 - - # dataset - with temp_file() as path: - if method == "to_netcdf": - ds.earthkit.to_netcdf(path) - elif method == "to_target_on_obj": - ds.earthkit.to_target("file", path, encoder="netcdf") - elif method == "to_target_func": - to_target("file", path, data=ds, encoder="netcdf") - - r = xr.open_dataset(path, engine="netcdf4", decode_times=False) - - for name in ["t", "r"]: - assert name in r.data_vars +def test_xr_write_forecast_per_month(): + ds_ek = from_source( + "url", earthkit_remote_test_data_file("xr_engine/date/2_months_6_hourly.grib") + ) - for name in ["latitude", "longitude"]: - assert name in r.coords + ds = ds_ek.to_xarray(time_dim_mode="valid_time") + r = ds.earthkit.to_fieldlist() + assert len(r) == 236 + + ref = [] + start = np.datetime64("1979-01-01T06:00:00", "ns") + end = np.datetime64("1979-03-01T00:00:00", "ns") + while start <= end: + base_date, base_time = datetime_to_grib(start) + ref.append([base_date, base_time, 0, "0", 0, 0, base_date, base_time]) + start += np.timedelta64(6, "h") + + keys = [ + "dataDate", + "dataTime", + "step", + "stepRange", + "startStep", + "endStep", + "validityDate", + "validityTime", + ] - assert np.allclose(ref_t_vals + 1.0, r["t"].isel(step=1, level=0).to_numpy().flatten()) - assert np.allclose(ref_r_vals + 1.0, r["r"].isel(step=1, level=0).to_numpy().flatten()) + for f, f_ref in zip(r, ref): + assert f.metadata(keys) == f_ref, f"Expected: {f_ref}\nGot: {f.metadata(keys)}" From c6d275c02e2f76c8ccdb1dbc4e52cda7c2273b76 Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 1 Jul 2025 14:32:09 +0100 Subject: [PATCH 2/6] Fix writing temporally averaged Xarray data to GRIB --- docs/guide/xarray/grib_to_netcdf.rst | 66 ++++++++++++++++++++++++++++ docs/guide/xarray/index.rst | 3 ++ docs/guide/xarray/to_grib.rst | 65 +++++++++++++++++++++++++++ 3 files changed, 134 insertions(+) create mode 100644 docs/guide/xarray/grib_to_netcdf.rst create mode 100644 docs/guide/xarray/to_grib.rst diff --git a/docs/guide/xarray/grib_to_netcdf.rst b/docs/guide/xarray/grib_to_netcdf.rst new file mode 100644 index 000000000..31bf63988 --- /dev/null +++ b/docs/guide/xarray/grib_to_netcdf.rst @@ -0,0 +1,66 @@ + +.. _xr_grib_to_netcdf: + + +Converting GRIB to NetCDF +---------------------------- + +To convert GRIB data to NetCDF first we need to convert GRIB to Xarray with :py:meth:`~data.readers.grib.index.GribFieldList.to_xarray` then generate NetCDF from it with :py:meth:`xarray.Dataset.to_netcdf`. We have 3 options to do this: + +Using the earthkit accessor +++++++++++++++++++++++++++++ + +By default, the earthkit Xarray engine attaches some special attributes to the generated Xarray dataset that cannot be written to NetCDF. In order to make ``to_netcdf()`` work we need to invoke it on the ``earthkit`` accessor and not directly on the Xarray dataset. + +.. code-block:: python + + import earthkit.data as ekd + + ds_fl = ekd.from_source("sample", "pl.grib") + ds_xr = ds_fl.to_xarray() + ds_xr.earthkit.to_netcdf("_from_grib.nc") + +Using the ``add_earthkit_attrs=False`` option +++++++++++++++++++++++++++++++++++++++++++++++++++ + +Alternatively, we can use the ``add_earthkit_attrs=False`` option in :py:meth:`~data.readers.grib.index.GribFieldList.to_xarray`. With this the earthkit attributes are not added to the generated dataset and it is safe to call :py:meth:`to_netcdf ` directly on it. + +.. code-block:: python + + import earthkit.data as ekd + + ds_fl = ekd.from_source("sample", "pl.grib") + ds_xr = ds_fl.to_xarray(add_earthkit_attrs=False) + ds_xr.to_netcdf("_from_grib.nc") + +Using to_target +++++++++++++++++ + +The third option is to use the :func:`to_target` method to convert GRIB directly to NetCDF. This method will generate an Xarray dataset and write it to a NetCDF file in one step. + +.. code-block:: python + + import earthkit.data as ekd + + ds_fl = ekd.from_source("sample", "pl.grib") + ds.fl.to_target("file", "_from_grib.nc") + + +To control the Xarray conversion we can pass options to the earthkit Xarray engine with ``earthkit_to_xarray_kwargs``. In this case ``add_earthkit_attrs=False`` is always enforced. + +.. code-block:: python + + import earthkit.data as ekd + + ds_fl = ekd.from_source("sample", "pl.grib") + ds.fl.to_target( + "file", "_from_grib.nc", earthkit_to_xarray_kwargs={"flatten_values": True} + ) + + +Examples +++++++++++++ + +For further details see the following notebook: + +- :ref:`/examples/grib_to_netcdf.ipynb` diff --git a/docs/guide/xarray/index.rst b/docs/guide/xarray/index.rst index 0460a825d..0c5ee4fb0 100644 --- a/docs/guide/xarray/index.rst +++ b/docs/guide/xarray/index.rst @@ -5,6 +5,9 @@ Xarray engine :maxdepth: 1 overview + dim + to_grib + grib_to_netcdf profile mars_profile grib_profile diff --git a/docs/guide/xarray/to_grib.rst b/docs/guide/xarray/to_grib.rst new file mode 100644 index 000000000..26b5c0d17 --- /dev/null +++ b/docs/guide/xarray/to_grib.rst @@ -0,0 +1,65 @@ + +.. _xr_to_grib: + + +Converting Xarray to GRIB +------------------------- + +.. warning:: + + This is an experimental feature and it is not yet fully supported. + +By default, ``add_earthkit_attrs=True`` in :py:meth:`~data.readers.grib.index.GribFieldList.to_xarray` and some special earthkit attributes are added to the dataset. This is needed for the Xarray to GRIB conversion. For this reason, if the Xarray is modified we must ensure the variable attributes are copied to the new Xarray dataset. By default, variable attributes are not kept in Xarray computations so we need to set the global Xarray ``keep_attrs`` option to enable it. See the examples below for details. + +.. note:: + + When Xarray generated with ``time_dim_mode="valid_time"`` is converted to GRIB, the "stepRange"/"step" keys in all the resulting GRIB fields will be set to 0, regardless of wether the original GRIB data contained forecast or not. + +Using to_target +++++++++++++++++ + +It is possible to directly write the Xarray dataset created with the earthkit engine into a GRIB file with :func:`to_target`. This is a memory efficient way to write GRIB to disk since only one field is loaded into memory at a time. We can call :func:`to_target` either on the ``earthkit`` accessor or as a top level function. + +.. code-block:: python + + # ensure attributes are kept + import xarray as xr + + xr.set_options(keep_attrs=True) + + # ds_xr is an Xarray dataset created with the earthkit engine, we modify it + ds_xr += 1 + + # option1: writing to GRIB file using the accessor + ds_xr.earthkit.to_target("file", "_from_xr_1.grib") + + # option: 2writing to GRIB file using the top level function + to_target("file", "_from_xr_2.grib", data=ds_xr) + + +Using to_fieldlist() +++++++++++++++++++++ + +We can also convert the Xarray dataset into a GRIB fieldlist by using :py:meth:`~data.utils.xarray.engine.XarrayEarthkit.to_fieldlist` on the ``earthkit`` accessor of the Xarray object. Please note that this will generate a fieldlist entirely stored in memory. + +.. code-block:: python + + >>> import xarray as xr + >>> xr.set_options(keep_attrs=True) + >>> ds_xr += 1 + >>> ds_fl1 = ds_xr.earthkit.to_fieldlist() + >>> ds_fl1[0] + ArrayField(r,500,20240603,0,0,0) + +The generated GRIB fieldlist can be saved to disk using the :func:`to_target` method. + +.. code-block:: python + + ds_fl1.to_target("file", "_from_xr_3.grib") + +Examples +++++++++++++ + +For further details see the following notebook: + +- :ref:`/examples/xarray_engine_to_grib.ipynb` From 3cbe8df9a73529f483105bb92ba1acfd63946bb0 Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 1 Jul 2025 14:49:53 +0100 Subject: [PATCH 3/6] Fix writing temporally averaged Xarray data to GRIB --- src/earthkit/data/utils/xarray/check.py | 28 ++++++++++++------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/src/earthkit/data/utils/xarray/check.py b/src/earthkit/data/utils/xarray/check.py index f75371974..6f2998fc4 100644 --- a/src/earthkit/data/utils/xarray/check.py +++ b/src/earthkit/data/utils/xarray/check.py @@ -48,17 +48,20 @@ def __init__(self, tensor): self.tensor = tensor def first_diff(self, coord_keys): - for i, f in enumerate(self.tensor.source): - t_coords = self.tensor._index_to_coords_value(i, self.tensor) - f_coords = f.metadata(coord_keys) - print(f"Checking field[{i}] with coords {t_coords} vs {f_coords}") - diff = ListDiff.diff(t_coords, f_coords) - if not diff.same: - name = "" - if diff.diff_index != -1: - name = coord_keys[diff.diff_index] - - return i, f, t_coords, f_coords, name, diff + if coord_keys: + for i, f in enumerate(self.tensor.source): + t_coords = self.tensor._index_to_coords_value(i, self.tensor) + f_coords = f.metadata(coord_keys) + try: + diff = ListDiff.diff(t_coords, f_coords) + if not diff.same: + name = "" + if diff.diff_index != -1: + name = coord_keys[diff.diff_index] + + return i, f, t_coords, f_coords, name, diff + except Exception: + pass def neighbour_field(self, field_num, index): f_other = None @@ -104,9 +107,6 @@ def check(self, details=False): f"Dimensions: \n {dims}" ) - print(text_num) - print(coord_keys) - if not details: raise ValueError(text_num) else: From f3d6694ff53f55a31d822fef25bb8ea7eddc862c Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 9 Sep 2025 11:19:58 +0100 Subject: [PATCH 4/6] Docs --- docs/guide/xarray/grib_to_netcdf.rst | 66 ---------------------------- docs/guide/xarray/index.rst | 3 -- docs/guide/xarray/to_grib.rst | 65 --------------------------- tests/xr_engine/test_xr_write.py | 10 +---- 4 files changed, 1 insertion(+), 143 deletions(-) delete mode 100644 docs/guide/xarray/grib_to_netcdf.rst delete mode 100644 docs/guide/xarray/to_grib.rst diff --git a/docs/guide/xarray/grib_to_netcdf.rst b/docs/guide/xarray/grib_to_netcdf.rst deleted file mode 100644 index 31bf63988..000000000 --- a/docs/guide/xarray/grib_to_netcdf.rst +++ /dev/null @@ -1,66 +0,0 @@ - -.. _xr_grib_to_netcdf: - - -Converting GRIB to NetCDF ----------------------------- - -To convert GRIB data to NetCDF first we need to convert GRIB to Xarray with :py:meth:`~data.readers.grib.index.GribFieldList.to_xarray` then generate NetCDF from it with :py:meth:`xarray.Dataset.to_netcdf`. We have 3 options to do this: - -Using the earthkit accessor -++++++++++++++++++++++++++++ - -By default, the earthkit Xarray engine attaches some special attributes to the generated Xarray dataset that cannot be written to NetCDF. In order to make ``to_netcdf()`` work we need to invoke it on the ``earthkit`` accessor and not directly on the Xarray dataset. - -.. code-block:: python - - import earthkit.data as ekd - - ds_fl = ekd.from_source("sample", "pl.grib") - ds_xr = ds_fl.to_xarray() - ds_xr.earthkit.to_netcdf("_from_grib.nc") - -Using the ``add_earthkit_attrs=False`` option -++++++++++++++++++++++++++++++++++++++++++++++++++ - -Alternatively, we can use the ``add_earthkit_attrs=False`` option in :py:meth:`~data.readers.grib.index.GribFieldList.to_xarray`. With this the earthkit attributes are not added to the generated dataset and it is safe to call :py:meth:`to_netcdf ` directly on it. - -.. code-block:: python - - import earthkit.data as ekd - - ds_fl = ekd.from_source("sample", "pl.grib") - ds_xr = ds_fl.to_xarray(add_earthkit_attrs=False) - ds_xr.to_netcdf("_from_grib.nc") - -Using to_target -++++++++++++++++ - -The third option is to use the :func:`to_target` method to convert GRIB directly to NetCDF. This method will generate an Xarray dataset and write it to a NetCDF file in one step. - -.. code-block:: python - - import earthkit.data as ekd - - ds_fl = ekd.from_source("sample", "pl.grib") - ds.fl.to_target("file", "_from_grib.nc") - - -To control the Xarray conversion we can pass options to the earthkit Xarray engine with ``earthkit_to_xarray_kwargs``. In this case ``add_earthkit_attrs=False`` is always enforced. - -.. code-block:: python - - import earthkit.data as ekd - - ds_fl = ekd.from_source("sample", "pl.grib") - ds.fl.to_target( - "file", "_from_grib.nc", earthkit_to_xarray_kwargs={"flatten_values": True} - ) - - -Examples -++++++++++++ - -For further details see the following notebook: - -- :ref:`/examples/grib_to_netcdf.ipynb` diff --git a/docs/guide/xarray/index.rst b/docs/guide/xarray/index.rst index 0c5ee4fb0..0460a825d 100644 --- a/docs/guide/xarray/index.rst +++ b/docs/guide/xarray/index.rst @@ -5,9 +5,6 @@ Xarray engine :maxdepth: 1 overview - dim - to_grib - grib_to_netcdf profile mars_profile grib_profile diff --git a/docs/guide/xarray/to_grib.rst b/docs/guide/xarray/to_grib.rst deleted file mode 100644 index 26b5c0d17..000000000 --- a/docs/guide/xarray/to_grib.rst +++ /dev/null @@ -1,65 +0,0 @@ - -.. _xr_to_grib: - - -Converting Xarray to GRIB -------------------------- - -.. warning:: - - This is an experimental feature and it is not yet fully supported. - -By default, ``add_earthkit_attrs=True`` in :py:meth:`~data.readers.grib.index.GribFieldList.to_xarray` and some special earthkit attributes are added to the dataset. This is needed for the Xarray to GRIB conversion. For this reason, if the Xarray is modified we must ensure the variable attributes are copied to the new Xarray dataset. By default, variable attributes are not kept in Xarray computations so we need to set the global Xarray ``keep_attrs`` option to enable it. See the examples below for details. - -.. note:: - - When Xarray generated with ``time_dim_mode="valid_time"`` is converted to GRIB, the "stepRange"/"step" keys in all the resulting GRIB fields will be set to 0, regardless of wether the original GRIB data contained forecast or not. - -Using to_target -++++++++++++++++ - -It is possible to directly write the Xarray dataset created with the earthkit engine into a GRIB file with :func:`to_target`. This is a memory efficient way to write GRIB to disk since only one field is loaded into memory at a time. We can call :func:`to_target` either on the ``earthkit`` accessor or as a top level function. - -.. code-block:: python - - # ensure attributes are kept - import xarray as xr - - xr.set_options(keep_attrs=True) - - # ds_xr is an Xarray dataset created with the earthkit engine, we modify it - ds_xr += 1 - - # option1: writing to GRIB file using the accessor - ds_xr.earthkit.to_target("file", "_from_xr_1.grib") - - # option: 2writing to GRIB file using the top level function - to_target("file", "_from_xr_2.grib", data=ds_xr) - - -Using to_fieldlist() -++++++++++++++++++++ - -We can also convert the Xarray dataset into a GRIB fieldlist by using :py:meth:`~data.utils.xarray.engine.XarrayEarthkit.to_fieldlist` on the ``earthkit`` accessor of the Xarray object. Please note that this will generate a fieldlist entirely stored in memory. - -.. code-block:: python - - >>> import xarray as xr - >>> xr.set_options(keep_attrs=True) - >>> ds_xr += 1 - >>> ds_fl1 = ds_xr.earthkit.to_fieldlist() - >>> ds_fl1[0] - ArrayField(r,500,20240603,0,0,0) - -The generated GRIB fieldlist can be saved to disk using the :func:`to_target` method. - -.. code-block:: python - - ds_fl1.to_target("file", "_from_xr_3.grib") - -Examples -++++++++++++ - -For further details see the following notebook: - -- :ref:`/examples/xarray_engine_to_grib.ipynb` diff --git a/tests/xr_engine/test_xr_write.py b/tests/xr_engine/test_xr_write.py index 8f98d547e..56756e3aa 100644 --- a/tests/xr_engine/test_xr_write.py +++ b/tests/xr_engine/test_xr_write.py @@ -257,8 +257,6 @@ def test_xr_write_bits_per_value(): assert r[0].metadata("bitsPerValue") == 8 - - @pytest.mark.cache @pytest.mark.parametrize("method", ["to_grib", "to_target_on_obj", "to_target_func"]) @pytest.mark.parametrize( @@ -337,7 +335,6 @@ def test_xr_write_to_grib_file_dataarray(method, kwargs): assert r.metadata("valid_datetime") == ref_valid - @pytest.mark.cache @pytest.mark.parametrize("method", ["to_grib", "to_target_on_obj", "to_target_func"]) @pytest.mark.parametrize( @@ -379,14 +376,9 @@ def test_xr_write_to_grib_file_dataset(method, kwargs): assert sorted(r.metadata("valid_datetime")) == sorted(ds_ek.metadata("valid_datetime")) - - - @pytest.mark.cache def test_xr_write_forecast_per_month(): - ds_ek = from_source( - "url", earthkit_remote_test_data_file("xr_engine/date/2_months_6_hourly.grib") - ) + ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/date/2_months_6_hourly.grib")) ds = ds_ek.to_xarray(time_dim_mode="valid_time") r = ds.earthkit.to_fieldlist() From e6e066659e26f49449fba6f8cc14e5dc5d536c2b Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 9 Sep 2025 11:32:25 +0100 Subject: [PATCH 5/6] Fix test --- tests/xr_engine/test_xr_time.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/xr_engine/test_xr_time.py b/tests/xr_engine/test_xr_time.py index 82baf8b68..4394df89b 100644 --- a/tests/xr_engine/test_xr_time.py +++ b/tests/xr_engine/test_xr_time.py @@ -578,9 +578,7 @@ def test_xr_time_step_range_2(kwargs, dims, step_units): @pytest.mark.cache def test_xr_time_forecast_per_month(): - ds_ek = from_source( - "url", earthkit_remote_test_data_file("test-data/xr_engine/date/2_months_6_hourly.grib") - ) + ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/date/2_months_6_hourly.grib")) ds = ds_ek.to_xarray(time_dim_mode="valid_time") From 7f913377082d9fcdb181e74f7bc7be3ce0d2ec9e Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 9 Sep 2025 11:36:13 +0100 Subject: [PATCH 6/6] Test --- tests/xr_engine/test_xr_write.py | 84 ++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/tests/xr_engine/test_xr_write.py b/tests/xr_engine/test_xr_write.py index 56756e3aa..0cee0ceac 100644 --- a/tests/xr_engine/test_xr_write.py +++ b/tests/xr_engine/test_xr_write.py @@ -376,6 +376,90 @@ def test_xr_write_to_grib_file_dataset(method, kwargs): assert sorted(r.metadata("valid_datetime")) == sorted(ds_ek.metadata("valid_datetime")) +@pytest.mark.cache +@pytest.mark.parametrize("method", ["to_netcdf", "to_target_on_obj", "to_target_func"]) +@pytest.mark.parametrize( + "kwargs", + [ + {"profile": "mars", "time_dim_mode": "raw"}, + ], +) +def test_xr_write_to_netcdf_file_dataarray(method, kwargs): + ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/level/pl.grib")) + ds_ek = ds_ek.sel(param=["t"], level=[500, 850]) + + ref_t_vals = ds_ek.sel(param="t", step=6, level=500).to_numpy().flatten() + + import xarray as xr + + xr.set_options(keep_attrs=True) + + ds = ds_ek.to_xarray(**kwargs) + ds += 1 + + # data-array + with temp_file() as path: + if method == "to_netcdf": + ds["t"].earthkit.to_netcdf(path) + elif method == "to_target_on_obj": + ds["t"].earthkit.to_target("file", path, encoder="netcdf") + elif method == "to_target_func": + to_target("file", path, data=ds["t"], encoder="netcdf") + + r = xr.open_dataset(path, engine="netcdf4", decode_times=False) + + for name in ["t"]: + assert name in r.data_vars + + for name in ["latitude", "longitude"]: + assert name in r.coords + + assert np.allclose(ref_t_vals + 1.0, r["t"].isel(step=1, level=0).to_numpy().flatten()) + + +@pytest.mark.cache +@pytest.mark.parametrize("method", ["to_netcdf", "to_target_on_obj", "to_target_func"]) +@pytest.mark.parametrize( + "kwargs", + [ + {"profile": "mars", "time_dim_mode": "raw"}, + ], +) +def test_xr_write_to_netcdf_file_dataset(method, kwargs): + ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/level/pl.grib")) + ds_ek = ds_ek.sel(param=["t", "r"], level=[500, 850]) + + ref_t_vals = ds_ek.sel(param="t", step=6, level=500).to_numpy().flatten() + ref_r_vals = ds_ek.sel(param="r", step=6, level=500).to_numpy().flatten() + + import xarray as xr + + xr.set_options(keep_attrs=True) + + ds = ds_ek.to_xarray(**kwargs) + ds += 1 + + # dataset + with temp_file() as path: + if method == "to_netcdf": + ds.earthkit.to_netcdf(path) + elif method == "to_target_on_obj": + ds.earthkit.to_target("file", path, encoder="netcdf") + elif method == "to_target_func": + to_target("file", path, data=ds, encoder="netcdf") + + r = xr.open_dataset(path, engine="netcdf4", decode_times=False) + + for name in ["t", "r"]: + assert name in r.data_vars + + for name in ["latitude", "longitude"]: + assert name in r.coords + + assert np.allclose(ref_t_vals + 1.0, r["t"].isel(step=1, level=0).to_numpy().flatten()) + assert np.allclose(ref_r_vals + 1.0, r["r"].isel(step=1, level=0).to_numpy().flatten()) + + @pytest.mark.cache def test_xr_write_forecast_per_month(): ds_ek = from_source("url", earthkit_remote_test_data_file("xr_engine/date/2_months_6_hourly.grib"))