Filtering when reading a Grid file, avoid time leaking into the Grid - #1667
Filtering when reading a Grid file, avoid time leaking into the Grid#1667dylannelson wants to merge 9 commits into
time leaking into the Grid#1667Conversation
Time was getting into grid data, causing issues when subsetting down the line. This should clean the data in a way that shouldn't happen in the future
pre-commit run --files uxarray/io/_ugrid.py Seemed to turn up a few issues with spaces and line lengths
erogluorhan
left a comment
erogluorhan
left a comment
There was a problem hiding this comment.
Maybe too early for a review, but I'd like to share some thoughts:
-
timewas shown to exist in a grid from a user reading in a grid and data from a single fileA little more clarity could be helpful here. IIRC,
timeis not shown directly in theGridobject but in theGrid._ds. -
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 throughxarray.Dataset, I'd say this dropping at the level of_dsmight break that. -
Maybe we'd want to fix it at the
Griditself rather than altering_ds, but would it be possible? -
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?
#1444 — Solution historyThis 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 —
|
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.timeis dropped as the grid is read; Path B never enters_read_ugrid, so itssubgrid_*_indicessurvive, and its sliced grid (copied from the already-cleangrid._ds) comes out clean for free. - Filter in
Grid.__init__: both paths. It would stripsubgrid_face_indicesfrom Path B's sliced grid →_slice_from_gridhits aKeyError. 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 lineEdit 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) ...
Sorry saw this after I was typing all the above content. Edit: Here's an update, more focused on # 4 to start This is a great question and I wish I had realized sooner. Trying to find a solution that was very high level, like in So I had to look at 2 questions:
Looking at # 1, I performed a similar check like before, read in a demo dataset, add time, save, reopen, check if there's time, and see if it crashes when there's a subset. Here's what I got
So basically from what I'm seeing right now, we have 3 different outcomes
I'll keep looking for a solution for # 2, but worried at it's impact, especially since every file type seems to have a unique way to handle the data when opening. |
|
The table in your latest comment showing no leaks for UGRID confused me a bit (From your other words, it sounds like UGRID has the time leaking). Could you clarify? |
|
@erogluorhan "the
|
It's no leak for UGRID in my current version, which fixed it. |
|
and for
that's effectively what I had originally, via The issue with this is that it fixes a symptom, without addressing the cause. We discussed this early on and it was turned down when we originally discussed it. It's basically a patch in the |
You're right; "without altering |
|
Other readers like MPAS, ESMF, Exodus, ICON, and FESOM2 still leave |
|
This new push should have everything we have wanted from the last few discussions extra_dims = set(grid_ds.dims) - set(DIM_NAMES)
grid_ds = grid_ds.drop_dims(extra_dims)The This solution is different in that it is dropping non-grid coordinates by name, which catches the scalar one that we've been having trouble with. Here's what this solution looks like instead: grid_coord_names = set(ugrid.SPHERICAL_COORD_NAMES) | set(ugrid.CARTESIAN_COORD_NAMES)
stray = [c for c in ds.coords if c not in grid_coord_names]
return ds.drop_vars(stray, errors="ignore")TestingBefore: Notebook Preview: Downstream impactPerformed a few downstream tests with claude code to test what affects it may have, here's the results of what it looked at:
|
dylannelson
added
run-benchmark
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
This looks like a relatively clean fix overall, with lots of tests run by hand to confirm things seem to be working!
I left a few minor comments inline, and I have a couple larger comments here:
- I had a suspicion that extra dimensions (not just extra coordinates) on input datasets would be an issue, but I tested a few things locally and found that it wasn't a problem (for some reason…). For example, if you had a time dimension on the input dataset, instead of just a time coordinate, it doesn't seem to lead to a crash later (although, it does leak into
Grid._ds). It isn't really obvious to me why it wouldn't crash…. Can you clarify why this might be the case, and/or add a test to ensure subsetting still works even after adding a time dimension to the grid? - Would you also be able to create a file with a time coordinate somewhere in the test suite (in meshfiles folder), and test that, too, instead of only having a test where the time coordinate gets assigned after loading the dataset? That would help me to be more convinced that this fully closes the original issue. It's a bit tricky to know for certain because the original issue report doesn't contain a complete verifiable example. But, the workflow there seemed to be "load from file, then subset (which crashes)" so ideally a test here could include that complete workflow as well.
| import uxarray.grid.slice as slice_module | ||
| from uxarray.grid.slice import _remap_dense, _remap_kernel, _remap_searchsorted | ||
|
|
||
| import numpy as np |
There was a problem hiding this comment.
remove repeated numpy import
There was a problem hiding this comment.
No clue how I missed that twice, thanks for the catch, should be gone
| from uxarray.core.dataarray import UxDataArray | ||
|
|
||
|
|
||
| def _drop_non_grid_coords(ds): |
There was a problem hiding this comment.
Relocate to uxarray/grid/utils.py or some other file, rather than defining this helper function directly in the main grid.py file
There was a problem hiding this comment.
Got it, should be moved there now
I agree with this, I think we shouldn't let extra dimensions either into the
I don't think we should expose a grid file with time coord/dim into our testing suite as it is uncommon to have such a presence and it would be misleading in the future. |
There was a problem hiding this comment.
Thanks very much for this fix! @Sevans711 has already made great points inline and overall ( see (1) here, but I believe (2) should be avoided), so I am approving it to make way once they're addressed.
Ah, for (2), I was going based on the file at That was leading me to think it would be reasonable enough to include a second version of this file which is extremely similar but just has a scalar time coordinate instead of a Time dimension, and then call Either way, I do still want to see evidence of the full original workflow being reproduced somewhere before I am ready to approve this as fully closing the original issue. Ideally this would involve an example of "load from file, then subset" crashing on |
|
@Sevans711
The summary of a path deep dive through many files and functions basically came down to:
Which happens about here it xarray: xarray/core/coordinates.py#L1214 Filtering dimensions seemed to be distinct from this problem (filtering coordinates) in my mind. Like I don't think we saw anything that indicates a time dimension could appear naturally like this. The downstream affects of filtering dimensions is different than coordinates too. If my understanding seems wrong, and a dimension filter here seems good, maybe we can scope this out over a call because there's lots of threads that have felt tangled while exploring this, and I don't want to overcompensate in this solution because it might introduce new problems. Maybe of note, that uxarray also won't report any Dims that it has that aren't in DIM_NAMES - which is a static list (in grid.py): @property
def dims(self) -> set:
"""Names of all unstructured grid dimensions."""
from uxarray.conventions.ugrid import DIM_NAMES
return set([dim for dim in DIM_NAMES if dim in self._ds.dims])so it shouldn't ever really report anything outside of:
|
|
@Sevans711 I think you're looking for the notebook I made in this comment here: There's a link in the middle to a notebook. It should be what I did and referenced above, and what I think what you're suggesting you wanted to see |
rajeeja
left a comment
rajeeja
left a comment
There was a problem hiding this comment.
This also silently drops ICON's clon/clat/vlon/vlat/elon/elat and FESOM2's lon/lat — since only 0-d coords collide, can you narrow it to and ds[coord].ndim == 0? That still strips the scalar time from #1444 but leaves every 1-d native coordinate in place, so ICON and FESOM2 grids round-trip unchanged.
Ah, thank you! This is much closer to what I wanted. Though, I do think it isn't fully there yet (but I also might not be fully understanding). I'm hoping that the full workflow implied in #1444 (comment) can be followed as closely as possible. I understand the workflow to be something like:
The tests you pointed to are subtly different; you manually attached the uxgrid to an entirely different Dataset, rather than getting the values for both from the same source. But, as per comment linked above, it sounds like the original issue used the same exact data for the grid and for the underlying dataset data. Are you still able to reproduce the issue when utilizing the same exact data in both cases? I believe your clarifications to my original question (1) make sense; it sounds like the original issue could maybe only occur specifically from a time coordinate, but not from a time dimension. Is that correct? Assuming it is true, and there is separately a desire to avoid time dimension leaking into the Grid (as per #1667 (comment)), I would lean towards scoping the "strip extra dimensions from Grid" fix towards a separate PR to avoid too many things getting tangled here. |
|
@rajeeja Good catch, I implemented the change and everything works the same on my end so looks good. Pushed the changes. |
Rajeev, this is being called after our I/O modules (e.g. |
| # internal xarray dataset for storing grid variables. | ||
| # drop stray coordinates (e.g. a `time` carried in from the source file) so they | ||
| # can't leak onto the grid and collide during subsetting (see #1444). | ||
| # drop stray scalar coordinates (e.g. a `time` carried in from the source file) |
There was a problem hiding this comment.
What if the dataset file ( that has the grid definition embedded into it) has a non-scalar time?
4 participants
Closes #1444 (
Overview
timewas shown to exist in a grid from a user reading in a grid and data from a single file.subsetthere was a crash related to time being different on the grid vs the datatimeshouldn'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 dataGrid.__init__or_read_ugrid, both were tested extensivelyGrid.__init__is run in multiple scenarios, including when a file is read from disk, but_slice_from_gridit can cause issues.timebecause it wasn't a dimension_read_ugridoccurs primarily when a file is read from disk_read_ugrid(ds)only solved the the issue for that file typeGrid.__init__as it solved the issue in all data types, but required a different approachtimegone from grid coords and grid variablesnode_lon,node_lat,face_node_connectivitypreservedn_node,n_face> 0)bounding_boxruns, no errortimepreservedsubgrid_face_indicessel(time=...)still workstimeaxis + coord strippedChanges Made:
_drop_non_grid_coords_drop_non_grid_coordsfor cleaning stray coordinates_drop_non_grid_coordsinGrid.__init__Expected Usage
This cleaning will occur during any and all grid creations
PR Checklist
General
Testing & Benchmarking
Documentation
docs/api.rst_)AI Disclosure
AI Usage: