Filtering when reading a Grid file, avoid time leaking into the Grid

[dylannelson]

Closes #1444 (and others (?) planned to update)

Overview

1. time was shown to exist in a grid from a user reading in a grid and data from a single file
2. Later when the user tried to use .subset there was a crash related to time being different on the grid vs the data
3. time shouldn't realistically be on the grid to start with, so this was a concern, and was causing downstream issues not addressed when reading in the data
4. There was 2 main ideas of where to place filtering to avoid this, Grid.__init__ or _read_ugrid, both were tested extensively
   • Grid.__init__ is run in multiple scenarios, including when a file is read from disk, but in other scenarios like _slice_from_grid which can occur when subsetting
   • the path to _read_ugrid occurs primarily when a file is read from disk, like:
     • ux.open_grid(..) -> Grid.from_dataset(..) -> _read_ugrid(ds)
     • _read_ugrid(ds) main function is: "Parses an unstructured grid dataset and encodes it in the UGRID" so this seems like the right place for cleaning
5. This has been tested in multiple ways and is shown to:
   • Scalar time gone from grid coords and grid variables
   • node_lon, node_lat, face_node_connectivity preserved
   • Grid non-empty (n_node, n_face > 0)
   • bounding_box runs, no error
   • Data time preserved
   • Subset returned faces
   • Sliced grid keeps subgrid_face_indices
   • sel(time=...) still works
   • Dimensional time axis + coord stripped
   • Leaked data variable stripped

Expected Usage

Changes are all located in uxarray\uxarray\io\_ugrid.py

Test these changes on an old version and the new version. This test likely fails in old uxarray and succeeds in this branch

import pandas as pd, xarray as xr, uxarray as ux

# This file exists in uxarray repo, update with your local path added if needed
path_to_nc = "uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc" 

# Opens a grid, adds scalar time, saves to local
(xr.open_dataset(path_to_nc)
   .assign_coords(time=pd.Timestamp("2016-10-01"))
   .to_netcdf("grid_with_time.nc"))

uxgrid = ux.open_grid("grid_with_time.nc")

print("grid coords:", list(uxgrid._ds.coords))
print("time removed?", "time" not in uxgrid._ds.coords)   # True = fix working

PR Checklist

General

• An issue is created and linked
• Added appropriate labels (if your uxarray repo permissions allow it)
• Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

• Adequate tests are created if there is new functionality
• Tests are not too basic (such as simply calling a function and nothing else)
• Tests cover all major paths in your new functions
• If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation

• Docstrings have been added to all new functions
• Docstrings have been updated with any function changes
• User (public) functions have been added to docs/api.rst
• Internal (private) function names start with an underscore (_)

Examples

• All notebook examples cleared the output of all cells before committing
• New notebook examples added to appropriate folder (gallery: docs/examples/; guide: docs/user-guide/; quickstart: docs/getting-started/)
• New notebook examples referenced in appropriate .rst file (gallery: docs/gallery.rst; guide: docs/userguide.rst; quickstart: docs/quickstart.rst)
• New notebook gallery examples added entry in docs/gallery.yml with appropriate thumbnail photo in docs/_static/thumbnails/

AI Disclosure

AI Usage:

• I take responsibility for all AI-generated content in my PR.
• I have tested all AI-generated content in my PR.

[erogluorhan]

Maybe too early for a review, but I'd like to share some thoughts:

1. time was shown to exist in a grid from a user reading in a grid and data from a single file

   A little more clarity could be helpful here. IIRC, time is not shown directly in the Grid object but in the Grid._ds.

2. That said, I'd like to know what actual purpose we had with that object attribute, i.e. _ds. If it was to keep complete track of the original file content through xarray.Dataset, I'd say this dropping at the level of _ds might break that.

3. Maybe we'd want to fix it at the Grid itself rather than altering _ds, but would it be possible?

4. Finally, ugrid (i.e. _ugrid.py) is only one of the several formats we support. What about other formats in our I/O, e.g. _mpas.py etc. Their grid files can come in a similar way. Maybe we will need to do this fix at the Grid level rather than particular I/O modules?

[dylannelson]

#1444 — Solution history

This is a bit about each version of the idea, why it's located where it is, and what the final code looks like.

---

Version 1 — .values guard in _slice_from_grid

