Repository navigation
Conversation
_set_desired_longitude_range decided whether to wrap by asking
lon.max() > 180. On a dask-backed coordinate that reduction is a compute,
and Grid.__init__ calls it -- so opening a chunked grid read and reduced
every longitude array before the caller had asked for anything, and did it
again on every isel and every copy(), each of which builds a new Grid.
The wrap is now xr.where((lon > 180) | (lon < -180), (lon + 180) % 360 - 180,
lon): elementwise, lazy, chunk-parallel, no reduction.
Measured on a synthetic 4M-node UGRID file, chunked at 500k nodes, best of 5:
dask computes wall tracemalloc peak
open_grid (before) 3 23.3ms 12.2 MB
open_grid (after) 2 16.2ms 1.0 MB
isel (before) 2
isel (after) 1
The computes that remain are a separate site on connectivity rather than
coordinates: _standardize_connectivity's conn.isnull().any() in io/_ugrid.py,
reached twice on the read path, and _slice_face_indices in grid/slice.py
materializing the connectivity it slices by. Neither is touched here.
Three behavioural differences, all from doing this per element rather than
per array.
* In-range longitudes are now left exactly alone. (lon + 180) - 180 does
not round-trip, so wrapping the whole array perturbed values that were
already in range by up to 3e-14 degrees -- in outCSne30, 2.1182935e-14
became 2.8421709e-14. On the elements that do need wrapping the two forms
are bit-identical.
* Longitudes below -180 are normalized. The old test was on the maximum
alone, so it reached the negative tail only when the same array also held
a value above 180.
* Both endpoints are kept, so the interval is the closed [-180, 180].
Folding 180.0 to -180.0 would match _xyz_to_lonlat_deg, which wraps
unconditionally into the half-open interval, but it breaks
antimeridian_face_indices: that reads a face as crossing from the span of
its longitudes, and a face with one vertex at 180 and the rest near -170
goes from a span of 350 to a span of 10 and disappears. Caught by
test_antimeridian_point_on and
test_to_geodataframe_preserves_antimeridian_faces.
Each variable is wrapped at most once, keyed on the xr.Variable object.
Without that, edge_lat -- which calls this on every property access,
outside its populate guard -- would stack a where layer onto the graph per
access. Keying on the Variable makes the memo
self-invalidating: assigning into _ds replaces that object, so a repopulated
or user-assigned coordinate is wrapped again.
The two Exodus round-trip tests compared a grid against its own
lon -> xyz -> lon reload with assert_allclose(rtol=1e-8), and passed only
because the old whole-array wrap applied to the original the identical
perturbation the reload applies. With the original left alone, the reload's
own error is exposed, and rtol is the wrong instrument for it twice over: a
longitude near zero has no magnitude for a relative tolerance to measure
against (outCSne30 nodes 4372, 4749, 7e-15 degrees apart), and longitude is
periodic, so a node on or one ulp short of the antimeridian reads 180.0 on
the original and -180.0 on the reload -- the same meridian, scored as a
360-degree error (179 nodes of outRLL1deg; outCSne30 nodes 3966, 5155).
Those two assertions now compare the difference modulo 360 to an absolute
tolerance, still ERROR_TOLERANCE. The helper still catches a 1e-7 shift and
still rejects an antipode.
Tier 0.2 of the chunked refactor plan.
Test suite: 962 passed, 1 skipped. test_plot_with_features fails identically
before and after (matplotlib figure size, unrelated).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The elementwise wrap is the right shape for a dask-backed coordinate and the
wrong one for an array already in memory. There was never a compute to defer
on that path, only a scan, and elementwise costs a pass plus ~18 bytes per
node of temporaries -- two bool masks, the arithmetic, the result -- where a
reduction costs a pass and allocates nothing.
That matters because every open_grid/open_dataset call in the benchmark suite
is eager; none pass chunks=. So the previous commit, measured there, was a
regression and nothing else.
Measured on a synthetic 4M-node UGRID file, best of 5. The wrap in isolation:
in range 0..360
old reduction 4.6ms 0MB 36.8ms 64MB
elementwise only 20.9ms 72MB 38.1ms 72MB
elementwise + guard 9.3ms 0MB 42.7ms 72MB
and through eager open_grid, where the file read dominates and the peak does
not move at all (180.0 MB in every arm):
in range 0..360
base (cmd/nogil) 145.5ms 181.9ms
elementwise only 164.4ms 186.1ms
elementwise + guard 147.9ms 185.9ms
_lon_within_range is a guard, not a decision: when it is true the wrap is the
identity on every element, so skipping it cannot change a value.
test_eager_fast_path_agrees_with_the_where_element_for_element asserts that
directly against the unguarded expression rather than assuming it, over five
inputs including both endpoints, the negative tail and NaN.
The reductions are the thing that made the old code compute, so they are
allowed only where there is nothing to defer -- da.chunks is None. A
dask-backed array skips the guard entirely, which
test_guard_is_skipped_for_dask_backed_arrays pins by asserting the guard does
compute when handed one. The chunked numbers are unchanged: 2 graph
executions, 16.0ms, 1.0 MB peak.
`and` short-circuits, so an array that does need wrapping usually pays a
single max -- the same reduction the old code paid -- before falling through.
The 0..360 column above is that extra max: ~4.6ms on a 190ms open_grid.
Test suite: 973 passed, 1 skipped. test_plot_with_features fails identically
before and after (matplotlib figure size, unrelated).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Every other open_grid in the suite is eager, so none of it can see what
Grid construction does to a dask-backed grid -- which is the path the lazy
longitude wrap changed, and the path the rest of the chunked refactor will
keep changing. OpenGridChunked adds time_open_grid and
track_peakmem_open_grid, parametrized over the oQU 480km and 120km meshes
already registered in helpers/_fixtures.py.
Chunks are held at N_CHUNKS=4 per grid dimension, so the graph is the same
shape at both resolutions and only the data under it grows. The file is read
directly rather than through CachedFixtures, because reading it is the
subject.
On MPAS this is where the lazy wrap matters most. node_lon, edge_lon and
face_lon all exist at construction, so the old max() > 180 check ran three
computes per open; the branch runs none. Measured with this benchmark's own
setup, best of 15, against HEAD's tree with coordinates.py taken from
cmd/nogil:
time tracemalloc peak
base branch base branch
480km 60.5ms 56.5ms 3.27MB 3.27MB
120km 42.3ms 38.0ms 3.53MB 4.02MB
The 120km peak reads higher on the branch, and it is not data. With gc
disabled, building the where graph allocates ~1.9 MB of transient objects,
nearly all in inspect.signature via dask/xarray op dispatch -- the same at
both resolutions and with a single chunk, so it does not scale with the
grid. At 480km it sits under the HDF5 read's own high-water mark; at 120km
the netCDF3 read is cheap enough that it becomes the peak. After a
gc.collect() the branch retains ~30 kB more than base, which is the extra
graph layers. Worth knowing for later steps: this benchmark sees
graph-construction overhead, not only bytes read.
Three things the numbers above depend on:
* Compare each resolution to its own history, not to the other. The two
files are different formats -- oQU480.grid.nc is netCDF4/HDF5,
oQU120.grid.nc is netCDF3 -- and the HDF5 open costs more, so 480km reads
slower than 120km despite a sixteenth of the data.
* Each open gets a fresh copy of the chunks dict. match_chunks_to_ugrid
(core/utils.py) writes the source-format dimension names into the
dict it is handed, so a reused one gives every sample after the first a
different argument.
* The warning "The specified chunks separate the stored chunks" is
filtered. oQU480 stores layerThickness, ssh and zMid as one chunk of all
1791 cells, so any n_face chunking splits them and xarray warns once per
open -- about data variables the grid reader drops.
Checked by calling the class the way asv does (setup(param), then the
time_/track_ methods) under -W error::UserWarning. Not run through asv
itself: its discovery subprocess cannot import uxarray from the uxarray
conda env, where the package is not installed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
cmdupuis3
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
cmdupuis3
marked this pull request as ready for review
September 28, 2026 21:02
Member
|
While hunting for where it may perform differently from main other than performance I think I found an example and a slowdown
Notebooks side by side (left is this branch, right is main): Code to try if you want to see how it works on your end/env/os: grid = ux.open_grid(MESH + r"\exodus\outCSne8\outCSne8.g")
section = grid.subset.constant_longitude(90.0)
box = grid.subset.bounding_box(lon_bounds=(-50, 50), lat_bounds=(-85, 85))
print("faces on the section at lon = 90:", section.n_face)
print("faces in the box lon -50..50: ", box.n_face)and grid = ux.open_grid(MESH + r"\mpas\QU\oQU480.231010.nc", chunks={"n_node": 1000})
for step in range(100):
grid.edge_lat
start = time.perf_counter()
grid.node_lon.values
elapsed = time.perf_counter() - start
print("dask graph layers behind node_lon:", len(grid.node_lon.data.dask.layers))
print(f"time to load node_lon: {1000 * elapsed:.0f} ms") |
Sevans711
requested changes
Oct 6, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
The benchmarks seem reasonable and show improvements which is a good sign that this is working as intended. (I see that the lazy_grid_construction.OpenGridChunked.time_open_grid('120km') benchmark shows ratio of 0.91, even though it wasn't listed directly in the "improved" section.)
Looks reasonable overall, just left some minor comments in-line!
Comment on lines
+144
to
+151
| if normalize: | ||
| x, y, z = _normalize_xyz(x, y, z) | ||
|
|
||
| lon = (lon + 180) % 360 - 180 | ||
| return lon, lat | ||
| _, lat_rad = _xyz_to_lonlat_rad(x, y, z, normalize=False) | ||
| # arctan2 is already in [-180, 180]; shifting through [0, 360) loses precision | ||
| lon = np.rad2deg(np.arctan2(y, x)) | ||
| lon = np.where(np.abs(z) > 1.0 - ERROR_TOLERANCE, 0.0, lon) | ||
| return lon, np.rad2deg(lat_rad) |
Collaborator
There was a problem hiding this comment.
If touching these lines is needed for this PR:
I would suggest to either write a new helper like _xyz_to_lonlat_rad_minus_pi_to_pi or repeat the other relevant internal logic from _xyz_to_lonlat_rad here. The current implementation computes relevant longitude pieces, like arctan2(y, x), multiple times (once here and once in _xyz_to_lonlat_rad).
Either way, please add to the docstring of both _xyz_to_lonlat_rad and _xyz_to_lonlat_deg a clarification about the range of outputs. It took me a while of looking at these to understand that rad and deg methods are not simply np.rad2deg apart, but rather the former returns in [0, 2*pi) while the latter returns in [-180, 180)!
Comment on lines
+761
to
+764
| memo = getattr(uxgrid, "_wrapped_lon_vars", None) | ||
| if memo is None: | ||
| memo = uxgrid._wrapped_lon_vars = {} | ||
|
|
Collaborator
There was a problem hiding this comment.
Minor style preference: either assume this exists (just write memo = uxgrid._wrapped_lon_vars) or don't include it inside Grid.__init__ at all.
The former matches existing style (e.g. Grid.__init__ defines self._kdtrees = {} then later there is grid_obj._kdtrees without checking whether "_kdtrees" attribute exists), while the latter keeps this implementation detail fully self-contained within this piece of code so it should be easier to maintain long-term.
Comment on lines
+738
to
+739
| Wraps elementwise with ``xr.where``, so dask-backed coordinates stay lazy and | ||
| in-range values pass through bit-for-bit. |
Collaborator
There was a problem hiding this comment.
Is xr.where actually the thing that is making dask-backed coordinates stay lazy? I would think that they would still stay lazy even if, instead of using xr.where below, you just used (da + 180) % 360 - 180.
I actually would guess that using xr.where like this is simply less efficient, because you eventually (whenever the dask graph resolves) compute the entire (da + 180) % 360 - 180 as part of it (the whole array gets computed, then passed into xr.where), while also needing to allocate an extra boolean array (out_of_range), producing a result which should be equivalent to (da + 180) % 360 - 180 anyway.
However, the point about in-range values passing through bit-by-bit is true, and may be worth the extra cost here. I think my suggestion here is just to update the docstring to clarify, the reason dask-backed coordinates stay lazy is just that .max() is being avoided in the dask case; there's nothing special about xr.where() in terms of keeping dask-backed data lazy.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1792
Overview
_set_desired_longitude_rangedoesif da.max() > 180on a lazy reduction, three times (node_lon, edge_lon, face_lon). Re-fires on everyGridconstruction, so everyiseland everycopy()pays it again.The basic idea is to use a memoized
xr.whereso that we don't need to materialize the whole dimension to view the values that are available per chunk.PR Checklist
General
Testing & Benchmarking
Documentation and Examples
docs/api.rst; internal (private) function names start with an underscore (_)AI Disclosure
AI Usage: Claude Opus 5.5