anemoi-datasets 0.4.5__py3-none-any.whl → 0.5.5__py3-none-any.whl

anemoi/datasets/create/functions/sources/grib.py

@@ -11,9 +11,87 @@
 import glob
 
 from earthkit.data import from_source
+from earthkit.data.indexing.fieldlist import FieldArray
 from earthkit.data.utils.patterns import Pattern
 
 
+def _load(context, name, record):
+    ds = None
+
+    param = record["param"]
+
+    if "path" in record:
+        context.info(f"Using {name} from {record['path']} (param={param})")
+        ds = from_source("file", record["path"])
+
+    if "url" in record:
+        context.info(f"Using {name} from {record['url']} (param={param})")
+        ds = from_source("url", record["url"])
+
+    ds = ds.sel(param=param)
+
+    assert len(ds) == 1, f"{name} {param}, expected one field, got {len(ds)}"
+    ds = ds[0]
+
+    return ds.to_numpy(flatten=True), ds.metadata("uuidOfHGrid")
+
+
+class Geography:
+    """This class retrieve the latitudes and longitudes of unstructured grids,
+    and checks if the fields are compatible with the grid.
+    """
+
+    def __init__(self, context, latitudes, longitudes):
+
+        latitudes, uuidOfHGrid_lat = _load(context, "latitudes", latitudes)
+        longitudes, uuidOfHGrid_lon = _load(context, "longitudes", longitudes)
+
+        assert (
+            uuidOfHGrid_lat == uuidOfHGrid_lon
+        ), f"uuidOfHGrid mismatch: lat={uuidOfHGrid_lat} != lon={uuidOfHGrid_lon}"
+
+        context.info(f"Latitudes: {len(latitudes)}, Longitudes: {len(longitudes)}")
+        assert len(latitudes) == len(longitudes)
+
+        self.uuidOfHGrid = uuidOfHGrid_lat
+        self.latitudes = latitudes
+        self.longitudes = longitudes
+        self.first = True
+
+    def check(self, field):
+        if self.first:
+            # We only check the first field, for performance reasons
+            assert (
+                field.metadata("uuidOfHGrid") == self.uuidOfHGrid
+            ), f"uuidOfHGrid mismatch: {field.metadata('uuidOfHGrid')} != {self.uuidOfHGrid}"
+            self.first = False
+
+
+class AddGrid:
+    """An earth-kit.data.Field wrapper that adds grid information."""
+
+    def __init__(self, field, geography):
+        self._field = field
+
+        geography.check(field)
+
+        self._latitudes = geography.latitudes
+        self._longitudes = geography.longitudes
+
+    def __getattr__(self, name):
+        return getattr(self._field, name)
+
+    def __repr__(self) -> str:
+        return repr(self._field)
+
+    def grid_points(self):
+        return self._latitudes, self._longitudes
+
+    @property
+    def resolution(self):
+        return "unknown"
+
+
 def check(ds, paths, **kwargs):
     count = 1
     for k, v in kwargs.items():
@@ -34,9 +112,13 @@ def _expand(paths):
             yield path
 
 
-def execute(context, dates, path, *args, **kwargs):
+def execute(context, dates, path, latitudes=None, longitudes=None, *args, **kwargs):
     given_paths = path if isinstance(path, list) else [path]
 
+    geography = None
+    if latitudes is not None and longitudes is not None:
+        geography = Geography(context, latitudes, longitudes)
+
     ds = from_source("empty")
     dates = [d.isoformat() for d in dates]
 
@@ -53,7 +135,10 @@ def execute(context, dates, path, *args, **kwargs):
             s = s.sel(valid_datetime=dates, **kwargs)
             ds = ds + s
 
-    if kwargs:
+    if kwargs and not context.partial_ok:
         check(ds, given_paths, valid_datetime=dates, **kwargs)
 
+    if geography is not None:
+        ds = FieldArray([AddGrid(_, geography) for _ in ds])
+
     return ds

anemoi/datasets/create/functions/sources/hindcasts.py

@@ -6,7 +6,6 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 #
-import datetime
 import logging
 
 from earthkit.data.core.fieldlist import MultiFieldList
