Repository navigation
Conversation
and pass through **kw_apply_ufunc
and other improvements (will put more detailed description in PR overview).
Sevans711
added
new feature
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have stayed the same:
Benchmarks that have got worse:
|
|
Re: ASV benchmarks, I think it looks okay. I am not expecting any performance changes, and I'm seeing no changes for the most part. The only exceptions are a small (~20%) slowdown in the 480km gradient benchmark (but no significant slowdown in the 120km gradient benchmark, so I think this might just be noise?), and a small increase to track_peakmem_gradient (~10%) but that also seems to occur randomly on different PRs which aren't touching gradient at all, so I think it's okay to ignore here? Requesting reviews because things mostly look okay. The readthedocs build isn't building right now but it's not this PR's fault, readthedocs just had a random issue. If reviewers can review the non-readthedocs parts for now (i.e. everything except the vector_calculus.ipynb changes) and then check the readthedocs parts in a later review, that would be awesome! |
cmdupuis3
left a comment
cmdupuis3
left a comment
There was a problem hiding this comment.
I have two speed enhancements on the cmd/2d+calculus branch.
Other issues:
- Results now inherit the input's attrs via apply_ufunc/binary ops (keep_attrs on in current xarray), so gradient/curl/divergence outputs carry the input's long_name/standard_name, which mislabels the derivative.
- Unit issues:
divergence(units=...)is now ignored. Also, curl and divergence units are missing their parentheses e.g. m/s/m vs. (m/s)/m
| # Create two components for vector field (using same data for simplicity in test) | ||
| u_component = uxds['t2m'] | ||
| v_component = uxds['t2m'] | ||
| v_component = uxds['t2m'].rename("vname") |
There was a problem hiding this comment.
This is kind of strange. Normally u and v are individual variables, so you'd load them like uxds['u'] and uxds['v'].
There was a problem hiding this comment.
Yes, but this is a quick way to get two variables, and there's nothing invalid about a vector field whose two components are equal to each other, right? Also, the only thing I changed here was to add the call to .rename(); getting both components from the "t2m" variable is already on main. I added a comment to clarify this: # (used .rename() to check below the result name/attrs when input names differ)
There was a problem hiding this comment.
I guess so. It's just kind of odd, I would think that if t2m is available, u and v would probably be available as well.
There was a problem hiding this comment.
Ah, this test isn't on a "real" dataset though, it's just on the "quad-hexagon" tutorial example data, which only has "t2m". It's not a good demonstration of how to use the feature but it should be okay for the test suite, I think!
There was a problem hiding this comment.
Okay, yeah if it isn't in the test dataset, we don't need to be worried. Maybe you can make v negative or something though, I sense that having them be exactly equal could lead to some confusing bugs in the future.
There was a problem hiding this comment.
I'm still not really understanding why you would prefer v == -u instead of v == u. Both are perfectly valid physically, especially on a local mesh made of just 4 hexagons, right? They both correspond to flows at a 45 degree angle; changing the sign of v just changes the direction from northeast/southwest to northwest/southeast.
Also, the tests we're talking about in this thread just test the output format, not the correctness of the computations (that is covered by the other tests in the file). If we were creating these tests from scratch in this PR I wouldn't mind changing the values, but since the tests already existed, I'd prefer to stick with the existing values to avoid making unnecessary changes.
There was a problem hiding this comment.
Because if v == u, then f(v) == f(u), etc, so it's harder to tell if f() is working correctly when something goes wrong down the line.
There was a problem hiding this comment.
There is now a negative sign on the v_component in the test_..._output_format tests as requested; for more details see #1823 (comment) below.
| *xarrays, | ||
| input_core_dims=input_core_dims, | ||
| output_core_dims=output_core_dims, | ||
| vectorize=True, |
There was a problem hiding this comment.
This vectorize=True overrides the one this function receives. This should just be vectorize by itself.
There was a problem hiding this comment.
The vectorize=True is intentional, as explained in the docstring. Added a comment at this line to clarify that it is intentional, pointing to the docstring for more details.
There was a problem hiding this comment.
No, that still doesn't make sense to me. We shouldn't be providing optional arguments in the API if they only have one valid value, unless the plan is to implement it in the future.
There was a problem hiding this comment.
Fair enough, and it's a style choice so I will go with your preference. Moved the clarification from documentation of the vectorize parameter in the docstring into a note at the bottom of the docstring, and updated the function accordingly.
| input_core_dims=input_core_dims, | ||
| output_core_dims=output_core_dims, | ||
| vectorize=True, | ||
| **kw_apply_ufunc, |
There was a problem hiding this comment.
Also, dask= is never set, so it defaults to dask="forbidden". This results in a regression where the calculus operators crash when they're running on chunked data.
Consider adding this snippet before result:
xarrays = [arr.to_xarray() for arr in uxarrays]
if any(x.chunks is not None for x in xarrays):
# f needs every grid element at once, so grid_dim must be a single chunk;
# other dims keep their chunks and are processed block-by-block.
xarrays = [
x.chunk({grid_dim: -1}) if x.chunks is not None else x for x in xarrays
]
kw_apply_ufunc.setdefault("dask", "parallelized")There was a problem hiding this comment.
Fixed! Didn't handle it originally because using dask="parallelized" requires setting output_dtypes (otherwise there is a crash). Now, the helper function _apply_1dfunc_with_grid_core_dim uses dask="parallelized" by default, and raises a helpful warning if output_dtypes was not provided. Meanwhile, _compute_gradient() now provides output_dtypes when calling the helper; see also my comment on your next point.
Added some tests which would be sensitive to this issue, however they crashed locally on my end due to issue #1798. So, I commented them out, with a TODO comment noting that they should be uncommented once that issue gets fixed (i.e. either as part of #1797, or as part of this PR if 1797 merges into main first). I confirmed locally that copy-pasting the @parallel_njit decorator from 1797, onto the _compute_gradients_on_faces causes these currently-commented-out tests to pass.
There was a problem hiding this comment.
I didn't automate chunking into a single chunk along the grid dimension because it feels a bit wrong to me to do that silently? The error that gets raised if n_face is split into multiple chunks looks like this:
ValueError: dimension n_face on 0th function argument to apply_ufunc with dask='parallelized' consists of multiple chunks, but is also a core dimension. To fix, either rechunk into a single array chunk along this dimension, i.e., ``.chunk(dict(n_face=-1))``, or pass ``allow_rechunk=True`` in ``dask_gufunc_kwargs`` but beware that this may significantly increase memory usage.
I think it's okay for that error to get raised all the way to the user, allowing them to rechunk their data before calling the relevant methods? That way if it is extra slow or has much larger memory usage, they will understand it is a dask/chunking issue, not just a uxarray issue.
|
|
||
| def _compute_gradient(data, scale_by_radius=True): | ||
| from uxarray import UxDataArray | ||
| def _compute_gradient(data, scale_by_radius=True, **kw_apply_ufunc): |
There was a problem hiding this comment.
kw_apply_ufunc.setdefault("output_dtypes", [np.float64, np.float64]) would help dask to know what the _compute_gradients_on_ types are, which prevents it from needing to compile them with numba on dummy args. As it is, those functions end up as np.float64 anyway, but maybe there's a way to be flexible with the precision.
There was a problem hiding this comment.
Fixed, using output_dtype = np.result_type(data.dtype, np.float64). The function uses arithmetic with grid values like face_lon which are already np.float64, so I don't think it's possible to end up with less precision than that.
But, using this command allows to support complex numbers! 😄 I doubt anyone will be using complex numbers anytime soon, but it's very simple to support them here, just needed this command and to pass output_dtype into the numba method so that the numba method's output array buffers get initialized with the correct types. Added some tests to ensure complex numbers are supported and give correct results, too!
also adds tests for what the "units" attr should contain.
See discussion in #1573 for details. Basically, DeprecationWarning is not visible to users by default, but FutureWarning is.
|
The readthedocs build is now ready for review too! The animation at the end misses the last frame but that's a minor issue and it has been fixed upstream (see holoviz/holoviews#7113, holoviz/panel#8860); I think it's okay to keep as-is for now and allow it to resolve itself when the next panel release comes out (it looks like panel releases have been occurring every couple of months or so, and the last one was in August). Regarding @cmdupuis3's point (1): I can't reproduce the inputs' Regarding point (2): the divergence(units=...) being ignored is intentional (see point (4) in original post), but the fact that it was being silently ignored is not intentional. Swapped DeprecationWarning to FutureWarning because DeprecationWarnings are often hidden by default (as discussed in #1573). The missing parentheses (e.g. "(m/s)/m" not "m/s/m") was due to an inconsistency between curl/divergence units (which always put parentheses on everything, e.g. "(K)/m") and gradient units (which never put parentheses on anything, e.g. "K/s/m"). I've updated this so that parentheses are added only when there is a "/" in the input units, e.g. "m/s" --> "(m/s)/m", while "K" --> "K/m". Updated the Regarding the speedups on the I believe all points have been addressed and this is now ready for re-review! |
|
Nice feature — one helper routing the grid dim through Before I dig in, please close out Chris's remaining points. On the |
@rajeeja thank you, and, done! I still don't personally understand your reasoning for the sign flip; the relevant tests are completely insensitive to the actual values of the data. Nothing in the tests in question (the |
3 participants
Closes #1723
Overview
Adds support for applying calculus operators to 2D+ arrays (more dims than just the grid dimension), by creating the new
_apply_1dfunc_with_grid_core_dim()helper method and using it inside_compute_gradient(). The operators which now support 2D+ arrays are gradient(), curl(), divergence(), and scalardotgradient().Tests note: the tests which prove that these changes are compatible with dask arrays remain commented out for now; they should be uncommented when #1798 merges. See #1823 (comment) for more details.
While working on this I made numerous miscellaneous improvements related to calculus methods, expanding the scope of this PR as follows.
Small expansions of scope (examples):
xr.set_options(display_expand_data=False)). The corresponding icon still shows/hides, the default state is just "hidden" instead of shown.~with\~in one markdown cell (starts with "Pure rotation is incompressible...") which was rendering locally with strikethrough for some text instead of tildes.Sidenote about the notebook: there are a few cells in the middle which I didn't touch but the diff shows they were edited; this was done by ruff formatting.
Small expansions of scope (functionality):
Tiny expansions of scope (functionality):
grad_u_zonalandgrad_v_meridionalinstead of the confusingu = grad_u_zonalandv = grad_v_meridional.Expected Usage
Exactly the same syntax as the operators have currently, but they can now be applied to arrays with more dimensions:
PR Checklist
General
Testing & Benchmarking
Documentation and Examples
docs/api.rst; internal (private) function names start with an underscore (_)AI Disclosure
AI Usage: Claude for discussions and some code suggestions; GitHub Copilot's inline code suggestions