Repository navigation
Conversation
subset.bounding_box asked for Grid.bounds, which builds a per-face box for every face in the mesh -- materializing the node coordinate and connectivity arrays to do it -- and only then tested which faces fall inside the region. The cost was proportional to the mesh rather than the crop, so cropping a few hundred faces out of a large grid needed memory proportional to the whole grid. On a 300M-face SCRIP mesh that exceeds the machine, and the subset is not slow but impossible (#1778). Face latitudes are already resident, so they now discard faces that cannot be in the result before any bounds are computed; exact bounds are built for the survivors only. Measured on a 12-corner SCRIP mesh, crop ~0.007% of the faces: 1,600,000 faces 0.84 GiB, 2.6 s -> 0.00 GiB, 0.8 s 3,200,000 faces 1.48 GiB, 5.1 s -> 0.00 GiB, 1.6 s Two things measurement settled that reasoning got wrong. The first attempt filtered longitude on face centres too, with a margin sized from the mean face area; that is unsafe near the poles, where a face is recorded as spanning every longitude and its bounds sit up to 354 degrees from its own centre on a HEALPix z3 grid. No finite margin covers that, so longitude is left to the exact stage. The second attempt kept a latitude margin for the same reason -- and it is unnecessary, because faces_within_lat_bounds keeps a face only when its bounds are fully contained in the interval, despite the docstring saying "overlap". A contained face has a contained centre, so the filter is conservative by equality rather than by a fudge factor. The exact path is unchanged below a million faces and whenever bounds are already computed, since the filter cannot repay itself there. Tests compare the filtered and exact paths across three resolutions and eight boxes including antimeridian-spanning and polar cases, and pin the two measurements the design rests on: that longitude bounds reach far beyond a face centre, and that the latitude test is containment rather than overlap. If either stops being true the pre-filter is wrong and these fail.
The latitude prefilter alone leaves a global belt: every face at the crop's latitude, at every longitude. On a HEALPix mesh a Texas-sized box still kept 8.3% of the faces. Longitude needs a margin where latitude did not, and the margin has to widen towards the poles because meridians converge -- a fixed-area face spans more longitude the higher it sits. Normalising the measured reach by 1/cos(lat) collapses it to a constant 4.39 face widths across zooms 4, 5 and 6, which is what makes the bound derivable rather than tuned. Eight face widths, latitude-corrected, had zero violations from zoom 3 through 7. Two kinds of face cannot be filtered on a centre at all and are always kept: one whose bounds are stored min > max crosses the antimeridian, so its centre is not between them, and one containing a pole is recorded as spanning every longitude. Both are under 1% of faces by zoom 6. Texas candidates fall from 8.3% to 0.5% of the mesh, improving with resolution. 44 of 44 boxes return exactly what the unfiltered path returns, including antimeridian-spanning and polar queries.
Sevans711
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
There was a problem hiding this comment.
This is a great idea. It definitely still needs some work before it is ready to merge; see inline comments. Two more notes here:
(1) This doesn't fully close the original issue; #1778 is scoped to include the face_node_connectivity issue too. The original message in this PR should be updated accordingly, to avoid the "Closes #XXX" pattern. (EDIT: actually, upon a more careful read through of 1778, it's a bit unclear to me what the intent was. The title seems to indicate it is scoped just to the bounding-box subset issue, but the text seems to include the connectivity. If you meant to scope it only to the former, please clarify, and then feel free to keep this PR marked as "would close 1778"!)
(2) The ASV benchmarks suite does not show any relevant improvements. Can benchmarks be added which show improvement from this PR? Ideally, at least one showing memory improvement, and at least one showing speed improvement. That would help prevent performance degradation in the future.
| def test_a_precomputed_bounds_array_is_used_as_is(self): | ||
| """A grid that already has ``bounds`` should not pay for the filter. | ||
| The saving comes from *not* building the whole-mesh index. Once it | ||
| exists the exact path is free, and taking the filter anyway would add | ||
| work for nothing. |
There was a problem hiding this comment.
I don't think this test is actually doing what its docstring suggests? It doesn't look like there is any check here which confirms that the grid has not paid for the filter; the code just tests whether the result is equal to what was expected (which is redundant with the test_prefilter_matches_the_exact_path test above).
There was a problem hiding this comment.
Replaced. test_bounds_are_computed_for_candidates_only now counts the faces the bounds kernel receives and asserts bounds is never populated; every screened call also asserts that.
| def test_an_empty_region_returns_no_faces(self): | ||
| """A box over open ocean far from any face must come back empty, not | ||
| raise -- the candidate array is empty before the exact stage runs.""" |
There was a problem hiding this comment.
Again, this test is not doing what its name and docstring suggest. The reason the result has length 0 here is that the bounds are tiny, so no face is fully contained within them; there is no empty region anywhere here at all (healpix grids span the entire sphere).
There was a problem hiding this comment.
Replaced with test_region_outside_a_partial_grid_is_empty, which uses the quad-hexagon grid and a box that is genuinely far from it.
| def test_a_polar_face_is_not_excluded_by_its_centre(self): | ||
| """Longitude must not be filtered on face centres. | ||
| A face containing a pole is recorded as spanning every longitude, so | ||
| its bounds sit up to 354 degrees from its own centre on a HEALPix z3 | ||
| grid. An early version of the pre-filter tested longitude centres | ||
| with a generous margin and still would have dropped those faces for | ||
| a narrow box -- no finite margin can cover a 354 degree gap. The | ||
| equivalence tests above did not catch it, because the boxes they use | ||
| happen not to need any face whose centre lies outside them. | ||
| This pins the measurement directly rather than through a subset, so | ||
| the reason survives even if the filter is rewritten. |
There was a problem hiding this comment.
Again, this test is not doing what its name and docstring suggest. Please fix name and docstring, or fix test.
There was a problem hiding this comment.
Dropped. The centre filter and its margin are gone, so there is nothing for this test to pin.
| """The assumption the pre-filter rests on, checked rather than trusted. | ||
| ``faces_within_lat_bounds`` reads as an overlap test -- its docstring | ||
| says "overlap" -- but the implementation keeps a face only when its |
There was a problem hiding this comment.
That function's docstring should just be updated as part of this PR, instead of explaining this as a quirk in multiple places.
There was a problem hiding this comment.
Done, both docstrings now say containment. The test is gone with the centre filter; test_partially_contained_faces_are_excluded now checks straddling faces directly.
| straddlers = 0 | ||
| rng = np.random.default_rng(0) | ||
| for _ in range(40): | ||
| la = rng.uniform(-70, 60) | ||
| lat_bounds = [la, la + 4] | ||
| kept = grid.get_faces_between_latitudes(lat_bounds) | ||
| centre_inside = np.flatnonzero( | ||
| (face_lat >= lat_bounds[0]) & (face_lat <= lat_bounds[1]) | ||
| ) | ||
| straddlers += len(np.setdiff1d(kept, centre_inside)) | ||
|
|
||
| assert straddlers == 0, ( | ||
| f"{straddlers} faces were kept whose centre lies outside the " | ||
| "query interval, so the latitude filter is now an overlap test " | ||
| "and the centre pre-filter would silently drop them" | ||
| ) |
There was a problem hiding this comment.
I don't think this actually confirms that you've tested any relevant points? I would have expected a test like this to instead choose lat bounds based on face bounds, checking what happens to faces which are only-partially contained within the lat bounds. I think this is another case of the test contents not actually corresponding to the docstring.
There was a problem hiding this comment.
Rewritten as you suggest: the box edge is placed through the middle of a face taken from the reference bounds, every straddling face must be excluded, and moving the edge to that face's own bound must bring exactly it in.
| # Longitude needs a margin, unlike latitude, and the margin has to widen | ||
| # towards the poles. Meridians converge, so a face of fixed area spans | ||
| # more longitude the higher its latitude -- the same cell that covers a | ||
| # degree at the equator covers many near the pole. Normalising the | ||
| # measured reach by ``1 / cos(lat)`` collapses it to a constant 4.39 face | ||
| # widths across HEALPix zooms 4, 5 and 6, which is what makes a bound | ||
| # derivable rather than tuned: 8 face widths, latitude-corrected, had zero | ||
| # violations across zooms 3 through 7. | ||
| # | ||
| # Two cases cannot be filtered on a centre at all and are always kept. | ||
| # A face whose bounds are stored min > max crosses the antimeridian, so | ||
| # its "centre" is not between them. A face containing a pole is recorded | ||
| # as spanning every longitude. Both are rare -- under 1% of faces by zoom | ||
| # 6 -- so keeping them unconditionally costs almost nothing, and the exact | ||
| # stage discards them if they do not belong. |
There was a problem hiding this comment.
Removed along with the face-width margin.
| face_width = np.degrees(np.sqrt(4.0 * np.pi / max(int(uxgrid.n_face), 1))) | ||
| margin = 8.0 * face_width / np.maximum(np.cos(np.radians(face_lat)), _MIN_COS_LAT) |
There was a problem hiding this comment.
What?
This function can be run with any grid, not just healpix. I think these are healpix estimates for face width. I could easily make a grid with a different face width. Right?
There was a problem hiding this comment.
You were right, that was HEALPix-specific. The new screen uses no face width at all: a face is a candidate iff all of its nodes lie in the box, which holds for any mesh.
| bounds_lon = np.asarray(uxgrid.face_bounds_lon.values) | ||
|
|
||
| unfilterable = (bounds_lon[:, 0] > bounds_lon[:, 1]) | ( | ||
| (bounds_lon[:, 0] <= -179.999) & (bounds_lon[:, 1] >= 179.999) |
There was a problem hiding this comment.
What is this actually doing? Why <= 179.999 and not <= 179.9999999 or <= 179.9 or < 180?
There was a problem hiding this comment.
Gone. Pole faces are no longer special-cased; their bounds are [-180, 180], so the exact test on candidates handles them.
| np.asarray(uxgrid.face_node_connectivity.values)[candidates], | ||
| np.asarray(uxgrid.n_nodes_per_face.values)[candidates], | ||
| np.asarray(uxgrid.node_x.values), | ||
| np.asarray(uxgrid.node_y.values), | ||
| np.asarray(uxgrid.node_z.values), | ||
| np.asarray(uxgrid.node_lon.values), | ||
| np.asarray(uxgrid.node_lat.values), |
There was a problem hiding this comment.
I don't think any of these asarray calls are necessary? It probably won't hurt to include (except in special case of subclass, where np.asarray converts to numpy array while np.asanyarray does not), but these aren't included in other parts of uxarray from what I can tell, and it feels like better style to keep them out unless there is any plausible reason to expect that these .values objects won't be numpy arrays already.
| keep_lon = faces_within_lon_bounds(lon_bounds, bounds_lon) | ||
| keep_lat = faces_within_lat_bounds(lat_bounds, bounds_lat) | ||
|
|
||
| return candidates[np.intersect1d(keep_lon, keep_lat)] |
There was a problem hiding this comment.
I didn't fully understand the longitude-handling logic in this function. That doesn't mean it is wrong, just leaving a reminder to myself to take a closer look and try to understand it before approving things.
There was a problem hiding this comment.
The longitude logic is now the same as faces_within_lon_bounds: a node is in the interval when lon_min <= lon <= lon_max, or, for a box with lon_min > lon_max, when it is in either half. Nodes are normalized to [-180, 180) as the bounds are.
There was a problem hiding this comment.
What do you mean by "either half"?
There was a problem hiding this comment.
Sorry, that was unclear. A box with lon_min > lon_max is the union of two intervals, [lon_min, 180) and [-180, lon_max] — the two "halves" I meant. A node is inside when it lies in either one, so equivalently it is outside only when lon_max < lon < lon_min, the gap between them. That is what the lon < lon_min and lon > lon_max test is checking.
For [170, -170]: 175 and -175 are in, 0 is out.
The docstring now spells this out with that example, and the branch has a one-line comment.
Each edge contributed either its interior latitude extreme or its first node, never both, so a face whose edges all bulge poleward lost its lowest node from its own bounds. geoflow-small has 70 such faces.
A face's bounds contain each of its nodes, so a face can only lie inside the box if all its nodes do. Bounds are then computed for those candidates alone, on the nodes they use, so the cost follows the region rather than the mesh. Replaces the centre-based filter and its tuned margins.
rajeeja
changed the title
|
@Sevans711 Reworked in 7e577f7..dd351b0: the centre filter is replaced by a node-coordinate screen with no margins, the bounds kernel now includes every node (it was dropping some), tests rewritten to check what they say, and |
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
Getting much closer! Some small inline comments.
The ASV benchmarks look good and show the relevant changes, which addresses my previous comment (2).
For my previous comment (1), your updated description is actually even more confusing to me. Do you think the connectivity point is in scope or out of scope for the original issue? (in scope --> this PR shouldn't close the issue. out of scope --> this PR should close the issue.)
| from uxarray.constants import INT_FILL_VALUE | ||
| from uxarray.grid.bounds import _construct_face_bounds_array | ||
| from uxarray.grid.coordinates import _lonlat_rad_to_xyz |
There was a problem hiding this comment.
Tiny style comment: please move these imports to top of file, unless you have a strong preference to put them here (or unless they somehow cause circular imports, but that would surprise me).
There was a problem hiding this comment.
Moved to the top of the file — no circular import, you were right to expect that.
One knock-on worth flagging: test_bounds_are_computed_for_candidates_only and test_cached_bounds_are_used monkeypatched uxarray.grid.bounds, which stops working once the names are bound in grid_accessor at import time. They now patch uxarray.subset.grid_accessor instead. I checked the counting test still fails if I bypass the screen, so it is still load-bearing rather than silently passing.
| actual = self.screened(make(gridpath), lon_bounds, lat_bounds) | ||
| np.testing.assert_array_equal(actual, np.sort(expected)) | ||
|
|
||
| def test_accessor_matches_whole_mesh_path(self, gridpath): |
There was a problem hiding this comment.
Could you add at least a small docstring for this test and the other tests being added here? Even one or two sentences for each test would make it much easier to quickly understand what the intent is for all of them.
I think I have the right understanding for them all because of the review but I'm not 100% sure; if you add docstrings I will be able to confirm whether I agree the tests are doing what you expect them to be doing.
There was a problem hiding this comment.
Added, one or two sentences on each of the nine tests in TestBoundingBoxSubset, saying what the test pins and why rather than restating the code. The class docstring already covers the shared invariant (screened result == whole-mesh result).
Please do check them against what you expected — a couple are subtle, particularly test_nodes_on_the_box_edge (nodes sit exactly on the box edge but the arcs between them bulge past it, so the face fits its own bounds and not the box) and test_pole_face (bounds are [-180, 180], so only a full-width box contains it).
| def test_face_bounds_contain_every_node(gridpath): | ||
| """A face's latitude bounds must not be narrower than its own nodes. | ||
|
|
||
| Each edge used to contribute either its interior latitude extreme or its |
There was a problem hiding this comment.
What do you mean by "used to"? Is this relative to main, or relative to an earlier commit in this PR?
There was a problem hiding this comment.
Relative to main — the bounds kernel on main has this bug today. Each edge contributed either its interior latitude extreme or its first node, never both, so a face whose edges all bulge poleward lost its lowest node (70 faces in geoflow-small). Fixed in this PR's first commit, 7e577f7.
It is in scope here because the screen relies on a face's bounds containing its nodes; without the fix the screen could drop a face the exact test keeps. Docstring now says all of that explicitly.
…ian logic - Move bounds/intersections/coordinates imports to module scope; no circular import. Monkeypatch targets in the tests move with them. - Add a docstring to each test in TestBoundingBoxSubset saying what it pins and why. - Spell out the lon_min > lon_max case in _faces_with_nodes_within_box: the box is the union of two intervals, so a node is outside only in the gap between them. - test_face_bounds_contain_every_node: say the behaviour it pins is relative to main, fixed in this PR's first commit.
|
Updated with latest Locally: @Sevans711 ready for another look when you have time. |
2 participants
Closes #1778 (its first point; the
face_node_connectivitycost is a reader problem and stays open in #1778's discussion).Overview
subset.bounding_boxcomputedGrid.boundsfor every face before testing which fall inside the box, so a small crop of a large mesh cost memory and time proportional to the mesh.It now screens faces on their node coordinates first: a face's bounds contain each of its nodes, so a face can only lie inside the box if all its nodes do. Bounds are computed for those candidates only, on the nodes they use. The screen needs no margin and makes no assumption about face size or shape. The result is identical to the whole-mesh path, and cached bounds are used when present.
Along the way, the bounds kernel dropped a node from its own face's bounds when both adjacent edges bulge poleward (70 faces in
geoflow-small); fixed in the first commit, since the screen relies on bounds containing the nodes.faces_within_lon_bounds/faces_within_lat_boundsdocstrings now say what they do (containment, not overlap).Measured
12-corner SCRIP mesh from the #1778 MCVE, crop of ~0.007% of faces:
HEALPix zoom 8 (786k faces), Texas box: 88 MB / 973 ms → 13 MB / 11 ms.
Testing
TestBoundingBoxSubsetcompares the screened path against the whole-mesh path on 7 grids × 12 boxes (polar, antimeridian-crossing, whole globe, empty), assertsboundsis never populated, counts the faces the bounds kernel sees, checks cached bounds are used, that faces straddling a box edge are excluded, nodes on the box edge, antimeridian and pole faces explicitly, and an empty region on a partial grid.test_face_bounds_contain_every_nodepins the kernel fix.New
benchmarks/subset.py:time_andtrack_peakmem_forbounding_boxon fresh oQU and HEALPix grids.PR Checklist
General
Testing & Benchmarking
benchmarks/subset.py)Documentation and Examples
docs/api.rstAI Disclosure
AI Usage: Claude