# uxarray/core/dataarray.py — _slice_from_grid()
n_face=sliced_grid._ds["subgrid_face_indices"].values   # was: [...["subgrid_face_indices"]]

• Location: within _slice_from_grid in core/dataarray.py.
• Idea: Pass the face indexer as .values (a plain array) so xarray does positional indexing with no coordinate alignment. It fixes the crash for every subset method, but only guards the symptom — the stray time stays on the grid, so it never addresses why the grid carries a time.

---

Version 2 — Name-based clean build in _read_ugrid

# uxarray/io/_ugrid.py — end of _read_ugrid()
ds = _keep_only_grid_vars(ds)   # keep only recognized grid vars/coords, drop the rest

• Location: at the end of _read_ugrid in io/_ugrid.py.
• Idea: Filter the parsed dataset down to recognized grid variables, dropping extras where the grid is built. This removes the stray time at its source (the leaky reader), so nothing downstream can collide. It's essentially the final answer — a simpler-looking alternative was proposed next, which prompted the comparison below.

---

Version 3 — drop_dims in Grid.__init__ (team's idea)

# uxarray/grid/grid.py — Grid.__init__
extra_dims = set(grid_ds.dims) - set(DIM_NAMES)
grid_ds = grid_ds.drop_dims(extra_dims)

• Location: at the top of Grid.__init__ in grid/grid.py.
• Idea: Drop any dimension not in DIM_NAMES, centrally for every reader, to remove stray non-grid data by axis. It captured a dimensional time, but the leaked time here is a scalar (0-d) coord with no dimension — so drop_dims finds nothing to remove and the bug survives.

---

Version 4 — "Best of both" (drop_dims + 0-d/name drop)

grid_ds = grid_ds.drop_dims(set(grid_ds.dims) - set(DIM_NAMES))              # dimensional extras
drop = [n for n, v in grid_ds.variables.items() if n not in KEEP and v.ndim == 0]
grid_ds = grid_ds.drop_vars(drop, errors="ignore")                          # scalar extras

• Location: at the top of Grid.__init__ in grid/grid.py.
• Idea: Keep V3's drop_dims and add a second pass dropping 0-d non-grid vars, catching the scalar time it missed. It works, but uses two mechanisms for what one can do — V5's name-based keep-list covers scalar, dimensional, and data-var leaks in a single pass.

---

Version 5 — Single name-based keep-list in _read_ugrid (final)

# uxarray/io/_ugrid.py — end of _read_ugrid()
def _keep_only_grid_vars(ds):
    keep = {"grid_topology", *SPHERICAL_COORD_NAMES, *CARTESIAN_COORD_NAMES,
            *CONNECTIVITY_NAMES, *DESCRIPTOR_NAMES}
    return ds.drop_vars([n for n in ds.variables if n not in keep], errors="ignore")

• Location: at the end of _read_ugrid in io/_ugrid.py — different location (see next section).
• Idea: One name-based keep-list that drops every kind of non-grid extra (scalar, dimensional, data vars) in a single pass, collapsing V2–V4 into the simplest complete fix. Could be paired with the .values option, but didn't seem necessary.

---

Why V5 lives in _read_ugrid, not Grid.__init__

Two code paths build a Grid. They share Grid.__init__ but not _read_ugrid, so why is it in _read_ugrid?

Path A — opening a grid from a file (where the time leaks in):

ux.open_grid(...)                   core/api.py
  └─ _open_dataset_with_fallback()  # full xr.Dataset — stray `time` included
  └─ Grid.from_dataset(...)         grid.py
       └─ no source_grid_spec
       └─ _parse_grid_type()
       └─ _read_ugrid(...)          io/_ugrid.py   ← renames in place, returns EVERYTHING (leak enters here)
  └─ Grid.__init__(...)             grid.py   # self._ds = grid_ds

Path B — building a sliced grid during a subset (trying not to disturb this path):

uxda.subset.bounding_box(...)
  └─ Grid.isel(...)                    grid.py
       └─ _slice_face_indices(...)     grid/slice.py
            ds = grid._ds.isel(...)    # copy of cleaned grid
            ds["subgrid_face_indices"] = ...
       └─ Grid.from_dataset(...)       slice.py
            └─ grid_ds = ds            (_read_ugrid is skipped)
  └─ Grid.__init__(...)                grid.py   # self._ds holds subgrid_*_indices

	Through _read_ugrid?	Through Grid.__init__?	_ds holds subgrid_*_indices?
Path A (file open)	Yes	Yes	No
Path B (subset)	No	Yes	Yes (needed by _slice_from_grid)

• Filter in _read_ugrid (V5): Path A only. time is dropped as the grid is read; Path B never enters _read_ugrid, so its subgrid_*_indices survive, and its sliced grid (copied from the already-clean grid._ds) comes out clean for free.
• Filter in Grid.__init__: both paths. It would strip subgrid_face_indices from Path B's sliced grid → _slice_from_grid hits a KeyError. Fixes the leak but breaks subsetting.

---

Final code preview

Basically came down to two edits, both in uxarray/io/_ugrid.py.

Edit 1 — add one import at the top of the file (ugrid is already imported; only DESCRIPTOR_NAMES is new):

import numpy as np
import xarray as xr

import uxarray.conventions.ugrid as ugrid
from uxarray.constants import INT_DTYPE, INT_FILL_VALUE
from uxarray.grid.connectivity import _replace_fill_values
from uxarray.conventions.descriptors import DESCRIPTOR_NAMES  # <---- this line

Edit 2 — filter at the end of _read_ugrid, then define the helper directly after it:

def _read_ugrid(ds):
    """Parses an unstructured grid dataset and encodes it in the UGRID
    conventions."""

    # ... (topology parse, coord/connectivity renames, dim swaps — unchanged) ...

    dim_dict[ds["face_node_connectivity"].dims[1]] = ugrid.N_MAX_FACE_NODES_DIM

    ds = ds.swap_dims(dim_dict)

    # ===== NEW =====
    # Strip non-grid extras (e.g. a stray scalar `time` coordinate, or unrelated data variables)
    ds = _keep_only_grid_vars(ds)
    # ===== end NEW =====

    return ds, dim_dict


# ===== NEW helper — place directly after _read_ugrid =====
def _keep_only_grid_vars(ds):
    """Return ``ds`` with only recognized UGRID grid variables/coordinates.

    Anything else on the dataset (a stray scalar ``time`` coordinate, unrelated
    data variables, etc.) is dropped so it cannot leak onto ``grid._ds``.

    Runs on the file-read path only.
    """
    # uxarray's own canonical grid-variable names (the same lists Grid filters against)
    keep = {"grid_topology"}
    keep.update(ugrid.SPHERICAL_COORD_NAMES)   # node/edge/face lon-lat
    keep.update(ugrid.CARTESIAN_COORD_NAMES)   # node/edge/face x-y-z
    keep.update(ugrid.CONNECTIVITY_NAMES)      # face_node_connectivity, edge_node_connectivity, ...
    keep.update(DESCRIPTOR_NAMES)              # n_nodes_per_face, face_areas, boundary_*_indices, ...

    # drop_vars removes variables/coords by name (never bare dimensions), so grid
    # dims survive with their variables; errors="ignore" tolerates absent names.
    drop = [name for name in ds.variables if name not in keep]
    return ds.drop_vars(drop, errors="ignore")
# ===== end NEW helper =====


# ===== existing code =====
def _encode_ugrid(ds):
    """Encodes an unstructured grid represented under a ``Grid`` object as a
    ``xr.Dataset`` with an updated grid topology variable."""

    if "grid_topology" in ds:
        ds = ds.drop_vars(["grid_topology"])
    # ... (rest of _encode_ugrid, unchanged) ...

[dylannelson]

Maybe too early for a review, but I'd like to share some thoughts:

1. > `time` was shown to exist in a grid from a user reading in a grid and data from a single file
   
   
   A little more clarity could be helpful here. IIRC, `time` is not shown directly in the `Grid` object but in the `Grid._ds`.

2. That said, I'd like to know what actual purpose we had with that object attribute, i.e. `_ds`. If it was to keep complete track of the original file content through `xarray.Dataset`, I'd say this dropping at the level of `_ds` might break that.

3. Maybe we'd want to fix it at the `Grid` itself rather than altering `_ds`, but would it be possible?

4. Finally, ugrid (i.e. `_ugrid.py`) is only one of the several formats we support. What about other formats in our I/O, e.g. _mpas.py etc. Their grid files can come in a similar way. Maybe we will need to do this fix at the Grid level rather than particular I/O modules?

Sorry saw this after I was typing all the above content. I'll think about this and we can discuss tomorrow?