diff --git a/docs/api/tiles.md b/docs/api/tiles.md index dec6c782..b262d17e 100644 --- a/docs/api/tiles.md +++ b/docs/api/tiles.md @@ -5,9 +5,9 @@ # xdas.virtual.tiles Lazy tile-backed virtual arrays: the only backend of the formats that -HDF5 virtual datasets cannot serve (Silixa TDMS, MiniSEED), the default -one for Febus, and available on request from every other engine -(`vtype="tiles"`). +HDF5 virtual datasets cannot serve (Silixa TDMS, and everything ObsPy +reads), the default one for Febus, and available on request from every +other engine (`vtype="tiles"`). ## TileArray @@ -27,7 +27,7 @@ Attributes TileArray.chunks TileArray.ntiles TileArray.engine - TileArray.attrs + TileArray.root ``` Methods @@ -40,7 +40,7 @@ Methods TileArray.from_variable TileArray.to_dataset TileArray.create_variable - TileArray.finalize_save + TileArray.sibling_datasets TileArray.concat TileArray.expand_dims TileArray.squeeze diff --git a/docs/release-notes.md b/docs/release-notes.md index b5cfc0e4..eba1ea92 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -3,25 +3,33 @@ ## 0.2.9 (unreleased) ### New Features -- **Tile-backed virtual arrays.** The new `xdas.virtual.tiles` module exposes file archives as one lazy `TileArray`. Slicing (any step, including negative), integer indexing, `np.newaxis`, concatenation, and the numpy manipulation routines (the `transpose`, `flip`, `split`, `stack` and `atleast` families, `expand_dims`, `squeeze`, `roll`, `tile`, `delete`, `append`/`insert`) all stay lazy; whole-array reductions (`sum`, `mean`, `min`, `max`, …) stream one tile row at a time; reads touch only the tiles the selection overlaps (@atrabattoni). -- **`vtype="tiles"` on every HDF5 engine.** The open functions with `vtype="tiles"` return tile-backed arrays for the asn, febus, terra15, apsensing, prodml and native xdas engines. Silixa and MiniSEED always emit them now (replacing the serialized-dask-graph fallback, with time-axis push-down for Silixa), and Febus defaults to them — one tile per file, where the HDF5 backing needed one virtual mapping per data block. Custom engines add support by implementing `Engine.load_tile(path, selection, **params)` (@atrabattoni). -- Tile-backed arrays round-trip through the native xdas netCDF format: the manifest is stored as a compact `__tiles__` sibling group, relocatable by editing its single root path and directly readable by the 0.3 line (@atrabattoni). -- **Explicit engine configuration.** The open functions declare `engine`, `vtype` and `ctype`, and `engine` also accepts a configured `xdas.io.Engine` instance. Format-specific parameters (`overlaps`/`offset` for febus, `ignore_last_sample` for miniseed, `swapped_dims` for prodml, `tz` for terra15, `group` for the native format) are engine constructor parameters, validated up front (@atrabattoni). -- `open_mfdataarray` no longer caps the number of files when the resolved vtype consolidates its scan results, which `tiles` does; the 100 000 ceiling remains for `hdf5`, which builds one virtual mapping per file and so cannot be fused into anything smaller (@atrabattoni). -- **Streamed multi-file combining.** `open_mfdataarray` now fuses scan results every 100 000 files instead of holding one data array per file until the end, so memory no longer grows with the archive: results are accumulated without coordinate simplification (lossless in any arrival order) and sorted once at the end, giving the same result as before whatever the file naming. Acquisitions interleaved in time now group by compatibility, one array per acquisition, instead of splitting at each alternation. Since the batch size is also the ceiling, anything that opened in one call before still takes the single-batch path unchanged (@atrabattoni). -- **`xdas.sortby`.** Sort a tile- or stack-backed data array along a dimension by coordinate value, lazily: the blocks are permuted through the manifest without reading any data. This is how the streamed combine orders shuffled archives, exposed for standalone use (@atrabattoni). -- `simplify` runs in linear time whatever the number of gaps: the reduce stage is now a one-pass sleeve instead of Douglas-Peucker, which degenerated quadratically on gap-rich coordinates (a 100 000-file gappy archive simplified in minutes; now milliseconds). The deviation guarantee is unchanged — dropped tie points stay within `tolerance` of the curve, surviving values never move — though the surviving tie-point selection may differ slightly on jittery axes (@atrabattoni). +- **Tile-backed virtual arrays.** The new `xdas.virtual.tiles` module exposes a file archive as one lazy `TileArray`: slicing (any step), integer indexing, `np.newaxis`, concatenation and the numpy manipulation routines stay lazy, reductions stream one tile row at a time, and a read touches only the tiles the selection overlaps. Select it with `vtype="tiles"` on any HDF5 engine — Febus defaults to it, Silixa and the ObsPy formats always use it. Tile-backed arrays round-trip through the native netCDF format as a compact `__tiles__` group, relocatable by editing the single root path of its header. Custom engines opt in by implementing `Engine.load_tile(path, selection, **params)` (@atrabattoni). +- **The `obspy` engine**, named for the library rather than for a format: decoding is `obspy.read`, so miniSEED, SAC, GSE2, SEG-2 and everything else ObsPy supports goes through it. Each contiguous `Trace` becomes one lazy `DataArray` and the collection mirrors the `Stream`, nested `network / station / location / channel`; files the miniseed engine rejected (two sampling rates, duplicated ids, interleaved acquisitions) are now readable (@atrabattoni). +- **`xdas.trim_overlaps`.** Resolve the overlaps of a data array or collection by dropping the duplicated samples, keeping the later copy by default (ObsPy's `merge(method=1, interpolation_samples=0)`) or the earlier one with `keep="first"`. Trimming stays at the manifest level, so a lazy array stays lazy; `xdas.split(da, "overlaps")` still keeps every copy (@atrabattoni). +- **`DataCollection.select`**, with `obspy.Stream.select` semantics: `dc.select(station="SX00*", channel="HH?")`. It is an alias of `query`, which now applies an indexer wherever its level sits in the tree rather than only at the root (@atrabattoni). -### Breaking Changes -- Python 3.10 support is dropped and the numpy requirement is raised to 2.3: the tile manifests use `np.strings` routines introduced in numpy 2.3, which itself requires Python 3.11+. Python 3.10 reaches end of life in October 2026 (@atrabattoni). -- Passing a bare read function as `engine` now raises a `TypeError`: subclass `xdas.io.Engine` instead (see the data-formats documentation) (@atrabattoni). -- Misspelled or unsupported keyword arguments passed next to an engine name now raise a `TypeError` instead of being silently ignored, and combining `vtype`, `ctype` or engine keywords with an already configured engine instance raises a `ValueError` (@atrabattoni). +### Improvements +- **Scanning scales to archives of any size.** `open_mfdataarray` fuses scan results every 100 000 files instead of holding one data array per file, so memory no longer grows with the archive. With `vtype="tiles"` the file-count ceiling is lifted and constant tile geometry costs one element instead of one per tile: a 23-million-tile archive opens in 1.11 GB instead of 1.67 GB (@atrabattoni). +- **Explicit engine configuration.** Every open function declares `engine`, `vtype` and `ctype`, and `engine` also accepts a configured `xdas.io.Engine` instance. Format-specific parameters (`overlaps`/`offset` for febus, `ignore_last_sample` for miniseed, `swapped_dims` for prodml, `tz` for terra15, `group` for the native format) are engine constructor arguments, validated up front — a misspelled keyword now raises instead of being silently ignored (@atrabattoni). +- **`xdas.sortby`.** Sort a tile- or stack-backed data array along a dimension by coordinate value, lazily: the blocks are permuted through the manifest without reading any data (@atrabattoni). +- `xdas.concat` opening a *new* dimension now checks that the inputs agree on their other coordinates and promotes the scalar ones that vary. Stacking the components of a station is `xd.concat(traces, "channel")`, lazily, with `channel` becoming a real coordinate (@atrabattoni). +- Acquisitions interleaved in time now group by compatibility, one array each, instead of splitting at every alternation (@atrabattoni). +- Saving and opening a data collection is linear in its size again (#81): saving 1300 events went from ~53 min to ~35 s (@atrabattoni). +- `simplify` runs in linear time whatever the number of gaps. The deviation guarantee is unchanged, though the surviving tie points may differ slightly on jittery axes (@atrabattoni). +- When no engine can open a file, the error now lists what each engine that recognised it said, instead of only reporting that none succeeded (@atrabattoni). -### Deprecations -- Writing dask-backed virtual arrays is deprecated and emits a `FutureWarning`; existing files still open, but no engine emits them any more (@atrabattoni). +### Breaking Changes +- Python 3.10 support is dropped and the numpy requirement is raised to 2.3 (@atrabattoni). +- **Dask virtualization is removed**, reader and writer alike, along with the `xdas.dask` module: a `__dask_array__` graph can no longer be read. A Dask array remains valid `DataArray` data — it is now computed on write like any other eager array, and `virtual=True` rejects it (@atrabattoni). +- **Opening a seismological file without naming an engine returns a nested collection, not a stacked array.** `xd.open(file)` on a three-component file used to return a `(3, 100)` array by guessing the traces were synchronized; it now returns the `network / station / location / channel / acquisition` tree. `xd.concat(traces, "channel")` is the one-liner back, and `engine="miniseed"` still gives the old shape. More generally, `xd.open` now combines whether it opened one file or many, so the shape it returns no longer depends on the file count (@atrabattoni). +- `DataCollection.query` raises a `KeyError` on an indexer naming no level of the collection, instead of silently returning everything unchanged. `dc.query(time=slice(0, 5))`, which used to be a no-op, now raises: use `sel` to trim inside the leaves (@atrabattoni). +- Custom engines must subclass `xdas.io.Engine`: passing a bare read function as `engine` now raises a `TypeError` (see the data-formats documentation) (@atrabattoni). ### Bug Fixes -- The miniseed `ctype` argument is now honored: it previously routed to an unused attribute and the reader always built interpolated time coordinates. The default is unchanged (@atrabattoni). +- Fix a STEIM-compressed `int32` miniSEED file being scanned as `float64`: the element type now comes from the file's encoding rather than from the empty array `headonly=True` returns (@atrabattoni). +- Fix the miniseed `ctype` argument being ignored: the reader always built interpolated time coordinates. The default is unchanged (@atrabattoni). +- Fix a data collection keyed by zero-padded codes — a SEED location such as `"00"`, say — reading back from netCDF as a sequence with its keys lost (@atrabattoni). +- Fix miniSEED scans being forced to a single process (@atrabattoni). ## 0.2.8 diff --git a/docs/user-guide/faq.md b/docs/user-guide/faq.md index b6cfb217..0aecd8c9 100644 --- a/docs/user-guide/faq.md +++ b/docs/user-guide/faq.md @@ -76,7 +76,7 @@ chunk boundaries automatically when used with {py:func}`~xdas.processing.process ## Can I use xdas with seismic data that is not DAS? Yes. The data model is generic: a {py:class}`~xdas.DataArray` can represent any -labeled N-dimensional array. The [](io/miniseed.md) page shows a complete example with +labeled N-dimensional array. The [](io/obspy.md) page shows a complete example with a large-N seismic array stored as miniSEED files. All signal processing routines in {py:mod}`xdas.signal` and {py:mod}`xdas.fft` work on any DataArray regardless of the physical quantity it represents. diff --git a/docs/user-guide/io/data-formats.md b/docs/user-guide/io/data-formats.md index 29f01c8f..0b617da0 100644 --- a/docs/user-guide/io/data-formats.md +++ b/docs/user-guide/io/data-formats.md @@ -34,13 +34,32 @@ Xdas support the following DAS formats: | SINTELA | ONYX | `"sintela"` | HDF5, tiles | `hdf5` | | Terra15 | Treble | `"terra15"` | HDF5, tiles | `hdf5` | -It also implements its own format and support ProdML and miniSEED: +It also implements its own format, supports ProdML, and reads every format +ObsPy reads — miniSEED, SAC, GSE2, SEG-2 and the rest — through a single +engine named after the library rather than after a format: | Format | `engine` argument | Virtualization | Default | |:-----------------:|:-----------------:|:-----------------:|:---------:| | Xdas | `None` | HDF5, tiles | `hdf5` | | ProdML | `"prodml"` | HDF5, tiles | `hdf5` | -| miniSEED | `"miniseed"` | tiles | `tiles` | +| ObsPy formats | `"obspy"` | tiles | `tiles` | +| miniSEED (legacy) | `"miniseed"` | tiles | `tiles` | + +The `"obspy"` engine is the only one that describes a file as a *collection* +rather than a single array: it emits one lazy data array per ObsPy `Trace`, +nested on the SEED hierarchy. See [](obspy.md). + +```{note} +`"miniseed"` is the engine `"obspy"` replaced, kept so that views written by +it keep decoding and code written against it keeps running. It describes a +whole file as one tile of stacked channels, which it classifies as +*synchronized* or *unsynchronized* and refuses anything else — a file holding +two sampling rates, for instance. Both engines read the same files, so +registration order settles auto-detection: `"obspy"` is tried first, and +`"miniseed"` is reached only through {py:func}`xdas.open_dataarray`, which +asks for the single stacked array the new engine cannot produce. Prefer +`"obspy"` in new code. +``` ```{note} A Febus file stores a stack of overlapping blocks rather than one contiguous array. The diff --git a/docs/user-guide/io/index.md b/docs/user-guide/io/index.md index f1f05fbf..2c0d1b96 100644 --- a/docs/user-guide/io/index.md +++ b/docs/user-guide/io/index.md @@ -7,5 +7,5 @@ This section covers reading and writing data with *xdas*. data-formats virtual-datasets -miniseed +obspy ``` diff --git a/docs/user-guide/io/miniseed.md b/docs/user-guide/io/miniseed.md deleted file mode 100644 index e599b10a..00000000 --- a/docs/user-guide/io/miniseed.md +++ /dev/null @@ -1,131 +0,0 @@ ---- -file_format: mystnb -kernelspec: - name: python3 ---- - -```{code-cell} -:tags: [remove-cell] - -import os -os.chdir("../../_data") - -import warnings -warnings.filterwarnings("ignore") - -import obspy -import numpy as np - -np.random.seed(0) - -network = "NX" -stations = ["SX001", "SX002", "SX003", "SX004", "SX005", "SX006", "SX007"] -location = "00" -channels = ["HHZ", "HHN", "HHE"] - -nchunk = 5 -chunk_duration = 60 -starttimes = [ - obspy.UTCDateTime("2024-01-01T00:00:00") + idx * chunk_duration - for idx in range(nchunk) -] -delta = 0.01 -failure = 0.1 - -for station in stations: - for starttime in starttimes: - if np.random.rand() < failure: - continue - for channel in channels: - data = np.random.randn(round(chunk_duration / delta)) - header = { - "delta": delta, - "starttime": starttime, - "network": network, - "station": station, - "location": location, - "channel": channel, - } - tr = obspy.Trace(data, header) - endtime = starttime + chunk_duration - dirpath = f"{network}/{station}" - if not os.path.exists(dirpath): - os.makedirs(dirpath) - fname = f"{network}.{station}.{location}.{channel}__{starttime}_{endtime}.mseed" - path = os.path.join(dirpath, fname) - tr.write(path) - -``` - -# Working with Large-N Seismic Arrays - -The virtualization capabilities of Xdas make it a good candidate for working with the large datasets produced by large-N seismic arrays. - -In this section, we will present several examples of how to handle long and numerous time series. - -```{note} -This part encourages experimenting with seismic data. Depending on the most common use cases users find, this could lead to changes in development direction. -``` - -## Exploring a dataset - -We will start by exploring a synthetic dataset composed of 7 stations, each with 3 channels (HHZ, HHN, HHE). Each trace for each station and channel is stored in multiple one-minute-long files. Some files are missing, resulting in data gaps. The sampling rate is 100 Hz. The data is organized in a directory structure that groups files by station. - -To open the dataset, we will provide a pattern that describes the directory structure and file names to the `xdas.open` function. The pattern is a string containing placeholders for the network, station, location, channel, start time, and end time. Placeholders for named fields are enclosed in curly braces, while simple brackets are used for varying parts of the file name that will be concatenated into different acquisitions (meant for changes in acquisition parameters). - -Next, we will plot the availability of the dataset. - -```{code-cell} -:tags: [remove-output] - -import xdas as xd - -pattern = "NX/{station}/NX.{station}.00.{channel}__[acquisition].mseed" -dc = xd.open(pattern, engine="miniseed") -xd.plot_availability(dc, dim="time") -``` -```{code-cell} -:tags: [remove-input] -from IPython.display import HTML -fig = xd.plot_availability(dc, dim="time") -HTML(fig.to_html()) -``` - -We can see that indeed some data is missing. Yet, as often, the different channels are synchronized. We can therefore reorganize the data by concatenating the channels of each station. - -```{code-cell} -:tags: [remove-output] - -dc = xd.DataCollection( - { - station: xd.concatenate(objs, dim="channel") - for station in dc - for objs in zip(*[dc[station][channel] for channel in dc[station]]) - }, - name="station", -) -xd.plot_availability(dc, dim="time") -``` -```{code-cell} -:tags: [remove-input] -from IPython.display import HTML -fig = xd.plot_availability(dc, dim="time") -HTML(fig.to_html()) -``` - -In our case, all stations are synchronized to GPS time. By selecting a time range where no data is missing, we can concatenate the stations to obtain an N-dimensional array representation of the dataset. - -```{code-cell} -dc = dc.sel(time=slice("2024-01-01T00:01:00", "2024-01-01T00:02:59.99")) -da = xd.concatenate((dc[station] for station in dc), dim="station") -da -``` - -This is useful for performing array analysis. In this example, we simply stack the energy. - -```{code-cell} -trace = np.square(da).mean("channel").mean("station") -trace.plot(ylim=(0, 3)) -``` - -All the processing capabilities of Xdas can be applied to the dataset. We encourage readers to explore the various possibilities. diff --git a/docs/user-guide/io/obspy.md b/docs/user-guide/io/obspy.md new file mode 100644 index 00000000..b43796fc --- /dev/null +++ b/docs/user-guide/io/obspy.md @@ -0,0 +1,194 @@ +--- +file_format: mystnb +kernelspec: + name: python3 +--- + +```{code-cell} +:tags: [remove-cell] + +import os +os.chdir("../../_data") + +import warnings +warnings.filterwarnings("ignore") + +import obspy +import numpy as np + +np.random.seed(0) + +network = "NX" +stations = ["SX001", "SX002", "SX003", "SX004", "SX005", "SX006", "SX007"] +location = "00" +channels = ["HHZ", "HHN", "HHE"] + +nchunk = 5 +chunk_duration = 60 +starttimes = [ + obspy.UTCDateTime("2024-01-01T00:00:00") + idx * chunk_duration + for idx in range(nchunk) +] +delta = 0.01 +failure = 0.1 + +for station in stations: + for starttime in starttimes: + if np.random.rand() < failure: + continue + for channel in channels: + data = np.random.randn(round(chunk_duration / delta)) + header = { + "delta": delta, + "starttime": starttime, + "network": network, + "station": station, + "location": location, + "channel": channel, + } + tr = obspy.Trace(data, header) + endtime = starttime + chunk_duration + dirpath = f"{network}/{station}" + if not os.path.exists(dirpath): + os.makedirs(dirpath) + fname = f"{network}.{station}.{location}.{channel}__{starttime}_{endtime}.mseed" + path = os.path.join(dirpath, fname) + tr.write(path) + +``` + +# Working with miniSEED and other ObsPy formats + +*Xdas* reads seismological data through ObsPy, with the engine named `"obspy"` +after the library rather than after any one format: decoding is +{py:func}`obspy.read`, so miniSEED, SAC, GSE2, SEG-2 and everything else ObsPy +supports goes through the same path. `engine="miniseed"` still works as an +alias. + +The engine mirrors {py:func}`obspy.read` exactly: **one contiguous ObsPy +`Trace` becomes one lazy `DataArray`**, and the collection mirrors the +`Stream`, nested on the four levels of the SEED hierarchy. Nothing is decoded +at this point — the scan only records where each trace lives. + +```{note} +This part encourages experimenting with seismic data. Depending on the most common use cases users find, this could lead to changes in development direction. +``` + +## Reading + +Our synthetic dataset holds 7 stations of 3 channels each, cut into +one-minute files, with some files missing. Point {py:func}`xdas.open` at them — +the directory layout does not have to be described, since the SEED identifiers +inside the files already say where each trace belongs. + +```{code-cell} +import xdas as xd + +dc = xd.open("NX/*/*.mseed", engine="obspy") +dc +``` + +Each level is named, and the leaves are `acquisition` sequences: `xdas.open` +combines what it scanned, so contiguous traces have been fused into a single +lazy array, gaps have moved *into* the time coordinate, and a new element +appears only where something genuinely changed — a different sampling rate, a +different data type. That is why the level is no longer called `trace`. + +## Selecting + +Because the levels are named, {py:meth}`~xdas.DataCollection.select` gives the +semantics of `obspy.Stream.select`, with shell-style globbing on the keys: + +```{code-cell} +dc.select(station="SX00[123]", channel="HH?") +``` + +`select` chooses *which* leaves are kept; {py:meth}`~xdas.DataCollection.sel` +trims *inside* each leaf by coordinate label. Indexing works too, and reads +like the seed id it is: + +```{code-cell} +dc["NX"]["SX001"]["00"]["HHZ"][0] +``` + +```{note} +A blank location code, which ObsPy spells `""`, becomes `"--"` — the FDSN +convention, and the only spelling that can be a group name when the collection +is written to netCDF. +``` + +## Availability + +```{code-cell} +:tags: [remove-output] + +xd.plot_availability(dc.select(channel="HHZ"), dim="time") +``` +```{code-cell} +:tags: [remove-input] +from IPython.display import HTML +fig = xd.plot_availability(dc.select(channel="HHZ"), dim="time") +HTML(fig.to_html()) +``` + +Some data is missing. The gaps are not holes in the collection — they live in +each channel's time coordinate, and {py:func}`xdas.split` recovers the original +contiguous segments exactly: + +```{code-cell} +da = dc["NX"]["SX003"]["00"]["HHZ"][0] +[part.sizes["time"] for part in xd.split(da, "gaps")] +``` + +## Overlapping data + +Files often share a sample at their seam, or an acquisition restarts slightly +before it stopped. Those overlaps are visible as backward steps of the +coordinate, and {py:func}`xdas.trim_overlaps` resolves them by dropping the +duplicated samples — never resampling, never filling, always on a sample +boundary: + +```python +dc = xd.trim_overlaps(dc) # the later data wins, as obspy's merge(method=1) +dc = xd.trim_overlaps(dc, keep="first") +``` + +It recurses over a collection, preserving the tree. If instead you want to keep +every copy and look at them, `xd.split(da, "overlaps")` cuts them apart. + +## Stacking channels and stations + +As often, the different channels of a station are synchronized. They can be +stacked into a two-dimensional array with {py:func}`xdas.concat`, which stays +lazy: the identifiers that vary along the new dimension become a coordinate, +the ones that do not stay scalar. + +```{code-cell} +def stack(node, dim): + return xd.concat([node[key][0] for key in sorted(node)], dim) + +da = stack(dc["NX"]["SX001"]["00"], "channel") +da +``` + +All stations here are synchronized to GPS time, so once a time range without +missing data is selected they can be stacked in turn into an N-dimensional +array, ready for array analysis: + +```{code-cell} +sub = dc.sel(time=slice("2024-01-01T00:01:00", "2024-01-01T00:02:59.99")) +da = xd.concat( + [stack(sub["NX"][station]["00"], "channel") for station in sorted(sub["NX"])], + "station", +) +da +``` + +In this example, we simply stack the energy. + +```{code-cell} +trace = np.square(da).mean("channel").mean("station") +trace.plot(ylim=(0, 3)) +``` + +All the processing capabilities of Xdas can be applied to the dataset. We encourage readers to explore the various possibilities. diff --git a/docs/user-guide/io/virtual-datasets.md b/docs/user-guide/io/virtual-datasets.md index 2a205161..e7477e30 100644 --- a/docs/user-guide/io/virtual-datasets.md +++ b/docs/user-guide/io/virtual-datasets.md @@ -23,7 +23,7 @@ To deal with large multi-file dataset, *Xdas* uses the concept of virtual datase Which types an engine offers is declared by its `_supported_vtypes` attribute; the first one listed is the default. See [](#choosing-a-virtualization-backend) for how to pick, and [](data-formats.md) for what each engine supports and defaults to. -A third backing, [Dask arrays](https://docs.Dask.org/en/stable/array.html), is deprecated and no longer used by any engine — see [](#dask-virtualization). +A third backing, [Dask arrays](https://docs.Dask.org/en/stable/array.html), was removed in 0.2.9: it was used by no engine, and tile virtualization describes the same mapping as plain array data instead of as a computation graph. A Dask array is still valid data for a {py:class}`xdas.DataArray`; it is simply computed on write rather than serialized. ## HDF5 Virtualization @@ -201,21 +201,3 @@ decimated reads matter. The larger the archive, the stronger the case for `tiles the only one of the two whose write, open and read costs do not all grow with the number of files. ``` - -(dask-virtualization)= -## Dask Virtualization (deprecated) - -```{deprecated} 0.2.9 -Dask virtualization is no longer used by any engine and writing it is deprecated. The -formats that once relied on it — those HDF5 virtual datasets cannot serve — now use -[tile virtualization](#tile-virtualization) instead. Existing files that store a Dask -graph can still be read, so nothing on disk is lost, but new datasets should not be -written this way. -``` - -Formats that HDF5 could not virtualize used to be loaded as Dask arrays: an -N-dimensional stack of chunks, each with a task attached that produces its values, -serialized into the native *Xdas* netCDF format as a computation graph. Tiles replace it -with a manifest that describes the same mapping as plain array data, which is both more -compact and far quicker to build, and which does not go sluggish once the graph reaches -millions of tasks. diff --git a/pyproject.toml b/pyproject.toml index 321a3ec6..9cf5df22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,12 +10,10 @@ authors = [ { name = "Alister Trabattoni", email = "alister.trabattoni@gmail.com" }, ] dependencies = [ - "dask<2025.4.0", "h5netcdf", "h5py", "hdf5plugin", "loky", - "msgpack", "numba", "numpy>=2.3", "obspy", @@ -42,7 +40,7 @@ docs = [ "sphinx-copybutton", "sphinx", ] -tests = ["dascore", "psutil", "seisbench", "torch"] +tests = ["dascore", "dask<2025.4.0", "psutil", "seisbench", "torch"] # Single source of truth for the version: xdas/__init__.py [tool.setuptools.dynamic] diff --git a/tests/coordinates/test_interp.py b/tests/coordinates/test_interp.py index 610925e9..e7adcb2e 100644 --- a/tests/coordinates/test_interp.py +++ b/tests/coordinates/test_interp.py @@ -776,6 +776,22 @@ def test_to_dataset_datetime(self): assert "time_indices" in dataset assert dataset["time_values"].dtype == np.dtype("datetime64[ns]") + def test_collect_legacy_spelling(self): + # the pre-break grammar, written without any interpolation variable + dataset = xr.Dataset( + { + "x_indices": ("x_points", np.array([0, 8])), + "x_values": ("x_points", np.array([100.0, 900.0])), + "__values__": ( + ("x",), + np.zeros(9), + {"coordinate_interpolation": "x: x_indices x_values"}, + ), + } + ) + recovered = InterpCoordinate._collect_from_dataset(dataset, "__values__") + assert np.allclose(recovered["x"].tie_values, [100.0, 900.0]) + class TestInterpCoordinateRegular: """Tests for InterpCoordinate with an enforced sampling_interval (regular mode).""" diff --git a/tests/coordinates/test_sampled.py b/tests/coordinates/test_sampled.py index d53d27f3..86f55d74 100644 --- a/tests/coordinates/test_sampled.py +++ b/tests/coordinates/test_sampled.py @@ -1,6 +1,7 @@ import tempfile import numpy as np +import numpy.testing as npt import pandas as pd import pytest @@ -813,6 +814,62 @@ def test_to_netcdf_and_back(self): result = xd.open(file.name) assert result.equals(expected) + def test_collect_legacy_spelling(self): + # the spelling that predates the CF-shaped grammar: the group named the + # coordinate, the mapping listed both tie point variables, and the + # interval was the sampling variable's own value + import xarray as xr + + dataset = xr.Dataset( + { + "time_values": ("time_points", np.array([0, 1_000_000_000])), + "time_lengths": ("time_points", np.array([100, 100])), + "time_sampling": ( + (), + 8, + { + "tie_point_mapping": "time: time_values time_lengths", + "dtype": "timedelta64[ns]", + "units": "milliseconds", + }, + ), + "__values__": ( + ("time",), + np.zeros(200), + {"coordinate_sampling": "time: time_sampling"}, + ), + } + ) + recovered = SampledCoordinate._collect_from_dataset(dataset, "__values__") + coord = recovered["time"] + assert coord.dim == "time" + assert coord.sampling_interval == np.timedelta64(8, "ms") + npt.assert_array_equal(coord.tie_lengths, [100, 100]) + + def test_collect_legacy_spelling_numeric(self): + # the same, on a numeric axis: no units/dtype attributes to decode, the + # sampling variable's value is the interval as it stands + import xarray as xr + + dataset = xr.Dataset( + { + "distance_values": ("distance_points", np.array([0.0])), + "distance_lengths": ("distance_points", np.array([30])), + "distance_sampling": ( + (), + 2.5, + {"tie_point_mapping": "distance: distance_values distance_lengths"}, + ), + "__values__": ( + ("distance",), + np.zeros(30), + {"coordinate_sampling": "distance: distance_sampling"}, + ), + } + ) + recovered = SampledCoordinate._collect_from_dataset(dataset, "__values__") + assert recovered["distance"].sampling_interval == 2.5 + class TestGetSplitIndices: def test_no_tolerance(self): diff --git a/tests/dask/test_core.py b/tests/dask/test_core.py deleted file mode 100644 index 1ecf2b3b..00000000 --- a/tests/dask/test_core.py +++ /dev/null @@ -1,80 +0,0 @@ -import dask -import numpy as np - -from xdas.dask import dumps, loads -from xdas.dask.core import from_dict, fuse, iskey, to_dict - - -class TestIsKey: - def test_valid(self): - keys = [("a", 0), ("a", 0, 1), "name-s0d9us-df63ij"] - for key in keys: - assert iskey(key) - - def test_invalid(self): - keys = ["", (sum, 0, 1), ("a",), ("a", "b")] - for key in keys: - assert not iskey(key) - - -class TestFuse: - def test_simple(self): - graph = { - "a": "b", - "b": (np.sum, [1, 2, 3]), - } - assert fuse(graph) == {"a": (np.sum, [1, 2, 3])} - - def test_recursive(self): - graph = { - "a": "b", - "b": "c", - "c": (np.sum, [1, 2, 3]), - } - assert fuse(graph) == {"a": (np.sum, [1, 2, 3])} - - def test_tuple(self): - graph = { - ("a", 0): ("b", 1), - ("b", 1): (np.sum, [1, 2, 3]), - } - assert fuse(graph) == {("a", 0): (np.sum, [1, 2, 3])} - - def test_ignore(self): - graph = { - "a": (sum, 1, 2), - "b": (sum, 3), - "c": (sum, "a", "b"), - } - assert fuse(graph) == graph - - -class TestIO: - def generate(self, tmpdir): - expected = np.random.rand(3, 10) - chunks = np.split(expected, 5, axis=1) - for idx, chunk in enumerate(chunks): - np.save(tmpdir / f"chunk_{idx}.npy", chunk) - paths = sorted(tmpdir.glob("*.npy")) - chunks = [dask.delayed(np.load)(str(path)) for path in paths] - chunks = [ - dask.array.from_delayed(chunk, shape=(3, 2), dtype=expected.dtype) - for chunk in chunks - ] - data = dask.array.concatenate(chunks, axis=1) - assert np.array_equal(data.compute(), expected) - return expected, data - - def test_dict(self, tmp_path): - expected, data = self.generate(tmp_path) - result = from_dict(to_dict(data)) - assert np.array_equal(result.compute(), expected) - sliced = result[:, 0] - assert np.array_equal(sliced.compute(), expected[:, 0]) - - def test_serial(self, tmp_path): - expected, data = self.generate(tmp_path) - result = loads(dumps(data)) - assert np.array_equal(result.compute(), expected) - sliced = result[:, 0] - assert np.array_equal(sliced.compute(), expected[:, 0]) diff --git a/tests/dask/test_serial.py b/tests/dask/test_serial.py deleted file mode 100644 index 4e96bc72..00000000 --- a/tests/dask/test_serial.py +++ /dev/null @@ -1,82 +0,0 @@ -import pytest -from dask.utils import itemgetter, methodcaller - -import xdas as xd -from xdas.dask.serial import dumps, loads - - -def test_tuple(): - objs = [ - (1, 2, 3), - [1, 2, 3], - (1, (2, 3)), - [1, [2, 3]], - (1, [2, 3]), - [1, (2, 3)], - ] - for obj in objs: - assert loads(dumps(obj)) == obj - - -def test_slice(): - objs = [ - slice(1, 2, 3), - slice(None), - ] - for obj in objs: - assert loads(dumps(obj)) == obj - - -def test_callable(): - objs = [ - dumps, - loads, - xd.DataArray, - xd.open_dataarray, - ] - for obj in objs: - assert loads(dumps(obj)) is obj - - -def test_keys(): - obj = {("a", 0, 0): ("b", 0, 0)} - assert loads(dumps(obj)) == obj - - -def test_mixed_structure(): - obj = { - "a": (1, 2, 3), - ("b", 0, 0): [1, 2, 3], - "c": (None, slice(1, 2, 3), slice(None)), - ("d", 1, 1): (dumps, "data"), - "e": (xd.DataArray, "path"), - } - assert loads(dumps(obj)) == obj - - -def test_methdocaller(): - obj = methodcaller("method") - assert loads(dumps(obj)) == obj - - -def test_itemgetter(): - obj = itemgetter(1) - assert loads(dumps(obj)) == obj - - -def test_unknown_type(): - with pytest.raises( - TypeError, match="Cannot encode object of type " - ): - dumps(object()) - - -def test_decode_unknown_code(): - import msgpack - - from xdas.dask.serial import decode - - # Call decode with an extension code not in the codes dict - data = msgpack.dumps(None) - with pytest.raises(ValueError, match="Unknown code"): - decode(99, data) diff --git a/tests/io/test_generic.py b/tests/io/test_generic.py index 28f44a95..b7194eca 100644 --- a/tests/io/test_generic.py +++ b/tests/io/test_generic.py @@ -37,6 +37,16 @@ def test_auto_engine_all_fail_raises_value_error(self): with pytest.raises(ValueError, match="no engine could open"): AutoEngine().open_dataarray("/definitely/nonexistent_file.hdf5") + def test_auto_engine_fail_message_lists_every_refusal(self, tmp_path): + # no single refusal is *the* reason, so the message reports them all: + # a file an engine recognised and then rejected must be able to + # explain itself rather than be reported as merely unreadable + fake = tmp_path / "fake.h5" + fake.write_bytes(b"not a valid hdf5 file") + with pytest.raises(ValueError, match="no engine could open") as excinfo: + AutoEngine().open_dataarray(str(fake)) + assert "\n xdas: " in str(excinfo.value) + def test_auto_engine_fail_message_includes_ctype(self, tmp_path): fake = tmp_path / "fake.h5" fake.write_bytes(b"not a valid hdf5 file") diff --git a/tests/io/test_miniseed.py b/tests/io/test_miniseed.py index 71c10f7a..d2918420 100644 --- a/tests/io/test_miniseed.py +++ b/tests/io/test_miniseed.py @@ -220,9 +220,14 @@ def test_miniseed_tile_backend(tmp_path): def test_miniseed_helpers(tmp_path): - # get_band_code with out-of-range sampling rate + # the stream converters and the band-code table moved to the obspy engine + # and are re-exported here, so imports written against this module still + # resolve + from xdas.io import obspy as obspy_engine + + assert to_stream is obspy_engine.to_stream + assert get_band_code is obspy_engine.get_band_code assert get_band_code(0.0) == "X" - assert get_band_code(6000.0) == "X" # to_stream raises on non-2D data da_3d = xd.DataArray(np.zeros((2, 3, 4)), dims=("a", "b", "c")) @@ -230,6 +235,88 @@ def test_miniseed_helpers(tmp_path): to_stream(da_3d) +def test_compressed_element_type_comes_from_the_encoding(tmp_path): + # `headonly=True` leaves `tr.data` an empty float64 array whatever the file + # holds, so a STEIM-compressed int32 file used to be scanned as float64 and + # the tile array then rejected the int32 it decoded + path = str(tmp_path / "steim.mseed") + tr = obspy.Trace(np.arange(100, dtype="int32"), make_header(1, "Z", 0)) + obspy.Stream([tr]).write(path, format="MSEED", encoding="STEIM2", reclen=512) + assert obspy.read(path)[0].stats.mseed["encoding"] == "STEIM2" + + da = xd.open_dataarray(path, engine="miniseed") + assert da.dtype == np.int32 + npt.assert_array_equal(da.values, tr.data) + + +def test_obspy_engine_is_preferred_by_auto_detection(tmp_path): + # both engines read the same files, so registration order settles which + # one auto-detection reaches first + from xdas.io import Engine + + names = list(Engine._registry) + assert names.index("obspy") < names.index("miniseed") + + make_network(tmp_path, samples=100) + path = min(tmp_path.glob("*00.mseed")) + # `open` asks for a collection first, which only the obspy engine provides + dc = xd.open(path) + assert dc.fields == ("network", "station", "location", "channel", "acquisition") + # `open_dataarray` asks for a single array, which the obspy engine cannot + # give for a three-component file; the legacy engine stacks it + da = xd.open_dataarray(path) + assert da.shape == (3, 100) + assert da.data.engine["name"] == "miniseed" + + +def test_pre_rename_manifest_still_decodes(tmp_path): + from xdas.virtual import TileArray + + make_network(tmp_path, samples=100) + path = str(min(tmp_path.glob("*00.mseed"))) + expected = np.array(obspy.read(path)) + + data = TileArray.from_tiles( + path, + (3, 100), + np.dtype("float64"), + {"name": "miniseed", "method": "synchronized", "ignore_last_sample": False}, + ) + npt.assert_allclose(np.asarray(data), expected) + + trimmed = TileArray.from_tiles( + path, + (3, 99), + np.dtype("float64"), + {"name": "miniseed", "method": "synchronized", "ignore_last_sample": True}, + ) + npt.assert_allclose(np.asarray(trimmed), expected[:, :-1]) + + +def test_pre_rename_unsynchronized_manifest(tmp_path): + from xdas.virtual import TileArray + + # one channel cut in two segments: the old engine folded the channel axis + # out of the scanned shape and could drop each segment's last sample + path = tmp_path / "gap.mseed" + st = obspy.Stream( + [ + obspy.Trace(np.arange(50.0), make_header(1, "Z", 0)), + obspy.Trace(np.arange(100.0, 150.0), make_header(1, "Z", 100)), + ] + ) + st.write(str(path), format="MSEED") + expected = np.concatenate([tr.data for tr in obspy.read(str(path))]) + + data = TileArray.from_tiles( + str(path), + (99,), + np.dtype("float64"), + {"name": "miniseed", "method": "unsynchronized", "ignore_last_sample": True}, + ) + npt.assert_allclose(np.asarray(data), expected[:-1]) + + def test_miniseed_unsynchronized_traces(tmp_path): path = tmp_path / "unsync.mseed" st = obspy.Stream() diff --git a/tests/io/test_obspy.py b/tests/io/test_obspy.py new file mode 100644 index 00000000..c2892d73 --- /dev/null +++ b/tests/io/test_obspy.py @@ -0,0 +1,439 @@ +import numpy as np +import numpy.testing as npt +import obspy +import pytest + +import xdas as xd +from xdas.coordinates import Coordinate +from xdas.io.obspy import ObsPyEngine, get_band_code, to_stream +from xdas.virtual import TileArray + + +def header(station="CH001", channel="HHZ", starttime=0.0, delta=0.01, location="00"): + return { + "delta": delta, + "starttime": obspy.UTCDateTime(starttime), + "network": "DX", + "station": station, + "location": location, + "channel": channel, + } + + +def write(path, traces, **kwargs): + st = obspy.Stream( + [obspy.Trace(np.asarray(data, dtype=np.float64), head) for head, data in traces] + ) + st.write(str(path), format="MSEED", **kwargs) + return st + + +class TestScan: + def test_one_trace_per_obspy_trace(self, tmp_path): + path = tmp_path / "three.mseed" + write( + path, + [ + (header(channel=f"HH{component}"), np.random.rand(100)) + for component in "ZNE" + ], + ) + dc = ObsPyEngine().open_datacollection(path) + assert dc.fields == ("network", "station", "location", "channel", "trace") + assert list(dc) == ["DX"] + assert list(dc["DX"]) == ["CH001"] + assert list(dc["DX"]["CH001"]) == ["00"] + assert sorted(dc["DX"]["CH001"]["00"]) == ["HHE", "HHN", "HHZ"] + da = dc["DX"]["CH001"]["00"]["HHZ"][0] + assert da.dims == ("time",) + assert da.shape == (100,) + assert da["network"].values == "DX" + assert da["station"].values == "CH001" + assert da["location"].values == "00" + assert da["channel"].values == "HHZ" + assert da["time"][0].values == np.datetime64("1970-01-01T00:00:00") + assert da["time"][-1].values == np.datetime64("1970-01-01T00:00:00.990") + + def test_gaps_become_separate_traces(self, tmp_path): + path = tmp_path / "gap.mseed" + write( + path, + [ + (header(), np.random.rand(50)), + (header(starttime=1.0), np.random.rand(40)), + ], + ) + dc = ObsPyEngine().open_datacollection(path) + traces = dc["DX"]["CH001"]["00"]["HHZ"] + assert len(traces) == 2 + assert [da.sizes["time"] for da in traces] == [50, 40] + + def test_mixed_sampling_rates_are_read(self, tmp_path): + # the old engine refused these outright + path = tmp_path / "mixed.mseed" + write( + path, + [ + (header(channel="HHZ", delta=0.01), np.random.rand(100)), + (header(channel="LHZ", delta=1.0), np.random.rand(10)), + ], + ) + dc = ObsPyEngine().open_datacollection(path) + channels = dc["DX"]["CH001"]["00"] + assert sorted(channels) == ["HHZ", "LHZ"] + assert channels["HHZ"][0].sizes["time"] == 100 + assert channels["LHZ"][0].sizes["time"] == 10 + + def test_duplicated_id_groups_into_one_sequence(self, tmp_path): + path = tmp_path / "dup.mseed" + write( + path, + [ + (header(), np.random.rand(50)), + (header(starttime=10.0), np.random.rand(50)), + (header(starttime=20.0), np.random.rand(50)), + ], + ) + dc = ObsPyEngine().open_datacollection(path) + assert len(dc["DX"]["CH001"]["00"]["HHZ"]) == 3 + + def test_ctype_drives_the_time_coordinate(self, tmp_path): + path = tmp_path / "one.mseed" + write(path, [(header(), np.random.rand(100))]) + dc = ObsPyEngine(ctype="dense").open_datacollection(path) + da = dc["DX"]["CH001"]["00"]["HHZ"][0] + assert isinstance(da["time"], Coordinate["dense"]) + + def test_traces_sharing_everything_are_refused(self, tmp_path): + path = tmp_path / "twins.mseed" + write(path, [(header(), np.zeros(50)), (header(), np.ones(50))]) + with pytest.raises(ValueError, match="nothing content-free separates"): + ObsPyEngine().open_datacollection(path) + + def test_dtype_comes_from_the_encoding(self, tmp_path): + # headonly leaves `tr.data` an empty float64 array whatever the file + # holds, so a STEIM-compressed file would scan with the wrong dtype + path = tmp_path / "steim.mseed" + st = obspy.Stream([obspy.Trace(np.arange(100, dtype=np.int32), header())]) + st.write(str(path), format="MSEED", encoding="STEIM2") + dc = ObsPyEngine().open_datacollection(path) + da = dc["DX"]["CH001"]["00"]["HHZ"][0] + assert da.dtype == np.int32 + npt.assert_array_equal(da.values, np.arange(100)) + + def test_open_dataarray_needs_a_single_trace(self, tmp_path): + path = tmp_path / "one.mseed" + write(path, [(header(), np.random.rand(50))]) + da = ObsPyEngine().open_dataarray(path) + assert da.dims == ("time",) + assert da.shape == (50,) + + path = tmp_path / "two.mseed" + write( + path, + [ + (header(channel="HHZ"), np.zeros(50)), + (header(channel="HHN"), np.ones(50)), + ], + ) + with pytest.raises(ValueError, match="holds 2 traces"): + ObsPyEngine().open_dataarray(path) + + +class TestBlankLocation: + def test_round_trips_through_netcdf_and_to_stream(self, tmp_path): + path = tmp_path / "blank.mseed" + write(path, [(header(location=""), np.random.rand(50))]) + dc = ObsPyEngine().open_datacollection(path) + # "" cannot be a netCDF group name; "--" is the FDSN convention + assert list(dc["DX"]["CH001"]) == ["--"] + da = dc["DX"]["CH001"]["--"]["HHZ"][0] + assert da["location"].values == "--" + npt.assert_allclose(da.values, obspy.read(str(path))[0].data) + + dc.to_netcdf(tmp_path / "blank.nc") + reopened = xd.open_datacollection(tmp_path / "blank.nc") + assert list(reopened["DX"]["CH001"]) == ["--"] + + stacked = xd.DataArray( + da.values[None], + {"space": [0.0], "time": da["time"]}, + ) + dim = {"space": "time"} + assert to_stream(stacked, location="--", dim=dim)[0].stats.location == "" + assert to_stream(stacked, location="00", dim=dim)[0].stats.location == "00" + + +class TestLoadTile: + def test_pointer_survives_resegmentation(self, tmp_path): + # the same samples written with different record boundaries: the trace + # count and the segmentation differ, the pointer must not + data = np.random.rand(400) + coarse = tmp_path / "coarse.mseed" + write(coarse, [(header(), data)], reclen=4096) + fine = tmp_path / "fine.mseed" + write(fine, [(header(), data)], reclen=512) + assert len(obspy.read(str(coarse))) == len(obspy.read(str(fine))) + for path in (coarse, fine): + da = ObsPyEngine().open_dataarray(path) + npt.assert_allclose(da.values, data) + npt.assert_allclose(da.isel(time=slice(10, 30)).values, data[10:30]) + + def test_split_records_still_resolve(self, tmp_path): + # two files holding the same span, one written as a single trace and + # one as two abutting traces that `join_contiguous` must fuse back + data = np.random.rand(200) + whole = tmp_path / "whole.mseed" + write(whole, [(header(), data)]) + split = tmp_path / "split.mseed" + write( + split, + [(header(), data[:120]), (header(starttime=1.20), data[120:])], + ) + # the split file scans as one trace: obspy already rejoins abutting + # records of the same channel + da = ObsPyEngine().open_dataarray(split) + npt.assert_allclose(da.values, data) + assert ObsPyEngine().open_dataarray(whole).equals(da) + + def test_traces_split_by_data_quality_are_rejoined(self, tmp_path): + # libmseed hands back one trace per data quality flag; the two are + # sample-exact contiguous, so each pointer must still resolve against + # the joined run + data = np.arange(200.0) + parts = [] + for index, (values, start, quality) in enumerate( + [(data[:120], 0.0, "D"), (data[120:], 1.20, "R")] + ): + tr = obspy.Trace(values, header(starttime=start)) + tr.stats.mseed = {"dataquality": quality} + path = tmp_path / f"part_{index}.mseed" + obspy.Stream([tr]).write(str(path), format="MSEED") + parts.append(path) + path = tmp_path / "quality.mseed" + path.write_bytes(b"".join(part.read_bytes() for part in parts)) + assert len(obspy.read(str(path))) == 2 + + dc = ObsPyEngine().open_datacollection(path) + traces = dc["DX"]["CH001"]["00"]["HHZ"] + assert [da.sizes["time"] for da in traces] == [120, 80] + npt.assert_allclose(traces[0].values, data[:120]) + npt.assert_allclose(traces[1].values, data[120:]) + + def test_legacy_unsynchronized_manifest_trimmed_and_squeezed(self, tmp_path): + # the old engine folded a single-channel axis out of the scanned shape + # and could drop the last sample of each segment + path = tmp_path / "gap.mseed" + st = write( + path, + [ + (header(), np.arange(50.0)), + (header(starttime=1.0), np.arange(100.0, 150.0)), + ], + ) + data = TileArray.from_tiles( + str(path), + (99,), + np.dtype("float64"), + { + "name": "miniseed", + "method": "unsynchronized", + "ignore_last_sample": True, + }, + ) + expected = np.concatenate([tr.data for tr in st]) + npt.assert_allclose(np.asarray(data), np.delete(expected, -1)) + + def test_same_start_different_length(self, tmp_path): + path = tmp_path / "twins.mseed" + short = np.arange(10.0) + long = np.arange(100.0, 130.0) + write(path, [(header(), short), (header(), long)]) + dc = ObsPyEngine().open_datacollection(path) + traces = dc["DX"]["CH001"]["00"]["HHZ"] + assert len(traces) == 2 + assert sorted(da.sizes["time"] for da in traces) == [10, 30] + for da in traces: + expected = short if da.sizes["time"] == 10 else long + npt.assert_allclose(da.values, expected) + + def test_missing_run_raises(self, tmp_path): + path = tmp_path / "one.mseed" + write(path, [(header(), np.random.rand(50))]) + da = ObsPyEngine().open_dataarray(path) + # point at a channel the file does not hold + with pytest.raises(ValueError, match="0 contiguous runs"): + ObsPyEngine.load_tile( + str(path), + (slice(0, 50),), + network="DX", + station="CH001", + location="00", + channel="HHN", + starttime=int(da["time"][0].values.astype("datetime64[ns]").view("i8")), + endtime=int(da["time"][-1].values.astype("datetime64[ns]").view("i8")), + ) + + +class TestLegacy: + def test_the_legacy_engine_is_separate(self): + from xdas.io import Engine + from xdas.io.miniseed import MiniSEEDEngine + + # `engine="miniseed"` names the engine this one replaced, kept for the + # views it wrote; see tests/io/test_miniseed.py + assert Engine["miniseed"] is MiniSEEDEngine + assert Engine["obspy"] is ObsPyEngine + + +class TestOpenRouting: + def make_network(self, dirpath, stations=3, channels="ZNE", chunks=2): + samples = 50 + for index in range(1, stations + 1): + for chunk in range(chunks): + traces = [ + ( + header( + station=f"CH{index:03d}", + channel=f"HH{component}", + starttime=chunk * samples * 0.01, + ), + np.random.rand(samples), + ) + for component in channels + ] + write(dirpath / f"CH{index:03d}_{chunk}.mseed", traces) + + def test_one_file_glob_and_list_agree(self, tmp_path): + self.make_network(tmp_path, stations=1, chunks=1) + paths = sorted(str(path) for path in tmp_path.glob("*.mseed")) + one = xd.open(paths[0], engine="obspy") + glob = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + listed = xd.open(paths, engine="obspy") + for dc in (one, glob, listed): + assert dc.fields == ( + "network", + "station", + "location", + "channel", + "acquisition", + ) + assert glob.equals(listed) + assert one.equals(listed) + + def test_contiguous_files_fuse_into_one_array(self, tmp_path): + self.make_network(tmp_path, stations=1, channels="Z", chunks=3) + dc = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + sequence = dc["DX"]["CH001"]["00"]["HHZ"] + assert len(sequence) == 1 + da = sequence[0] + assert da.sizes["time"] == 150 + assert isinstance(da.data, TileArray) + + def test_gaps_live_in_the_coordinate(self, tmp_path): + path = tmp_path / "gap.mseed" + write( + path, + [ + (header(), np.random.rand(50)), + (header(starttime=10.0), np.random.rand(40)), + ], + ) + dc = xd.open(path, engine="obspy") + da = dc["DX"]["CH001"]["00"]["HHZ"][0] + assert da.sizes["time"] == 90 + parts = xd.split(da, "gaps") + assert [part.sizes["time"] for part in parts] == [50, 40] + + def test_rate_change_arrives_as_two_elements(self, tmp_path): + for index, delta in enumerate([0.01, 0.02]): + write( + tmp_path / f"chunk_{index}.mseed", + [(header(starttime=index * 10.0, delta=delta), np.random.rand(50))], + ) + dc = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + sequence = dc["DX"]["CH001"]["00"]["HHZ"] + assert len(sequence) == 2 + assert sequence.name == "acquisition" + + def test_select_globs_like_obspy(self, tmp_path): + self.make_network(tmp_path, stations=3, chunks=1) + dc = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + result = dc.select(station="CH00[12]", channel="HH?") + assert sorted(result["DX"]) == ["CH001", "CH002"] + assert sorted(result["DX"]["CH001"]["00"]) == ["HHE", "HHN", "HHZ"] + assert dc.query(station="CH001").equals(dc.select(station="CH001")) + + def test_concat_along_channel_folds_the_other_columns(self, tmp_path): + self.make_network(tmp_path, stations=1, chunks=1) + dc = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + channels = dc["DX"]["CH001"]["00"] + da = xd.concat([channels[key][0] for key in sorted(channels)], "channel") + assert da.dims == ("channel", "time") + assert da.shape == (3, 50) + assert da["channel"].values.tolist() == ["HHE", "HHN", "HHZ"] + assert isinstance(da.data, TileArray) + # only the varying column unfolds; the other three stay 0-d + manifest = da.data.to_dataset() + assert manifest["channel"].ndim == 1 + for column in ("network", "station", "location"): + assert manifest[column].ndim == 0 + npt.assert_allclose( + da.values, + np.stack([channels[key][0].values for key in sorted(channels)]), + ) + + def test_netcdf_round_trip_of_a_per_trace_view(self, tmp_path): + self.make_network(tmp_path, stations=1, chunks=2) + dc = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + dc.to_netcdf(tmp_path / "view.nc") + reopened = xd.open_datacollection(tmp_path / "view.nc") + da = reopened["DX"]["CH001"]["00"]["HHZ"][0] + assert isinstance(da.data, TileArray) + npt.assert_allclose(da.values, dc["DX"]["CH001"]["00"]["HHZ"][0].values) + + def test_auto_detection(self, tmp_path): + # naming the engine is not required: `AutoEngine` reaches + # `open_datacollection` too, so the shape does not depend on it + self.make_network(tmp_path, stations=1, chunks=2) + named = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + assert xd.open(str(tmp_path / "*.mseed")).equals(named) + one = xd.open(str(tmp_path / "CH001_0.mseed")) + assert one.fields == named.fields + + def test_parallel_scan(self, tmp_path): + self.make_network(tmp_path, stations=2, chunks=2) + serial = xd.open(str(tmp_path / "*.mseed"), engine="obspy", parallel=1) + parallel = xd.open(str(tmp_path / "*.mseed"), engine="obspy", parallel=2) + assert parallel.equals(serial) + + def test_sac_opens_through_the_same_path(self, tmp_path): + path = tmp_path / "trace.sac" + data = np.arange(50, dtype=np.float32) + obspy.Trace(data, header()).write(str(path), format="SAC") + dc = xd.open(path, engine="obspy") + da = dc["DX"]["CH001"]["00"]["HHZ"][0] + assert da.sizes["time"] == 50 + assert da.dtype == np.float32 + npt.assert_allclose(da.values, data) + + +class TestHelpers: + def test_get_band_code_out_of_range(self): + assert get_band_code(0.0) == "X" + assert get_band_code(6000.0) == "X" + + def test_to_stream_requires_2d(self): + da = xd.DataArray(np.zeros((2, 3, 4)), dims=("a", "b", "c")) + with pytest.raises(ValueError, match="2D"): + to_stream(da) + + def test_stream_round_trip(self, tmp_path): + TestOpenRouting().make_network(tmp_path, stations=1, chunks=1) + dc = xd.open(str(tmp_path / "*.mseed"), engine="obspy") + channels = dc["DX"]["CH001"]["00"] + da = xd.concat([channels[key][0] for key in sorted(channels)], "channel") + st = da.to_stream(dim={"channel": "time"}) + assert len(st) == 3 + result = xd.DataArray.from_stream(st) + npt.assert_allclose(result.values, da.values) diff --git a/tests/io/test_tiles_vtype.py b/tests/io/test_tiles_vtype.py index f6e9c3ff..86a0cdf4 100644 --- a/tests/io/test_tiles_vtype.py +++ b/tests/io/test_tiles_vtype.py @@ -286,6 +286,7 @@ def test_default_vtypes(): "asn": "hdf5", "febus": "tiles", "miniseed": "tiles", + "obspy": "tiles", "prodml": "hdf5", "silixa": "tiles", "terra15": "hdf5", diff --git a/tests/io/test_xdas_io.py b/tests/io/test_xdas_io.py index ce0d8888..68381eb0 100644 --- a/tests/io/test_xdas_io.py +++ b/tests/io/test_xdas_io.py @@ -172,6 +172,49 @@ def test_create_dirs_no_dirname(self, tmp_path): os.chdir(orig) +class TestSaveDatamappingEmpty: + def test_empty_mapping_writes_nothing(self, tmp_path): + path = str(tmp_path / "empty.nc") + save_datamapping(DataMapping({}), path) + assert not os.path.exists(path) + + +class TestOpenDatamappingOnVariable: + def test_variable_group_raises(self, tmp_path): + da = make_da() + dc = xd.DataCollection({"a": da}) + path = str(tmp_path / "dc.nc") + dc.to_netcdf(path) + with pytest.raises(ValueError, match="data array as a data collection"): + open_datacollection(path, group="collection/a/time_indices") + + +class TestNestedCollections: + def test_nested_mapping_round_trip(self, tmp_path): + da = make_da() + dc = xd.DataCollection({"n1": {"a": da, "b": da}, "n2": {"c": da}}) + path = str(tmp_path / "nested.nc") + dc.to_netcdf(path) + result = open_datacollection(path) + assert result.equals(dc) + + def test_nested_non_sequential_integer_keys_stay_mapping(self, tmp_path): + da = make_da() + dc = xd.DataCollection({"n1": {"2": da, "5": da}}) + path = str(tmp_path / "nested.nc") + dc.to_netcdf(path) + result = open_datacollection(path) + assert result.equals(dc) + + def test_open_with_explicit_group(self, tmp_path): + da = make_da() + dc = xd.DataCollection({"a": da, "b": da}) + path = str(tmp_path / "dc.nc") + dc.to_netcdf(path) + result = open_datacollection(path, group="collection") + assert result.equals(dc) + + class TestOpenSaveDatasequence: def test_open_datasequence(self, tmp_path): da = make_da() diff --git a/tests/test_dataarray.py b/tests/test_dataarray.py index 9cebcf01..559cd7f9 100644 --- a/tests/test_dataarray.py +++ b/tests/test_dataarray.py @@ -526,7 +526,8 @@ def test_io_with_zfp_compression(self, tmp_path): _da = xd.DataArray.from_netcdf(tmpfile_compressed) assert np.abs(da - _da).max().values < 0.001 - def test_io_dask(self, tmp_path): + def test_io_dask_writes_eagerly(self, tmp_path): + """A dask-backed array is computed on write: no graph is stored.""" values = np.random.rand(3, 10) chunks = np.split(values, 5, axis=1) for idx, chunk in enumerate(chunks): @@ -547,13 +548,15 @@ def test_io_dask(self, tmp_path): fname = tmp_path / "tmp.nc" expected.to_netcdf(fname) result = xd.open_dataarray(fname) - assert isinstance(result.data, dask.array.Array) + assert not isinstance(result.data, dask.array.Array) assert np.array_equal(expected.values, result.values) assert expected.dtype == result.dtype assert expected.coords.equals(result.coords) assert expected.dims == result.dims assert expected.name == result.name assert expected.attrs == result.attrs + with pytest.raises(ValueError, match="virtual array as data"): + expected.to_netcdf(tmp_path / "virtual.nc", virtual=True) def test_io_non_dimensional(self, tmp_path): expected = xd.DataArray(coords={"dim": 0}, dims=()) diff --git a/tests/test_datacollection.py b/tests/test_datacollection.py index 1317a0ce..c34d6bec 100644 --- a/tests/test_datacollection.py +++ b/tests/test_datacollection.py @@ -125,14 +125,69 @@ def test_query(self): assert result.equals(expected) result = dc.query(instrument="das*") assert result.equals(dc) + # an indexer applies wherever its level sits, not only at the root: + # das2 holds three acquisitions and keeps the first two result = dc.query(acquisition=slice(0, 2)) - assert result.equals(dc) + assert [len(result[key]) for key in result] == [2, 2] + assert result["das1"].equals(dc["das1"]) def test_fields(self): da = xd.testing.dummy() dc = self.nest(da) assert dc.fields == ("instrument", "acquisition") + def test_fields_recursive(self): + da = xd.testing.dummy() + dc = xd.DataCollection( + { + "DX": xd.DataCollection( + { + "CH001": xd.DataCollection( + { + "00": xd.DataCollection( + {"HHZ": xd.DataCollection([da], "acquisition")}, + "channel", + ) + }, + "location", + ) + }, + "station", + ) + }, + "network", + ) + assert dc.fields == ( + "network", + "station", + "location", + "channel", + "acquisition", + ) + + def test_query_is_strict(self): + da = xd.testing.dummy() + dc = self.nest(da) + with pytest.raises(KeyError, match="do not name any level"): + dc.query(nonexistent="das1") + # a dimension name is not a level name: `sel` trims inside leaves, + # `query` chooses leaves + with pytest.raises(KeyError, match="do not name any level"): + dc.query(time=slice(0, 5)) + + def test_select_is_query(self): + da = xd.testing.dummy() + dc = self.nest(da) + assert dc.select(instrument="das1").equals(dc.query(instrument="das1")) + assert dc.select({"instrument": "das1"}).equals(dc.query(instrument="das1")) + + def test_query_does_not_mutate_indexers(self): + da = xd.testing.dummy() + dc = self.nest(da) + indexers = {"instrument": "das1"} + dc.query(indexers, acquisition=0) + assert indexers == {"instrument": "das1"} + def test_map(self): da = xd.testing.dummy() dc = self.nest(da) @@ -278,7 +333,7 @@ def test_sequence_from_netcdf(self, tmp_path): def test_query_invalid_key_in_sequence(self): da = xd.testing.dummy() dc = xd.DataCollection([da, da], "seq") - with pytest.raises(ValueError, match="query must be a string"): + with pytest.raises(ValueError, match="query must be an integer or a slice"): dc.query(seq="bad_string_key") def test_query_invalid_key_in_mapping(self): @@ -299,6 +354,20 @@ def test_from_netcdf_non_sequential_int_keys(self, tmp_path): # Keys 0 and 2 are not a sequential range → returns as-is DataMapping assert isinstance(result, xd.DataCollection) + def test_zero_padded_keys_survive_every_read_path(self, tmp_path): + # a SEED location such as "00" is a mapping key, not a position: every + # way in must compare the canonical decimal spelling, not parse ints + da = xd.testing.dummy() + dc = xd.DataCollection({"00": da, "01": da}, "location") + path = tmp_path / "padded.nc" + dc.to_netcdf(path) + for result in ( + xd.open_datacollection(path), + xd.DataCollection.from_netcdf(path), + xd.open_datacollection(path, engine="xdas"), + ): + assert list(result) == ["00", "01"] + def test_sequence_from_netcdf_direct(self, tmp_path): from xdas.core.datacollection import DataSequence diff --git a/tests/test_routines.py b/tests/test_routines.py index 8d46f041..9afca094 100644 --- a/tests/test_routines.py +++ b/tests/test_routines.py @@ -1,4 +1,6 @@ import numpy as np +import numpy.testing as npt +import obspy import pytest import xdas as xd @@ -507,6 +509,22 @@ def test_invalid_engine_type_raises(self, tmp_path): xd.open_dataarray(path, engine=42) +class TestOpenDatacollection: + def test_engine_arguments_need_an_engine(self, tmp_path): + da = xd.testing.dummy(shape=(10, 5)) + path = str(tmp_path / "dc.nc") + xd.DataCollection([da, da]).to_netcdf(path) + with pytest.raises(ValueError, match="require naming an engine"): + xd.open_datacollection(path, vtype="tiles") + + def test_group_is_a_native_parameter(self, tmp_path): + da = xd.testing.dummy(shape=(10, 5)) + path = str(tmp_path / "dc.nc") + xd.DataCollection([da, da]).to_netcdf(path) + with pytest.raises(ValueError, match="native-format parameter"): + xd.open_datacollection(path, group="whatever", engine="xdas") + + class TestOpenMFDatacollectionEdgeCases: def test_nonexistent_path_in_list_raises(self, tmp_path): with pytest.raises(FileNotFoundError): @@ -665,6 +683,76 @@ def test_mixed_empty_and_nonempty_uses_nonempty(self): assert result.equals(da) +class TestConcatNewDim: + def trace(self, channel, station="CH001", values=None): + time = {"tie_indices": [0, 4], "tie_values": [0.0, 4.0]} + if values is None: + values = np.arange(5.0) + return xd.DataArray( + values, + { + "network": (None, "DX"), + "station": (None, station), + "channel": (None, channel), + "time": time, + }, + ) + + def test_varying_scalar_is_promoted(self): + objs = [self.trace(channel) for channel in ("HHZ", "HHN", "HHE")] + da = xd.concat(objs, "channel") + assert da.dims == ("channel", "time") + assert da.shape == (3, 5) + # `channel` is the concat coordinate: `expand_dims` promotes it and + # `concat_coords` sorts it + assert sorted(da["channel"].values.tolist()) == ["HHE", "HHN", "HHZ"] + # the constant scalars stay scalar + assert da["network"].dim is None + assert da["station"].dim is None + + def test_other_varying_scalar_is_promoted_along_the_new_dim(self): + objs = [self.trace("HHZ", station=f"CH{idx:03d}") for idx in (1, 2, 3)] + da = xd.concat(objs, "component") + assert da.dims == ("component", "time") + assert da["station"].dim == "component" + assert da["station"].values.tolist() == ["CH001", "CH002", "CH003"] + assert da["network"].dim is None + + def test_promotion_follows_the_concat_order(self): + # `concat` sorts by the concat coordinate; a promoted scalar must be + # gathered in that same order, not in input order + objs = [ + self.trace(channel, station=station) + for channel, station in [("HHZ", "C"), ("HHE", "A"), ("HHN", "B")] + ] + da = xd.concat(objs, "channel") + assert da["channel"].values.tolist() == ["HHE", "HHN", "HHZ"] + assert da["station"].values.tolist() == ["A", "B", "C"] + + def test_unequal_non_scalar_coord_raises(self): + da1 = self.trace("HHZ") + da2 = self.trace("HHN") + da2["time"] = {"tie_indices": [0, 4], "tie_values": [10.0, 14.0]} + with pytest.raises(ValueError, match="'time' differs"): + xd.concat([da1, da2], "channel") + + def test_missing_coord_raises(self): + da1 = self.trace("HHZ") + da2 = self.trace("HHN").drop_coords("network") + with pytest.raises(ValueError, match="must share their coordinates"): + xd.concat([da1, da2], "channel") + + def test_concat_along_existing_dim_is_unchanged(self): + # the promotion machinery only runs when a new dimension is opened + da1 = self.trace("HHZ") + da2 = self.trace("HHN") + da2["time"] = {"tie_indices": [0, 4], "tie_values": [5.0, 9.0]} + da = xd.concat([da1, da2], "time") + assert da.dims == ("time",) + assert da.shape == (10,) + assert da["channel"].values == "HHZ" + + class TestConcatCoordsEdgeCases: def test_tolerance_with_dense_coord_is_noop(self): # Dense coordinates now implement a (degenerate) `simplify`, so passing a @@ -746,6 +834,220 @@ def test_invalid_type_raises(self): _get_timeline_dataframe("not_valid") +class TestTrimOverlaps: + delta = 0.01 + + def segments(self, spans, dtype=float): + """Build ``(obspy.Stream, DataArray)`` from ``(start_second, values)`` pairs.""" + st = obspy.Stream() + objs = [] + for start, values in spans: + values = np.asarray(values, dtype=dtype) + st.append( + obspy.Trace( + values.copy(), + { + "delta": self.delta, + "starttime": obspy.UTCDateTime(start), + "network": "DX", + "station": "CH001", + "location": "00", + "channel": "HHZ", + }, + ) + ) + t0 = np.datetime64(round(start * 1e9), "ns") + dt = np.timedelta64(round(self.delta * 1e9), "ns") + objs.append( + xd.DataArray( + values, + { + "time": { + "tie_indices": [0, len(values) - 1], + "tie_values": [t0, t0 + (len(values) - 1) * dt], + } + }, + ) + ) + return st, xd.concat(objs, "time", tolerance=False) + + def test_keep_last_matches_obspy_merge(self): + for spans in [ + [(0.0, np.arange(5.0)), (0.03, np.arange(100.0, 105.0))], + [(0.0, np.arange(5.0)), (0.04, np.arange(100.0, 105.0))], + [ + (0.0, np.arange(5.0)), + (0.03, np.arange(100.0, 105.0)), + (0.06, np.arange(200.0, 205.0)), + ], + ]: + st, da = self.segments(spans) + st.merge(method=1, interpolation_samples=0) + result = xd.trim_overlaps(da) + npt.assert_array_equal(result.values, np.asarray(st[0].data)) + assert result["time"][0].values == np.datetime64( + str(st[0].stats.starttime.datetime), "ns" + ) + + def test_keep_first_is_the_mirror(self): + # the later segment's head goes instead of the earlier one's tail + _, da = self.segments([(0.0, np.arange(5.0)), (0.03, np.arange(100.0, 105.0))]) + npt.assert_array_equal( + xd.trim_overlaps(da, keep="first").values, + [0.0, 1.0, 2.0, 3.0, 4.0, 102.0, 103.0, 104.0], + ) + npt.assert_array_equal( + xd.trim_overlaps(da, keep="last").values, + [0.0, 1.0, 2.0, 100.0, 101.0, 102.0, 103.0, 104.0], + ) + + def test_replaces_ignore_last_sample(self): + # the old flag dropped the last sample of every segment; the shared + # sample only, and only where it is genuinely shared, is enough + _, da = self.segments([(0.0, np.arange(5.0)), (0.04, np.arange(100.0, 105.0))]) + result = xd.trim_overlaps(da) + npt.assert_array_equal( + result.values, [0.0, 1.0, 2.0, 3.0, 100.0, 101.0, 102.0, 103.0, 104.0] + ) + # a clean seam is left untouched + _, clean = self.segments( + [(0.0, np.arange(5.0)), (0.05, np.arange(100.0, 105.0))] + ) + assert xd.trim_overlaps(clean).equals(clean) + + def test_no_overlap_is_a_noop(self): + da = xd.testing.dummy(dims=("time",), shape=(10,), step=0.01) + assert xd.trim_overlaps(da).equals(da) + + def test_wholly_covered_part_is_dropped(self): + # the middle segment is entirely covered by the last one; the first + # must still be trimmed against that last one, not against the middle + _, da = self.segments( + [ + (0.0, np.arange(10.0)), + (0.05, np.arange(100.0, 103.0)), + (0.04, np.arange(200.0, 210.0)), + ] + ) + result = xd.trim_overlaps(da) + # sorted by start: [0.00-0.09], [0.04-0.13], [0.05-0.07]; keeping the + # last, the 0.04-0.13 segment survives only outside 0.05-0.07 + npt.assert_array_equal( + result.values, + # 0.00-0.03 from the first, 0.04 from the third, 0.05-0.07 from the + # second, 0.08-0.13 from the third again + [0.0, 1.0, 2.0, 3.0, 200.0, 100.0, 101.0, 102.0] + + [204.0, 205.0, 206.0, 207.0, 208.0, 209.0], + ) + assert result["time"].get_split_indices("overlaps").size == 0 + + def test_enveloped_part_keeps_both_sides(self): + # a short high-precedence segment inside a long one: the long one must + # keep a run on each side of it, not lose everything past the overlap + _, da = self.segments([(0.0, np.arange(20.0)), (0.05, np.arange(100.0, 103.0))]) + result = xd.trim_overlaps(da) + expected = np.concatenate( + [np.arange(5.0), np.arange(100.0, 103.0), np.arange(8.0, 20.0)] + ) + npt.assert_array_equal(result.values, expected) + assert result.sizes["time"] == 20 + + def test_chain_of_three_mutual_overlaps(self): + _, da = self.segments( + [ + (0.0, np.arange(10.0)), + (0.05, np.arange(100.0, 110.0)), + (0.10, np.arange(200.0, 210.0)), + ] + ) + result = xd.trim_overlaps(da) + npt.assert_array_equal( + result.values, + np.concatenate( + [np.arange(5.0), np.arange(100.0, 105.0), np.arange(200.0, 210.0)] + ), + ) + assert result["time"].get_split_indices("overlaps").size == 0 + + def test_sub_tolerance_jitter_is_not_trimmed(self): + t0 = np.datetime64("2024-01-01T00:00:00.000000000") + dt = np.timedelta64(10_000_000, "ns") + # the second segment starts one microsecond early: jitter, not an overlap + coord = { + "tie_indices": [0, 4, 5, 9], + "tie_values": [ + t0, + t0 + 4 * dt, + t0 + 5 * dt - np.timedelta64(1000, "ns"), + t0 + 9 * dt - np.timedelta64(1000, "ns"), + ], + } + da = xd.DataArray(np.arange(10.0), {"time": coord}) + result = xd.trim_overlaps(da, tolerance=0.001) + npt.assert_array_equal(result.values, np.arange(10.0)) + + def test_stays_lazy(self, tmp_path): + from xdas.virtual import TileArray + + objs = [] + for index, start in enumerate([0, 8]): + da = xd.testing.dummy(dims=("time",), shape=(10,), step=0.01) + da["time"] = da["time"] + np.timedelta64(start * 10_000_000, "ns") + path = tmp_path / f"chunk_{index}.nc" + da.to_netcdf(path) + objs.append(xd.open_dataarray(path, engine="xdas", vtype="tiles")) + da = xd.concat(objs, "time", tolerance=False) + result = xd.trim_overlaps(da) + assert isinstance(result.data, TileArray) + assert result.sizes["time"] == 18 + assert result["time"].get_split_indices("overlaps").size == 0 + + def test_recurses_over_a_collection(self): + _, da = self.segments([(0.0, np.arange(5.0)), (0.03, np.arange(100.0, 105.0))]) + dc = xd.DataCollection( + {"CH001": xd.DataCollection([da, da], "acquisition")}, "station" + ) + result = xd.trim_overlaps(dc) + assert result.fields == ("station", "acquisition") + assert list(result) == ["CH001"] + assert len(result["CH001"]) == 2 + for element in result["CH001"]: + npt.assert_array_equal( + element.values, [0.0, 1.0, 2.0, 100.0, 101.0, 102.0, 103.0, 104.0] + ) + + def test_disjoint_claims(self): + # a stored view need not be in start order: here the second part lies + # entirely below the first, so the claims cannot merge into one span + # and the third part is trimmed against both + coord = { + "tie_indices": [0, 12, 13, 23, 24, 28], + "tie_values": [0.0, 12.0, 10.0, 20.0, 1.0, 5.0], + } + da = xd.DataArray( + np.concatenate( + [np.arange(13.0), np.arange(100.0, 111.0), np.arange(200.0, 205.0)] + ), + {"time": coord}, + ) + result = xd.trim_overlaps(da) + npt.assert_array_equal(result["time"].values, np.arange(21.0)) + npt.assert_array_equal( + result.values, + # 0 from the first part, 1-5 from the last, 6-9 from the first + # again, 10-20 from the second + [0.0] + + list(np.arange(200.0, 205.0)) + + [6.0, 7.0, 8.0, 9.0] + + list(np.arange(100.0, 111.0)), + ) + + def test_invalid_keep_raises(self): + da = xd.testing.dummy(dims=("time",), shape=(10,), step=0.01) + with pytest.raises(ValueError, match="`keep` must be"): + xd.trim_overlaps(da, keep="both") + + class TestSortby: def make_archive(self, tmp_path, vtype, pairs=((0, 2), (1, 3))): """Save 4 time chunks and fuse them losslessly as two runs. diff --git a/tests/virtual/test_core.py b/tests/virtual/test_core.py index a63b071a..b7a9b941 100644 --- a/tests/virtual/test_core.py +++ b/tests/virtual/test_core.py @@ -42,9 +42,9 @@ def test_base_is_abstract(self): with pytest.raises(TypeError, match="abstract"): VirtualBackend() - def test_finalize_save_defaults_to_nothing(self): + def test_sibling_datasets_defaults_to_nothing(self): source = VirtualSource("path.h5", "data", (2, 3), np.dtype("f8")) - assert source.finalize_save("path.nc") is None + assert source.sibling_datasets() == {} def test_derived_properties_shared_by_both_backends(self): source = VirtualSource("path.h5", "data", (2, 3), np.dtype("f8")) diff --git a/tests/virtual/test_tiles.py b/tests/virtual/test_tiles.py index 12de5034..5ca8daf7 100644 --- a/tests/virtual/test_tiles.py +++ b/tests/virtual/test_tiles.py @@ -1,9 +1,9 @@ """The tile-backed virtual array and its integration in the DataArray and native format.""" +import json import math import os -import dask.array as da_ import h5py import numpy as np import numpy.testing as npt @@ -12,7 +12,7 @@ import xdas as xd from xdas.io import Engine -from xdas.virtual.tiles import TileArray +from xdas.virtual.tiles import TILES_GROUP, TileArray, _Edges NX = 5 @@ -60,14 +60,10 @@ def stack(tmp_path): sizes.append(useful) parts.append(data[1:-1]) row += useful - manifest = TileArray.from_tiles( - paths, (sizes, NX), "float64", ENGINE, attrs={"units": "strain"} - ) + manifest = TileArray.from_tiles(paths, (sizes, NX), "float64", ENGINE) # per-tile source origins are view state: assigned through the manifest manifest = TileArray( - manifest.dataset.assign(starts_0=("tile_0", np.array([1, 1, 1]))), - manifest.dtype, - manifest.engine, + manifest.dataset.assign(starts_0=("tile_0", np.array([1, 1, 1]))) ) return manifest, np.concatenate(parts) @@ -99,11 +95,26 @@ def windowed(tmp_path): manifest = TileArray.from_tiles( paths, (sizes, NX), "float64", {"name": "h5py", "dataset": "data"} ) - manifest = TileArray( - manifest.dataset.assign(starts_0=("tile_0", np.array(starts))), - manifest.dtype, - manifest.engine, - ) + manifest = TileArray(manifest.dataset.assign(starts_0=("tile_0", np.array(starts)))) + return manifest, np.concatenate(parts) + + +@pytest.fixture +def uniform(tmp_path): + """Four files of identical height: the size column is constant. + + What a scanned acquisition looks like — ``combine`` starts a new one + whenever the geometry changes, so every file contributes the same + number of rows. + """ + paths, parts = [], [] + for k in range(4): + path = str(tmp_path / f"uni{k}.h5") + data = 100.0 * k + np.arange(6 * NX).reshape(6, NX) + _tile_file(path, data) + paths.append(path) + parts.append(data) + manifest = TileArray.from_tiles(paths, ([6] * 4, NX), "float64", ENGINE) return manifest, np.concatenate(parts) @@ -139,7 +150,7 @@ def _with_starts(manifest, *starts): for k, entry in enumerate(starts) if entry is not None } - return TileArray(manifest.dataset.assign(assign), manifest.dtype, manifest.engine) + return TileArray(manifest.dataset.assign(assign)) def _random_key(rng, shape, max_step=1): @@ -198,7 +209,7 @@ def test_shape_and_geometry(self, stack): manifest, reference = stack assert manifest.shape == reference.shape assert manifest.ntiles == 3 - npt.assert_array_equal(manifest._edges[0], [0, 10, 17, 29]) + npt.assert_array_equal(list(manifest._edges[0]), [0, 10, 17, 29]) def test_reads_across_sources(self, stack): manifest, reference = stack @@ -209,13 +220,19 @@ def test_reads_across_sources(self, stack): def test_dataset_model(self, stack, tmp_path): manifest, _ = stack dataset = manifest.dataset + # a varying column sizes its own axis, and says nothing assert tuple(dataset["sizes_0"].dims) == ("tile_0",) - assert tuple(dataset["sizes_1"].dims) == ("tile_1",) + assert dataset["sizes_0"].attrs == {} + # one tile spans the trailing axis: a constant column, folded; + # the header gives the tile count its dimension no longer does + assert tuple(dataset["sizes_1"].dims) == () + assert json.loads(dataset.attrs["header"])["ntiles"] == [3, 1] # per-file paths vary along tile_0 only: the trailing axis folds assert tuple(dataset["paths"].dims) == ("tile_0",) - # the common directory splits off: 0-d root, root-relative paths - assert dataset["root"].ndim == 0 - assert os.fsdecode(dataset["root"].values[()]) == str(tmp_path) + # the common directory splits off into the header, the stored + # paths staying relative to it + assert json.loads(dataset.attrs["header"])["root"] == str(tmp_path) + assert "root" not in dataset # strings are held as fixed-width bytes, not str objects assert dataset["paths"].dtype.kind == "S" assert dataset["paths"].values.tolist() == [b"src0.h5", b"src1.h5", b"src2.h5"] @@ -256,31 +273,22 @@ def test_validation(self, stack): TileArray.from_tiles("a", (5, NX), "f8", ENGINE, starts_0=[0]) with pytest.raises(ValueError, match="reserved"): TileArray.from_tiles("a", (5, NX), "f8", ENGINE, root=["r"]) - bad_root = manifest.dataset.copy() - bad_root["root"] = (("tile_0",), np.array(["a", "b", "c"], dtype=object)) - with pytest.raises(ValueError, match="0-d"): - TileArray(bad_root, manifest.dtype, manifest.engine) + bad_root = _reheader(manifest, root=["a", "b", "c"]) + with pytest.raises(ValueError, match="`root` must be a string"): + TileArray(bad_root) dataset = manifest.dataset.copy() with pytest.raises(ValueError, match="`sizes_0`"): - TileArray( - dataset.drop_vars(["sizes_0", "sizes_1"]), - manifest.dtype, - manifest.engine, - ) + TileArray(dataset.drop_vars(["sizes_0", "sizes_1"])) with pytest.raises(ValueError, match="`paths`"): - TileArray(dataset.drop_vars("paths"), manifest.dtype, manifest.engine) + TileArray(dataset.drop_vars("paths")) bad_starts = dataset.assign(starts_0=("tile_0", np.array([-1, 0, 0]))) with pytest.raises(ValueError, match="non-negative"): - TileArray(bad_starts, manifest.dtype, manifest.engine) + TileArray(bad_starts) def test_extra_variables_are_params(self, stack): """Any non-geometry manifest variable is a per-tile engine parameter.""" manifest, _ = stack - arr = TileArray( - manifest.dataset.assign(record=(("tile_0",), np.arange(3))), - manifest.dtype, - manifest.engine, - ) + arr = TileArray(manifest.dataset.assign(record=(("tile_0",), np.arange(3)))) assert arr._params == ("record",) def test_string_params_decode_to_str(self, tmp_path): @@ -360,9 +368,292 @@ def test_relative_paths_are_anchored(self, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path.parent) npt.assert_array_equal(np.asarray(manifest), data) - def test_attrs(self, stack): + def test_the_manifest_carries_nothing_but_the_header(self, stack): + """A tile array is a duck array: user metadata is the DataArray's.""" + manifest, _ = stack + assert set(manifest.dataset.attrs) == {"header"} + assert not hasattr(manifest, "attrs") + assert all( + not variable.attrs for variable in manifest.dataset.variables.values() + ) + + def test_header_records_the_engine_and_the_dtype(self, stack): + """What the columns cannot hold travels in the header, and must agree.""" manifest, _ = stack - assert manifest.attrs == {"units": "strain"} + header = json.loads(manifest.dataset.attrs["header"]) + assert header["engine"] == ENGINE + assert header["dtype"] == "float64" + # an omitted argument is read off the header + assert TileArray(manifest.dataset).equals(manifest) + # a given one must say the same thing + with pytest.raises(ValueError, match="records the engine"): + TileArray(manifest.dataset, engine="h5py") + with pytest.raises(ValueError, match="records the dtype"): + TileArray(manifest.dataset, dtype="f4") + # a manifest without a header (the scan-time assembly) needs both + headless = manifest.dataset.drop_attrs() + with pytest.raises(ValueError, match="records no engine"): + TileArray(headless, dtype="f8") + with pytest.raises(ValueError, match="records no dtype"): + TileArray(headless, engine=ENGINE) + + def test_an_engine_of_one_name_stores_bare(self): + """The header spells a constant-less engine as its name alone.""" + arr = TileArray.from_tiles("a", (5, NX), "f8", "h5py") + assert json.loads(arr.dataset.attrs["header"])["engine"] == "h5py" + assert TileArray(arr.dataset).engine == {"name": "h5py"} + + +def _reheader(manifest, **updates): + """Return the dataset of *manifest* with its header entries overridden.""" + header = json.loads(manifest.dataset.attrs["header"]) + header.update(updates) + return manifest.dataset.assign_attrs(header=json.dumps(header)) + + +def _is_collapsed(values): + """Whether a geometry column is held as a zero-stride broadcast view.""" + return values.strides == (0,) + + +class TestConstantGeometry: + """Constant geometry columns cost O(1), and behave exactly like full ones.""" + + def test_absent_columns_are_broadcast_views(self, uniform): + manifest, _ = uniform + assert "starts_0" not in manifest.dataset + assert "steps_0" not in manifest.dataset + for column, default in ((manifest._starts[0], 0), (manifest._steps[0], 1)): + assert _is_collapsed(column) + npt.assert_array_equal(column, [default] * 4) + + def test_stored_constant_column_collapses(self, uniform): + manifest, reference = uniform + # the column is read (it is in the file), then found constant + assert "sizes_0" in manifest.dataset + assert _is_collapsed(manifest._sizes[0]) + npt.assert_array_equal(manifest._sizes[0], [6] * 4) + assert manifest._edges[0].size == 6 + assert manifest._edges[0].values is None + assert manifest.shape == reference.shape + assert manifest.chunks == ((6, 6, 6, 6), (NX,)) + + def test_varying_column_stays_eager(self, stack): + manifest, _ = stack + assert not _is_collapsed(manifest._sizes[0]) + assert manifest._edges[0].size is None + npt.assert_array_equal(manifest._edges[0].values, [0, 10, 17, 29]) + + def test_truncated_last_tile_keeps_the_column_varying(self, tmp_path): + paths, parts = [], [] + for k, height in enumerate([6, 6, 4]): + path = str(tmp_path / f"trunc{k}.h5") + data = 100.0 * k + np.arange(height * NX).reshape(height, NX) + _tile_file(path, data) + paths.append(path) + parts.append(data) + manifest = TileArray.from_tiles(paths, ([6, 6, 4], NX), "float64", ENGINE) + assert not _is_collapsed(manifest._sizes[0]) + assert manifest._edges[0].size is None + npt.assert_array_equal(np.asarray(manifest), np.concatenate(parts)) + + @pytest.mark.parametrize("sizes", [[6, 6, 6, 6], [1, 1, 1], [5], [3, 3]]) + def test_closed_form_edges_match_the_eager_ones(self, sizes): + count, size = len(sizes), sizes[0] + closed = _Edges(np.broadcast_to(np.int64(size), (count,))) + eager = _Edges(np.asarray(sizes, dtype=np.int64)) + assert closed.size == size and eager.size is None + assert closed.total == eager.total + assert list(closed) == [int(edge) for edge in eager] + for i in range(count + 1): + assert closed.at(i) == eager.at(i) + for value in range(-2, closed.total + 3): + for side in ("left", "right"): + assert closed.searchsorted(value, side) == eager.searchsorted( + value, side + ) + for i0 in range(count + 1): + for i1 in range(i0, count + 1): + npt.assert_array_equal(closed.span(i0, i1), eager.span(i0, i1)) + + @pytest.mark.parametrize("seed", range(8)) + def test_uniform_slicing_matches_numpy(self, uniform, seed): + manifest, reference = uniform + rng = np.random.default_rng(seed) + key = _random_key(rng, manifest.shape, max_step=3) + npt.assert_array_equal(np.asarray(manifest[key]), reference[key]) + + def test_tile_aligned_slice_stays_collapsed(self, uniform): + manifest, reference = uniform + view = manifest[6:18] + assert _is_collapsed(view._sizes[0]) + assert view._edges[0].size == 6 + npt.assert_array_equal(np.asarray(view), reference[6:18]) + + def test_partial_slice_falls_back_to_the_eager_form(self, uniform): + manifest, reference = uniform + view = manifest[4:20] + assert not _is_collapsed(view._sizes[0]) + assert view._edges[0].size is None + npt.assert_array_equal(np.asarray(view), reference[4:20]) + + def test_reversal_and_streaming_on_a_uniform_grid(self, uniform): + manifest, reference = uniform + npt.assert_array_equal(np.asarray(manifest[::-1]), reference[::-1]) + # the streaming reductions walk the boundaries pairwise + npt.assert_allclose(np.mean(manifest), np.mean(reference)) + npt.assert_array_equal(np.max(manifest, axis=1), np.max(reference, axis=1)) + + def test_uniform_concat_stays_collapsed(self, uniform): + manifest, reference = uniform + fused = TileArray.concat([manifest, manifest]) + assert _is_collapsed(fused._sizes[0]) + assert fused._edges[0].size == 6 + npt.assert_array_equal( + np.asarray(fused), np.concatenate([reference, reference]) + ) + + def test_uniform_round_trip(self, uniform, tmp_path): + manifest, reference = uniform + path = str(tmp_path / "uniform.nc") + wrap(manifest).to_netcdf(path) + reopened = xd.open_dataarray(path) + assert reopened.data.equals(manifest) + assert _is_collapsed(reopened.data._sizes[0]) + npt.assert_array_equal(reopened.values, reference) + + def test_flipped_uniform_round_trip(self, uniform, tmp_path): + manifest, reference = uniform + view = np.flip(manifest, axis=0)[:, ::2] + assert _is_collapsed(view._sizes[0]) + path = str(tmp_path / "flipped.nc") + wrap(view).to_netcdf(path) + reopened = xd.open_dataarray(path) + assert reopened.data.equals(view) + npt.assert_array_equal(reopened.values, reference[::-1, ::2]) + + def test_constant_column_is_stored_as_one_element(self, uniform): + manifest, _ = uniform + # a constant column costs one element, not one per tile + assert manifest.dataset["sizes_0"].dims == () + assert int(manifest.dataset["sizes_0"].values[()]) == 6 + + def test_constant_non_default_starts_are_stored_as_one_element(self, stack): + manifest, _ = stack + assert manifest.dataset["starts_0"].dims == () + npt.assert_array_equal(manifest._starts[0], [1] * 3) + + def test_every_constant_column_folds(self, uniform): + manifest, _ = uniform + # nothing carries `tile_1` any more, and nothing has to: the + # header states how many tiles every axis holds + assert manifest.dataset["sizes_1"].dims == () + assert "tile_1" not in manifest.dataset.dims + assert json.loads(manifest.dataset.attrs["header"])["ntiles"] == [4, 1] + assert manifest.dataset["sizes_1"].attrs == {} + assert manifest.shape[1] == NX + assert manifest.chunks[1] == (NX,) + + def test_varying_column_is_stored_expanded(self, stack): + manifest, _ = stack + assert manifest.dataset["sizes_0"].dims == ("tile_0",) + + def test_expanded_stored_column_folds_on_open(self, uniform): + manifest, reference = uniform + expanded = manifest.dataset.assign( + sizes_0=("tile_0", np.full(4, 6, dtype="int64")), + starts_0=("tile_0", np.zeros(4, dtype="int64")), + steps_0=("tile_0", np.ones(4, dtype="int64")), + ) + reopened = TileArray(expanded) + # a column written per tile folds on the way in + assert reopened.dataset["sizes_0"].dims == () + assert "starts_0" not in reopened.dataset + assert "steps_0" not in reopened.dataset + assert reopened.equals(manifest) + npt.assert_array_equal(np.asarray(reopened), reference) + + def test_folded_column_is_not_written_per_tile(self, uniform, tmp_path): + manifest, _ = uniform + path = str(tmp_path / "folded.nc") + wrap(manifest).to_netcdf(path) + with h5py.File(path) as file: + assert file[f"/{TILES_GROUP}/sizes_0"].shape == () + reopened = xd.open_dataarray(path) + assert reopened.data.dataset["sizes_0"].dims == () + + def test_wholly_folded_grid_keeps_its_shape(self, uniform, tmp_path): + manifest, reference = uniform + # one tile of one file, concatenated with itself: every column + # is constant, so only the grid still counts the tiles + path = os.fsdecode(manifest._full_paths().item(0)) + single = TileArray.from_tiles(path, ([6], NX), "float64", ENGINE) + doubled = TileArray.concat([single, single]) + assert doubled.shape == (12, NX) + assert not any( + dim.startswith("tile_") for dim in map(str, doubled.dataset.dims) + ) + assert json.loads(doubled.dataset.attrs["header"])["ntiles"] == [2, 1] + path = str(tmp_path / "folded_grid.nc") + wrap(doubled).to_netcdf(path) + reopened = xd.open_dataarray(path) + expected = np.concatenate([reference[:6], reference[:6]]) + assert reopened.data.shape == (12, NX) + assert reopened.data.equals(doubled) + npt.assert_array_equal(reopened.values, expected) + # and it stays sliceable with no tile dimension to index + npt.assert_array_equal(np.asarray(reopened.data[3:9]), expected[3:9]) + + def test_tile_counts_follow_a_view(self, uniform): + manifest, _ = uniform + counts = lambda arr: json.loads(arr.dataset.attrs["header"])["ntiles"] + # a tile-aligned slice folds, and restates what it leaves behind + assert counts(manifest[6:18]) == [2, 1] + # a synthetic axis is one tile it never had to be told about + assert counts(np.expand_dims(manifest, 0)) == [4, 1, 1] + + def test_a_hand_built_manifest_implies_its_header(self, uniform): + """No header: the declared dimensions give the counts, the rest defaults.""" + manifest, reference = uniform + paths = manifest._full_paths().ravel() + built = xr.Dataset( + { + "sizes_0": ("tile_0", np.full(4, 6, dtype="int64")), + "sizes_1": ((), np.int64(NX)), + "paths": ("tile_0", paths), + } + ) + reopened = TileArray(built, manifest.dtype, manifest.engine) + header = json.loads(reopened.dataset.attrs["header"]) + assert header["ntiles"] == [4, 1] + assert "root" not in header and "axes" not in header + assert reopened.root == "" # paths stored whole + assert reopened.equals(manifest) + npt.assert_array_equal(np.asarray(reopened), reference) + + @pytest.mark.parametrize( + "ntiles, match", + [ + ([4], "one count per geometry axis"), + ("many", "one count per geometry axis"), + ([4, 0], "positive integers"), + ([3, 1], "disagrees with the size of `tile_0`"), + ], + ) + def test_malformed_tile_counts_raise(self, uniform, ntiles, match): + manifest, _ = uniform + with pytest.raises(ValueError, match=match): + TileArray(_reheader(manifest, ntiles=ntiles)) + + def test_hidden_and_synthetic_axes_round_trip(self, uniform, tmp_path): + manifest, reference = uniform + view = np.expand_dims(manifest[:, 2], 1) + assert view.shape == (24, 1) + path = str(tmp_path / "mapped.nc") + wrap(view).to_netcdf(path) + reopened = xd.open_dataarray(path) + assert reopened.data.equals(view) + npt.assert_array_equal(reopened.values, reference[:, 2][:, None]) class TestSourcePaths: @@ -377,7 +668,7 @@ def stored(self, manifest): return manifest.to_dataset()["paths"].values.ravel().tolist() def round_trip(self, manifest): - return TileArray(manifest.to_dataset(), manifest.dtype, manifest.engine) + return TileArray(manifest.to_dataset()) def test_root_splits_off(self, tmp_path): manifest = self.make(tmp_path / "sources" / "f.h5") @@ -400,19 +691,19 @@ def test_stored_paths_read(self, tmp_path): npt.assert_array_equal(np.asarray(restored), data[1:]) def test_rootless_manifest_reads(self, tmp_path): - """Manifests without a `root` (the pre-split stored form) still work.""" + """A manifest whose header states no root uses its paths as stored.""" data = np.arange(5 * NX, dtype="= 1 @@ -2126,16 +2386,231 @@ def test_eager_save_writes_values(self, stack, tmp_path): assert not isinstance(reopened.data, TileArray) npt.assert_array_equal(reopened.values, reference) - def test_dask_write_deprecated(self, tmp_path): - import dask - data = da_.from_delayed(dask.delayed(np.zeros)((4, NX)), (4, NX), np.float64) - da = xd.DataArray(data, dims=DIMS) - path = str(tmp_path / "dask.nc") - with pytest.warns(FutureWarning, match="dask-backed"): - da.to_netcdf(path, virtual=True) +class TestStoredForm: + """What each kind of view spends on disk, and that it comes back whole.""" + + def columns(self, manifest): + """The manifest variables, as ``{name: dimensions}``.""" + return { + name: tuple(map(str, manifest.dataset[name].dims)) + for name in map(str, manifest.dataset.data_vars) + } + + def header(self, manifest): + return json.loads(manifest.dataset.attrs["header"]) + + def round_trip(self, manifest, tmp_path, name): + """Write *manifest* to the native format and read it back.""" + path = str(tmp_path / f"{name}.nc") + dims = tuple(f"dim_{k}" for k in range(manifest.ndim)) + xd.DataArray(manifest, dims=dims).to_netcdf(path) + reopened = xd.open_dataarray(path) + assert isinstance(reopened.data, TileArray) + assert reopened.data.equals(manifest) + assert self.header(reopened.data) == self.header(manifest) + return reopened + + def test_scanned_acquisition(self, uniform, tmp_path): + """Equal-length files: every geometry column is one element.""" + manifest, reference = uniform + assert self.columns(manifest) == { + "sizes_0": (), + "sizes_1": (), + "paths": ("tile_0",), + } + assert self.header(manifest) == { + "ntiles": [4, 1], + "engine": ENGINE, + "dtype": "float64", + "root": str(tmp_path), + } + npt.assert_array_equal( + self.round_trip(manifest, tmp_path, "scan").values, reference + ) + + def test_truncated_last_file(self, uniform, tmp_path): + """A short last file is the one thing that keeps a size column varying.""" + manifest, reference = uniform + view = manifest[:21] + assert self.columns(view)["sizes_0"] == ("tile_0",) + npt.assert_array_equal(view.dataset["sizes_0"].values, [6, 6, 6, 3]) + assert "starts_0" not in view.dataset and "steps_0" not in view.dataset + npt.assert_array_equal( + self.round_trip(view, tmp_path, "truncated").values, reference[:21] + ) + + def test_decimated_window(self, uniform, tmp_path): + """A stepped slice: sizes and origins vary, the step is one element.""" + manifest, reference = uniform + view = manifest[3:22:2] + assert self.columns(view)["sizes_0"] == ("tile_0",) + assert self.columns(view)["starts_0"] == ("tile_0",) + assert self.columns(view)["steps_0"] == () + assert int(view.dataset["steps_0"].values[()]) == 2 + npt.assert_array_equal( + self.round_trip(view, tmp_path, "decimated").values, reference[3:22:2] + ) + + def test_reversal(self, uniform, tmp_path): + """Flipping negates one folded step and moves one folded origin.""" + manifest, reference = uniform + view = manifest[::-1] + assert self.columns(view)["starts_0"] == () + assert self.columns(view)["steps_0"] == () + assert int(view.dataset["starts_0"].values[()]) == 5 + assert int(view.dataset["steps_0"].values[()]) == -1 + npt.assert_array_equal( + self.round_trip(view, tmp_path, "reversed").values, reference[::-1] + ) + + def test_transposed_with_a_new_axis(self, uniform, tmp_path): + """The arrangement is one list; the inserted axis is a null in it.""" + manifest, reference = uniform + view = np.transpose(manifest)[:, np.newaxis] + assert self.header(view)["ntiles"] == [4, 1, 1] + assert self.header(view)["axes"] == [1, None, 0] + assert self.columns(view)["sizes_2"] == () + reopened = self.round_trip(view, tmp_path, "transposed") + npt.assert_array_equal(reopened.values, reference.T[:, np.newaxis]) + + def test_pinned_axis_is_hidden(self, uniform, tmp_path): + """Integer indexing a source axis leaves it in the grid, unpresented.""" + manifest, reference = uniform + view = manifest[:, 2] + assert self.header(view)["axes"] == [0] + assert self.columns(view)["sizes_1"] == () + assert self.columns(view)["starts_1"] == () + assert int(view.dataset["sizes_1"].values[()]) == 1 + assert int(view.dataset["starts_1"].values[()]) == 2 + reopened = self.round_trip(view, tmp_path, "pinned") + npt.assert_array_equal(reopened.values, reference[:, 2]) + + def test_stacked_axis(self, uniform, tmp_path): + """A stacked axis is synthetic but a full grid citizen: paths vary along it.""" + manifest, reference = uniform + paths = [os.fsdecode(path) for path in manifest._full_paths().ravel()] + other = TileArray.from_tiles(paths[::-1], ([6] * 4, NX), "float64", ENGINE) + view = np.stack([manifest, other]) + assert self.header(view)["ntiles"] == [4, 1, 2] + assert self.header(view)["axes"] == [None, 0, 1] + assert self.columns(view)["paths"] == ("tile_0", "tile_2") + reopened = self.round_trip(view, tmp_path, "stacked") + flipped = np.concatenate(np.split(reference, 4)[::-1]) + npt.assert_array_equal(reopened.values, np.stack([reference, flipped])) + + def test_pinning_a_stacked_axis_deletes_it(self, uniform, tmp_path): + """A pinned synthetic axis is contentless: it leaves the grid entirely.""" + manifest, reference = uniform + view = np.stack([manifest, manifest])[1] + # back to the plain two-axis grid, identity arrangement + assert "axes" not in self.header(view) + assert self.header(view)["ntiles"] == [4, 1] + assert "sizes_2" not in view.dataset + assert view.equals(manifest) + npt.assert_array_equal( + self.round_trip(view, tmp_path, "unstacked").values, reference + ) + + def test_two_synthetic_axes_transposed_past_each_other(self, uniform, tmp_path): + """Reordering synthetic axes renumbers the grid, not the map.""" + manifest, reference = uniform + paths = [os.fsdecode(path) for path in manifest._full_paths().ravel()] + others = [ + TileArray.from_tiles( + paths[k:] + paths[:k], ([6] * 4, NX), "float64", ENGINE + ) + for k in range(1, 4) + ] + rotate = lambda k: np.concatenate( + np.split(reference, 4)[k:] + np.split(reference, 4)[:k] + ) + stacked = np.stack( + [np.stack([manifest, others[0]]), np.stack([others[1], others[2]])] + ) + expected = np.stack( + [np.stack([reference, rotate(1)]), np.stack([rotate(2), rotate(3)])] + ) + assert stacked.shape == expected.shape + view = np.swapaxes(stacked, 0, 1) + assert self.header(view)["axes"] == [None, None, 0, 1] + reopened = self.round_trip(view, tmp_path, "swapped") + npt.assert_array_equal(reopened.values, np.swapaxes(expected, 0, 1)) + + def test_engine_parameters(self, uniform, tmp_path): + """Per-tile parameters are columns like any other: constants fold.""" + manifest, reference = uniform + paths = [os.fsdecode(path) for path in manifest._full_paths().ravel()] + + class ShiftEngine(Engine, name="shift"): + @staticmethod + def load_tile(path, selection, *, offset, gain): + with h5py.File(path, "r") as file: + return gain * file["data"][selection] + offset + + try: + view = TileArray.from_tiles( + paths, ([6] * 4, NX), "float64", "shift", offset=[0, 1, 2, 3], gain=1.0 + ) + assert self.columns(view)["offset"] == ("tile_0",) + assert self.columns(view)["gain"] == () + reopened = self.round_trip(view, tmp_path, "params") + expected = reference + np.repeat([0, 1, 2, 3], 6)[:, None] + npt.assert_array_equal(reopened.values, expected) + finally: + del Engine._registry["shift"] + + def test_non_utf8_root_survives_the_header(self, tmp_path): + """A root JSON cannot spell verbatim rides on escaped surrogates.""" + directory = os.fsdecode(os.fsencode(str(tmp_path)) + b"/caf\xe9") + os.mkdir(directory) + data = np.arange(4.0 * NX).reshape(4, NX) + _tile_file(os.path.join(directory, "f.h5"), data) + manifest = TileArray.from_tiles( + os.path.join(directory, "f.h5"), (4, NX), "float64", ENGINE + ) + assert manifest.root == directory + # the JSON document itself stays pure ASCII + manifest.dataset.attrs["header"].encode("ascii") + reopened = self.round_trip(manifest, tmp_path, "surrogate") + assert reopened.data.root == directory + npt.assert_array_equal(reopened.values, data) + + def test_the_placeholder_is_a_projection(self, uniform, tmp_path): + """It costs nothing, states the shape and the type, and points home.""" + manifest, _ = uniform + path = str(tmp_path / "placeholder.nc") + wrap(manifest).to_netcdf(path) + with h5py.File(path) as file: + placeholder = file["__values__"] + assert placeholder.shape == manifest.shape + assert placeholder.dtype == manifest.dtype + assert placeholder.id.get_storage_size() == 0 + assert placeholder.attrs["__tiling__"] == TILES_GROUP + assert "__tile_array__" not in placeholder.attrs + assert file.attrs["Conventions"] == "CF-1.13" + + def test_a_placeholder_of_another_type_is_refused(self, uniform, tmp_path): + """The projection must project: a divergent type declares a cast.""" + manifest, _ = uniform + path = str(tmp_path / "mistyped.nc") + wrap(manifest).to_netcdf(path) + with h5py.File(path, "a") as file: + header = json.loads(file[TILES_GROUP].attrs["header"]) + header["dtype"] = "float32" + file[TILES_GROUP].attrs["header"] = json.dumps(header) + with pytest.raises(ValueError, match="placeholder is float64"): + xd.open_dataarray(path, engine="xdas") + + def test_dunder_attributes_stay_out_of_the_data_array(self, uniform, tmp_path): + """The `__*__` namespace is the format's; the rest is the user's.""" + manifest, _ = uniform + da = wrap(manifest) + da.attrs = {"instrument": "OptoDAS"} + path = str(tmp_path / "attrs.nc") + da.to_netcdf(path) reopened = xd.open_dataarray(path) - npt.assert_array_equal(reopened.values, np.zeros((4, NX))) + assert reopened.attrs == {"instrument": "OptoDAS"} class TestPermuteTiles: diff --git a/xdas/__init__.py b/xdas/__init__.py index a1966f7d..0145ba61 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -57,6 +57,7 @@ "plot_availability", "sortby", "split", + "trim_overlaps", ] from . import ( @@ -109,5 +110,6 @@ routines, sortby, split, + trim_overlaps, ) from .core.methods import * diff --git a/xdas/coordinates/interp.py b/xdas/coordinates/interp.py index 1b19aec2..5105c791 100644 --- a/xdas/coordinates/interp.py +++ b/xdas/coordinates/interp.py @@ -222,7 +222,10 @@ def _get_indexer(self, value, method=None): "jitter in the tie values, consider smoothing the coordinate by " "including some tolerance. This can be done by " "`da[dim] = da[dim].simplify(tolerance)`, or by specifying a " - "tolerance when opening multiple files." + "tolerance when opening multiple files. If the overlaps are " + "genuine, resolve them with `xdas.trim_overlaps(da)`, which " + "drops the duplicated samples, or cut them apart with " + "`xdas.split(da, 'overlaps')`, which keeps every copy." ) else: # pragma: no cover raise @@ -309,7 +312,9 @@ def _concat(self, other): @override def _to_dataset(self, dataset, attrs): - mapping = f"{self.name}: {self.name}_indices {self.name}_values" + # CF-1.13: a group names its tie point coordinate variables and + # ends with the interpolation variable describing them + mapping = f"{self.name}_values: {self.name}_interpolation" if "coordinate_interpolation" in attrs: attrs["coordinate_interpolation"] += " " + mapping else: @@ -322,7 +327,10 @@ def _to_dataset(self, dataset, attrs): ) interp_attrs = { "interpolation_name": "linear", - "tie_points_mapping": f"{self.name}_points: {self.name}_indices {self.name}_values", + # interpolated dimension: index variable, subsampled dimension + "tie_point_mapping": f"{self.dim}: {self.name}_indices {self.name}_points", + # xdas reconstructs in float64 (int64 nanoseconds for datetimes) + "computational_precision": "64", } if self.sampling_interval is not None: interp_attrs.update( @@ -345,18 +353,24 @@ def _collect_from_dataset(cls, dataset, name): coords = {} mapping = dataset[name].attrs.pop("coordinate_interpolation", None) if mapping is not None: - for dim, indices, values in re.findall(r"(\w+): (\w+) (\w+)", mapping): + for coord, dim, indices, values in _parse_interpolation(mapping, dataset): data = { "tie_indices": dataset[indices].values, "tie_values": dataset[values].values, } - interp_attrs = dataset[f"{dim}_interpolation"].attrs + # the oldest files spelled the mapping without writing an + # interpolation variable at all + interp_attrs = ( + dataset[f"{coord}_interpolation"].attrs + if f"{coord}_interpolation" in dataset + else {} + ) if "sampling_interval" in interp_attrs: data["sampling_interval"] = decode_delta( "sampling_interval", interp_attrs ) data["tolerance"] = decode_delta("tolerance", interp_attrs) - coords[dim] = Coordinate(data, dim) + coords[coord] = Coordinate(data, dim) return coords def __add__(self, other): @@ -673,6 +687,58 @@ def _continuous_segments(self): return num[mask], den[mask] +def _parse_interpolation(mapping, dataset): + """ + Yield ``(name, dim, indices, values)`` per group of *mapping*. + + Reads a ``coordinate_interpolation`` attribute in either spelling. + CF-1.13 words each group ``tie_point_coordinate_variable: [...] + interpolation_variable``, the coordinate name and the interpolated + dimension then coming from the interpolation variable and its + ``tie_point_mapping``; xdas wrote ``dimension: index_variable + value_variable`` before the format break. Only the CF spelling ends + a group with a variable carrying ``interpolation_name``, which is + what tells the two apart. The tie point coordinate variable name is + taken from the group as written, whatever it is. + + Parameters + ---------- + mapping : str + The attribute value to parse. + dataset : xarray.Dataset + The dataset the named variables live in. + + Yields + ------ + tuple of str + Coordinate name, interpolated dimension, tie point index + variable and tie point coordinate variable. + """ + groups, tie_points = [], [] + for word in mapping.split(): + if word.endswith(":"): + tie_points.append(word[:-1]) + else: + groups.append((tie_points, word)) + tie_points = [] + if all( + len(tie_points) == 1 + and word in dataset + and "interpolation_name" in dataset[word].attrs + for tie_points, word in groups + ): + for (values,), interpolation in groups: + name = interpolation.removesuffix("_interpolation") + dim, indices, _ = re.match( + r"(\w+): (\w+) (\w+)", + dataset[interpolation].attrs["tie_point_mapping"], + ).groups() + yield name, dim, indices, values + else: + for dim, indices, values in re.findall(r"(\w+): (\w+) (\w+)", mapping): + yield dim, dim, indices, values + + def _sleeve(x, y, epsilon): """ Reduce the piecewise-linear curve *(x, y)* with a one-pass greedy sleeve. diff --git a/xdas/coordinates/sampled.py b/xdas/coordinates/sampled.py index 8faff535..25aee5dc 100644 --- a/xdas/coordinates/sampled.py +++ b/xdas/coordinates/sampled.py @@ -10,10 +10,11 @@ from typing_extensions import override from .core import ( - CODE_TO_UNITS, UNITS_TO_CODE, AxisCoordinate, Coordinate, + decode_delta, + encode_delta, is_monotonic_increasing, parse_data_dim, parse_scalar_delta, @@ -322,7 +323,13 @@ def _concat(self, other): @override def _to_dataset(self, dataset, attrs): - mapping = f"{self.name}: {self.name}_sampling" + # Variation on CF-1.13 coordinate subsampling: a group names its + # tie point coordinate variable and ends with the sampling + # variable describing it. The sampling variable is a container + # whose mapping puts the segment length variable in the tie point + # index variable's slot, and whose spacing travels as attributes, + # like the regular metadata of an interpolated coordinate. + mapping = f"{self.name}_values: {self.name}_sampling" if "coordinate_sampling" in attrs: attrs["coordinate_sampling"] += " " + mapping else: @@ -332,25 +339,16 @@ def _to_dataset(self, dataset, attrs): if np.issubdtype(self.tie_values.dtype, np.datetime64) else self.tie_values ) - tie_lengths = self.tie_lengths - interp_attrs = { - "tie_point_mapping": f"{self.dim}: {self.name}_values {self.name}_lengths", + sampling_attrs = { + # interpolated dimension: segment length variable, subsampled dimension + "tie_point_mapping": f"{self.dim}: {self.name}_lengths {self.name}_points", + **encode_delta("sampling_interval", self.sampling_interval), } - - # timedelta - if np.issubdtype(self.sampling_interval.dtype, np.timedelta64): - code, count = np.datetime_data(self.sampling_interval.dtype) - interp_attrs["dtype"] = "timedelta64[ns]" - interp_attrs["units"] = CODE_TO_UNITS[code] - sampling_interval = count * self.sampling_interval.astype(int) - else: - sampling_interval = self.sampling_interval - dataset.update( { - f"{self.name}_sampling": ((), sampling_interval, interp_attrs), + f"{self.name}_sampling": ((), np.nan, sampling_attrs), + f"{self.name}_lengths": (f"{self.name}_points", self.tie_lengths), f"{self.name}_values": (f"{self.name}_points", tie_values), - f"{self.name}_lengths": (f"{self.name}_points", tie_lengths), } ) return dataset, attrs @@ -361,29 +359,35 @@ def _collect_from_dataset(cls, dataset, name): coords = {} mapping = dataset[name].attrs.pop("coordinate_sampling", None) if mapping is not None: - matches = re.findall(r"(\w+): (\w+)", mapping) - for match in matches: - name, sampling = match - dim, values, lengths = re.match( - r"(\w+): (\w+) (\w+)", dataset[sampling].attrs["tie_point_mapping"] + for first, sampling in re.findall(r"(\w+): (\w+)", mapping): + sampling_attrs = dataset[sampling].attrs + dim, second, third = re.match( + r"(\w+): (\w+) (\w+)", sampling_attrs["tie_point_mapping"] ).groups() + if "sampling_interval" in sampling_attrs: + coord, values, lengths = ( + sampling.removesuffix("_sampling"), + first, + second, + ) + interval = decode_delta("sampling_interval", sampling_attrs) + else: + # the spelling that predates the CF-shaped grammar: the + # group named the coordinate rather than its tie point + # variable, the mapping listed both tie point variables, + # and the interval was the sampling variable's own value + coord, values, lengths = first, second, third + interval = dataset[sampling].values[()] + if "units" in sampling_attrs and "dtype" in sampling_attrs: + interval = np.timedelta64( + interval, UNITS_TO_CODE[sampling_attrs["units"]] + ).astype(sampling_attrs["dtype"]) data = { "tie_values": dataset[values].values, "tie_lengths": dataset[lengths].values, - "sampling_interval": dataset[sampling].values[()], + "sampling_interval": interval, } - - # timedelta - if ( - "dtype" in dataset[sampling].attrs - and "units" in dataset[sampling].attrs - ): - data["sampling_interval"] = np.timedelta64( - data["sampling_interval"], - UNITS_TO_CODE[dataset[sampling].attrs.pop("units")], - ).astype(dataset[sampling].attrs.pop("dtype")) - - coords[name] = Coordinate(data, dim) + coords[coord] = Coordinate(data, dim) return coords def __add__(self, other): diff --git a/xdas/core/__init__.py b/xdas/core/__init__.py index d988bb16..79e23617 100644 --- a/xdas/core/__init__.py +++ b/xdas/core/__init__.py @@ -28,6 +28,7 @@ "plot_availability", "sortby", "split", + "trim_overlaps", ] from .dataarray import DataArray @@ -51,4 +52,5 @@ plot_availability, sortby, split, + trim_overlaps, ) diff --git a/xdas/core/dataarray.py b/xdas/core/dataarray.py index 1ed99f1a..717ba22f 100644 --- a/xdas/core/dataarray.py +++ b/xdas/core/dataarray.py @@ -11,7 +11,6 @@ import numpy as np import xarray as xr -from dask.array import Array as DaskArray from numpy.lib.mixins import NDArrayOperatorsMixin from ..coordinates import AxisCoordinate, Coordinates @@ -105,7 +104,8 @@ def __repr__(self): data_repr = np.array2string( self.data, precision=precision, threshold=0, edgeitems=edgeitems ) - elif isinstance(self.data, DaskArray): + elif type(self.data).__module__.startswith("dask."): + # duck-typed: dask is no longer a dependency, only valid data data_repr = f"DaskArray: {_to_human(self.data.nbytes)} ({self.data.dtype})" else: data_repr = repr(self.data) @@ -842,7 +842,7 @@ def to_stream( the obspy stream version of the data array. """ - from ..io.miniseed import to_stream + from ..io.obspy import to_stream return to_stream(self, network, station, location, channel, dim) @@ -868,7 +868,7 @@ def from_stream(cls, st, dims=("channel", "time")): DataArray: The consolidated data array. """ - from ..io.miniseed import from_stream + from ..io.obspy import from_stream return from_stream(st, dims) diff --git a/xdas/core/datacollection.py b/xdas/core/datacollection.py index d62bcc2b..c8bfe798 100644 --- a/xdas/core/datacollection.py +++ b/xdas/core/datacollection.py @@ -75,6 +75,18 @@ def empty(self): """``True`` if the collection contains no elements.""" return len(self) == 0 + @property + def fields(self): + """Ordered, deduplicated tuple of the node names of the whole subtree.""" + values = self.values() if self.ismapping() else self + out = (self.name,) + tuple( + name + for value in values + if isinstance(value, DataCollection) + for name in value.fields + ) + return uniquifiy(out) + def query(self, indexers=None, **indexers_kwargs): """ Query a given subset from a data collection. @@ -82,6 +94,11 @@ def query(self, indexers=None, **indexers_kwargs): The data collection is walked through, if any node name corresponds to a key of the `indexers`, the corresponding value is used to select a subset of that node. + Each indexer must name a level of the collection, i.e. be one of `fields`. + This is what distinguishes querying from `sel`: `query` chooses *which* + leaves are kept by their position in the hierarchy, while `sel` trims + *inside* each leaf by coordinate label. + Parameters ---------- indexers : dict, optional @@ -95,6 +112,11 @@ def query(self, indexers=None, **indexers_kwargs): DataCollection: The queried data. + Raises + ------ + KeyError + If an indexer does not name any level of the collection. + Examples -------- >>> import xdas as xd @@ -114,48 +136,68 @@ def query(self, indexers=None, **indexers_kwargs): 0: """ - if indexers is None: - indexers = {} + indexers = {} if indexers is None else dict(indexers) indexers.update(indexers_kwargs) - if self.name in indexers: - key = indexers[self.name] - if self.issequence(): + fields = self.fields + unknown = [key for key in indexers if key not in fields] + if unknown: + raise KeyError( + f"{unknown} do not name any level of the collection; " + f"available: {list(fields)}" + ) + return self._query(indexers) + + def select(self, indexers=None, **indexers_kwargs): + """ + Select a given subset from a data collection. + + Alias of `query`, named after `obspy.Stream.select`. See `query`. + """ + return self.query(indexers, **indexers_kwargs) + + def _query(self, indexers): + """Recursive half of `query`, with the indexers already validated. + + Every level is walked, whether or not it is named in *indexers*: an + indexer applies wherever its level sits in the tree, not only at the + root. + """ + key = indexers.get(self.name) + if self.issequence(): + data = list(self) + if self.name in indexers: if isinstance(key, int): - data = [self[key]] + data = [data[key]] elif isinstance(key, slice): - data = self[key] + data = data[key] else: - raise ValueError(f"{self.name} query must be a string") - data = [ - ( - value.query(indexers) - if isinstance(value, DataCollection) - else value - ) - for value in data - ] - elif self.ismapping(): + raise ValueError(f"{self.name} query must be an integer or a slice") + data = [ + (value._query(indexers) if isinstance(value, DataCollection) else value) + for value in data + ] + elif self.ismapping(): + data = dict(self) + if self.name in indexers: if isinstance(key, str): data = { name: value - for name, value in self.items() + for name, value in data.items() if fnmatch(name, key) } else: raise ValueError(f"{self.name} query must be a string") - data = { - name: ( - value.query(indexers) - if isinstance(value, DataCollection) - else value - ) - for name, value in data.items() - } - else: # pragma: no cover - raise TypeError("unknown type of data collection") - return DataCollection(data, self.name) - else: - return self + data = { + name: ( + value._query(indexers) + if isinstance(value, DataCollection) + else value + ) + for name, value in data.items() + } + else: # pragma: no cover + raise TypeError("unknown type of data collection") + return DataCollection(data, self.name) def issequence(self): """Return ``True`` if this is a :class:`DataSequence`.""" @@ -185,15 +227,7 @@ def from_netcdf(cls, fname, group=None): """ if isinstance(fname, Path): fname = str(fname) - self = DataMapping.from_netcdf(fname, group) - try: - keys = [int(key) for key in self.keys()] - if keys == list(range(len(keys))): - return DataSequence.from_mapping(self) - else: - return self - except ValueError: - return self + return as_sequence_if_positional(DataMapping.from_netcdf(fname, group)) class DataMapping(DataCollection, dict): @@ -238,14 +272,6 @@ def __repr__(self): def __reduce__(self): return self.__class__, (dict(self), self.name) - @property - def fields(self): - """Ordered, deduplicated tuple of node names at this level and its immediate children.""" - out = (self.name,) + tuple( - value.name for value in self.values() if isinstance(value, DataCollection) - ) - return uniquifiy(out) - def to_netcdf( self, fname, @@ -433,14 +459,6 @@ def __repr__(self): def __reduce__(self): return self.__class__, (list(self), self.name) - @property - def fields(self): - """Ordered, deduplicated tuple of node names at this level and its immediate children.""" - out = (self.name,) + tuple( - value.name for value in self if isinstance(value, DataCollection) - ) - return uniquifiy(out) - def to_mapping(self): """Convert to an integer-keyed :class:`DataMapping`.""" return DataMapping(dict(enumerate(self)), self.name) @@ -626,6 +644,19 @@ def parse(data, name=None): return data, name +def as_sequence_if_positional(dm): + """Return :class:`DataMapping` *dm* as a sequence if its keys are its positions. + + A sequence is written under the canonical decimal spelling of its + positions, so that is what is compared: parsing the keys as integers + instead would read a mapping keyed by a zero-padded code — a SEED + location, say — back as a sequence, losing the keys. + """ + if list(dm) == [str(index) for index in range(len(dm))]: + return DataSequence.from_mapping(dm) + return dm + + def get_depth(group): """Return the maximum nesting depth of an HDF5 *group* by counting ``"/"`` separators.""" if not isinstance(group, h5py.Group): diff --git a/xdas/core/routines.py b/xdas/core/routines.py index 844feaa3..a24ba149 100644 --- a/xdas/core/routines.py +++ b/xdas/core/routines.py @@ -167,6 +167,23 @@ def open( return open_datacollection(paths) except Exception: # noqa: BLE001, S110 - fall back to dataarray pass + try: + dc = _resolve_engine( + engine, vtype, ctype, engine_kwargs + ).open_datacollection(paths) + except NotImplementedError: + pass # the engine describes a file as one array + else: + # combine whether one file was opened or many, so the returned + # shape never depends on the file count + return combine_by_field( + [dc], + dim, + tolerance, + False if squeeze is None else squeeze, + None, + verbose, + ) return open_dataarray( paths, engine=engine, vtype=vtype, ctype=ctype, **engine_kwargs ) @@ -181,8 +198,20 @@ def open( parallel=parallel, verbose=verbose, ) - except Exception: # noqa: BLE001, S110 - fall back to mfdataarray + except Exception: # noqa: BLE001, S110 - not native collections pass + try: + return open_mfdatacollection( + paths, + dim, + tolerance, + squeeze=False if squeeze is None else squeeze, + parallel=parallel, + verbose=verbose, + engine=_resolve_engine(engine, vtype, ctype, engine_kwargs), + ) + except NotImplementedError: + pass # the engine describes a file as one array return open_mfdataarray( paths, dim, @@ -211,7 +240,16 @@ def open( def open_mfdatacollection( - paths, dim="first", tolerance=None, squeeze=False, verbose=False, parallel=None + paths, + dim="first", + tolerance=None, + squeeze=False, + verbose=False, + parallel=None, + engine=None, + vtype=None, + ctype=None, + **engine_kwargs, ): """ Open a multiple file DataCollection. @@ -244,6 +282,18 @@ def open_mfdatacollection( global xdas configuration. Default to None. verbose: bool Whether to display a progress bar. Default to False. + engine: str or Engine, optional + The file format engine to use, given by name or as a configured + :class:`~xdas.io.Engine` instance. Default to the native format. + vtype : str, optional + The virtualization type to use. If None, the engine default is used. + Only valid when `engine` is given by name. + ctype : str or dict, optional + The coordinate type(s) to use. If None, the engine defaults are used. + Only valid when `engine` is given by name. + **engine_kwargs + Format-specific engine parameters forwarded to the engine constructor. + Only valid when `engine` is given by name. Returns ------- @@ -252,6 +302,8 @@ def open_mfdatacollection( """ paths = _ensure_str_paths(paths) + if engine is not None: + engine = _resolve_engine(engine, vtype, ctype, engine_kwargs) if isinstance(paths, str): paths = sorted(glob(paths)) @@ -277,10 +329,12 @@ def open_mfdatacollection( iterator = tqdm(paths, desc="Fetching metadata from files") else: iterator = paths - objs = [open_datacollection(path) for path in iterator] + objs = [open_datacollection(path, engine=engine) for path in iterator] else: executor = get_reusable_executor(max_workers) - futures = [executor.submit(open_datacollection, path) for path in paths] + futures = [ + executor.submit(open_datacollection, path, engine=engine) for path in paths + ] if verbose: iterator = tqdm( as_completed(futures), @@ -290,7 +344,10 @@ def open_mfdatacollection( else: iterator = as_completed(futures) objs = [future.result() for future in iterator] - return combine_by_field(objs, dim, tolerance, squeeze, True, verbose) + # the native format stacks hdf5 sources; the engines that describe a file + # as a collection are tile-backed, and let `concat` pick + virtual = True if engine is None else None + return combine_by_field(objs, dim, tolerance, squeeze, virtual, verbose) def open_mfdatatree( @@ -668,7 +725,7 @@ def consume(da): runs.extend(combine_by_coords(objs, dim, False, False)) objs.clear() - if (max_workers == 1) or (engine.name == "miniseed"): # TODO: dirty miniseed fix + if max_workers == 1: iterator = ( tqdm(paths, desc="Fetching metadata from files") if verbose else paths ) @@ -761,7 +818,7 @@ def _combine_runs(runs, dim, tolerance, squeeze): else da[dim].values ) ) - collection = DataCollection(results) + collection = DataCollection(results, "acquisition") if squeeze and len(collection) == 1: return collection[0] return collection @@ -814,7 +871,9 @@ def open_dataarray(fname, engine=None, vtype=None, ctype=None, **engine_kwargs): return engine.open_dataarray(fname) -def open_datacollection(fname, group=None): +def open_datacollection( + fname, group=None, engine=None, vtype=None, ctype=None, **engine_kwargs +): """ Open a DataCollection from a file. @@ -822,6 +881,21 @@ def open_datacollection(fname, group=None): ---------- fname : str The path of the DataCollection. + group : str, optional + The location of the data collection within the file. Root by default. + Only meaningful for the native format. + engine: str or Engine, optional + The file format engine to use, given by name or as a configured + :class:`~xdas.io.Engine` instance. Default to the native format. + vtype : str, optional + The virtualization type to use. If None, the engine default is used. + Only valid when `engine` is given by name. + ctype : str or dict, optional + The coordinate type(s) to use. If None, the engine defaults are used. + Only valid when `engine` is given by name. + **engine_kwargs + Format-specific engine parameters forwarded to the engine constructor. + Only valid when `engine` is given by name. Returns ------- @@ -832,11 +906,27 @@ def open_datacollection(fname, group=None): ------ FileNotFound If no file can be found. + NotImplementedError + If the engine does not describe a file as a collection. """ fname = _ensure_str_paths(fname) if not os.path.exists(fname): raise FileNotFoundError("no file to open") - return DataCollection.from_netcdf(fname, group) + if engine is None: + if vtype is not None or ctype is not None or engine_kwargs: + raise ValueError( + "`vtype`, `ctype` and engine keyword arguments require naming an " + "engine; the native format reads a collection as it was written" + ) + return DataCollection.from_netcdf(fname, group) + if group is not None: + raise ValueError( + "`group` is a native-format parameter; pass it as an engine keyword " + "argument instead" + ) + return _resolve_engine(engine, vtype, ctype, engine_kwargs).open_datacollection( + fname + ) def asdataarray(obj, tolerance=None): @@ -914,9 +1004,10 @@ def combine_by_field( nodes = [dc for dc in objs if isinstance(dc, dict)] if leaves and not nodes: objs = [da for dc in leaves for da in dc] - dc = combine_by_coords(objs, dim, tolerance, squeeze, virtual, verbose) - dc.name = leaves[0].name - return dc + # the level is named for what its elements are, and combining changes + # that: whatever the inputs held, each output element is one + # acquisition epoch. `combine_by_coords` names it. + return combine_by_coords(objs, dim, tolerance, squeeze, virtual, verbose) elif nodes and not leaves: (name,) = {dc.name for dc in nodes} keys = sorted(set.union(*[set(dc.keys()) for dc in nodes])) @@ -1002,9 +1093,12 @@ def combine_by_coords( bag.append(da) bags.append(bag) - # concatenate each bag + # concatenate each bag. `Bag` splits on sampling rate, dtype and non-concat + # coordinates, and gaps land inside the coordinate, so every element of the + # result is one acquisition epoch — which is what the level is named for. collection = DataCollection( - [concatenate(bag, dim, tolerance, virtual, verbose) for bag in bags] + [concatenate(bag, dim, tolerance, virtual, verbose) for bag in bags], + "acquisition", ) # squeeze if possible @@ -1107,6 +1201,54 @@ def check_sampling_interval(self, da): raise CompatibilityError("sampling intervals are not compatible") +def _get_promoted_coords(objs, dim): + """Check the non-concat coordinates of *objs* and report the varying scalars. + + Called when *dim* opens a new dimension, where :func:`concat` would + otherwise keep the first element's coordinates and silently discard the + others'. Coordinates that every element shares are kept as they are; + scalar ones that differ are what the new dimension is made of, and are + returned for promotion to a coordinate along it — the ``channel`` of a + stack of seismic traces, say. Anything else is a genuine incompatibility. + + Parameters + ---------- + objs : list of DataArray + The data arrays about to be concatenated, before ``expand_dims``. + dim : str + The name of the new dimension. A coordinate of that name is left + alone: ``expand_dims`` promotes it. + + Returns + ------- + dict + Mapping from coordinate name to the list of per-element scalar values, + in the order of *objs*. + """ + names = [name for name in objs[0].coords if name != dim] + for da in objs[1:]: + other = [name for name in da.coords if name != dim] + if set(other) != set(names): + raise ValueError( + "objects to concatenate along the new dimension " + f"{dim!r} must share their coordinates; got {sorted(names)} " + f"and {sorted(other)}" + ) + promoted = {} + for name in names: + coord = objs[0].coords[name] + if all(da.coords[name].equals(coord) for da in objs[1:]): + continue + if not all(da.coords[name].dim is None for da in objs): + raise ValueError( + f"coordinate {name!r} differs across the objects to concatenate " + f"along the new dimension {dim!r}; only scalar coordinates may " + "vary, and are then promoted to a coordinate along that dimension" + ) + promoted[name] = [da.coords[name].values for da in objs] + return promoted + + def concat( objs, dim="first", @@ -1150,7 +1292,16 @@ def concat( ------- DataArray The concatenated dataarray. Coordinates along axes other than *dim* are - taken from the first element; no compatibility check is performed on ``objs[1:]``. + taken from the first element. When *dim* opens a new dimension, the + other elements must carry the same non-concat coordinates, except for + scalar ones, which are promoted to a coordinate along the new dimension + when they vary. + + Raises + ------ + ValueError + If *dim* opens a new dimension and the elements do not agree on their + non-concat coordinates other than varying scalars. """ objs = list(objs) @@ -1166,16 +1317,18 @@ def concat( axis = objs[0].get_axis_num(dim) dim = objs[0].dims[axis] # ensure not "first" or "last" dims = objs[0].dims + promoted = {} else: axis = 0 dims = (dim, *objs[0].dims) + promoted = _get_promoted_coords(objs, dim) objs = [da.expand_dims(dim) for da in objs] - # TODO: check that objs[1:] have the same non-concat coords as objs[0] coords = objs[0].coords.drop_dims(dim) name = objs[0].name attrs = objs[0].attrs dim_has_coords = dim in objs[0].coords + order = list(range(len(objs))) if dim_has_coords: coord, order = concat_coords( @@ -1189,6 +1342,9 @@ def concat( objs = [objs[idx] for idx in order] coords[dim] = coord + for coord_name, values in promoted.items(): + coords[coord_name] = (dim, [values[idx] for idx in order]) + iterator = ( tqdm(objs, desc="Linking dataarray") if verbose else objs ) # TODO : remove tqdm? @@ -1441,6 +1597,128 @@ def split(da, indices_or_sections="discontinuities", dim="first", tolerance=None ) +def trim_overlaps(obj, keep="last", dim="first", tolerance=None): + """ + Remove the overlapping samples of a data array, keeping one copy of each. + + An overlap is a place where the coordinate steps backwards: two segments + describe the same span of time (or distance), typically because two files + share a sample at their seam, or because an acquisition was restarted + slightly before it stopped. This routine cuts the data array at its + overlaps, drops the duplicated samples from all but one segment, and + concatenates what is left back into a single data array. + + Trimming lands on a sample boundary, never between two: nothing is ever + resampled, interpolated or filled. Sub-sample misalignment therefore + survives as a discontinuity of the coordinate. Everything is done at the + manifest level, so a lazy data array stays lazy and no data is read. + + Parameters + ---------- + obj : DataArray or DataCollection + The data to trim. A data collection is trimmed leaf by leaf, its tree + preserved. + keep : {"last", "first"}, optional + Which copy of an overlapping span to keep. ``"last"`` (default) gives + the later segment precedence, on the assumption that the following + data carries the more correct time — this is ObsPy's + ``Stream.merge(method=1, interpolation_samples=0)``. ``"first"`` is + the mirror image. + dim : str, optional + The dimension along which to look for overlaps. Default to "first". + tolerance : float or timedelta64, optional + The magnitude below which a backward step is not considered an + overlap. For time coordinates, numeric values are considered as + seconds. By default only exactly zero-magnitude steps are ignored. + Note that jitter is usually better handled upstream, by spending a + tolerance when combining or by `da[dim] = da[dim].simplify(tolerance)`. + + Returns + ------- + DataArray or DataCollection + The data with its overlaps resolved, of the same type as *obj*. + + See Also + -------- + split : Cut a data array at its overlaps, keeping every copy (1 to N). + + Examples + -------- + >>> import numpy as np + >>> import xdas as xd + + Two segments of five samples overlapping by two: + + >>> coord = {"tie_indices": [0, 4, 5, 9], "tie_values": [0.0, 4.0, 3.0, 7.0]} + >>> da = xd.DataArray(np.arange(10.0), {"time": coord}) + >>> xd.trim_overlaps(da).values + array([0., 1., 2., 5., 6., 7., 8., 9.]) + >>> xd.trim_overlaps(da, keep="first").values + array([0., 1., 2., 3., 4., 7., 8., 9.]) + + """ + if keep not in ("last", "first"): + raise ValueError(f"`keep` must be either 'last' or 'first', got {keep!r}") + if isinstance(obj, DataCollection): + return obj.map(lambda da: trim_overlaps(da, keep, dim, tolerance)) + axis = obj.get_axis_num(dim) + dim = obj.dims[axis] + parts = split(obj, "overlaps", dim, tolerance) + if len(parts) == 1: + return obj + + # The parts are walked in order of decreasing precedence, each keeping only + # what no higher-precedence part already claimed. Claims accumulate as a set + # of spans rather than a single watermark, because a part may be *enveloped* + # in a lower-precedence one — the covering part then keeps a run on each + # side of it, and a watermark, which can only trim an end, would drop the + # far side along with the overlap. This also resolves a part wholly covered + # by a neighbour (it contributes nothing, while the part beyond it is still + # compared against that neighbour) and chains of mutually overlapping parts. + kept = [] + claimed = [] + for part in reversed(parts) if keep == "last" else parts: + coord = part[dim] + for start, stop in _uncovered(coord, claimed): + kept.append(part.isel({dim: slice(start, stop)})) + claimed = _claim(claimed, coord[0].values, coord[-1].values) + return concat(kept, dim, tolerance) + + +def _uncovered(coord, claimed): + """Index ranges of *coord* that no span of *claimed* covers. + + *claimed* is a list of disjoint ``(first, last)`` value pairs in ascending + order, both bounds inclusive. Bounds are resolved through the coordinate's + own label look-up, so nothing is materialised. + """ + runs = [] + cursor = 0 + for first, last in claimed: + # `to_index` clamps: a bound past either end of the coordinate resolves + # to the full length or to zero rather than raising + stop = coord.to_index(slice(None, first), endpoint=False).stop + if stop > cursor: + runs.append((cursor, stop)) + cursor = max(cursor, coord.to_index(slice(None, last)).stop) + if cursor < len(coord): + runs.append((cursor, len(coord))) + return runs + + +def _claim(claimed, first, last): + """Add the span ``(first, last)`` to the disjoint, ascending list *claimed*.""" + out = [] + for start, stop in claimed: + if stop < first or start > last: + out.append((start, stop)) + else: + first, last = min(first, start), max(last, stop) + out.append((first, last)) + out.sort() + return out + + def align(*objs): """ Given any number of data arrays, returns new objects with aligned dimensions. diff --git a/xdas/dask/__init__.py b/xdas/dask/__init__.py deleted file mode 100644 index 04046f48..00000000 --- a/xdas/dask/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -""" -Dask integration helpers for xdas HDF5 files. - -Serializes and deserializes dask arrays inside xdas HDF5 files. -""" - -__all__ = ["create_variable", "dumps", "loads"] - -from .core import create_variable, dumps, loads diff --git a/xdas/dask/core.py b/xdas/dask/core.py deleted file mode 100644 index ab6800df..00000000 --- a/xdas/dask/core.py +++ /dev/null @@ -1,90 +0,0 @@ -""" -Functions to store and restore dask arrays as HDF5 variables. - -Uses msgpack serialization of the dask task graph. -""" - -import numpy as np -from dask.array import Array - -from . import serial - - -def create_variable(arr, file, name, dims=None, dtype=None): - """ - Serialize *arr* and store it as an HDF5 variable attribute. - - Parameters - ---------- - arr : dask.array.Array - Dask array to persist. - file : netCDF4-like file handle - Open file in which to create the variable. - name : str - Variable name inside the file. - dims : sequence of str, optional - Dimension names for the variable. - dtype : dtype-like, optional - Data type for the variable. - - Returns - ------- - variable - The newly created file variable. - """ - variable = file.create_variable(name, dims, dtype) - variable.attrs.update({"__dask_array__": np.frombuffer(dumps(arr), "uint8")}) - return variable - - -def dumps(arr): - """Serialize a dask array.""" - return serial.dumps(to_dict(arr)) - - -def loads(data): - """Deserialize a dask array.""" - return from_dict(serial.loads(data)) - - -def to_dict(arr): - """Convert a dask array to a dictionary.""" - graph = arr.__dask_graph__().cull(arr.__dask_keys__()) - graph = fuse(graph) - return { - "dask": graph, # TODO: fuse then encode fails... - "name": arr.name, - "chunks": arr.chunks, - "dtype": str(arr.dtype), - } - - -def from_dict(dct): - """Convert a dictionary to a dask array.""" - return Array(**dct) - - -def fuse(graph): - """Simpligy a graph by grouping intermediate empty computations.""" - dsk = {} - ignore = set() - for key, computation in graph.items(): - if key in ignore: - continue - while iskey(computation) and computation in graph: - ignore.add(computation) - computation = graph[computation] - dsk[key] = computation - return dsk - - -def iskey(obj): - """Return ``True`` if *obj* looks like a dask graph key (string or ``(str, int…)`` tuple).""" - if isinstance(obj, str): - return len(obj) > 0 - return ( - isinstance(obj, tuple) - and len(obj) > 1 - and isinstance(obj[0], str) - and all(isinstance(index, int) for index in obj[1:]) - ) diff --git a/xdas/dask/serial.py b/xdas/dask/serial.py deleted file mode 100644 index cf03ead6..00000000 --- a/xdas/dask/serial.py +++ /dev/null @@ -1,97 +0,0 @@ -""" -msgpack-based serialization for dask task graphs. - -Handles tuples, slices, callables, ``methodcaller``, and ``itemgetter`` -objects. -""" - -import importlib - -import msgpack -from dask.utils import itemgetter, methodcaller - -codes = { - "tuple": 1, - "slice": 2, - "callable": 3, - "methodcaller": 4, - "itemgetter": 5, -} - - -def encode(obj): - """ - Msgpack *default* hook — encode non-native types as :class:`msgpack.ExtType`. - - Handles ``tuple``, ``slice``, ``callable``, :class:`methodcaller`, and - :class:`itemgetter`. - - Parameters - ---------- - obj : object - Object to encode. - - Returns - ------- - msgpack.ExtType - """ - if isinstance(obj, tuple): - code = codes["tuple"] - obj = list(obj) - elif isinstance(obj, slice): - code = codes["slice"] - obj = {"start": obj.start, "stop": obj.stop, "step": obj.step} - elif isinstance(obj, methodcaller): - code = codes["methodcaller"] - obj = obj.method - elif isinstance(obj, itemgetter): - code = codes["itemgetter"] - obj = obj.index - elif callable(obj): - code = codes["callable"] - obj = {"module": obj.__module__, "name": obj.__name__} - else: - raise TypeError(f"Cannot encode object of type {type(obj)}") - data = dumps(obj) - return msgpack.ExtType(code, data) - - -def decode(code, data): - """ - Msgpack *ext_hook* — decode an :class:`msgpack.ExtType` back to the original object. - - Parameters - ---------- - code : int - Extension type code (one of the values in :data:`codes`). - data : bytes - Raw msgpack bytes for the payload. - - Returns - ------- - object - The decoded Python object. - """ - obj = loads(data) - if code == codes["tuple"]: - return tuple(obj) - elif code == codes["slice"]: - return slice(obj["start"], obj["stop"], obj["step"]) - elif code == codes["callable"]: - return getattr(importlib.import_module(obj["module"]), obj["name"]) - elif code == codes["methodcaller"]: - return methodcaller(obj) - elif code == codes["itemgetter"]: - return itemgetter(obj) - else: - raise ValueError(f"Unknown code {code}") - - -def dumps(obj): - """Serialize *obj* to msgpack bytes, encoding extension types via :func:`encode`.""" - return msgpack.dumps(obj, default=encode, strict_types=True) - - -def loads(obj): - """Deserialize msgpack *obj* bytes, restoring extension types via :func:`decode`.""" - return msgpack.loads(obj, strict_map_key=False, ext_hook=decode) diff --git a/xdas/io/__init__.py b/xdas/io/__init__.py index 81bff4fe..95622685 100644 --- a/xdas/io/__init__.py +++ b/xdas/io/__init__.py @@ -1,8 +1,9 @@ """ I/O subsystem: plugin-based :class:`Engine` registry and concrete engines. -Supports xdas native, ASN, APSensing, Febus, MiniSEED, ProdML, Silixa, and -Terra15 formats. +Supports xdas native, ASN, APSensing, Febus, ProdML, Silixa and Terra15 +formats, plus everything ObsPy reads (MiniSEED, SAC, GSE2, ...). The legacy +`"miniseed"` engine is kept for stored views written by it. """ __all__ = [ @@ -13,11 +14,23 @@ "febus", "get_free_port", "miniseed", + "obspy", "prodml", "silixa", "terra15", "xdas", ] -from . import apsensing, asn, febus, miniseed, prodml, silixa, terra15, xdas +from . import apsensing, asn, febus, prodml, silixa, terra15, xdas from .core import AutoEngine, Engine, get_free_port + +# isort: split +# The two obspy-based engines are imported last, and `obspy` before +# `miniseed`. `AutoEngine` tries engines in registration order: these two open +# far more than the format-specific ones, so they must come after them — a +# promiscuous engine placed early would shadow the others — and the new engine +# must be reached before the legacy one it replaces. +from . import obspy + +# isort: split +from . import miniseed diff --git a/xdas/io/core.py b/xdas/io/core.py index 59ee5273..8de7afe2 100644 --- a/xdas/io/core.py +++ b/xdas/io/core.py @@ -17,7 +17,7 @@ class Engine: The Engine class provides a plugin architecture for reading and writing various file formats. Each Engine subclass corresponds to a specific file format (e.g., - "xdas", "asn", "miniseed") and implements methods to open and save DataArray or + "xdas", "asn", "obspy") and implements methods to open and save DataArray or DataCollection objects. Engines are registered in a class-level registry using the `__init_subclass__` hook, @@ -197,6 +197,10 @@ class AutoEngine(Engine): - The first engine that successfully opens the file is used - If all engines fail, an informative error message is raised + Registration order therefore settles which engine wins when several read + the same file: `"obspy"` is registered before the legacy `"miniseed"`, and + both after the format-specific engines. + Parameters ---------- vtype : str, optional @@ -232,21 +236,59 @@ class AutoEngine(Engine): def open_dataarray(self, fname): """Try each registered engine in order and return the first successful result.""" - for engine in self._ordered_engines(): + opened, out = self._try_engines("open_dataarray", fname) + if opened: + return out + raise ValueError(self._failure_message(fname, out)) + + def open_datacollection(self, fname): + """Try each registered engine in order and return the first collection. + + Raises :exc:`NotImplementedError` when no engine describes *fname* as a + collection, so that callers fall back to opening it as a data array the + same way they do for a named engine. + """ + opened, out = self._try_engines("open_datacollection", fname) + if opened: + return out + raise NotImplementedError( + self._failure_message(fname, out, " as a data collection") + ) + + def _try_engines(self, method, fname): + """Return ``(True, result)`` from the first engine that opens *fname*. + + On failure, returns ``(False, refusals)``: the ``{engine: error}`` of + every engine that recognised the file far enough to say something about + it. An engine that offers neither this shape nor the asked vtype or + ctype says nothing about the file and is left out. No single refusal is + *the* reason — the engines are tried in a cache-warmed order, not in + order of likelihood — so they are all reported. + """ + refusals = {} + for name in self._ordered_engines(): try: - out = Engine[engine](vtype=self.vtype, ctype=self.ctype).open_dataarray( - fname - ) - AutoEngine._last_successful_engine = engine - return out - except Exception: # noqa: BLE001, S112 - try the next engine + engine = Engine[name](vtype=self.vtype, ctype=self.ctype) + out = getattr(engine, method)(fname) + except NotImplementedError: continue + except Exception as exc: # noqa: BLE001 - try the next engine + refusals[name] = exc + continue + AutoEngine._last_successful_engine = name + return True, out + return False, refusals + + def _failure_message(self, fname, refusals=None, suffix=""): message = f"no engine could open the file '{fname}'" if self.ctype is not None: message += f" with ctype '{self.ctype}'" if self.vtype is not None: message += f" with vtype '{self.vtype}'" - raise ValueError(message) + message += suffix + for name, error in (refusals or {}).items(): + message += f"\n {name}: {type(error).__name__}: {error}" + return message def _ordered_engines(self): return [self._last_successful_engine] + [ diff --git a/xdas/io/miniseed.py b/xdas/io/miniseed.py index 7c9546f3..8ceede1d 100644 --- a/xdas/io/miniseed.py +++ b/xdas/io/miniseed.py @@ -1,20 +1,45 @@ -"""I/O engine for MiniSEED files via ObsPy (:class:`MiniSEEDEngine`).""" +"""Legacy MiniSEED engine (:class:`MiniSEEDEngine`), kept for stored views. + +This is the engine `engine="miniseed"` named before :mod:`xdas.io.obspy` +replaced it, preserved verbatim so that manifests written by it keep decoding +and code written against it keeps running. It describes a whole file as one +tile of stacked channels, which it classifies at scan time as *synchronized* +(all traces share one time coordinate) or *unsynchronized* (the time axis is +the concatenation of the first channel's segments, every other channel assumed +to match), and refuses anything else — a file holding two sampling rates, for +instance. + +New code should use the `"obspy"` engine, which reads all of those and emits +one lazy data array per ObsPy trace. Both engines take part in format +auto-detection, `"obspy"` first: this one is reached only for a file the new +engine cannot describe as a single data array, which is exactly the shape +:func:`xdas.open_dataarray` was asked for. +""" from typing import ClassVar import numpy as np import obspy -from ..coordinates import ( - AxisCoordinate, - Coordinate, - Coordinates, - get_sampling_interval, -) +from ..coordinates import AxisCoordinate, Coordinate, Coordinates from ..core import DataArray, concat_coords from ..virtual import TileArray from .core import Engine +# the stream converters and the band-code table were never miniSEED-specific +# and now live with the obspy engine; re-exported so that imports from this +# module keep working +from .obspy import from_stream, get_band_code, get_dtype, to_stream + +__all__ = [ + "MiniSEEDEngine", + "from_stream", + "get_band_code", + "get_time_coord", + "to_stream", + "uniquifiy", +] + class MiniSEEDEngine(Engine, name="miniseed"): """ @@ -59,7 +84,10 @@ def read_header(self, path): ctype = self.ctype["time"] st = obspy.read(path, headonly=True) - dtype = uniquifiy(tr.data.dtype for tr in st) + # from the encoding, not from `tr.data`, which `headonly=True` leaves + # an empty float64 array: a STEIM-compressed file decodes to int32 and + # the tile array checks the scanned type against every decoded tile + dtype = uniquifiy(get_dtype(tr) for tr in st) if not isinstance(dtype, np.dtype): # pragma: no cover raise ValueError("All traces must have the same dtype") @@ -147,84 +175,6 @@ def load_tile(path, selection, *, method="synchronized", ignore_last_sample=Fals return data[selection] -def to_stream( - da, - network="NET", - station="DAS{:05}", - location="00", - channel="{:1}N1", - dim=None, -): - """ - Convert a 2-D :class:`DataArray` to an :class:`obspy.Stream`. - - Parameters - ---------- - da : DataArray - 2-D array with one time and one distance/channel dimension. - network, station, location, channel : str - SEED identifiers. *station* and *channel* may contain ``{:...}`` - format specs that are filled with the channel index. - dim : dict, optional - ``{distance_dim: time_dim}`` mapping. Defaults to ``{"last": "first"}``. - - Returns - ------- - obspy.Stream - """ - if dim is None: - dim = {"last": "first"} - dimdist, dimtime = dim.copy().popitem() - if not da.ndim == 2: - raise ValueError("the data array must be 2D") - starttime = obspy.UTCDateTime(str(da[dimtime][0].values)) - delta = get_sampling_interval(da, dimtime) - band_code = get_band_code(1.0 / delta) - if "{" in channel and "}" in channel: - channel = channel.format(band_code) - header = { - "network": network, - "location": location, - "channel": channel, - "starttime": starttime, - "delta": delta, - } - return obspy.Stream( - [ - obspy.Trace( - data=np.ascontiguousarray(da.isel({dimdist: idx}).values), - header=header | {"station": station.format(idx + 1)}, - ) - for idx in range(len(da[dimdist])) - ] - ) - - -def from_stream(st, dims=("channel", "time")): - """ - Convert an :class:`obspy.Stream` to a :class:`DataArray`. - - Parameters - ---------- - st : obspy.Stream - Homogeneous stream (all traces must share start time and sample rate). - dims : tuple of str, optional - Dimension names for the output array. - - Returns - ------- - DataArray - """ - data = np.stack([tr.data for tr in st]) - channel = [tr.id for tr in st] - # Regular by construction from the stream's own sample rate, at ns - # resolution so a `to_stream` round trip preserves the coordinate. - t0 = np.datetime64(st[0].stats.starttime.datetime) - dt = np.rint(1e6 * st[0].stats.delta).astype("m8[us]").astype("m8[ns]") - time = Coordinate["interpolated"].from_block(t0, st[0].stats.npts, dt, dim=dims[1]) - return DataArray(data, {dims[0]: channel, dims[1]: time}) - - def get_time_coord(tr, ignore_last_sample, ctype): """Build a :class:`Coordinate` for the time axis of trace *tr*.""" t0 = np.datetime64(tr.stats.starttime) @@ -241,14 +191,3 @@ def uniquifiy(seq): return seq[0] else: return seq - - -def get_band_code(sampling_rate): - """Return the SEED band code character for *sampling_rate* (Hz).""" - band_code = ["T", "P", "R", "U", "V", "L", "M", "B", "H", "C", "F"] - limits = [0.000001, 0.00001, 0.0001, 0.001, 0.01, 0.1, 1, 10, 80, 250, 1000, 5000] - index = np.searchsorted(limits, sampling_rate, "right") - 1 - if index < 0 or index >= len(band_code): - return "X" - else: - return band_code[index] diff --git a/xdas/io/obspy.py b/xdas/io/obspy.py new file mode 100644 index 00000000..4da8318e --- /dev/null +++ b/xdas/io/obspy.py @@ -0,0 +1,387 @@ +"""I/O engine for the formats ObsPy reads (:class:`ObsPyEngine`). + +The engine is named for the library, not for a format: decoding is +:func:`obspy.read`, so every format ObsPy supports — MiniSEED, SAC, GSE2, +SEG-2 and the rest — goes through it. It replaces +:mod:`xdas.io.miniseed`, which is kept alongside it for the views it wrote. +""" + +from typing import ClassVar + +import numpy as np +import obspy + +from ..coordinates import Coordinate, get_sampling_interval +from ..core import DataArray, DataCollection +from ..virtual import TileArray +from .core import Engine + +#: The blank location code, as FDSN spells it. ObsPy returns ``""``, which +#: cannot be a netCDF group name, and ObsPy's own FDSN client performs the +#: same mapping. +BLANK_LOCATION = "--" + +#: The levels of the SEED hierarchy, outermost first. +LEVELS = ("network", "station", "location", "channel") + +#: Element type each MiniSEED encoding *decodes to*, which is not the type it +#: was written from: libmseed unpacks every integer encoding to ``int32``, +#: whatever its on-disk width. Needed because ``headonly=True`` leaves +#: ``tr.data`` an empty ``float64`` array, so the decoded type cannot be read +#: off it. +MSEED_DTYPES = { + "ASCII": np.dtype("S1"), + "INT16": np.dtype("int32"), + "INT32": np.dtype("int32"), + "FLOAT32": np.dtype("float32"), + "FLOAT64": np.dtype("float64"), + "STEIM1": np.dtype("int32"), + "STEIM2": np.dtype("int32"), + "GEOSCOPE24": np.dtype("float32"), + "GEOSCOPE16_3": np.dtype("float32"), + "GEOSCOPE16_4": np.dtype("float32"), + "CDSN": np.dtype("int32"), + "SRO": np.dtype("int32"), + "DWWSSN": np.dtype("int32"), +} + + +class ObsPyEngine(Engine, name="obspy"): + """ + Engine for the file formats ObsPy reads, as lazy tile-backed data arrays. + + The engine mirrors :func:`obspy.read`: each contiguous + :class:`obspy.Trace` becomes one lazy one-dimensional + :class:`~xdas.DataArray`, and the collection mirrors the + :class:`obspy.Stream`, nested on the four levels of the SEED hierarchy: + + .. code-block:: text + + network -> station -> location -> channel -> trace -> DataArray + + Merging contiguous traces, separating acquisition epochs and moving gaps + into the time coordinate are not the engine's job: + :func:`~xdas.combine_by_coords` does all three, and + :func:`~xdas.open` calls it. Overlaps are resolved by + :func:`~xdas.trim_overlaps`. + + Parameters + ---------- + vtype : str, optional + The virtualization type to use. Default to "tiles". + ctype : str or dict, optional + The coordinate type to use for the time axis. Default to "interpolated". + + """ + + _supported_vtypes: ClassVar[list] = ["tiles"] + _supported_ctypes: ClassVar[dict] = { + "time": ["interpolated", "sampled", "dense"], + } + + def open_datacollection(self, fname): + """Return the traces of *fname* as a collection nested on the SEED hierarchy.""" + st = obspy.read(fname, headonly=True) + # libmseed's trace-list assembly returns traces in no useful order; + # `sort` puts them in (id, starttime) order + st.sort() + tree = {} + pointers = set() + for tr in st: + pointer = get_pointer(tr) + key = tuple(pointer.values()) + if key in pointers: + raise ValueError( + f"{fname} holds two traces sharing every identifier and both " + f"time bounds ({pointer}); nothing content-free separates them" + ) + pointers.add(key) + branch = tree + for level in LEVELS[:-1]: + branch = branch.setdefault(pointer[level], {}) + branch.setdefault(pointer["channel"], []).append( + self._open_trace(fname, tr, pointer) + ) + return nest(tree, LEVELS) + + def open_dataarray(self, fname): + """Return the unique trace of *fname* as a lazy tile-backed data array.""" + st = obspy.read(fname, headonly=True) + if len(st) != 1: + raise ValueError( + f"{fname} holds {len(st)} traces, not one; open it with " + "`open_datacollection` (or `open`, which combines the result)" + ) + return self._open_trace(fname, st[0], get_pointer(st[0])) + + def _open_trace(self, fname, tr, pointer): + """Build the lazy data array of the single trace *tr* of *fname*.""" + data = TileArray.from_tiles( + str(fname), + (tr.stats.npts,), + get_dtype(tr), + {"name": "obspy"}, + **pointer, + ) + coords = {level: (None, pointer[level]) for level in LEVELS} + coords["time"] = get_time_coord(tr, self.ctype["time"]) + # the id ObsPy prints, with the location code normalized as the tree + # keys have it + name = ".".join(pointer[level] for level in LEVELS) + return DataArray(data, coords, name=name) + + @staticmethod + def load_tile( + path, + selection, + *, + network, + station, + location, + channel, + starttime, + endtime, + ): + """Read a source selection of *path*, decoding with :func:`obspy.read`. + + The tile is addressed by the data's own address — the four SEED + identifiers plus both time bounds — never by a position in the stream: + ObsPy's trace *count* comes from libmseed's segmentation policy, so an + index would designate a different trace after any re-segmenting version + bump, silently. The contiguous runs of the selected channel are joined + before the span is looked up, so the pointer resolves whatever record + boundaries the reader drew. + + Manifests written before the engine was renamed name the `"miniseed"` + engine instead, and are decoded by + :meth:`~xdas.io.miniseed.MiniSEEDEngine.load_tile`. + """ + st = obspy.read(path).select( + network=network, + station=station, + location="" if location == BLANK_LOCATION else location, + channel=channel, + ) + runs = join_contiguous(st) + # a run may legitimately be longer than the tile — a re-segmenting + # reader joins the same span from different records, and two traces may + # share a start and differ in length. The smallest covering run is the + # one the pointer named; ties mean genuine duplicates. + covering = [ + run + for run in runs + if run["start"] - run["delta"] // 2 <= starttime + and endtime <= run["end"] + run["delta"] // 2 + ] + if covering: + shortest = min(len(run["data"]) for run in covering) + covering = [run for run in covering if len(run["data"]) == shortest] + if len(covering) != 1: + raise ValueError( + f"{len(covering)} contiguous runs of " + f"{network}.{station}.{location}.{channel} cover " + f"[{np.datetime64(starttime, 'ns')}, {np.datetime64(endtime, 'ns')}] " + f"in {path}; exactly one is required" + ) + (run,) = covering + offset = round((starttime - run["start"]) / run["delta"]) + npts = round((endtime - starttime) / run["delta"]) + 1 + return run["data"][offset : offset + npts][selection] + + +def nest(tree, levels): + """Wrap the nested dict *tree* of trace lists into named collection levels.""" + name, *rest = levels + if rest: + data = {key: nest(value, rest) for key, value in tree.items()} + else: + # one element per ObsPy `Trace`: this is the faithful `obspy.read` + # mirror. Once combined, each element is an acquisition epoch instead + # and `combine_by_coords` renames the level accordingly. + data = {key: DataCollection(value, "trace") for key, value in tree.items()} + return DataCollection(data, name) + + +def get_pointer(tr): + """Return the columns that address *tr* in its file. + + The four SEED identifiers are kept apart rather than joined into one + ``"NET.STA.LOC.CHA"`` string so that the tile manifest folds each + independently — a `concat` along `channel` unfolds only that field — and so + that they map one to one onto :meth:`obspy.Stream.select`. + + ``starttime`` alone does not identify a trace: two traces can share all + four identifiers *and* a start time and still hold different data. + ``endtime`` completes the key. Neither is recoverable from the tile + geometry, which describes the *view* once the array is sliced while the + pointer must keep naming the source. + """ + return { + "network": tr.stats.network, + "station": tr.stats.station, + "location": tr.stats.location or BLANK_LOCATION, + "channel": tr.stats.channel, + "starttime": tr.stats.starttime.ns, + "endtime": tr.stats.endtime.ns, + } + + +def get_dtype(tr): + """Return the element type *tr* decodes to. + + Under ``headonly=True`` ``tr.data`` is an empty ``float64`` array whatever + the file holds, so the MiniSEED encoding is authoritative when present. + The fallback is right for the formats that ignore ``headonly`` and decode + fully anyway; those simply pay a slower scan. + """ + encoding = getattr(tr.stats, "mseed", {}).get("encoding") + if encoding in MSEED_DTYPES: + return MSEED_DTYPES[encoding] + return tr.data.dtype + + +def get_time_coord(tr, ctype): + """Build the time :class:`~xdas.Coordinate` of trace *tr*. + + Regular by construction: ObsPy already splits a trace at its gaps. + """ + t0 = np.datetime64(tr.stats.starttime.ns, "ns") + dt = np.rint(1e6 * tr.stats.delta).astype("m8[us]").astype("m8[ns]") + return Coordinate[ctype].from_block(t0, tr.stats.npts, dt, dim="time") + + +def join_contiguous(traces): + """Group *traces* into sample-exact contiguous runs. + + The legitimate half of :meth:`obspy.Stream._cleanup`, implemented here + rather than called through that private method: traces that continue each + other to the sample are concatenated and nothing else is touched — no gap + filling, no overlap arbitration, no masking. + + Returns + ------- + list of dict + One entry per run, with its ``start`` and ``end`` in nanoseconds, its + sampling ``delta`` in nanoseconds, and its ``data``. + """ + runs = [] + for tr in sorted(traces, key=lambda tr: (tr.stats.starttime.ns, tr.stats.npts)): + delta = round(tr.stats.delta * 1e9) + start = tr.stats.starttime.ns + if runs and runs[-1]["delta"] == delta and start == runs[-1]["stop"]: + runs[-1]["stop"] += tr.stats.npts * delta + runs[-1]["chunks"].append(tr.data) + else: + runs.append( + { + "start": start, + "stop": start + tr.stats.npts * delta, + "delta": delta, + "chunks": [tr.data], + } + ) + return [ + { + "start": run["start"], + "end": run["stop"] - run["delta"], + "delta": run["delta"], + "data": ( + run["chunks"][0] + if len(run["chunks"]) == 1 + else np.concatenate(run["chunks"]) + ), + } + for run in runs + ] + + +def to_stream( + da, + network="NET", + station="DAS{:05}", + location="00", + channel="{:1}N1", + dim=None, +): + """ + Convert a 2-D :class:`DataArray` to an :class:`obspy.Stream`. + + Parameters + ---------- + da : DataArray + 2-D array with one time and one distance/channel dimension. + network, station, location, channel : str + SEED identifiers. *station* and *channel* may contain ``{:...}`` + format specs that are filled with the channel index. The blank + location code is spelled ``"--"`` in xdas and ``""`` in ObsPy; either + is accepted here. + dim : dict, optional + ``{distance_dim: time_dim}`` mapping. Defaults to ``{"last": "first"}``. + + Returns + ------- + obspy.Stream + """ + if dim is None: + dim = {"last": "first"} + dimdist, dimtime = dim.copy().popitem() + if not da.ndim == 2: + raise ValueError("the data array must be 2D") + starttime = obspy.UTCDateTime(str(da[dimtime][0].values)) + delta = get_sampling_interval(da, dimtime) + band_code = get_band_code(1.0 / delta) + if "{" in channel and "}" in channel: + channel = channel.format(band_code) + header = { + "network": network, + "location": "" if location == BLANK_LOCATION else location, + "channel": channel, + "starttime": starttime, + "delta": delta, + } + return obspy.Stream( + [ + obspy.Trace( + data=np.ascontiguousarray(da.isel({dimdist: idx}).values), + header=header | {"station": station.format(idx + 1)}, + ) + for idx in range(len(da[dimdist])) + ] + ) + + +def from_stream(st, dims=("channel", "time")): + """ + Convert an :class:`obspy.Stream` to a :class:`DataArray`. + + Parameters + ---------- + st : obspy.Stream + Homogeneous stream (all traces must share start time and sample rate). + dims : tuple of str, optional + Dimension names for the output array. + + Returns + ------- + DataArray + """ + data = np.stack([tr.data for tr in st]) + channel = [tr.id for tr in st] + # Regular by construction from the stream's own sample rate, at ns + # resolution so a `to_stream` round trip preserves the coordinate. From + # `.ns`, not `.datetime`, which truncates the start time to microseconds. + t0 = np.datetime64(st[0].stats.starttime.ns, "ns") + dt = np.rint(1e6 * st[0].stats.delta).astype("m8[us]").astype("m8[ns]") + time = Coordinate["interpolated"].from_block(t0, st[0].stats.npts, dt, dim=dims[1]) + return DataArray(data, {dims[0]: channel, dims[1]: time}) + + +def get_band_code(sampling_rate): + """Return the SEED band code character for *sampling_rate* (Hz).""" + band_code = ["T", "P", "R", "U", "V", "L", "M", "B", "H", "C", "F"] + limits = [0.000001, 0.00001, 0.0001, 0.001, 0.01, 0.1, 1, 10, 80, 250, 1000, 5000] + index = np.searchsorted(limits, sampling_rate, "right") - 1 + if index < 0 or index >= len(band_code): + return "X" + else: + return band_code[index] diff --git a/xdas/io/xdas.py b/xdas/io/xdas.py index a2e0858b..5a23d431 100644 --- a/xdas/io/xdas.py +++ b/xdas/io/xdas.py @@ -2,11 +2,17 @@ I/O engine for the native xdas HDF5/NetCDF4 format (:class:`XdasEngine`). Supports :class:`DataArray`, :class:`DataSequence`, and :class:`DataMapping`. + +Whatever the shape of the object, the file is opened once in each +direction: reading walks a single :func:`xarray.open_datatree`, writing +lands every group's metadata through a single +:meth:`xarray.DataTree.to_netcdf` and the data variables through one +writable `h5netcdf` handle. One handle matters on write: every writable +`h5netcdf` open walks the whole file to find the next free dimension +id, so per-array opens made saving a collection quadratic in its size. """ -import json import os -import warnings from pathlib import Path from typing import ClassVar @@ -14,13 +20,12 @@ import h5py import hdf5plugin # noqa import xarray as xr -from dask.array import Array as DaskArray from ..coordinates import Coordinates from ..core import DataArray, DataCollection, DataMapping, DataSequence -from ..dask import create_variable, loads +from ..core.datacollection import as_sequence_if_positional from ..virtual import TileArray, VirtualBackend -from ..virtual.tiles import TILES_GROUP +from ..virtual.tiles import TILING from .core import Engine @@ -105,46 +110,74 @@ def open_dataarray(fname, group=None, vtype=None): if isinstance(fname, Path): fname = str(fname) - # read metadata - with xr.open_dataset( - fname, group=group, engine="h5netcdf", decode_timedelta=False, phony_dims="sort" - ) as dataset: - # check file format - if not ( - "Conventions" in dataset.attrs and "CF" in dataset.attrs["Conventions"] - ): - raise TypeError( - "file format not recognized. please provide the file format " - "with the `engine` keyword argument" - ) + # one open covers the data array and any tile manifest beside it. + # "access" is xarray's own default and silences its warning; "sort" + # would rescan every group of the file on each open. + with xr.open_datatree( + fname, + group=group, + engine="h5netcdf", + decode_timedelta=False, + phony_dims="access", + ) as node: + return _read_dataarray(node, fname, group, vtype) + + +def _read_dataarray(node, fname, group=None, vtype=None): + """Build a :class:`DataArray` from the open tree *node* holding it. + + Parameters + ---------- + node : xarray.DataTree + The open node holding the data array (and its tile manifest as a + child, if any). + fname : str + Path of the file the node was opened from, reopened by the hdf5 + virtual backend. + group : str, optional + Location of *node* within the file, needed by the same backend. + vtype : str, optional + Virtualization backing of the returned data (see + :func:`open_dataarray`). + """ + dataset = node.dataset - # identify the "main" data array - if len(dataset) == 1: - name = next(iter(dataset.keys())) + # check file format + if not ("Conventions" in dataset.attrs and "CF" in dataset.attrs["Conventions"]): + raise TypeError( + "file format not recognized. please provide the file format " + "with the `engine` keyword argument" + ) + + # identify the "main" data array + if len(dataset) == 1: + name = next(iter(dataset.keys())) + else: + data_vars = { + key: var + for key, var in dataset.items() + if any("coordinate" in attr for attr in var.attrs) + } + if len(data_vars) == 1: + name = next(iter(data_vars.keys())) else: - data_vars = { - key: var - for key, var in dataset.items() - if any("coordinate" in attr for attr in var.attrs) - } - if len(data_vars) == 1: - name = next(iter(data_vars.keys())) - else: - raise ValueError("several possible data arrays detected") + raise ValueError("several possible data arrays detected") - # read coordinates - coords = Coordinates._from_dataset(dataset, name) + # read coordinates + coords = Coordinates._from_dataset(dataset, name) # read data - if "__tile_array__" in dataset[name].attrs: - spec = json.loads(dataset[name].attrs.pop("__tile_array__")) - location = TILES_GROUP if group is None else f"{group}/{TILES_GROUP}" - with xr.open_dataset(fname, group=location, engine="h5netcdf") as manifest: - manifest = manifest.load() - # the placeholder variable carries the dtype; the spec only the engine - data = TileArray(manifest, dataset[name].dtype, spec["engine"]) - elif "__dask_array__" in dataset[name].attrs: - data = loads(dataset[name].attrs.pop("__dask_array__")) + attrs = dataset[name].attrs + if TILING in attrs: + # the placeholder points at the group describing its tiling and + # projects the array it stands for: its type must be that array's + manifest = node[attrs[TILING]].to_dataset(inherit=False).load() + data = TileArray(manifest) + if data.dtype != dataset[name].dtype: + raise ValueError( + f"the placeholder is {dataset[name].dtype} where its manifest " + f"records {data.dtype}" + ) else: with h5py.File(fname) as file: if group: @@ -154,14 +187,16 @@ def open_dataarray(fname, group=None, vtype=None): variable ) + # the file has two reservation registries: CF's plain words, which + # the coordinates have consumed, and xdas's dunder-fenced ones + attrs = { + key: value + for key, value in attrs.items() + if not (key.startswith("__") and key.endswith("__")) + } + # pack everything - return DataArray( - data, - coords, - dataset[name].dims, - name, - None if dataset[name].attrs == {} else dataset[name].attrs, - ) + return DataArray(data, coords, dataset[name].dims, name, attrs or None) def save_dataarray( @@ -189,18 +224,41 @@ def save_dataarray( """ if isinstance(fname, Path): fname = str(fname) + _save_tree({group: da}, fname, mode, virtual, encoding, create_dirs) - if virtual is None: - virtual = isinstance(da.data, (VirtualBackend, DaskArray)) - # initialize - dataset = xr.Dataset(attrs={"Conventions": "CF-1.9"}) - variable_attrs = {} if da.attrs is None else da.attrs - variable_name = "__values__" if da.name is None else da.name +def _save_tree(leaves, fname, mode, virtual, encoding, create_dirs): + """Write *leaves* (``{location or None: DataArray}``) in two passes. - # prepare metadata - for coord in da.coords.values(): - dataset, variable_attrs = coord._to_dataset(dataset, variable_attrs) + First every group's metadata — coordinates and tile manifests — as + one :class:`xarray.DataTree`, then every data variable through a + single writable `h5netcdf` handle, since xarray cannot write the + virtual ones. *mode* applies to the first pass; the second appends. + """ + # prepare metadata: one tree node per group, plus per-leaf variable + # attributes and virtual-ness for the second pass + nodes = {} + entries = [] + for location, da in leaves.items(): + isvirtual = isinstance(da.data, VirtualBackend) if virtual is None else virtual + if isvirtual: + if encoding is not None: + raise ValueError("cannot use `encoding` in virtual mode") + if not isinstance(da.data, VirtualBackend): + raise ValueError( + "can only use `virtual=True` with a virtual array as data" + ) + dataset = xr.Dataset(attrs={"Conventions": "CF-1.13"}) + attrs = {} if da.attrs is None else dict(da.attrs) + for coord in da.coords.values(): + dataset, attrs = coord._to_dataset(dataset, attrs) + nodes["/" if location is None else location] = dataset + if isvirtual: + for relpath, sibling in da.data.sibling_datasets().items(): + nodes[relpath if location is None else f"{location}/{relpath}"] = ( + sibling + ) + entries.append((location, da, isvirtual, attrs)) # create parent directories if needed if create_dirs: @@ -208,66 +266,47 @@ def save_dataarray( if dirname: os.makedirs(dirname, exist_ok=True) - # write data - with h5netcdf.File(fname, mode=mode) as file: - # group - if group is not None and group not in file: - file.create_group(group) - file = file if group is None else file[group] - - # dims - file.dimensions.update(da.sizes) - - # variable - if not virtual: - encoding = {} if encoding is None else encoding - variable = file.create_variable( - variable_name, da.dims, da.dtype, data=da.values, **encoding + # write metadata, one public-API call for every group + xr.DataTree.from_dict(nodes).to_netcdf(fname, mode=mode, engine="h5netcdf") + + # write data variables, one writable open for the whole file + with h5netcdf.File(fname, mode="a") as file: + for location, da, isvirtual, attrs in entries: + target = file if location is None else file[location] + + # dims the metadata pass did not create (those carrying no + # coordinate variable) + target.dimensions.update( + { + dim: size + for dim, size in da.sizes.items() + if dim not in target.dimensions + } ) - else: - if encoding is not None: - raise ValueError("cannot use `encoding` with in virtual mode") - if isinstance(da.data, VirtualBackend): - variable = da.data.create_variable( - file, variable_name, da.dims, da.dtype - ) - elif isinstance(da.data, DaskArray): - warnings.warn( - "writing dask-backed virtual arrays is deprecated; the " - "tile-backed engines (xdas.virtual.tiles) replace them", - FutureWarning, - ) - variable = create_variable( - da.data, file, variable_name, da.dims, da.dtype - ) + + # variable + variable_name = "__values__" if da.name is None else da.name + if isvirtual: + variable = da.data.create_variable(target, variable_name, da.dims) else: - raise ValueError( - "can only use `virtual=True` with a virtual array as data" + # anything not virtual is materialized here, a dask array + # computed like any other lazy value + variable = target.create_variable( + variable_name, + da.dims, + da.dtype, + data=da.values, + **({} if encoding is None else encoding), ) - # attrs - if variable_attrs: - variable.attrs.update(variable_attrs) - - # write metadata - dataset.to_netcdf(fname, mode="a", group=group, engine="h5netcdf") - - # append what of the stored form outlives the variable (the tile manifest) - if virtual and isinstance(da.data, VirtualBackend): - da.data.finalize_save(fname, group) + # attrs + if attrs: + variable.attrs.update(attrs) def open_datacollection(fname, group=None): """Read a :class:`DataCollection` from *fname*, auto-detecting sequence vs. mapping.""" - dc = open_datamapping(fname, group) - try: - keys = [int(key) for key in dc] - except ValueError: - return dc - if set(keys) == set(range(len(keys))): - return DataSequence([dc[str(key)] for key in range(len(keys))], dc.name) - else: - return dc + return as_sequence_if_positional(open_datamapping(fname, group)) def save_datacollection( @@ -287,30 +326,40 @@ def open_datamapping(fname, group=None): if isinstance(fname, Path): fname = str(fname) - with h5py.File(fname, "r") as file: - if group is None: - group = file[next(iter(file.keys()))] - else: - group = file[group] - name = group.name.split("/")[-1] - if isinstance(group, h5py.Dataset): + # the whole collection in one open — walking the already open tree + # replaces one targeted reopen per data array + with xr.open_datatree( + fname, + engine="h5netcdf", + decode_timedelta=False, + phony_dims="access", + ) as tree: + node = tree if group is None else tree[group] + if group is None and not node.dataset.data_vars: + node = next(iter(node.children.values())) + # a collection node holds only groups; finding variables on it (or + # being handed a variable path) means the file is something else + if isinstance(node, xr.DataArray) or node.dataset.data_vars: raise ValueError( "it looks like you are trying to open a data array as a data collection." ) + return _read_datamapping(node, fname) + + +def _read_datamapping(node, fname): + """Build a :class:`DataMapping` from the open tree *node* holding it. + + A child node holding variables is a data array; one holding only + groups is a nesting level whose single child is a named collection. + """ + name = node.name + dm = DataMapping({}, name=None if name == "collection" else name) + for key, child in node.children.items(): + if child.dataset.data_vars: + dm[key] = _read_dataarray(child, fname, group=child.path) else: - if not isinstance(group, h5py.Group): # pragma: no cover - raise RuntimeError( - "something went wrong while opening the data collection." - ) - keys = list(group.keys()) - dm = DataMapping({}, name=None if name == "collection" else name) - for key in keys: - subgroup = group[key] - if _get_depth(subgroup) == 0: - dm[key] = DataArray.from_netcdf(fname, subgroup.name) - else: - subgroup = subgroup[next(iter(subgroup.keys()))] - dm[key] = DataCollection.from_netcdf(fname, subgroup.name) + subnode = next(iter(child.children.values())) + dm[key] = as_sequence_if_positional(_read_datamapping(subnode, fname)) return dm @@ -318,20 +367,30 @@ def save_datamapping( dm, fname, mode="w", group=None, virtual=None, encoding=None, create_dirs=False ): """Write :class:`DataMapping` *dm* to *fname*, writing each key as a separate group.""" + if isinstance(fname, Path): + fname = str(fname) if mode == "w" and group is None and os.path.exists(fname): os.remove(fname) - for key in dm: - name = dm.name if dm.name is not None else "collection" + leaves = _collect_leaves(dm, group) + if leaves: + _save_tree(leaves, fname, "a", virtual, encoding, create_dirs) + + +def _collect_leaves(dc, group): + """Flatten collection *dc* into ``{location: DataArray}`` under *group*.""" + if isinstance(dc, DataSequence): + dc = dc.to_mapping() + name = dc.name if dc.name is not None else "collection" + leaves = {} + for key in dc: location = "/".join([name, str(key)]) if group is not None: location = f"{group}/{location}" - if create_dirs: - dirname = os.path.dirname(fname) - if dirname: - os.makedirs(dirname, exist_ok=True) - dm[key].to_netcdf( - fname, mode="a", group=location, virtual=virtual, encoding=encoding - ) + if isinstance(dc[key], DataArray): + leaves[location] = dc[key] + else: + leaves.update(_collect_leaves(dc[key], location)) + return leaves def open_datasequence(fname, group=None): @@ -346,23 +405,3 @@ def save_datasequence( """Write :class:`DataSequence` *ds* to *fname* by converting to a mapping first.""" dm = ds.to_mapping() save_datamapping(dm, fname, mode, group, virtual, encoding, create_dirs) - - -def _get_depth(group): - """Nesting depth of *group*, ignoring any tile manifest it contains. - - A tile-backed data array keeps its manifest in a `TILES_GROUP` - sibling of its variables. That group must not count, or the data - array would look one level deeper than it is and be mistaken for a - nested collection. - """ - if not isinstance(group, h5py.Group): - raise ValueError("not a group") - depths = [0] - - def visit(name): - if TILES_GROUP not in name.split("/"): - depths.append(name.count("/")) - - group.visit(visit) - return max(depths) diff --git a/xdas/virtual/core.py b/xdas/virtual/core.py index 6bab281c..e08a240b 100644 --- a/xdas/virtual/core.py +++ b/xdas/virtual/core.py @@ -18,9 +18,9 @@ class VirtualBackend(ABC): - ``__getitem__`` selects lazily and ``__array__`` materializes; - :meth:`from_variable` wraps one stored HDF5 variable, so open paths dispatch ``VirtualBackend[vtype].from_variable(...)``; - - the :meth:`create_variable`/:meth:`finalize_save` pair writes - the backend's stored form, so save paths need no per-backend - branch. + - the :meth:`create_variable`/:meth:`sibling_datasets` pair + describes the backend's stored form, so save paths need no + per-backend branch. Beyond the contract the backends share no implementation and differ freely (blocking, concatenation). Subclasses register by @@ -110,25 +110,28 @@ def from_variable(cls, variable): """ @abstractmethod - def create_variable(self, file, name, dims=None, dtype=None): + def create_variable(self, file, name, dims=None): """ Write this array as variable *name* of an open h5netcdf *file*. The first half of the persistence contract — each backend writes its own stored form: an HDF5 virtual dataset for the - hdf5 backend, a placeholder variable carrying the engine - specification for the tiles backend. + hdf5 backend, a placeholder variable pointing at its manifest + for the tiles backend. The element type is the array's own. """ - def finalize_save(self, fname, group=None): + def sibling_datasets(self): """ - Append what of the stored form outlives the variable. - - The second half of the persistence contract, called once the - file handle of :meth:`create_variable` is closed. Default: the - variable is the whole stored form, nothing to append — the - tiles backend appends its manifest as a sibling group. + Return what of the stored form lives beside the variable. + + The second half of the persistence contract: a mapping from + group name, relative to the variable's group, to the + :class:`xarray.Dataset` to store there. The writer lands them + with the rest of the file's metadata. Default: the variable is + the whole stored form, nothing beside it — the tiles backend + returns its manifest under a ``__tiles__`` sibling group. """ + return {} # --- derived from the contract, shared by every backend --- diff --git a/xdas/virtual/hdf5.py b/xdas/virtual/hdf5.py index a044ab83..5993b38d 100644 --- a/xdas/virtual/hdf5.py +++ b/xdas/virtual/hdf5.py @@ -48,7 +48,7 @@ def from_variable(cls, variable): """ return VirtualSource(variable) - def create_variable(self, file, name, dims=None, dtype=None): + def create_variable(self, file, name, dims=None): """ Write this virtual array into *file* and register it as a named variable. @@ -60,8 +60,6 @@ def create_variable(self, file, name, dims=None, dtype=None): Variable name to create inside *file*. dims : sequence of str, optional Dimension names for the variable. - dtype : dtype-like, optional - Override data type for the variable. Returns ------- diff --git a/xdas/virtual/tiles.py b/xdas/virtual/tiles.py index 301c8800..25012c90 100644 --- a/xdas/virtual/tiles.py +++ b/xdas/virtual/tiles.py @@ -6,53 +6,80 @@ :class:`xarray.Dataset` (used as a container only) with one dimension ``tile_k`` per data axis: -- 1-D *geometry* variables place the grid along each axis ``k``: +- *geometry* variables place the grid along each axis ``k``: ``sizes_k`` (samples contributed by each tile, required), ``starts_k`` (origin inside the decoded source, default 0) and ``steps_k`` (source stride, the per-tile decimation, default 1), reading as ``virtual[pos : pos + size] = source[start : start + size * step : step]`` - with ``pos`` the running sum of the previous sizes along the axis; + with ``pos`` the running sum of the previous sizes along the axis. + A column that varies is 1-D over ``tile_k``; one holding the same + value everywhere folds to a 0-d variable (and one holding its kind's + default is left out), so an acquisition of equal-length files spends + one element instead of one per tile; - N-D *parameter* variables: ``paths`` (the source file of each tile, required) plus any format-private per-tile values forwarded to the engine as keyword arguments. A parameter carries only the tile dimensions along which it actually varies — a constant folds to a 0-d - variable — and broadcasts over the grid at read time; -- an optional 0-d ``root`` variable: the common directory of the - sources, split off so the per-tile ``paths`` stay root-relative — - one shared constant instead of a per-tile repeat. Absent (or empty), - the paths are used as stored; -- an optional *axis map*: the geometry axes are stored in *source* - order (``sizes_g`` describes source axis ``g``), and a 1-D ``axes`` - variable lists the geometry axis each virtual axis presents, in - virtual order — how transposes and inserted axes stay lazy. Axes - beyond the 0-d ``source_ndim`` are *synthetic* (inserted, one sample - wide, backed by no source data); geometry axes left out of the map - are *hidden* (pinned to a single sample, read but not presented — - how integer indexing stays lazy). Both variables are absent from - freshly scanned manifests: the default is the identity, every - geometry axis a source axis presented in stored order. - -String variables (``paths``, ``root`` and any string parameter) are -held as fixed-width bytes (numpy ``S`` kind, filesystem encoding): one + variable — and broadcasts over the grid at read time. + +The variables are therefore exactly the per-tile columns; everything +else the tiling needs is one JSON document, the ``header`` dataset +attribute: + +- ``ntiles``: the grid shape, and the single authority on it — an axis + whose columns have all folded has no dimension left to measure; +- ``engine``: the engine specification, a registered name + (``xdas.io.Engine[name]``) or an object pairing that ``name`` with the + constants to pass its ``load_tile``; +- ``dtype``: the element type of the sources as the engine decodes + them. Not a decode target: engines decode into the element type of + their sources, and the array records that type at scan time so + laziness holds — a virtual array answers ``dtype`` without touching + its sources — and verifies it against every decoded tile. Casting is + an explicit extra step (:meth:`TileArray.astype`), outside the tiles + machinery; +- ``root``, optional: the common directory of the sources, split off so + the per-tile ``paths`` stay root-relative — one shared constant + instead of a per-tile repeat. Absent, the paths are used as stored; +- ``axes``, optional: the axis map (below). + +The geometry axes are stored in *source* order (``sizes_g`` describes +source axis ``g``), and ``axes`` is the full virtual arrangement in +numpy indexing-key notation: the geometry axis each virtual axis +presents, with ``null`` where the axis is *synthetic* (inserted, one +sample wide, backed by no source data) — ``[1, null, 0]`` reads as +``transpose(a)[:, np.newaxis]``. Absent means the identity: no +transposition, no insertion, every geometry axis a source axis +presented in stored order. Everything else derives from the list: the +source rank is the grid rank less the number of ``null``s, and geometry +axes missing from it are *hidden* (pinned to a single sample, read but +not presented — how integer indexing stays lazy). Two invariants keep +manifests canonical: synthetic grid axes stay trailing and are numbered +in their virtual order (a transpose that reorders them renumbers the +grid rather than the map), and pinning a synthetic axis deletes it from +the grid rather than hiding it (a hidden synthetic axis would be +contentless once the pin has selected into ``paths``). + +String variables (``paths`` and any string parameter) are held as +fixed-width bytes (numpy ``S`` kind, filesystem encoding): one contiguous block instead of a heap of str objects, comparisons that vectorize, and a stored form that lands on disk as netCDF char arrays with no encoding step. The constructor recodes str-valued variables -(hand-built manifests, legacy stored forms) on entry; values decode -back to str only when handed to the engine. - -What arrays cannot carry — the ``dtype`` and the ``engine`` -specification — lives on the array itself, by value, and travels -beside the dataset when the array is persisted (see -:func:`xdas.io.xdas.save_dataarray`). The dtype is not a decode -target: engines decode into the element type of their sources, and -the array records that type at scan time so laziness holds — a -virtual array answers ``dtype`` without touching its sources — and -verifies it against every decoded tile. Casting is an explicit -extra step (:meth:`TileArray.astype`), outside the tiles machinery. -Every dataset attribute is a user attribute. - -Geometry loads eagerly at construction (tiny); parameters stay folded — -a constant occupies one element whatever the grid size — and broadcast +(hand-built manifests) on entry; values decode back to str only when +handed to the engine. The ``root``, being a header value, travels +through JSON instead: non-UTF-8 directory names survive as escaped +surrogates, exactly as :func:`os.fsdecode` spells them. + +A tile array carries no user attributes of its own — it is a duck +array, like the numpy array it stands in for. Metadata belongs to the +enclosing :class:`xdas.DataArray`. + +Geometry loads eagerly at construction; a column holding one value +everywhere (the absent ``starts_k`` and ``steps_k``, and the ``sizes_k`` +of an acquisition of equal-length files) is detected and kept as a +broadcast view, so it costs one element instead of one per tile and its +tile boundaries stay a closed form. Parameters stay folded — a constant +occupies one element whatever the grid size — and broadcast over the grid only as tiles are read. Slicing (any nonzero step, negative ones reversing lazily), integer indexing and ``np.newaxis`` fold into the geometry and the axis map and return a new @@ -111,7 +138,13 @@ """Prefix of the tile-grid dimensions of a manifest dataset.""" TILES_GROUP = "__tiles__" -"""Name of the sibling group holding a stored tile array's manifest.""" +"""Conventional name of the group holding a stored tile array's manifest.""" + +HEADER = "header" +"""Manifest attribute holding the JSON document of everything but the columns.""" + +TILING = "__tiling__" +"""Attribute of a placeholder variable naming the group describing its tiling.""" _UNITS = ("B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB") """Decimal byte units, as xarray spells them in its ``Size:`` header.""" @@ -318,6 +351,135 @@ def _bounding_key(key, shape): return tuple(box), tuple(residual), empty +def _as_engine(spec): + """Return engine specification *spec* as a normalized dict. + + A plain string is shorthand for ``{"name": spec}``; the json round + trip deep-copies and normalizes (tuples become lists), so equality + survives a store round trip. + """ + if isinstance(spec, str): + spec = {"name": spec} + spec = json.loads(json.dumps(spec)) + if not isinstance(spec, dict) or "name" not in spec: + raise ValueError("the engine specification must have a `name` key") + return spec + + +def _encode_axes(axes, source_ndim, ngrid): + """Return the header ``axes`` of a map, or None when it is the identity.""" + encoded = [None if g >= source_ndim else g for g in axes] + return None if encoded == list(range(ngrid)) else encoded + + +def _decode_axes(encoded, ngrid): + """Return ``(axes, source_ndim)`` from the header ``axes`` entry. + + The synthetic axes are the ``null`` entries — they trail the source + axes in the grid, numbered in the order the list presents them — + so the source rank and the geometry axis of every virtual axis both + fall out of the one list. + """ + if encoded is None: + return tuple(range(ngrid)), ngrid + if not isinstance(encoded, list): + raise ValueError("`axes` must be a list") + if not encoded: + raise ValueError("a tile array needs at least one visible axis") + source_ndim = ngrid - sum(g is None for g in encoded) + if source_ndim < 1: + raise ValueError("`axes` must leave at least one source axis") + axes, synthetic = [], source_ndim + for g in encoded: + if g is None: + axes.append(synthetic) + synthetic += 1 + elif isinstance(g, int) and 0 <= g < source_ndim: + axes.append(g) + else: + raise ValueError( + f"`axes` must name source geometry axes below {source_ndim}" + ) + if len(set(axes)) != len(axes): + raise ValueError("`axes` must name distinct geometry axes") + return tuple(axes), source_ndim + + +def _make_header(counts, dtype, engine, root, axes, source_ndim): + """Return the header describing a manifest, defaults left out. + + The engine collapses to its bare name when it carries no constants, + and the axis map and the root are absent at their defaults (the + identity arrangement, and paths stored whole). + """ + header = {"ntiles": [int(count) for count in counts]} + header["engine"] = engine["name"] if set(engine) == {"name"} else dict(engine) + header["dtype"] = str(np.dtype(dtype)) + encoded = _encode_axes(axes, source_ndim, len(counts)) + if encoded is not None: + header["axes"] = encoded + if root: + header["root"] = root + return header + + +def _read_header(dataset): + """Return the parsed header of manifest *dataset*.""" + return json.loads(dataset.attrs[HEADER]) + + +def _write_header(dataset, header): + """Return *dataset* carrying *header* as its JSON header attribute.""" + return dataset.assign_attrs({HEADER: json.dumps(header)}) + + +def _canonical(dataset, header, axes, source_ndim): + """Bring the grid of *dataset* back to its canonical numbering. + + The two invariants of the axis map, enforced in one place: a + synthetic axis absent from *axes* has been pinned, and is + contentless once the pin has selected into the columns, so it + leaves the grid; the synthetic axes that remain trail the source + axes in the order they are presented. Mutates the ``ntiles`` of + *header* and returns the rewritten dataset and axis map — writing + the header back is the caller's. + """ + ngrid = len(header["ntiles"]) + keep = list(range(source_ndim)) + [g for g in axes if g >= source_ndim] + if keep == list(range(ngrid)): + return dataset, axes + dropped = [g for g in range(ngrid) if g not in keep] + dataset = dataset.isel( + {f"{TILE_PREFIX}{g}": 0 for g in dropped if f"{TILE_PREFIX}{g}" in dataset.dims} + ).drop_vars( + [ + f"{kind}_{g}" + for g in dropped + for kind in ("sizes", "starts", "steps") + if f"{kind}_{g}" in dataset + ] + ) + renames = {} + for new, old in enumerate(keep): + if new == old: + continue + if f"{TILE_PREFIX}{old}" in dataset.dims: + renames[f"{TILE_PREFIX}{old}"] = f"{TILE_PREFIX}{new}" + for kind in ("sizes", "starts", "steps"): + if f"{kind}_{old}" in dataset: + renames[f"{kind}_{old}"] = f"{kind}_{new}" + if renames: + # two passes through temporaries: a renumbering that swaps two + # axes would otherwise rename a name onto a live one + dataset = dataset.rename({old: f"_{new}" for old, new in renames.items()}) + dataset = dataset.rename({f"_{new}": new for new in renames.values()}) + # renaming leaves the columns in their old dimension order + order = [f"{TILE_PREFIX}{k}" for k in range(len(keep))] + dataset = dataset.transpose(*order, missing_dims="ignore") + header["ntiles"] = [header["ntiles"][g] for g in keep] + return dataset, tuple(keep.index(g) for g in axes) + + def _assign_geometry(dataset, assign): """Apply geometry *assign* to *dataset*, folding all-default columns away. @@ -337,6 +499,114 @@ def _assign_geometry(dataset, assign): return dataset.drop_vars([name for name in drop if name in dataset]) +def _isel(dataset, indexers): + """Index *dataset* along the tile dimensions it still carries. + + An axis whose every variable folded has no dimension left to index: + what it holds is one value for the whole axis, which selecting + among its tiles leaves untouched. Only the header's ``ntiles`` + records how many tiles such a selection kept, so it is restated + here — for every indexed axis, dimension or not. + """ + header = _read_header(dataset) + counts = header["ntiles"] + for dim, key in indexers.items(): + g = int(dim[len(TILE_PREFIX) :]) + counts[g] = ( + len(range(*key.indices(counts[g]))) + if isinstance(key, slice) + else len(np.asarray(key)) + ) + dataset = dataset.isel( + {dim: key for dim, key in indexers.items() if dim in dataset.dims} + ) + return _write_header(dataset, header) + + +def _fold_geometry(dataset, sizes, starts, steps): + """Rewrite the geometry of *dataset* in its canonical stored form. + + A column holding one value everywhere is stored as a 0-d variable — + one element whatever the tile count — and one holding the kind's + default is not stored at all. No column has to stay expanded to + keep the grid measurable: the header's ``ntiles`` states how many + tiles an axis holds, the count no dimension is left to give. + """ + columns = { + f"{kind}_{g}": values + for kind, per_axis in (("sizes", sizes), ("starts", starts), ("steps", steps)) + for g, values in enumerate(per_axis) + } + assign, drop = {}, [] + for name, values in columns.items(): + kind = name.split("_")[0] + default = {"sizes": None, "starts": 0, "steps": 1}[kind] + if not len(values) or not bool((values == values[0]).all()): + continue # varying: worth one value per tile, and it sizes the axis + if default is not None and int(values[0]) == default: + if name in dataset: + drop.append(name) + continue + stored = dataset.get(name) + if stored is None or stored.dims != (): + assign[name] = xr.Variable((), np.asarray(values[0], np.int64)) + if assign: + dataset = dataset.assign(assign) + return dataset.drop_vars(drop) if drop else dataset + + +class _Edges: + """The running sample offsets of the tiles along one geometry axis. + + Either the eager ``cumsum`` of a varying size column, or — when + every tile is the same width — that width alone, every boundary + being a multiple of it. Constant columns arrive as zero-stride + broadcast views (see :meth:`TileArray._geometry`), which is what + picks the closed forms: they hold no array and turn locating a + tile into a division. + """ + + def __init__(self, sizes): + self.count = len(sizes) + self.size = int(sizes[0]) if self.count and sizes.strides == (0,) else None + self.values = ( + None if self.size is not None else np.concatenate(([0], np.cumsum(sizes))) + ) + + @property + def total(self): + """int: the extent the axis spans, every tile summed.""" + if self.size is None: + return int(self.values[-1]) + return self.count * self.size + + def at(self, index): + """int: the sample offset where tile *index* begins.""" + if self.size is None: + return int(self.values[index]) + return int(index) * self.size + + def searchsorted(self, value, side): + """int: locate *value*, exactly as :func:`numpy.searchsorted` would.""" + if self.size is None: + return int(np.searchsorted(self.values, value, side)) + # boundaries are 0, size, ..., count * size: count them by division + index = -(-value // self.size) if side == "left" else value // self.size + 1 + return int(min(max(index, 0), self.count + 1)) + + def span(self, start, stop): + """Return the offsets of tiles *start* to *stop* (excluded), as an array.""" + if self.size is None: + return self.values[start:stop] + return np.arange(start, stop, dtype=np.int64) * self.size + + def __iter__(self): + """Iterate over the boundaries, first to last.""" + if self.size is None: + return iter(self.values) + return iter(range(0, (self.count + 1) * self.size, self.size)) + + def _materialize(value): """Read any :class:`TileArray` in *value*, descending one level.""" if isinstance(value, TileArray): @@ -373,15 +643,17 @@ class TileArray(VirtualBackend, np.lib.mixins.NDArrayOperatorsMixin, vtype="tile Parameters ---------- dataset : xarray.Dataset - The manifest dataset: the 1-D geometry and the N-D per-tile - parameter variables described in the module docstring. Every - dataset attribute is a user attribute. - dtype : str or numpy.dtype + The manifest dataset described in the module docstring: the + per-tile columns, plus the ``header`` attribute holding + everything else the tiling needs. + dtype : str or numpy.dtype, optional Element type of the sources as the engine decodes them (little-endian or single-byte), recorded so the lazy array can report it without reading. Verified against every decoded - tile, never used to cast. - engine : str or dict + tile, never used to cast. Read off the header when omitted; + given, it must agree with the header when there is one (a + hand-built dataset carries none, and then it is required). + engine : str or dict, optional The engine specification, stored by value with the array: the key ``"name"`` selects a registered engine (``xdas.io.Engine[name]``); the remaining keys are passed to @@ -389,16 +661,17 @@ class TileArray(VirtualBackend, np.lib.mixins.NDArrayOperatorsMixin, vtype="tile shorthand for ``{"name": engine}``. Never an :class:`~xdas.io.Engine` instance: the specification must reproduce the decode with no instance alive, so open-time - settings (``vtype``, ``ctype``) do not belong in it. + settings (``vtype``, ``ctype``) do not belong in it. Read off + the header when omitted, on the same terms as *dtype*. """ # manifest concatenation fuses scan products into one compact array, # so multi-file scans can drain batches (see VirtualBackend) consolidates = True - def __init__(self, dataset, dtype, engine): + def __init__(self, dataset, dtype=None, engine=None): # canonical string dtype is fixed-width bytes: str-valued - # variables (hand-built or legacy stored manifests) recode here + # variables (hand-built manifests) recode here recode = { name: xr.Variable(dataset[name].dims, _as_bytes(dataset[name].values)) for name in map(str, dataset.data_vars) @@ -406,29 +679,48 @@ def __init__(self, dataset, dtype, engine): } if recode: dataset = dataset.assign(recode) + ngrid = 0 + while f"sizes_{ngrid}" in dataset: + ngrid += 1 + if ngrid == 0: + raise ValueError("a tile array needs a `sizes_0` geometry variable") + if HEADER in dataset.attrs: + header = _read_header(dataset) + else: + # a hand-assembled dataset: the columns are all it says — the + # tile counts come from the dimensions it declares, everything + # else is the default — and the tiles machinery owns the whole + # manifest, stray attributes included + header = { + "ntiles": [ + int(dataset.sizes[dim]) if dim in dataset.dims else 1 + for dim in (f"{TILE_PREFIX}{k}" for k in range(ngrid)) + ] + } + dataset = dataset.drop_attrs(deep=True) self.dataset = dataset self._cache = None - if isinstance(engine, str): - engine = {"name": engine} - # the json round trip deep-copies and normalizes (tuples become - # lists), so equality survives a store round trip - engine = json.loads(json.dumps(engine)) - if not isinstance(engine, dict) or "name" not in engine: - raise ValueError("the engine specification must have a `name` key") + # the header is authoritative; an explicit argument fills in for + # a manifest that carries none (the scan-time assembly, and + # hand-built datasets) and must otherwise agree + if engine is None and "engine" not in header: + raise ValueError("the manifest records no engine") + self._engine = _as_engine(header["engine"] if engine is None else engine) + if "engine" in header and self._engine != _as_engine(header["engine"]): + raise ValueError(f"the manifest records the engine {header['engine']!r}") # imported here: xdas.io imports this module at package init from ..io.core import Engine - Engine[engine["name"]] # fail fast on unregistered engines - self._engine = engine - self._dtype = np.dtype(dtype) + Engine[self._engine["name"]] # fail fast on unregistered engines + if dtype is None and "dtype" not in header: + raise ValueError("the manifest records no dtype") + self._dtype = np.dtype(header["dtype"] if dtype is None else dtype) + if "dtype" in header and self._dtype != np.dtype(header["dtype"]): + raise ValueError(f"the manifest records the dtype {header['dtype']!r}") if self._dtype.byteorder == ">": raise ValueError("only little-endian or single-byte dtypes are supported") - ngrid = 0 - while f"sizes_{ngrid}" in dataset: - ngrid += 1 - if ngrid == 0: - raise ValueError("a tile array needs a `sizes_0` geometry variable") self.dims = dims = tuple(f"{TILE_PREFIX}{g}" for g in range(ngrid)) + self._counts = self._tile_counts(header) self._sizes = self._geometry("sizes", None) self._starts = self._geometry("starts", 0) self._steps = self._geometry("steps", 1) @@ -444,32 +736,11 @@ def __init__(self, dataset, dtype, engine): reach = self._starts[g] + (self._sizes[g] - 1) * self._steps[g] if np.any(reach < 0): raise ValueError(f"`steps_{g}` walks out of the source") - self._edges = tuple( - np.concatenate(([0], np.cumsum(sizes))) for sizes in self._sizes - ) + self._edges = tuple(_Edges(sizes) for sizes in self._sizes) # the axis map: which geometry axis each virtual axis presents. - # absent variables mean the identity — every geometry axis is a + # an absent entry means the identity — every geometry axis is a # source axis, presented in stored (= source) order - if "source_ndim" in dataset: - if tuple(dataset["source_ndim"].dims) != (): - raise ValueError("`source_ndim` must be a 0-d variable") - self._source_ndim = int(dataset["source_ndim"].values[()]) - else: - self._source_ndim = ngrid - if not 0 < self._source_ndim <= ngrid: - raise ValueError("`source_ndim` must be between 1 and the geometry rank") - if "axes" in dataset: - if len(dataset["axes"].dims) != 1 or dataset["axes"].dims[0] in dims: - raise ValueError("`axes` must be 1-D over its own dimension") - self._axes = tuple(int(g) for g in dataset["axes"].values) - else: - self._axes = tuple(range(ngrid)) - if not self._axes: - raise ValueError("a tile array needs at least one visible axis") - if len(set(self._axes)) != len(self._axes) or not all( - 0 <= g < ngrid for g in self._axes - ): - raise ValueError(f"`axes` must name distinct geometry axes below {ngrid}") + self._axes, self._source_ndim = _decode_axes(header.get("axes"), ngrid) # synthetic axes (beyond the source rank) and hidden axes (absent # from the map) have no extent to offer: their tiles are one # sample wide, and a hidden axis holds a single tile — several @@ -481,15 +752,12 @@ def __init__(self, dataset, dtype, engine): raise ValueError(f"axis {g} is synthetic or hidden: sizes must be 1") if g not in self._axes and len(self._sizes[g]) != 1: raise ValueError(f"hidden axis {g} must hold a single tile") - self._shape = tuple(int(self._edges[g][-1]) for g in self._axes) + self._shape = tuple(self._edges[g].total for g in self._axes) if "paths" not in dataset: raise ValueError("a tile array needs a `paths` variable") - if "root" in dataset: - if tuple(dataset["root"].dims) != (): - raise ValueError("`root` must be a 0-d variable") - self.root = os.fsdecode(dataset["root"].values[()]) - else: - self.root = "" + self.root = header.get("root", "") + if not isinstance(self.root, str): + raise ValueError("`root` must be a string") geometry = { f"{kind}_{g}" for kind in ("sizes", "starts", "steps") for g in range(ngrid) } @@ -497,8 +765,7 @@ def __init__(self, dataset, dtype, engine): sorted( name for name in map(str, dataset.data_vars) - if name not in geometry - and name not in ("paths", "root", "axes", "source_ndim") + if name not in geometry and name != "paths" ) ) for name in ("paths", *self._params): @@ -507,9 +774,26 @@ def __init__(self, dataset, dtype, engine): raise ValueError( f"`{name}` dimensions must be an ordered subset of {dims}" ) + # canonical stored form: constant geometry columns hold one + # element whatever the tile count, and the header restates + # everything the columns no longer measure + dataset = _fold_geometry(self.dataset, self._sizes, self._starts, self._steps) + header = json.dumps( + _make_header( + self._counts, + self._dtype, + self._engine, + self.root, + self._axes, + self._source_ndim, + ) + ) + if dataset.attrs.get(HEADER) != header: + dataset = dataset.assign_attrs({HEADER: header}) + self.dataset = dataset @classmethod - def from_tiles(cls, paths, sizes, dtype, engine, *, attrs=None, **params): + def from_tiles(cls, paths, sizes, dtype, engine, **params): """Build a tile array from per-tile descriptions of fresh sources. The scan-time encoder: every tile is read from the origin of @@ -528,7 +812,7 @@ def from_tiles(cls, paths, sizes, dtype, engine, *, attrs=None, **params): the working directory cannot be trusted later, as reads are lazy and stored views outlive the session. The common directory of the absolute paths is then split off into the - 0-d ``root`` variable, the stored per-tile paths staying + header's ``root``, the stored per-tile paths staying root-relative. sizes : sequence of int or 1-D array-like One entry per axis (this defines the rank): the samples each @@ -545,8 +829,6 @@ def from_tiles(cls, paths, sizes, dtype, engine, *, attrs=None, **params): (``xdas.io.Engine[name]``); the remaining keys are passed to its ``load_tile`` as keyword parameters. A plain string is shorthand for ``{"name": engine}``. - attrs : dict, optional - User attributes of the virtual array. **params : array-like Per-tile engine parameters, broadcast over the grid: each read passes the tile's value to the engine as a keyword @@ -585,8 +867,6 @@ def from_tiles(cls, paths, sizes, dtype, engine, *, attrs=None, **params): f"`paths` shape {paths.shape} does not match the grid {counts}" ) data["paths"] = _fold_param(paths, counts, dims) - if root: - data["root"] = ((), np.asarray(root)) reserved = {"paths", "root"} | { f"{kind}_{k}" for kind in ("sizes", "starts", "steps") for k in range(ndim) } @@ -594,8 +874,15 @@ def from_tiles(cls, paths, sizes, dtype, engine, *, attrs=None, **params): if name in reserved: raise ValueError(f"parameter name {name!r} is reserved") data[name] = _fold_param(np.asarray(values), counts, dims) - dataset = xr.Dataset(data, attrs=dict(attrs or {})) - return cls(dataset, dtype, engine) + header = _make_header( + counts, + dtype, + _as_engine(engine), + os.fsdecode(root), + tuple(range(ndim)), + ndim, + ) + return cls(_write_header(xr.Dataset(data), header)) @classmethod def from_variable(cls, variable): @@ -624,12 +911,9 @@ def from_variable(cls, variable): def to_dataset(self): """Encode this tile array as its manifest dataset. - The stored form — a copy of the wrapped dataset carrying the - user attributes. What the dataset cannot hold, the by-value - ``dtype`` and ``engine``, the caller reads off the properties - and stores beside it. Source paths are stored exactly as the - array holds them: relative to the 0-d ``root`` variable - carrying their common directory. + The stored form, whole: a copy of the wrapped dataset, columns + and header alike. Source paths are stored exactly as the array + holds them, relative to the header's ``root``. Returns ------- @@ -638,14 +922,16 @@ def to_dataset(self): """ return self.dataset.copy() - def create_variable(self, file, name, dims=None, dtype=None): + def create_variable(self, file, name, dims=None): """ Create the placeholder variable of the stored form. - The placeholder records the dtype (as any typed store would, - e.g. zarr array metadata) and carries the engine specification - by value in a ``__tile_array__`` attribute; the manifest itself - is appended by :meth:`finalize_save` once *file* is closed. + The placeholder is a projection of the array — its dtype and + its dimensions, no values (an unwritten contiguous dataset + occupies no bytes) — pointing at the group that describes the + tiling through a ``__tiling__`` attribute, in the fashion of + CF's ``grid_mapping``. The manifest itself is stored there, as + :meth:`sibling_datasets` describes. Parameters ---------- @@ -655,68 +941,91 @@ def create_variable(self, file, name, dims=None, dtype=None): Variable name to create inside *file*. dims : sequence of str, optional Dimension names for the variable. - dtype : dtype-like, optional - Element type of the placeholder. Default to :attr:`dtype`. Returns ------- variable The newly created file variable. """ - variable = file.create_variable( - name, dims, self.dtype if dtype is None else dtype - ) - variable.attrs["__tile_array__"] = json.dumps({"engine": self.engine}) + variable = file.create_variable(name, dims, self.dtype) + variable.attrs[TILING] = TILES_GROUP return variable - def finalize_save(self, fname, group=None): + def sibling_datasets(self): """ - Append the manifest as a sibling group of the stored variable. + Return the manifest, stored beside the placeholder variable. The second half of the stored form, written natively by xarray - once the file handle of :meth:`create_variable` is closed: the - manifest dataset lands in a ``__tiles__`` group next to the - placeholder variable. + with the rest of the file's metadata: the manifest dataset + lands in the ``__tiles__`` group the placeholder points at. - Parameters - ---------- - fname : str - Path of the file holding the placeholder variable. - group : str, optional - Group of the placeholder variable within the file. + Returns + ------- + dict + ``{TILES_GROUP: manifest}``. """ manifest = self.to_dataset() # strings are fixed-width bytes by construction and land on disk # as char arrays; only stale open-time encodings need clearing for name in list(manifest.variables): manifest[name].encoding.clear() - location = TILES_GROUP if group is None else f"{group}/{TILES_GROUP}" - manifest.to_netcdf(fname, mode="a", group=location, engine="h5netcdf") + return {TILES_GROUP: manifest} def _geometry(self, kind, default): - """Load the eager 1-D ``{kind}_k`` arrays (*default* where absent).""" + """Load the eager 1-D ``{kind}_k`` arrays (*default* where absent). + + A column holding one value everywhere — always so when the + variable is absent or stored 0-d — is kept as a zero-stride + broadcast view instead of a full-length array: nothing + downstream writes to the geometry, and :class:`_Edges` keys its + closed forms off that view. Detected, never assumed — a + truncated last tile makes a size column vary. + """ arrays = [] for k, dim in enumerate(self.dims): name = f"{kind}_{k}" - if name in self.dataset: - if tuple(self.dataset[name].dims) != (dim,): - raise ValueError(f"`{name}` must have dimensions ({dim!r},)") - arrays.append(np.asarray(self.dataset[name].values, dtype=np.int64)) - else: - count = int(self.dataset.sizes[dim]) - arrays.append(np.full(count, default, dtype=np.int64)) + stored = self.dataset[name].dims if name in self.dataset else None + if stored == (dim,): + values = np.asarray(self.dataset[name].values, dtype=np.int64) + if ( + values.strides != (0,) + and len(values) + and (values == values[0]).all() + ): + values = np.broadcast_to(values[0], values.shape) + arrays.append(values) + continue + if stored is not None and stored != (): + raise ValueError(f"`{name}` must have dimensions ({dim!r},) or ()") + # a folded column names one value, the grid the tile count + value = np.int64( + default if stored is None else self.dataset[name].values[()] + ) + arrays.append(np.broadcast_to(value, (self._counts[k],))) return tuple(arrays) + def _tile_counts(self, header): + """Return how many tiles each geometry axis holds. + + The header's ``ntiles`` is the single authority: an axis whose + columns have all folded has no dimension left to measure. The + dimensions the manifest still declares must agree with it. + """ + counts = header["ntiles"] + if not isinstance(counts, list) or len(counts) != len(self.dims): + raise ValueError("`ntiles` must hold one count per geometry axis") + for k, dim in enumerate(self.dims): + if not (isinstance(counts[k], int) and counts[k] >= 1): + raise ValueError("`ntiles` must hold positive integers") + if dim in self.dataset.dims and int(self.dataset.sizes[dim]) != counts[k]: + raise ValueError(f"`ntiles` disagrees with the size of `{dim}`") + return tuple(int(count) for count in counts) + @property def engine(self): """dict: the engine specification (``"name"`` plus its parameters).""" return self._engine - @property - def attrs(self): - """dict: the user attributes of the virtual array.""" - return dict(self.dataset.attrs) - @property def chunks(self): """Tuple of tuple of int: the tiling, as per-axis tile extents. @@ -725,22 +1034,23 @@ def chunks(self): """ return tuple(tuple(int(size) for size in self._sizes[g]) for g in self._axes) - def _assign_axes(self, dataset, axes, ngrid): - """Set the axis map of *dataset* to *axes*, dropping what is derivable. + def _assign_axes(self, dataset, axes): + """Set the axis map of *dataset* to *axes*, canonicalizing the grid. - The identity parts stay absent: `axes` is stored only when it - differs from the stored geometry order, `source_ndim` only when - synthetic axes exist. The source rank is this array's. + The header states the arrangement only when it is not the + identity; hidden synthetic axes leave the grid and the ones + that remain are renumbered into virtual order first (see + :func:`_canonical`). The source rank is this array's — no + operation of the grid creates or destroys a source axis. """ - drop = [name for name in ("axes", "source_ndim") if name in dataset] - if drop: - dataset = dataset.drop_vars(drop) - assign = {} - if axes != tuple(range(ngrid)): - assign["axes"] = ("axis", np.asarray(axes, dtype=np.int64)) - if self._source_ndim != ngrid: - assign["source_ndim"] = ((), np.asarray(self._source_ndim, dtype=np.int64)) - return dataset.assign(assign) if assign else dataset + header = _read_header(dataset) + dataset, axes = _canonical(dataset, header, axes, self._source_ndim) + encoded = _encode_axes(axes, self._source_ndim, len(header["ntiles"])) + if encoded is None: + header.pop("axes", None) + else: + header["axes"] = encoded + return _write_header(dataset, header) @property def shape(self): @@ -801,7 +1111,8 @@ def _fold(self, key): Slices of any nonzero step trim the geometry (a negative step folds as its ascending twin, then the axis flips); an integer - pins the geometry axis at one sample and hides it from the map; + pins the geometry axis at one sample and hides it from the map + (a synthetic axis, contentless once pinned, leaves the grid); ``None`` inserts a synthetic axis. Raises :class:`_Unfoldable` for other entries, for empty selections (a grid needs at least one tile) and for all-integer keys (a scalar is not a tile @@ -840,9 +1151,9 @@ def _fold(self, key): if (lo, hi, s) == (0, extent, 1): continue edges = self._edges[g] - i0 = int(np.searchsorted(edges, lo, "right")) - 1 - i1 = int(np.searchsorted(edges, hi, "left")) - pos = edges[i0:i1] + i0 = edges.searchsorted(lo, "right") - 1 + i1 = edges.searchsorted(hi, "left") + pos = edges.span(i0, i1) size = self._sizes[g][i0:i1] start = self._starts[g][i0:i1] step = self._steps[g][i0:i1] @@ -858,11 +1169,11 @@ def _fold(self, key): assign[f"steps_{g}"] = (dim, (step * s)[keep]) if not new_axes: raise _Unfoldable("an all-integer key selects a scalar, not a grid") - dataset = _assign_geometry(self.dataset.isel(indexers), assign) + dataset = _assign_geometry(_isel(self.dataset, indexers), assign) new_axes = tuple(new_axes) if new_axes != self._axes: - dataset = self._assign_axes(dataset, new_axes, len(self.dims)) - result = type(self)(dataset, self.dtype, self.engine) + dataset = self._assign_axes(dataset, new_axes) + result = type(self)(dataset) for axis in flips: result = result._flip(axis) # np.newaxis entries insert synthetic axes at their output position @@ -889,8 +1200,8 @@ def _permute_tiles(self, order, axis=0): order = np.asarray(order) if not np.array_equal(np.sort(order), np.arange(len(self._sizes[g]))): raise ValueError(f"`order` must be a permutation of axis {axis} tiles") - dataset = self.dataset.isel({self.dims[g]: order}) - return type(self)(dataset, self.dtype, self.engine) + dataset = _isel(self.dataset, {self.dims[g]: order}) + return type(self)(dataset) def _flip(self, axis): """Reverse the array along virtual *axis*, staying virtual. @@ -904,13 +1215,13 @@ def _flip(self, axis): starts = self._starts[g] + (self._sizes[g] - 1) * self._steps[g] steps = -self._steps[g] dataset = _assign_geometry( - self.dataset.isel({dim: slice(None, None, -1)}), + _isel(self.dataset, {dim: slice(None, None, -1)}), { f"starts_{g}": (dim, starts[::-1]), f"steps_{g}": (dim, steps[::-1]), }, ) - return type(self)(dataset, self.dtype, self.engine) + return type(self)(dataset) @classmethod def concat(cls, arrays, dim=0): @@ -967,6 +1278,8 @@ def concat(cls, arrays, dim=0): ): raise ValueError("can only concatenate compatible tile arrays") data = {} + counts = list(first._counts) + counts[gaxis] = sum(array._counts[gaxis] for array in arrays) for kind, per_axis, default in ( ("sizes", [array._sizes for array in arrays], None), ("starts", [array._starts for array in arrays], 0), @@ -980,10 +1293,6 @@ def concat(cls, arrays, dim=0): if default is not None and bool((values == default).all()): continue data[f"{kind}_{g}"] = (dims[g], values) - if first._axes != tuple(range(ngrid)): - data["axes"] = ("axis", np.asarray(first._axes, dtype=np.int64)) - if first._source_ndim != ngrid: - data["source_ndim"] = ((), np.asarray(first._source_ndim, dtype=np.int64)) # the fused root is the deepest directory containing every input # root ("" when an input has none or no common directory exists) roots = [array.root for array in arrays] @@ -991,8 +1300,6 @@ def concat(cls, arrays, dim=0): root = os.path.commonpath(roots) if all(roots) else "" except ValueError: root = "" - if root: - data["root"] = ((), np.asarray(os.fsencode(root))) for name in ("paths", *first._params): if name == "paths": variables = [array._rebased_paths(root) for array in arrays] @@ -1011,13 +1318,15 @@ def concat(cls, arrays, dim=0): ) parts = [] for variable, array in zip(variables, arrays): - counts = dict(zip(dims, (len(sizes) for sizes in array._sizes))) - parts.append(variable.set_dims({dim: counts[dim] for dim in union})) + sizes = dict(zip(dims, array._counts)) + parts.append(variable.set_dims({dim: sizes[dim] for dim in union})) axis_pos = union.index(dims[gaxis]) values = np.concatenate([part.values for part in parts], axis=axis_pos) data[name] = xr.Variable(union, values) - dataset = xr.Dataset(data, attrs=first.attrs) - return cls(dataset, first.dtype, first.engine) + header = _make_header( + counts, first.dtype, first.engine, root, first._axes, first._source_ndim + ) + return cls(_write_header(xr.Dataset(data), header)) def _full_paths(self): """Return the full source byte path of every tile, root joined, over the grid.""" @@ -1109,7 +1418,7 @@ def _read(self): widths.append(size) selection, widths = tuple(selection), tuple(widths) dest = tuple( - slice(int(self._edges[g][index[g]]), int(self._edges[g][index[g] + 1])) + slice(self._edges[g].at(index[g]), self._edges[g].at(index[g] + 1)) for g in self._axes ) kwargs = dict(spec) @@ -1139,10 +1448,10 @@ def _read(self): def equals(self, other): """Whether *other* describes the same tiling (not elementwise). - Compares the engine, dtype, geometry, parameters and user - attributes; ``==`` stays elementwise, as on any numpy-like - array. Paths compare joined: two arrays naming the same source - files are equal however each splits its ``root``. + Compares the engine, dtype, geometry and parameters; ``==`` + stays elementwise, as on any numpy-like array. Paths compare + joined: two arrays naming the same source files are equal + however each splits its ``root``. """ if not isinstance(other, TileArray): return False @@ -1150,7 +1459,6 @@ def equals(self, other): self.engine != other.engine or self.dtype != other.dtype or self.shape != other.shape - or self.attrs != other.attrs or self._params != other._params or self._axes != other._axes or self._source_ndim != other._source_ndim @@ -1311,17 +1619,19 @@ def expand_dims(self, axis=0): dataset = self.dataset.assign( {f"sizes_{ngrid}": (f"{TILE_PREFIX}{ngrid}", np.ones(1, np.int64))} ) + header = _read_header(dataset) + header["ntiles"] = [*header["ntiles"], 1] axes = self._axes[:axis] + (ngrid,) + self._axes[axis:] - dataset = self._assign_axes(dataset, axes, ngrid + 1) - return type(self)(dataset, self.dtype, self.engine) + return type(self)(self._assign_axes(_write_header(dataset, header), axes)) def squeeze(self, axis=None): """Drop unit axes, staying virtual. - The dropped axes leave the virtual order (real ones stay in the - grid as hidden, one-sample reads); squeezing every axis away - materializes the value as a 0-d array, as a grid needs at least - one visible axis. + The dropped axes leave the virtual order — a source axis stays + in the grid as a hidden, one-sample read, a synthetic one + leaves it altogether; squeezing every axis away materializes + the value as a 0-d array, as a grid needs at least one visible + axis. Parameters ---------- @@ -1351,12 +1661,16 @@ def squeeze(self, axis=None): if len(drop) == self.ndim: return np.asarray(self).reshape(()) axes = tuple(g for k, g in enumerate(self._axes) if k not in drop) - dataset = self._assign_axes(self.dataset, axes, len(self.dims)) - return type(self)(dataset, self.dtype, self.engine) + return type(self)(self._assign_axes(self.dataset, axes)) def transpose(self, order=None): """Permute the axes, staying virtual (reversed order by default). + Only the axis map moves — unless the permutation reorders + synthetic axes past each other, which renumbers them (and the + columns carrying them) so they keep trailing the source axes in + virtual order. + Parameters ---------- order : sequence of int, optional @@ -1372,8 +1686,7 @@ def transpose(self, order=None): if sorted(order) != list(range(self.ndim)): raise ValueError(f"axes {order} do not permute a {self.ndim}-d array") axes = tuple(self._axes[a] for a in order) - dataset = self._assign_axes(self.dataset, axes, len(self.dims)) - return type(self)(dataset, self.dtype, self.engine) + return type(self)(self._assign_axes(self.dataset, axes)) def astype(self, dtype, **kwargs): """Materialize and cast the values to *dtype*.""" @@ -1381,7 +1694,7 @@ def astype(self, dtype, **kwargs): def __deepcopy__(self, memo): """Copy without the read cache; the dataset is immutable.""" - return type(self)(self.dataset, self.dtype, self.engine) + return type(self)(self.dataset) def __repr__(self): """Summarize the array on one line, as the data of a data array.