Clean up and vectorize the face_areas API#1577
Open
rajeeja wants to merge 8 commits into
Open
Conversation
Un-deprecate compute_face_areas() as the public quadrature-capable entry point (areas by default, return_jacobian and as_dataarray flags), rename the internal worker to _compute_face_areas_and_jacobian, privatize the low-level area.py helpers, and document the quadrature options. Vectorize get_all_face_area_from_coords with prange and hoist the per-face quadrature setup out of the loop for a 6.3x speedup with identical areas. Closes #1571.
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Sevans711
added
documentation
Sevans711
self-requested a review
July 15, 2026 23:16
Collaborator
|
Thank you for these changes! I added run-benchmarks label but then realized I'm not sure if there are any benchmarks which actually rely on face_areas. If there's no visible impact, then would you be able to add at least one face_areas benchmark? If the impact is visible already, I don't feel strongly about needing a face_areas specific benchmark. I'll come back to this soon (probably tomorrow), and if the benchmark impact is visible then I will review the rest! |
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Sevans711
requested changes
Jul 16, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Looking through the plan of action discussed in #1571:
- (Looks good to me) Un-deprecate and make
compute_face_areas()the public, kwargs-capable entry point. Returnface_areasonly by default, withreturn_jacobian=False(numpy-style) for callers that also want the jacobian. - (Looks good to me) Rename
_compute_face_areas→_compute_face_areas_and_jacobian - (Not completed as planned; see notes below) Add an optional
as_dataarray=Falseflag to return aUxDataArray(result.uxgrid = self) so per-face areas are easy to plot. - (Looks good to me) Keep
calculate_total_face_area()but document it as equivalent tocompute_face_areas().sum(). - (Looks good to me) Mark the low-level
area.pyroutines private - (Mostly good but I have more notes below) Vectorize the core computation.
- (Mostly good but I have more notes below) Fix the user guide (
area_calc.ipynb). - (Looks good, I did a small commit to help) improve docstrings for
face_areasproperty andcompute_face_areas().
More notes / requests / suggestions:
- (3) Originally we discussed returning a UxDataArray object, not an xarray.DataArray. The current version here returns an xarray.DataArray. Request: make the optional flag return UxDataArray instead. I think UxDataArray is the way to go; it is more convenient, and I think the whole point for this optional flag is to return something more convenient.
- For example, plotting result from current version looks something like:
ux.UxDataArray(uxds.uxgrid.compute_face_areas(as_dataarray=True), uxgrid=uxds.uxgrid).plot(). But, with this requested change, it would become simpler:uxds.uxgrid.compute_face_areas(as_dataarray=True).plot(). - I wonder if
as_dataarrayis the wrong name here... Request/suggestion: rename the flag toas_uxarrayoras_uxdataarrayto clarify it is a UxDataArray instead? Of the three names, my preferences would beas_uxarray(concise & clear) >as_uxdataarray(verbose & clear) >as_dataarray(sounds like xarray not uxarray). (Let me know if you want me to take on the task of setting this up!)
- For example, plotting result from current version looks something like:
- (7a) Request: plot face_areas somewhere in area_calc.ipynb, using the new optional flag. Either in a new section, or inside an existing section, e.g. maybe it makes sense to add in the "Calculate Area from Multiple Faces in Spherical Coordinates" section?
- (7b) In the healpix.ipynb guide, there is a section comparing geometric face area calculations versus theoretical healpix areas. This has been updated to now be equivalent to grid.compute_face_faces() / grid.face_areas. There is also some numerical analysis of how the results vary across the grid. Request: plot this face_areas ratio across the grid. After making changes in response to point (3) above, it should be as simple as:
(grid.compute_face_faces(as_uxarray=True) / grid.face_areas).plot(). - (6) I also measured roughly 6x faster face_areas computations on main versus on this branch. This is a great improvement! Some of the ASV benchmarks show improvements in reducing peak memory usage, but none of them show this significant speedup. Request: add at least one benchmark that gets improved by these changes.
erogluorhan
changed the title
Collaborator
|
@rajeeja thank you for adding the total face area error check. I don't want to interrupt you if you're still working on things, but wanted to share my thoughts now while I am thinking about it:
|
compute_face_areas() gains as_uxarray (renamed from as_dataarray) returning a UxDataArray paired with the grid; plot face areas in area_calc.ipynb and the geometric/theoretical ratio in healpix.ipynb using it; add a FaceAreas asv benchmark capturing the vectorization speedup.
Contributor
Author
|
@Sevans711 thanks for the review — fixed the items: |
Contributor
Author
|
@Sevans711 Good catches on both. You're right the check misfires on grids whose |
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 #1571.
compute_face_areas()as the public quadrature-capable entry point — returns areas by default, withreturn_jacobianandas_uxarrayflags (the latter returns aUxDataArraypaired with the grid)._compute_face_areas_and_jacobianand privatize the low-levelarea.pyhelpers.quadrature_rule/order/latitude_adjusted_areaoptions; notecalculate_total_face_areais equivalent tocompute_face_areas().sum().get_all_face_area_from_coordswithprangeand hoist the per-face quadrature setup out of the loop — ~6x faster with identical areas; adds aFaceAreasasv benchmark.area_calc.ipynband the geometric/theoretical ratio inhealpix.ipynbusing the newas_uxarrayflag.