anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__py3-none-any.whl

anemoi/datasets/grids.py

@@ -8,15 +8,11 @@
 # nor does it submit to any jurisdiction.
 
 
-import base64
 import logging
 from typing import Any
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 import numpy as np
+from anemoi.utils.grids import latlon_to_xyz
 from numpy.typing import NDArray
 
 LOG = logging.getLogger(__name__)
@@ -92,71 +88,6 @@ def plot_mask(
         plt.savefig(path + "-global-zoomed.png")
 
 
-# TODO: Use the one from anemoi.utils.grids instead
-# from anemoi.utils.grids import ...
-def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> Tuple[NDArray[Any], NDArray[Any]]:
-    """Convert Cartesian coordinates to latitude and longitude.
-
-    Parameters
-    ----------
-    x : NDArray[Any]
-        X coordinates.
-    y : NDArray[Any]
-        Y coordinates.
-    z : NDArray[Any]
-        Z coordinates.
-
-    Returns
-    -------
-    Tuple[NDArray[Any], NDArray[Any]]
-        Latitude and longitude coordinates.
-    """
-    return (
-        np.rad2deg(np.arcsin(np.minimum(1.0, np.maximum(-1.0, z)))),
-        np.rad2deg(np.arctan2(y, x)),
-    )
-
-
-# TODO: Use the one from anemoi.utils.grids instead
-# from anemoi.utils.grids import ...
-def latlon_to_xyz(
-    lat: NDArray[Any], lon: NDArray[Any], radius: float = 1.0
-) -> Tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
-    """Convert latitude and longitude to Cartesian coordinates.
-
-    Parameters
-    ----------
-    lat : NDArray[Any]
-        Latitude coordinates.
-    lon : NDArray[Any]
-        Longitude coordinates.
-    radius : float, optional
-        Radius of the sphere. Defaults to 1.0.
-
-    Returns
-    -------
-    Tuple[NDArray[Any], NDArray[Any], NDArray[Any]]
-        X, Y, and Z coordinates.
-    """
-    # https://en.wikipedia.org/wiki/Geographic_coordinate_conversion#From_geodetic_to_ECEF_coordinates
-    # We assume that the Earth is a sphere of radius 1 so N(phi) = 1
-    # We assume h = 0
-    #
-    phi = np.deg2rad(lat)
-    lda = np.deg2rad(lon)
-
-    cos_phi = np.cos(phi)
-    cos_lda = np.cos(lda)
-    sin_phi = np.sin(phi)
-    sin_lda = np.sin(lda)
-
-    x = cos_phi * cos_lda * radius
-    y = cos_phi * sin_lda * radius
-    z = sin_phi * radius
-
-    return x, y, z
-
-
 class Triangle3D:
     """A class to represent a 3D triangle and perform intersection tests with rays."""
 
@@ -272,8 +203,8 @@ def cutout_mask(
     global_lons: NDArray[Any],
     cropping_distance: float = 2.0,
     neighbours: int = 5,
-    min_distance_km: Optional[Union[int, float]] = None,
-    plot: Optional[str] = None,
+    min_distance_km: int | float | None = None,
+    plot: str | None = None,
 ) -> NDArray[Any]:
     """Return a mask for the points in [global_lats, global_lons] that are inside of [lats, lons].
 
@@ -465,7 +396,7 @@ def thinning_mask(
     return np.array([i for i in indices])
 
 
-def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> List[int]:
+def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> list[int]:
     """Find the outline of the grid points.
 
     Parameters
@@ -513,92 +444,6 @@ def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> List
     return outside
 
 
-def deserialise_mask(encoded: str) -> NDArray[Any]:
-    """Deserialise a mask from a base64 encoded string.
-
-    Parameters
-    ----------
-    encoded : str
-        Base64 encoded string.
-
-    Returns
-    -------
-    NDArray[Any]
-        Deserialised mask array.
-    """
-    import pickle
-    import zlib
-
-    packed = pickle.loads(zlib.decompress(base64.b64decode(encoded)))
-
-    mask = []
-    value = False
-    for count in packed:
-        mask.extend([value] * count)
-        value = not value
-    return np.array(mask, dtype=bool)
-
-
-def _serialise_mask(mask: NDArray[Any]) -> str:
-    """Serialise a mask to a base64 encoded string.
-
-    Parameters
-    ----------
-    mask : NDArray[Any]
-        Mask array.
-
-    Returns
-    -------
-    str
-        Base64 encoded string.
-    """
-    import pickle
-    import zlib
-
-    assert len(mask.shape) == 1
-    assert len(mask)
-
-    packed = []
-    last = mask[0]
-    count = 1
-
-    for value in mask[1:]:
-        if value == last:
-            count += 1
-        else:
-            packed.append(count)
-            last = value
-            count = 1
-
-    packed.append(count)
-
-    # We always start with an 'off' value
-    # So if the first value is 'on', we need to add a zero
-    if mask[0]:
-        packed.insert(0, 0)
-
-    return base64.b64encode(zlib.compress(pickle.dumps(packed))).decode("utf-8")
-
-
-def serialise_mask(mask: NDArray[Any]) -> str:
-    """Serialise a mask and ensure it can be deserialised.
-
-    Parameters
-    ----------
-    mask : NDArray[Any]
-        Mask array.
-
-    Returns
-    -------
-    str
-        Base64 encoded string.
-    """
-    result = _serialise_mask(mask)
-    # Make sure we can deserialise it
-    assert np.all(mask == deserialise_mask(result))
-    return result
-
-
 def nearest_grid_points(
     source_latitudes: NDArray[Any],
     source_longitudes: NDArray[Any],
@@ -644,29 +489,3 @@ def nearest_grid_points(
     else:
         distances, indices = cKDTree(source_points).query(target_points, k=k, distance_upper_bound=max_distance)
     return distances, indices
-
-
-if __name__ == "__main__":
-    global_lats, global_lons = np.meshgrid(
-        np.linspace(90, -90, 90),
-        np.linspace(-180, 180, 180),
-    )
-    global_lats = global_lats.flatten()
-    global_lons = global_lons.flatten()
-
-    lats, lons = np.meshgrid(
-        np.linspace(50, 40, 100),
-        np.linspace(-10, 15, 100),
-    )
-    lats = lats.flatten()
-    lons = lons.flatten()
-
-    mask = cutout_mask(lats, lons, global_lats, global_lons, cropping_distance=5.0)
-
-    import matplotlib.pyplot as plt
-
-    fig = plt.figure(figsize=(10, 5))
-    plt.scatter(global_lons, global_lats, s=0.01, marker="o", c="r")
-    plt.scatter(global_lons[mask], global_lats[mask], s=0.1, c="k")
-    # plt.scatter(lons, lats, s=0.01)
-    plt.savefig("cutout.png")

anemoi/datasets/schemas/recipe.json

@@ -0,0 +1,131 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "type": "object",
+    "$id": "https://ecmwf.int/anemoi-datasets-recipe.schema.json",
+    "title": "Product",
+    "description": "Anemoi datasets recipe configuration",
+    "additionalProperties": false,
+    "$defs": {
+        "source-or-filter": {
+            "type": "object",
+            "minProperties": 1,
+            "maxProperties": 1
+        },
+        "pipe": {
+            "type": "array",
+            "items": {
+                "$ref": "#/$defs/input-object"
+            }
+        },
+        "join": {
+            "type": "array",
+            "items": {
+                "$ref": "#/$defs/input-object"
+            }
+        },
+        "concat": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "minProperties": 2,
+                "maxProperties": 2,
+                "required": [
+                    "dates"
+                ]
+            }
+        },
+        "input-object": {
+            "oneOf": [
+                {
+                    "$ref": "#/$defs/pipe"
+                },
+                {
+                    "$ref": "#/$defs/join"
+                },
+                {
+                    "$ref": "#/$defs/concat"
+                },
+                {
+                    "$ref": "#/$defs/source-or-filter"
+                }
+            ]
+        }
+    },
+    "properties": {
+        "env": {
+            "type": "object"
+        },
+        "description": {
+            "type": "string"
+        },
+        "name": {
+            "type": "string"
+        },
+        "licence": {
+            "type": "string"
+        },
+        "attribution": {
+            "type": "string"
+        },
+        "dates": {
+            "type": "object",
+            "required": [
+                "start",
+                "end"
+            ],
+            "properties": {
+                "start": {
+                    "type": "string",
+                    "format": "date"
+                },
+                "end": {
+                    "type": "string",
+                    "format": "date"
+                },
+                "frequency": {
+                    "type": [
+                        "integer",
+                        "string"
+                    ]
+                },
+                "group_by": {
+                    "type": [
+                        "integer",
+                        "string"
+                    ]
+                }
+            }
+        },
+        "input": {
+            "$ref": "#/$defs/input-object"
+        },
+        "data_sources": {
+            "type": "object",
+            "patternProperties": {
+                "^[a-zA-Z_][a-zA-Z0-9_]*$": {
+                    "$ref": "#/$defs/input-object"
+                }
+            },
+            "additionalProperties": false
+        },
+        "output": {
+            "type": "object"
+        },
+        "statistics": {
+            "type": "object"
+        },
+        "build": {
+            "type": "object"
+        },
+        "common": {
+            "type": "object"
+        },
+        "platform": {
+            "type": "object"
+        }
+    },
+    "required": [
+        "dates",
+        "input"
+    ]
+}

