Add new Grid class for handling multiple grids without preferencing polygons

Author: mx-moth

docs/api/conventions/arakawa-c.rst

@@ -14,14 +14,11 @@ Arakawa C grid datasets have four :ref:`grids <grids>`:
 *face*, *left*, *back*, and *node*.
 :class:`ArakawaCGridKind` represents this.
 Each grid is :ref:`indexed <indexing>` by a grid kind and two integers *i* and *j*.
-The convention native index type is :data:`ArakawaCIndex`.
 
 .. autoclass:: ArakawaCGridKind
    :members:
    :undoc-members:
 
-.. autodata:: ArakawaCIndex
-
 Topology
 ========
 

docs/api/conventions/grid.rst

@@ -17,14 +17,11 @@ Indexing
 CF grid datasets have one :ref:`grid <grids>`: *face*.
 :class:`CFGridKind` represents this.
 Each face is :ref:`indexed <indexing>` by two integers *x* and *y*.
-The convention native index type is :data:`CFGridIndex`.
 
 .. autoclass:: CFGridKind
    :members:
    :undoc-members:
 
-.. autodata:: CFGridIndex
-
 Topology
 ========
 

docs/api/conventions/interface.rst

@@ -21,8 +21,12 @@ the :class:`~emsarray.conventions.Convention` interface.
 .. autoclass:: emsarray.conventions.DimensionConvention
 
    .. autoattribute:: grid_dimensions
-   .. automethod:: unpack_index
-   .. automethod:: pack_index
+
+.. autoclass:: emsarray.conventions.Grid
+   :members:
+
+.. autoclass:: emsarray.conventions.DimensionGrid
+   :members: dimensions
 
 .. type:: GridKind
 

docs/api/conventions/ugrid.rst

@@ -13,14 +13,11 @@ Indexing
 UGRID datasets have three :ref:`grids <grids>`: *face*, *edge* and *node*.
 :class:`UGridKind` represents this.
 Each grid is :ref:`indexed <indexing>` by a single integer.
-The convention native index type is :data:`UGridIndex`.
 
 .. autoclass:: UGridKind
    :members:
    :undoc-members:
 
-.. autodata:: UGridIndex
-
 Topology
 ========
 

src/emsarray/cli/commands/export_geometry.py

@@ -3,17 +3,14 @@
 from collections.abc import Callable
 from pathlib import Path
 
-import xarray
-
 import emsarray
 from emsarray.cli import BaseCommand, CommandException
 from emsarray.operations import geometry
-from emsarray.types import Pathish
 
 logger = logging.getLogger(__name__)
 
-Writer = Callable[[xarray.Dataset, Pathish], None]
-format_writers: dict[str, Writer] = {
+
+format_writers: dict[str, Callable] = {
     'geojson': geometry.write_geojson,
     'shapefile': geometry.write_shapefile,
     'wkt': geometry.write_wkt,
@@ -46,6 +43,14 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
                 "The output format will be guessed using the output extension by default"
             ))
 
+        parser.add_argument(
+            "-g", "--grid-kind", type=str,
+            default=None,
+            help=(
+                "Grid kind to export. Will export the default grid if not specified."
+            ),
+        )
+
     def guess_format(self, output_path: Path) -> str:
         extension = output_path.suffix
         if extension in {'.json', '.geojson'}:
@@ -69,13 +74,31 @@ def handle(self, options: argparse.Namespace) -> None:
             output_format = self.guess_format(output_path)
             logger.debug("Guessed output format as %r", output_format)
 
-        count = dataset.ems.polygons[dataset.ems.mask].size
-        logger.debug("Dataset contains %d polygons", count)
+        if options.grid_kind is None:
+            grid_kind = dataset.ems.default_grid_kind
+        else:
+            grid_kind_names = {str(grid_kind): grid_kind for grid_kind in dataset.ems.grid_kinds}
+            try:
+                grid_kind = grid_kind_names[options.grid_kind]
+            except KeyError:
+                grid_kind_choices = ", ".join(grid_kind_names.keys())
+                raise CommandException(
+                    f"Unknown grid kind {options.grid_kind!r}. "
+                    f"Valid choices are: {grid_kind_choices}."
+                )
 
         try:
             writer = format_writers[output_format]
         except KeyError:
             raise CommandException(f"Unknown output format {output_format!r}")
 
+        grid = dataset.ems.grids[grid_kind]
+        geometries = grid.geometry[grid.mask]
+        count = geometries.size
+
+        logger.debug("Grid kind: %s", grid_kind)
+        logger.debug("Geometry type: %s", grid.geometry_type.__name__)
+        logger.debug("Geometry count: %d", count)
         logger.debug("Exporting geometry as %r", output_format)
-        writer(dataset, output_path)
+
+        writer(dataset, output_path, grid_kind=grid_kind)

src/emsarray/conventions/__init__.py

@@ -13,7 +13,9 @@
 Convention instances can be instantiated directly in this case.
 Refer to each Convention implementation for details.
 """
-from ._base import Convention, DimensionConvention, Specificity
+from ._base import (
+    Convention, DimensionConvention, DimensionGrid, Grid, Specificity
+)
 from ._registry import get_dataset_convention, register_convention
 from ._utils import open_dataset
 from .arakawa_c import ArakawaC
@@ -23,6 +25,7 @@
 
 __all__ = [
     "Convention", "DimensionConvention",
+    "Grid", "DimensionGrid",
     "Specificity",
     "get_dataset_convention", "register_convention",
     "open_dataset",