diff --git a/HOW_TO_RELEASE.md b/HOW_TO_RELEASE.md index 8c7dc963e2d..88115013a2f 100644 --- a/HOW_TO_RELEASE.md +++ b/HOW_TO_RELEASE.md @@ -28,7 +28,7 @@ upstream https://github.com/pydata/xarray (push) Then run ```sh - pixi run release-contributors + pixi run -e release release-contributors ``` and copy the output. diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst index df851eec990..0305f7f733f 100644 --- a/doc/api-hidden.rst +++ b/doc/api-hidden.rst @@ -396,6 +396,7 @@ plot.imshow plot.pcolormesh plot.scatter + plot.lines plot.surface CFTimeIndex.all diff --git a/doc/api/plotting.rst b/doc/api/plotting.rst index e2f6068351a..06977562d37 100644 --- a/doc/api/plotting.rst +++ b/doc/api/plotting.rst @@ -10,6 +10,7 @@ Dataset :toctree: ../generated/ :template: autosummary/accessor_method.rst + Dataset.plot.lines Dataset.plot.scatter Dataset.plot.quiver Dataset.plot.streamplot @@ -32,6 +33,7 @@ DataArray DataArray.plot.hist DataArray.plot.imshow DataArray.plot.line + DataArray.plot.lines DataArray.plot.pcolormesh DataArray.plot.step DataArray.plot.scatter diff --git a/doc/user-guide/io.md b/doc/user-guide/io.md index 3a9aa5155d6..c6a8a648af5 100644 --- a/doc/user-guide/io.md +++ b/doc/user-guide/io.md @@ -1244,6 +1244,35 @@ ds.to_zarr( The number of chunks on Tair matches our dask chunks, while there is now only a single chunk in the directory stores of each coordinate. +(io.zarr.rectilinear-chunks)= + +### Variable-sized (rectilinear) chunks + +Zarr v3 supports _rectilinear_ chunk grids, where chunk sizes vary along one +or more dimensions. This is useful when natural data boundaries (yearly +chunks of a daily time series, per-tile spatial extents) don't align to a +regular grid. Reading such arrays requires `zarr-python >= 3.2`, with the +experimental feature enabled: + +```python +import zarr + +with zarr.config.set({"array.rectilinear_chunks": True}): + roundtrip = xr.open_zarr("rectilinear.zarr", zarr_format=3) + roundtrip.chunks["x"] # e.g. (10, 20, 30) +``` + +```{note} +xarray can currently only *read* rectilinear-chunked Zarr V3 arrays, not +write them. Rectilinear arrays must be created with `zarr-python` directly +(or another tool, such as Icechunk) before being opened with `xr.open_zarr`. +This also means a Dataset opened from a rectilinear-chunked store cannot be +written back with {py:meth}`Dataset.to_zarr`, including writing to a +`region` of the store -- both raise a `TypeError` naming the affected +variable. Write support is tracked in +[GH11279](https://github.com/pydata/xarray/pull/11279). +``` + ### Groups Nested groups in zarr stores can be represented by loading the store as a diff --git a/doc/user-guide/plotting.md b/doc/user-guide/plotting.md index 020412d03be..52c8d5979ce 100644 --- a/doc/user-guide/plotting.md +++ b/doc/user-guide/plotting.md @@ -689,6 +689,59 @@ ds.plot.scatter(x="A", y="B", z="z", hue="y", markersize="x", row="x", col="w"); For more advanced scatter plots, we recommend converting the relevant data variables to a pandas DataFrame and using the extensive plotting capabilities of `seaborn`. +### Lines + +{py:func}`xarray.plot.lines` calls matplotlib.collections.LineCollection under the hood, +allowing multiple lines being drawn efficiently. It uses similar arguments as +{py:func}`xarray.plot.scatter`. + +Let's return to the air temperature dataset: + +```{code-cell} +airtemps = xr.tutorial.open_dataset("air_temperature") +air = airtemps.air - 273.15 +air.attrs = airtemps.air.attrs +air.attrs["units"] = "deg C" + +air.isel(lon=10).plot.lines(x="time", hue="lat"); +``` + +Make it a little more transparent: + +```{code-cell} +air.isel(lon=10).plot.lines(x="time", hue="lat", alpha=0.2); +``` + +Zoom in a little on the x-axis, and compare a few latitudes and longitudes, +group them using `hue` and `linewidth`. The `linewidth` kwarg works in +a similar way as `markersize` kwarg for scatter plots, it lets you vary the +line's size by variable value. + +```{code-cell} +:tags: [remove-stderr] + +air_zoom = air.isel(time=slice(1200, 1500), lat=[5, 10, 15], lon=[10, 15]) +air_zoom.plot.lines(x="time", hue="lat", linewidth="lon", add_colorbar=False); +``` + +Lines can modify the linestyle but does not allow markers. Instead combine {py:func}`xarray.plot.lines` +with {py:func}`xarray.plot.scatter`: + +```{code-cell} +air.isel(lat=10, lon=10)[:200].plot.lines(x="time", color="k", linestyle="dashed") +air.isel(lat=10, lon=10)[:200].plot.scatter(x="time", color="k", marker="^"); +``` + +Switching to another dataset with more variables we can analyse in similar +fashion as {py:func}`xarray.plot.scatter`: + +```{code-cell} +:tags: [remove-stderr] + +ds = xr.tutorial.scatter_example_dataset(seed=42) +ds.plot.lines(x="A", y="B", hue="y", linewidth="x", row="x", col="w"); +``` + ### Quiver Visualizing vector fields is supported with quiver plots: diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 864b344c4ab..0bcda67edff 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -6,18 +6,59 @@ What's New ========== -.. _whats-new.2026.07.1: +.. _whats-new.2026.09.0: -v2026.07.1 (unreleased) ------------------------ +v2026.09.0 (Sep 29, 2026) +------------------------- + +This release adds read support for Zarr V3 rectilinear (variable-sized) chunks, +Arrow PyCapsule export for :py:class:`DataArray`, a new +:py:meth:`DataArray.plot.lines` plotting method, :py:class:`DataTree` support +in :py:func:`apply_ufunc`, and flox-accelerated groupby medians. Bottleneck is +now disabled by default, the remaining zarr-python 2 compatibility code has been +removed, and the minimum ``h5netcdf`` version is now 1.8. It also includes many +bug fixes. + +.. warning:: + + This is the last xarray release that will support Python 3.11. Future + releases will require Python 3.12 or later. + +Thanks to the 45 contributors to this release: +ANIRUDDHA ADAK, Ahmet Kamer Çivi, Albert Yau, Andrew Scherer, Anirban Mandal, +Aryan Singh K., Asish Kumar, Chandan P, Charles Turner, Deepak Cherian, Devraj +Pal, Dipak Chaudhari, Eltsefon Mark, Evan Lyall, Illviljan, Joe Hamman, Johnson +K C, Jules Chéron, Justus Magin, Kropiunig, Marcus Campbell, Mark Harfouche, +Matt Van Horn, Matthew Rocklin, Matthias Schabel, Michael Niklas, Nick +Hodgskin, Patrick N. Raanes, Peter Hron, Puneet Dixit, QinXi, Samuel Le +Meur-Diebolt, Spencer Clark, Stanley C, Stephan Hoyer, Thomas Kluyver, Tom +Nicholas, dsk0425-sketch, genrichez, hushen, imam, omsatpute61-afk, stepit, whn +and ✨Sarah Z✨ New Features ~~~~~~~~~~~~ +- Support reading Zarr V3 arrays with rectilinear (variable-sized) chunk + grids. Using this feature needs ``zarr-python >= 3.2`` with + ``zarr.config.set({"array.rectilinear_chunks": True})``; xarray's minimum + supported zarr version is unchanged. Writing rectilinear chunks from + xarray is not yet supported (:pull:`11592`, extracted from :pull:`11279`). + By `Max Jones `_ and + `Tom Nicholas `_. + - Added `PyArrowCapsule interface `_ to :py:class:`DataArray` (``__arrow_c_schema__`` and ``__arrow_c_stream__``), enabling near zero-copy - export to pyarrow, polars or duckdb. + export to pyarrow, polars or duckdb (:pull:`11338`). By `Jules Chéron `_. +- Added new plot method :py:meth:`DataArray.plot.lines` which allows creating line plots efficiently in + a similar manner to :py:meth:`DataArray.plot.scatter`, also available for datasets. (:pull:`7173`) + By `Jimmy Westling `_. + +- :py:func:`apply_ufunc` now accepts :py:class:`DataTree` inputs, applying + ``func`` to the datasets at each node and returning trees with the same + structure (:pull:`11552`). + By `QinXi `_. + - The ``h5netcdf`` backend now reports compression and filter settings in the variable ``encoding`` consistently with the ``netCDF4`` backend (using ``h5netcdf``'s ``Variable.filters()``). This means data compressed with codecs @@ -29,7 +70,8 @@ New Features Breaking Changes ~~~~~~~~~~~~~~~~ - Disable using bottleneck by default, as certain operations are less numerically - stable than the equivalent numpy functions. + stable than the equivalent numpy functions. Use + ``xr.set_options(use_bottleneck=True)`` to opt back in (:pull:`11461`). By `Thomas Kluyver `_. - All remaining zarr-python 2.x compatibility code has been removed from the zarr backend, following the bump of the minimum ``zarr`` version to 3.0. @@ -48,11 +90,38 @@ Breaking Changes Deprecations ~~~~~~~~~~~~ +- Passing ``keepdims`` to groupby, resample or rolling reductions (e.g. + ``ds.rolling(time=12).mean(keepdims=True)``) now emits a ``FutureWarning``. + The argument was previously silently ignored or produced unexpected shapes + (:issue:`11518`, :pull:`11521`). + By `Charles Turner `_. Bug Fixes ~~~~~~~~~ +- Treat a full ``MultiIndex`` key with tuple-valued levels as scalar selection, + so ``.sel`` no longer preserves a length-1 dimension for nested tuple keys + that identify a single row (:issue:`11341`, :pull:`11348`). +- Warn when tuple-style ``DataArray`` coordinates are renamed by explicitly + provided dimension names (:issue:`11234`, :pull:`11292`). + By `Asish Kumar `_. +- Preserve non-grouped coordinates in fallback ``groupby`` reductions when grouping + by a non-leading dimension reorders the underlying variable dimensions + (:issue:`11188`, :pull:`11290`). + By `Sarthak `_. +- Fix :py:meth:`Dataset.chunk` and :py:meth:`DataArray.chunk` raising + ``ZeroDivisionError`` when using ``"auto"`` chunks on an object that has + a zero-length dimension (:pull:`11486`). + By `Charles Turner `_. +- Fix the shape of a :py:class:`~xarray.indexes.CoordinateTransformIndex`-backed + coordinate after ``.transpose()`` for non-square arrays (:pull:`11513`). + By `Samuel Le Meur-Diebolt `_. +- Fix a bug in the ``scipy`` backend where mixing non-adjacent scalar and + array indexers in :py:meth:`Dataset.sel` could silently transpose + dimension sizes when reading from a closed file object + (:issue:`10338`, :pull:`11638`). + By `Anirban Mandal `_. - Fixed ``DataArray.str.replace`` replacing every occurrence instead of none when ``n=0``. ``re.sub`` treats ``count=0`` as "replace all", so the regex code path collapsed ``n=0`` onto ``n=-1``, while the ``regex=False`` path already handled @@ -60,11 +129,12 @@ Bug Fixes By `Alexander Kropiunig `_. - ``min`` and ``max`` of object arrays with ``skipna=False`` now return ``NaN`` for slices containing missing values, instead of a result that depended on - the position of the missing value (:issue:`11501`, :pull:`#11627`). + the position of the missing value (:issue:`11501`, :pull:`11627`). By `Michael Niklas `_. - Avoid pandas' deprecated ``Series.values`` when creating a :py:class:`~xarray.Variable` from a timezone-aware :py:class:`pandas.Series` - (:issue:`11501`, :pull:`#11627`). + (:issue:`11501`, :pull:`11627`). + By `Michael Niklas `_. - Fix deadlocks when reading and writing netCDF files with dask at the same time. Combined locks now always acquire their locks in the same order, which previously depended on memory addresses, and a failed non-blocking acquire, @@ -87,24 +157,23 @@ Bug Fixes index instead of an empty index when the interval direction conflicts with the sign of ``step`` (:pull:`11623`). By `Ahmet Kamer Çivi `_. -- Fix async zarr tests using ``wraps`` with ``autospec=True`` on async methods, - which caused ``AsyncMock`` objects to leak through instead of real array data - (:pull:`11232`). - By `Joe Hamman `_. +- Preserve NumPy ``StringDType`` variables and coordinates in Zarr format 3 + round trips (:issue:`11466`, :pull:`11474`). + By `stanbot8 `_. - Fixed dask-backed bottleneck rolling reductions declaring a dtype that could differ from the dtype returned by the matching numpy-backed bottleneck path, - notably ``object`` instead of ``float64`` for boolean inputs. + notably ``object`` instead of ``float64`` for boolean inputs (:pull:`11449`). By `Matthew Rocklin `_. - Fix :py:func:`~xarray.plot.utils.label_from_attrs` breaking LaTeX axis labels when ``textwrap.wrap`` splits the string between adjacent ``$...$`` blocks, producing invalid ``$$`` sequences that matplotlib cannot render - (:issue:`11452`). + (:issue:`11452`, :pull:`11476`). By `Gen Richez `_. - Fixed :py:meth:`Dataset.stack` raising ``KeyError`` when a stacked dimension - has a falsy but valid name such as ``""``, ``False`` or ``0`` (:issue:`9969`). + has a falsy but valid name such as ``""``, ``False`` or ``0`` (:issue:`9969`, :pull:`11477`). By `JOhnsonKC201 `_. - :py:func:`polyval` now propagates ``NaN`` for ``NaT`` entries in ``timedelta64`` - coordinates instead of returning a large sentinel value (:issue:`11462`). + coordinates instead of returning a large sentinel value (:issue:`11462`, :pull:`11478`). By `Dipak Chaudhari `_. - The zarr backend now writes boolean arrays with native ``bool`` dtype instead of converting them to ``int8``. Zarr supports ``bool`` natively, so the @@ -115,20 +184,20 @@ Bug Fixes - No longer emit a ``SerializationWarning`` about a missing ``_FillValue`` when encoding a CF coordinate variable (a 1D variable named after its dimension) to an integer dtype. CF forbids missing values in coordinate variables, so a - ``_FillValue`` is not expected there (:issue:`10305`). + ``_FillValue`` is not expected there (:issue:`10305`, :pull:`11524`). By `NoiceHax `_. - Raise an informative ``TypeError`` when a :py:class:`~xarray.Coordinates` object is passed as a coordinate value, e.g. ``ds.assign_coords({"x": coords})``, instead of silently creating a broken coordinate. Pass the object directly with - ``ds.assign_coords(coords)`` (:issue:`10194`). - By `NoiceHex `_. + ``ds.assign_coords(coords)`` (:issue:`10194`, :pull:`11523`). + By `NoiceHax `_. - Fixed two bugs affecting a :py:class:`~xarray.indexes.CoordinateTransformIndex`-backed coordinate after ``.transpose()``: :py:meth:`~xarray.indexes.CoordinateTransformIndex.create_variables` (used by, e.g., ``.copy()`` and ``.reindex_like()``-based alignment) discarded the transposed dims order and silently reverted to the transform's original order; and ``xr.align(..., join="exact")`` raised a spurious ``AlignmentError`` for two objects sharing an equal multi-dimensional index whose associated coordinate variables - simply had a different dims order (:issue:`11530`). + simply had a different dims order (:issue:`11530`, :pull:`11532`). By `Samuel Le Meur-Diebolt `_. - :py:meth:`DataArray.to_series` and :py:meth:`Dataset.to_dataframe` no longer call ``.todense()`` on ``sparse.COO``-backed variables, which could raise @@ -136,20 +205,20 @@ Bug Fixes ``Dataset.to_dataframe``) crash outright since ``sparse.COO`` refuses to densify implicitly. ``to_series`` now returns only the array's stored entries; ``to_dataframe`` indexes by the union of stored entries across all - sparse variables sharing the same dims (:issue:`4007`). + sparse variables sharing the same dims (:issue:`4007`, :pull:`11528`). By `patnr `_. - Following `pandas-dev/pandas#64793`_, ensure that resampling an array to a ``Day`` frequency along a :py:class:`xarray.CFTimeIndex` produces the same results as resampling to an equivalent ``Hour`` frequency, including with the - use of ``origin`` and ``offset`` options (:pull:`11546`). This effectively + use of ``origin`` and ``offset`` options (:pull:`11547`). This effectively rolls back the resample-related changes introduced in :pull:`10650`. By `Spencer Clark `_. - Fix regression where accessing an object-dtype index eagerly attempts to - import cftime, slowing down operations. + import cftime, slowing down operations (:pull:`11558`). By `Peter Hron `_. - Fixed :py:meth:`DataArray.coarsen()` and :py:meth:`Dataset.coarsen()` raising a type error when applying reduction methods, due to the reduction methods being - dynamically generated (:issue:`8136`). + dynamically generated (:issue:`8136`, :pull:`11556`). By `Andrew Scherer `_. - Fixed a bug that caused rechunking a multi-dimensional cftime array along a subset of its dimensions to raise an error (:issue:`11567`, :pull:`11576`). @@ -157,7 +226,7 @@ Bug Fixes - :py:meth:`Dataset.diff` and :py:meth:`DataArray.diff` now raise a ``ValueError`` when ``dim`` is not an existing dimension, instead of silently returning the object unchanged. This matches the behavior of other methods such as - :py:meth:`Dataset.differentiate` and reductions like ``mean`` (:issue:`7748`). + :py:meth:`Dataset.differentiate` and reductions like ``mean`` (:issue:`7748`, :pull:`11628`). By `imam2004i `_. - Fixed indexing with an empty indexer array. An empty indexer array is now always turned into an empty slice for the backend, so that the in-memory part of the @@ -182,9 +251,21 @@ Documentation (:issue:`7517`, :pull:`11618`). By `Anirban Mandal `_. +Performance +~~~~~~~~~~~ +- Add flox support for :py:meth:`DataArray.groupby().median`, + :py:meth:`Dataset.groupby().median`, :py:meth:`DataArray.resample().median`, and + :py:meth:`Dataset.resample().median`. This significantly speeds up median reductions + when flox is installed by using flox's blockwise implementation, including + rechunking when needed. (:issue:`11238`, :pull:`11239`). By `Samuel Le Meur-Diebolt + `_. Internal Changes ~~~~~~~~~~~~~~~~ +- Fix async zarr tests using ``wraps`` with ``autospec=True`` on async methods, + which caused ``AsyncMock`` objects to leak through instead of real array data + (:pull:`11232`). + By `Joe Hamman `_. .. _whats-new.2026.07.0: @@ -236,9 +317,6 @@ Deprecations Bug Fixes ~~~~~~~~~ -- Treat a full ``MultiIndex`` key with tuple-valued levels as scalar selection, - so ``.sel`` no longer preserves a length-1 dimension for nested tuple keys - that identify a single row (:issue:`11341`). - :py:meth:`Dataset.drop_encoding` and :py:meth:`DataArray.drop_encoding` no longer copy the underlying data, avoiding excessive memory use on large datasets (:issue:`11390`, :pull:`11394`). @@ -438,13 +516,6 @@ Bug Fixes By `Emmanuel Ferdman `_. - :func:`combine_by_coords` no longer returns an empty dataset when a generator is passed as ``data_objects`` (:issue:`10114`, :pull:`11265`). By `Amartya Anand `_. -- Warn when tuple-style ``DataArray`` coordinates are renamed by explicitly - provided dimension names (:issue:`11234`, :pull:`11292`). - By `Asish Kumar `_. -- Preserve non-grouped coordinates in fallback ``groupby`` reductions when grouping - by a non-leading dimension reorders the underlying variable dimensions - (:issue:`11188`, :pull:`11290`). - By `Sarthak `_. - Fix h5netcdf backend module detection and ros3 tests (:issue:`11243`, :pull:`11274`). By `Kai Mühlbauer `_. @@ -475,12 +546,6 @@ Internal Changes runtime behavior. This enables CI integration for type stub validation and helps prevent type annotation regressions (:issue:`11086`). By `Kristian Kollsgård `_. -- Add flox support for :py:meth:`DataArray.groupby().median`, - :py:meth:`Dataset.groupby().median`, :py:meth:`DataArray.resample().median`, and - :py:meth:`Dataset.resample().median`. This significantly speeds up median reductions - when flox is installed by using flox's blockwise implementation, including - rechunking when needed. (:issue:`11238`, :pull:`11239`). By `Samuel Le Meur-Diebolt - `_. - Remove ``setup.py`` file (:pull:`11261`). By `Nick Hodgskin `_. diff --git a/pyproject.toml b/pyproject.toml index 6a011031765..56dd2a3a540 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -447,6 +447,8 @@ slowy = "slowy" Commun = "Commun" Noice = "Noice" Hax = "Hax" +Yau = "Yau" +whn = "whn" # Tests Ome = "Ome" diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index b134921419f..619a468a4b8 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -235,20 +235,25 @@ def __init__(self, zarr_array): self._array = zarr_array self.shape = self._array.shape - # preserve vlen string object dtype (GH 7328) + # preserve vlen string object dtype (GH 7328), but keep a native StringDType + dtype = self._array.dtype if ( - self._array.serializer + dtype.kind != "T" + and self._array.serializer and self._array.serializer.to_dict()["name"] == "vlen-utf8" ): dtype = coding.strings.create_vlen_dtype(str) - else: - dtype = self._array.dtype self.dtype = dtype def get_array(self): return self._array + def _preserve_string_dtype(self, value): + if self.dtype.kind == "T" and isinstance(value, str): + return np.asarray(value, dtype=self.dtype) + return value + def _oindex(self, key): return self._array.oindex[key] @@ -288,13 +293,15 @@ def __getitem__(self, key): method = self._vindex elif isinstance(key, indexing.OuterIndexer): method = self._oindex - return indexing.explicit_indexing_adapter( + value = indexing.explicit_indexing_adapter( key, array.shape, indexing.IndexingSupport.VECTORIZED, method ) # if self.ndim == 0: # could possibly have a work-around for 0d data here + return self._preserve_string_dtype(value) + async def async_getitem(self, key): array = self._array if isinstance(key, indexing.BasicIndexer): @@ -303,9 +310,24 @@ async def async_getitem(self, key): method = self._async_vindex elif isinstance(key, indexing.OuterIndexer): method = self._async_oindex - return await indexing.async_explicit_indexing_adapter( + value = await indexing.async_explicit_indexing_adapter( key, array.shape, indexing.IndexingSupport.VECTORIZED, method ) + return self._preserve_string_dtype(value) + + +def _rectilinear_encoding_error(key: str, value, name) -> TypeError: + """Error for a rectilinear (variable-sized) chunk or shard spec, which xarray + can read but not yet write.""" + return TypeError( + f"encoding[{key!r}]={value!r} for variable {name!r} is a rectilinear " + "(variable-sized) grid, e.g. read from a store opened with " + "zarr.config.set({'array.rectilinear_chunks': True}). xarray can read " + "rectilinear grids but not yet write them, including with `region` or " + "`append_dim`. To write a regular grid to a new array instead, clear " + f"the encoding (`del ds[{name!r}].encoding[{key!r}]`) and make sure the " + "variable's own chunks are uniform (e.g. `ds.chunk({dim: size})`)." + ) def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): @@ -368,6 +390,8 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): for x in enc_chunks_tuple: if not isinstance(x, int): + if isinstance(x, list | tuple): + raise _rectilinear_encoding_error("chunks", enc_chunks_tuple, name) raise TypeError( "zarr chunk sizes specified in `encoding['chunks']` " "must be an int or a tuple of ints. " @@ -383,6 +407,27 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): return enc_chunks_tuple +def _compact_chunk_sizes( + chunk_sizes: tuple[tuple[int, ...], ...], +) -> tuple[int | tuple[int, ...], ...]: + """Replace a dask-style listing of chunk sizes with a single int along any + dimension where it describes a regular grid, e.g. ((10, 10, 5),) -> (10,). + + Gives a consistent representation across zarr-python versions. + """ + compacted: list[int | tuple[int, ...]] = [] + for sizes in chunk_sizes: + sizes = tuple(sizes) + if not sizes: + # zero-length dimension: dask rejects an empty tuple here + compacted.append((0,)) + elif len(set(sizes[:-1])) <= 1 and sizes[-1] <= sizes[0]: + compacted.append(sizes[0]) + else: + compacted.append(sizes) + return tuple(compacted) + + def _get_zarr_dims_and_attrs(zarr_obj, dimension_key, try_nczarr): # Check for attributes and dimension name metadata as discussed in the Zarr encoding # specification https://docs.xarray.dev/en/stable/internals/zarr-encoding-spec.html @@ -507,12 +552,26 @@ def extract_zarr_variable_encoding( if chunks is None: chunks = "auto" encoding["chunks"] = chunks + + # Reject rectilinear shards (a sequence of sequences) like we do for chunks. + # Zarr also accepts an int, tuple of ints, "auto" or a dict here. + shards = encoding.get("shards") + if isinstance(shards, list | tuple) and any( + isinstance(x, list | tuple) for x in shards + ): + raise _rectilinear_encoding_error("shards", shards, name) + if isinstance(shards, integer_types): + # Expand to a tuple: zarr-python 3.2.x crashes on an int shard spec. + encoding["shards"] = variable.ndim * (int(shards),) + return encoding # Function below is copied from conventions.encode_cf_variable. # The only change is to raise an error for object dtypes. -def encode_zarr_variable(var, needs_copy=True, name=None): +def encode_zarr_variable( + var, needs_copy=True, name=None, *, zarr_format: ZarrFormat | None = None +): """ Converts a Variable into another Variable which follows some of the CF conventions: @@ -535,6 +594,13 @@ def encode_zarr_variable(var, needs_copy=True, name=None): var = conventions.encode_cf_variable(var, name=name, coders=ZARR_CODERS) var = ensure_dtype_not_object(var, name=name) + if ( + zarr_format == 3 + and var.dtype.kind == "T" + and var.dtype == np.dtypes.StringDType() + ): + return var + # zarr allows unicode, but not variable-length strings, so it's both # simpler and more compact to always encode as UTF-8 explicitly. # TODO: allow toggling this explicitly via dtype in encoding. @@ -854,23 +920,46 @@ def ds(self): def open_store_variable(self, name): zarr_array = self.members[name] - data = indexing.LazilyIndexedArray(ZarrArrayWrapper(zarr_array)) try_nczarr = self._mode == "r" dimensions, attributes = _get_zarr_dims_and_attrs( zarr_array, DIMENSION_KEY, try_nczarr ) attributes = dict(attributes) + array_wrapper = ZarrArrayWrapper(zarr_array) + data = indexing.LazilyIndexedArray(array_wrapper) + + try: + chunks = tuple(zarr_array.chunks) + except NotImplementedError: + # Rectilinear chunk grid (zarr-python >= 3.2): `.chunks` raises, so + # read the per-chunk sizes instead, e.g. ((10, 20, 30),). + chunks = zarr_array.read_chunk_sizes + # Normalise to an int per regular dim and a tuple per rectilinear dim, + # since what zarr-python returns above varies between versions. + chunks = _compact_chunk_sizes( + tuple(x if isinstance(x, tuple) else (x,) for x in chunks) + ) + preferred_chunks = dict(zip(dimensions, chunks, strict=True)) encoding = { - "chunks": zarr_array.chunks, - "preferred_chunks": dict(zip(dimensions, zarr_array.chunks, strict=True)), + "chunks": chunks, + "preferred_chunks": preferred_chunks, } + if array_wrapper.dtype.kind == "T": + encoding["dtype"] = array_wrapper.dtype + + try: + shards = zarr_array.shards + except NotImplementedError: + # Rectilinear shard grid: `.shards` raises, so read the per-shard + # (i.e. outer/storage chunk) sizes instead, e.g. ((1, 2),). + shards = _compact_chunk_sizes(zarr_array.write_chunk_sizes) encoding.update( { "compressors": zarr_array.compressors, "filters": zarr_array.filters, - "shards": zarr_array.shards, + "shards": shards, } ) if self.zarr_group.metadata.zarr_format == 3: @@ -930,7 +1019,8 @@ def set_attributes(self, attributes): _put_attrs(self.zarr_group, attributes) def encode_variable(self, variable, name=None): - variable = encode_zarr_variable(variable, name=name) + zarr_format = self.zarr_group.metadata.zarr_format + variable = encode_zarr_variable(variable, name=name, zarr_format=zarr_format) return variable def encode_attribute(self, a): @@ -1198,6 +1288,21 @@ def set_variables( else: del v.encoding["_FillValue"] + # We need to do this for both new and existing variables to ensure we're not + # writing to a partial chunk, even though we don't use the `encoding` value + # when writing to an existing variable. See + # https://github.com/pydata/xarray/issues/8371 for details. + # Note: Ideally there should be two functions, one for validating the chunks and + # another one for extracting the encoding. + # Must run before any resize below, so a rejected encoding can't + # leave behind a resized array that was never written to. + encoding = extract_zarr_variable_encoding( + v, + raise_on_invalid=vn in check_encoding_set, + name=vn, + zarr_format=3 if is_zarr_v3_format else 2, + ) + zarr_shape = None write_region = self._write_region if self._write_region is not None else {} write_region = {dim: write_region.get(dim, slice(None)) for dim in dims} @@ -1223,19 +1328,6 @@ def set_variables( zarr_shape = zarr_array.shape region = tuple(write_region[dim] for dim in dims) - # We need to do this for both new and existing variables to ensure we're not - # writing to a partial chunk, even though we don't use the `encoding` value - # when writing to an existing variable. See - # https://github.com/pydata/xarray/issues/8371 for details. - # Note: Ideally there should be two functions, one for validating the chunks and - # another one for extracting the encoding. - encoding = extract_zarr_variable_encoding( - v, - raise_on_invalid=vn in check_encoding_set, - name=vn, - zarr_format=3 if is_zarr_v3_format else 2, - ) - # When shards are specified, dask chunks must align with shard boundaries # (not just zarr chunk boundaries) to avoid data corruption during # parallel writes. See https://github.com/pydata/xarray/issues/10831 diff --git a/xarray/coding/variables.py b/xarray/coding/variables.py index 42976533819..15017f076d3 100644 --- a/xarray/coding/variables.py +++ b/xarray/coding/variables.py @@ -680,10 +680,12 @@ def encode(self): raise NotImplementedError def decode(self, variable: Variable, name: T_Name = None) -> Variable: - if variable.dtype.kind == "T": + if ( + variable.dtype.kind == "T" + and variable.encoding.get("dtype") != variable.dtype + ): return variable.astype(object) - else: - return variable + return variable class NativeEnumCoder(VariableCoder): diff --git a/xarray/core/options.py b/xarray/core/options.py index d7955025812..0fcdb2b7f5a 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -321,7 +321,7 @@ class set_options: netCDF files with ``open_dataset()`` and ``to_netcdf()`` if ``engine`` is not explicitly specified. May be any permutation or subset of ``['netcdf4', 'h5netcdf', 'scipy']``. - use_bottleneck : bool, default: True + use_bottleneck : bool, default: False Whether to use ``bottleneck`` to accelerate 1D reductions and 1D rolling reduction operations. use_flox : bool, default: True diff --git a/xarray/plot/__init__.py b/xarray/plot/__init__.py index 49f12e13bfc..58721ab7e0e 100644 --- a/xarray/plot/__init__.py +++ b/xarray/plot/__init__.py @@ -13,6 +13,7 @@ hist, imshow, line, + lines, pcolormesh, plot, step, @@ -28,6 +29,7 @@ "hist", "imshow", "line", + "lines", "pcolormesh", "plot", "scatter", diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index 2b4c28a9027..43726f6bc5d 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -135,6 +135,130 @@ def line( def line(self, *args, **kwargs) -> list[Line3D] | FacetGrid[DataArray]: return dataarray_plot.line(self._da, *args, **kwargs) + @overload + def lines( # type: ignore[misc,unused-ignore] # None is hashable :( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: None = None, # no wrap -> primitive + col: None = None, # no wrap -> primitive + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + **kwargs: Any, + ) -> LineCollection: ... + + @overload + def lines( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable | None = None, + col: Hashable, # wrap -> FacetGrid + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + **kwargs: Any, + ) -> FacetGrid[DataArray]: ... + + @overload + def lines( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + **kwargs: Any, + ) -> FacetGrid[DataArray]: ... + + @functools.wraps(dataarray_plot.lines) + def lines(self, *args, **kwargs) -> LineCollection | FacetGrid[DataArray]: + return dataarray_plot.lines(self._da, *args, **kwargs) + @overload def step( # type: ignore[misc,unused-ignore] # None is hashable :( self, @@ -923,6 +1047,130 @@ def __call__(self, *args, **kwargs) -> NoReturn: "an explicit plot method, e.g. ds.plot.scatter(...)" ) + @overload + def lines( # type: ignore[misc,unused-ignore] # None is hashable :( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: None = None, # no wrap -> primitive + col: None = None, # no wrap -> primitive + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + **kwargs: Any, + ) -> LineCollection: ... + + @overload + def lines( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable | None = None, + col: Hashable, # wrap -> FacetGrid + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + **kwargs: Any, + ) -> FacetGrid[DataArray]: ... + + @overload + def lines( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + **kwargs: Any, + ) -> FacetGrid[DataArray]: ... + + @functools.wraps(dataset_plot.lines) + def lines(self, *args, **kwargs) -> LineCollection | FacetGrid[DataArray]: + return dataset_plot.lines(self._ds, *args, **kwargs) + @overload def scatter( # type: ignore[misc,unused-ignore] # None is hashable :( self, diff --git a/xarray/plot/dataarray_plot.py b/xarray/plot/dataarray_plot.py index 9fead9934d6..554a89fa5e2 100644 --- a/xarray/plot/dataarray_plot.py +++ b/xarray/plot/dataarray_plot.py @@ -21,6 +21,7 @@ _guess_coords_to_plot, _infer_interval_breaks, _infer_xy_labels, + _line, _Normalize, _process_cmap_cbar_kwargs, _rescale_imshow_rgb, @@ -36,7 +37,7 @@ if TYPE_CHECKING: from matplotlib.axes import Axes - from matplotlib.collections import PathCollection, QuadMesh + from matplotlib.collections import LineCollection, PathCollection, QuadMesh from matplotlib.colors import Colormap, Normalize from matplotlib.container import BarContainer from matplotlib.contour import QuadContourSet @@ -186,24 +187,38 @@ def _prepare_plot1d_data( """ # If there are more than 1 dimension in the array than stack all the # dimensions so the plotter can plot anything: - if darray.ndim > 1: + if darray.ndim >= 2: # When stacking dims the lines will continue connecting. For floats # this can be solved by adding a nan element in between the flattening # points: - dims_T = [] - if np.issubdtype(darray.dtype, np.floating): - for v in ["z", "x"]: - dim = coords_to_plot.get(v, None) - if (dim is not None) and (dim in darray.dims): - darray_nan = np.nan * darray.isel({dim: -1}) + dims_T: list[Hashable] = [] + if plotfunc_name == "lines" and np.issubdtype(darray.dtype, np.floating): + i = 0 + for v in ("z", "x"): + coord = coords_to_plot.get(v, None) + if coord is not None: + if coord in darray.dims: + # Dimension coordinate: + d = coord + else: + # Coordinate with multiple dimensions: + c = darray[coord] + dims_filt = dict.fromkeys(c.dims) + for k in dims_filt.keys() & set(dims_T): + dims_filt.pop(k) + + d = tuple(dims_filt.keys())[i] + + darray_nan = np.nan * darray.isel({d: -1}) darray = concat( [darray, darray_nan], - dim=dim, + dim=d, coords="minimal", compat="override", join="exact", ) - dims_T.append(coords_to_plot[v]) + dims_T.append(d) + # i += 1 # Lines should never connect to the same coordinate when stacked, # transpose to avoid this as much as possible: @@ -480,6 +495,10 @@ def line( primitive : list of Line3D or FacetGrid When either col or row is given, returns a FacetGrid, otherwise a list of matplotlib Line3D objects. + + See also + -------- + Use :py:func:`xarray.plot.lines` for efficient plotting of many lines. """ # Handle facetgrids first if row or col: @@ -1047,7 +1066,7 @@ def newplotfunc( ) if add_legend_: - if plotfunc.__name__ in ["scatter", "line"]: + if plotfunc.__name__ in ["scatter", "lines"]: _add_legend( ( hueplt_norm @@ -1117,6 +1136,177 @@ def _add_labels( _set_concise_date(ax, axis=axis) +@overload +def lines( # type: ignore[misc,unused-ignore] # None is hashable :( + darray: DataArray, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: None = None, # no wrap -> primitive + col: None = None, # no wrap -> primitive + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs, +) -> LineCollection: ... + + +@overload +def lines( + darray: T_DataArray, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable | None = None, + col: Hashable, # wrap -> FacetGrid + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs, +) -> FacetGrid[DataArray]: ... + + +@overload +def lines( + darray: T_DataArray, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs, +) -> FacetGrid[DataArray]: ... + + +@_plot1d +def lines( + xplt: DataArray | None, + yplt: DataArray | None, + ax: Axes, + add_labels: bool | Iterable[bool] = True, + **kwargs, +) -> LineCollection: + """ + Line plot of DataArray values. + + Wraps :func:`matplotlib:matplotlib.collections.LineCollection` which allows + efficient plotting of many lines in a similar fashion to + :py:func:`xarray.plot.scatter`. + """ + if "u" in kwargs or "v" in kwargs: + raise ValueError("u, v are not allowed in lines plots.") + + zplt: DataArray | None = kwargs.pop("zplt", None) + hueplt: DataArray | None = kwargs.pop("hueplt", None) + sizeplt: DataArray | None = kwargs.pop("sizeplt", None) + + if hueplt is not None: + kwargs.update(c=hueplt.to_numpy().ravel()) + + if sizeplt is not None: + kwargs.update(s=sizeplt.to_numpy().ravel()) + + plts_or_none = (xplt, yplt, zplt) + _add_labels(add_labels, plts_or_none, ("", "", ""), ax) + + xplt_np = None if xplt is None else xplt.to_numpy().ravel() + yplt_np = None if yplt is None else yplt.to_numpy().ravel() + zplt_np = None if zplt is None else zplt.to_numpy().ravel() + plts_np = tuple(p for p in (xplt_np, yplt_np, zplt_np) if p is not None) + + if len(plts_np) == 3: + import mpl_toolkits + + assert isinstance(ax, mpl_toolkits.mplot3d.axes3d.Axes3D) + return _line(ax, *plts_np, **kwargs) + + if len(plts_np) == 2: + return _line(ax, *plts_np, **kwargs) + + raise ValueError("At least two variables required for a lines plot.") + + @overload def scatter( # type: ignore[misc,unused-ignore] # None is hashable :( darray: DataArray, diff --git a/xarray/plot/dataset_plot.py b/xarray/plot/dataset_plot.py index bc51d1eee80..37ad379ae83 100644 --- a/xarray/plot/dataset_plot.py +++ b/xarray/plot/dataset_plot.py @@ -921,3 +921,180 @@ def scatter( da = _temp_dataarray(ds, y, locals_) return da.plot.scatter(*locals_.pop("args", ()), **locals_) + + +@overload +def lines( # type: ignore[misc,unused-ignore] # None is hashable :( + ds: Dataset, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: None = None, # no wrap -> primitive + col: None = None, # no wrap -> primitive + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs: Any, +) -> LineCollection: ... + + +@overload +def lines( + ds: Dataset, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable | None = None, + col: Hashable, # wrap -> FacetGrid + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs: Any, +) -> FacetGrid[DataArray]: ... + + +@overload +def lines( + ds: Dataset, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs: Any, +) -> FacetGrid[DataArray]: ... + + +@_update_doc_to_dataset(dataarray_plot.lines) +def lines( + ds: Dataset, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | None = None, + row: Hashable | None = None, + col: Hashable | None = None, + col_wrap: int | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs: Any, +) -> LineCollection | FacetGrid[DataArray]: + """ + Line plot Dataset data variables against each other. + + Wraps :func:`matplotlib:matplotlib.collections.LineCollection` which allows + efficient plotting of many lines in a similar fashion to + :py:func:`xarray.plot.scatter`. + """ + locals_ = locals() + del locals_["ds"] + locals_.update(locals_.pop("kwargs", {})) + da = _temp_dataarray(ds, y, locals_) + + return da.plot.lines(*locals_.pop("args", ()), **locals_) diff --git a/xarray/plot/facetgrid.py b/xarray/plot/facetgrid.py index 8f94c3236b8..5bd730507c7 100644 --- a/xarray/plot/facetgrid.py +++ b/xarray/plot/facetgrid.py @@ -526,7 +526,7 @@ def map_plot1d( func_kwargs["add_title"] = False add_labels_ = np.zeros(self.axs.shape + (3,), dtype=bool) - if kwargs.get("z") is not None: + if coords_to_plot["z"] is not None: # 3d plots looks better with all labels. 3d plots can't sharex either so it # is easy to get lost while rotating the plots: add_labels_[:] = True @@ -539,10 +539,10 @@ def map_plot1d( # Set up the lists of names for the row and column facet variables: if self._single_group: full = tuple( - {self._single_group: x} - for x in range(self.data[self._single_group].size) + {self._single_group: v} + for v in range(self.data[self._single_group].size) ) - empty = tuple(None for x in range(self._nrow * self._ncol - len(full))) + empty = (None,) * (self._nrow * self._ncol - len(full)) name_d = full + empty else: rowcols = itertools.product( @@ -565,8 +565,8 @@ def map_plot1d( subset = self.data.isel(d) mappable = func( subset, - x=x, - y=y, + x=coords_to_plot["x"], + z=coords_to_plot["z"], ax=ax, hue=hue, _size=size_, diff --git a/xarray/plot/utils.py b/xarray/plot/utils.py index c08b014bfd2..95fd0ea5ba9 100644 --- a/xarray/plot/utils.py +++ b/xarray/plot/utils.py @@ -13,7 +13,7 @@ ) from datetime import date, datetime from inspect import getfullargspec -from typing import TYPE_CHECKING, Any, Literal, cast, overload +from typing import TYPE_CHECKING, Any, Literal, overload import numpy as np import pandas as pd @@ -38,8 +38,13 @@ if TYPE_CHECKING: from matplotlib.axes import Axes - from matplotlib.colors import Normalize + from matplotlib.collections import LineCollection + from matplotlib.colorizer import Colorizer + from matplotlib.colors import Colormap, Normalize + from matplotlib.lines import Line2D from matplotlib.ticker import FuncFormatter + from matplotlib.typing import ColorType, DrawStyleType, LineStyleType + from mpl_toolkits.mplot3d.art3d import Line3DCollection from numpy.typing import ArrayLike from xarray.core.dataarray import DataArray @@ -1075,9 +1080,9 @@ def _get_color_and_size(value): elif prop == "sizes": if isinstance(self, mpl.collections.LineCollection): - arr = self.get_linewidths() + arr = np.ma.asarray(self.get_linewidths()) else: - arr = self.get_sizes() + arr = np.ma.asarray(self.get_sizes()) _color = kwargs.pop("color", "k") def _get_color_and_size(value): @@ -1090,7 +1095,7 @@ def _get_color_and_size(value): ) # Get the unique values and their labels: - values = np.unique(arr) + values = np.unique(arr[~arr.mask]) label_values = np.asarray(func(values)) label_values_are_numeric = np.issubdtype(label_values.dtype, np.number) @@ -1733,17 +1738,20 @@ def _add_legend( # values correctly. Order might be different because # legend_elements uses np.unique instead of pd.unique, # FacetGrid.add_legend might have troubles with this: - hdl, lbl = [], [] + hdl: list[Line2D] = [] + lbl: list[str] = [] for p in primitive: hdl_, lbl_ = legend_elements(p, prop, num="auto", func=huesizeplt.func) hdl += hdl_ lbl += lbl_ - # Only save unique values: - u, ind = np.unique(lbl, return_index=True) - ind = np.argsort(ind) - lbl = cast(list, u[ind].tolist()) - hdl = cast(list, np.array(hdl)[ind].tolist()) + # Only save unique values, don't sort values as it was already sort in + # legend_elements: + lbl_ = np.array(lbl) + _, ind = np.unique(lbl_, return_index=True) + ind = np.sort(ind) + lbl = lbl_[ind].tolist() + hdl = np.array(hdl)[ind].tolist() # Add a subtitle: hdl, lbl = _legend_add_subtitle(hdl, lbl, label_from_attrs(huesizeplt.data)) @@ -1846,6 +1854,255 @@ def _guess_coords_to_plot( return coords_to_plot +@overload +def _line( + self, # Axes, + x: float | ArrayLike, + y: float | ArrayLike, + z: None = ..., + s: float | ArrayLike | None = ..., + c: Sequence[ColorType] | ColorType | None = ..., + *, + linestyle: LineStyleType | None = ..., + cmap: str | Colormap | None = ..., + norm: str | Normalize | None = ..., + vmin: float | None = ..., + vmax: float | None = ..., + alpha: float | None = ..., + edgecolors: Literal["face", "none"] | ColorType | Sequence[ColorType] | None = ..., + plotnonfinite: bool = ..., + data=..., + **kwargs, +) -> LineCollection: ... + + +@overload +def _line( + self, # Axes3D, + x: float | ArrayLike, + y: float | ArrayLike, + z: float | ArrayLike = ..., + s: float | ArrayLike | None = ..., + c: Sequence[ColorType] | ColorType | None = ..., + *, + linestyle: LineStyleType | None = ..., + cmap: str | Colormap | None = ..., + norm: str | Normalize | None = ..., + vmin: float | None = ..., + vmax: float | None = ..., + alpha: float | None = ..., + edgecolors: Literal["face", "none"] | ColorType | Sequence[ColorType] | None = ..., + plotnonfinite: bool = ..., + data=..., + drawstyle: DrawStyleType = ..., + **kwargs, +) -> Line3DCollection: ... + + +def _line( + self, # Axes | Axes3D + x: float | ArrayLike, + y: float | ArrayLike, + z: float | ArrayLike | None = None, + s: float | ArrayLike | None = None, + c: ArrayLike | Sequence[ColorType] | ColorType | None = None, + *, + linestyle: LineStyleType | None = None, + cmap: str | Colormap | None = None, + norm: str | Normalize | None = None, + vmin: float | None = None, + vmax: float | None = None, + alpha: float | None = None, + edgecolors: Literal["face", "none"] | ColorType | Sequence[ColorType] | None = None, + colorizer: Colorizer | None = None, + plotnonfinite: bool = False, + data=None, + drawstyle: DrawStyleType = "default", + **kwargs, +) -> LineCollection | Line3DCollection: + """ + ax.scatter-like wrapper for LineCollection. + + This function helps the handling of datetimes since Linecollection doesn't + support it directly, just like PatchCollection doesn't either. + + The function attempts to be as similar to the scatter version as possible. + """ + import matplotlib.collections as mcoll + import matplotlib.pyplot as plt + from matplotlib import _api, cbook + + rcParams = plt.matplotlib.rcParams + + def _parse_lines_color_args( + self, c, edgecolors, kwargs, xsize, get_next_color_func + ): + if edgecolors is None: + # Use "face" instead of rcParams['scatter.edgecolors'] + edgecolors = "face" + + c, colors, edgecolors = self._parse_scatter_color_args( + c, + edgecolors, + kwargs, + x_.size, + get_next_color_func=self._get_patches_for_fill.get_next_color, + ) + + return c, colors, edgecolors + + linewidths = s # Can be different in scatter, but same in line plots. + + # add edgecolors and linewidths to kwargs so they + # can be processed by normalize_kwargs + if edgecolors is not None: + kwargs.update({"edgecolors": edgecolors}) + if linewidths is not None: + kwargs.update({"linewidths": linewidths}) + + kwargs = cbook.normalize_kwargs(kwargs, mcoll.Collection) + # re direct linewidth and edgecolor so it can be + # further processed by the rest of the function + linewidths = kwargs.pop("linewidth", None) + edgecolors = kwargs.pop("edgecolor", None) + + # Process **kwargs to handle aliases, conflicts with explicit kwargs: + x_: np.ndarray + y_: np.ndarray + x_, y_ = self._process_unit_info( + [("x", x), ("y", y)], kwargs + ) # type ignore[union-attr] + + # Handle z inputs: + if z is not None: + from mpl_toolkits.mplot3d.art3d import Line3DCollection + + LineCollection_ = Line3DCollection + add_collection_ = self.add_collection3d + auto_scale = self.auto_scale_xyz + auto_scale_args: tuple[Any, ...] = (x_, y_, z, self.has_data()) + else: + LineCollection_ = plt.matplotlib.collections.LineCollection + add_collection_ = self.add_collection + auto_scale = self._request_autoscale_view + auto_scale_args = tuple() + + if s is None: + s = np.array([rcParams["lines.linewidth"]]) + + s_: np.ndarray = np.ma.ravel(s) + if len(s_) not in (1, x_.size) or ( + not np.issubdtype(s_.dtype, np.floating) + and not np.issubdtype(s_.dtype, np.integer) + ): + raise ValueError( + "s must be a scalar, or float array-like with the same size as x and y" + ) + + # get the original edgecolor the user passed before we normalize + orig_edgecolor = edgecolors + if edgecolors is None: + orig_edgecolor = kwargs.get("edgecolor", None) + c, colors, edgecolors = _parse_lines_color_args( + self, + c, + edgecolors, + kwargs, + x_.size, + get_next_color_func=self._get_patches_for_fill.get_next_color, + ) + + if plotnonfinite and colors is None: + c = np.ma.masked_invalid(c) + ( + x_, + y_, + s_, + edgecolors, + linewidths, + ) = cbook._combine_masks( # type: ignore[attr-defined] # non-public? + x_, y_, s_, edgecolors, linewidths + ) + else: + ( + x_, + y_, + s_, + c, + colors, + edgecolors, + linewidths, + ) = cbook._combine_masks( # type: ignore[attr-defined] # non-public? + x_, y_, s_, c, colors, edgecolors, linewidths + ) + + # Unmask edgecolors if it was actually a single RGB or RGBA. + if ( + x_.size in (3, 4) + and isinstance(edgecolors, np.ma.MaskedArray) + and not np.ma.is_masked(orig_edgecolor) + ): + edgecolors = edgecolors.data # type: ignore[assignment,unused-ignore] + + # load default linestyle from rcParams + if linestyle is None: + linestyle = rcParams["lines.linestyle"] + + if drawstyle == "default": + # Draw linear lines: + xyz = list(v for v in (x_, y_, z) if v is not None) + else: + # Draw stepwise lines: + from matplotlib.cbook import STEP_LOOKUP_MAP + + step_func = STEP_LOOKUP_MAP[drawstyle] + xyz = step_func(*tuple(v for v in (x_, y_, z) if v is not None)) + + # Broadcast arrays to correct format: + # https://stackoverflow.com/questions/42215777/matplotlib-line-color-in-3d + points = np.stack(np.broadcast_arrays(*xyz), axis=-1).reshape(-1, 1, len(xyz)) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + + collection = LineCollection_( + segments, + linewidths=s_, + linestyles=linestyle, + facecolors=colors, + edgecolors=edgecolors, + alpha=alpha, + # offset_transform=kwargs.pop("transform", self.transData), + ) + # collection.set_transform(plt.matplotlib.transforms.IdentityTransform()) + collection.update(kwargs) + + if colors is None: + if colorizer: + collection._set_colorizer_check_keywords( + colorizer, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax + ) + else: + collection.set_cmap(cmap) + collection.set_norm(norm) + collection.set_array(c) + collection._scale_norm(norm, vmin, vmax) + else: + extra_kwargs = {"cmap": cmap, "norm": norm, "vmin": vmin, "vmax": vmax} + extra_keys = [k for k, v in extra_kwargs.items() if v is not None] + if any(extra_keys): + keys_str = ", ".join(f"'{k}'" for k in extra_keys) + _api.warn_external( + "No data for colormapping provided via 'c'. " + f"Parameters {keys_str} will be ignored" + ) + collection._internal_update(kwargs) + + add_collection_(collection) + + auto_scale(*auto_scale_args) + + return collection + + def _set_concise_date(ax: Axes, axis: Literal["x", "y", "z"] = "x") -> None: """ Use ConciseDateFormatter which is meant to improve the diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index f50e662af64..18b0898f2ed 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -164,6 +164,9 @@ def get_dask_chunkmanager(): requires_zarr_v3 = requires_zarr has_zarr_v3_dtypes, requires_zarr_v3_dtypes = _importorskip("zarr", "3.1.0") has_zarr_v3_async_oindex, requires_zarr_v3_async_oindex = _importorskip("zarr", "3.1.2") +has_zarr_rectilinear_chunks, requires_zarr_rectilinear_chunks = _importorskip( + "zarr", "3.2.0" +) if has_zarr: import zarr diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 6879b6a9b96..25371da451a 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -102,8 +102,10 @@ requires_scipy, requires_scipy_or_netCDF4, requires_zarr, + requires_zarr_rectilinear_chunks, requires_zarr_v3, requires_zarr_v3_async_oindex, + requires_zarr_v3_dtypes, ) from xarray.tests.test_coding_times import ( _ALL_CALENDARS, @@ -4127,6 +4129,48 @@ def test_region_write(self) -> None: self.check_requests(expected, patches) +@requires_zarr_v3_dtypes +@pytest.mark.skipif(not HAS_STRING_DTYPE, reason="requires StringDType") +def test_roundtrip_stringdtype_zarr_v3() -> None: + dtype = np.dtypes.StringDType() + data = np.array(["a", "bb", "ccc"], dtype=dtype) + expected = Dataset( + { + "data": ("dim", data.copy()), + "scalar": np.array("a", dtype=dtype), + }, + coords={ + "dim": ("dim", data.copy()), + "nondim": ("dim", data.copy()), + }, + ) + store = zarr.storage.MemoryStore({}, read_only=False) + + with assert_no_warnings(): + expected.to_zarr(store, zarr_format=3, consolidated=False) + actual = xr.open_zarr(store, consolidated=False).load() + + for name in expected.variables: + assert zarr.open_array(store=store, path=str(name), mode="r").dtype == dtype + assert actual[name].dtype == dtype + assert_identical(expected, actual) + + +@requires_zarr_v3_dtypes +@pytest.mark.skipif(not HAS_STRING_DTYPE, reason="requires StringDType") +def test_stringdtype_with_na_object_uses_the_compatibility_path() -> None: + dtype = np.dtypes.StringDType(na_object=np.nan) + data = np.array(["a", "bb", "ccc"], dtype=dtype) + expected = Dataset({"data": ("dim", data.copy())}) + store = zarr.storage.MemoryStore({}, read_only=False) + + expected.to_zarr(store, zarr_format=3, consolidated=False) + actual = xr.open_zarr(store, consolidated=False).load() + + assert zarr.open_array(store=store, path="data", mode="r").dtype.kind != "T" + assert (actual["data"].values == data.astype(object)).all() + + @requires_zarr class TestZarrDictStore(ZarrBase): @contextlib.contextmanager @@ -7279,6 +7323,237 @@ def test_extract_zarr_variable_encoding() -> None: ) +@requires_zarr_rectilinear_chunks +class TestZarrRectilinearChunksRead: + """Reading rectilinear (variable-sized) zarr chunks. + + xarray can't write these yet, so the stores are created with zarr directly. + """ + + @staticmethod + def create_zarr_array( + store_path, shape, chunks, dimension_names, dtype, shards=None + ): + import zarr + + root = zarr.open_group(store_path, mode="w", zarr_format=3) + return root.create( + "var", + shape=shape, + # older zarr stubs (<3.2) don't include rectilinear chunk types + chunks=chunks, # type: ignore[arg-type, unused-ignore] + shards=shards, # type: ignore[arg-type, unused-ignore] + dtype=dtype, + dimension_names=dimension_names, + ) + + # expected_chunks is what xarray puts in encoding["chunks"]: an int per + # regular dim, a tuple of sizes per rectilinear dim (zarr-python itself + # differs between versions here). expected_dask_chunks is the expanded form. + cases = pytest.mark.parametrize( + "shape,chunks,dimension_names,dtype,expected_chunks,expected_dask_chunks", + [ + pytest.param( + (60,), + ((10, 20, 30),), + ("x",), + "float32", + ((10, 20, 30),), + ((10, 20, 30),), + id="1d-rectilinear", + ), + pytest.param( + (6, 20), + (2, (5, 10, 5)), + ("x", "y"), + "float64", + (2, (5, 10, 5)), + ((2, 2, 2), (5, 10, 5)), + id="mixed-regular-and-rectilinear", + ), + ], + ) + + @cases + def test_read( + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_chunks, + expected_dask_chunks, + ) -> None: + import zarr + + data = np.arange(np.prod(shape), dtype=dtype).reshape(shape) + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, shape, chunks, dimension_names, dtype + ) + arr[:] = data + + roundtrip = xr.open_zarr( + store_path, zarr_format=3, consolidated=False, chunks=None + ) + assert roundtrip["var"].encoding["chunks"] == expected_chunks + assert roundtrip["var"].encoding["preferred_chunks"] == dict( + zip(dimension_names, expected_chunks, strict=True) + ) + assert isinstance(roundtrip["var"].data, np.ndarray) + np.testing.assert_array_equal(roundtrip["var"].values, data) + + @cases + @requires_dask + def test_read_dask( + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_chunks, + expected_dask_chunks, + ) -> None: + """Dask arrays get the exact variable-sized chunks.""" + import zarr + + data = np.arange(np.prod(shape), dtype=dtype).reshape(shape) + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, shape, chunks, dimension_names, dtype + ) + arr[:] = data + + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert isinstance(roundtrip["var"].data, dask_array_type) + assert roundtrip["var"].data.chunks == expected_dask_chunks + np.testing.assert_array_equal(roundtrip["var"].values, data) + + @requires_dask + def test_read_zero_length_dim(self, tmp_path) -> None: + """zarr reports no chunk sizes at all along a zero-length dimension, + which must not be passed to dask as an empty tuple.""" + import zarr + + store_path = tmp_path / "source.zarr" + with zarr.config.set({"array.rectilinear_chunks": True}): + self.create_zarr_array( + store_path, + shape=(0, 20), + chunks=(2, (5, 10, 5)), + dimension_names=("x", "y"), + dtype="float32", + ) + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert roundtrip["var"].shape == (0, 20) + assert roundtrip["var"].encoding["chunks"][1] == (5, 10, 5) + assert roundtrip["var"].data.chunks == ((0,), (5, 10, 5)) + + def test_read_rectilinear_shards(self, tmp_path) -> None: + """Regular inner chunks with rectilinear shards. The regular chunks + must be reported as an int regardless of zarr-python version.""" + import zarr + + data = np.array([1.0, 2.0, 3.0], dtype="float32") + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, + shape=(3,), + chunks=(1,), + shards=((1, 2),), + dimension_names=("x",), + dtype="float32", + ) + arr[:] = data + + roundtrip = xr.open_zarr( + store_path, zarr_format=3, consolidated=False, chunks=None + ) + assert roundtrip["var"].encoding["chunks"] == (1,) + assert roundtrip["var"].encoding["shards"] == ((1, 2),) + np.testing.assert_array_equal(roundtrip["var"].values, data) + + def test_write_after_read_gives_helpful_error(self, tmp_path) -> None: + """Writing isn't supported yet; the error should say so, not just + "must be an int".""" + import zarr + + data = np.arange(60, dtype="float32") + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, + shape=(60,), + chunks=((10, 20, 30),), + dimension_names=("x",), + dtype="float32", + ) + arr[:] = data + + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + with pytest.raises(TypeError, match=r"rectilinear"): + roundtrip.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + + def test_append_does_not_resize_before_erroring(self, tmp_path) -> None: + """A failed append must not leave the existing array resized.""" + import zarr + + data = np.arange(60, dtype="float32") + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, + shape=(60,), + chunks=((10, 20, 30),), + dimension_names=("x",), + dtype="float32", + ) + arr[:] = data + + ds = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + with pytest.raises(TypeError, match=r"rectilinear"): + ds.to_zarr( + store_path, append_dim="x", zarr_format=3, consolidated=False + ) + + assert zarr.open_array(store_path / "var").shape == (60,) + + def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: + """Rectilinear shards must be rejected on write, like rectilinear chunks.""" + data = np.arange(60, dtype="float32") + ds = xr.Dataset({"var": ("x", data)}) + ds["var"].encoding["chunks"] = (10,) + ds["var"].encoding["shards"] = ((20, 10, 30),) + + import zarr + + with zarr.config.set({"array.rectilinear_chunks": True}): + with pytest.raises(TypeError, match=r"rectilinear"): + ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + + @pytest.mark.parametrize("shards", [20, (20,), "auto"], ids=repr) + def test_write_regular_shards_still_works(self, tmp_path, shards) -> None: + """Every non-rectilinear shard spec zarr accepts must still write.""" + data = np.arange(60, dtype="float32") + ds = xr.Dataset({"var": ("x", data)}) + ds["var"].encoding["chunks"] = 10 + ds["var"].encoding["shards"] = shards + ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + assert ( + xr.open_zarr(tmp_path / "dest.zarr")["var"].encoding["shards"] is not None + ) + + @pytest.fixture def fsspec_memory_zarr_stores(): """Write two zarr stores to fsspec's global in-memory filesystem.""" diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index 22d737ba7af..0a078a18242 100644 --- a/xarray/tests/test_plot.py +++ b/xarray/tests/test_plot.py @@ -3033,6 +3033,57 @@ def test_legend_labels_facetgrid(self) -> None: ) assert actual == expected + def test_legend_labels_facegrid2(self) -> None: + ds = xr.tutorial.scatter_example_dataset(seed=42) + + g = ds.plot.scatter( + x="A", y="B", hue="y", markersize="x", row="x", col="w", add_colorbar=False + ) + + legend = g.figlegend + assert legend is not None + actual_text = [t.get_text() for t in legend.texts] + expected_text = [ + "y [yunits]", + "$\\mathdefault{0.0}$", + "$\\mathdefault{0.1}$", + "$\\mathdefault{0.2}$", + "$\\mathdefault{0.3}$", + "$\\mathdefault{0.4}$", + "$\\mathdefault{0.5}$", + "$\\mathdefault{0.6}$", + "$\\mathdefault{0.7}$", + "$\\mathdefault{0.8}$", + "$\\mathdefault{0.9}$", + "$\\mathdefault{1.0}$", + "x [xunits]", + "$\\mathdefault{0}$", + "$\\mathdefault{1}$", + "$\\mathdefault{2}$", + ] + assert actual_text == expected_text + + actual_size = [v.get_markersize() for v in legend.get_lines()] + expected_size = [ + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 4.242640687119285, + 6.708203932499369, + 8.48528137423857, + ] + np.testing.assert_allclose(expected_size, actual_size) + def test_add_legend_by_default(self) -> None: sc = self.ds.plot.scatter(x="A", y="B", hue="hue") fig = sc.figure @@ -3402,8 +3453,9 @@ def test_maybe_gca() -> None: @requires_matplotlib +@pytest.mark.parametrize("plotfunc", ["scatter", "lines"]) @pytest.mark.parametrize( - "x, y, z, hue, markersize, row, col, add_legend, add_colorbar", + "x, y, z, hue, _size, row, col, add_legend, add_colorbar", [ ("A", "B", None, None, None, None, None, None, None), ("B", "A", None, "w", None, None, None, True, None), @@ -3412,30 +3464,33 @@ def test_maybe_gca() -> None: ("B", "A", "z", "w", None, None, None, True, None), ("A", "B", "z", "y", "x", None, None, True, True), ("A", "B", "z", "y", "x", "w", None, True, True), + ("A", "B", "z", "y", "x", "w", "x", True, True), ], ) -def test_datarray_scatter( - x, y, z, hue, markersize, row, col, add_legend, add_colorbar +def test_plot1d_functions( + x: Hashable, + y: Hashable, + z: Hashable, + hue: Hashable, + _size: Hashable, + row: Hashable, + col: Hashable, + add_legend: bool | None, + add_colorbar: bool | None, + plotfunc: str, ) -> None: - """Test datarray scatter. Merge with TestPlot1D eventually.""" - ds = xr.tutorial.scatter_example_dataset() - - extra_coords = [v for v in [x, hue, markersize] if v is not None] - - # Base coords: - coords = dict(ds.coords) - - # Add extra coords to the DataArray: - coords.update({v: ds[v] for v in extra_coords}) - - darray = xr.DataArray(ds[y], coords=coords) + """Test plot1d function. Merge with TestPlot1D eventually.""" + ds = xr.tutorial.scatter_example_dataset(seed=42) with figure_context(): - darray.plot.scatter( + getattr(ds.plot, plotfunc)( x=x, + y=y, z=z, hue=hue, - markersize=markersize, + _size=_size, + row=row, + col=col, add_legend=add_legend, add_colorbar=add_colorbar, ) @@ -3561,6 +3616,107 @@ def test_plot1d_filtered_nulls() -> None: assert expected == actual +@requires_matplotlib +@pytest.mark.parametrize("plotfunc", ["lines"]) +def test_plot1d_lines_color(plotfunc: str, x="z", color="b") -> None: + from matplotlib.colors import to_rgba_array + + ds = xr.tutorial.scatter_example_dataset(seed=42) + + darray = ds.A.sel(x=0, y=0) + + with figure_context(): + _, ax = plt.subplots() + getattr(darray.plot, plotfunc)(x=x, color=color) + coll = ax.collections[0] + + # Make sure color is respected: + expected_color = np.asarray(to_rgba_array(color)) + actual_color = np.asarray(coll.get_edgecolor()) + np.testing.assert_allclose(expected_color, actual_color) + + +@requires_matplotlib +@pytest.mark.parametrize("plotfunc", ["lines"]) +def test_plot1d_lines_linestyle(plotfunc: str, x="z", linestyle="dashed") -> None: + # TODO: Is there a public function that converts linestyle to dash pattern? + from matplotlib.lines import ( # type: ignore[attr-defined] + _get_dash_pattern, + _scale_dashes, + ) + + ds = xr.tutorial.scatter_example_dataset(seed=42) + + darray = ds.A.sel(x=0, y=0) + + with figure_context(): + _, ax = plt.subplots() + getattr(darray.plot, plotfunc)(x=x, linestyle=linestyle) + coll = ax.collections[0] + + # Make sure linestyle is respected: + w = np.atleast_1d(coll.get_linewidth())[0] + expected_linestyle = [_scale_dashes(*_get_dash_pattern(linestyle), w)] + actual_linestyle = coll.get_linestyle() + assert expected_linestyle == actual_linestyle + + +@requires_matplotlib +def test_plot1d_lines_facetgrid_legend() -> None: + # asserts that order is correct, only unique values, no nans/masked values. + + ds = xr.tutorial.scatter_example_dataset(seed=42) + + with figure_context(): + g = ds.plot.lines( + x="A", y="B", hue="y", linewidth="x", row="x", col="w", add_colorbar=False + ) + + legend = g.figlegend + assert legend is not None + actual_text = [t.get_text() for t in legend.texts] + expected_text = [ + "y [yunits]", + "$\\mathdefault{0.0}$", + "$\\mathdefault{0.1}$", + "$\\mathdefault{0.2}$", + "$\\mathdefault{0.3}$", + "$\\mathdefault{0.4}$", + "$\\mathdefault{0.5}$", + "$\\mathdefault{0.6}$", + "$\\mathdefault{0.7}$", + "$\\mathdefault{0.8}$", + "$\\mathdefault{0.9}$", + "$\\mathdefault{1.0}$", + "x [xunits]", + "$\\mathdefault{0}$", + "$\\mathdefault{1}$", + "$\\mathdefault{2}$", + ] + assert expected_text == actual_text + + actual_size = [v.get_linewidth() for v in legend.get_lines()] + expected_size = [ + 1.5, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 6.0, + 1.5, + 1.224744871391589, + 1.9364916731037085, + 2.449489742783178, + ] + np.testing.assert_allclose(expected_size, actual_size) + + @requires_matplotlib def test_9155() -> None: # A test for types from issue #9155