@@ -14,7 +13,6 @@ from earthkit.data.core.fieldlist import MultiFieldList
 from anemoi.datasets.create.functions.sources.mars import mars
 
 LOGGER = logging.getLogger(__name__)
-DEBUG = True
 
 
 def _to_list(x):
@@ -23,91 +21,34 @@ def _to_list(x):
     return [x]
 
 
-class HindcastCompute:
-    def __init__(self, base_times, available_steps, request):
-        self.base_times = base_times
-        self.available_steps = available_steps
-        self.request = request
-
-    def compute_hindcast(self, date):
-        result = []
-        for step in sorted(self.available_steps):  # Use the shortest step
-            start_date = date - datetime.timedelta(hours=step)
-            hours = start_date.hour
-            if hours in self.base_times:
-                r = self.request.copy()
-                r["date"] = start_date
-                r["time"] = f"{start_date.hour:02d}00"
-                r["step"] = step
-                result.append(r)
-
-        if not result:
-            raise ValueError(
-                f"Cannot find data for {self.request} for {date} (base_times={self.base_times}, "
-                f"available_steps={self.available_steps})"
-            )
-
-        if len(result) > 1:
-            raise ValueError(
-                f"Multiple requests for {self.request} for {date} (base_times={self.base_times}, "
-                f"available_steps={self.available_steps})"
-            )
-
-        return result[0]
-
-
-def use_reference_year(reference_year, request):
-    request = request.copy()
-    hdate = request.pop("date")
-
-    if hdate.year >= reference_year:
-        return None, False
+def hindcasts(context, dates, **request):
 
-    try:
-        date = datetime.datetime(reference_year, hdate.month, hdate.day)
-    except ValueError:
-        if hdate.month == 2 and hdate.day == 29:
-            return None, False
-        raise
+    from anemoi.datasets.dates import HindcastsDates
 
-    request.update(date=date.strftime("%Y-%m-%d"), hdate=hdate.strftime("%Y-%m-%d"))
-    return request, True
+    provider = context.dates_provider
+    assert isinstance(provider, HindcastsDates)
 
+    context.trace("H️", f"hindcasts {len(dates)=}")
 
-def hindcasts(context, dates, **request):
     request["param"] = _to_list(request["param"])
-    request["step"] = _to_list(request["step"])
+    request["step"] = _to_list(request.get("step", 0))
     request["step"] = [int(_) for _ in request["step"]]
 
-    if request.get("stream") == "enfh" and "base_times" not in request:
-        request["base_times"] = [0]
-
-    available_steps = request.pop("step")
-    available_steps = _to_list(available_steps)
-
-    base_times = request.pop("base_times")
-
-    reference_year = request.pop("reference_year")
+    context.trace("H️", f"hindcast {request}")
 
-    context.trace("H️", f"hindcast {request} {base_times} {available_steps} {reference_year}")
-
-    c = HindcastCompute(base_times, available_steps, request)
     requests = []
     for d in dates:
-        req = c.compute_hindcast(d)
-        req, ok = use_reference_year(reference_year, req)
-        if ok:
-            requests.append(req)
-
-    # print("HINDCASTS requests", reference_year, base_times, available_steps)
-    # print("HINDCASTS dates", compress_dates(dates))
+        r = request.copy()
+        hindcast = provider.mapping[d]
+        r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
+        r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
+        r["time"] = hindcast.refdate.strftime("%H")
+        r["step"] = hindcast.step
+        requests.append(r)
 
     if len(requests) == 0:
-        # print("HINDCASTS no requests")
         return MultiFieldList([])
 