anemoi/datasets/testing.py

@@ -12,19 +12,17 @@
 
 import logging
 from typing import Any
-from typing import List
-from typing import Optional
 
 LOG = logging.getLogger(__name__)
 
 
 def assert_field_list(
-    fs: List[Any],
-    size: Optional[int] = None,
-    start: Optional[Any] = None,
-    end: Optional[Any] = None,
+    fs: list[Any],
+    size: int | None = None,
+    start: Any | None = None,
+    end: Any | None = None,
     constant: bool = False,
-    skip: Optional[Any] = None,
+    skip: Any | None = None,
 ) -> None:
     """Asserts various properties of a list of fields.
 
@@ -85,3 +83,91 @@ def assert_field_list(
     assert south >= -90, south
     assert east <= 360, east
     assert west >= -180, west
+
+
+class IndexTester:
+    """Class to test indexing of datasets."""
+
+    def __init__(self, ds: Any) -> None:
+        """Initialise the IndexTester.
+
+        Parameters
+        ----------
+        ds : Any
+            Dataset.
+        """
+        self.ds = ds
+        self.np = ds[:]  # Numpy array
+
+        assert self.ds.shape == self.np.shape, (self.ds.shape, self.np.shape)
+        assert (self.ds == self.np).all()
+
+    def __getitem__(self, index: Any) -> None:
+        """Test indexing.
+
+        Parameters
+        ----------
+        index : Any
+            Index.
+        """
+        LOG.info("IndexTester: %s", index)
+        if self.ds[index] is None:
+            assert False, (self.ds, index)
+
+        if not (self.ds[index] == self.np[index]).all():
+            assert (self.ds[index] == self.np[index]).all()
+
+
+def default_test_indexing(ds):
+
+    t = IndexTester(ds)
+
+    t[0:10, :, 0]
+    t[:, 0:3, 0]
+    # t[:, :, 0]
+    t[0:10, 0:3, 0]
+    t[:, :, :]
+
+    if ds.shape[1] > 2:  # Variable dimension
+        t[:, (1, 2), :]
+        t[:, (1, 2)]
+
+    t[0]
+    t[0, :]
+    t[0, 0, :]
+    t[0, 0, 0, :]
+
+    if ds.shape[2] > 1:  # Ensemble dimension
+        t[0:10, :, (0, 1)]
+
+    for i in range(3):
+        t[i]
+        start = 5 * i
+        end = len(ds) - 5 * i
+        step = len(ds) // 10
+
+        t[start:end:step]
+        t[start:end]
+        t[start:]
+        t[:end]
+        t[::step]
+
+
+class Trace:
+
+    def __init__(self, ds):
+        self.ds = ds
+        self.f = open("trace.txt", "a")
+
+    def __getattr__(self, name: str) -> Any:
+
+        print(name, file=self.f, flush=True)
+        return getattr(self.ds, name)
+
+    def __len__(self) -> int:
+        print("__len__", file=self.f, flush=True)
+        return len(self.ds)
+
+    def __getitem__(self, index: Any) -> Any:
+        print("__getitem__", file=self.f, flush=True)
+        return self.ds[index]