Feature/new datasets (#99)

* main changes

* bugfix

* few bugs and add unit tests

* work with more planetary computer ds

* add optional dependencies

Author: mariahpope

pyproject.toml

@@ -97,6 +97,9 @@ optional-dependencies.tests = [
 optional-dependencies.xarray = [
   "gcsfs",
   "kerchunk",
+  "pandas", 
+  "pystac_client",
+  "planetary_computer",
 ]
 
 urls.Documentation = "https://anemoi-datasets.readthedocs.io/"

src/anemoi/datasets/create/__init__.py

@@ -7,6 +7,7 @@
 # nor does it submit to any jurisdiction.
 #
 
+import cftime
 import datetime
 import json
 import logging
@@ -66,6 +67,19 @@ def json_tidy(o):
     if isinstance(o, datetime.timedelta):
         return frequency_to_string(o)
 
+    if isinstance(o, cftime.DatetimeJulian):
+        import pandas as pd
+        
+        o = pd.Timestamp(
+            o.year,
+            o.month,
+            o.day,
+            o.hour,
+            o.minute,
+            o.second,
+        )
+        return o.isoformat()
+
     raise TypeError(repr(o) + " is not JSON serializable")
 
 

src/anemoi/datasets/create/functions/sources/xarray/__init__.py

@@ -46,6 +46,12 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
 
     if isinstance(dataset, str) and ".zarr" in dataset:
         data = xr.open_zarr(name_to_zarr_store(dataset), **options)
+    elif "planetarycomputer" in dataset:
+        store = name_to_zarr_store(dataset)
+        if "store" in store:
+            data = xr.open_zarr(**store)
+        if "filename_or_obj" in store:
+            data = xr.open_dataset(**store)
     else:
         data = xr.open_dataset(dataset, **options)
 

src/anemoi/datasets/create/functions/sources/xarray/fieldlist.py

@@ -60,6 +60,15 @@ def from_xarray(cls, ds, flavour=None):
                     flavour = yaml.safe_load(f)
                 else:
                     flavour = json.load(f)
+                    
+        if isinstance(flavour, dict):
+            flavour_coords = [coords["name"] for coords in flavour["rules"].values()]
+            ds_dims = [dim for dim in ds._dims]
+            for dim in ds_dims:
+                    if dim in flavour_coords and dim not in ds._coord_names:
+                        ds = ds.assign_coords({dim:ds[dim]})
+                    else:
+                        pass
 
         guess = CoordinateGuesser.from_flavour(ds, flavour)
 

src/anemoi/datasets/create/functions/sources/xarray/variable.py

@@ -117,7 +117,7 @@ def sel(self, missing, **kwargs):
 
         variable = Variable(
             ds=self.ds,
-            var=self.variable.isel({k: i}),
+            variable=self.variable.isel({k: i}),
             coordinates=coordinates,
             grid=self.grid,
             time=self.time,

src/anemoi/datasets/data/stores.py

@@ -85,6 +85,41 @@ def __getitem__(self, key):
         return response["Body"].read()
 
 
+class PlanetaryComputerStore(ReadOnlyStore):
+    """We write our own Store to access catalogs on Planetary Computer,
+    as it requires some extra arguements to use xr.open_zarr.
+    """
+
+    def __init__(self, data_catalog_id):
+        self.data_catalog_id = data_catalog_id
+
+    def __getitem__(self):
+        import pystac_client
+        import planetary_computer
+
+        catalog = pystac_client.Client.open(
+            "https://planetarycomputer.microsoft.com/api/stac/v1/",
+            modifier=planetary_computer.sign_inplace,
+        )
+        collection = catalog.get_collection(self.data_catalog_id)
+
+        asset = collection.assets["zarr-abfs"]
+
+        if "xarray:storage_options" in asset.extra_fields:
+            store = {
+                "store": asset.href,
+                "storage_options": asset.extra_fields["xarray:storage_options"],
+                **asset.extra_fields["xarray:open_kwargs"],
+            }
+        else:
+            store = {
+                "filename_or_obj": asset.href,
+                **asset.extra_fields["xarray:open_kwargs"],
+            }
+
+        return store
+
+
 class DebugStore(ReadOnlyStore):
     """A store to debug the zarr loading."""
 
@@ -121,6 +156,9 @@ def name_to_zarr_store(path_or_url):
         if len(bits) == 5 and (bits[1], bits[3], bits[4]) == ("s3", "amazonaws", "com"):
             s3_url = f"s3://{bits[0]}{parsed.path}"
             store = S3Store(s3_url, region=bits[2])
+        elif store.startswith("https://planetarycomputer.microsoft.com/"):
+            data_catalog_id = store.rsplit('/', 1)[-1]
+            store = PlanetaryComputerStore(data_catalog_id).__getitem__()
         else:
             store = HTTPStore(store)
 

tests/xarray/test_zarr.py

@@ -12,6 +12,7 @@
 
 from anemoi.datasets.create.functions.sources.xarray import XarrayFieldList
 from anemoi.datasets.testing import assert_field_list
+from anemoi.datasets.data.stores import name_to_zarr_store
 
 
 def test_arco_era5_1():
@@ -91,6 +92,56 @@ def test_inca_one_date():
     print(fs[0].datetime())
 
 
+def test_noaa_replay():
+    ds = xr.open_zarr(
+        "gs://noaa-ufs-gefsv13replay/ufs-hr1/1.00-degree/03h-freq/zarr/fv3.zarr",
+        storage_options={"token": "anon"},
+    )
+
+    flavour = {
+        "rules": {
+            "latitude": {"name": "grid_yt"},
+            "longitude": {"name": "grid_xt"},
+            "time": {"name": "time"},
+            "level": {"name": "pfull"},
+        },
+        "levtype": "pl",
+    }
+
+    fs = XarrayFieldList.from_xarray(ds, flavour)
+
+    assert_field_list(
+        fs,
+        36956954,
+        "1993-12-31T18:00:00",
+        "1999-06-13T03:00:00",
+    )
+
+
+def test_planetary_computer_conus404():
+    url = "https://planetarycomputer.microsoft.com/api/stac/v1/collections/conus404"
+    ds = xr.open_zarr(**name_to_zarr_store(url))
+
+    flavour = {
+        "rules": {
+            "latitude": {"name": "lat"},
+            "longitude": {"name": "lon"},
+            "x": {"name": "west_east"},
+            "y": {"name": "south_north"},
+            "time": {"name": "time"},
+        },
+    }
+
+    fs = XarrayFieldList.from_xarray(ds, flavour)
+
+    assert_field_list(
+        fs,
+        74634912,
+        "1979-10-01T00:00:00",
+        "2022-09-30T23:00:00",
+    )
+
+
 if __name__ == "__main__":
     # test_arco_era5_2()
     # exit()