-    # print("HINDCASTS requests", requests)
-
     return mars(
         context,
         dates,

anemoi/datasets/create/functions/sources/mars.py

@@ -203,16 +203,22 @@ def mars(context, dates, *requests, request_already_using_valid_datetime=False,
         request_already_using_valid_datetime=request_already_using_valid_datetime,
         date_key=date_key,
     )
+
+    requests = list(requests)
+
     ds = from_source("empty")
+    context.trace("✅", f"{[str(d) for d in dates]}")
+    context.trace("✅", f"Will run {len(requests)} requests")
+    for r in requests:
+        r = {k: v for k, v in r.items() if v != ("-",)}
+        context.trace("✅", f"mars {r}")
+
     for r in requests:
         r = {k: v for k, v in r.items() if v != ("-",)}
 
         if context.use_grib_paramid and "param" in r:
             r = use_grib_paramid(r)
 
-        if DEBUG:
-            context.trace("✅", f"from_source(mars, {r}")
-
         for k, v in r.items():
             if k not in MARS_KEYS:
                 raise ValueError(

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

@@ -12,6 +12,7 @@ import logging
 from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.data.stores import name_to_zarr_store
+from anemoi.datasets.utils.fields import NewMetadataField as NewMetadataField
 
 from .. import iterate_patterns
 from .fieldlist import XarrayFieldList
@@ -49,7 +50,11 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
         data = xr.open_dataset(dataset, **options)
 
     fs = XarrayFieldList.from_xarray(data, flavour)
-    result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
+
+    if len(dates) == 0:
+        return fs.sel(**kwargs)
+    else:
+        result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 
     if len(result) == 0:
         LOG.warning(f"No data found for {dataset} and dates {dates} and {kwargs}")

anemoi/datasets/create/functions/sources/xarray/coordinates.py

@@ -56,6 +56,8 @@ class Coordinate:
     is_step = False
     is_date = False
     is_member = False
+    is_x = False
+    is_y = False
 
     def __init__(self, variable):
         self.variable = variable
@@ -66,10 +68,11 @@ class Coordinate:
         return 1 if self.scalar else len(self.variable)
 
     def __repr__(self):
-        return "%s[name=%s,values=%s]" % (
+        return "%s[name=%s,values=%s,shape=%s]" % (
             self.__class__.__name__,
             self.variable.name,
             self.variable.values if self.scalar else len(self),
+            self.variable.shape,
         )
 
     def reduced(self, i):
@@ -225,11 +228,13 @@ class LatitudeCoordinate(Coordinate):
 
 class XCoordinate(Coordinate):
     is_grid = True
+    is_x = True
     mars_names = ("x",)
 
 
 class YCoordinate(Coordinate):
     is_grid = True
+    is_y = True
     mars_names = ("y",)
 
 

anemoi/datasets/create/functions/sources/xarray/field.py

@@ -7,6 +7,7 @@
 # nor does it submit to any jurisdiction.
 #
 
+import datetime
 import logging
 
 from earthkit.data.core.fieldlist import Field
@@ -71,13 +72,18 @@ class XArrayField(Field):
     def shape(self):
         return self._shape
 
-    def to_numpy(self, flatten=False, dtype=None):
-        values = self.selection.values
+    def to_numpy(self, flatten=False, dtype=None, index=None):
+        if index is not None:
+            values = self.selection[index]
+        else:
+            values = self.selection
 
         assert dtype is None
+
         if flatten:
-            return values.flatten()
-        return values.reshape(self.shape)
+            return values.values.flatten()
+
+        return values  # .reshape(self.shape)
 
     def _make_metadata(self):
         return XArrayMetadata(self)
@@ -103,7 +109,16 @@ class XArrayField(Field):
 
     @property
     def forecast_reference_time(self):
-        return self.owner.forecast_reference_time
+        date, time = self.metadata("date", "time")
+        assert len(time) == 4, time
+        assert len(date) == 8, date
+        yyyymmdd = int(date)
+        time = int(time) // 100
+        return datetime.datetime(yyyymmdd // 10000, yyyymmdd // 100 % 100, yyyymmdd % 100, time)
 
     def __repr__(self):
         return repr(self._metadata)
+
+    def _values(self):
+        # we don't use .values as this will download the data
+        return self.selection

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

@@ -70,10 +70,10 @@ class XarrayFieldList(FieldList):
                 skip.update(attr_val.split(" "))
 
         for name in ds.data_vars:
-            v = ds[name]
-            _skip_attr(v, "coordinates")
-            _skip_attr(v, "bounds")
-            _skip_attr(v, "grid_mapping")
+            variable = ds[name]
+            _skip_attr(variable, "coordinates")
+            _skip_attr(variable, "bounds")
+            _skip_attr(variable, "grid_mapping")
 
         # Select only geographical variables
         for name in ds.data_vars:
@@ -81,14 +81,14 @@ class XarrayFieldList(FieldList):
             if name in skip:
                 continue
 
-            v = ds[name]
+            variable = ds[name]
             coordinates = []
 
-            for coord in v.coords:
+            for coord in variable.coords:
 
                 c = guess.guess(ds[coord], coord)
                 assert c, f"Could not guess coordinate for {coord}"
-                if coord not in v.dims:
+                if coord not in variable.dims:
                     c.is_dim = False
                 coordinates.append(c)
 
@@ -98,17 +98,17 @@ class XarrayFieldList(FieldList):
             if grid_coords < 2:
                 continue
 
-            variables.append(
-                Variable(
-                    ds=ds,
-                    var=v,
-                    coordinates=coordinates,
-                    grid=guess.grid(coordinates),
-                    time=Time.from_coordinates(coordinates),
-                    metadata={},
-                )
+            v = Variable(
+                ds=ds,
+                variable=variable,
+                coordinates=coordinates,
+                grid=guess.grid(coordinates, variable),
+                time=Time.from_coordinates(coordinates),
+                metadata={},
             )
 
+            variables.append(v)
+
         return cls(ds, variables)
 
     def sel(self, **kwargs):

anemoi/datasets/create/functions/sources/xarray/flavour.py

@@ -8,6 +8,8 @@
 #
 
 
+import logging
+
 from .coordinates import DateCoordinate
 from .coordinates import EnsembleCoordinate
 from .coordinates import LatitudeCoordinate
@@ -18,8 +20,13 @@ from .coordinates import StepCoordinate
 from .coordinates import TimeCoordinate
 from .coordinates import XCoordinate
 from .coordinates import YCoordinate
+from .coordinates import is_scalar
 from .grid import MeshedGrid
+from .grid import MeshProjectionGrid
 from .grid import UnstructuredGrid
+from .grid import UnstructuredProjectionGrid
+
+LOG = logging.getLogger(__name__)
 
 
 class CoordinateGuesser:
@@ -155,31 +162,138 @@ class CoordinateGuesser:
             f" {long_name=}, {standard_name=}, units\n\n{c}\n\n{type(c.values)} {c.shape}"
         )
 
-    def grid(self, coordinates):
+    def grid(self, coordinates, variable):
         lat = [c for c in coordinates if c.is_lat]
         lon = [c for c in coordinates if c.is_lon]
 
-        if len(lat) != 1:
-            raise NotImplementedError(f"Expected 1 latitude coordinate, got {len(lat)}")
+        if len(lat) == 1 and len(lon) == 1:
+            return self._lat_lon_provided(lat, lon, variable)
+
+        x = [c for c in coordinates if c.is_x]
+        y = [c for c in coordinates if c.is_y]
+
+        if len(x) == 1 and len(y) == 1:
+            return self._x_y_provided(x, y, variable)
+
+        raise NotImplementedError(f"Cannot establish grid {coordinates}")
+
+    def _check_dims(self, variable, x_or_lon, y_or_lat):
+
+        x_dims = set(x_or_lon.variable.dims)
+        y_dims = set(y_or_lat.variable.dims)
+        variable_dims = set(variable.dims)
 
-        if len(lon) != 1:
-            raise NotImplementedError(f"Expected 1 longitude coordinate, got {len(lon)}")
+        if not (x_dims <= variable_dims) or not (y_dims <= variable_dims):
+            raise ValueError(
+                f"Dimensions do not match {variable.name}{variable.dims} !="
+                f" {x_or_lon.name}{x_or_lon.variable.dims} and {y_or_lat.name}{y_or_lat.variable.dims}"
+            )
+
+        variable_dims = tuple(v for v in variable.dims if v in (x_dims | y_dims))
+        if x_dims == y_dims:
+            # It's unstructured
+            return variable_dims, True
 
+        if len(x_dims) == 1 and len(y_dims) == 1:
+            # It's a mesh
+            return variable_dims, False
+
+        raise ValueError(
+            f"Cannot establish grid for {variable.name}{variable.dims},"
+            f" {x_or_lon.name}{x_or_lon.variable.dims},"
+            f" {y_or_lat.name}{y_or_lat.variable.dims}"
+        )
+
+    def _lat_lon_provided(self, lat, lon, variable):
         lat = lat[0]
         lon = lon[0]
 
-        if (lat.name, lon.name) in self._cache:
-            return self._cache[(lat.name, lon.name)]
+        dim_vars, unstructured = self._check_dims(variable, lon, lat)
+
+        if (lat.name, lon.name, dim_vars) in self._cache:
+            return self._cache[(lat.name, lon.name, dim_vars)]
 
-        assert len(lat.variable.shape) == len(lon.variable.shape), (lat.variable.shape, lon.variable.shape)
-        if len(lat.variable.shape) == 1:
-            grid = MeshedGrid(lat, lon)
+        if unstructured:
+            grid = UnstructuredGrid(lat, lon, dim_vars)
         else:
-            grid = UnstructuredGrid(lat, lon)
+            grid = MeshedGrid(lat, lon, dim_vars)
 
-        self._cache[(lat.name, lon.name)] = grid
+        self._cache[(lat.name, lon.name, dim_vars)] = grid
         return grid
 
+    def _x_y_provided(self, x, y, variable):
+        x = x[0]
+        y = y[0]
+
+        _, unstructured = self._check_dims(variable, x, y)
+
+        if x.variable.dims != y.variable.dims:
+            raise ValueError(f"Dimensions do not match {x.name}{x.variable.dims} != {y.name}{y.variable.dims}")
+
+        if (x.name, y.name) in self._cache:
+            return self._cache[(x.name, y.name)]
+
+        if (x.name, y.name) in self._cache:
+            return self._cache[(x.name, y.name)]
+
+        assert len(x.variable.shape) == len(y.variable.shape), (x.variable.shape, y.variable.shape)
+
+        grid_mapping = variable.attrs.get("grid_mapping", None)
+
+        if grid_mapping is None:
+            LOG.warning(f"No 'grid_mapping' attribute provided for '{variable.name}'")
+            LOG.warning("Trying to guess...")
+
+            PROBE = {
+                "prime_meridian_name",
+                "reference_ellipsoid_name",
+                "crs_wkt",
+                "horizontal_datum_name",
+                "semi_major_axis",
+                "spatial_ref",
+                "inverse_flattening",
+                "semi_minor_axis",
+                "geographic_crs_name",
+                "GeoTransform",
+                "grid_mapping_name",
+                "longitude_of_prime_meridian",
+            }
+            candidate = None
+            for v in self.ds.variables:
+                var = self.ds[v]
+                if not is_scalar(var):
+                    continue
+
+                if PROBE.intersection(var.attrs.keys()):
+                    if candidate:
+                        raise ValueError(f"Multiple candidates for 'grid_mapping': {candidate} and {v}")
+                    candidate = v
+
+            if candidate:
+                LOG.warning(f"Using '{candidate}' as 'grid_mapping'")
+                grid_mapping = candidate
+            else:
+                LOG.warning("Could not fine a candidate for 'grid_mapping'")
+
+        if grid_mapping is None:
+            if "crs" in self.ds[variable].attrs:
+                grid_mapping = self.ds[variable].attrs["crs"]
+                LOG.warning(f"Using CRS {grid_mapping} from variable '{variable.name}' attributes")
+
+        if grid_mapping is None:
+            if "crs" in self.ds.attrs:
+                grid_mapping = self.ds.attrs["crs"]
+                LOG.warning(f"Using CRS {grid_mapping} from global attributes")
+
+        if grid_mapping is not None:
+            if unstructured:
+                return UnstructuredProjectionGrid(x, y, grid_mapping)
+            else:
+                return MeshProjectionGrid(x, y, grid_mapping)
+
+        LOG.error("Could not fine a candidate for 'grid_mapping'")
+        raise NotImplementedError(f"Unstructured grid {x.name} {y.name}")
+
 
 class DefaultCoordinateGuesser(CoordinateGuesser):
     def __init__(self, ds):