Conversation
github-actions
bot
added
📝 Add / update documentation
nstarman
force-pushed
the
xarray
branch
2 times, most recently
from
d1c0f93 to
e71531b
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #498 +/- ##
==========================================
- Coverage 96.30% 96.11% -0.20%
==========================================
Files 57 46 -11
Lines 2547 2417 -130
Branches 128 122 -6
==========================================
- Hits 2453 2323 -130
+ Misses 71 70 -1
- Partials 23 24 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
This PR initializes the unxt-xarray package, which provides seamless integration between unxt's JAX-based physical quantities and xarray's labeled multi-dimensional arrays through .unxt accessors.
Key Changes:
- Introduces the new
unxt-xarrayworkspace package with quantify/dequantify functionality - Adds xarray accessor registration for DataArray and Dataset objects
- Includes comprehensive test coverage with unit, conversion, and property-based tests
Reviewed changes
Copilot reviewed 20 out of 22 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| pyproject.toml | Added unxt-xarray to workspace dependencies and test groups |
| packages/unxt-xarray/pyproject.toml | Package configuration for unxt-xarray |
| packages/unxt-xarray/src/unxt_xarray/init.py | Main package entry point, registers accessors |
| packages/unxt-xarray/src/unxt_xarray/_src/init.py | Internal module exports |
| packages/unxt-xarray/src/unxt_xarray/_src/accessors.py | Implements .unxt accessor for DataArray and Dataset |
| packages/unxt-xarray/src/unxt_xarray/_src/conversion.py | Core unit conversion functions for xarray objects |
| packages/unxt-xarray/tests/unit/test_accessors.py | Tests for xarray accessor functionality |
| packages/unxt-xarray/tests/unit/test_conversion.py | Tests for conversion functions |
| packages/unxt-xarray/tests/unit/test_properties.py | Property-based tests using hypothesis |
| packages/unxt-xarray/docs/index.md | Package documentation overview |
| packages/unxt-xarray/docs/xarray-guide.md | Comprehensive usage guide |
| packages/unxt-xarray/docs/_data/README.md | Documentation for sample data files |
| packages/unxt-xarray/docs/_data/generate_sample_data.py | Script to generate sample NetCDF data |
| packages/unxt-xarray/README.md | Package README with quick start |
| packages/unxt-hypothesis/docs/testing-guide.md | Minor header formatting update |
| noxfile.py | Added xarray package support to nox sessions |
| docs/packages/unxt-xarray | Symlink to package documentation |
| docs/index.md | Added unxt-xarray to documentation index |
| .github/workflows/cd-unxt-xarray.yml | CD workflow for unxt-xarray package |
| .github/copilot-instructions.md | Added documentation formatting guidelines |
|
@adrn I'd love some input on design decisions. Not sure this is the API we want. |
| return DataArray( | ||
| data=new_data, | ||
| coords=new_coords, | ||
| dims=obj.dims, | ||
| name=obj.name, | ||
| attrs=obj.attrs, | ||
| ) |
There was a problem hiding this comment.
These constructors pass through obj.attrs by reference. Since dequantify() later mutates attrs (and coord attrs), this can unexpectedly mutate the original xarray object's metadata if xarray doesn't defensively copy. Prefer passing a shallow copy (e.g., attrs=dict(obj.attrs)) and similarly copying coord.attrs when building Variables to avoid shared-metadata side effects.
| return DataArray( | ||
| data=new_data, | ||
| coords=new_coords, | ||
| dims=obj.dims, | ||
| name=obj.name, | ||
| attrs=obj.attrs, | ||
| ) |
There was a problem hiding this comment.
These constructors pass through obj.attrs by reference. Since dequantify() later mutates attrs (and coord attrs), this can unexpectedly mutate the original xarray object's metadata if xarray doesn't defensively copy. Prefer passing a shallow copy (e.g., attrs=dict(obj.attrs)) and similarly copying coord.attrs when building Variables to avoid shared-metadata side effects.
| new_coords = {} | ||
| for name, coord in obj.coords.items(): | ||
| unit = units.get(name) | ||
| new_coords[name] = ( | ||
| coord | ||
| if unit is None | ||
| else Variable(coord.dims, u.Q(coord.data, unit), coord.attrs) | ||
| ) |
There was a problem hiding this comment.
Coordinate handling duplicates unit-attachment logic and bypasses _array_attach_units(), which already centralizes parsing/None handling. Consider using the helper (and then wrapping into Variable) to keep unit parsing behavior consistent between data variables and coordinates and reduce divergence over time.
| @st.composite | ||
| def dataarray_with_quantity(draw): | ||
| """Generate a DataArray containing a Quantity.""" | ||
| # Generate a quantity - ensure it's at least 1D | ||
| q = draw(ust.quantities()) | ||
|
|
||
| # If scalar, make it 1D by wrapping in an array | ||
| if q.ndim == 0: | ||
| q = u.Quantity(jnp.array([q.value]), q.unit) | ||
|
|
||
| # Create DataArray | ||
| da = xr.DataArray(q, dims=["x"]) | ||
|
|
||
| return da, u.unit_of(q) |
There was a problem hiding this comment.
dataarray_with_quantity() is currently unused in this test module. Either remove it to keep the property test suite focused, or add a test that consumes it. If you do keep it, note that ust.quantities() can produce arrays with ndim > 1, which will make dims=['x'] invalid—consider constraining the strategy to 1D quantities or generating dims based on q.ndim.
| # Format units as strings for attributes | ||
| unit_strs: dict[Hashable, str] = {} | ||
| for name, unit in units.items(): | ||
| if unit is not None: | ||
| if format is not None: | ||
| unit_strs[name] = format.format(unit) | ||
| else: | ||
| unit_strs[name] = str(unit) |
There was a problem hiding this comment.
The format branch in dequantify() introduces behavior that isn't covered by the new tests (they exercise the default str(unit) path). Add a unit test that calls dequantify(format=...) and asserts the produced attribute string matches the expected formatted output (or at least differs from the default) to prevent regressions.
|
|
||
| ::: | ||
|
|
||
| :::{tab-item} uv |
There was a problem hiding this comment.
The uv installation block has mismatched / nested fence markers (bash + ```bash + closing ), which is likely to break MyST/Sphinx rendering. Use a single fenced code block inside the tab item (matching the working pattern used elsewhere, e.g. in xarray-guide.md).
packages/unxt-xarray/README.md
Outdated
| # Create a DataArray with unit attributes | ||
| da = xr.DataArray( | ||
| data=[1.0, 2.0, 3.0], | ||
| dims=["x"], | ||
| coords={"x": [0.0, 1.0, 2.0]}, | ||
| attrs={"units": "m"}, | ||
| ) | ||
|
|
||
| # Quantify: convert to unxt Quantities | ||
| quantified = da.unxt.quantify() | ||
| print(quantified) | ||
| # <xarray.DataArray (x: 3)> | ||
| # <Quantity([1. 2. 3.], 'm')> | ||
| # Coordinates: | ||
| # * x (x) <Quantity([0. 1. 2.], 's')> |
There was a problem hiding this comment.
This example/output is internally inconsistent: the coordinate x is created without a units attribute, but the printed output shows it as a Quantity with unit 's'. Also, elsewhere in the docs you note that dimension coordinates are typically unwrapped by xarray, so showing a dimension coordinate as a Quantity is misleading. Update the example to (a) set coordinate unit metadata if you want it quantified, and (b) use a non-dimension coordinate if the intent is to demonstrate preserved Quantity coordinates.
| ```` | ||
|
|
||
| ::: | ||
|
|
||
| :::{tab-item} uv | ||
|
|
||
| ````bash | ||
|
|
||
| ```bash | ||
| uv add unxt | ||
| ```` | ||
|
|
||
| ::: | ||
|
|
||
| :::: | ||
|
|
||
| ``` | ||
|
|
||
| ``` | ||
|
|
There was a problem hiding this comment.
The documented tab-set installation snippet includes malformed nested code fences for the uv example (bash … ```bash … ). Since this file is guidance for contributors, it should show a known-good MyST pattern (a single fenced block inside each tab-item) to avoid propagating broken docs formatting into package docs.
| ```` | |
| ::: | |
| :::{tab-item} uv | |
| ````bash | |
| ```bash | |
| uv add unxt | |
| ```` | |
| ::: | |
| :::: | |
| ``` | |
| ``` | |
| ::: | |
| :::{tab-item} uv | |
| ```bash | |
| uv add unxt |
:::
::::
```
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
No description provided.