diff --git a/docs/api/atoms.md b/docs/api/atoms.md index c0b26cc3..d2685e61 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -32,7 +32,12 @@ Methods Atom.initialize Atom.initialize_from_state Atom.call + Atom.flush Atom.reset + Atom.process + Atom.iter_chunks + Atom.gather + Atom.merge Atom.save_state Atom.set_state Atom.load_state @@ -75,13 +80,65 @@ Methods Partial.get_state ``` -## Decorators +## Decorators and composition ```{eval-rst} .. autosummary:: :toctree: ../_autosummary + as_function atomized + compose +``` + +## Task atoms + +Public processing vocabulary with physical parameters only. + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + Decimate + Differentiate + Filter + Integrate + Resample + STFT +``` + +## Function forms + +Every atom has a function form exported at the top level of `xdas`: called on +data it applies eagerly, called on `...` it returns the atom. + +```{eval-rst} +.. currentmodule:: xdas +``` + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + annotate + decimate + detrend + differentiate + filter + hilbert + integrate + medfilt + pick + rechunk + resample + sliding_mean_removal + stft + taper + trigger +``` + +```{eval-rst} +.. currentmodule:: xdas.atoms ``` ## Signal processing @@ -90,13 +147,41 @@ Methods .. autosummary:: :toctree: ../_autosummary - DownSample FIRFilter IIRFilter - LFilter - MLPicker ResamplePoly - SOSFilter +``` + +## Detection and picking + +`Annotate` runs a SeisBench model window by window, `Trigger` turns the +characteristic function it produces into a pick table, and `Picker` is the +whole pipeline a weight set describes — its own filter, its own sampling rate, +its own per-phase thresholds. Each has a lowercase functional twin at the top +level of `xdas` (`xdas.annotate`, `xdas.trigger`, `xdas.pick`). + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + Annotate Trigger + Picker +``` + +## Kernel atoms + +Expert layer (`xdas.atoms.kernel`): exact stateful primitives with machine +parameters, designed by the task atoms from the data at the first call. + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + DownSample + LFilter + Polyphase + Rechunk + SOSFilter UpSample -``` \ No newline at end of file +``` diff --git a/docs/api/processing.md b/docs/api/processing.md index e4774ae1..f0f7b4f8 100644 --- a/docs/api/processing.md +++ b/docs/api/processing.md @@ -11,6 +11,9 @@ :toctree: ../_autosummary process + watch + get_source + get_writer ``` ## Loaders @@ -43,6 +46,7 @@ StreamWriter ZMQPublisher ZMQSubscriber + ResultWriter ``` ### DataArrayWriter @@ -66,4 +70,16 @@ ZMQPublisher.submit ZMQPublisher.write ZMQPublisher.result + ZMQPublisher.wait_for_subscribers + ZMQPublisher.close +``` + +### ZMQSubscriber + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + ZMQSubscriber.wait_until_subscribed + ZMQSubscriber.close ``` \ No newline at end of file diff --git a/docs/api/testing.md b/docs/api/testing.md index d26c7f79..d0f966f2 100644 --- a/docs/api/testing.md +++ b/docs/api/testing.md @@ -8,5 +8,7 @@ .. autosummary:: :toctree: ../_autosummary + assert_chunk_invariant dummy + inject_gaps ``` diff --git a/docs/api/xdas.md b/docs/api/xdas.md index 8012fb94..40ed6d73 100644 --- a/docs/api/xdas.md +++ b/docs/api/xdas.md @@ -34,6 +34,8 @@ get_sampling_interval sortby split + stack + trim_overlaps plot_availability ``` @@ -153,6 +155,7 @@ Methods .. autosummary:: :toctree: ../_autosummary + DataCollection.select DataCollection.query DataCollection.issequence DataCollection.ismapping diff --git a/docs/release-notes.md b/docs/release-notes.md index eba1ea92..5c70ebb7 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -2,34 +2,60 @@ ## 0.2.9 (unreleased) -### New Features -- **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). +### Tiles Backend +- **Tile-backed virtual arrays.** The new `xdas.virtual.tiles` module exposes a file archive as one lazy `TileArray`: slicing, integer indexing, concatenation and the numpy manipulation routines stay lazy, reductions stream, 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 — and custom engines opt in by implementing `Engine.load_tile` (@atrabattoni). +- 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 (@atrabattoni). +- Scanning scales to archives of any size: memory no longer grows with the file count and constant tile geometry is stored once — a 23-million-file archive opens in about 1 GB (@atrabattoni). +- **`xdas.sortby`** sorts a virtual data array along a dimension by coordinate value without reading any data (@atrabattoni). + +### Atoms +- **Every atom has a user-friendly functional form, and every function an atom form.** This is the norm now rather than the exception. The function is the face one writes — `xdas.resample(da, 50.0)` applies straight away — and seeding it with `...` gives the atom behind it, `xdas.resample(..., 50.0)`, ready to compose with `>>` and to be streamed. So the same code runs on a slice in memory and, chunk by chunk, on an archive that does not fit in one. Ordinary numpy expressions join in under the same seed: `20 * np.log10(np.abs(atom))` extends the pipeline instead of computing (@atrabattoni). +- **A new default processing vocabulary, in physical units.** The top-level functions no longer take the parameters of the SciPy routine underneath — an output sample count, a decimation factor, a normalised frequency — but the quantities of the measurement: `xdas.resample(da, 50.0)` names a target rate in Hz, `xdas.filter(da, (1.0, 10.0))` its corner frequencies, `xdas.medfilt(da, {"time": 0.5})` its kernel in seconds. A parameter then keeps its meaning when the sampling rate changes, which is what lets one pipeline serve whatever it is given. The new set is `xdas.filter`, `xdas.resample`, `xdas.decimate`, `xdas.integrate`, `xdas.differentiate`, `xdas.stft`, `xdas.detrend`, `xdas.taper`, `xdas.hilbert`, `xdas.medfilt`, `xdas.sliding_mean_removal`, `xdas.rechunk`, and — for picking — `xdas.annotate`, `xdas.trigger` and `xdas.pick`. The SciPy-shaped functions stay in `xdas.signal` (@atrabattoni). +- **Gap-aware chunked processing.** Stateful atoms judge the seams of their input: state carries across continuous chunks and is flushed and restarted at gaps and rate changes, and eager calls split gappy input the same way — so a result no longer depends on how the data was chunked. The new `flush()` lifecycle drains buffered tails, atoms that cannot answer chunk by chunk now raise instead of answering wrong, and `xdas.testing.assert_chunk_invariant` asserts the whole property on a pipeline of your own (@atrabattoni). +- **`process()` on every atom, with source and sink auto-dispatch.** `pipeline.process(source, out=...)` infers both ends — an array, a virtual array, a path, a `DataCollection`, `xdas.watch(dir)` or a ZeroMQ address in; a directory, a `.csv`, an address, a writer or `None` out — so one line covers everything from a slice in memory to a growing archive. Walking a collection, each leaf's result is labelled with its tree path, and an eager call or an accumulation beyond the `"memory_limit"` configuration entry — a quarter of the machine's memory by default, or of what a container or scheduler allows the process — raises rather than filling the machine (@atrabattoni). +- **Picking, end to end.** `Annotate` (replacing `MLPicker`) drives a SeisBench model with everything its weight set declares — window overlap, stacking, blinding, preprocessing — and overlaps GPU compute with transfers; `Trigger` gains per-phase thresholds and no longer loses the last pick of a record; `Picker(model)` assembles the whole chain from the weight set, so `xdas.pick(dc, model)` turns a network tree of waveforms into one flat pick table (@atrabattoni). +- **The exact machine-parameter atoms become an expert layer.** `LFilter`, `SOSFilter`, `DownSample`, `UpSample` and the new fused `Polyphase` move to `xdas.atoms.kernel`. They remain public and importable from `xdas.atoms`, but they are no longer what one reaches for: the vocabulary above designs them from the data at the first call. Resampling rides `Polyphase`, which is 2.6–8.7× faster than the chain it replaces (@atrabattoni). + +### Seismological Data - **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). +- **`xdas.stack`** collapses a level of a collection into an array dimension — `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array — lazily on virtual arrays, with `tolerance` snapping near-identical sampling grids and `join="inner"/"outer"` handling leaves that disagree (@atrabattoni). +- **`xdas.trim_overlaps`** resolves the overlaps of a data array or collection by dropping the duplicated samples, keeping the later copy by default or the earlier one with `keep="first"`, at the manifest level so lazy arrays stay lazy (@atrabattoni). +- **`DataCollection.select`**, with `obspy.Stream.select` semantics: `dc.select(station="SX00*", channel="HH?")` (@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). +- **Explicit engine configuration.** Every open function declares `engine`, `vtype` and `ctype`, and `engine` also accepts a configured `xdas.io.Engine` instance. Format-specific parameters are engine constructor arguments, validated up front — a misspelled keyword now raises instead of being silently ignored (@atrabattoni). +- **Process pools for chunk ingress and egress.** `DataArrayLoader` and `DataArrayWriter` accept `pool="processes"`, which reads and writes chunks in worker processes instead of threads — on compressed archives, an order of magnitude faster. Chunks travel through shared memory instead of being copied from one process to the other, roughly 6 times faster on the way in and 30 times on the way out, and arrive read-only. It needs no third-party dependency, and `pool="threads"` remains the default (@atrabattoni). +- **Subscribing to a stream no longer races it.** A publisher drops what it sends to a subscriber it has not registered yet, and being connected is not being subscribed — which is why streaming code is so often found sleeping and hoping. It now answers each new subscription with a greeting, in passing as it streams and without ever waiting for anyone, and `ZMQSubscriber.wait_until_subscribed()` returns when that greeting arrives: proof that nothing published from then on will be missed. A subscriber can therefore join a real-time flux at any point — the ASN one is greeted with the header describing the stream, and skips ahead to it if it arrived before the first packet, where it used to read whatever came first and fail to make sense of it. A `timeout` makes both subscribers raise rather than wait forever on a stream that has gone quiet. Replaying a recording is the one case that needs the other end to wait, since nothing a subscriber does can hold back a replay already under way: `ZMQPublisher.wait_for_subscribers()` does that, and `nsubscribers` tells how many are listening (@atrabattoni). +- `xdas.concat` can open a *new* dimension, checking that the inputs agree on their other coordinates and promoting the scalar ones that vary: stacking the components of a station is `xd.concat(traces, "channel")` (@atrabattoni). +- `sel` works on string and categorical coordinates: exact labels, lists and reordering no longer require a sorted axis (@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). +- `simplify` runs in linear time whatever the number of gaps (@atrabattoni). +- When no engine can open a file, the error now lists what each engine that recognised it said (@atrabattoni). + +### Deprecations +- `MLPicker` and `xdas.mlpicker` are deprecated in favour of `Annotate` and `Picker`; they remain as aliases until 0.4 (@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). +- **Dask virtualization is removed**, along with the `xdas.dask` module. A Dask array remains valid `DataArray` data — it is now computed on write — but `virtual=True` rejects it (@atrabattoni). +- **Opening a seismological file without naming an engine returns a nested collection, not a stacked array.** `xd.open(file)` no longer guesses that the traces of a file are synchronized; `xd.concat(traces, "channel")` is the one-liner back, and `engine="miniseed"` still gives the old shape (@atrabattoni). +- `DataCollection.query` raises a `KeyError` on an indexer naming no level of the collection, instead of silently returning everything unchanged (@atrabattoni). ### Bug Fixes -- 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). +- **ZeroMQ publishers and subscribers release their sockets.** Both ends now have `close()` and work as context managers, `process()` closes a publisher it opened itself from a `"tcp://..."` spec, and one that is simply dropped is closed by the garbage collector. Until now every publisher and subscriber ever built held its socket and its context's I/O thread for the life of the process (@atrabattoni). +- Fix a file handle leaking on every auto-detection attempt: probing a file that is not TDMS left it open, since a reader that fails to build is never handed to the `with` that would have closed it (@atrabattoni). +- Fix a STEIM-compressed `int32` miniSEED file being scanned as `float64`, the miniseed `ctype` argument being ignored, and miniSEED scans being forced to a single process (@atrabattoni). +- Fix chunked `DownSample` dropping its trailing samples when the stream length is not a multiple of the factor (@atrabattoni). +- Fix resampling losing track of the *other* coordinates of the dimension it resamples: decimating a DAS acquisition left its `station` coordinate at full length, labelling every lane with the code of the lane at its own index. Labels now follow the samples they name (@atrabattoni). +- Fix `DataCollection` coercing a `pandas.DataFrame` leaf into a broken `DataArray`: a table is now a leaf of its own kind (@atrabattoni). +- Fix a data collection keyed by zero-padded codes — a SEED location such as `"00"` — reading back from netCDF with its keys lost (@atrabattoni). +- Fix a directory sink joining its chunks along the wrong dimension when the pipeline's output does not lead with the chunked one (@atrabattoni). +- Fix the numpy dispatch overriding explicitly passed arguments with its registered defaults: `np.cumsum(da, 0)` accumulated along the last axis whatever the caller said (@atrabattoni). + +### Refactoring +- Custom engines must subclass `xdas.io.Engine`: passing a bare read function as `engine` now raises a `TypeError` (@atrabattoni). +- `Trigger` moved to `xdas.atoms.detect` with a lowercase twin `xdas.trigger`; the `xdas.trigger` module remains importable and `find_picks` is unchanged. Its `dim` default is now `"time"` instead of `"last"` (@atrabattoni). ## 0.2.8 diff --git a/docs/user-guide/faq.md b/docs/user-guide/faq.md index 0aecd8c9..6726d370 100644 --- a/docs/user-guide/faq.md +++ b/docs/user-guide/faq.md @@ -68,10 +68,11 @@ output samples. When you split data into chunks and apply the filter independent each chunk, the state is re-initialised at every boundary and the transient response distorts the result near each chunk edge. -Use the stateful atom equivalents from {py:mod}`xdas.atoms` (e.g. -{py:class}`~xdas.atoms.IIRFilter`, {py:class}`~xdas.atoms.LFilter`) inside a -{py:class}`~xdas.atoms.Sequential` pipeline. These atoms carry the filter state across -chunk boundaries automatically when used with {py:func}`~xdas.processing.process`. +Build the pipeline out of atoms — `xd.filter(..., (1.0, 10.0))` rather than +`xs.filter(da, ...)` — and run it with {py:meth}`~xdas.atoms.Atom.process`. The atoms +carry the filter state across chunk boundaries, and flush and restart it wherever the +input actually has a gap or changes sampling rate, so the chunked answer is the eager +one. {py:func}`xdas.testing.assert_chunk_invariant` checks that on your own pipeline. ## Can I use xdas with seismic data that is not DAS? diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index db3240d5..2919ee9e 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -29,7 +29,8 @@ access large multi-file acquisitions as a single seamless array. :link: pipeline/index :link-type: doc Building processing sequences with {py:class}`~xdas.atoms.Atom` objects and applying -them chunk-by-chunk on datasets larger than memory, including real-time streaming. +them chunk-by-chunk on datasets larger than memory, including real-time streaming and +phase picking with SeisBench models. ``` ```{grid-item-card} How-To Guides diff --git a/docs/user-guide/io/obspy.md b/docs/user-guide/io/obspy.md index b43796fc..7db7e059 100644 --- a/docs/user-guide/io/obspy.md +++ b/docs/user-guide/io/obspy.md @@ -62,8 +62,7 @@ for station in stations: *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. +supports goes through the same path. The engine mirrors {py:func}`obspy.read` exactly: **one contiguous ObsPy `Trace` becomes one lazy `DataArray`**, and the collection mirrors the @@ -71,7 +70,9 @@ The engine mirrors {py:func}`obspy.read` exactly: **one contiguous ObsPy 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. +The legacy `"miniseed"` engine is a different reader, kept for the code written +against it: it returns one stacked-channel array per file rather than a +collection. See [](data-formats.md). ``` ## Reading @@ -82,6 +83,7 @@ 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 numpy as np import xdas as xd dc = xd.open("NX/*/*.mseed", engine="obspy") @@ -103,7 +105,7 @@ semantics of `obspy.Stream.select`, with shell-style globbing on the keys: dc.select(station="SX00[123]", channel="HH?") ``` -`select` chooses *which* leaves are kept; {py:meth}`~xdas.DataCollection.sel` +`select` chooses *which* leaves are kept; {py:meth}`~xdas.DataMapping.sel` trims *inside* each leaf by coordinate label. Indexing works too, and reads like the seed id it is: @@ -159,15 +161,12 @@ 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. +collapsed into a dimension of a two-dimensional array with +{py:func}`xdas.stack`, which stays lazy: the level's keys become the new +coordinate, and the identifiers that do not vary 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 = xd.stack(dc["NX"]["SX001"]["00"], "channel")[0] da ``` @@ -177,10 +176,7 @@ 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 = xd.stack(xd.stack(sub["NX"], "channel"), "station")["00"][0] da ``` diff --git a/docs/user-guide/pipeline/atoms.md b/docs/user-guide/pipeline/atoms.md index ad0e82f1..0a9f9f9d 100644 --- a/docs/user-guide/pipeline/atoms.md +++ b/docs/user-guide/pipeline/atoms.md @@ -14,73 +14,150 @@ os.chdir("../../_data") # Composing a processing sequence -The xdas library provides various routines from NumPy, SciPy, and ObsPy that have been optimized for DAS DataArray objects, and which can be incorporated in a processing pipeline. See [](processing) for an explanation of the xdas processing workflows, e.g. for bigger-than-RAM datasets. Higher-level operations (FK-filters, STA/LTA detector, etc.) can be constructed from a sequence of the elementary operations implemented in xdas. To facilitate this and other user-defined operations, xdas offers a convenient framework to create and execute a (nested) sequences of atomic operations. By using sequences, built-in and user-defined processing tasks mesh seamlessly with the optimization and IO-infrastructure that xdas offers, improving the robustness and reproducibility of complex processing pipelines. +*Xdas* ships its processing vocabulary — filtering, resampling, integration, +spectra, machine-learning picking — in two interchangeable faces. The +**function** is the one you write day to day; the **atom** is the same +operation as an object, which composes into a pipeline and streams. Every +function has an atom form and every atom a function form, so a pipeline built +here runs unchanged on an array in memory and, chunk by chunk, on an archive +that does not fit in one (see [](processing.md)). -## Chaining elementary operations (atoms) +The parameters are physical throughout — a rate in hertz, a corner frequency, +a window in seconds — rather than the arguments of the SciPy routine +underneath (a decimation factor, a sample count, a normalised frequency). That +is what lets one pipeline be defined once and applied to whatever it is given: +nothing in it silently means something else at another sampling rate. The +SciPy-shaped functions are still there, in {py:mod}`xdas.signal`. -There are three "flavours" declaring the atoms that can be used to compose a sequence, illustrated by the following example: +## Applying and composing -```{code-cell} +Called on data, the function applies: + +```{code-cell} import numpy as np -import xdas -import xdas.signal as xs -from xdas.atoms import Partial, Sequential, IIRFilter - -sequence = Sequential( - [ - xs.taper(..., dim="time"), - Partial(np.square), - IIRFilter(order=4, cutoff=1.5, btype="highpass", dim="time"), - ] +import xdas as xd + +da = xd.synthetics.wavelet_wavefronts() +xd.resample(da, 25.0, dim="time") +``` + +Called on `...` — the placeholder standing for the data to come — the same +function returns the atom instead, and atoms compose with `>>`: + +```{code-cell} +pipeline = ( + xd.taper(..., dim="time") + >> xd.filter(..., (5.0, None), dim="time") + >> xd.resample(..., 25.0, dim="time") ) -sequence +pipeline ``` -In the snippet above, we define our `sequence` as an instance of the `Sequential` class, which contains three operations. The first operation applies a Tukey taper along the time dimension, encoded by the xdas implementation of the SciPy library routines (`xdas.signal`). Since this functions takes a data array as the first argument, we use `...` as a placeholder. +Nothing in that pipeline names the data it will be given: `25.0` is the rate to +land on, not a decimation factor, so the same object takes a 50 Hz record and a +1 kHz one to 25 Hz — resampling by a rational ratio where it has to. Calling it +applies it: -The second operation in this sequence is defined by the `square` operation built into NumPy. Since this function is not imported directly from xdas, using `...` as a placeholder won't work. This is where `Partial` comes in: wrapping `Partial` around `np.square` would be equivalent to `np.square(...)`, effectively converting an arbitrary routine into an xdas routine and inserting a placeholder as the first argument (to be substituted with a data array later). +```{code-cell} +result = pipeline(da) +result.plot(yincrease=False) +``` -The last operation, `IIRFilter`, instantiates a specific class dedicated to chunked execution. It inherits from the `Atom` class, which handles the logic of initialising and passing around state objects (like the filter state). This allows us to process our data one chunk at a time, without explicitly having to handle state updates and transfer. +The same pipeline can be reused: it is defined once and carries no data. -## Executing a sequence +Ordinary NumPy expressions compose too. Under the `...` seed they are *traced* +— appended to the pipeline rather than computed — so an expression reads as +mathematics: -Once the processing sequence has been defined, it can operate on data in memory by simply calling the sequence with the data array as the argument: +```{code-cell} +energy = 20 * np.log10(np.abs(xd.resample(..., 25.0, dim="time"))) +energy +``` -```{code-cell} -from xdas.synthetics import wavelet_wavefronts +## The vocabulary + +| | | +| --- | --- | +| {py:func}`~xdas.filter` | band, low- or high-pass, from a corner pair in Hz | +| {py:func}`~xdas.resample` | to a target rate, by any rational ratio | +| {py:func}`~xdas.decimate` | to a target rate, when the ratio is an integer | +| {py:func}`~xdas.integrate`, {py:func}`~xdas.differentiate` | in the coordinate's own units | +| {py:func}`~xdas.stft` | window and hop in seconds | +| {py:func}`~xdas.detrend`, {py:func}`~xdas.taper` | whole-record shaping | +| {py:func}`~xdas.hilbert` | analytic signal | +| {py:func}`~xdas.medfilt`, {py:func}`~xdas.sliding_mean_removal` | kernels in seconds or meters | +| {py:func}`~xdas.rechunk` | a streaming-cadence knob, not an operation | +| {py:func}`~xdas.annotate`, {py:func}`~xdas.trigger`, {py:func}`~xdas.pick` | machine-learning picking (see [](picking.md)) | + +Each has an atom behind it — {py:class}`~xdas.atoms.Filter`, +{py:class}`~xdas.atoms.Resample`, and so on — reached by seeding with `...`. + +Below them sits an expert layer, {py:mod}`xdas.atoms.kernel`, holding the exact +primitives these design from the data at the first call: +{py:class}`~xdas.atoms.LFilter`, {py:class}`~xdas.atoms.SOSFilter`, +{py:class}`~xdas.atoms.DownSample` and friends, which take machine parameters — +filter coefficients, integer factors — rather than physical ones. Reach for +them when you need to say exactly what runs; otherwise you should never have to +meet them. + +## Wrapping your own functions + +Any callable taking a data array as its first argument becomes an atom by +composition — `>>` wraps it: -da = wavelet_wavefronts() -result = sequence(da) -result.plot(yincrease=False) +```{code-cell} +pipeline = xd.taper(..., dim="time") >> np.square +pipeline ``` -The same sequence can be re-used, so it only needs to be defined once. +`Partial` does the same explicitly, and is what to reach for when the extra +arguments have to be given at definition time: + +```{code-cell} +from xdas.atoms import Partial -For executing a sequence on chunked data (e.g., larger-than-memory data sets), see the next section: [](processing.md). +Partial(np.percentile, ..., 90.0, axis=0) +``` ## Defining custom atoms -The `Partial` method is a convenient wrapper for simple functions that take an xdas DataArray as the first argument, which covers a lot of cases. However, more complex routines, particularly those that rely on a state, will require a more explicit treatment. Such operations can be subclassed from the `Atom` base class, and adhere to the following structure: +An operation that carries a *state* from one chunk to the next — a recursive +filter, a running mean, a detector — is written as a subclass of `Atom`. +`call` maps one input chunk to zero or more output chunks, and `flush` emits +whatever remains buffered at the end of a run: -```{code-cell} +```{code-cell} from xdas.atoms import Atom, State class MyStatefulRoutine(Atom): - def __init__(self, a, b, c=10): - super().__init__() - # Set class-specific parameters - self.a = a - self.b = b - self.c = c - # Define the state variable (if needed) - self.state = State(...) - - def initialize(self, da, **kwargs): - # Initialize state based on DataArray ``da`` - ... - - def call(self, da, **kwargs): - # Apply routine to DataArray ``da`` - ... + def __init__(self, a, dim="time"): + super().__init__() + # Configuration: kept as-is, shared by clones of this atom + self.a = a + self.dim = dim + # State: reset between runs, carried across the chunks of one run + self.buffer = State(...) + + def initialize(self, da, **flags): + # Called on the first chunk of a run, to size the state from the data + self.buffer = ... + + def call(self, da, **flags): + # Applied to every chunk; may return nothing, one chunk, or several + ... + + def flush(self): + # Called at the end of a run: emit what is still buffered + return [] ``` + +*Xdas* handles the rest: state is carried across chunk boundaries, flushed and +reset at every gap or sampling-rate change of the input, and `flush` is called +at the end of the stream. {py:func}`xdas.testing.assert_chunk_invariant` checks +that a pipeline gives the very same answer eagerly and chunk by chunk, gaps +included — the thing worth verifying before trusting a custom atom on an +archive. + +For executing a pipeline on chunked data, see the next section: +[](processing.md). diff --git a/docs/user-guide/pipeline/index.md b/docs/user-guide/pipeline/index.md index 14f99308..d79f4cee 100644 --- a/docs/user-guide/pipeline/index.md +++ b/docs/user-guide/pipeline/index.md @@ -1,8 +1,9 @@ # Pipeline Processing This section covers how to build and execute processing pipelines in *xdas*, from -composing atomic operations to applying them on larger-than-memory datasets and -streaming them over a network. +composing atomic operations to applying them on larger-than-memory datasets, +streaming them over a network, and picking seismic phases with machine-learning +models. ```{toctree} :maxdepth: 1 @@ -10,4 +11,5 @@ streaming them over a network. atoms processing streaming +picking ``` diff --git a/docs/user-guide/pipeline/picking.md b/docs/user-guide/pipeline/picking.md new file mode 100644 index 00000000..72105f9f --- /dev/null +++ b/docs/user-guide/pipeline/picking.md @@ -0,0 +1,302 @@ +# Picking seismic phases + +*Xdas* runs SeisBench models as atoms, so picking a whole network is +{py:func}`xdas.open` followed by {py:func}`xdas.pick`: + +```python +import seisbench.models as sbm +import xdas as xd + +model = sbm.PhaseNet.from_pretrained("original") + +dc = xd.open("20May2026_LabuanBajo/*.mseed") +picks = xd.pick(dc, model) +``` + +`picks` is one flat `pandas.DataFrame` for the whole network — a single array +works just as well. Everything the model needs — the filter its weights ship, +the sampling rate it was trained at, how its components are ordered, which +labels it emits and at what threshold each is picked — is read off the weight +set, so nothing above has to be repeated. + +```{note} +Unlike the rest of this guide, the code on this page is not executed when the +documentation is built: it needs SeisBench and its downloaded weights. Every +output shown was produced by running the code as written — on the day of data +described below, or, for the DAS example, on the synthetic collection built +there. +``` + +## The data + +One day of a temporary network in eastern Indonesia: 8 short-period stations, +three `SH?` components each, six of them recorded at 40 Hz and two (`LBFI`, +`LEMFI`) at 50 Hz — PhaseNet wants 100. {py:func}`xdas.open` gives the SEED +tree, without decoding anything (see [](../io/obspy.md)): + +```python +>>> dc.select(station="DBNFM") +Network: + IA: + Station: + DBNFM: + Location: + --: + Channel: + SHE: + Acquisition: + 0: + SHN: + Acquisition: + 0: + SHZ: + Acquisition: + 0: +``` + +One station, `OMBFM`, started at 00:43 and lost a second of data twelve minutes +later. That gap is not a hole in the collection — it lives in the time +coordinate, and the pipeline treats what it separates as two continuous runs, +each windowed, filtered and triggered on its own: + +```python +>>> da = dc["IA"]["OMBFM"]["--"]["SHZ"][0] +>>> [part.sizes["time"] for part in xd.split(da, "gaps")] +[28160, 3130760] +``` + +## What `xd.pick` runs + +{py:func}`xdas.pick` is the eager face of the {py:class}`~xdas.atoms.Picker` +atom, which assembles the pipeline `model.classify(stream)` runs in SeisBench — +the filter the weight set declares, resampling to its rate, +{py:class}`~xdas.atoms.Annotate`, then {py:class}`~xdas.atoms.Trigger` — and +nothing else. The stages come from the weights, so two pickers over the same +model class differ: + +```python +>>> from xdas.atoms import Picker +>>> picker = Picker(model) +>>> [type(stage).__name__ for stage in picker] +['Resample', 'Annotate', 'Trigger'] +>>> picker[-1].thresh +{'P': 0.3, 'S': 0.3} +``` + +```python +>>> obs = Picker(sbm.PhaseNet.from_pretrained("obs")) +>>> [type(stage).__name__ for stage in obs] +['_ChannelFilter', 'Resample', 'Annotate', 'Trigger'] +>>> obs[-1].thresh +{'P': 0.2, 'S': 0.1} +``` + +Three stages for one weight set and four for the other, from one architecture: +`obs` declares a 0.5 Hz highpass on its hydrophone channel (`??H`) and +`original` declares no filter at all. The target rate is the model's own — +`diting` runs at 50 Hz, not 100 — and the thresholds are read per phase off the +weight set, the noise class deliberately getting no entry so that it is never +picked. Each stage is droppable: `resample=False`, `filter=False`, an explicit +`thresh=` overriding the weight set. + +Nothing had to say that `SHE`, `SHN` and `SHZ` are the three components of one +instrument either. Walking the collection, the picker is offered each level +before the walk descends into it and claims the `channel` one, stacking it into +the component dimension the model expects (this is {py:func}`xdas.stack`, which +you can also call yourself); the station level, whose keys are not channel +codes, is walked past untouched. `components=False` turns the recognition off +and picks leaf by leaf. + +## The pick table + +```python +>>> picks.head() + network station location acquisition phase time value +0 IA DBNFM -- 0 P 2026-05-20 00:22:43.446860 0.410896 +1 IA DBNFM -- 0 P 2026-05-20 00:58:37.956860 0.544187 +2 IA DBNFM -- 0 P 2026-05-20 02:09:35.936860 0.683518 +3 IA DBNFM -- 0 P 2026-05-20 02:16:24.486860 0.733029 +4 IA DBNFM -- 0 P 2026-05-20 02:50:49.036860 0.342554 +>>> len(picks) +1918 +>>> picks["phase"].value_counts() +phase +P 985 +S 933 +Name: count, dtype: int64 +``` + +Each named level of the tree contributes a column, filled with the key the leaf +was reached by, and the columns are ordered identity first and measurement +last: the tree path (or the scalar coordinates of a single array), then the +other dimension coordinates, then the dimension the picks were found along, +then the value. A pick table therefore reads the same however the collection +was nested and whatever order the input's dimensions came in. The `channel` +level contributes nothing, because it became a dimension before the walk +reached a leaf — a pick belongs to an instrument, not to one of its components. + +The rows are not sorted: they come out lane by lane and leaf by leaf, so a +station's `P` picks precede its `S` picks. Sort the frame if the order matters. + +## Taking the pipeline apart + +The stages are ordinary atoms, so the same thing can be run one piece at a +time. Two minutes of one station, its three components stacked into an array: + +```python +>>> da = xd.stack(dc["IA"]["DBNFM"]["--"], "channel")[0] +>>> sub = da.sel(time=slice("2026-05-20T00:22:00", "2026-05-20T00:24:00")).load() +>>> sub + +[[-1827 -1832 -1783 ... -1788 -2109 -2192] + [ 688 725 623 ... 656 742 828] + [ 2990 2976 2972 ... 2867 2906 2755]] +Coordinates: + network: 'IA' + station: 'DBNFM' + location: '--' + * time (time): 2026-05-20T00:22:00.021 to 2026-05-20T00:23:59.996 + * channel (channel): ['SHE' ... 'SHZ'] +``` + +{py:class}`~xdas.atoms.Annotate` consumes the component dimension and appends +its classes as a `phase` dimension, laying the samples out last so that the +characteristic function of one phase is contiguous: + +```python +>>> cft = xd.annotate(xd.resample(sub, 100.0), model) +>>> cft + +[[nan nan nan ... nan nan nan] + [nan nan nan ... nan nan nan] + [nan nan nan ... nan nan nan]] +Coordinates: + network: 'IA' + station: 'DBNFM' + location: '--' + * phase (phase): ['N' ... 'S'] + * time (time): 2026-05-20T00:21:59.771 to 2026-05-20T00:23:59.761 +``` + +The `nan` values at both ends of each row are the ends of the record, blinded +by the model's own `annotate_batch_post` as SeisBench blinds them; the values +in between are the characteristic function. + +The `phase` coordinate carries the model's own labels, so a class is selected +by name — `cft.sel(phase="P")` — and never by position. That matters more than +it looks: the label order is a property of the weight set, `NPS` here and `PSN` +for most other PhaseNet weights, so anything positional silently addresses the +wrong phase on the next set of weights. {py:class}`~xdas.atoms.Trigger` keys its +thresholds the same way: + +```python +>>> xd.trigger(cft, thresh={"P": 0.3, "S": 0.3}) + network station location phase time value +0 IA DBNFM -- P 2026-05-20 00:22:43.461860 0.711655 +1 IA DBNFM -- P 2026-05-20 00:23:53.941860 0.578763 +``` + +A label the mapping does not name is never triggered, which is how the noise +class keeps being computed and carried without ever producing a pick. Both the +timing and the list differ a little from the whole-day run above — the first +pick moves by 15 ms and a second, weaker one appears — because the model, and +before it the resampler, see two minutes of record here instead of a day. + +## The same walk at two scales + +Nothing above is specific to seismometers: a DAS acquisition is a lane per +channel, and a DAS archive is a collection like any other. Take a synthetic +cable pair, each recorded as two consecutive acquisitions: + +```python +>>> from xdas.synthetics import randn_wavefronts +>>> da = randn_wavefronts().isel(distance=slice(None, None, 200)) +>>> das = xd.DataCollection( +... { +... cable: xd.DataCollection( +... xd.split(da.sel(distance=slice(*bounds)), 2, dim="time"), "acquisition" +... ) +... for cable, bounds in {"east": (0, 40000), "west": (60000, 100000)}.items() +... }, +... "cable", +... ) +>>> das +Cable: + east: + Acquisition: + 0: + 1: + west: + Acquisition: + 0: + 1: +``` + +Calling the picker walks that tree in memory: + +```python +>>> picker = xd.pick(..., model) +>>> picker(das) + cable acquisition distance phase time value +0 east 0 0.0 S 2024-01-01 00:00:53.560 0.346529 +1 east 0 20000.0 P 2024-01-01 00:00:39.170 0.390122 +2 east 0 40000.0 P 2024-01-01 00:00:35.640 0.402109 +3 east 0 40000.0 S 2024-01-01 00:00:39.930 0.371556 +4 west 0 60000.0 P 2024-01-01 00:00:35.560 0.305741 +5 west 0 60000.0 S 2024-01-01 00:00:39.880 0.330890 +6 west 0 80000.0 P 2024-01-01 00:00:38.960 0.452030 +``` + +`process()` walks the very same tree, but streams each leaf in chunks instead +of loading it — which is the only form left once a leaf is an archive rather +than an array. With `out=None` the results are accumulated and merged, and the +two answers are the same table: + +```python +>>> streamed = picker.process(das, chunks={"time": 5000}, out=None) +>>> streamed.equals(picker(das)) +True +``` + +The state carries across the chunks and across the acquisitions of a sequence, +and the tables are labelled as each leaf is produced, so a sink can be given +the rows directly. A `*.csv` destination is *shared* — every leaf appends to +one table, the `cable` and `acquisition` columns keeping the rows apart: + +```python +>>> picker.process(das, chunks={"time": 5000}, out="picks.csv") +``` + +A directory destination fans out instead, one subdirectory per leaf mirroring +the tree path. See [](processing.md) for the rest of the source and sink +vocabulary, and [](streaming.md) for picking a stream as it arrives. + +## How close is this to SeisBench? + +Close enough to be worth stating precisely. Stage by stage, against +`model.classify(stream)` over the 17 cached PhaseNet weight sets and the 8 +stations above: + +- the preprocessing filter is **bit-identical** — a maximum absolute difference + of exactly 0 over 9.3 million samples; +- annotation is bit-identical too, once both sides use the same batch size: + exactly 0 on all 17 weight sets, the ~1e-6 seen otherwise being float32 + convolution non-associativity between one-window and 256-window batches; +- triggering differs in exactly one place, a sample whose value equals the + threshold exactly — ObsPy's `trigger_onset` turns on at `>= thresh`, xdas at + `> thresh`. It never happened on this dataset; +- fed the *same* resampled data, `xd.pick` and `model.classify` produced 1979 + picks each over the reference day, **every one on the same sample**. + +Which leaves the resampler as the one real deviation, and it is a deviation of +passband rather than of care. SeisBench resamples with ObsPy's +`Trace.resample`, which defaults to `window="hann"` and applies that window *in +the frequency domain*: unity at DC, zero at Nyquist, and **half the amplitude +at half the input Nyquist**. *Xdas* resamples with a polyphase FIR — the only +form that can run chunk by chunk — which is flat across that band. On this +40 Hz data the two differ by 6.3 % relative RMS (median over the sweep), enough +to move most picks: 1918 picks instead of 1979 over the reference day. +Reproducing SeisBench here would mean reproducing a worse resampler, so *xdas* +does not. `resample=False` drops the stage — pass it when the data is already +at the model's rate, or to compare the two implementations on the same +waveforms. diff --git a/docs/user-guide/pipeline/processing.md b/docs/user-guide/pipeline/processing.md index 2de8e131..b5dc7e60 100644 --- a/docs/user-guide/pipeline/processing.md +++ b/docs/user-guide/pipeline/processing.md @@ -16,62 +16,66 @@ os.chdir("../../_data") ## Chunked processing: basic concepts -Given the sheer size of DAS data, it is often impossible to process an entire data set directly in memory. Hence, chunked-based processing is a necessity that requires an additional layer of computational logistics. A naive approach to chunked processing would be to load a chunk of data, apply a `Sequential` pipeline to it (see [*Composing a processing sequence*](atoms.md)), and write the resulting data to disk. Assuming that disk I/O is the limiting factor, this scenario would leave the CPU mostly idle as it has to wait for new data to be read and processed data to be written to disk. - -To maximise the pipeline throughput, xdas applies a staggered protocol of reading, processing, and writing data in parallel, as illustrated in the figure below: +Given the sheer size of DAS data, it is often impossible to process an entire +data set directly in memory. Hence, chunk-based processing is a necessity that +requires an additional layer of computational logistics. A naive approach would +be to load a chunk of data, apply a pipeline to it (see +[*Composing a processing sequence*](atoms.md)), and write the result to disk. +Assuming that disk I/O is the limiting factor, this scenario would leave the +CPU mostly idle as it has to wait for new data to be read and processed data to +be written to disk. + +To maximise the pipeline throughput, xdas applies a staggered protocol of +reading, processing, and writing data in parallel, as illustrated in the figure +below: ![](/_static/processing.svg) -With this approach, execution time is determined by the slowest of the three steps (reading, processing, writing) rather than by the sum of the three, a concept known as *latency hiding*. If, for example, reading and writing a chunk of data takes 2 seconds, and processing takes 1 second, then the total execution time per chunk is 2 seconds instead of 5. +With this approach, execution time is determined by the slowest of the three +steps (reading, processing, writing) rather than by the sum of the three, a +concept known as *latency hiding*. If, for example, reading and writing a chunk +of data takes 2 seconds, and processing takes 1 second, then the total +execution time per chunk is 2 seconds instead of 5. -A second feature of xdas, is that it automatically handles state updates and transfer. Many types of filters (e.g. recursive filters and STA/LTA algorithms) rely on some kind of memory of previously seen data, known as the *state* of the filter. The state of each filter needs to be preserved and transferred from one chunk to the next. Moreover, if the computation pipeline gets interrupted and needs to be restarted, the states need to be properly initialised for a seamless continuation. xdas offers optimised filters that handle state updates internally. +A second feature of xdas is that it handles the *state* of the pipeline. Many +operations (recursive filters, decimation, STA/LTA detectors) carry a memory of +the data already seen, which must be transferred from one chunk to the next. +Xdas does this for you, and it does it *knowing where the runs are*: state is +carried across chunks that follow one another, and flushed and reset wherever +the input has a gap or changes sampling rate — so a chunked run answers exactly +what the same pipeline answers in one piece. ## Example -The following example shows how to apply a simple processing pipeline to a large dataset. -First, build and validate the pipeline on a small in-memory subset: +Build and validate the pipeline on a small in-memory subset: ```{code-cell} -:tags: [remove-output] - import numpy as np import xdas as xd -import xdas.signal as xs -from xdas.atoms import Sequential, Partial, LFilter -from xdas.processing import process, DataArrayLoader, DataArrayWriter -from scipy.signal import iirfilter da = xd.synthetics.wavelet_wavefronts() -b, a = iirfilter(4, 0.1, btype="high") - -atom = Sequential( - [ - Partial(xs.decimate, 2, ftype="fir", dim="distance"), - LFilter(b, a, dim="time"), - Partial(np.square), - ] +pipeline = ( + xd.decimate(..., 0.02, dim="distance") + >> xd.filter(..., (5.0, None), dim="time") + >> np.square ) -monolithic = atom(da) +monolithic = pipeline(da) ``` -Then apply the same pipeline chunk-by-chunk using {py:func}`~xdas.processing.process`. -The {py:class}`~xdas.processing.DataArrayLoader` splits the input into fixed-size chunks -along a given dimension, while {py:class}`~xdas.processing.DataArrayWriter` collects and -writes each processed chunk to a directory on disk: +Then run the very same pipeline chunk by chunk with +{py:meth}`~xdas.atoms.Atom.process`, which infers what to read from the source +it is given and what to write from the `out` destination: ```{code-cell} :tags: [remove-output] -import os -os.makedirs("output", exist_ok=True) - -dl = DataArrayLoader(da, chunks={"time": 100}) -dw = DataArrayWriter("output") -chunked = process(atom, dl, dw) +chunked = pipeline.process(da, chunks={"time": 100}, out="output") +``` -assert chunked.equals(monolithic) +```{code-cell} +chunked.equals(monolithic) ``` ```{code-cell} @@ -81,7 +85,61 @@ import shutil shutil.rmtree("output") ``` -The result is identical to the monolithic run but can scale to datasets that do not fit in -memory. The loader and writer can be swapped for other variants — for example, -{py:class}`~xdas.processing.ZMQPublisher` to stream results over a network (see -[](streaming.md)). +The result is identical to the monolithic run but scales to datasets that do +not fit in memory. + +## Sources and destinations + +`process` dispatches on what it is given, so the same pipeline serves every +scale: + +| `source` | what happens | +| --- | --- | +| a `DataArray` | applied in one piece, or in chunks with `chunks=` | +| a virtual array | streamed, `chunks="auto"` following the storage layout | +| a path, a directory or a glob | opened, then streamed | +| a `DataCollection` | walked leaf by leaf, each leaf streamed | +| `xdas.watch(dir)` | a directory followed as files arrive (see [](streaming.md)) | +| `"tcp://..."` | subscribed to over ZeroMQ | + +and on the destination it is given: + +| `out` | what happens | +| --- | --- | +| `None` | the output chunks are accumulated and returned | +| a directory | written there, joined along the chunked dimension | +| a `*.csv` file | appended to, for pipelines that emit tables | +| `"tcp://..."` | published over ZeroMQ | +| a writer instance | used as configured | + +`out=None` is the convenient form and the dangerous one: the result must fit in +memory. Beyond the `"memory_limit"` configuration entry it raises rather than +filling the machine. The default is a quarter of the memory the process can +use — of the machine's, or of what a container or a batch scheduler allows it, +whichever is smaller — so it scales with where the code runs: + +```{code-cell} +xd.config.get("memory_limit") / 2**30 # GiB +``` + +Raise it, or lower it, with `xd.config.set("memory_limit", 64 * 2**30)`. + +The explicit form remains available and is what to reach for to configure the +ends themselves — a process pool, a compression, a writer of another kind: + +```{code-cell} +:tags: [remove-output] + +from xdas.processing import DataArrayLoader, DataArrayWriter, process + +os.makedirs("output", exist_ok=True) +dl = DataArrayLoader(da, chunks={"time": 100}) +dw = DataArrayWriter("output") +chunked = process(pipeline, dl, dw) +``` + +```{code-cell} +:tags: [remove-cell] + +shutil.rmtree("output") +``` diff --git a/docs/user-guide/pipeline/streaming.md b/docs/user-guide/pipeline/streaming.md index e2f5d16a..0063551f 100644 --- a/docs/user-guide/pipeline/streaming.md +++ b/docs/user-guide/pipeline/streaming.md @@ -26,7 +26,6 @@ In this section, we will mimic the use of several machine by using multithreadin ```{code-cell} import threading -import time import xdas as xd from xdas.processing import ZMQPublisher, ZMQSubscriber @@ -39,43 +38,67 @@ da = xd.testing.dummy() packets = xd.split(da, 5) ``` -We then publish the packets on machine 1. +We then publish the packets on machine 1. A publisher sends to whoever is +subscribed at that instant: anything it publishes before machine 2 has +subscribed is lost. Here we replay a finite recording and want all of it, so +machine 1 waits for its subscriber with +{py:meth}`~xdas.processing.ZMQPublisher.wait_for_subscribers`. ```{code-cell} address = f"tcp://localhost:{xd.io.get_free_port()}" publisher = ZMQPublisher(address) def publish(): + publisher.wait_for_subscribers() for packet in packets: publisher.submit(packet) - # give a chance to the subscriber to connect in time and to get the last packet - time.sleep(0.1) machine1 = threading.Thread(target=publish) machine1.start() ``` -Let's receive the packets on machine 2. +Let's receive the packets on machine 2. The subscriber is an infinite iterator, +so here we stop it once the whole stream has been received. ```{code-cell} subscriber = ZMQSubscriber(address) -packets = [] +received = [] def subscribe(): for packet in subscriber: - packets.append(packet) + received.append(packet) + if len(received) == len(packets): + break machine2 = threading.Thread(target=subscribe) machine2.start() ``` -Now we wait for machine 1 to finish sending its packet and see if everything went well. +Now we wait for both machines to be done and see if everything went well. ```{code-cell} machine1.join() -print(f"We received {len(packets)} packets!") -assert xd.concatenate(packets).equals(da) +machine2.join() +print(f"We received {len(received)} packets!") +assert xd.concatenate(received).equals(da) +``` + +Both ends hold a socket for as long as they live. Closing them releases it, +along with the background thread ZeroMQ runs underneath: + +```{code-cell} +subscriber.close() +publisher.close() +``` + +Where a publisher or a subscriber has a scope, using it as a context manager +says the same thing and cannot be forgotten: + +```python +with ZMQPublisher(address) as publisher: + for packet in packets: + publisher.submit(packet) ``` ## Using encoding @@ -90,10 +113,67 @@ import hdf5plugin address = f"tcp://localhost:{xd.io.get_free_port()}" encoding = {"chunks": (10, 10), **hdf5plugin.Zfp(accuracy=1e-6)} publisher = ZMQPublisher(address, encoding) # Add encoding here, the rest is the same +publisher.close() +``` + +## Real-time streams + +An instrument is the other kind of publisher: it streams what it measures +whether or not anyone listens, and it never waits — so it is normal, and not an +error, for a subscriber to miss whatever was published before it connected. +Subscribing to a live flux is the same two lines as above, minus the waiting: + +```python +with ZMQSubscriber("tcp://interrogator:5555") as subscriber: + for packet in subscriber: + ... ``` -{py:class}`~xdas.io.asn.ZMQSubscriber` +The waiting moves to the receiving end, where it costs the stream nothing. A +publisher answers each new subscription with a greeting, in passing as it +streams, and {py:meth}`~xdas.processing.ZMQSubscriber.wait_until_subscribed` +returns when it arrives — proof that the publisher has you on its list and that +nothing it sends from then on will be missed: + +```python +subscriber = ZMQSubscriber("tcp://interrogator:5555", timeout=10.0) +subscriber.wait_until_subscribed() +``` + +Only a publisher that publishes can acknowledge anybody, so waiting on a stream +that has gone quiet raises once `timeout` (in seconds) has passed, rather than +hanging. The same `timeout` bounds every packet read. ```{note} -Xdas also implements the ZeroMQ protocol used by the OptoDAS interrogators by ASN. Equivalent {py:class}`~xdas.io.asn.ZMQPublisher` and {py:class}`~xdas.io.asn.ZMQSubscriber` can be found in {py:mod}`xdas.io.asn`. This can be useful get data in real-time from one instrument of that kind. Note that compression is not available with that protocol yet. +Xdas also implements the ZeroMQ protocol used by the OptoDAS interrogators by ASN. Equivalent {py:class}`~xdas.io.asn.ZMQPublisher` and {py:class}`~xdas.io.asn.ZMQSubscriber` can be found in {py:mod}`xdas.io.asn`. This can be useful to get data in real-time from one instrument of that kind. Note that compression is not available with that protocol yet. +``` + +## Processing a stream + +A pipeline consumes and produces a stream by naming the address on either end, +so nothing about the pipeline itself changes between replaying an archive and +following an instrument: + +```python +pipeline.process("tcp://localhost:5556", out="tcp://*:5557") ``` + +A directory that is still being filled is the other unbounded source: +{py:func}`xdas.watch` follows it as files arrive, where a bare directory path +means "process what is there and stop". + +```python +pipeline.process(xd.watch("/incoming", engine="febus"), out="results/") +``` + +Unbounded sources are processed until they are stopped. Interrupting with +`Ctrl-C` flushes the pipeline, closes the destination cleanly and returns what +was written; `until=` stops on its own at a coordinate value: + +```python +pipeline.process(xd.watch("/incoming"), out="results/", until="2026-05-20T12:00:00") +``` + +Gaps are announced as they arrive rather than upfront — a stream cannot be +inspected ahead of time — and each one flushes and restarts the state of every +stage, exactly as it does on an archive. diff --git a/tests/conftest.py b/tests/conftest.py index 969d28e8..e2e4c3a0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,5 @@ +from contextlib import ExitStack + import pytest import xdas @@ -7,6 +9,32 @@ def pytest_configure(config): xdas.config.set("n_workers", 1) +@pytest.fixture +def opened(): + """ + Register a ZMQ endpoint to be closed when the test ends, and return it. + + Endpoints are context managers, but a test that hands one to a thread, or + lets one raise on the way in, reads better flat than nested. Wrapping the + call in this keeps the socket from outliving the test either way. + """ + with ExitStack() as stack: + yield stack.enter_context + + +@pytest.fixture +def fake_model(): + """ + Return the :func:`tests.fakemodel.fake_model` factory. + + Import the module directly instead when the models are needed at collection + time, e.g. to build a :func:`pytest.mark.parametrize` table. + """ + from tests.fakemodel import fake_model + + return fake_model + + def pytest_addoption(parser): parser.addoption( "--skip-slow", action="store_true", default=False, help="skip slow tests" diff --git a/tests/coordinates/test_dense.py b/tests/coordinates/test_dense.py index cc8cd0ac..113ce048 100644 --- a/tests/coordinates/test_dense.py +++ b/tests/coordinates/test_dense.py @@ -3,6 +3,7 @@ import pytest import xarray as xr +import xdas as xd from xdas.coordinates import DenseCoordinate, ScalarCoordinate @@ -195,6 +196,38 @@ def test_is_monotonic_increasing(self): ) assert not DenseCoordinate(times_bad)._is_monotonic_increasing() + def test_is_monotonic_increasing_duplicates(self): + # the check is strict: repeated values are not monotonic increasing. + assert not DenseCoordinate([1, 2, 2, 3])._is_monotonic_increasing() + + def test_is_monotonic_increasing_strings(self): + # string dtypes have no `subtract` loop, so the check cannot be + # arithmetic; it must still work on labels. + assert DenseCoordinate(["N", "P", "S"])._is_monotonic_increasing() + assert not DenseCoordinate(["P", "N", "S"])._is_monotonic_increasing() + assert not DenseCoordinate(["N", "P", "P"])._is_monotonic_increasing() + + @pytest.mark.parametrize( + "values", + [ + [1, 2, 3], + [1, 2, 2, 3], + [3, 2, 1], + [1], + [], + np.array([0, 1, 2], dtype="datetime64[s]"), + np.array([0, 1, 1, 2], dtype="datetime64[s]"), + np.array([2, 1, 0], dtype="datetime64[s]"), + ], + ) + def test_is_monotonic_increasing_matches_arithmetic(self, values): + # the pandas-based check must agree with the arithmetic one it replaced + # wherever the latter was defined. + coord = DenseCoordinate(values) + zero = np.timedelta64(0) if np.issubdtype(coord.dtype, np.datetime64) else 0 + expected = bool(np.all(np.diff(coord.values) > zero)) + assert coord._is_monotonic_increasing() == expected + def test_add(self): coord = DenseCoordinate([1.0, 2.0, 3.0], "x") result = coord + 1.0 @@ -232,6 +265,151 @@ def test_collect_from_dataset_object_dtype(self): assert "x" in result +class TestDenseCoordinateStringSelection: + @staticmethod + def dataarray(): + return xd.DataArray( + np.arange(12).reshape(4, 3), + {"time": [0.0, 1.0, 2.0, 3.0], "phase": ["N", "P", "S"]}, + ) + + def test_sel_scalar(self): + result = self.dataarray().sel(phase="P") + assert np.array_equal(result.values, [1, 4, 7, 10]) + assert result.coords["phase"].values == "P" + + def test_sel_list(self): + result = self.dataarray().sel(phase=["P", "S"]) + assert np.array_equal(result.values, [[1, 2], [4, 5], [7, 8], [10, 11]]) + assert np.array_equal(result.coords["phase"].values, ["P", "S"]) + + def test_sel_list_reorders(self): + result = self.dataarray().sel(phase=["S", "P"]) + assert np.array_equal(result.values, [[2, 1], [5, 4], [8, 7], [11, 10]]) + assert np.array_equal(result.coords["phase"].values, ["S", "P"]) + + def test_sel_slice(self): + result = self.dataarray().sel(phase=slice("N", "P")) + assert np.array_equal(result.values, [[0, 1], [3, 4], [6, 7], [9, 10]]) + assert np.array_equal(result.coords["phase"].values, ["N", "P"]) + + def test_sel_missing_label(self): + with pytest.raises(KeyError): + self.dataarray().sel(phase="Q") + + +class TestDenseCoordinateUnsortedSelection: + """ + Label selection on an axis whose labels are not in sorted order. + + A categorical axis is unordered by nature — a SeisBench phase axis is + ``"PSN"`` on 14 of the 17 cached ``PhaseNet`` weight sets and ``"NPS"`` on + the other three — yet every label still designates exactly one position, so + an exact look-up is well defined. Ordered look-ups (a slice, or ``method``) + are not, and stay guarded. + """ + + @staticmethod + def dataarray(): + return xd.DataArray( + np.arange(12).reshape(4, 3), + {"time": [0.0, 1.0, 2.0, 3.0], "phase": ["P", "S", "N"]}, + ) + + def test_the_axis_is_not_monotonic_increasing(self): + assert not self.dataarray()["phase"]._is_monotonic_increasing() + + def test_sel_scalar(self): + result = self.dataarray().sel(phase="P") + assert np.array_equal(result.values, [0, 3, 6, 9]) + assert result.coords["phase"].values == "P" + + def test_sel_list(self): + result = self.dataarray().sel(phase=["P", "S"]) + assert np.array_equal(result.values, [[0, 1], [3, 4], [6, 7], [9, 10]]) + assert np.array_equal(result.coords["phase"].values, ["P", "S"]) + + def test_sel_list_preserves_the_requested_order(self): + result = self.dataarray().sel(phase=["S", "P"]) + assert np.array_equal(result.values, [[1, 0], [4, 3], [7, 6], [10, 9]]) + assert np.array_equal(result.coords["phase"].values, ["S", "P"]) + + def test_sel_missing_label(self): + with pytest.raises(KeyError): + self.dataarray().sel(phase="Q") + + def test_sel_slice_is_still_refused(self): + # a slice resolves its bounds by searching, so the guard still catches + # it and sends it down the split-on-overlaps path, which cannot order + # these labels either: the selection fails rather than returning + # something arbitrary. The `TypeError` is numpy's, from differencing + # strings while looking for the overlaps. + with ( + pytest.raises(TypeError), + pytest.warns(match="not monotonic increasing"), + ): + self.dataarray().sel(phase=slice("P", "S")) + + def test_sel_with_method_is_still_refused(self): + # a neighbour search is an ordered look-up too. + with pytest.raises(NotImplementedError, match="overlaps"): + self.dataarray().sel(phase="P", method="nearest") + + +class TestDenseCoordinateUnsortedNumericSelection: + """Exact selection needs no order on numeric axes either.""" + + @staticmethod + def dataarray(): + return xd.DataArray(np.arange(4), {"x": [30.0, 10.0, 40.0, 20.0]}) + + def test_sel_scalar(self): + assert self.dataarray().sel(x=40.0).values == 2 + + def test_sel_list_preserves_the_requested_order(self): + result = self.dataarray().sel(x=[40.0, 10.0]) + assert np.array_equal(result.values, [2, 1]) + assert np.array_equal(result.coords["x"].values, [40.0, 10.0]) + + def test_sel_datetime_scalar(self): + da = xd.DataArray( + np.arange(3), + {"time": np.array(["2000-01-03", "2000-01-01", "2000-01-02"], "M8[s]")}, + ) + assert da.sel(time="2000-01-01").values == 1 + + def test_sorted_axes_are_untouched(self): + da = xd.DataArray(np.arange(4), {"x": [10.0, 20.0, 30.0, 40.0]}) + assert da.sel(x=30.0).values == 2 + assert np.array_equal(da.sel(x=slice(20.0, 30.0)).values, [1, 2]) + assert da.sel(x=24.0, method="nearest").values == 1 + + +class TestDenseCoordinateDuplicatedSelection: + """ + Duplicated labels are ambiguous, whatever the guard does. + + The ``is_unique`` half of the monotonicity check keeps routing slices + through the split-on-overlaps path, and an exact look-up is refused by + ``pandas`` because it cannot name a single position. + """ + + @staticmethod + def dataarray(): + return xd.DataArray(np.arange(4), {"x": [10.0, 20.0, 20.0, 30.0]}) + + def test_the_axis_is_not_monotonic_increasing(self): + assert not self.dataarray()["x"]._is_monotonic_increasing() + + def test_sel_slice_still_takes_the_guarded_path(self): + with pytest.warns(match="not monotonic increasing"): + self.dataarray().sel(x=slice(10.0, 30.0)) + + def test_sel_scalar_is_refused(self): + with pytest.raises(pd.errors.InvalidIndexError): + self.dataarray().sel(x=20.0) + + class TestDenseCoordinateToRegular: def test_never_regular(self): coord = DenseCoordinate([0.0, 1.0, 2.0], "x") diff --git a/tests/fakemodel.py b/tests/fakemodel.py new file mode 100644 index 00000000..df1b0b83 --- /dev/null +++ b/tests/fakemodel.py @@ -0,0 +1,259 @@ +""" +Offline stand-ins for SeisBench weight sets, for the machine-learning tests. + +Every interesting property of a SeisBench picker belongs to the *weight set*, +not to the architecture: ``component_order``, ``in_channels``, the *order* of +``labels``, ``sampling_rate``, which keys ``default_args`` declares, and whether +the weights ship their own preprocessing filter. Surveying the 17 cached +``PhaseNet`` weight sets, all of these vary — see :data:`WEIGHT_SETS` for the +five archetypes that between them cover the whole range. + +:class:`FakeModel` is a real :class:`seisbench.models.WaveformModel` subclass +whose forward pass is a fixed, trivial function of its input, so tests can +assert exact values without downloading weights or touching the network. It is +deliberately not a picker that works: it is a picker-shaped contract. + +Import it either directly, for :func:`pytest.mark.parametrize` tables:: + + from tests.fakemodel import WEIGHT_SETS, fake_model + +or through the ``fake_model`` fixture declared in ``tests/conftest.py``. +""" + +import numpy as np +import torch +from seisbench.models import WaveformModel + +#: Archetypal weight sets, one per combination the real ``PhaseNet`` sets show. +#: +#: The values mirror the cached metadata of the weight set each is named after, +#: except that ``blinding`` and ``overlap`` are scaled to the toy window of +#: :func:`fake_model` (``in_samples=8`` rather than 3001). Pass a different +#: ``in_samples`` and you must pass a matching ``default_args`` with it. +WEIGHT_SETS = { + # ENZ, NPS labels, declares both overlap and blinding, no thresholds. + "original": { + "component_order": "ENZ", + "in_channels": 3, + "labels": "NPS", + "sampling_rate": 100, + "default_args": {"overlap": 0.5, "blinding": (1, 1)}, + }, + # ZNE, NPS labels, and the one weight set that does not run at 100 Hz. + "diting": { + "component_order": "ZNE", + "in_channels": 3, + "labels": "NPS", + "sampling_rate": 50, + "default_args": { + "P_threshold": 0.3, + "S_threshold": 0.3, + "blinding": (1, 1), + }, + }, + # ZNE, PSN labels, thresholds far from the 0.3 fallback and far apart. + "geofon": { + "component_order": "ZNE", + "in_channels": 3, + "labels": "PSN", + "sampling_rate": 100, + "default_args": { + "P_threshold": 0.5704745853696115, + "S_threshold": 0.07349645833964447, + "blinding": (1, 1), + }, + }, + # Z12H with a fourth (hydrophone) channel and a per-channel filter, and no + # blinding key at all — the combination that breaks assumptions. + "obs": { + "component_order": "Z12H", + "in_channels": 4, + "labels": "PSN", + "sampling_rate": 100, + "default_args": {"P_threshold": 0.2, "S_threshold": 0.1}, + "filter_args": {"??H": ["highpass"]}, + "filter_kwargs": {"??H": {"freq": 0.5}}, + }, + # No blinding either. The flat filter is *invented*: `obs`'s per-channel + # one is the only filter any cached PhaseNet weight set declares, and the + # flat form still has to work, so one preset carries it. + "volpick": { + "component_order": "ZNE", + "in_channels": 3, + "labels": "PSN", + "sampling_rate": 100, + "default_args": {"P_threshold": 0.39, "S_threshold": 0.34}, + "filter_args": ("highpass",), + "filter_kwargs": {"freq": 1.0}, + }, +} + + +class FakeModel(WaveformModel): + """ + A ``WaveformModel`` with no weights and a closed-form forward pass. + + The forward pass is + + .. math:: y_{b,k,t} = k + \\sum_c (c + 1) \\, x_{b,c,t} + + so the output reveals both which input slot each component landed in (the + ``c + 1`` weights) and the order of the classes (the ``k`` offset). + Preprocessing is a peak normalisation scaled by the ``scale`` annotate + argument, which makes the result depend on the whole window and therefore + pins the sliding-window stitching rather than just the arithmetic. + + Parameters + ---------- + component_order : str, optional + Component letters in the order the model's input slots expect them, + e.g. ``"ENZ"``, ``"ZNE"`` or ``"Z12H"``. Defaults to ``"ZNE"``. + in_channels : int, optional + Number of input slots. Defaults to 3; the ``obs`` archetype uses 4. + labels : str or list of str, optional + Output class labels, *in order*. Defaults to ``"PSN"``. + classes : int, optional + Number of output classes. Defaults to ``len(labels)``. + sampling_rate : float, optional + Sampling rate the weights were trained at. Defaults to 100. + in_samples : int, optional + Model window length in samples. Defaults to 8, small enough that a + whole test record fits in a golden array. + default_args : dict, optional + Exactly what the weight set declares. Left empty by default, so tests + can add or omit ``blinding``, ``overlap`` and ``*_threshold`` keys one + at a time. + annotate_args : dict, optional + Per-instance overrides of the class-level ``_annotate_args`` defaults, + given as ``{key: default_value}``. Use this to move the *fallback* a + weight set falls back to, as opposed to what it declares. + filter_args, filter_kwargs : optional + Preprocessing filter the weights ship, flat (a tuple plus a dict) or + per channel pattern (two dicts keyed identically). + piggyback : float, optional + When given, ``annotate_batch_pre`` returns the ``(batch, piggyback)`` + pair rather than a bare tensor, and ``annotate_batch_post`` multiplies + by it. ``None`` (default) keeps the bare-tensor form. + + Attributes + ---------- + seen_argdicts : list of dict + A copy of the argdict every ``annotate_batch_pre`` call received, in + order. Empty until the model is driven. + seen_piggybacks : list + The piggyback every ``annotate_batch_post`` call received, in order. + """ + + _annotate_args = WaveformModel._annotate_args.copy() + _annotate_args["*_threshold"] = ("Detection threshold for the provided phase", 0.3) + _annotate_args["blinding"] = ( + "Number of prediction samples to discard on each side of each window", + (0, 0), + ) + _annotate_args["scale"] = ("Gain applied by annotate_batch_pre", 1.0) + + def __init__( + self, + component_order="ZNE", + in_channels=3, + labels="PSN", + classes=None, + sampling_rate=100, + in_samples=8, + default_args=None, + annotate_args=None, + filter_args=None, + filter_kwargs=None, + piggyback=None, + ): + super().__init__( + citation="fake model, for tests only", + component_order=component_order, + sampling_rate=sampling_rate, + output_type="array", + default_args=dict(default_args or {}), + in_samples=in_samples, + pred_sample=(0, in_samples), + labels=labels, + filter_args=filter_args, + filter_kwargs=filter_kwargs, + ) + self.in_channels = in_channels + self.classes = len(labels) if classes is None else classes + self.piggyback = piggyback + self.seen_argdicts = [] + self.seen_piggybacks = [] + # `_annotate_args` is a class attribute in SeisBench; shadow it per + # instance so one test can move a fallback without leaking to the next. + self._annotate_args = dict(type(self)._annotate_args) + for key, value in (annotate_args or {}).items(): + doc = self._annotate_args.get(key, ("Fake annotate argument", None))[0] + self._annotate_args[key] = (doc, value) + + def forward(self, x): + """Return ``(batch, classes, samples)`` as a fixed function of *x*.""" + weights = torch.arange( + 1, self.in_channels + 1, dtype=x.dtype, device=x.device + ).reshape(1, -1, 1) + pooled = (x * weights).sum(dim=-2) + offsets = torch.arange(self.classes, dtype=x.dtype, device=x.device).reshape( + 1, -1, 1 + ) + return pooled.unsqueeze(-2) + offsets + + def annotate_batch_pre(self, batch, argdict): + """Peak-normalise *batch*, recording the argdict it was given.""" + self.seen_argdicts.append(dict(argdict)) + scale = self._argdict_get_with_default(argdict, "scale") + peak = batch.abs().amax(dim=-1, keepdim=True) + normalized = scale * batch / (peak + 1e-10) + if self.piggyback is None: + return normalized + return normalized, self.piggyback + + def annotate_batch_post(self, batch, piggyback, argdict): + """Transpose to ``(batch, samples, classes)`` and blind the edges.""" + self.seen_piggybacks.append(piggyback) + batch = torch.transpose(batch, -1, -2) + if piggyback is not None: + batch = batch * piggyback + prenan, postnan = self._argdict_get_with_default(argdict, "blinding") + if prenan > 0: + batch[..., :prenan, :] = np.nan + if postnan > 0: + batch[..., -postnan:, :] = np.nan + return batch + + +def fake_model(name=None, **overrides): + """ + Build a :class:`FakeModel`, optionally from a :data:`WEIGHT_SETS` preset. + + Parameters + ---------- + name : str, optional + Key of :data:`WEIGHT_SETS` to start from. ``None`` (default) starts + from :class:`FakeModel`'s own defaults. + **overrides + Passed to :class:`FakeModel`, overriding the preset key by key. + + Returns + ------- + FakeModel + A model in evaluation mode on the CPU. + + Examples + -------- + >>> from tests.fakemodel import fake_model + >>> model = fake_model("obs") + >>> model.component_order, model.in_channels, list(model.labels) + ('Z12H', 4, ['P', 'S', 'N']) + >>> "blinding" in model.default_args + False + >>> model = fake_model("original", in_samples=16) + >>> model.component_order, list(model.labels), model.in_samples + ('ENZ', ['N', 'P', 'S'], 16) + """ + kwargs = dict(WEIGHT_SETS[name]) if name is not None else {} + kwargs.update(overrides) + return FakeModel(**kwargs).eval() diff --git a/tests/io/test_asn.py b/tests/io/test_asn.py index 624c5930..c468e136 100644 --- a/tests/io/test_asn.py +++ b/tests/io/test_asn.py @@ -1,6 +1,5 @@ import json import threading -import time import h5py import numpy as np @@ -16,6 +15,10 @@ def get_free_local_address(): return f"tcp://localhost:{port}" +TIMEOUT = 60.0 +"""Seconds any wait is given before failing, so that no test can hang.""" + + da_float32 = xd.testing.dummy(shape=(100, 10), step=(0.1, 10.0), dtype="float32") da_int16 = xd.testing.dummy(shape=(100, 10), step=(0.1, 10.0), dtype="int16") @@ -41,11 +44,16 @@ def test_roi_end_before_sensor_distances(self): class TestASNEnginePublisher: + @pytest.fixture(autouse=True) + def _close_endpoints(self, opened): + """Close whatever this test opens, so that no socket outlives it.""" + self.opened = opened + def test_write_method(self): from xdas.io.asn import ZMQPublisher as ASNZMQPublisher address = get_free_local_address() - pub = ASNZMQPublisher(address) + pub = self.opened(ASNZMQPublisher(address)) pub.write(da_float32) @@ -187,6 +195,11 @@ def test_read_keeps_unevenly_decimated_rois_irregular(self, tmp_path): class TestZMQPublisher: + @pytest.fixture(autouse=True) + def _close_endpoints(self, opened): + """Close whatever this test opens, so that no socket outlives it.""" + self.opened = opened + def test_get_header(self): header = ZMQPublisher._get_header(da_float32) assert header["bytesPerPackage"] == 40 @@ -203,35 +216,29 @@ def test_get_header(self): def test_init_conect_set_header(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) pub.submit(da_float32) - time.sleep(0.01) assert pub.header == ZMQPublisher._get_header(da_float32) def test_send_header(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) pub.submit(da_float32) - time.sleep(0.01) - socket = self.get_socket(address) + socket = self.get_socket(pub) pub.submit(da_float32) # a packet must be sent once subscriber is connected - time.sleep(0.01) assert socket.recv() == json.dumps(pub.header).encode("utf-8") def test_send_data(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) pub.submit(da_float32) - time.sleep(0.01) - socket = self.get_socket(address) + socket = self.get_socket(pub) pub.submit(da_float32) # a packet must be sent once subscriber is connected - time.sleep(0.01) socket.recv() # header message = socket.recv() assert message[:8] == da_float32["time"][0].values.astype("M8[ns]").tobytes() assert message[8:] == da_float32.data.tobytes() pub.submit(da_int16) - time.sleep(0.01) socket.recv() # header message = socket.recv() assert message[:8] == da_int16["time"][0].values.astype("M8[ns]").tobytes() @@ -239,14 +246,12 @@ def test_send_data(self): def test_send_chunks(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 10) pub.submit(chunks[0]) - time.sleep(0.01) - socket = self.get_socket(address) + socket = self.get_socket(pub) for chunk in chunks[1:]: pub.submit(chunk) - time.sleep(0.01) assert socket.recv() == json.dumps(pub.header).encode("utf-8") for chunk in chunks[1:]: # first was sent before subscriber connected message = socket.recv() @@ -255,18 +260,15 @@ def test_send_chunks(self): def test_several_subscribers(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 10) pub.submit(chunks[0]) - time.sleep(0.01) - socket1 = self.get_socket(address) + socket1 = self.get_socket(pub) for chunk in chunks[1:5]: pub.submit(chunk) - time.sleep(0.01) - socket2 = self.get_socket(address) + socket2 = self.get_socket(pub, 2) for chunk in chunks[5:]: pub.submit(chunk) - time.sleep(0.01) assert socket1.recv() == json.dumps(pub.header).encode("utf-8") for chunk in chunks[1:]: # first was sent before subscriber connected message = socket1.recv() @@ -280,18 +282,15 @@ def test_several_subscribers(self): def test_change_header(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 10) pub.submit(chunks[0]) - time.sleep(0.01) - socket = self.get_socket(address) + socket = self.get_socket(pub) for chunk in chunks[1:5]: pub.submit(chunk) header1 = pub.header - time.sleep(0.01) for chunk in chunks[5:]: pub.submit(chunk.isel(distance=slice(0, 5))) - time.sleep(0.01) header2 = pub.header assert socket.recv() == json.dumps(header1).encode("utf-8") for chunk in chunks[1:5]: # first was sent before subscriber connected @@ -304,21 +303,46 @@ def test_change_header(self): assert message[:8] == chunk["time"][0].values.astype("M8[ns]").tobytes() assert message[8:] == chunk.isel(distance=slice(0, 5)).data.tobytes() - def get_socket(self, address): - socket = zmq.Context().socket(zmq.SUB) - socket.connect(address) + def get_socket(self, pub, nsubscribers=1): + """Subscribe to *pub* and hand back the socket, once *pub* knows of it.""" + context = self.opened(zmq.Context()) + socket = self.opened(context.socket(zmq.SUB)) + socket.setsockopt(zmq.RCVTIMEO, round(1000 * TIMEOUT)) + socket.connect(pub.address) socket.setsockopt(zmq.SUBSCRIBE, b"") - time.sleep(0.01) + pub.wait_for_subscribers(nsubscribers, TIMEOUT) return socket + def test_wait_for_subscribers_times_out(self): + address = get_free_local_address() + pub = self.opened(ZMQPublisher(address)) + assert pub.nsubscribers == 0 + with pytest.raises(TimeoutError, match="0 of the 1 subscriber"): + pub.wait_for_subscribers(timeout=0.1) + + def test_unsubscribing_is_accounted_for(self): + address = get_free_local_address() + pub = self.opened(ZMQPublisher(address)) + socket = self.get_socket(pub) + assert pub.nsubscribers == 1 + socket.setsockopt(zmq.UNSUBSCRIBE, b"") + assert pub._read_subscriptions(TIMEOUT) # blocks until the cancellation + assert pub.nsubscribers == 0 + class TestZMQSubscriber: + @pytest.fixture(autouse=True) + def _close_endpoints(self, opened): + """Close whatever this test opens, so that no socket outlives it.""" + self.opened = opened + def test_one_chunk(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = [da_float32] - threading.Thread(target=self.publish, args=(pub, chunks)).start() - sub = ZMQSubscriber(address) + thread = threading.Thread(target=self.publish, args=(pub, chunks)) + thread.start() + sub = self.get_subscriber(address) assert sub.address == address assert sub.packet_size == 4008 assert sub.shape == (100, 10) @@ -331,8 +355,10 @@ def test_one_chunk(self): assert sub.delta == np.timedelta64(100, "ms") result = next(sub) assert result.equals(da_float32) + thread.join() # the socket only ever belongs to one thread at a time chunks = [da_int16] - threading.Thread(target=self.publish, args=(pub, chunks)).start() + thread = threading.Thread(target=self.publish, args=(pub, chunks)) + thread.start() result = next(sub) assert sub.packet_size == 2008 assert sub.dtype == np.int16 @@ -340,10 +366,10 @@ def test_one_chunk(self): def test_several_chunks(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 5) threading.Thread(target=self.publish, args=(pub, chunks)).start() - sub = ZMQSubscriber(address) + sub = self.get_subscriber(address) assert sub.packet_size == 808 assert sub.shape == (20, 10) assert sub.dtype == np.float32 @@ -359,15 +385,15 @@ def test_several_chunks(self): def test_several_subscribers(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 5) thread = threading.Thread(target=self.publish, args=(pub, chunks[:2])) thread.start() - sub1 = ZMQSubscriber(address) + sub1 = self.get_subscriber(address) thread.join() - thread = threading.Thread(target=self.publish, args=(pub, chunks[2:])) + thread = threading.Thread(target=self.publish, args=(pub, chunks[2:], 2)) thread.start() - sub2 = ZMQSubscriber(address) + sub2 = self.get_subscriber(address) for chunk in chunks: result = next(sub1) @@ -378,21 +404,21 @@ def test_several_subscribers(self): def test_change_header(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 5) chunks = [chunk.isel(distance=slice(0, 5)) for chunk in chunks[:2]] + chunks[2:] threading.Thread(target=self.publish, args=(pub, chunks)).start() - sub = ZMQSubscriber(address) + sub = self.get_subscriber(address) for chunk in chunks: result = next(sub) assert result.equals(chunk) def test_roiDec(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = [da_float32] threading.Thread(target=self.publish, args=(pub, chunks)).start() - sub = ZMQSubscriber(address) + sub = self.get_subscriber(address) message = ( b"{\n" b' "bytesPerPackage": 64008,\n' @@ -450,16 +476,79 @@ def test_roiDec(self): def test_iter(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 5) threading.Thread(target=self.publish, args=(pub, chunks)).start() - sub = ZMQSubscriber(address) + sub = self.get_subscriber(address) sub = (chunk for _, chunk in zip(range(5), sub)) result = xd.concat(list(sub)) assert result.equals(da_float32) - def publish(self, pub, chunks): - time.sleep(0.01) + def get_subscriber(self, address): + """A subscriber that raises rather than waiting forever on a lost stream.""" + return self.opened(ZMQSubscriber(address, timeout=TIMEOUT)) + + def publish(self, pub, chunks, nsubscribers=1): + """ + Replay *chunks* once the subscribers are there to get them. + + These tests check a recording arrives whole, so the replay waits. A + real-time publisher does not, which is what + :meth:`test_subscriber_joins_a_real_time_flux` covers. + """ + pub.wait_for_subscribers(nsubscribers, TIMEOUT) for chunk in chunks: pub.submit(chunk) - time.sleep(0.01) + + def test_subscriber_joins_a_real_time_flux(self): + # An interrogator streams what it measures whether or not anyone is + # listening. Whoever connects gets the welcome header and picks the + # stream up wherever it happens to land — never the whole of it. + address = get_free_local_address() + pub = self.opened(ZMQPublisher(address)) + chunks = xd.split(da_float32, 10) + stop = threading.Event() + + def flux(): + while not stop.is_set(): + for chunk in chunks: + if stop.is_set(): + return + pub.submit(chunk) + + thread = threading.Thread(target=flux) + thread.start() + try: + sub = self.get_subscriber(address) + received = [next(sub) for _ in range(5)] + finally: + stop.set() + thread.join() + + assert sub.shape == (10, 10) + for chunk in received: + assert any(chunk.equals(published) for published in chunks) + + def test_init_skips_data_until_the_header(self): + # A subscriber that beat the first publication to the socket gets no + # welcome message, and can be handed data before the header arrives. + address = get_free_local_address() + pub = self.opened(ZMQPublisher(address)) + header = json.dumps(ZMQPublisher._get_header(da_float32)).encode("utf-8") + + def publish(): + pub.wait_for_subscribers(timeout=TIMEOUT) + pub._send_data(da_float32) + pub._send_message(header) + pub._send_data(da_float32) + + threading.Thread(target=publish).start() + sub = self.get_subscriber(address) + assert sub.shape == (100, 10) + assert next(sub).equals(da_float32) + + def test_timeout(self): + address = get_free_local_address() + self.opened(ZMQPublisher(address)) # binds, never publishes anything + with pytest.raises(TimeoutError, match="no message received"): + ZMQSubscriber(address, timeout=0.1) diff --git a/tests/test_atoms.py b/tests/test_atoms.py index 05266c1b..b54e003c 100644 --- a/tests/test_atoms.py +++ b/tests/test_atoms.py @@ -11,14 +11,16 @@ DownSample, FIRFilter, IIRFilter, - MLPicker, + LFilter, Partial, + Polyphase, ResamplePoly, Sequential, UpSample, + atomized, ) +from xdas.atoms.core import _whole_record from xdas.signal import lfilter -from xdas.synthetics import randn_wavefronts class TestAbstractAtom: @@ -205,6 +207,8 @@ def test_upsample(self): result = xd.concat([atom(chunk, chunk_dim="time") for chunk in chunks], "time") assert result.equals(expected) + assert UpSample(1, dim="time")(da).equals(da) + def test_firfilter(self): da = xd.testing.dummy() chunks = xd.split(da, 6, "time") @@ -213,7 +217,10 @@ def test_firfilter(self): expected["time"] -= np.timedelta64(10, "ms") * 5 atom = FIRFilter(11, 10.0, "lowpass", dim="time") result = atom(da) - assert result.equals(expected) + # The polyphase form accumulates the taps in a different order than + # `lfilter`, so the two agree to rounding rather than exactly. + assert np.allclose(result.values, expected.values, atol=1e-16, rtol=1e-11) + assert result.coords.equals(expected.coords) result = xd.concat([atom(chunk, chunk_dim="time") for chunk in chunks], "time") assert np.allclose(result.values, expected.values, atol=1e-16, rtol=1e-11) @@ -256,64 +263,96 @@ def test_nothing_to_do(self): assert result.equals(da) -class TestMLPicker: - @pytest.mark.slow - def test_picker(self): - from seisbench.models import PhaseNet +class TestPolyphase: + """The fused kernel against the upsample/filter/downsample chain it replaces.""" - model = PhaseNet.from_pretrained("diting") - picker = MLPicker(model, "time", device="cpu", component_strategy="Z") - da = randn_wavefronts() - # da = da.isel(time=slice(0, 5000)) TODO: why not faster ? - expected = picker(da) - chunks = xd.split(da, 4, "time") - result = xd.concat([picker(chunk, chunk_dim="time") for chunk in chunks]) - assert result.equals(expected) + @staticmethod + def chain(taps, up, down, dim="time"): + """The unfused formulation, kept here as the reference.""" + + def apply(da): + # UpSample already carries the `up` energy scaling of the taps. + da = UpSample(up, dim=dim)(da) if up > 1 else da + da = LFilter(taps, [1.0], dim)(da) + da[dim] -= xd.get_sampling_interval(da, dim, cast=False) * ( + (len(taps) - 1) // 2 + ) + return DownSample(down, dim)(da) if down > 1 else da + + return apply + + @pytest.mark.parametrize("up, down", [(1, 1), (1, 2), (1, 5), (2, 1), (2, 5)]) + def test_equals_the_explicit_chain(self, up, down): + da = xd.testing.dummy(shape=(101, 5)) + taps = sp.firwin(20 * max(up, down) + 1, 0.4 / max(up, down)) + expected = self.chain(taps, up, down)(da) + result = Polyphase(up * taps, up, down, "time")(da) + assert result.sizes["time"] == expected.sizes["time"] + np.testing.assert_allclose(result.values, expected.values, atol=1e-15) + assert result.coords.equals(expected.coords) - @pytest.mark.slow - def test_compare_with_seisbench(self): - import obspy - from seisbench.models import PhaseNet - - model = PhaseNet.from_pretrained("original") # works at 100 Hz - model.to_preferred_device() - picker = MLPicker(model, "time", component_strategy="clone") - - # generate one trace - da = randn_wavefronts() # 100 Hz - da = da.isel(distance=slice(0, 1)) - - # xdas - result = picker(da) - - # convert to one stream with clonning - st = da.to_stream() - tr = st[0] - st = obspy.Stream() - for component in model.component_order: - _tr = tr.copy() - _tr.stats.component = component - st.append(_tr) - - # seisbench - expected = model.annotate(st) - expected = xd.DataArray.from_stream(expected) - - # align because of different overlap managment - _result = result.sel(time=slice(expected["time"][0].values, None)) - _result = _result.isel(distance=0) - _expected = expected.sel(time=slice(None, result["time"][-1].values)) - _expected = _expected.transpose("time", "channel") - - # remove unfinished end part - _result = _result[:-1000] - _expected = _expected[:-1000] - - # check equal by removing the + @pytest.mark.parametrize("up, down", [(1, 2), (2, 5), (3, 10)]) + def test_pinned_against_upfirdn(self, up, down): + # An eager call emits ceil(size * up / down) samples, the leading ones + # of scipy's own upfirdn output; the group delay moves the coordinate, + # never the values. + da = xd.testing.dummy(shape=(101, 5)) + taps = sp.firwin(21, 0.4 / max(up, down)) + result = Polyphase(taps, up, down, "time")(da) + full = sp.upfirdn(taps, da.values, up, down, axis=0) + assert result.sizes["time"] == -(-101 * up // down) np.testing.assert_allclose( - _result.values, _expected.values, rtol=1e-5, atol=1e-7 + result.values, full[: result.sizes["time"]], atol=1e-15 ) - np.testing.assert_array_max_ulp(_result.values, _expected.values, maxulp=300) + + @pytest.mark.parametrize("up, down", [(1, 2), (2, 5), (2, 3)]) + @pytest.mark.parametrize("nchunk", [3, 7]) + def test_chunked_equals_eager(self, up, down, nchunk): + da = xd.testing.dummy(shape=(101, 5)) + taps = sp.firwin(21, 0.4 / max(up, down)) + eager = Polyphase(taps, up, down, "time")(da) + atom = Polyphase(taps, up, down, "time") + outs = [atom(chunk, chunk_dim="time") for chunk in xd.split(da, nchunk, "time")] + chunked = xd.concat(outs, "time") + np.testing.assert_allclose(chunked.values, eager.values, atol=1e-15) + assert chunked.coords.equals(eager.coords) + + def test_keeps_the_data_precision(self): + da = xd.testing.dummy(shape=(101, 5), dtype="float32") + taps = sp.firwin(21, 0.2) + assert taps.dtype == np.float64 + result = Polyphase(taps, 1, 2, "time")(da) + assert result.dtype == np.float32 + expected = Polyphase(taps, 1, 2, "time")(da.copy(data=da.values.astype(float))) + np.testing.assert_allclose(result.values, expected.values, rtol=1e-6) + + def test_rate_the_coordinate_cannot_represent_exactly(self): + # 100 Hz resampled by 3/10 is 10/3 nanoseconds per output sample: the + # truncated step must be declared as jitter, not silently drift. + da = xd.testing.dummy(shape=(101, 5)) + taps = sp.firwin(31, 0.4 / 10) + result = Polyphase(3 * taps, 3, 10, "time")(da) + coord = result.coords["time"] + assert coord.isregular() + assert coord.tolerance > np.timedelta64(0, "ns") + + def test_too_few_taps(self): + da = xd.testing.dummy(shape=(20, 5)) + atom = Polyphase(sp.firwin(3, 0.4), 5, 1, "time") + with pytest.raises(ValueError, match="at least 5 taps"): + atom(da) + + def test_empty_input_emits_nothing(self): + da = xd.testing.dummy(shape=(101, 5)) + atom = Polyphase(sp.firwin(21, 0.2), 1, 2, "time") + assert atom(da.isel(time=slice(0, 0))) == xd.DataCollection([]) + + def test_upsample_single_sample(self): + # a one-sample chunk has a single tie: the upsampled block still + # spans `factor` samples, which takes a second tie to say. + da = xd.testing.dummy(shape=(101, 5)) + result = UpSample(3, dim="time")(da.isel(time=slice(0, 1))) + assert result.sizes["time"] == 3 class TestAtomCoreMissingBranches: @@ -383,11 +422,13 @@ def test_atomized_two_atom_args_raises(self): xs.integrate(atom1, atom2) def test_atomized_sequential_input(self): + # Composition has value semantics: the input pipeline is not mutated. atom = xs.integrate(...) seq = Sequential([atom]) - initial_len = len(seq) - xs.integrate(seq) - assert len(seq) == initial_len + 1 + result = xs.integrate(seq) + assert isinstance(result, Sequential) + assert len(result) == 2 + assert len(seq) == 1 def test_set_state_nested_atom(self): from xdas.atoms.core import Atom, State @@ -507,7 +548,7 @@ def stream(sequence, nchunks=4): np.testing.assert_array_equal(first.values, second.values) -class TestMLPickerMissingBranches: +class TestLazyModule: def test_lazy_module_import_error(self): from xdas.atoms.ml import LazyModule @@ -515,9 +556,275 @@ def test_lazy_module_import_error(self): with pytest.raises(ImportError, match="is not installed by default"): _ = mod.something - def test_mlpicker_invalid_component_strategy(self): - import seisbench.models as sbm - model = sbm.PhaseNet.from_pretrained("geofon") - with pytest.raises(ValueError, match="component_strategy must be one of"): - MLPicker(model, dim="time", component_strategy="invalid") +class TestCompose: + def test_rshift_atoms(self): + pipeline = xs.detrend(...) >> xs.integrate(...) + assert isinstance(pipeline, Sequential) + assert len(pipeline) == 2 + + def test_rshift_value_semantics(self): + head = xs.detrend(...) + pipeline = head >> xs.integrate(...) + longer = pipeline >> np.square + assert isinstance(head, Partial) + assert len(pipeline) == 2 + assert len(longer) == 3 + + def test_irshift(self): + pipeline = xs.detrend(...) + pipeline >>= xs.integrate(...) + pipeline >>= np.square + assert isinstance(pipeline, Sequential) + assert len(pipeline) == 3 + + def test_rshift_callable_wraps(self): + pipeline = xs.detrend(...) >> np.square + assert isinstance(pipeline[-1], Partial) + assert pipeline[-1].func is np.square + + def test_rrshift_callable_prepends(self): + pipeline = np.square >> xs.detrend(...) + assert isinstance(pipeline, Sequential) + assert pipeline[0].func is np.square + + def test_rrshift_applies_to_data(self): + da = xd.testing.dummy() + result = da >> Partial(np.square) + assert np.allclose(result.values, np.square(da.values)) + + def test_rrshift_applies_pipeline_to_data(self): + da = xd.testing.dummy() + pipeline = Partial(np.abs) >> Partial(np.square) + result = da >> pipeline + assert np.allclose(result.values, np.square(np.abs(da.values))) + + def test_named_sequential_kept_nested_on_right(self): + named = Sequential([Partial(np.square)], name="named") + pipeline = xs.detrend(...) >> named + assert len(pipeline) == 2 + assert pipeline[1] is named + + def test_named_sequential_extended_keeps_name(self): + named = Sequential([Partial(np.square)], name="named") + pipeline = named >> Partial(np.abs) + assert pipeline.name == "named" + assert len(pipeline) == 2 + assert len(named) == 1 + + def test_unnamed_sequentials_flatten(self): + left = Partial(np.abs) >> Partial(np.square) + right = Partial(np.sqrt) >> Partial(np.abs) + pipeline = left >> right + assert len(pipeline) == 4 + + def test_rshift_with_data_on_right_raises(self): + with pytest.raises(TypeError): + xs.detrend(...) >> 1.0 + + +class TestTracing: + def test_ufunc_appends(self): + atom = xs.detrend(...) + traced = np.square(atom) + assert isinstance(traced, Sequential) + assert len(traced) == 2 + + def test_expression_matches_eager(self): + da = xd.testing.dummy() + atom = xs.detrend(...) + traced = 20 * np.log10(np.abs(atom) + 1e-12) + expected = 20 * np.log10(np.abs(xs.detrend(da)) + 1e-12) + assert np.allclose(traced(da).values, expected.values) + + def test_reflected_scalar(self): + da = xd.testing.dummy() + traced = 2.0 * xs.detrend(...) + assert np.allclose(traced(da).values, 2.0 * xs.detrend(da).values) + + def test_untraceable_attribute_raises(self): + atom = xs.detrend(...) + with pytest.raises(AttributeError): + _ = atom.values + + def test_fan_in_raises(self): + atom1 = xs.detrend(...) + atom2 = xs.detrend(...) + with pytest.raises(TypeError, match="fan-in"): + np.add(atom1, atom2) + + def test_same_atom_twice_raises(self): + atom = xs.detrend(...) + with pytest.raises(TypeError, match="fan-in"): + np.add(atom, atom) + + def test_equality_is_identity(self): + atom1 = xs.detrend(...) + atom2 = xs.detrend(...) + alias = atom1 + assert atom1 == alias + assert atom1 != atom2 + assert len({atom1, atom2}) == 2 + + def test_inplace_operator_traces_out_of_place(self): + da = xd.testing.dummy() + atom = xs.detrend(...) + atom *= 2.0 + assert isinstance(atom, Sequential) + assert np.allclose(atom(da).values, 2.0 * xs.detrend(da).values) + + def test_out_to_another_atom_raises(self): + atom1 = xs.detrend(...) + atom2 = xs.detrend(...) + with pytest.raises(TypeError): + np.multiply(atom1, 2.0, out=atom2) + + def test_out_to_a_data_array_raises(self): + # writing a traced expression into an array cannot be honoured: there + # is nothing to write until the pipeline runs. + atom = xs.detrend(...) + with pytest.raises(TypeError): + np.multiply(atom, 2.0, out=xd.testing.dummy()) + + def test_non_call_ufunc_method_raises(self): + atom = xs.detrend(...) + with pytest.raises(TypeError): + np.add.reduce(atom) + + def test_right_shift_as_data_traces(self): + # np.right_shift with the atom on the *left* is an ordinary traced + # ufunc, not the `da >> atom` application path. + traced = np.right_shift(xs.detrend(...) >> Partial(np.abs), 1) + assert isinstance(traced, Sequential) + + +class TestWholeRecordRefusal: + """Whole-record functions carry their own guard at the definition site.""" + + @staticmethod + def whole_record_atom(*args, **kwargs): + @atomized + @_whole_record() + def whole_record(da, dim="time"): + return da + + return whole_record(*args, **kwargs) + + def test_chunked_along_dim_raises(self): + da = xd.testing.dummy() + atom = self.whole_record_atom(...) + with pytest.raises(ValueError, match="whole record"): + atom(da, chunk_dim="time") + + def test_chunked_along_other_dim_passes(self): + da = xd.testing.dummy() + atom = self.whole_record_atom(...) + atom(da, chunk_dim="distance") + + def test_unchunked_passes(self): + da = xd.testing.dummy() + atom = self.whole_record_atom(...) + atom(da) + + def test_positional_dim_resolved(self): + da = xd.testing.dummy() + atom = self.whole_record_atom(..., "distance") + atom(da, chunk_dim="time") + with pytest.raises(ValueError, match="whole record"): + atom(da, chunk_dim="distance") + + def test_alias_dim_resolved_against_the_data(self): + da = xd.testing.dummy() # dims ("time", "distance") + atom = self.whole_record_atom(..., "last") + atom(da, chunk_dim="time") + with pytest.raises(ValueError, match="whole record"): + atom(da, chunk_dim="distance") + + def test_streaming_class_unaffected(self): + da = xd.testing.dummy() + chunks = xd.split(da, 3, "time") + atom = IIRFilter(4, 10.0, "lowpass", dim="time") + for chunk in chunks: + atom(chunk, chunk_dim="time") + + +class TestFresh: + def test_fresh_is_stateless_and_config_shared(self): + sos = sp.iirfilter(4, 0.1, btype="lowpass", output="sos") + atom = Partial(xs.sosfilt, sos, ..., dim="time", zi=...) + da = xd.testing.dummy() + atom(da, chunk_dim="time") + assert atom.initialized + clone = atom.fresh() + assert not clone.initialized + assert clone.func is atom.func + assert atom.initialized # the original is untouched + assert clone(da).equals(Partial(xs.sosfilt, sos, ..., dim="time", zi=...)(da)) + + def test_fresh_recurses_into_sequences(self): + sos = sp.iirfilter(4, 0.1, btype="lowpass", output="sos") + seq = Sequential( + [Partial(xs.sosfilt, sos, ..., dim="time", zi=...), Partial(np.square)], + name="energy", + ) + da = xd.testing.dummy() + seq(da, chunk_dim="time") + clone = seq.fresh() + assert not clone.initialized + assert clone.name == "energy" + assert len(clone) == len(seq) + assert clone[0] is not seq[0] + assert clone[0].func is seq[0].func + + def test_fresh_keeps_the_type_of_a_self_assembling_sequence(self): + # a pipeline that builds its own stages (a `Picker`) cannot be rebuilt + # by calling its constructor with them. + class Assembled(Sequential): + def __init__(self, factor): + super().__init__([Partial(np.square)] * factor, name="assembled") + self.factor = factor + + seq = Assembled(2) + clone = seq.fresh() + assert isinstance(clone, Assembled) + assert clone.factor == 2 + assert len(clone) == 2 + assert clone[0] is not seq[0] + + +class TestFreshNested: + def test_fresh_recurses_into_nested_class_atoms(self): + da = xd.testing.dummy() + atom = xd.atoms.Filter((1.0, 10.0)) + atom(da, chunk_dim="time") + clone = atom.fresh() + assert not clone.initialized + assert clone.filter is not atom.filter + assert atom.initialized + + +class TestRefusalHelper: + def test_alias_without_data_is_conservative(self): + atom = Partial(np.square) + atom._refuse_chunked_along("distance", "time", None) # distinct: passes + with pytest.raises(ValueError, match="whole record"): + atom._refuse_chunked_along("last", "time", None) + + def test_no_chunking_passes(self): + Partial(np.square)._refuse_chunked_along("time", None, None) + + def test_kernel_dict_checks_its_keys(self): + da = xd.testing.dummy() + atom = Partial(np.square) + atom._refuse_chunked_along({"distance": 5}, "time", da) + with pytest.raises(ValueError, match="whole record"): + atom._refuse_chunked_along({"time": 5}, "time", da) + + +class TestInitializedRecurses: + def test_a_fresh_nested_atom_reports_uninitialized(self): + sos = sp.iirfilter(4, 0.1, btype="lowpass", output="sos") + seq = Sequential([Partial(xs.sosfilt, sos, ..., dim="time", zi=...)]) + assert not seq.initialized + seq(xd.testing.dummy(), chunk_dim="time") + assert seq.initialized diff --git a/tests/test_atoms_detect.py b/tests/test_atoms_detect.py new file mode 100644 index 00000000..5ad2464a --- /dev/null +++ b/tests/test_atoms_detect.py @@ -0,0 +1,457 @@ +""" +Tests for the detection atoms of `xdas.atoms.detect`. + +`Trigger` joined the public layer with the rest of the task vocabulary, so it +now defaults to `dim="time"` and has a lowercase twin, `xd.trigger`. +""" + +import numpy as np +import pandas as pd +import pytest + +import xdas as xd +from xdas.atoms import Trigger + + +def generate(): + return xd.DataArray( + data=[[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2]], + coords={ + "space": [0.0], + "time": { + "tie_indices": [0, 9], + "tie_values": [0.0, 9.0], + "sampling_interval": 1.0, + }, + }, + ) + + +expected = pd.DataFrame({"space": [0.0, 0.0], "time": [2.0, 7.0], "value": [0.9, 0.7]}) + + +def test_trigger(): + cft = generate() + + # test monolithic processing + picks = Trigger(thresh=0.5, dim="time")(cft) + assert picks.equals(expected) + + # test chunked processing + atom = Trigger(thresh=0.5, dim="time") + chunks = xd.split(cft, 3, dim="time") + result = [] + for chunk in chunks: + picks = atom(chunk, chunk_dim="time") + result.append(picks) + result = pd.concat(result, ignore_index=True) + assert result.equals(expected) + + +class TestTwin: + def test_is_a_function_and_the_module_still_imports(self): + # `xdas.trigger` stays importable as the compat module, but the + # attribute is the lowercase twin and importing the module does not + # shadow it. + assert callable(xd.trigger) + from xdas.trigger import Trigger as compat + + assert compat is Trigger + import xdas.trigger # noqa: F401 + + assert callable(xd.trigger) + + def test_eager_call(self): + assert xd.trigger(generate(), thresh=0.5).equals(expected) + + def test_seed_returns_the_atom(self): + atom = xd.trigger(..., thresh=0.5) + assert isinstance(atom, Trigger) + assert atom(generate()).equals(expected) + + def test_extends_a_pipeline(self): + pipeline = xd.taper(...) >> xd.trigger(..., thresh=0.5) + assert isinstance(pipeline(generate()), pd.DataFrame) + + def test_dim_defaults_to_time(self): + assert Trigger(thresh=0.5).dim == "time" + assert xd.trigger(generate(), thresh=0.5).equals(expected) + + +class TestTriggerCoords: + def generate(self): + cft = xd.DataArray( + data=[ + [0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2], + [0.0, 0.0, 0.1, 0.1, 0.0, 0.0, 0.8, 0.9, 0.1, 0.0], + ], + coords={ + "space": [0.0, 100.0], + "time": { + "tie_indices": [0, 9], + "tie_values": [0.0, 9.0], + "sampling_interval": 1.0, + }, + }, + ) + return cft.assign_coords(station=("space", ["ST01", "ST02"])) + + def test_annotates_with_a_non_dimensional_coordinate(self): + cft = self.generate() + picks = Trigger(thresh=0.5, dim="time", coords=["time", "station"])(cft) + expected = pd.DataFrame( + { + "time": [2.0, 7.0, 7.0], + "station": ["ST01", "ST01", "ST02"], + "value": [0.9, 0.7, 0.9], + } + ) + assert picks.equals(expected) + + def test_defaults_to_the_dimension_coordinates(self): + cft = self.generate() + picks = Trigger(thresh=0.5, dim="time")(cft) + assert list(picks.columns) == ["space", "time", "value"] + + def test_selects_and_orders_the_requested_columns(self): + cft = self.generate() + picks = Trigger(thresh=0.5, dim="time", coords=["station", "space", "time"])( + cft + ) + assert list(picks.columns) == ["station", "space", "time", "value"] + + def test_chunked_annotation_matches_monolithic(self): + cft = self.generate() + coords = ["time", "station"] + expected = Trigger(thresh=0.5, dim="time", coords=coords)(cft) + atom = Trigger(thresh=0.5, dim="time", coords=coords) + picks = [atom(chunk, chunk_dim="time") for chunk in xd.split(cft, 3, "time")] + result = pd.concat(picks, ignore_index=True) + assert result.sort_values(coords, ignore_index=True).equals( + expected.sort_values(coords, ignore_index=True) + ) + + def test_annotates_with_a_coordinate_along_the_picked_dimension(self): + # A label riding on the picked dimension is indexed by absolute sample + # number, so it is read off the labels kept for the whole run rather + # than off the chunk in hand: a pick found in a chunk already gone + # still names the right sample. + cft = self.generate().assign_coords( + sample=("time", [f"s{n}" for n in range(10)]) + ) + coords = ["time", "sample", "station"] + expected = Trigger(thresh=0.5, dim="time", coords=coords)(cft) + assert list(expected["sample"]) == ["s2", "s7", "s7"] + atom = Trigger(thresh=0.5, dim="time", coords=coords) + picks = [atom(chunk, chunk_dim="time") for chunk in xd.split(cft, 3, "time")] + result = pd.concat(picks, ignore_index=True) + assert result.sort_values(coords, ignore_index=True).equals( + expected.sort_values(coords, ignore_index=True) + ) + + def test_unknown_coordinate_raises(self): + cft = self.generate() + with pytest.raises(KeyError, match="not a coordinate"): + Trigger(thresh=0.5, dim="time", coords=["time", "elevation"])(cft) + + +def test_trigger_1d(): + """1D input (no spatial dimension) covers the coords=() branch in _call_numeric.""" + cft = xd.DataArray( + data=[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2], + coords={ + "time": { + "tie_indices": [0, 9], + "tie_values": [0.0, 9.0], + "sampling_interval": 1.0, + }, + }, + ) + picks = Trigger(thresh=0.5, dim="time")(cft) + assert len(picks) == 2 + assert list(picks["time"]) == [2.0, 7.0] + + +class TestThresholdMapping: + """W6: thresholds keyed on the `phase` labels rather than on their position.""" + + def generate(self, labels=("N", "P", "S")): + data = { + "N": [0.0, 0.9, 0.0, 0.0, 0.0], + "P": [0.0, 0.0, 0.4, 0.0, 0.0], + "S": [0.0, 0.0, 0.0, 0.8, 0.0], + } + return xd.DataArray( + data=[data[label] for label in labels], + coords={ + "phase": list(labels), + "time": { + "tie_indices": [0, 4], + "tie_values": [0.0, 4.0], + "sampling_interval": 1.0, + }, + }, + ) + + def test_one_threshold_per_label(self): + picks = Trigger(thresh={"P": 0.3, "S": 0.5})(self.generate()) + assert list(picks["phase"]) == ["P", "S"] + assert list(picks["time"]) == [2.0, 3.0] + assert list(picks["value"]) == [0.4, 0.8] + + def test_a_lane_below_its_own_threshold_does_not_trigger(self): + picks = Trigger(thresh={"P": 0.5, "S": 0.5})(self.generate()) + assert list(picks["phase"]) == ["S"] + + def test_unlisted_labels_never_trigger(self): + # `N` is the loudest lane of the three and carries no entry. + picks = Trigger(thresh={"P": 0.3, "S": 0.5})(self.generate()) + assert "N" not in set(picks["phase"]) + + def test_the_label_order_is_irrelevant(self): + # The order of a model's labels is a property of its weight set: the + # same mapping must give the same picks whatever order they come in. + thresh = {"P": 0.3, "S": 0.5} + expected = Trigger(thresh=thresh)(self.generate()) + flipped = Trigger(thresh=thresh)(self.generate(("S", "P", "N"))) + columns = ["phase", "time", "value"] + assert flipped.sort_values(columns, ignore_index=True)[columns].equals( + expected.sort_values(columns, ignore_index=True)[columns] + ) + + def test_a_scalar_still_applies_to_every_lane(self): + picks = Trigger(thresh=0.3)(self.generate()) + assert list(picks["phase"]) == ["N", "P", "S"] + + def test_bytes_labels_are_decoded(self): + cft = self.generate().assign_coords(phase=np.array([b"N", b"P", b"S"])) + picks = Trigger(thresh={"P": 0.3, "S": 0.5})(cft) + assert list(picks["time"]) == [2.0, 3.0] + + def test_numeric_labels_are_keyed_by_their_string_form(self): + # A weight set declaring no `labels` gets positional ones, so the + # `phase` coordinate is integer-valued and the mapping still keys on it. + cft = self.generate().assign_coords(phase=[0, 1, 2]) + picks = Trigger(thresh={1: 0.3, 2: 0.5})(cft) + assert list(picks["phase"]) == [1, 2] + + def test_chunked_matches_monolithic(self): + cft = self.generate() + thresh = {"P": 0.3, "S": 0.5} + expected = Trigger(thresh=thresh)(cft) + atom = Trigger(thresh=thresh) + picks = [atom(chunk, chunk_dim="time") for chunk in xd.split(cft, 3, "time")] + picks += atom.flush() + result = pd.concat(picks, ignore_index=True) + assert result.equals(expected) + + def test_without_a_phase_coordinate_raises(self): + with pytest.raises(ValueError, match="'phase' coordinate, which the data"): + Trigger(thresh={"P": 0.3})(generate()) + + def test_along_the_phase_dimension_raises(self): + with pytest.raises(ValueError, match="it is the dimension"): + Trigger(thresh={"P": 0.3}, dim="phase")(self.generate()) + + def test_an_unknown_label_raises(self): + with pytest.raises(KeyError, match=r"\['Pg'\]"): + Trigger(thresh={"Pg": 0.3})(self.generate()) + + +class TestScalarCoords: + """W6: 0-d coordinates become constant columns.""" + + def generate(self): + return generate().assign_coords(station="ST01", depth=1000.0) + + def test_named_explicitly(self): + picks = Trigger(thresh=0.5, coords=["time", "station"])(self.generate()) + assert list(picks.columns) == ["time", "station", "value"] + assert list(picks["station"]) == ["ST01", "ST01"] + + def test_auto_leads_with_them(self): + # identity first, measurement last: the scalar coordinates lead, then + # the other dimension coordinates, then the picked dimension. The tree + # path of a collection walk takes the same leading position, so a pick + # table reads the same whichever source its identity came from. + picks = Trigger(thresh=0.5, coords="auto")(self.generate()) + assert list(picks.columns) == ["station", "depth", "space", "time", "value"] + assert list(picks["depth"]) == [1000.0, 1000.0] + + def test_auto_does_not_depend_on_the_input_dimension_order(self): + cft = self.generate() + expected = Trigger(thresh=0.5)(cft) + transposed = Trigger(thresh=0.5)(cft.transpose("time", "space")) + assert list(transposed.columns) == list(expected.columns) + + def test_auto_is_the_default(self): + expected = Trigger(thresh=0.5, coords="auto")(self.generate()) + assert Trigger(thresh=0.5)(self.generate()).equals(expected) + + def test_none_keeps_the_dimension_coordinates_only(self): + picks = Trigger(thresh=0.5, coords=None)(self.generate()) + assert list(picks.columns) == ["space", "time", "value"] + + def test_a_constant_column_survives_an_empty_chunk(self): + cft = self.generate() + atom = Trigger(thresh=0.5, coords=["station"]) + picks = [atom(chunk, chunk_dim="time") for chunk in xd.split(cft, 5, "time")] + assert picks[0].empty + result = pd.concat(picks + atom.flush(), ignore_index=True) + assert list(result["station"]) == ["ST01", "ST01"] + + def test_an_unknown_coords_string_raises(self): + with pytest.raises(ValueError, match="must be 'auto', None or a sequence"): + Trigger(thresh=0.5, coords="all") + + +class TestFlush: + """W6: a trigger still open at the end of a run is closed, not lost.""" + + def generate(self): + return xd.DataArray( + data=[[0.0, 0.1, 0.9, 0.8, 0.7]], + coords={ + "space": [0.0], + "time": { + "tie_indices": [0, 4], + "tie_values": [0.0, 4.0], + "sampling_interval": 1.0, + }, + }, + ) + + def test_the_eager_call_closes_the_run(self): + picks = Trigger(thresh=0.5)(self.generate()) + assert isinstance(picks, pd.DataFrame) + assert list(picks["time"]) == [2.0] + assert list(picks["value"]) == [0.9] + + def test_chunk_invariance(self): + cft = self.generate() + expected = Trigger(thresh=0.5)(cft) + atom = Trigger(thresh=0.5) + picks = [atom(chunk, chunk_dim="time") for chunk in xd.split(cft, 3, "time")] + picks += atom.flush() + assert pd.concat(picks, ignore_index=True).equals(expected) + + def test_nothing_open_emits_nothing(self): + atom = Trigger(thresh=0.5) + atom(generate(), chunk_dim="time") + assert atom.flush() == [] + + def test_flushing_twice_emits_once(self): + atom = Trigger(thresh=0.5) + atom(self.generate(), chunk_dim="time") + assert len(atom.flush()) == 1 + assert atom.flush() == [] + + def test_before_initialization_emits_nothing(self): + assert Trigger(thresh=0.5).flush() == [] + + def test_each_run_of_a_gappy_record_is_closed(self): + tail = xd.DataArray( + data=[[0.0, 0.1, 0.7, 0.6, 0.6]], + coords={ + "space": [0.0], + "time": { + "tie_indices": [0, 4], + "tie_values": [10.0, 14.0], + "sampling_interval": 1.0, + }, + }, + ) + picks = Trigger(thresh=0.5)(xd.concat([self.generate(), tail], "time")) + assert isinstance(picks, pd.DataFrame) + assert list(picks["time"]) == [2.0, 12.0] + + def test_iter_chunks_flushes(self): + chunks = list(xd.split(self.generate(), 3, "time")) + atom = Trigger(thresh=0.5) + picks = pd.concat(atom.iter_chunks(chunks, "time"), ignore_index=True) + assert list(picks["time"]) == [2.0] + + def test_one_lane_open_among_several(self): + cft = xd.DataArray( + data=[[0.0, 0.9, 0.0], [0.0, 0.9, 0.8]], + coords={ + "space": [0.0, 100.0], + "time": { + "tie_indices": [0, 2], + "tie_values": [0.0, 2.0], + "sampling_interval": 1.0, + }, + }, + ) + picks = Trigger(thresh=0.5)(cft) + assert list(picks["space"]) == [0.0, 100.0] + assert list(picks["time"]) == [1.0, 1.0] + + +class TestChunkSemantics: + """ + `Trigger` carries its open triggers across chunks along `dim`, elementwise + across. + + Chunking along a dimension the atom does not pick along must change + nothing. It used to be false: `annotations` froze the *first* chunk's + lane coordinates, and `offset`, `coord`, `status`, `index` and `value` + accumulated on every call, so a run chunked along ``distance`` labelled + every later chunk's picks with the first chunk's lanes and the wrong + times. + """ + + def cft(self, nlanes=8, nsamples=40): + rng = np.random.default_rng(42) + template = xd.testing.dummy( + dims=("time", "distance"), + shape=(nsamples, nlanes), + datetime=False, + step=1.0, + ) + values = rng.random((nsamples, nlanes)) + return xd.DataArray(values, dict(template.coords), ("time", "distance")) + + @pytest.mark.parametrize("size", [7, 13, 40]) + def test_chunking_along_the_picked_dimension_is_invariant(self, size): + xd.testing.assert_chunk_invariant( + Trigger(thresh=0.8), self.cft(), {"time": size} + ) + + def test_a_single_lane_is_invariant(self): + xd.testing.assert_chunk_invariant( + Trigger(thresh=0.8), self.cft(nlanes=1), {"time": 7} + ) + + @pytest.mark.parametrize("size", [1, 3, 8]) + def test_chunking_along_another_dimension_is_invariant(self, size): + # Regression: the lanes of a later chunk used to be annotated with the + # first chunk's `distance` values, on a time axis that kept growing. + xd.testing.assert_chunk_invariant( + Trigger(thresh=0.8), self.cft(), {"distance": size} + ) + + def test_each_lane_keeps_its_own_identity_and_time_base(self): + # Two lanes picking at different samples: the second chunk used to be + # labelled with the first chunk's `distance` value, and its index to be + # shifted by the first chunk's length. + cft = xd.DataArray( + data=[[0.0, 0.9], [0.9, 0.0], [0.0, 0.0], [0.0, 0.0]], + coords={ + "time": { + "tie_indices": [0, 3], + "tie_values": [0.0, 3.0], + "sampling_interval": 1.0, + }, + "distance": [0.0, 100.0], + }, + ) + eager = Trigger(thresh=0.5)(cft) + assert sorted(zip(eager["distance"], eager["time"])) == [ + (0.0, 1.0), + (100.0, 0.0), + ] + chunked = Trigger(thresh=0.5).process(cft, chunks={"distance": 1}) + assert sorted(zip(chunked["distance"], chunked["time"])) == sorted( + zip(eager["distance"], eager["time"]) + ) diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py new file mode 100644 index 00000000..bcb21178 --- /dev/null +++ b/tests/test_atoms_ml.py @@ -0,0 +1,2041 @@ +""" +Tests for the machine-learning atom and for the fake model the ML suite runs on. + +READ THIS BEFORE CHANGING ``xdas/atoms/ml.py`` +---------------------------------------------- +The two classes at the bottom of this module are a *characterization pin* of +``MLPicker`` as it behaves today, taken deliberately before the ``Annotate`` / +``Picker`` rework (plan §7, W2: "the values are unchanged against the current +``MLPicker`` on the same input"). The rework changes some of that behaviour on +purpose, so the pin is split in two, and the split is the point: + +``TestMLPickerValuePin`` + What the rework must **not** change: the numbers. Every assertion here + reaches the array through dimension *names* — ``transpose("time", lane, + "phase")`` before comparing, ``isel`` with a name, never a positional axis + — so it holds whatever order the output dimensions come back in. If one of + these fails after the rework, the characteristic function moved and that is + a regression. + +``TestMLPickerBehaviourTheReworkChanges`` + What the rework is **expected** to change, isolated so it can be updated on + its own: + + - ``test_output_is_laid_out_sample_last`` — the old ``MLPicker`` forced the + sample dimension first (``chunk.transpose(self.dim, ...)``) whatever went + in. ``Annotate`` keeps the input's order among the other dimensions but + lays the output out sample-last, so ``("distance", "time")`` in gives + ``("distance", "phase", "time")`` out. Do not touch + ``TestMLPickerValuePin``, which reaches the array by name and holds + either way. + - ``test_output_stops_at_the_last_complete_window`` — today the trailing + samples that no full window covers are dropped. W1 emits the + end-aligned final window at ``flush()``, so the output grows to span the + input. ``TestMLPickerValuePin`` only ever looks at the leading samples + the two agree on, so it survives that too. +""" + +import pickle + +import numpy as np +import numpy.testing as npt +import obspy +import pandas as pd +import pytest +import torch +from seisbench.models import WaveformModel + +import xdas as xd +from tests.fakemodel import WEIGHT_SETS, FakeModel, fake_model +from xdas.atoms import Annotate, Filter, MLPicker, Picker, Resample, Sequential +from xdas.atoms.ml import _ChannelFilter, _model_filter, _model_thresholds + +#: A short, exactly representable signal: two lanes that are neither equal nor +#: proportional, so a bug that mixes lanes shows up in the golden. +PIN_SIGNAL = np.array( + [3, -1, 4, -1, 5, -9, 2, -6, 5, 3, -5, 8, 9, -7, 9, 3, -2, 3, -8, 4, 6, 2, -6, 4], + float, +) + +#: Number of samples the current implementation emits for :func:`pin_array`: +#: ``range(0, 24 - 8, 4)`` windows of ``step = in_samples // 2 = 4`` samples. +PIN_SAMPLES = 16 + + +def pin_array(dims, n_lanes=2): + """Build the pin input, laid out over *dims* (which must contain ``time``).""" + values = np.stack( + [np.roll(PIN_SIGNAL, 5 * j) * (j + 1) for j in range(n_lanes)], axis=-1 + ) + lane = next(dim for dim in dims if dim != "time") + template = xd.testing.dummy( + dims=("time", lane), shape=(len(PIN_SIGNAL), n_lanes), ctype="interpolated" + ) + coords = dict(template.coords) + if lane == "id": + coords["id"] = np.array([f"A{index:02d}" for index in range(n_lanes)]) + da = xd.DataArray(values, coords, ("time", lane)) + return da if dims == ("time", lane) else da.transpose(*dims) + + +# --------------------------------------------------------------------------- +# The fake model (plan §3 as an executable contract) +# --------------------------------------------------------------------------- + + +class TestFakeModelSpansTheWeightSetAxes: + """Every property §3 lists as per-weights is actually varied by the presets.""" + + @staticmethod + def spread(key, default=None): + return { + weights.get(key, default) + if not isinstance(weights.get(key, default), dict) + else "dict" + for weights in WEIGHT_SETS.values() + } + + def test_component_order(self): + assert self.spread("component_order") == {"ENZ", "ZNE", "Z12H"} + + def test_in_channels(self): + assert self.spread("in_channels") == {3, 4} + + def test_label_order_flips(self): + assert self.spread("labels") == {"NPS", "PSN"} + + def test_sampling_rate(self): + assert self.spread("sampling_rate") == {50, 100} + + def test_blinding_present_and_absent(self): + declared = {"blinding" in w["default_args"] for w in WEIGHT_SETS.values()} + assert declared == {True, False} + + def test_overlap_present_and_absent(self): + declared = {"overlap" in w["default_args"] for w in WEIGHT_SETS.values()} + assert declared == {True, False} + + def test_per_phase_thresholds_present_and_absent(self): + declared = { + any(key.endswith("_threshold") for key in w["default_args"]) + for w in WEIGHT_SETS.values() + } + assert declared == {True, False} + + def test_filters_absent_flat_and_per_channel(self): + kinds = set() + for weights in WEIGHT_SETS.values(): + args = weights.get("filter_args") + kinds.add("absent" if args is None else type(args).__name__) + assert kinds == {"absent", "tuple", "dict"} + obs = fake_model("obs") + assert obs.filter_args == {"??H": ["highpass"]} + assert obs.filter_kwargs == {"??H": {"freq": 0.5}} + + +class TestFakeModel: + """The contract ``MLPicker`` and its successors drive the model through.""" + + def test_is_a_waveform_model(self): + model = fake_model() + assert isinstance(model, WaveformModel) + assert isinstance(model, torch.nn.Module) + assert not model.training # the factory returns it in eval mode + assert model.to("cpu") is model + + @pytest.mark.parametrize("name", sorted(WEIGHT_SETS)) + def test_every_preset_exposes_what_the_atom_reads(self, name): + model = fake_model(name) + assert isinstance(model.in_samples, int) + assert model.classes == len(model.labels) + assert len(model.component_order) == model.in_channels + assert set(model.default_args) == set(WEIGHT_SETS[name]["default_args"]) + + def test_overrides_win_over_the_preset(self): + model = fake_model("obs", sampling_rate=50, in_samples=16) + assert model.component_order == "Z12H" + assert model.sampling_rate == 50 + assert model.in_samples == 16 + + def test_forward_is_a_fixed_function_of_the_input(self): + model = fake_model() # ZNE, 3 channels, PSN + batch = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4) + out = model(batch) + assert out.shape == (2, model.classes, 4) + # y[b, k, t] = k + sum_c (c + 1) * x[b, c, t] + pooled = (batch * torch.tensor([1.0, 2.0, 3.0]).reshape(1, -1, 1)).sum(dim=-2) + for k in range(model.classes): + assert torch.equal(out[:, k, :], pooled + k) + + def test_annotate_batch_pre_reads_the_argdict(self): + model = fake_model() + batch = torch.tensor([[1.0, -2.0, 4.0]]) + plain = model.annotate_batch_pre(batch, {}) + scaled = model.annotate_batch_pre(batch, {"scale": 10.0}) + np.testing.assert_allclose(plain.numpy(), [[0.25, -0.5, 1.0]], rtol=1e-6) + np.testing.assert_allclose(scaled.numpy(), [[2.5, -5.0, 10.0]], rtol=1e-6) + assert model.seen_argdicts == [{}, {"scale": 10.0}] + + def test_annotate_batch_post_transposes_and_blinds(self): + model = fake_model() + batch = torch.arange(1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 5) + out = model.annotate_batch_post(batch.clone(), None, {"blinding": (1, 2)}) + assert out.shape == (1, 5, 3) # (batch, samples, classes) + assert torch.isnan(out[0, 0]).all() + assert torch.isnan(out[0, -2:]).all() + assert torch.equal(out[0, 1:3], batch[0, :, 1:3].T) + + def test_blinding_falls_back_to_the_class_default_when_undeclared(self): + # `obs` declares no blinding; the fallback is (0, 0), i.e. blind nothing. + model = fake_model("obs") + assert "blinding" not in model.default_args + batch = torch.ones(1, model.classes, 5) + out = model.annotate_batch_post(batch, None, {}) + assert not torch.isnan(out).any() + + def test_annotate_args_moves_the_fallback_per_instance(self): + model = fake_model(annotate_args={"blinding": (2, 0)}) + batch = torch.ones(1, model.classes, 5) + out = model.annotate_batch_post(batch, None, {}) + assert torch.isnan(out[0, :2]).all() + assert not torch.isnan(out[0, 2:]).any() + assert FakeModel._annotate_args["blinding"][1] == (0, 0) # class untouched + + def test_piggyback_pair_form(self): + model = fake_model(piggyback=3.0) + pair = model.annotate_batch_pre(torch.tensor([[1.0, 2.0]]), {}) + assert isinstance(pair, tuple) and len(pair) == 2 + piggyback = pair[1] + assert piggyback == 3.0 + out = model.annotate_batch_post( + torch.ones(1, model.classes, 2), piggyback, {"blinding": (0, 0)} + ) + assert torch.equal(out, torch.full((1, 2, model.classes), 3.0)) + assert model.seen_piggybacks == [3.0] + + @pytest.mark.parametrize("name", ["original", "diting", "geofon"]) + def test_the_atom_can_drive_every_preset_that_declares_blinding(self, name): + model = fake_model(name) + picker = MLPicker(model, "time", device="cpu") + result = picker(pin_array(("time", "distance"))) + assert result.sizes["phase"] == model.classes + assert list(result.coords["phase"].values) == list(model.labels) + + @pytest.mark.parametrize("name", ["obs", "volpick"]) + def test_undeclared_blinding_falls_back_instead_of_raising(self, name): + # Neither weight set declares `blinding`. Before W1 the atom read + # `default_args["blinding"]` directly and raised `KeyError` here (plan + # §3, consequence 3); now blinding is the model's own business, applied + # by `annotate_batch_post`, and its fallback is (0, 0) — blind nothing. + picker = Annotate(fake_model(name), "time", device="cpu") + result = picker(pin_array(("time", "distance"))) + assert not np.isnan(result.values).any() + + +# --------------------------------------------------------------------------- +# Characterization pin — see the module docstring +# --------------------------------------------------------------------------- + +GOLDEN_CLONE = np.array( + [[[ np.nan, np.nan, np.nan], + [ np.nan, np.nan, np.nan]], + [[-0.6666667, 0.3333333, 1.3333333], + [ 6. , 7. , 8. ]], + [[ 2.6666667, 3.6666667, 4.666667 ], + [ 2. , 3. , 4. ]], + [[-0.6666667, 0.3333333, 1.3333333], + [-6. , -5. , -4. ]], + [[ 3.3333335, 4.3333335, 5.3333335], + [ 4. , 5. , 6. ]], + [[-6. , -5. , -4. ], + [ 2.5 , 3.5 , 4.5 ]], + [[ 1.3333334, 2.3333335, 3.3333335], + [-0.8333334, 0.1666667, 1.1666666]], + [[-4. , -3. , -2. ], + [ 2.6666667, 3.6666667, 4.666667 ]], + [[ 3.3333335, 4.3333335, 5.3333335], + [-0.6666667, 0.3333333, 1.3333333]], + [[ 2. , 3. , 4. ], + [ 3.3333335, 4.3333335, 5.3333335]], + [[-3.3333335, -2.3333335, -1.3333335], + [-6. , -5. , -4. ]], + [[ 5.3333335, 6.3333335, 7.3333335], + [ 1.3333334, 2.3333335, 3.3333335]], + [[ 6. , 7. , 8. ], + [-4. , -3. , -2. ]], + [[-4.666667 , -3.666667 , -2.666667 ], + [ 3.3333335, 4.3333335, 5.3333335]], + [[ 6. , 7. , 8. ], + [ 2. , 3. , 4. ]], + [[ 2. , 3. , 4. ], + [-3.3333335, -2.3333335, -1.3333335]]] +) # fmt: skip + +GOLDEN_Z = np.array( + [[[ np.nan, np.nan, np.nan], + [ np.nan, np.nan, np.nan]], + [[-0.3333333, 0.6666666, 1.6666666], + [ 3. , 4. , 5. ]], + [[ 1.3333334, 2.3333335, 3.3333335], + [ 1. , 2. , 3. ]], + [[-0.3333333, 0.6666666, 1.6666666], + [-3. , -2. , -1. ]], + [[ 1.6666667, 2.6666667, 3.6666667], + [ 2. , 3. , 4. ]], + [[-3. , -2. , -1. ], + [ 1.25 , 2.25 , 3.25 ]], + [[ 0.6666667, 1.6666667, 2.6666667], + [-0.4166667, 0.5833333, 1.5833333]], + [[-2. , -1. , 0. ], + [ 1.3333334, 2.3333335, 3.3333335]], + [[ 1.6666667, 2.6666667, 3.6666667], + [-0.3333333, 0.6666666, 1.6666666]], + [[ 1. , 2. , 3. ], + [ 1.6666667, 2.6666667, 3.6666667]], + [[-1.6666667, -0.6666667, 0.3333333], + [-3. , -2. , -1. ]], + [[ 2.6666667, 3.6666667, 4.666667 ], + [ 0.6666667, 1.6666667, 2.6666667]], + [[ 3. , 4. , 5. ], + [-2. , -1. , 0. ]], + [[-2.3333335, -1.3333335, -0.3333335], + [ 1.6666667, 2.6666667, 3.6666667]], + [[ 3. , 4. , 5. ], + [ 1. , 2. , 3. ]], + [[ 1. , 2. , 3. ], + [-1.6666667, -0.6666667, 0.3333333]]] +) # fmt: skip + + +def canonical(result, lane): + """ + Return *result*'s pinned values as ``(time, lane, phase)``, by dimension name. + + Never index a dimension positionally here: the whole point of the pin is + that it outlives the rework's change of output dimension order. Only the + leading :data:`PIN_SAMPLES` samples are taken, so it also outlives W1's + end-aligned final window extending the output. + """ + result = result.transpose("time", lane, "phase") + return result.isel({"time": slice(0, PIN_SAMPLES)}).values + + +class TestMLPickerValuePin: + """The numbers today's ``MLPicker`` produces — these must not move.""" + + @pytest.mark.parametrize( + "dims", [("time", "distance"), ("distance", "time"), ("time", "id")] + ) + def test_das_layouts_agree_with_the_golden(self, dims): + lane = next(dim for dim in dims if dim != "time") + picker = MLPicker(fake_model("original"), "time", device="cpu") + result = picker(pin_array(dims)) + np.testing.assert_allclose(canonical(result, lane), GOLDEN_CLONE, rtol=1e-6) + + def test_named_component_strategy_agrees_with_the_golden(self): + picker = MLPicker( + fake_model("original"), "time", device="cpu", component_strategy="Z" + ) + result = picker(pin_array(("time", "distance"))) + np.testing.assert_allclose(canonical(result, "distance"), GOLDEN_Z, rtol=1e-6) + + @pytest.mark.parametrize("indices", [[7, 13], [4, 8, 12, 16, 20], [1]]) + def test_streamed_chunks_agree_with_the_golden(self, indices): + picker = MLPicker(fake_model("original"), "time", device="cpu") + chunks = xd.split(pin_array(("time", "distance")), indices, "time") + result = xd.concat(list(picker.iter_chunks(chunks)), "time") + np.testing.assert_allclose( + canonical(result, "distance"), GOLDEN_CLONE, rtol=1e-6 + ) + + def test_the_phase_coordinate_follows_the_model_label_order(self): + # `original` labels NPS, `geofon` labels PSN — plan §3's order flip. + for name, labels in ( + ("original", ["N", "P", "S"]), + ("geofon", ["P", "S", "N"]), + ): + picker = MLPicker(fake_model(name), "time", device="cpu") + result = picker(pin_array(("time", "distance"))) + assert list(result.coords["phase"].values) == labels + + @pytest.mark.parametrize("name", ["original", "geofon"]) + def test_the_phases_can_be_selected_by_label(self, name): + # the point of labelling the axis: `sel` must work whichever order the + # weight set declares, since `isel` positions differ between the two. + picker = MLPicker(fake_model(name), "time", device="cpu") + cft = picker(pin_array(("time", "distance"))) + result = cft.sel(phase=["P", "S"]) + assert list(result.coords["phase"].values) == ["P", "S"] + for phase in ("P", "S"): + np.testing.assert_array_equal( + result.sel(phase=phase).values, cft.sel(phase=phase).values + ) + with pytest.raises(KeyError): + cft.sel(phase="Q") + + @pytest.mark.parametrize( + "dims", [("time", "distance"), ("distance", "time"), ("time", "id")] + ) + def test_the_non_sample_coordinates_pass_through_untouched(self, dims): + lane = next(dim for dim in dims if dim != "time") + da = pin_array(dims) + picker = MLPicker(fake_model("original"), "time", device="cpu") + result = picker(da) + np.testing.assert_array_equal( + result.coords[lane].values, da.coords[lane].values + ) + + def test_the_sample_coordinate_starts_at_the_first_input_sample(self): + da = pin_array(("time", "distance")) + picker = MLPicker(fake_model("original"), "time", device="cpu") + result = picker(da) + np.testing.assert_array_equal( + result.coords["time"].values[:PIN_SAMPLES], + da.coords["time"].values[:PIN_SAMPLES], + ) + + +class TestMLPickerBehaviourTheReworkChanges: + """ + Two assertions W1+W2 is expected to rewrite. See the module docstring. + + Nothing else in this file encodes either fact, so updating these two is + enough; if a value assertion in ``TestMLPickerValuePin`` fails as well, the + characteristic function moved and that is a regression, not the rename. + """ + + @pytest.mark.parametrize( + "dims", [("time", "distance"), ("distance", "time"), ("time", "id")] + ) + def test_output_is_laid_out_sample_last(self, dims): + # W1+W2: the other dimensions keep their order, then `phase`, then the + # samples — so the characteristic function of one lane is contiguous. + lane = next(dim for dim in dims if dim != "time") + picker = MLPicker(fake_model("original"), "time", device="cpu") + result = picker(pin_array(dims)) + assert result.dims == (lane, "phase", "time") + assert result.values.flags["C_CONTIGUOUS"] + + def test_output_stops_at_the_last_complete_window(self): + # W1: `flush()` emits the end-aligned final window, so the output spans + # the whole input. + da = pin_array(("time", "distance")) + picker = MLPicker(fake_model("original"), "time", device="cpu") + result = picker(da) + assert da.sizes["time"] == len(PIN_SIGNAL) + assert result.sizes["time"] == len(PIN_SIGNAL) + + +# --------------------------------------------------------------------------- +# W1 + W2 — `Annotate` +# --------------------------------------------------------------------------- + +#: Spike spacing, equal to the fake model's window, so that every window sees +#: exactly one spike of every component and the peak normalisation is a no-op. +SPIKE_PERIOD = 8 + + +def annotate_model(name="original", **overrides): + """A preset with a plain 50 % overlap and no blinding, for legible values.""" + overrides.setdefault("default_args", {"overlap": 0.5}) + return fake_model(name, **overrides) + + +def spikes(ncomp, n=24): + """Unit spikes: component *c* fires at every sample congruent to *c*.""" + values = np.zeros((n, ncomp)) + for index in range(ncomp): + values[index::SPIKE_PERIOD, index] = 1.0 + return values + + +def component_array(labels, dims=("time", "channel"), sample_dim="time", n=24): + """ + A record whose component *c* is a spike train naming itself. + + Because the fake model weights input slot *k* by ``k + 1``, the value the + characteristic function takes at sample *c* says which slot component *c* + was permuted into — see :func:`slot_of`. + """ + comp_dim = next(dim for dim in dims if dim != sample_dim) + template = xd.testing.dummy( + dims=(sample_dim, comp_dim), shape=(n, len(labels)), ctype="interpolated" + ) + coords = dict(template.coords) + coords[comp_dim] = np.array(labels) + da = xd.DataArray(spikes(len(labels), n), coords, (sample_dim, comp_dim)) + return da.transpose(*dims) + + +def trace_array(n=24, sample_dim="time", index=0): + """A single trace, spiking at every sample congruent to *index*.""" + template = xd.testing.dummy( + dims=(sample_dim,), shape=(n,), ctype="interpolated", step=0.01 + ) + values = spikes(index + 1, n)[:, index] + return xd.DataArray(values, dict(template.coords), (sample_dim,)) + + +def slot_of(result, index, sample_dim="time", **isel): + """Return the model input slot the component spiking at *index* landed in.""" + value = result.isel({sample_dim: index, "phase": 0, **isel}).values + return round(float(value)) - 1 + + +class TestAnnotateRenaming: + """W1: `MLPicker` becomes `Annotate`, the old names warn until 0.4.""" + + def test_mlpicker_is_a_deprecated_alias(self): + with pytest.warns(DeprecationWarning, match="removed in 0.4"): + picker = MLPicker(annotate_model(), "time", device="cpu") + assert isinstance(picker, Annotate) + + def test_the_mlpicker_twin_warns_too(self): + with pytest.warns(DeprecationWarning, match="removed in 0.4"): + atom = xd.mlpicker(..., annotate_model(), "time", device="cpu") + assert isinstance(atom, Annotate) + + def test_annotate_has_an_eager_twin(self): + da = pin_array(("time", "distance")) + model = annotate_model() + expected = Annotate(model, "time", device="cpu")(da) + assert xd.annotate(da, model, "time", device="cpu").equals(expected) + atom = xd.annotate(..., model, "time", device="cpu") + assert isinstance(atom, Annotate) + + +class TestAnnotateReadsTheModelArgdict: + """W1: the parameters SeisBench reads from the model are no longer invented.""" + + def test_argdict_is_the_weight_sets_defaults_plus_the_call_kwargs(self): + model = fake_model("diting") + picker = Annotate(model, "time", device="cpu", scale=2.0) + assert picker.argdict == model.default_args | {"scale": 2.0} + assert "blinding" in picker.argdict # what the weight set declares + + def test_every_batch_is_preprocessed_with_that_argdict(self): + model = fake_model("original") + Annotate(model, "time", device="cpu", scale=2.0)(pin_array(("time", "id"))) + assert model.seen_argdicts + assert all( + argdict == {"overlap": 0.5, "blinding": (1, 1), "scale": 2.0} + for argdict in model.seen_argdicts + ) + + def test_the_argdict_reaches_the_values(self): + da = pin_array(("time", "distance")) + plain = Annotate(annotate_model(), "time", device="cpu")(da) + scaled = Annotate(annotate_model(), "time", device="cpu", scale=2.0)(da) + # phase 0 carries no class offset, so the gain shows undiluted + np.testing.assert_allclose( + scaled.isel(phase=0).values, 2.0 * plain.isel(phase=0).values, rtol=1e-6 + ) + + def test_preprocessing_sees_the_filled_three_dimensional_batch(self): + # Trap 1: today `annotate_batch_pre` is handed the 2-D staging buffer, + # before the component slots are filled. + shapes = [] + + class Recording(FakeModel): + def annotate_batch_pre(self, batch, argdict): + shapes.append(tuple(batch.shape)) + return super().annotate_batch_pre(batch, argdict) + + model = Recording( + component_order="ENZ", in_channels=3, default_args={"overlap": 0.5} + ).eval() + Annotate(model, "time", device="cpu")(pin_array(("time", "distance"))) + assert shapes and all(shape == (2, 3, 8) for shape in shapes) + + def test_the_piggyback_pair_is_plumbed_through_to_post(self): + da = pin_array(("time", "distance")) + model = annotate_model(piggyback=3.0) + result = Annotate(model, "time", device="cpu")(da) + plain = Annotate(annotate_model(), "time", device="cpu")(da) + assert model.seen_piggybacks == [3.0] * len(model.seen_piggybacks) + assert len(model.seen_piggybacks) > 1 + np.testing.assert_allclose(result.values, 3.0 * plain.values, rtol=1e-6) + + def test_blinding_is_left_to_the_model(self): + model = annotate_model(default_args={"overlap": 0.5, "blinding": (2, 0)}) + result = Annotate(model, "time", device="cpu")(pin_array(("time", "distance"))) + # one-sided blinding: only the first two samples are covered by nothing + assert np.isnan(result.isel(time=slice(0, 2)).values).all() + assert not np.isnan(result.isel(time=slice(2, None)).values).any() + + def test_blinding_of_zero_blinds_nothing(self): + model = annotate_model(default_args={"overlap": 0.5, "blinding": (0, 0)}) + result = Annotate(model, "time", device="cpu")(pin_array(("time", "distance"))) + assert not np.isnan(result.values).any() + + +class TestAnnotateWindowing: + """W1: the overlap, the stacking and the end-aligned final window.""" + + @pytest.mark.parametrize( + "overlap, noverlap", [(0, 0), (0.25, 2), (0.5, 4), (3, 3), (7, 7)] + ) + def test_overlap_is_read_the_seisbench_way(self, overlap, noverlap): + model = annotate_model(default_args={"overlap": overlap}) + picker = Annotate(model, "time", device="cpu") + assert picker.noverlap == noverlap + assert picker.step == 8 - noverlap + + def test_overlap_falls_back_to_the_models_own_default(self): + model = annotate_model(default_args={}, annotate_args={"overlap": 4}) + assert Annotate(model, "time", device="cpu").noverlap == 4 + + def test_an_overlap_of_a_whole_window_is_refused(self): + model = annotate_model(default_args={"overlap": 8}) + with pytest.raises(ValueError, match="shorter than one model window"): + Annotate(model, "time", device="cpu") + + def test_a_zero_overlap_still_windows(self): + model = annotate_model(default_args={"overlap": 0}) + result = Annotate(model, "time", device="cpu")(pin_array(("time", "distance"))) + assert result.sizes["time"] == len(PIN_SIGNAL) + assert not np.isnan(result.values).any() + + def test_the_final_window_is_end_aligned(self): + # 23 samples: the stride leaves a remainder of 3 samples that no + # grid-aligned window covers, and the output still spans the input. + da = pin_array(("time", "distance")).isel(time=slice(0, 23)) + result = Annotate(annotate_model(), "time", device="cpu")(da) + assert result.sizes["time"] == 23 + np.testing.assert_array_equal( + result.coords["time"].values, da.coords["time"].values + ) + + def test_a_record_of_exactly_one_window_is_annotated(self): + da = pin_array(("time", "distance")).isel(time=slice(0, 8)) + result = Annotate(annotate_model(), "time", device="cpu")(da) + assert result.sizes["time"] == 8 + + def test_a_record_shorter_than_one_window_raises(self): + da = pin_array(("time", "distance")).isel(time=slice(0, 7)) + with pytest.raises(ValueError, match="shorter along"): + Annotate(annotate_model(), "time", device="cpu")(da) + + @pytest.mark.parametrize("indices", [[9, 11], [3, 6, 9], [23]]) + def test_a_chunk_completing_no_window_holds_everything_back(self, indices): + da = pin_array(("time", "distance")) + expected = Annotate(annotate_model(), "time", device="cpu")(da) + picker = Annotate(annotate_model(), "time", device="cpu") + chunks = list(picker.iter_chunks(xd.split(da, indices, "time"), "time")) + assert xd.concat(chunks, "time").equals(expected) + + def test_a_chunked_record_shorter_than_one_window_raises_at_the_end(self): + # Chunked, the shortfall is only known once the stream has ended: the + # chunks pile up in the buffer, and the flush that closes the run says + # so rather than answering with nothing, as the eager call does upfront. + da = pin_array(("time", "distance")).isel(time=slice(0, 7)) + picker = Annotate(annotate_model(), "time", device="cpu") + with pytest.raises(ValueError, match="shorter along"): + list(picker.iter_chunks(xd.split(da, 3, "time"), "time")) + + def test_flushing_before_any_window_emits_nothing(self): + assert Annotate(annotate_model(), "time", device="cpu").flush() == [] + + def test_stacking_max_takes_the_running_maximum(self): + da = pin_array(("time", "distance")) + average = Annotate(annotate_model(), "time", device="cpu")(da) + maximum = Annotate(annotate_model(), "time", device="cpu", stacking="max")(da) + assert maximum.dims == average.dims + # the leading samples are covered by a single window: nothing to stack + np.testing.assert_allclose( + maximum.isel(time=slice(0, 4)).values, + average.isel(time=slice(0, 4)).values, + rtol=1e-6, + ) + assert np.all(maximum.values >= average.values - 1e-6) + assert not np.allclose(maximum.values, average.values) + + def test_stacking_max_leaves_uncovered_samples_undefined(self): + model = annotate_model(default_args={"overlap": 0.5, "blinding": (2, 0)}) + result = Annotate(model, "time", device="cpu", stacking="max")( + pin_array(("time", "distance")) + ) + assert np.isnan(result.isel(time=slice(0, 2)).values).all() + assert not np.isnan(result.isel(time=slice(2, None)).values).any() + + def test_an_unknown_stacking_rule_is_refused(self): + with pytest.raises(ValueError, match="stacking must be"): + Annotate(annotate_model(), "time", device="cpu", stacking="median") + + +class TestAnnotateChunkSemantics: + """ + `Annotate` carries its window across chunks along `dim`, elementwise across. + + Chunking along a dimension the atom does not work along must change + nothing. It used to be false here: the tail buffer was allocated only + when the chunking followed `dim`, but refilled on every call, so a run + chunked along ``distance`` concatenated one chunk's *time* tail onto the + next chunk's *other lanes* and came out ragged. + """ + + @pytest.mark.parametrize("size", [8, 13, 16, 32]) + def test_chunking_along_the_sample_dimension_is_invariant(self, size): + da = xd.testing.dummy(dims=("time", "distance"), shape=(64, 3)) + atom = Annotate(annotate_model(), "time", device="cpu") + xd.testing.assert_chunk_invariant(atom, da, {"time": size}) + + @pytest.mark.parametrize("size", [1, 2, 3]) + def test_chunking_along_another_dimension_is_invariant(self, size): + # Regression: 4 lanes chunked 2 by 2 used to give a `DataSequence` of + # ragged pieces with overlapping time coordinates and NaN values. + da = pin_array(("time", "distance"), n_lanes=4) + atom = Annotate(annotate_model(), "time", device="cpu") + xd.testing.assert_chunk_invariant(atom, da, {"distance": size}) + + def test_chunking_along_another_dimension_leaves_no_tail_behind(self): + da = pin_array(("time", "distance"), n_lanes=4) + atom = Annotate(annotate_model(), "time", device="cpu") + streamed = atom.process(da, chunks={"distance": 2}) + assert isinstance(streamed, xd.DataArray) + assert streamed.sizes["time"] == da.sizes["time"] + assert not np.isnan(streamed.values).any() + + def test_a_record_of_exactly_one_window_chunked_across_is_invariant(self): + # The body completes no window: everything comes out of `flush`. + da = pin_array(("time", "distance"), n_lanes=4).isel(time=slice(0, 8)) + atom = Annotate(annotate_model(), "time", device="cpu") + xd.testing.assert_chunk_invariant(atom, da, {"distance": 2}) + + def test_a_record_shorter_than_a_window_still_raises_when_chunked_across(self): + da = pin_array(("time", "distance"), n_lanes=4).isel(time=slice(0, 7)) + atom = Annotate(annotate_model(), "time", device="cpu") + with pytest.raises(ValueError, match="shorter along"): + atom.process(da, chunks={"distance": 2}) + + def test_the_component_dimension_is_not_a_lane_axis(self): + # The exemption is about *lanes*: a chunk holding a subset of the + # components is not a record the model can read. + da = component_array(["SHE", "SHN", "SHZ"]) + atom = Annotate(annotate_model(), "time", device="cpu") + with pytest.raises(ValueError, match="component dimension"): + atom.process(da, chunks={"channel": 2}) + + +class TestAnnotatePostShapeIsNamed: + """ + W9: a model whose ``annotate_batch_post`` breaks the stacking contract. + + The atom adopts SeisBench's ``(batch, samples, classes)``, which is + ``PhaseNet``'s convention, not the ``WaveformModel`` default. Surveying + the shipped SeisBench models, four keep the base default (``CRED``, + ``GPD``, ``DPPDetector``, ``DPPPicker``) and ``CRED`` is an ``"array"`` + model, so it reaches the accumulation and used to get a bare + ``RuntimeError`` from the broadcast. None of the 17 cached ``PhaseNet`` + weight sets does this. + """ + + def test_the_base_waveform_model_order_is_named(self): + model = annotate_model() + model.annotate_batch_post = lambda batch, piggyback, argdict: batch + picker = Annotate(model, "time", device="cpu") + with pytest.raises(ValueError, match=r"not \(8, 3\) = \(in_samples, classes\)"): + picker(pin_array(("time", "distance"))) + + def test_a_window_prediction_of_another_length_is_named(self): + model = annotate_model() + model.annotate_batch_post = lambda batch, piggyback, argdict: torch.transpose( + batch, -1, -2 + )[..., :2, :] + picker = Annotate(model, "time", device="cpu") + with pytest.raises(ValueError, match=r"batch ending in \(2, 3\)"): + picker(pin_array(("time", "distance"))) + + +class TestAnnotateComponents: + """W2: finding the component dimension and permuting it into the model.""" + + def test_components_are_permuted_into_the_model_order(self): + # ENZ weights fed ZNE data: every component lands in its own slot. + result = Annotate(annotate_model("original"), "time", device="cpu")( + component_array(["SHZ", "SHN", "SHE"]) + ) + assert result.dims == ("phase", "time") + assert [slot_of(result, index) for index in range(3)] == [2, 1, 0] + + def test_the_four_component_obs_layout_needs_no_special_case(self): + result = Annotate(annotate_model("obs"), "time", device="cpu")( + component_array(["HHZ", "HH1", "HH2", "HHH"]) + ) + assert [slot_of(result, index) for index in range(4)] == [0, 1, 2, 3] + + def test_horizontal_components_are_matched_flexibly(self): + # ZNE data on Z12H weights: N feeds the `1` slot and E the `2` slot. + result = Annotate(annotate_model("obs"), "time", device="cpu")( + component_array(["HHZ", "HHN", "HHE"]) + ) + assert [slot_of(result, index) for index in range(3)] == [0, 1, 2] + + def test_flexible_horizontal_components_can_be_switched_off(self): + da = component_array(["HHZ", "HHN", "HHE"]) + picker = Annotate( + annotate_model("obs"), + "time", + components="channel", + device="cpu", + flexible_horizontal_components=False, + ) + with pytest.raises(ValueError, match="not labelled by components"): + picker(da) + + def test_detection_declines_rather_than_guessing(self): + # Same data, no explicit `components=`: the dimension simply is not + # recognised, so it stays a batch axis. + picker = Annotate( + annotate_model("obs"), + "time", + device="cpu", + flexible_horizontal_components=False, + ) + result = picker(component_array(["HHZ", "HHN", "HHE"])) + assert result.dims == ("channel", "phase", "time") + + def test_duplicated_orientations_raise(self): + picker = Annotate(annotate_model(), "time", device="cpu") + with pytest.raises(ValueError, match="repeats component orientations"): + picker(component_array(["SHZ", "HHZ"])) + + def test_two_candidate_dimensions_raise(self): + template = xd.testing.dummy( + dims=("time", "station", "channel"), shape=(24, 2, 3), step=1.0 + ) + coords = dict(template.coords) + coords["station"] = np.array(["N", "E"]) + coords["channel"] = np.array(["SHZ", "SHN", "SHE"]) + da = xd.DataArray(np.zeros((24, 2, 3)), coords, ("time", "station", "channel")) + picker = Annotate(annotate_model(), "time", device="cpu") + with pytest.raises(ValueError, match="several dimensions"): + picker(da) + + def test_components_false_disables_detection(self): + picker = Annotate(annotate_model(), "time", components=False, device="cpu") + result = picker(component_array(["SHZ", "SHN", "SHE"])) + assert result.dims == ("channel", "phase", "time") + # each lane is cloned into all three slots: 1 + 2 + 3 + assert round(float(result.isel(time=0, channel=0, phase=0).values)) == 6 + + def test_components_names_the_dimension_explicitly(self): + da = component_array(["SHZ", "SHN", "SHE"]) + named = Annotate(annotate_model(), "time", components="channel", device="cpu") + detected = Annotate(annotate_model(), "time", device="cpu") + assert named(da).equals(detected(da)) + + def test_components_must_name_a_dimension(self): + picker = Annotate(annotate_model(), "time", components="sensor", device="cpu") + with pytest.raises(ValueError, match="is not a dimension of the input"): + picker(component_array(["SHZ", "SHN", "SHE"])) + + def test_components_must_name_a_labelled_dimension(self): + picker = Annotate(annotate_model(), "time", components="distance", device="cpu") + with pytest.raises(ValueError, match="not labelled by components"): + picker(pin_array(("time", "distance"))) + + def test_byte_labels_are_read_like_string_ones(self): + da = component_array(["SHZ", "SHN", "SHE"]) + expected = Annotate(annotate_model(), "time", device="cpu")(da) + da.coords["channel"] = np.array([b"SHZ", b"SHN", b"SHE"]) + assert Annotate(annotate_model(), "time", device="cpu")(da).equals(expected) + + def test_a_dimension_without_a_coordinate_is_never_a_candidate(self): + da = component_array(["SHZ", "SHN", "SHE"]) + bare = xd.DataArray(da.values, {"time": da.coords["time"]}, da.dims) + result = Annotate(annotate_model(), "time", device="cpu")(bare) + assert result.dims == ("channel", "phase", "time") + + +class TestAnnotateLayouts: + """W2: one trace, one instrument, a grid of instruments.""" + + def test_a_single_trace(self): + result = Annotate(annotate_model(), "time", device="cpu")(trace_array()) + assert result.dims == ("phase", "time") + assert slot_of(result, 0) == 5 # cloned: 1 + 2 + 3 + + def test_one_instrument_either_way_round(self): + rowwise = component_array(["SHZ", "SHN", "SHE"], dims=("time", "channel")) + colwise = component_array(["SHZ", "SHN", "SHE"], dims=("channel", "time")) + picker = Annotate(annotate_model(), "time", device="cpu") + assert picker(rowwise).equals(picker(colwise)) + + def test_a_grid_of_instruments(self): + template = xd.testing.dummy( + dims=("station", "channel", "time"), + shape=(2, 3, 24), + step=1.0, + datetime=False, + ) + coords = dict(template.coords) + coords["station"] = np.array(["ALPHA", "BRAVO"]) + coords["channel"] = np.array(["SHZ", "SHN", "SHE"]) + values = np.stack([spikes(3).T * (index + 1) for index in range(2)]) + da = xd.DataArray(values, coords, ("station", "channel", "time")) + result = Annotate(annotate_model("original"), "time", device="cpu")(da) + assert result.dims == ("station", "phase", "time") + for station in range(2): + slots = [slot_of(result, index, station=station) for index in range(3)] + assert slots == [2, 1, 0] # the lane gain is normalised away + + +class TestComponentStrategy: + """W2: `"auto"` resolving two ways, and each explicit strategy.""" + + def test_auto_clones_when_there_is_no_component_dimension(self): + result = Annotate(annotate_model(), "time", device="cpu")(trace_array()) + assert slot_of(result, 0) == 5 # every slot filled: 1 + 2 + 3 + + def test_auto_pads_a_partial_component_set(self): + result = Annotate(annotate_model("original"), "time", device="cpu")( + component_array(["SHZ"]) + ) + assert slot_of(result, 0) == 2 # the Z slot of ENZ, the rest zeroed + + @pytest.mark.parametrize( + "name, first, other", [("original", "E", "Z"), ("diting", "Z", "E")] + ) + def test_pad_fills_slot_zero_positionally(self, name, first, other): + # SeisBench's `"pad"` counts to the first slot rather than naming it: + # that is `E` for ENZ weights and `Z` for ZNE ones. + da = trace_array() + + def annotate(strategy): + picker = Annotate( + annotate_model(name), + "time", + component_strategy=strategy, + device="cpu", + ) + return picker(da) + + assert annotate("pad").equals(annotate(first)) + assert not annotate("pad").equals(annotate(other)) + + def test_clone_needs_a_single_component(self): + picker = Annotate( + annotate_model(), "time", component_strategy="clone", device="cpu" + ) + with pytest.raises(ValueError, match="needs a single component"): + picker(component_array(["SHZ", "SHN", "SHE"])) + + def test_clone_accepts_a_component_dimension_of_one(self): + picker = Annotate( + annotate_model(), "time", component_strategy="clone", device="cpu" + ) + assert slot_of(picker(component_array(["SHZ"])), 0) == 5 + + def test_a_named_slot_overrides_the_label(self): + picker = Annotate( + annotate_model("original"), "time", component_strategy="E", device="cpu" + ) + assert slot_of(picker(component_array(["SHZ"])), 0) == 0 + + def test_strict_accepts_a_complete_component_set(self): + picker = Annotate( + annotate_model("original"), + "time", + component_strategy="strict", + device="cpu", + ) + da = component_array(["SHZ", "SHN", "SHE"]) + assert picker(da).equals(Annotate(annotate_model(), "time", device="cpu")(da)) + + def test_strict_refuses_a_partial_component_set(self): + picker = Annotate( + annotate_model(), "time", component_strategy="strict", device="cpu" + ) + with pytest.raises(ValueError, match="component_strategy is 'strict'"): + picker(component_array(["SHZ", "SHN"])) + + def test_strict_refuses_data_without_components(self): + picker = Annotate( + annotate_model(), "time", component_strategy="strict", device="cpu" + ) + with pytest.raises(ValueError, match="needs a component dimension"): + picker(trace_array()) + + def test_an_unknown_strategy_is_refused(self): + with pytest.raises(ValueError, match="component_strategy must be one of"): + Annotate(annotate_model(), "time", component_strategy="nope", device="cpu") + + +class TestAnnotateAssumesNoName: + """W2: nothing in the implementation may spell `time` or `channel`.""" + + def test_neither_dimension_is_conventionally_named(self): + da = component_array( + ["SHZ", "SHN", "SHE"], dims=("samples", "sensor"), sample_dim="samples" + ) + picker = Annotate(annotate_model("original"), dim="samples", device="cpu") + result = picker(da) + assert result.dims == ("phase", "samples") + slots = [slot_of(result, index, sample_dim="samples") for index in range(3)] + assert slots == [2, 1, 0] + + def test_the_component_dimension_can_also_be_named_explicitly(self): + da = component_array( + ["SHZ", "SHN", "SHE"], dims=("samples", "sensor"), sample_dim="samples" + ) + named = Annotate( + annotate_model(), dim="samples", components="sensor", device="cpu" + ) + assert named(da).dims == ("phase", "samples") + + @pytest.mark.parametrize( + "dims, alias", + [(("samples", "sensor"), "first"), (("sensor", "samples"), "last")], + ) + def test_the_first_and_last_aliases_resolve(self, dims, alias): + da = component_array(["SHZ", "SHN", "SHE"], dims=dims, sample_dim="samples") + picker = Annotate(annotate_model(), dim=alias, device="cpu") + assert picker(da).dims == ("phase", "samples") + + def test_an_unknown_sample_dimension_raises(self): + picker = Annotate(annotate_model(), dim="nope", device="cpu") + with pytest.raises(ValueError, match="is not a dimension of the input"): + picker(pin_array(("time", "distance"))) + + +# --------------------------------------------------------------------------- +# W3 — the filter the weight set ships +# --------------------------------------------------------------------------- + +#: Long enough for a 0.5 Hz highpass at 100 Hz to be a filter rather than an +#: edge effect, and random so that no channel is a multiple of another. +FILTER_SIGNAL = np.random.default_rng(0).standard_normal(2000).cumsum() + + +def waveform(labels=("BHZ", "BH1", "BH2", "BDH"), dim="channel", step=0.01): + """A record of *labels* channels, each a differently rolled random walk.""" + values = np.stack( + [np.roll(FILTER_SIGNAL, 137 * index) for index in range(len(labels))], axis=-1 + ) + template = xd.testing.dummy( + dims=("time", dim), + shape=values.shape, + step=(step, 1.0), + ctype="interpolated", + ) + coords = dict(template.coords) + coords[dim] = np.array(labels) + return xd.DataArray(values, coords, ("time", dim)) + + +def obspy_filtered(da, name, dim="channel", **kwargs): + """The same record filtered by obspy, the reference W3 has to reproduce.""" + fs = 1.0 / xd.get_sampling_interval(da, "time") + traces = [] + for index in range(da.sizes[dim]): + trace = obspy.Trace(np.asarray(da.isel({dim: index}).values).copy()) + trace.stats.sampling_rate = fs + traces.append(trace) + obspy.Stream(traces).filter(name, **kwargs) + return np.stack([trace.data for trace in traces], axis=-1) + + +class TestChannelFilter: + """W3: the per-channel form, which SeisBench's own DAS wrapper refuses.""" + + def test_only_the_matching_channels_are_filtered(self): + da = waveform() + result = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + assert result.dims == da.dims + untouched = result.isel(channel=slice(0, 3)).values + assert np.array_equal(untouched, np.asarray(da.values)[:, :3]) + assert not np.allclose(result.isel(channel=3).values, da.values[:, 3]) + + def test_the_matching_channel_is_filtered_exactly_as_obspy_does(self): + da = waveform() + result = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + expected = obspy_filtered(da, "highpass", freq=0.5) + assert np.max(np.abs(result.isel(channel=3).values - expected[:, 3])) == 0.0 + + def test_a_pattern_matching_nothing_is_a_no_operation(self): + # SeisBench's `stream.select(channel=...)` on a pattern no trace + # answers filters an empty stream, which is not an error. + da = waveform() + result = _ChannelFilter("??Q", (0.5, None), component_order="Z12H")(da) + assert result.equals(da) + + def test_data_without_a_channel_dimension_is_left_alone(self): + da = pin_array(("time", "distance")) + result = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + assert result.equals(da) + + def test_components_false_disables_the_stage(self): + da = waveform() + atom = _ChannelFilter( + "??H", (0.5, None), component_order="Z12H", components=False + ) + assert atom(da).equals(da) + + def test_the_channel_dimension_is_found_by_its_labels_not_its_name(self): + result = _ChannelFilter("??H", (0.5, None), component_order="Z12H")( + waveform(dim="sensor") + ) + expected = _ChannelFilter("??H", (0.5, None), component_order="Z12H")( + waveform() + ) + assert np.allclose(result.values, expected.values) + + def test_the_channel_dimension_can_be_named_explicitly(self): + da = waveform(dim="sensor") + atom = _ChannelFilter( + "??H", (0.5, None), component_order="Z12H", components="sensor" + ) + assert not np.allclose(atom(da).values[:, 3], da.values[:, 3]) + + def test_horizontal_components_are_matched_flexibly_when_detecting(self): + # `ZNE` labels against `Z12H` weights: N -> 1 and E -> 2, so the + # dimension is still recognised and the hydrophone still filtered. + da = waveform(("BHZ", "BHN", "BHE", "BDH")) + result = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + assert not np.allclose(result.values[:, 3], da.values[:, 3]) + + @pytest.mark.parametrize("indices", [(500,), (137, 900, 1500)]) + def test_chunked_processing_equals_eager_processing(self, indices): + da = waveform() + atom = _ChannelFilter("??H", (0.5, None), component_order="Z12H") + expected = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + chunked = xd.concat( + [atom(chunk, chunk_dim="time") for chunk in xd.split(da, indices, "time")], + "time", + ) + assert np.allclose(chunked.values, expected.values) + + def test_integer_data_is_promoted_rather_than_truncated(self): + da = waveform() + da = xd.DataArray(np.asarray(da.values).astype(np.int32), da.coords, da.dims) + result = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + assert result.dtype == np.float64 + assert np.array_equal(result.values[:, :3], np.asarray(da.values)[:, :3]) + + def test_the_parameters_of_filter_are_inherited(self): + # Being a `Filter`, it takes `zerophase` — which the wrapper this + # replaced could not express — and the subset still holds. + da = waveform() + atom = _ChannelFilter( + "??H", (0.5, None), component_order="Z12H", zerophase=True + ) + result = atom(da) + causal = _ChannelFilter("??H", (0.5, None), component_order="Z12H")(da) + assert np.array_equal(result.values[:, :3], np.asarray(da.values)[:, :3]) + assert not np.allclose(result.values[:, 3], causal.values[:, 3]) + + def test_a_fir_filter_is_refused(self): + # It shifts the coordinate to compensate its group delay, which would + # leave the channels the pattern does not match on the wrong samples. + with pytest.raises(ValueError, match="cannot filter a subset"): + _ChannelFilter("??H", (0.5, None), ftype="fir") + + +class TestModelFilter: + """W3: the stage is built from the weight set, and only when it declares one.""" + + @pytest.mark.parametrize("name", ["original", "diting", "geofon"]) + def test_a_weight_set_declaring_none_gets_no_stage(self, name): + assert _model_filter(fake_model(name)) is None + + def test_the_flat_form_filters_everything(self): + # `volpick`'s flat filter is invented (no cached weight set ships one), + # but the form has to work: a 1 Hz highpass on every channel. + stage = _model_filter(fake_model("volpick")) + assert isinstance(stage, Filter) + assert (stage.freq, stage.order, stage.ftype) == ((1.0, None), 4, "iir") + da = waveform(("BHZ", "BHN", "BHE")) + expected = obspy_filtered(da, "highpass", freq=1.0) + assert np.max(np.abs(stage(da).values - expected)) == 0.0 + + def test_the_per_channel_form_filters_the_hydrophone_alone(self): + stage = _model_filter(fake_model("obs")) + assert isinstance(stage, _ChannelFilter) + assert stage.pattern == "??H" + assert stage.freq == (0.5, None) + assert stage.component_order == "Z12H" + da = waveform() + result = stage(da) + expected = obspy_filtered(da, "highpass", freq=0.5) + assert np.array_equal(result.values[:, :3], np.asarray(da.values)[:, :3]) + assert np.max(np.abs(result.values[:, 3] - expected[:, 3])) == 0.0 + + def test_the_corner_count_defaults_to_the_obspy_one(self): + # `obs` declares only `freq`, so 4 corners is what obspy would have used + # — and it is `Filter`'s own default, which is why the match is exact. + assert _model_filter(fake_model("obs")).order == 4 + + def test_a_declared_corner_count_is_honoured(self): + model = fake_model( + filter_args=("highpass",), filter_kwargs={"freq": 1.0, "corners": 2} + ) + assert _model_filter(model).order == 2 + + @pytest.mark.parametrize( + "args, kwargs, freq", + [ + (("highpass",), {"freq": 2.0}, (2.0, None)), + (("lowpass",), {"freq": 20.0}, (None, 20.0)), + (("bandpass",), {"freqmin": 1.0, "freqmax": 20.0}, (1.0, 20.0)), + ], + ) + def test_every_obspy_band_that_has_a_filter_equivalent(self, args, kwargs, freq): + stage = _model_filter(fake_model(filter_args=args, filter_kwargs=kwargs)) + assert stage.freq == freq + da = waveform(("BHZ", "BHN", "BHE")) + expected = obspy_filtered(da, args[0], **kwargs) + assert np.max(np.abs(stage(da).values - expected)) == 0.0 + + def test_zerophase_warns_and_doubles_the_order(self): + # SeisBench's concession, and ours for the same reason: `filtfilt` has + # no causal streaming form. + model = fake_model( + filter_args=("highpass",), + filter_kwargs={"freq": 1.0, "zerophase": True}, + ) + with pytest.warns(UserWarning, match="no causal streaming form"): + stage = _model_filter(model) + assert stage.order == 8 + + def test_a_corner_above_half_the_nyquist_is_clamped(self): + # At 100 Hz the Nyquist is 50 and half of it 25, so a 40 Hz lowpass + # comes back clamped — the filter stays valid on data at a lower rate. + model = fake_model(filter_args=("lowpass",), filter_kwargs={"freq": 40.0}) + assert _model_filter(model).freq == (None, pytest.approx(25.0, rel=1e-5)) + + def test_the_clamp_follows_the_weight_sets_own_rate(self): + model = fake_model( + "diting", filter_args=("lowpass",), filter_kwargs={"freq": 40.0} + ) + assert _model_filter(model).freq == (None, pytest.approx(12.5, rel=1e-5)) + + def test_a_model_without_a_rate_is_not_clamped(self): + model = fake_model( + sampling_rate=None, filter_args=("lowpass",), filter_kwargs={"freq": 40.0} + ) + assert _model_filter(model).freq == (None, 40.0) + + def test_several_patterns_become_a_sequence_applied_in_order(self): + model = fake_model( + "obs", + filter_args={"??Z": ("highpass",), "??H": ("highpass",)}, + filter_kwargs={"??Z": {"freq": 1.0}, "??H": {"freq": 0.5}}, + ) + stage = _model_filter(model) + assert isinstance(stage, Sequential) + assert [atom.pattern for atom in stage] == ["??Z", "??H"] + da = waveform() + result = stage(da) + expected = obspy_filtered(da, "highpass", freq=1.0) + assert np.max(np.abs(result.values[:, 0] - expected[:, 0])) == 0.0 + assert np.array_equal(result.values[:, 1:3], np.asarray(da.values)[:, 1:3]) + + def test_overlapping_patterns_filter_a_channel_twice_as_obspy_would(self): + # Two `stream.select` calls hitting the same trace filter it twice, in + # declaration order; nothing deduplicates them, here or in SeisBench. + model = fake_model( + "obs", + filter_args={"??H": ("highpass",), "BD?": ("highpass",)}, + filter_kwargs={"??H": {"freq": 0.5}, "BD?": {"freq": 0.5}}, + ) + da = waveform() + once = _ChannelFilter("??H", (0.5, None), component_order="Z12H") + twice = _model_filter(model)(da) + assert not np.allclose(twice.values[:, 3], once(da).values[:, 3]) + assert np.allclose(twice.values[:, 3], once(once(da)).values[:, 3]) + + def test_an_empty_declaration_gets_no_stage(self): + assert _model_filter(fake_model(filter_args={}, filter_kwargs={})) is None + + def test_a_pattern_missing_from_the_kwargs_raises(self): + model = fake_model(filter_args={"??H": ("highpass",)}, filter_kwargs={}) + with pytest.raises(ValueError, match="in `filter_args` but not in"): + _model_filter(model) + + def test_a_declaration_naming_no_single_filter_type_raises(self): + model = fake_model(filter_args=("highpass", "lowpass"), filter_kwargs={}) + with pytest.raises(ValueError, match="exactly one obspy filter type"): + _model_filter(model) + + def test_a_band_with_no_filter_equivalent_raises(self): + # `Filter` has no bandstop, and silently skipping the filter would feed + # the network what it was not trained on — the very bug W3 fixes. + model = fake_model( + filter_args=("bandstop",), filter_kwargs={"freqmin": 1.0, "freqmax": 2.0} + ) + with pytest.raises(ValueError, match="'bandstop' filter, which has no"): + _model_filter(model) + + def test_the_stage_reads_the_dimension_names_it_is_given(self): + da = waveform(dim="sensor").rename({"time": "samples"}) + stage = _model_filter(fake_model("obs"), dim="samples", components="sensor") + assert not np.allclose(stage(da).values[:, 3], da.values[:, 3]) + + +class TestFilterStageOrder: + """W3: the model's filter runs *before* the resampling, as SeisBench does.""" + + def test_filtering_before_resampling_is_not_the_same_operation(self): + da = waveform(("BHZ", "BHN", "BHE"), step=0.02) # 50 Hz, model wants 100 + before = (_model_filter(fake_model("volpick")) >> Resample(100.0))(da) + after = (Resample(100.0) >> _model_filter(fake_model("volpick")))(da) + assert before.shape == after.shape + assert not np.allclose(before.values, after.values) + + def test_the_stage_composes_ahead_of_resample_and_annotate(self): + da = waveform(("BHZ", "BHN", "BHE"), step=0.02) + model = fake_model("volpick", default_args={"overlap": 0.5}) + pipeline = ( + _model_filter(model) + >> Resample(model.sampling_rate) + >> Annotate(model, "time", device="cpu") + ) + assert [type(stage).__name__ for stage in pipeline] == [ + "Filter", + "Resample", + "Annotate", + ] + result = pipeline(da) + assert result.dims == ("phase", "time") + assert np.isfinite(result.values).any() + + +def component_trace(index, start=0.0, n=24, step=0.01, sample_dim="time"): + """ + One component's trace on a grid that declares its own sampling interval. + + The leaf shape a collection of single-channel traces has, and what + :meth:`~xdas.atoms.Annotate.gather` collapses back into a component + dimension. *start* offsets the grid, which is how the sub-sample rounding + differences of real data are reproduced. + """ + return xd.DataArray( + spikes(index + 1, n)[:, index], + { + sample_dim: { + "tie_indices": [0, n - 1], + "tie_values": [start, start + step * (n - 1)], + "sampling_interval": step, + } + }, + (sample_dim,), + ) + + +def das_array(start=0.0, n=24, step=0.01, nchannels=3): + """One DAS acquisition: a section with a numeric spatial axis.""" + return xd.DataArray( + np.tile(spikes(1, n), (1, nchannels)), + { + "time": { + "tie_indices": [0, n - 1], + "tie_values": [start, start + step * (n - 1)], + "sampling_interval": step, + }, + "distance": 10.0 * np.arange(nchannels), + }, + ("time", "distance"), + ) + + +def component_collection(keys, level="channel", starts=None, **kwargs): + """A collection level holding one single-component trace per key.""" + starts = [0.0] * len(keys) if starts is None else starts + return xd.DataCollection( + { + key: component_trace(index, start=start, **kwargs) + for index, (key, start) in enumerate(zip(keys, starts)) + }, + level, + ) + + +def leaves(obj): + """Every data array of a walked collection, in walk order.""" + if isinstance(obj, xd.DataArray): + return [obj] + values = obj.values() if obj.ismapping() else obj + return [leaf for value in values for leaf in leaves(value)] + + +def assert_same_result(left, right): + """Assert two walked results hold the same arrays, dimension names included.""" + left, right = leaves(left), leaves(right) + assert len(left) == len(right) + for one, other in zip(left, right): + assert one.dims == other.dims + npt.assert_allclose(one.values, other.values) + + +class TestAnnotateGathersItsComponentLevel: + """ + W5: a collection keeping one leaf per channel is one input, not three. + + `xd.pick(dc, model)` must not require a prior `xd.stack`, so `Annotate` + implements the `gather` hook: the atom that knows the model's + `component_order` is the one that knows which leaves group together. + """ + + def test_a_channel_level_becomes_the_component_dimension(self): + dc = component_collection(["SHZ", "SHN", "SHE"]) + result = Annotate(annotate_model("original"), device="cpu")(dc) + assert isinstance(result, xd.DataArray) + assert result.dims == ("phase", "time") + assert [slot_of(result, index) for index in range(3)] == [2, 1, 0] + + @pytest.mark.parametrize("level", ["channel", "component"]) + def test_every_default_candidate_name_is_gathered(self, level): + dc = component_collection(["SHZ", "SHN", "SHE"], level=level) + result = Annotate(annotate_model(), device="cpu")(dc) + assert result.dims == ("phase", "time") + + def test_an_obs_station_gathers_despite_its_hydrophone(self): + # `BDH` is a pressure instrument, so the four keys share their band + # code and nothing more: requiring a common two-character stem would + # refuse exactly the layout the `Z12H` weights exist for. + dc = component_collection(["BHZ", "BH1", "BH2", "BDH"]) + result = Annotate(annotate_model("obs"), device="cpu")(dc) + assert result.dims == ("phase", "time") + assert [slot_of(result, index) for index in range(4)] == [0, 1, 2, 3] + + def test_a_level_named_otherwise_is_left_alone(self): + # the keys resolve, the name is not a candidate: three stations whose + # codes happen to end in a component letter stay three stations + dc = component_collection(["ABZ", "ABN", "ABE"], level="station") + result = Annotate(annotate_model(), device="cpu")(dc) + assert isinstance(result, xd.DataCollection) + assert list(result) == ["ABZ", "ABN", "ABE"] + assert all(leaf.dims == ("phase", "time") for leaf in leaves(result)) + + def test_keys_naming_no_component_map_over_the_leaves(self): + # several DAS formats call their spatial axis `channel`; such a level + # must walk straight past the gather rather than trip over it + dc = component_collection(["0", "1", "2"]) + result = Annotate(annotate_model(), device="cpu")(dc) + assert isinstance(result, xd.DataCollection) + assert len(leaves(result)) == 3 + + def test_station_codes_ending_in_a_horizontal_letter_do_not_resolve(self): + # the whole reason the key rule is not W2's last-letter matcher: the + # `1` <-> `N` and `2` <-> `E` flexibility makes `STA1`/`STA2` a clean + # pair of horizontals, and folding two stations into one instrument + # would be silent + dc = component_collection(["STA1", "STA2"]) + result = Annotate(annotate_model("diting"), device="cpu")(dc) + assert isinstance(result, xd.DataCollection) + assert len(leaves(result)) == 2 + + def test_partially_resolving_keys_raise(self): + dc = component_collection(["SHZ", "SHN", "foo"]) + with pytest.raises(ValueError, match="and other things"): + Annotate(annotate_model(), device="cpu")(dc) + + def test_repeated_orientations_raise(self): + dc = component_collection(["SHZ", "SHE", "HHZ"]) + with pytest.raises(ValueError, match="repeats component orientations"): + Annotate(annotate_model(), device="cpu")(dc) + + def test_keys_differing_by_more_than_their_orientation_raise(self): + dc = component_collection(["SHZ", "HHN", "HHE"]) + with pytest.raises(ValueError, match="does not name one instrument"): + Annotate(annotate_model(), device="cpu")(dc) + + def test_a_lone_orientation_letter_is_a_channel_code(self): + dc = component_collection(["Z", "N", "E"]) + result = Annotate(annotate_model("original"), device="cpu")(dc) + assert [slot_of(result, index) for index in range(3)] == [2, 1, 0] + + def test_mixing_lone_letters_with_full_codes_raises(self): + dc = component_collection(["Z", "SHN", "SHE"]) + with pytest.raises(ValueError, match="does not name one instrument"): + Annotate(annotate_model(), device="cpu")(dc) + + def test_an_empty_level_maps_over_its_nothing(self): + # `query` and `sel` can leave a level with no keys behind; there is + # nothing to identify as components, so it is not a component level + result = Annotate(annotate_model(), device="cpu")( + xd.DataCollection({}, "channel") + ) + assert isinstance(result, xd.DataCollection) + assert leaves(result) == [] + + def test_components_false_restores_the_leaf_by_leaf_walk(self): + dc = component_collection(["SHZ", "SHN", "SHE"]) + result = Annotate(annotate_model(), components=False, device="cpu")(dc) + assert isinstance(result, xd.DataCollection) + assert len(leaves(result)) == 3 + + def test_the_candidate_names_are_a_default_not_a_rule(self): + dc = component_collection(["SHZ", "SHN", "SHE"], level="orientation") + atom = Annotate(annotate_model(), components="orientation", device="cpu") + assert atom(dc).dims == ("phase", "time") + + def test_an_explicitly_named_level_naming_no_component_raises(self): + # asking for a level by name and getting nothing back is an error, + # where the same keys under a *default* candidate name merely map + dc = component_collection(["0", "1", "2"], level="orientation") + atom = Annotate(annotate_model(), components="orientation", device="cpu") + with pytest.raises(ValueError, match="could not identify the component level"): + atom(dc) + + def test_the_snapping_tolerance_reaches_stack(self): + # a thousandth of a sample apart: one grid, two roundings of it + dc = component_collection(["SHZ", "SHN", "SHE"], starts=[0.0, 1e-5, 0.0]) + assert Annotate(annotate_model(), device="cpu")(dc).dims == ("phase", "time") + strict = Annotate(annotate_model(), tolerance=False, device="cpu") + with pytest.raises(ValueError, match="coordinate 'time' differs"): + strict(dc) + + def test_a_das_collection_walks_past_the_gather_untouched(self): + acquisitions = [das_array(start) for start in (0.0, 100.0)] + dc = xd.DataCollection( + { + "N1": xd.DataCollection( + {"C1": xd.DataCollection(acquisitions, "acquisition")}, "cable" + ) + }, + "node", + ) + result = Annotate(annotate_model(), device="cpu")(dc) + assert result.name == "node" + assert result["N1"].name == "cable" + assert all( + leaf.dims == ("distance", "phase", "time") for leaf in leaves(result) + ) + + def test_a_das_cable_whose_spatial_level_is_named_channel_maps(self): + trace = trace_array(n=24) + dc = xd.DataCollection( + {"C1": xd.DataCollection({str(i): trace for i in range(3)}, "channel")}, + "cable", + ) + result = Annotate(annotate_model(), device="cpu")(dc) + assert result["C1"].name == "channel" + assert len(leaves(result)) == 3 + + +class TestGatherIsAHookOnTheWalk: + """W5: `gather` is consulted by the walk, and pipelines delegate it.""" + + def test_an_atom_that_declares_no_gather_maps_over_the_leaves(self): + dc = component_collection(["SHZ", "SHN", "SHE"]) + result = Filter((0.5, None), dim="time")(dc) + assert isinstance(result, xd.DataCollection) + assert list(result) == ["SHZ", "SHN", "SHE"] + + def test_a_pipeline_delegates_to_the_stage_that_knows_the_model(self): + dc = component_collection(["SHZ", "SHN", "SHE"]) + pipeline = Filter((0.5, None), dim="time") >> Annotate( + annotate_model(), device="cpu" + ) + assert pipeline(dc).dims == ("phase", "time") + + def test_a_pipeline_of_unclaiming_stages_gathers_nothing(self): + dc = component_collection(["SHZ", "SHN", "SHE"]) + pipeline = Filter((0.5, None), dim="time") >> Filter((None, 10.0), dim="time") + assert len(leaves(pipeline(dc))) == 3 + + def test_the_gather_happens_before_the_first_stage(self): + # W3's per-channel filter selects `??H` out of `Z12H`, which it can + # only do once the channels are a dimension: if the gather ran stage + # by stage the filter would see three separate one-channel leaves and + # match nothing at all + model = annotate_model("obs") + dc = component_collection(["BHZ", "BH1", "BH2", "BDH"], n=64) + filtered = (_model_filter(model) >> Annotate(model, device="cpu"))(dc) + plain = Annotate(model, device="cpu")(dc) + assert filtered.dims == plain.dims == ("phase", "time") + assert not np.allclose(filtered.values, plain.values) + + +class TestPreStackingAgrees: + """ + W5: both input shapes are accepted, and they agree. + + `annotate(dc)` equals `annotate(xd.stack(dc, "channel"))` equals + `annotate(xd.stack(dc, "channel", dim="component"))`: the gather turns a + channel *level* into a channel *dimension* and does nothing more, so a + pre-stacked input simply enters the pipeline one step further along. + """ + + def station(self): + return component_collection(["SHZ", "SHN", "SHE"]) + + def tree(self): + return xd.DataCollection( + {"IA": xd.DataCollection({"ABC": self.station()}, "station")}, "network" + ) + + @pytest.mark.parametrize("shape", ["station", "tree"]) + @pytest.mark.parametrize("dim", [None, "component"]) + def test_stacking_first_gives_the_same_answer(self, shape, dim): + dc = getattr(self, shape)() + gathered = Annotate(annotate_model("original"), device="cpu")(dc) + stacked = Annotate(annotate_model("original"), device="cpu")( + xd.stack(dc, "channel", dim=dim) + ) + assert_same_result(gathered, stacked) + + def test_the_gather_is_a_no_op_on_a_collection_of_stacked_arrays(self): + dc = xd.DataCollection({"ABC": xd.stack(self.station(), "channel")}, "station") + result = Annotate(annotate_model(), device="cpu")(dc) + assert list(result) == ["ABC"] + assert result["ABC"].dims == ("phase", "time") + + +class DetectionModel(FakeModel): + """ + A :class:`FakeModel` shaped like the ``EQTransformer`` family. + + Three of them, on both counts that matter: ``classes`` counts only the + picking heads, so it disagrees with the number of labels, and ``forward`` + returns one tensor per head rather than one stacked batch. The values are + the ones :class:`FakeModel` would produce for the same three labels, so the + two can be compared directly. + """ + + #: SeisBench's name for the *picked* subset of the labels — not the label + #: list, which is what `Annotate.phases` holds. + phases = "PS" + + def __init__(self, **kwargs): + super().__init__(labels=("Detection", "P", "S"), classes=2, **kwargs) + + def forward(self, x): + """Return one ``(batch, samples)`` tensor per label, as its heads do.""" + weights = torch.arange( + 1, self.in_channels + 1, dtype=x.dtype, device=x.device + ).reshape(1, -1, 1) + pooled = (x * weights).sum(dim=-2) + return tuple(pooled + index for index in range(len(self.labels))) + + def annotate_batch_post(self, batch, piggyback, argdict): + """Stack the heads before blinding them, as ``EQTransformer`` does.""" + batch = torch.stack(batch, dim=1) + return super().annotate_batch_post(batch, piggyback, argdict) + + +def detection_model(**overrides): + """:func:`annotate_model`'s counterpart for the detection-trace family.""" + overrides.setdefault("default_args", {"overlap": 0.5}) + return DetectionModel(**overrides).eval() + + +class TestAnnotateLabelsAreOptional: + """ + W6: `model.labels` is per-weights, and SeisBench lets it be absent. + + It lands with the label-keyed thresholds because the `phase` coordinate is + what those thresholds key on: a wrong or missing label there silently + thresholds the wrong class. + """ + + def test_absent_labels_fall_back_to_positional_ones(self): + # What `WaveformModel._predictions_to_stream` does with `labels=None`. + model = annotate_model("original", labels=None, classes=3) + picker = Annotate(model, "time", device="cpu") + assert picker.phases == [0, 1, 2] + result = picker(pin_array(("time", "distance"))) + assert list(result.coords["phase"].values) == [0, 1, 2] + + def test_callable_labels_are_refused_by_name(self): + model = annotate_model("original", labels=lambda stations: "P", classes=3) + picker = Annotate(model, "time", device="cpu") + with pytest.raises(TypeError, match="`labels` is a callable"): + _ = picker.phases + + def test_a_model_declaring_neither_labels_nor_classes_raises(self): + model = annotate_model("original", labels=None, classes=3) + del model.classes + picker = Annotate(model, "time", device="cpu") + with pytest.raises(ValueError, match="neither `labels` nor `classes`"): + picker(pin_array(("time", "distance"))) + + +class TestAnnotateDetectionModels: + """ + The `EQTransformer` family, whose `classes` counts fewer than its labels. + + Its detection trace is an output like any other but not a phase, so the + architecture leaves it out of `classes` and the weight set out of `phases`. + Sizing the buffers on `classes` used to make the whole family unrunnable. + """ + + def test_the_class_count_comes_from_the_labels(self): + atom = Annotate(detection_model(), "time", device="cpu") + assert atom.model.classes == 2 + assert atom.classes == 3 + + def test_every_label_is_annotated(self): + result = Annotate(detection_model(), "time", device="cpu")(trace_array()) + assert result.dims == ("phase", "time") + assert list(result.coords["phase"].values) == ["Detection", "P", "S"] + + def test_the_heads_agree_with_a_model_stacking_them_itself(self): + da = component_array(["SHZ", "SHN", "SHE"]) + reference = annotate_model(None, labels=("Detection", "P", "S")) + assert_same_result( + Annotate(detection_model(), "time", device="cpu")(da), + Annotate(reference, "time", device="cpu")(da), + ) + + def test_the_detection_trace_is_no_more_picked_than_noise(self): + assert _model_thresholds(detection_model()) == {"P": 0.3, "S": 0.3} + picks = Picker(detection_model(), device="cpu")( + component_array(["SHZ", "SHN", "SHE"]) + ) + assert set(picks["phase"]) == {"P", "S"} + + +# --------------------------------------------------------------------------- +# W7 — `Picker` +# --------------------------------------------------------------------------- + + +def picker_model(name="original", **overrides): + """A preset windowed like :func:`annotate_model`, for the picker tests.""" + default_args = dict(WEIGHT_SETS[name]["default_args"]) + default_args.setdefault("overlap", 0.5) + overrides.setdefault("default_args", default_args) + return fake_model(name, **overrides) + + +def station_tree(keys=("SHZ", "SHN", "SHE"), stations=("DBNFM", "LBFI")): + """A ``network / station / location / channel`` tree, one trace per channel.""" + return xd.DataCollection( + { + "IA": xd.DataCollection( + { + station: xd.DataCollection( + {"--": component_collection(list(keys))}, "location" + ) + for station in stations + }, + "station", + ) + }, + "network", + ) + + +def stage_names(pipeline): + """The class names of a pipeline's stages, in order.""" + return [type(stage).__name__ for stage in pipeline] + + +class TestPickerStages: + """ + W7: every stage is built from the weight set, and each one is droppable. + + Three stages for `original` and four for `obs`, from one model class: that + is §3's point made executable. + """ + + def test_a_weight_set_declaring_no_filter_is_three_stages(self): + picker = Picker(picker_model("original"), device="cpu") + assert stage_names(picker) == ["Resample", "Annotate", "Trigger"] + + def test_the_obs_weight_set_is_one_stage_longer(self): + picker = Picker(picker_model("obs"), device="cpu") + assert stage_names(picker) == [ + "_ChannelFilter", + "Resample", + "Annotate", + "Trigger", + ] + assert picker[0].pattern == "??H" + assert picker[0].freq == (0.5, None) + + def test_a_flat_filter_is_a_stage_too(self): + picker = Picker(picker_model("volpick"), device="cpu") + assert stage_names(picker) == ["Filter", "Resample", "Annotate", "Trigger"] + + def test_the_filter_leads_the_resampling(self): + # W3: `annotate_stream_pre` filters *before* resampling, and filtering + # after is a different operation. + names = stage_names(Picker(picker_model("obs"), device="cpu")) + assert names.index("_ChannelFilter") < names.index("Resample") + + @pytest.mark.parametrize("name, rate", [("original", 100), ("diting", 50)]) + def test_the_resampling_targets_the_weight_sets_own_rate(self, name, rate): + # It is not always 100: `diting` runs at 50. + picker = Picker(picker_model(name), device="cpu") + assert picker[0].target == rate + + def test_the_resampling_is_a_no_op_on_data_already_at_that_rate(self): + da = component_array(["SHZ", "SHN", "SHE"]) # 100 Hz, as the model wants + resample = Picker(picker_model("original"), device="cpu")[0] + npt.assert_array_equal(resample(da).values, da.values) + + def test_resample_false_drops_the_stage(self): + picker = Picker(picker_model("original"), resample=False, device="cpu") + assert stage_names(picker) == ["Annotate", "Trigger"] + + def test_filter_false_drops_the_declared_filter(self): + picker = Picker(picker_model("obs"), filter=False, device="cpu") + assert stage_names(picker) == ["Resample", "Annotate", "Trigger"] + + def test_a_weight_set_without_a_rate_needs_resample_false(self): + model = picker_model("original", sampling_rate=None) + with pytest.raises(ValueError, match="declares no `sampling_rate`"): + Picker(model, device="cpu") + assert stage_names(Picker(model, resample=False, device="cpu")) == [ + "Annotate", + "Trigger", + ] + + def test_the_sample_dimension_reaches_every_stage(self): + picker = Picker(picker_model("obs"), dim="t", device="cpu") + assert [stage.dim for stage in picker] == ["t"] * 4 + + @pytest.mark.parametrize("alias", ["first", "last"]) + def test_the_positional_dimension_aliases_are_refused(self, alias): + # The picks are annotated with coordinates, which are named, so the + # sample dimension must be nameable before any data is seen. + with pytest.raises(ValueError, match="needs its sample dimension by name"): + Picker(picker_model(), dim=alias, device="cpu") + + def test_the_annotate_kwargs_reach_the_model(self): + picker = Picker(picker_model("original"), overlap=0, device="cpu") + assert picker[1].noverlap == 0 + + def test_it_is_a_pipeline_like_any_other(self): + picker = Picker(picker_model(), device="cpu") + assert isinstance(picker, Sequential) + assert isinstance(picker, xd.atoms.Atom) + assert repr(picker).startswith("Picker:") + + def test_it_pickles(self): + picker = Picker(picker_model(), device="cpu") + assert stage_names(pickle.loads(pickle.dumps(picker))) == stage_names(picker) + + def test_it_composes(self): + pipeline = xd.filter(..., (1.0, None), dim="time") >> Picker( + picker_model(), device="cpu" + ) + assert isinstance(pipeline, Sequential) + assert stage_names(pipeline) == ["Filter", "Picker"] + + def test_it_inherits_the_collection_hooks_of_its_stages(self): + picker = Picker(picker_model(), device="cpu") + assert picker.merge is not None # from the trigger, the last stage + assert picker.gather(component_collection(["SHZ", "SHN", "SHE"])) is not None + + +class TestPickerThresholds: + """W7: the thresholds are the weight set's own, one per non-noise label.""" + + @pytest.mark.parametrize("name", list(WEIGHT_SETS)) + def test_the_noise_label_never_gets_an_entry(self, name): + model = fake_model(name) + assert "N" in model.labels # every preset carries a noise class + assert set(_model_thresholds(model)) == {"P", "S"} + + def test_the_thresholds_are_read_off_the_weight_set(self): + assert _model_thresholds(fake_model("geofon")) == { + "P": 0.5704745853696115, + "S": 0.07349645833964447, + } + + def test_an_undeclared_threshold_falls_back_to_the_documented_default(self): + assert _model_thresholds(fake_model("original")) == {"P": 0.3, "S": 0.3} + + def test_the_fallback_is_the_models_own_documentation(self): + model = fake_model("original", annotate_args={"*_threshold": 0.7}) + assert _model_thresholds(model) == {"P": 0.7, "S": 0.7} + + def test_a_model_documenting_one_phase_apart_is_honoured(self): + # `eqcct` documents an `S_threshold` of its own next to the catch-all. + model = fake_model("original", annotate_args={"S_threshold": 0.42}) + assert _model_thresholds(model) == {"P": 0.3, "S": 0.42} + + def test_a_model_documenting_nothing_at_all_falls_back_to_the_constant(self): + class Plain: + labels, default_args = "PSN", {} + + assert _model_thresholds(Plain()) == {"P": 0.3, "S": 0.3} + + def test_the_call_wins_over_the_weight_set(self): + model = picker_model("geofon") + assert _model_thresholds(model, S_threshold=0.5)["S"] == 0.5 + assert Picker(model, S_threshold=0.5, device="cpu")[-1].thresh["S"] == 0.5 + + def test_a_threshold_above_one_is_passed_through_faithfully(self): + # `iquique` declares P = 1.12, which simply never fires: that is the + # weight set's own metadata, not ours to clamp. + model = fake_model("original", default_args={"P_threshold": 1.12}) + assert _model_thresholds(model)["P"] == 1.12 + + def test_positional_labels_get_positional_thresholds(self): + model = fake_model("original", labels=None, classes=3) + assert _model_thresholds(model) == {0: 0.3, 1: 0.3, 2: 0.3} + + def test_the_label_order_of_the_weight_set_is_irrelevant(self): + # `original` labels `NPS` and `geofon` `PSN`; keying on the label is + # what makes a positional flip harmless. + assert list(_model_thresholds(fake_model("original"))) == ["P", "S"] + assert list(_model_thresholds(fake_model("geofon"))) == ["P", "S"] + + def test_the_picker_takes_them_by_default(self): + picker = Picker(picker_model("obs"), device="cpu") + assert picker[-1].thresh == {"P": 0.2, "S": 0.1} + + @pytest.mark.parametrize("thresh", [0.5, {"P": 0.5}]) + def test_thresh_overrides_the_weight_set(self, thresh): + picker = Picker(picker_model("geofon"), thresh=thresh, device="cpu") + assert picker[-1].thresh == thresh + + +class TestPickerPicks: + """W7: §1's two lines, and the table they promise.""" + + def test_one_station_gives_one_table(self): + da = component_array(["SHZ", "SHN", "SHE"]) + picks = Picker(picker_model("original"), device="cpu")(da) + assert isinstance(picks, pd.DataFrame) + assert list(picks.columns) == ["phase", "time", "value"] + assert list(picks["phase"]) == ["P", "S"] # the noise class is not picked + + def test_scalar_coordinates_annotate_the_picks(self): + da = component_array(["SHZ", "SHN", "SHE"]) + da = da.assign_coords(network="IA", station="DBNFM", location="--") + picks = Picker(picker_model("original"), device="cpu")(da) + assert list(picks.columns) == [ + "network", + "station", + "location", + "phase", + "time", + "value", + ] + assert set(picks["station"]) == {"DBNFM"} + + def test_a_whole_network_in_one_call(self): + # §1: the picker gathers the channel level itself and merges the + # per-leaf tables, so a collection answers with one flat table. + picks = xd.pick(station_tree(), picker_model("original"), device="cpu") + assert isinstance(picks, pd.DataFrame) + assert list(picks.columns) == [ + "network", + "station", + "location", + "phase", + "time", + "value", + ] + assert list(picks["station"]) == ["DBNFM", "DBNFM", "LBFI", "LBFI"] + assert list(picks["phase"]) == ["P", "S", "P", "S"] + + def test_pre_stacking_the_channels_agrees(self): + model = picker_model("original") + dc = component_collection(["SHZ", "SHN", "SHE"]) + gathered = xd.pick(dc, model, device="cpu") + stacked = xd.pick(xd.stack(dc, "channel"), model, device="cpu") + assert gathered.equals(stacked) + + def test_components_false_walks_leaf_by_leaf(self): + dc = component_collection(["SHZ", "SHN", "SHE"]) + picks = xd.pick(dc, picker_model("original"), components=False, device="cpu") + assert list(picks["channel"]) == ["SHZ", "SHZ", "SHN", "SHN", "SHE", "SHE"] + + def test_a_threshold_no_lane_reaches_gives_an_empty_table(self): + da = component_array(["SHZ", "SHN", "SHE"]) + picks = Picker(picker_model("original"), thresh=1e6, device="cpu")(da) + assert picks.empty + + def test_the_eager_and_streamed_forms_agree(self): + da = component_array(["SHZ", "SHN", "SHE"], n=64) + model = picker_model("original") + expected = Picker(model, device="cpu")(da) + streamed = Picker(model, device="cpu").process(da, chunks={"time": 13}) + assert streamed.equals(expected) + + def test_the_twin_returns_the_atom_on_an_ellipsis(self): + picker = xd.pick(..., picker_model("original"), device="cpu") + assert isinstance(picker, Picker) + + def test_the_twin_and_the_class_are_two_faces_of_one_thing(self): + da = component_array(["SHZ", "SHN", "SHE"]) + model = picker_model("original") + assert xd.pick(da, model, device="cpu").equals(Picker(model, device="cpu")(da)) + + def test_a_gap_splits_the_record_into_runs(self): + # `Atom.__call__` splits at coordinate discontinuities, so a gappy + # station is picked segment by segment, as SeisBench's grouping does. + da = component_array(["SHZ", "SHN", "SHE"], n=64) + gappy = xd.concat( + [da.isel(time=slice(0, 24)), da.isel(time=slice(40, 64))], "time" + ) + picks = Picker(picker_model("original"), device="cpu")(gappy) + assert len(picks) == 4 # one P and one S per run + + +DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +class TestAnnotateAsyncOutputs: + """ + The in-flight output queue changes *when* outputs appear, never what + they hold: max_buffers=0 is the fully synchronous reference, and the + default queue must answer identically, eagerly and chunked, on the CPU + (where transfers complete on arrival) and on CUDA (where they are + genuinely in flight). + """ + + @pytest.mark.parametrize("device", DEVICES) + def test_async_and_sync_agree_eagerly(self, device): + da = pin_array(("time", "distance")) + sync = Annotate(annotate_model(), "time", device=device, max_buffers=0) + queued = Annotate(annotate_model(), "time", device=device) + assert queued(da).equals(sync(da)) + + @pytest.mark.parametrize("device", DEVICES) + @pytest.mark.parametrize("indices", [[13], [4, 8, 12, 16, 20]]) + def test_async_and_sync_agree_chunked(self, device, indices): + da = pin_array(("time", "distance")) + sync = Annotate(annotate_model(), "time", device=device, max_buffers=0) + queued = Annotate(annotate_model(), "time", device=device) + expected = xd.concat( + list(sync.iter_chunks(xd.split(da, indices, "time"), "time")), "time" + ) + result = xd.concat( + list(queued.iter_chunks(xd.split(da, indices, "time"), "time")), "time" + ) + assert result.equals(expected) + + def test_chunk_invariance_holds_with_the_queue(self): + da = xd.testing.dummy(dims=("time", "distance"), shape=(64, 3)) + atom = Annotate(annotate_model(), "time", device="cpu") + xd.testing.assert_chunk_invariant(atom, da, {"time": 13}) + + @pytest.mark.parametrize("device", DEVICES) + def test_picks_are_identical_through_picker(self, device): + da = component_array(["SHZ", "SHN", "SHE"], n=64) + model = picker_model("original") + picks = Picker(model, resample=False, device=device)(da) + sync = Annotate(model, "time", device=device, max_buffers=0) + from xdas.atoms import Trigger + + expected = Trigger(_model_thresholds(model), dim="time")(sync(da)) + assert picks.equals(expected) + + def test_flush_drains_the_queue(self): + da = pin_array(("time", "distance")) + atom = Annotate(annotate_model(), "time", device="cpu") + chunks = [ + out + for chunk in xd.split(da, [13], "time") + for out in xd.atoms.core._aschunks(atom(chunk, chunk_dim="time")) + ] + chunks += atom.flush() + assert len(atom.inflight) == 0 + expected = Annotate(annotate_model(), "time", device="cpu", max_buffers=0)(da) + assert xd.concat(chunks, "time").equals(expected) + + +class TestAnnotateDeviceDefault: + def test_the_device_defaults_to_what_is_available(self): + atom = Annotate(annotate_model()) + expected = "cuda" if torch.cuda.is_available() else "cpu" + assert atom.device.type == expected diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py new file mode 100644 index 00000000..aa67ac88 --- /dev/null +++ b/tests/test_atoms_runs.py @@ -0,0 +1,502 @@ +""" +Run-semantics tests: seams, flush, collections and the commutation invariant. + +The testable invariant of the continuous-run model (plan §5.5) is that +splitting anywhere commutes with processing:: + + concat(atom(split(da, indices))) == atom(da) + +for *arbitrary* split points, including at discontinuities — continuous +elements carry state across, discontinuous ones reset, and tails are flushed. +""" + +import numpy as np +import pytest + +import xdas as xd +from xdas.atoms import ( + STFT, + Atom, + Decimate, + DownSample, + Filter, + Integrate, + Rechunk, + Sequential, + State, +) +from xdas.atoms.core import _aschunks +from xdas.testing import dummy + + +def collect(atom, chunks, dim="time"): + """Fold *chunks* through *atom*, drain it, and return all output chunks.""" + outs = [] + for chunk in chunks: + outs += _aschunks(atom(chunk, chunk_dim=dim)) + outs += atom.flush() + atom.reset() + return outs + + +def gappy(da, at=50, gap=10): + """Split *da* in two runs separated by a gap of *gap* samples.""" + left = da.isel(time=slice(0, at)) + right = da.isel(time=slice(at + gap, None)) + return left, right, xd.concat([left, right], "time") + + +@pytest.fixture +def da(): + # 101 samples: an awkward length that exercises the flushed tails. + return dummy(shape=(101, 5)) + + +class TestCommutation: + # Split points chosen to be awkward: single-sample chunks, unequal sizes. + splits = [[37, 61], [1, 100], [13, 14, 50, 87], 7] + + factories = [ + lambda: DownSample(3, dim="time"), + lambda: Rechunk({"time": 7}), + lambda: Filter((1.0, 10.0)), + lambda: Filter((None, 10.0), ftype="fir"), + lambda: Decimate(25.0), + lambda: STFT(0.16), # expander: 16-sample windows, 8-sample hops + lambda: Integrate(), + lambda: Sequential([Decimate(25.0), Filter((1.0, 10.0)), np.square]), + ] + + @pytest.mark.parametrize("split", splits) + @pytest.mark.parametrize("factory", factories) + def test_chunked_equals_eager(self, da, factory, split): + expected = factory()(da) + chunks = xd.split(da, split, "time") + result = xd.concat(collect(factory(), chunks), "time") + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values, atol=1e-15, rtol=1e-9) + + @pytest.mark.parametrize("factory", factories) + def test_collection_input_folds(self, da, factory): + expected = factory()(da) + collection = factory()(xd.split(da, [13, 50, 87], "time")) + assert isinstance(collection, xd.DataSequence) + result = xd.concat(list(collection), "time") + assert np.allclose(result.values, expected.values, atol=1e-15, rtol=1e-9) + + +class TestFlush: + def test_default_is_noop(self): + assert Atom().flush() == [] + + def test_downsample_emits_tail(self, da): + # 101 % 3 != 0: the strided remainder starts on an output sample. + atom = DownSample(3, dim="time") + expected = atom(da) + assert expected.sizes["time"] == 34 + outs = collect(atom, xd.split(da, 4, "time")) + assert xd.concat(outs, "time").equals(expected) + + def test_flush_empties_the_buffer(self, da): + atom = DownSample(3, dim="time") + atom(da.isel(time=slice(0, 50)), chunk_dim="time") + assert len(atom.flush()) == 1 + assert atom.flush() == [] + + def test_eager_call_is_complete(self, da): + # Single-call prototyping returns the full output: nothing left over. + atom = DownSample(3, dim="time") + atom(da) + assert atom.flush() == [] + + +class TestSeams: + def test_gap_resets_state(self, da): + # Integrate carries a cumulative offset: carrying it across the gap + # would corrupt the second run. + left, right, _ = gappy(da) + atom = Integrate() + expected = [Integrate()(left), Integrate()(right)] + result = collect(atom, [left, right]) + assert len(result) == 2 + for out, exp in zip(result, expected): + assert out.equals(exp) + + def test_gap_flushes_the_tail(self, da): + # The seam call returns the flushed tail of the old run before the + # fresh output of the new one. + left, right, _ = gappy(da, at=50) # 50 % 3 != 0: pending remainder + atom = DownSample(3, dim="time") + outs = _aschunks(atom(left, chunk_dim="time")) + outs += _aschunks(atom(right, chunk_dim="time")) + outs += atom.flush() + expected = xd.concat( + [DownSample(3, dim="time")(run) for run in (left, right)], "time" + ) + assert xd.concat(outs, "time").equals(expected) + + def test_rate_change_redesigns(self): + # Same filter atom, stream whose rate halves mid-way: each run must be + # filtered with coefficients designed for its own rate. + left = dummy(shape=(100, 5), step=(0.01, 10.0)) + start = left["time"].end + 5 * left["time"].sampling_interval + right = dummy(shape=(50, 5), step=(0.02, 10.0)) + right["time"] += start - right["time"].start + atom = Filter((1.0, 10.0)) + result = collect(atom, [left, right]) + expected = [Filter((1.0, 10.0))(left), Filter((1.0, 10.0))(right)] + assert len(result) == 2 + for out, exp in zip(result, expected): + assert np.allclose(out.values, exp.values) + + def test_overlap_raises(self, da): + atom = Filter((1.0, 10.0)) + atom(da.isel(time=slice(0, 50)), chunk_dim="time") + with pytest.raises(ValueError, match="overlap"): + atom(da.isel(time=slice(40, 80)), chunk_dim="time") + + def test_on_discontinuity_raise(self, da): + left, right, _ = gappy(da) + atom = Filter((1.0, 10.0)) + atom.on_discontinuity = "raise" + atom(left, chunk_dim="time") + with pytest.raises(ValueError, match="discontinuous"): + atom(right, chunk_dim="time") + + def test_invalid_policy(self, da): + left, right, _ = gappy(da) + atom = Filter((1.0, 10.0)) + atom.on_discontinuity = "ignore" + atom(left, chunk_dim="time") + with pytest.raises(ValueError, match="on_discontinuity"): + atom(right, chunk_dim="time") + + @pytest.mark.filterwarnings("ignore::FutureWarning") + def test_irregular_coordinate_raises(self): + coords = { + "time": {"tie_indices": [0, 49], "tie_values": [0.0, 49.0]}, + "distance": {"tie_indices": [0, 4], "tie_values": [0.0, 40.0]}, + } + da = xd.DataArray(np.random.randn(50, 5), coords) + atom = Integrate() + atom(da, chunk_dim="time") + with pytest.raises(ValueError, match="regular"): + atom(da, chunk_dim="time") + + +class TestEagerRuns: + def test_gappy_input_splits_into_runs(self, da): + # Filters never cross discontinuities: eager on a gappy record equals + # per-run processing, re-joined with the gap kept in the coords. + left, right, joined = gappy(da) + atom = Filter((1.0, 10.0)) + result = atom(joined) + assert isinstance(result, xd.DataArray) + expected = xd.concat( + [Filter((1.0, 10.0))(run) for run in (left, right)], "time" + ) + assert np.allclose(result.values, expected.values) + assert result["time"].get_split_indices().size == 1 + + def test_gapless_input_unchanged(self, da): + assert xd.filter(da, (1.0, 10.0)).sizes["time"] == da.sizes["time"] + + def test_chunked_internal_gap_splits(self, da): + # The ingress invariant: an internally gappy chunk is split into runs + # before the seam-aware call, so state never crosses a gap. + left, right, joined = gappy(da) + atom = Integrate() + result = _aschunks(atom(joined, chunk_dim="time")) + expected = [Integrate()(left), Integrate()(right)] + assert len(result) == 2 + for out, exp in zip(result, expected): + assert out.equals(exp) + + +class TestSeamEdgeCases: + def test_empty_chunk_is_skipped(self, da): + atom = Integrate() + atom(da.isel(time=slice(0, 50)), chunk_dim="time") + assert _aschunks(atom(da.isel(time=slice(0, 0)), chunk_dim="time")) == [] + + def test_single_sample_chunk_adopts_the_stream_rate(self, da): + # a one-sample chunk of a sampled coordinate declares no rate of its + # own: continuous with the stream, it inherits the seam's delta + # rather than forgetting it. + import scipy.signal as sp + + from xdas.atoms import Partial + from xdas.signal import sosfilt + + sampled = dummy(shape=(52, 5), ctype="sampled") + sos = sp.iirfilter(4, 0.1, btype="lowpass", output="sos") + atom = Partial(sosfilt, sos, ..., dim="time", zi=...) + expected = Partial(sosfilt, sos, ..., dim="time", zi=...)(sampled) + outs = collect(atom, xd.split(sampled, [50, 51], "time")) + result = xd.concat(outs, "time") + assert np.allclose(result.values, expected.values) + + def test_single_sample_chunk_does_not_hide_the_next_seam(self): + # the first chunk holds one sample, so the seam it leaves behind knows + # no rate: the gap that follows is judged on the incoming chunk's. + sampled = dummy(shape=(52, 5), ctype="sampled") + left, right = sampled.isel(time=slice(0, 1)), sampled.isel(time=slice(11, None)) + atom = DownSample(2, dim="time") + atom.on_discontinuity = "raise" + atom(left, chunk_dim="time") + with pytest.raises(ValueError, match="discontinuous"): + atom(right, chunk_dim="time") + + def test_stream_of_single_samples_has_nothing_to_judge(self): + sampled = dummy(shape=(3, 5), ctype="sampled") + atom = DownSample(2, dim="time") + outs = collect(atom, xd.split(sampled, [1, 2], "time")) + expected = DownSample(2, dim="time")(sampled) + assert np.allclose(xd.concat(outs, "time").values, expected.values) + + def test_alias_resolves_on_a_dimension_without_a_coordinate(self): + # the alias has to resolve against the dimensions, not the + # coordinates: a dimension carrying none is still chunkable, and the + # seam state has to be allocated for it. + import scipy.signal as sp + + from xdas.atoms import LFilter + + values = np.random.default_rng(0).normal(size=(24, 12)) + coords = {"time": {"tie_indices": [0, 23], "tie_values": [0.0, 2.3]}} + bare = xd.DataArray(values, coords, ("time", "distance")) + b, a = sp.butter(2, 0.4) + expected = LFilter(b, a, dim="last")(bare) + outs = collect( + LFilter(b, a, dim="last"), xd.split(bare, 4, "distance"), "distance" + ) + result = np.concatenate([out.values for out in outs], axis=1) + assert np.allclose(result, expected.values) + + def test_first_alias_resolves_on_eager_calls(self, da): + result = DownSample(2, dim="first")(da) + expected = DownSample(2, dim="time")(da) + assert result.equals(expected) + + def test_scalar_coordinate_is_not_judged(self, da): + # chunked along a dimension the chunk only knows as a scalar + # coordinate: there is no axis to judge a seam on. + chunk = da.isel(time=slice(0, 10)).mean("time") + chunk = chunk.assign_coords(time=da["time"][0].values) + atom = Integrate(dim="distance") + atom(chunk, chunk_dim="time") + atom(chunk, chunk_dim="time") # no seam judgement, no raise + + +class TestFoldEdgeCases: + def test_dimensionless_atom_maps_eagerly(self, da): + from xdas.atoms import Partial + + collection = xd.DataCollection([da, da]) + result = Partial(np.square)(collection) + assert len(result) == 2 + for out in result: + assert np.allclose(out.values, np.square(da.values)) + + def test_chunked_sequence_folds_through(self, da): + atom = Integrate() + elements = xd.DataCollection(xd.split(da, 3, "time")) + outs = _aschunks(atom(elements, chunk_dim="time")) + outs += atom.flush() + expected = Integrate()(da) + assert np.allclose(xd.concat(outs, "time").values, expected.values) + + def test_join_falls_back_on_unconcatenatable_chunks(self, da): + from xdas.atoms import Atom + + other = dummy(shape=(10, 3)) + result = Atom()._join([da, other], "time") + assert isinstance(result, xd.DataSequence) + + +class TestSplitAnnouncement: + def test_eager_call_announces_the_split_count(self, da): + left = da.isel(time=slice(0, 30)) + mid = da.isel(time=slice(40, 60)) + right = da.isel(time=slice(70, None)) + joined = xd.concat([left, mid, right], "time") + with pytest.warns(UserWarning, match="2 discontinuities along 'time'"): + xd.filter(joined, (1.0, 10.0)) + + def test_singular_wording(self, da): + _, _, joined = gappy(da) + with pytest.warns(UserWarning, match="1 discontinuity along 'time'"): + xd.filter(joined, (1.0, 10.0)) + + @pytest.mark.filterwarnings("error::UserWarning") + def test_gapless_input_is_silent(self, da): + xd.filter(da, (1.0, 10.0)) + + def test_every_leaf_of_a_collection_reports(self, da): + # The message names the source by its start, so two leaves with the + # same gap count are not deduplicated into one warning. + _, _, joined = gappy(da) + other = joined.copy() + other["time"] = other["time"] + np.timedelta64(1, "h") + collection = xd.DataCollection({"das1": joined, "das2": other}) + with pytest.warns(UserWarning) as record: + xd.filter(collection, (1.0, 10.0)) + messages = [str(w.message) for w in record if "discontinuit" in str(w.message)] + assert len(messages) == 2 + assert messages[0] != messages[1] + + +class TestCollections: + def test_mapping_maps_over_leaves(self, da): + collection = xd.DataCollection({"das1": da, "das2": da}) + result = xd.filter(collection, (1.0, 10.0)) + expected = xd.filter(da, (1.0, 10.0)) + assert result["das1"].equals(expected) + assert result["das2"].equals(expected) + + def test_chunked_mapping_raises(self, da): + collection = xd.DataCollection({"das1": da}) + with pytest.raises(NotImplementedError): + xd.filter(..., (1.0, 10.0))(collection, chunk_dim="time") + + def test_sequence_with_gap(self, da): + # Continuous elements carry state across, discontinuous ones reset. + left, right, _ = gappy(da) + elements = xd.split(left, 2, "time") + xd.split(right, 2, "time") + atom = Integrate() + collection = atom(xd.DataCollection(elements)) + result = xd.concat(list(collection), "time") + expected = xd.concat([Integrate()(left), Integrate()(right)], "time") + assert np.allclose(result.values, expected.values) + + +class TestIterChunks: + def test_matches_eager(self, da): + pipeline = Sequential([Decimate(25.0), Filter((1.0, 10.0))]) + outs = list(pipeline.iter_chunks(xd.split(da, 5, "time"))) + expected = Sequential([Decimate(25.0), Filter((1.0, 10.0))])(da) + result = xd.concat(outs, "time") + assert np.allclose(result.values, expected.values, atol=1e-15, rtol=1e-9) + + def test_explicit_chunk_dim(self, da): + atom = DownSample(2, dim="distance") + outs = list(atom.iter_chunks(xd.split(da, 2, "time"), chunk_dim="time")) + assert xd.concat(outs, "time").equals(DownSample(2, dim="distance")(da)) + + def test_resets_at_the_end(self, da): + atom = Integrate() + list(atom.iter_chunks(xd.split(da, 3, "time"))) + assert not atom.initialized + + +class StreamMean(Atom): + """Test reduction: accumulate per chunk, emit the single result at flush.""" + + def __init__(self, dim="time"): + super().__init__() + self.dim = dim + self.numerator = State(...) + self.denominator = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + if chunk_dim == self.dim: + self.numerator = State(0.0 * da.sum(self.dim)) + self.denominator = State(0) + else: + self.numerator = State(None) + self.denominator = State(None) + + def call(self, da, **flags): + if self.numerator is None: + return da.mean(self.dim) + self.numerator = State(self.numerator + da.sum(self.dim)) + self.denominator = State(self.denominator + da.sizes[self.dim]) + return None + + def flush(self): + if not isinstance(self.numerator, xd.DataArray): + return [] + out = self.numerator / self.denominator + self.numerator = State(0.0 * self.numerator) + self.denominator = State(0) + return [out] + + +class TestSTFTRuns: + def test_frames_never_span_gaps(self, da): + left, right, both = gappy(da) + streamed = collect(STFT(0.16), [left, right]) + expected = [STFT(0.16)(left), STFT(0.16)(right)] + assert len(streamed) == 2 + for out, exp in zip(streamed, expected): + assert out.coords.equals(exp.coords) + assert np.allclose(out.values, exp.values) + # eager on the gappy record splits into the same per-run frames + eager = STFT(0.16)(both) + assert np.allclose(xd.concat(streamed, "time").values, eager.values) + + def test_short_run_emits_nothing_when_streaming(self, da): + # 64-sample windows never fit in a 50-sample run: the buffered tail + # is dropped at flush, nothing is emitted and nothing raises. + outs = collect(STFT(0.64), [da.isel(time=slice(0, 50))]) + assert outs == [] + + def test_eager_short_record_raises(self, da): + with pytest.raises(ValueError, match="shorter"): + STFT(1.28)(da) + + +class TestReduction: + def test_streaming_equals_eager(self, da): + expected = StreamMean()(da) + outs = collect(StreamMean(), xd.split(da, [37, 61], "time")) + assert len(outs) == 1 + assert np.allclose(outs[0].values, expected.values) + + def test_call_returns_no_chunk(self, da): + atom = StreamMean() + out = atom(da.isel(time=slice(0, 50)), chunk_dim="time") + assert list(out) == [] + + +class TestRechunk: + def test_sizes(self, da): + atom = Rechunk({"time": 30}) + outs = collect(atom, xd.split(da, 7, "time")) + assert [out.sizes["time"] for out in outs] == [30, 30, 30, 11] + + def test_never_merges_across_gaps(self, da): + left, right, _ = gappy(da) + atom = Rechunk({"time": 40}) + outs = collect(atom, [left, right]) + # 50-sample run then 41-sample run: the partial buffer is flushed at + # the seam instead of being merged with the next run. + assert [out.sizes["time"] for out in outs] == [40, 10, 40, 1] + for out in outs: + assert out["time"].get_split_indices().size == 0 + + def test_eager_is_identity(self, da): + assert xd.rechunk(da, {"time": 30}).equals(da) + + def test_twin_seed(self): + atom = xd.rechunk(..., {"time": 30}) + assert isinstance(atom, Rechunk) + + def test_invalid_chunks(self): + with pytest.raises(TypeError): + Rechunk({"time": 10, "distance": 2}) + with pytest.raises(ValueError): + Rechunk({"time": 0}) + + +class TestProcess: + def test_end_of_stream_flush(self, da, tmp_path): + # The DownSample tail-drop bug: process() must drain the atom. + import xdas.processing as xp + + atom = DownSample(3, dim="time") + expected = DownSample(3, dim="time")(da) + loader = xp.DataArrayLoader(da, {"time": 25}) + writer = xp.DataArrayWriter(tmp_path) + result = xp.process(atom, loader, writer) + assert result.equals(expected) diff --git a/tests/test_atoms_tasks.py b/tests/test_atoms_tasks.py new file mode 100644 index 00000000..becd2f09 --- /dev/null +++ b/tests/test_atoms_tasks.py @@ -0,0 +1,476 @@ +import inspect + +import numpy as np +import numpy.testing as npt +import pytest + +import xdas as xd +import xdas.signal as xs +import xdas.spectral +from xdas.atoms import ( + STFT, + Decimate, + Differentiate, + Filter, + Integrate, + Partial, + ResamplePoly, + Sequential, +) +from xdas.synthetics import wavelet_wavefronts +from xdas.testing import dummy + + +def through_chunks(atom, da, nchunk=6, dim="time"): + chunks = xd.split(da, nchunk, dim) + return xd.concat([atom(chunk, chunk_dim=dim) for chunk in chunks], dim) + + +class TestFunctionForms: + def test_seed_returns_atom(self): + atom = xd.filter(..., (1.0, 10.0)) + assert isinstance(atom, Filter) + assert atom.freq == (1.0, 10.0) + + def test_seed_returns_partial_for_functions(self): + atom = xd.taper(...) + assert isinstance(atom, Partial) + + def test_atom_input_composes(self): + head = xd.decimate(..., 25.0) + pipeline = xd.filter(head, (1.0, 10.0)) + assert isinstance(pipeline, Sequential) + assert len(pipeline) == 2 + assert not isinstance(head, Sequential) + + def test_signature(self): + parameters = list(inspect.signature(xd.decimate).parameters) + assert parameters == ["da", "target", "window", "dim"] + + def test_docstring(self): + assert "Decimate" in xd.decimate.__doc__ + assert "target" in xd.decimate.__doc__ + + def test_names(self): + assert xd.filter.__name__ == "filter" + assert xd.sliding_mean_removal.__name__ == "sliding_mean_removal" + + def test_time_defaults(self): + assert xd.filter(..., (1.0, 10.0)).dim == "time" + assert xd.decimate(..., 25.0).dim == "time" + + def test_seed_idiom_matches_eager(self): + da = wavelet_wavefronts() + + def workflow(da): + da = xd.decimate(da, 25.0) + da = xd.filter(da, (1.0, 10.0)) + return np.square(da) + + preview = workflow(da) + result = workflow(...)(da) + assert np.allclose(result.values, preview.values) + + +class TestFilter: + def test_bandpass_matches_signal(self): + da = wavelet_wavefronts() + result = xd.filter(da, (1.0, 10.0)) + expected = xs.filter(da, (1.0, 10.0), "bandpass", corners=4, dim="time") + assert result.equals(expected) + + def test_lowpass_matches_signal(self): + da = wavelet_wavefronts() + result = xd.filter(da, (None, 10.0)) + expected = xs.filter(da, 10.0, "lowpass", corners=4, dim="time") + assert result.equals(expected) + + def test_highpass_matches_signal(self): + da = wavelet_wavefronts() + result = xd.filter(da, (1.0, None)) + expected = xs.filter(da, 1.0, "highpass", corners=4, dim="time") + assert result.equals(expected) + + def test_zerophase_matches_signal(self): + da = wavelet_wavefronts() + result = xd.filter(da, (1.0, 10.0), zerophase=True) + expected = xs.filter( + da, (1.0, 10.0), "bandpass", corners=4, zerophase=True, dim="time" + ) + assert result.equals(expected) + + def test_iir_chunked_equals_monolithic(self): + da = wavelet_wavefronts() + atom = Filter((1.0, 10.0)) + expected = atom(da) + result = through_chunks(atom, da) + assert result.equals(expected) + + def test_fir_chunked_equals_monolithic(self): + da = wavelet_wavefronts() + atom = Filter((None, 10.0), ftype="fir") + expected = atom(da) + result = through_chunks(atom, da) + assert np.allclose(result.values, expected.values, atol=1e-16, rtol=1e-11) + assert result.coords.equals(expected.coords) + + def test_fir_compensates_lag(self): + # The FIR filter is linear-phase with the group delay compensated on + # the coordinate: a lowpassed wavelet must not shift in time. + da = wavelet_wavefronts() + result = xd.filter(da, (None, 10.0), ftype="fir", transition=5.0) + reference = xd.filter(da, (None, 10.0), zerophase=True) + assert result.sizes == da.sizes + trace = result.isel(distance=100) + reference_trace = reference.isel(distance=100) + peak = trace["time"].values[int(np.argmax(trace.values))] + reference_peak = reference_trace["time"].values[ + int(np.argmax(reference_trace.values)) + ] + assert abs(peak - reference_peak) <= np.timedelta64(40, "ms") + + def test_zerophase_iir_chunked_raises(self): + da = wavelet_wavefronts() + atom = Filter((1.0, 10.0), zerophase=True) + chunk, *_ = xd.split(da, 6, "time") + with pytest.raises(ValueError, match="whole record"): + atom(chunk, chunk_dim="time") + + def test_zerophase_iir_chunked_along_other_dim_passes(self): + da = wavelet_wavefronts() + atom = Filter((1.0, 10.0), zerophase=True) + expected = atom(da) + result = through_chunks(atom, da, dim="distance") + assert result.equals(expected) + + def test_distance_dim(self): + da = wavelet_wavefronts() + result = xd.filter(da, (None, 0.005), dim="distance") + expected = xs.filter(da, 0.005, "lowpass", corners=4, dim="distance") + assert result.equals(expected) + + def test_scalar_freq_raises(self): + with pytest.raises(TypeError, match="pair of corner frequencies"): + Filter(10.0) + + def test_open_both_ends_raises(self): + with pytest.raises(ValueError, match="at least one corner"): + Filter((None, None)) + + def test_invalid_ftype_raises(self): + with pytest.raises(ValueError, match="ftype"): + Filter((1.0, 10.0), ftype="cheby") + + +class TestFilterState: + def test_initialize_from_state_is_noop_for_iir(self): + # Only the FIR path has a design to rebuild from restored state. + atom = Filter((1.0, 10.0)) + assert atom.initialize_from_state() is None + assert not atom.filter.initialized + + +class TestDecimate: + def test_matches_resample_poly(self): + # For an integer factor Decimate shares its design with ResamplePoly. + da = wavelet_wavefronts() + result = xd.decimate(da, 25.0) + expected = ResamplePoly(25.0, dim="time")(da) + assert result.equals(expected) + + def test_chunked_equals_monolithic(self): + # Chunk sizes must be multiples of the factor: draining the trailing + # remainder needs the flush() lifecycle. + da = wavelet_wavefronts() + atom = Decimate(25.0) + expected = atom(da) + result = through_chunks(atom, da) + assert np.allclose(result.values, expected.values, atol=1e-16, rtol=1e-11) + assert result.coords.equals(expected.coords) + + def test_non_integer_factor_raises(self): + da = wavelet_wavefronts() + with pytest.raises(ValueError, match="integer multiple"): + xd.decimate(da, 30.0) + + def test_upsampling_raises(self): + da = wavelet_wavefronts() + with pytest.raises(ValueError, match="integer multiple"): + xd.decimate(da, 60.0) + + def test_factor_one_is_identity(self): + da = wavelet_wavefronts() # already at 50 Hz + assert xd.decimate(da, 50.0).equals(da) + + +class TestResample: + def test_matches_resample_poly(self): + da = wavelet_wavefronts() + result = xd.resample(da, 20.0) + expected = ResamplePoly(20.0, dim="time")(da) + assert result.equals(expected) + + +class TestIntegrate: + def test_matches_signal(self): + da = wavelet_wavefronts() + result = xd.integrate(da) + expected = xs.integrate(da, dim="time") + assert result.equals(expected) + + def test_chunked_equals_monolithic(self): + da = wavelet_wavefronts() + atom = Integrate() + expected = atom(da) + result = through_chunks(atom, da) + assert np.allclose(result.values, expected.values) + assert result.coords.equals(expected.coords) + + def test_midpoints(self): + # Midpoints on the distance dim: xs.integrate cannot shift datetime + # coordinates by a float half-step, so time is not testable here. + da = wavelet_wavefronts() + result = xd.integrate(da, midpoints=True, dim="distance") + expected = xs.integrate(da, midpoints=True, dim="distance") + assert result.equals(expected) + + +class TestDifferentiate: + def test_matches_signal(self): + da = wavelet_wavefronts() + result = xd.differentiate(da) + expected = xs.differentiate(da, dim="time") + assert result.equals(expected) + + def test_chunked_equals_monolithic(self): + da = wavelet_wavefronts() + atom = Differentiate() + expected = atom(da) + result = through_chunks(atom, da) + assert np.allclose(result.values, expected.values) + assert result.coords.equals(expected.coords) + + +class TestWholeRecordFunctions: + def test_detrend_matches_signal(self): + da = wavelet_wavefronts() + assert xd.detrend(da).equals(xs.detrend(da, dim="time")) + + def test_taper_matches_signal(self): + da = wavelet_wavefronts() + assert xd.taper(da).equals(xs.taper(da, dim="time")) + + def test_hilbert_matches_signal(self): + da = wavelet_wavefronts() + assert xd.hilbert(da).equals(xs.hilbert(da, dim="time")) + + def test_sliding_mean_removal_matches_signal(self): + da = wavelet_wavefronts() + assert xd.sliding_mean_removal(da, 1.0).equals( + xs.sliding_mean_removal(da, 1.0, dim="time") + ) + + def test_medfilt_physical_units(self): + da = wavelet_wavefronts() + dt = xd.get_sampling_interval(da, "time") + dx = xd.get_sampling_interval(da, "distance") + result = xd.medfilt(da, {"time": 7 * dt, "distance": 5 * dx}) + expected = xs.medfilt(da, {"time": 7, "distance": 5}) + assert result.equals(expected) + + def test_chunked_raises(self): + da = wavelet_wavefronts() + chunk, *_ = xd.split(da, 6, "time") + for atom in [ + xd.detrend(...), + xd.taper(...), + xd.hilbert(...), + xd.sliding_mean_removal(..., 1.0), + xd.medfilt(..., {"time": 0.1}), + ]: + with pytest.raises(ValueError, match="whole record"): + atom(chunk, chunk_dim="time") + + def test_chunked_along_other_dim_passes(self): + da = wavelet_wavefronts() + atom = xd.taper(...) + expected = atom(da) + result = through_chunks(atom, da, dim="distance") + assert result.equals(expected) + + +class TestSTFT: + # dummy is sampled at 100 Hz: 0.32 s windows are 32 samples (a fast FFT + # size, so the target is not snapped) and 0.16 s hops are 16 samples. + + def test_matches_legacy_spectral(self): + da = dummy(shape=(200, 5)) + result = xd.stft(da, 0.32, hop=0.16) + expected = xdas.spectral.stft( + da, window="hann", nperseg=32, noverlap=16, dim={"time": "frequency"} + ) + assert np.allclose(result.values, expected.values) + assert np.array_equal(result["time"].values, expected["time"].values) + assert np.allclose(result["frequency"].values, expected["frequency"].values) + assert result["distance"].equals(expected["distance"]) + + def test_default_hop_is_half_window(self): + da = dummy(shape=(200, 5)) + result = xd.stft(da, 0.32) + expected = xd.stft(da, 0.32, hop=0.16) + assert np.allclose(result.values, expected.values) + + def test_wlen_snaps_to_fast_length(self): + # 1.27 s at 100 Hz is 127 samples, a prime: the next fast size is 128. + da = dummy(shape=(400, 5)) + result = xd.stft(da, 1.27) + assert result.sizes["frequency"] == 128 // 2 + 1 + + def test_psd_scaling_matches_legacy(self): + da = dummy(shape=(200, 5)) + result = xd.stft(da, 0.32, hop=0.16, scaling="psd") + expected = xdas.spectral.stft( + da, + window="hann", + nperseg=32, + noverlap=16, + scaling="psd", + dim={"time": "frequency"}, + ) + assert np.allclose(result.values, expected.values) + + def test_nfft_zero_padding(self): + da = dummy(shape=(200, 5)) + result = xd.stft(da, 0.32, hop=0.16, nfft=64) + expected = xdas.spectral.stft( + da, + window="hann", + nperseg=32, + noverlap=16, + nfft=64, + dim={"time": "frequency"}, + ) + assert result.sizes["frequency"] == 64 // 2 + 1 + assert np.allclose(result.values, expected.values) + assert np.allclose(result["frequency"].values, expected["frequency"].values) + + def test_nfft_smaller_than_window_raises(self): + da = dummy(shape=(200, 5)) + with pytest.raises(ValueError, match="nfft"): + xd.stft(da, 0.32, nfft=16) + + def test_complex_input_two_sided(self): + da = dummy(shape=(200, 5), dtype=complex) + result = xd.stft(da, 0.32, hop=0.16) + expected = xdas.spectral.stft( + da, + window="hann", + nperseg=32, + noverlap=16, + return_onesided=False, + dim={"time": "frequency"}, + ) + assert result.sizes["frequency"] == 32 + assert np.allclose(result.values, expected.values) + assert np.allclose(result["frequency"].values, expected["frequency"].values) + + def test_invalid_parameters(self): + with pytest.raises(ValueError, match="wlen"): + STFT(0.0) + with pytest.raises(ValueError, match="hop"): + STFT(1.0, hop=2.0) + with pytest.raises(ValueError, match="scaling"): + STFT(1.0, scaling="power") + + def test_record_shorter_than_window_raises(self): + da = dummy(shape=(50, 5)) + with pytest.raises(ValueError, match="shorter"): + xd.stft(da, 1.0) + + def test_chunked_along_other_dim(self): + da = dummy(shape=(200, 5)) + expected = xd.stft(da, 0.32, hop=0.16) + result = through_chunks(STFT(0.32, hop=0.16), da, 2, "distance") + assert np.allclose(result.values, expected.values) + + def test_non_dimensional_coords(self): + da = dummy(shape=(200, 5)) + da["latitude"] = ("distance", np.arange(5.0)) + da["quality"] = ("time", np.arange(200.0)) + result = xd.stft(da, 0.32, hop=0.16) + # coords along other dimensions are kept; those along the transformed + # dimension are dropped (TODO in STFT._transform, as in spectral.stft) + assert result["latitude"].equals(da["latitude"]) + assert "quality" not in result.coords + + def test_function_form_seed(self): + atom = xd.stft(..., 1.0) + assert isinstance(atom, STFT) + assert atom.dim == "time" + + def test_pipeline_chunk_invariant_over_cuts_and_gaps(self): + da = dummy(shape=(400, 5)) + pipeline = xd.filter(..., (None, 20.0)) >> xd.stft(..., 0.32, hop=0.16) + xd.testing.assert_chunk_invariant( + pipeline, da, {"time": 100}, cuts=2, gaps=2, atol=1e-12 + ) + + +class TestLabelsFollowTheSamples: + """ + A resampled dimension carries its other coordinates onto the output grid. + + A DAS channel is named by a non-dimensional ``station`` coordinate + attached to ``distance``, and picking answers with those names: a stage + that changes the number of samples has to say what became of them, or the + output is labelled with the wrong lanes. They follow the *samples* — + output ``k`` is drawn from input ``k * down / up`` — so unlike the + dimension coordinate they carry no group-delay shift, and the mapping + cannot depend on the chunking. + """ + + def labelled(self, n=120, dim="distance", step=(0.01, 10.0)): + da = dummy(dims=("time", "distance"), shape=(200, n), step=step) + labels = np.array([f"S{index:04d}" for index in range(da.sizes[dim])]) + return da.assign_coords(station=(dim, labels)) + + def test_decimation_subsamples_them(self): + da = self.labelled() + result = xd.decimate(da, target=1 / 20.0, dim="distance") + assert result.sizes["distance"] == 60 + npt.assert_array_equal(result["station"].values, da["station"].values[::2]) + + def test_rational_resampling_lands_on_the_output_grid(self): + da = self.labelled(n=8) + result = xd.resample(da, 1 / 25.0, dim="distance") # up=2, down=5 + assert result.sizes["distance"] == len(result["station"].values) + npt.assert_array_equal( + result["station"].values, ["S0000", "S0002", "S0005", "S0007"] + ) + + def test_upsampling_repeats_them(self): + from xdas.atoms import UpSample + + da = self.labelled(n=4) + result = UpSample(3, dim="distance")(da) + assert result.sizes["distance"] == 12 + npt.assert_array_equal( + result["station"].values[:6], ["S0000"] * 3 + ["S0001"] * 3 + ) + + def test_the_labels_do_not_depend_on_the_chunking(self): + # chunked along the very dimension being decimated: every chunk must + # label its output with the same input samples the eager call does. + da = dummy(dims=("time", "distance"), shape=(120, 3)) + labels = np.array([f"T{index:04d}" for index in range(120)]) + da = da.assign_coords(label=("time", labels)) + eager = xd.decimate(..., target=25.0, dim="time")(da) + atom = xd.decimate(..., target=25.0, dim="time") + chunks = list(atom.iter_chunks(xd.split(da, 7, "time"), "time")) + streamed = np.concatenate([chunk["label"].values for chunk in chunks]) + npt.assert_array_equal(streamed, eager["label"].values) + + def test_an_untouched_dimension_keeps_its_labels(self): + da = self.labelled() + result = xd.decimate(da, target=25.0, dim="time") + npt.assert_array_equal(result["station"].values, da["station"].values) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 00000000..d695a938 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,53 @@ +import os + +import xdas.config as xc + + +class TestTotalMemory: + def test_reads_the_machine(self): + # Whatever the machine, the answer must be a plausible byte count. + total = xc.total_memory() + assert isinstance(total, int) + assert total > 2**28 # no machine xdas runs on has under 256 MiB + + def test_a_cgroup_limit_wins_over_the_machine(self, monkeypatch, tmp_path): + limit = tmp_path / "memory.max" + limit.write_text("1073741824\n") # 1 GiB, far under any real machine + monkeypatch.setattr(xc, "CGROUP_LIMITS", (str(limit),)) + assert xc.total_memory() == 2**30 + + def test_an_unlimited_cgroup_loses_to_the_machine(self, monkeypatch, tmp_path): + # cgroup v2 spells an absent limit "max"; v1 writes a sentinel larger + # than any machine. Neither may become the answer. + for content in ("max\n", "9223372036854771712\n"): + limit = tmp_path / "memory.max" + limit.write_text(content) + monkeypatch.setattr(xc, "CGROUP_LIMITS", (str(limit),)) + physical = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") + assert xc.total_memory() == physical + + def test_missing_files_are_skipped(self, monkeypatch, tmp_path): + monkeypatch.setattr(xc, "CGROUP_LIMITS", (str(tmp_path / "absent"),)) + physical = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") + assert xc.total_memory() == physical + + def test_falls_back_when_nothing_can_be_read(self, monkeypatch): + def unavailable(name): + raise ValueError(name) + + monkeypatch.setattr(xc, "CGROUP_LIMITS", ()) + monkeypatch.setattr(xc.os, "sysconf", unavailable) + assert xc.total_memory() == xc.FALLBACK_MEMORY + + +class TestMemoryLimit: + def test_default_is_a_share_of_the_machine(self): + assert xc.get("memory_limit") == int(xc.MEMORY_FRACTION * xc.total_memory()) + + def test_set_overrides_it(self): + previous = xc.get("memory_limit") + try: + xc.set("memory_limit", 123) + assert xc.get("memory_limit") == 123 + finally: + xc.set("memory_limit", previous) diff --git a/tests/test_datacollection.py b/tests/test_datacollection.py index c34d6bec..7e09d287 100644 --- a/tests/test_datacollection.py +++ b/tests/test_datacollection.py @@ -1,4 +1,6 @@ import h5py +import numpy as np +import pandas as pd import pytest import xdas as xd @@ -511,3 +513,202 @@ def test_sequence_sel_all_elements_become_empty(self): dc = xd.DataCollection([da_near, da_far], "instrument") result = dc.sel(distance=slice(-100, -1)) assert len(result) == 0 + + +class TestDataFrameLeaves: + def test_a_dataframe_stays_a_dataframe(self): + # the rebuild used to wrap every non-DataArray leaf in DataArray(...), + # silently destroying a table. + df = pd.DataFrame({"time": [1.0, 2.0], "value": [0.5, 0.9]}) + dc = xd.DataCollection({"ST01": df, "ST02": df.copy()}, "station") + assert isinstance(dc["ST01"], pd.DataFrame) + pd.testing.assert_frame_equal(dc["ST01"], df) + + def test_repr_shows_the_table(self): + df = pd.DataFrame({"time": [1.0, 2.0], "value": [0.5, 0.9]}) + dc = xd.DataCollection( + {"das": xd.DataCollection({"ST01": df}, "station")}, "instrument" + ) + text = repr(dc) + assert "ST01" in text + assert "das" in text + + +class TestMergeCollectionResults: + """ + W8: the walk labels each leaf's result with its own tree path as the + result is produced, then folds the results through the atom's `merge` + hook on the way back up. + """ + + thresh = {"P": 0.5, "S": 0.5} + + def cft(self, t0=0.0, quiet=False, **scalars): + """A characteristic function peaking on P and S at ``t0 + 1``.""" + lane = [0.0, 0.0, 0.0] if quiet else [0.0, 0.8, 0.0] + da = xd.DataArray( + data=[[0.0, 0.0, 0.0], lane, lane], + coords={ + "phase": ["N", "P", "S"], + "time": { + "tie_indices": [0, 2], + "tie_values": [t0, t0 + 2.0], + "sampling_interval": 1.0, + }, + }, + ) + for name, value in scalars.items(): + da = da.assign_coords(**{name: value}) + return da + + def tree(self): + """A ``network / station / location`` collection of pick-able leaves.""" + return xd.DataCollection( + { + "IA": xd.DataCollection( + { + "DBNFM": xd.DataCollection({"--": self.cft()}, "location"), + "LBFI": xd.DataCollection({"00": self.cft(10.0)}, "location"), + }, + "station", + ) + }, + "network", + ) + + def test_leaves_carry_their_tree_path_and_merge_into_one_table(self): + result = xd.trigger(self.tree(), thresh=self.thresh) + assert isinstance(result, pd.DataFrame) + # identity leads, then the dimension coordinates, then the value + assert list(result.columns) == [ + "network", + "station", + "location", + "phase", + "time", + "value", + ] + assert list(result["network"]) == ["IA"] * 4 + assert list(result["station"]) == ["DBNFM", "DBNFM", "LBFI", "LBFI"] + assert list(result["location"]) == ["--", "--", "00", "00"] + assert list(result["time"]) == [1.0, 1.0, 11.0, 11.0] + assert list(result.index) == [0, 1, 2, 3] + + def test_annotation_happens_at_production_time(self): + # every leaf of the un-merged tree already carries the full path, so a + # streaming walk can hand a leaf straight to a sink and keep its + # identity + tree = xd.trigger(..., thresh=self.thresh)(self.tree(), merge=False) + leaf = tree["IA"]["DBNFM"]["--"] + assert isinstance(leaf, pd.DataFrame) + assert list(leaf.columns)[:3] == ["network", "station", "location"] + assert set(leaf["location"]) == {"--"} + + def test_merge_false_keeps_the_tree(self): + tree = xd.trigger(..., thresh=self.thresh)(self.tree(), merge=False) + assert isinstance(tree, xd.DataCollection) + assert tree.name == "network" + assert tree["IA"].name == "station" + assert list(tree["IA"]) == ["DBNFM", "LBFI"] + + def test_an_atom_without_a_merge_hook_rebuilds_the_tree(self): + da = xd.testing.dummy() + dc = xd.DataCollection({"das1": da, "das2": da}, "instrument") + atom = xd.atoms.Partial(np.square) + assert atom.merge is None + result = atom(dc) + assert isinstance(result, xd.DataCollection) + assert list(result) == ["das1", "das2"] + + def test_sequence_levels_contribute_their_position(self): + dc = xd.DataCollection( + {"node": xd.DataCollection([self.cft(0.0), self.cft(10.0)], "acquisition")}, + "cable", + ) + result = xd.trigger(dc, thresh=self.thresh) + assert list(result.columns)[:2] == ["cable", "acquisition"] + assert list(result["cable"]) == ["node"] * 4 + assert list(result["acquisition"]) == [0, 0, 1, 1] + + def test_an_unnamed_level_contributes_no_column(self): + dc = xd.DataCollection({"a": self.cft(), "b": self.cft()}) + result = xd.trigger(dc, thresh=self.thresh) + assert list(result.columns) == ["phase", "time", "value"] + assert len(result) == 4 + + def test_a_single_leaf_collection_still_merges(self): + dc = xd.DataCollection({"DBNFM": self.cft()}, "station") + result = xd.trigger(dc, thresh=self.thresh) + assert isinstance(result, pd.DataFrame) + assert list(result["station"]) == ["DBNFM"] * 2 + + def test_a_leaf_without_a_pick_contributes_no_row(self): + dc = xd.DataCollection( + {"DBNFM": self.cft(), "QUIET": self.cft(quiet=True)}, "station" + ) + result = xd.trigger(dc, thresh=self.thresh) + assert set(result["station"]) == {"DBNFM"} + + def test_a_collection_without_any_pick_gives_an_empty_table(self): + dc = xd.DataCollection( + {"A": self.cft(quiet=True), "B": self.cft(quiet=True)}, "station" + ) + result = xd.trigger(dc, thresh=self.thresh) + assert isinstance(result, pd.DataFrame) + assert len(result) == 0 + + def test_an_empty_collection_gives_an_empty_table(self): + result = xd.trigger(xd.DataCollection({}, "station"), thresh=self.thresh) + assert isinstance(result, pd.DataFrame) + assert result.empty + + @pytest.mark.filterwarnings("error::UserWarning") + def test_an_agreeing_scalar_coordinate_dedupes_silently(self): + # what the obspy engine produces: every leaf carries its four SEED + # identifiers as scalar coordinates, and the tree keys hold the very + # same values + dc = xd.DataCollection( + { + "DBNFM": self.cft(station="DBNFM"), + "LBFI": self.cft(station="LBFI"), + }, + "station", + ) + result = xd.trigger(dc, thresh=self.thresh) + assert list(result.columns) == ["station", "phase", "time", "value"] + assert list(result["station"]) == ["DBNFM", "DBNFM", "LBFI", "LBFI"] + + def test_a_disagreeing_scalar_coordinate_warns_and_the_tree_path_wins(self): + dc = xd.DataCollection({"DBNFM": self.cft(station="ST01")}, "station") + with pytest.warns(UserWarning, match="disagrees with the tree path"): + result = xd.trigger(dc, thresh=self.thresh) + assert list(result.columns) == ["station", "phase", "time", "value"] + assert list(result["station"]) == ["DBNFM", "DBNFM"] + + def test_sequential_delegates_to_its_last_merging_stage(self): + pipeline = xd.atoms.Partial(np.abs) >> xd.trigger(..., thresh=self.thresh) + assert pipeline.merge is not None + result = pipeline(self.tree()) + assert isinstance(result, pd.DataFrame) + assert list(result.columns)[:3] == ["network", "station", "location"] + + def test_sequential_without_a_merging_stage_has_none(self): + pipeline = xd.atoms.Partial(np.abs) >> xd.atoms.Partial(np.square) + assert pipeline.merge is None + + def test_unjoinable_leaf_chunks_are_annotated_one_by_one(self): + class TableAndArray(xd.atoms.Atom): + """Emits two chunks of different types, which cannot be joined.""" + + def initialize(self, x, **flags): + pass + + def call(self, x, **flags): + return [x, pd.DataFrame({"value": [1.0]})] + + dc = xd.DataCollection({"DBNFM": self.cft()}, "station") + leaf = TableAndArray()(dc)["DBNFM"] + assert isinstance(leaf, xd.DataCollection) + assert isinstance(leaf[0], xd.DataArray) + assert list(leaf[1].columns) == ["station", "value"] + assert list(leaf[1]["station"]) == ["DBNFM"] diff --git a/tests/test_fft.py b/tests/test_fft.py index 0c9977e5..24071fdc 100644 --- a/tests/test_fft.py +++ b/tests/test_fft.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import xdas as xd import xdas.fft as xfft @@ -82,3 +83,34 @@ def test_ifft(self): def test_irfft(self): da = xd.testing.dummy() assert xfft.irfft(da).equals(xfft.irfft(da, dim={"distance": "signal"})) + + +class TestChunkGuard: + """FFTs need the whole record along their dimension: they refuse chunked + execution along it but stay usable in pipelines chunked along another.""" + + def test_chunked_along_transform_dim_raises(self): + da = xd.testing.dummy() + chunk = da.isel(time=slice(0, 50)) + for func in (xfft.fft, xfft.rfft, xfft.ifft, xfft.irfft): + atom = func(..., dim={"time": "frequency"}) + with pytest.raises(ValueError, match="whole record"): + atom(chunk, chunk_dim="time") + + def test_default_dim_resolves_to_the_last_one(self): + # the default `dim` is not an unknown dimension: it is the last one, + # so it is refused along the last one and allowed along the others. + da = xd.testing.dummy() + atom = xfft.fft(...) + with pytest.raises(ValueError, match="whole record"): + atom(da.isel(distance=slice(0, 50)), chunk_dim="distance") + result = xfft.fft(...)(da.isel(time=slice(0, 50)), chunk_dim="time") + assert result.dims == ("time", "spectrum") + + def test_chunked_along_other_dim_commutes(self): + da = xd.testing.dummy() + atom = xfft.rfft(..., dim={"distance": "wavenumber"}) + chunks = [atom(chunk, chunk_dim="time") for chunk in xd.split(da, 4, "time")] + result = xd.concat(chunks, "time") + expected = xfft.rfft(da, dim={"distance": "wavenumber"}) + assert result.equals(expected) diff --git a/tests/test_numpy.py b/tests/test_numpy.py index f21c5079..eef585e6 100644 --- a/tests/test_numpy.py +++ b/tests/test_numpy.py @@ -99,3 +99,19 @@ def test_out(self): out = da.copy() np.cumsum(da, axis=-1, out=out) assert not out.equals(da) + + def test_explicit_axis_wins_over_registered_default(self): + # registration defaults such as `axis=-1` fill in only when the caller + # says nothing; an explicitly passed axis must survive, positional + # included. + da = xd.testing.dummy() + expected = np.cumsum(da.data, axis=0) + np.testing.assert_array_equal(np.cumsum(da, axis=0).data, expected) + np.testing.assert_array_equal(np.cumsum(da, 0).data, expected) + np.testing.assert_array_equal(np.cumsum(da).data, np.cumsum(da.data, axis=-1)) + + def test_explicit_axis_wins_on_a_reduction(self): + da = xd.testing.dummy() + result = np.sum(da, 0) + assert result.dims == ("distance",) + np.testing.assert_array_equal(result.data, np.sum(da.data, axis=0)) diff --git a/tests/test_pools.py b/tests/test_pools.py new file mode 100644 index 00000000..8799840b --- /dev/null +++ b/tests/test_pools.py @@ -0,0 +1,378 @@ +import gc +import os +import shutil +import sys +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor + +import cloudpickle +import numpy as np +import pytest +from loky import process_executor + +import xdas as xd +from xdas.processing.core import AUTO_CHUNK_NBYTES, get_pool +from xdas.processing.pools import ( + _ARENAS, + PREFIX, + Arena, + ProcessFuture, + ProcessPool, + ShmRef, + _directory, + _init_worker, + _offloadable, + _park, + _resolve, + _run, + _unlink, + attach, + sweep, + view, +) + +# The tasks below are sent to worker processes, which would have to import this +# module to unpickle them by reference -- and cannot: pytest imports the test +# suite without putting its directory on the path, so `tests` is a package only +# in the parent. Sending them by value keeps the worker from needing it. +cloudpickle.register_pickle_by_value(sys.modules[__name__]) + + +def _double(x): + """A task returning something that is not a chunk.""" + return 2 * x + + +def _raise(): + """A task that fails, to check errors propagate and slots come back.""" + raise ValueError("broken task") + + +def _identity(chunk): + """A task taking a chunk in, so the argument path gets exercised.""" + return chunk + + +def _make(length): + """A task producing a chunk without being sent one, so nothing is staged.""" + return xd.testing.dummy(shape=(length, 10)) + + +def _sum(chunk): + """A task reading a staged argument without sending an array back.""" + return float(chunk.values.sum()) + + +class TestArena: + def test_slots_are_distinct_and_reused(self): + arena = Arena(3, 1024) + try: + slots = [arena.reserve() for _ in range(3)] + offsets = [offset for _, offset, _ in slots] + assert sorted(offsets) == [0, 1024, 2048] + assert arena.reserve() is None # exhausted, callers fall back + arena.release(offsets[0]) + assert arena.reserve()[1] == offsets[0] + finally: + arena.close() + + def test_close_unlinks_but_keeps_mappings_valid(self): + arena = Arena(2, 1024) + _, offset, _ = arena.reserve() + ref = ShmRef(arena.path, offset, (4,), "> xd.filter(..., (1.0, 10.0)) >> np.square + + +def gappy(da, at=50, gap=10): + """Return *da* with a gap of *gap* samples at index *at*.""" + left = da.isel(time=slice(0, at)) + right = da.isel(time=slice(at + gap, None)) + return xd.concat([left, right], "time") + + +@pytest.fixture +def virtual(da, tmp_path): + for index, chunk in enumerate(xd.split(da, 4, "time")): + chunk.to_netcdf(tmp_path / f"{index:03d}.nc") + return xd.open_mfdataarray(str(tmp_path / "*.nc")) + + +class TestSourceDispatch: + def test_in_memory_eager(self, da, pipeline): + assert pipeline.process(da).equals(pipeline(da)) + + def test_in_memory_chunked(self, da, pipeline): + expected = pipeline(da) + result = pipeline.process(da, chunks={"time": 30}) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_in_memory_gappy_chunked(self, da, pipeline): + da = gappy(da) + expected = pipeline(da) + result = pipeline.process(da, chunks={"time": 30}) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_virtual_auto(self, da, pipeline, virtual): + source = xp.get_source(virtual) + assert isinstance(source, xp.DataArrayLoader) + result = pipeline.process(virtual) + expected = pipeline(da) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_glob_and_directory(self, da, pipeline, virtual, tmp_path): + expected = pipeline(da) + for spec in (str(tmp_path / "*.nc"), str(tmp_path)): + result = pipeline.process(spec) + assert np.allclose(result.values, expected.values) + + def test_iterable_passthrough(self, da, pipeline): + chunks = xd.split(da, 4, "time") + result = pipeline.process(iter(list(chunks))) + expected = pipeline(da) + assert np.allclose(result.values, expected.values) + + def test_loader_passthrough(self, da, pipeline): + loader = xp.DataArrayLoader(da, {"time": 30}) + assert xp.get_source(loader) is loader + result = pipeline.process(loader) + assert np.allclose(result.values, pipeline(da).values) + + def test_multi_acquisition_chains_loaders(self, da, tmp_path): + other = xd.testing.dummy(shape=(40, 3)) + da.to_netcdf(tmp_path / "000.nc") + other.to_netcdf(tmp_path / "001.nc") + collection = xd.open_mfdataarray(str(tmp_path / "*.nc")) + assert isinstance(collection, xd.DataSequence) + source = xp.get_source(collection) + assert isinstance(source, _ChainSource) + assert source.chunk_dim == "time" + assert source.nbytes == da.nbytes + other.nbytes + chunks = list(source) + assert xd.concat(chunks[:1]).equals(da) or len(chunks) >= 2 + + def test_unknown_scheme_raises(self): + with pytest.raises(ValueError, match="URL scheme"): + xp.get_source("ftp://somewhere") + + def test_invalid_source_raises(self): + with pytest.raises(TypeError, match="source"): + xp.get_source(42) + + def test_tcp_scheme(self, opened): + address = f"tcp://localhost:{xd.io.get_free_port()}" + source = opened(xp.get_source(address)) + assert isinstance(source, xp.ZMQSubscriber) + assert source.unbounded + assert source.chunk_dim == "time" + + +class TestAutoChunks: + def test_file_aligned(self, virtual, monkeypatch): + # A tiny budget aligns chunk boundaries to the per-file extents. + monkeypatch.setattr(xp.core, "AUTO_CHUNK_NBYTES", 1) + dim, divs = _auto_chunks(virtual) + assert dim == "time" + assert divs == [0, 26, 51, 76, 101] + + def test_merged_to_budget(self, virtual, monkeypatch): + # Two files fit the budget: boundaries merge pairwise. + nbytes_per_slice = virtual.nbytes // virtual.sizes["time"] + monkeypatch.setattr(xp.core, "AUTO_CHUNK_NBYTES", 52 * nbytes_per_slice) + _, divs = _auto_chunks(virtual) + assert divs == [0, 51, 101] + + def test_dense_fallback(self, da, monkeypatch): + nbytes_per_slice = da.nbytes // da.sizes["time"] + monkeypatch.setattr(xp.core, "AUTO_CHUNK_NBYTES", 30 * nbytes_per_slice) + dim, divs = _auto_chunks(da) + assert dim == "time" + assert divs == [0, 30, 60, 90, 101] + + def test_loader_accepts_auto(self, virtual, monkeypatch): + monkeypatch.setattr(xp.core, "AUTO_CHUNK_NBYTES", 1) + loader = xp.DataArrayLoader(virtual, "auto") + assert loader.chunk_size is None + assert len(loader) == 4 + assert xd.concat(list(loader), "time").equals(virtual.load()) + + def test_loader_rejects_bad_chunks(self, da): + with pytest.raises(TypeError, match="auto"): + xp.DataArrayLoader(da, "automatic") + + def test_tile_aligned(self, da, tmp_path, monkeypatch): + for index, chunk in enumerate(xd.split(da, 4, "time")): + chunk.to_netcdf(tmp_path / f"{index:03d}.nc") + tiled = xd.open_mfdataarray(str(tmp_path / "*.nc"), vtype="tiles") + monkeypatch.setattr(xp.core, "AUTO_CHUNK_NBYTES", 1) + dim, divs = _auto_chunks(tiled) + assert dim == "time" + assert divs == [0, 26, 51, 76, 101] + + def test_chained_loaders_with_explicit_chunks(self, da, tmp_path): + other = xd.testing.dummy(shape=(40, 3)) + da.to_netcdf(tmp_path / "000.nc") + other.to_netcdf(tmp_path / "001.nc") + collection = xd.open_mfdataarray(str(tmp_path / "*.nc")) + # The per-run chunk size is clipped to the smallest run. + source = xp.get_source(collection, {"time": 60}) + chunks = list(source) + assert [chunk.sizes["time"] for chunk in chunks] == [60, 41, 40] + + +class TestSinkDispatch: + def test_directory(self, da, pipeline, tmp_path): + result = pipeline.process(da, out=str(tmp_path / "out"), chunks={"time": 30}) + assert np.allclose(result.values, pipeline(da).values) + assert len(os.listdir(tmp_path / "out")) > 0 + + def test_dataarray_to_file_raises(self, da, pipeline, tmp_path): + with pytest.raises(ValueError, match="directory"): + pipeline.process(da, out=str(tmp_path / "out.nc"), chunks={"time": 30}) + + def test_dataframe_to_csv(self, da, tmp_path): + atom = Partial(lambda da: pd.DataFrame({"mean": [float(np.mean(da.values))]})) + path = tmp_path / "picks.csv" + result = atom.process(da, out=str(path), chunks={"time": 30}) + assert path.exists() + assert len(result) == 4 + + def test_dataframe_to_other_suffix_raises(self, da, tmp_path): + atom = Partial(lambda da: pd.DataFrame({"mean": [0.0]})) + with pytest.raises(ValueError, match="csv"): + atom.process(da, out=str(tmp_path / "picks.parquet"), chunks={"time": 30}) + + def test_stream_to_directory(self, tmp_path): + data = np.random.randint(-1000, 1000, size=(1000, 3), dtype=np.int32) + starttime = np.datetime64("2023-01-01T00:00:00") + da = xd.DataArray( + data=data, + coords={ + "time": { + "tie_indices": [0, data.shape[0] - 1], + "tie_values": [ + starttime, + starttime + np.timedelta64(10, "ms") * (data.shape[0] - 1), + ], + "sampling_interval": np.timedelta64(10, "ms"), + }, + "distance": 5.0 * np.arange(data.shape[1]), + }, + ) + atom = Partial( + lambda da: da.to_stream( + network="NT", + station="ST{:03}", + channel="HN1", + location="00", + dim={"distance": "time"}, + ) + ) + result = atom.process(da, out=str(tmp_path), chunks={"time": 100}) + assert len(result) == 3 + assert (tmp_path / "2023").exists() + + def test_writer_instance_passthrough(self, da, pipeline, tmp_path): + writer = xp.DataArrayWriter(tmp_path, create_dirs=True) + result = pipeline.process(da, out=writer, chunks={"time": 30}) + assert np.allclose(result.values, pipeline(da).values) + + def test_none_with_no_output_returns_none(self, da): + atom = Partial(lambda da: None) + assert atom.process(da, chunks={"time": 30}) is None + + def test_eager_with_no_output_returns_none(self, da, tmp_path): + atom = Partial(lambda da: None) + assert atom.process(da, out=str(tmp_path / "out")) is None + assert not (tmp_path / "out").exists() + + def test_empty_chunk_dropped_by_writer(self, da, tmp_path): + writer = xp.DataArrayWriter(tmp_path) + writer.submit(da.isel(time=slice(0, 0))) + assert len(os.listdir(tmp_path)) == 0 + + def test_unknown_chunk_type_raises(self, da): + atom = Partial(lambda da: object()) + with pytest.raises(TypeError, match="no writer"): + atom.process(da, out="somewhere", chunks={"time": 30}) + + def test_invalid_out_raises(self, da, pipeline): + with pytest.raises(TypeError, match="cannot infer"): + pipeline.process(da, out=42, chunks={"time": 30}) + + def test_unknown_sink_scheme_raises(self, da): + with pytest.raises(ValueError, match="URL scheme"): + xp.get_writer("ftp://somewhere", da) + + def test_writer_instance_in_get_writer(self, da, tmp_path): + writer = xp.DataArrayWriter(tmp_path) + assert xp.get_writer(writer, da) is writer + + def test_eager_gappy_with_out(self, da, pipeline, tmp_path): + # The eager result is a collection: each run is written as a chunk. + result = pipeline.process(gappy(da), out=str(tmp_path / "out")) + expected = pipeline(gappy(da)) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_tcp_sink(self, da, pipeline): + address = f"tcp://localhost:{xd.io.get_free_port()}" + result = pipeline.process(da, out=address, chunks={"time": 30}) + assert result is None + + def test_legacy_signature(self, da, pipeline, tmp_path): + loader = xp.DataArrayLoader(da, {"time": 30}) + writer = xp.DataArrayWriter(tmp_path) + result = xp.process(pipeline, loader, writer) + assert np.allclose(result.values, pipeline(da).values) + + def test_eager_with_out(self, da, pipeline, tmp_path): + # In-memory source, no chunks: eager call, then sink dispatch. + result = pipeline.process(da, out=str(tmp_path / "out")) + assert np.allclose(result.values, pipeline(da).values) + assert len(os.listdir(tmp_path / "out")) > 0 + + def test_none_accumulates_streams(self, tmp_path): + data = np.random.randint(-1000, 1000, size=(200, 3), dtype=np.int32) + starttime = np.datetime64("2023-01-01T00:00:00") + da = xd.DataArray( + data=data, + coords={ + "time": { + "tie_indices": [0, data.shape[0] - 1], + "tie_values": [ + starttime, + starttime + np.timedelta64(10, "ms") * (data.shape[0] - 1), + ], + "sampling_interval": np.timedelta64(10, "ms"), + }, + "distance": 5.0 * np.arange(data.shape[1]), + }, + ) + atom = Partial( + lambda da: da.to_stream( + network="NT", + station="ST{:03}", + channel="HN1", + location="00", + dim={"distance": "time"}, + ) + ) + result = atom.process(da, chunks={"time": 100}) + assert len(result) == 6 # 3 stations x 2 chunks, unmerged + + +class TestJoinChunks: + def test_empty_and_single(self, da): + assert _join_chunks([]) is None + assert _join_chunks([da]) is da + + def test_dataframes(self): + parts = [pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [2]})] + result = _join_chunks(parts) + assert list(result["a"]) == [1, 2] + + def test_unconcatenatable_falls_back_to_collection(self, da): + other = xd.testing.dummy(shape=(10, 3)) + result = _join_chunks([da, other], "time") + assert isinstance(result, xd.DataSequence) + + def test_no_dim_falls_back_to_collection(self, da): + assert isinstance(_join_chunks([da, da], None), xd.DataSequence) + + def test_mixed_types_fall_back_to_list(self, da): + parts = [da, pd.DataFrame({"a": [1]})] + assert isinstance(_join_chunks(parts, "time"), list) + + def test_to_human(self): + assert _to_human(1) == "1 B" + assert _to_human(5 * 2**20) == "5.0 MB" + assert _to_human(2**42) == "4.0 TB" + + +class TestGuards: + def test_eager_on_huge_virtual_raises(self, pipeline, virtual, monkeypatch): + monkeypatch.setitem(Config.config, "memory_limit", 1) + with pytest.raises(ValueError, match="process"): + pipeline(virtual) + + def test_process_streams_below_guard(self, da, pipeline, virtual, monkeypatch): + # Streaming stays legal with the same tiny limit on ingress, since + # chunks are loaded one at a time; only out=None accumulation trips. + monkeypatch.setitem(Config.config, "memory_limit", 1) + with pytest.raises(ValueError, match="memory_limit"): + pipeline.process(virtual) + + def test_accumulation_guard(self, da, pipeline, monkeypatch): + monkeypatch.setitem(Config.config, "memory_limit", 1) + with pytest.raises(ValueError, match="memory_limit"): + pipeline.process(da, chunks={"time": 30}) + + def test_disk_sink_ignores_guard(self, da, pipeline, tmp_path, monkeypatch): + monkeypatch.setitem(Config.config, "memory_limit", 1) + pipeline.process(da, out=str(tmp_path / "out"), chunks={"time": 30}) + + +class TestUnbounded: + class Source: + """A never-ending source that the user interrupts after 3 chunks.""" + + chunk_dim = "time" + unbounded = True + + def __init__(self, chunks): + self.chunks = chunks + self.stopped = False + + def __iter__(self): + yield from self.chunks + raise KeyboardInterrupt + + def stop(self): + self.stopped = True + + def test_keyboard_interrupt_flushes(self, da, pipeline): + chunks = list(xd.split(da, 4, "time")) + source = self.Source(chunks) + result = pipeline.process(source) + expected = pipeline(da) + assert source.stopped + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_keyboard_interrupt_propagates_when_bounded(self, da, pipeline): + class Bounded(self.Source): + unbounded = False + + with pytest.raises(KeyboardInterrupt): + pipeline.process(Bounded(list(xd.split(da, 4, "time")))) + + def test_until_truncates(self, da, pipeline): + until = da["time"][60].values + expected = pipeline(da.sel(time=slice(None, until))) + result = pipeline.process(da, chunks={"time": 30}, until=until) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_until_truncates_an_unchunked_source(self, da, pipeline): + # Nothing is chunked here, so `until` cuts the source itself; the cut + # is the same one, inclusive, as `sel` slices. + until = da["time"][60].values + result = pipeline.process(da, until=until) + expected = pipeline(da.sel(time=slice(None, until))) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_until_skips_late_chunks(self, da, pipeline): + until = da["time"][30].values + chunks = list(xd.split(da, [30], "time")) + result = pipeline.process(iter(chunks), until=until) + expected = pipeline(da.sel(time=slice(None, until))) + assert np.allclose(result.values, expected.values) + + def test_until_ignores_chunks_without_the_dim(self, da): + # Chunks that do not carry the chunked dimension are passed through. + chunk = xd.testing.dummy(dims=("distance",), shape=(10,), step=(10.0,)) + atom = Partial(np.square) + result = atom.process(iter([chunk]), until=np.datetime64("2024-05-21")) + assert np.allclose(result.values, np.square(chunk).values) + + def test_until_as_string(self, da, pipeline): + until = da["time"][60].values + result = pipeline.process(da, chunks={"time": 30}, until=str(until)) + expected = pipeline(da.sel(time=slice(None, until))) + assert np.allclose(result.values, expected.values) + + def test_until_inside_gap_breaks(self, da, pipeline): + # Chunks split exactly at the gap: the first chunk ends before + # `until`, the next one starts after it and is skipped entirely. + left = da.isel(time=slice(0, 50)) + right = da.isel(time=slice(60, None)) + until = left["time"][-1].values + np.timedelta64(50, "ms") + result = pipeline.process(iter([left, right]), until=until) + expected = pipeline(left) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + def test_realtime_seam_warns_per_seam(self, da): + left = da.isel(time=slice(0, 50)) + right = da.isel(time=slice(60, None)) + source = self.Source([left, right]) + atom = Partial(np.square) + with pytest.warns(UserWarning, match="realtime source has a discontinuity"): + atom.process(source) + + @pytest.mark.filterwarnings("error::UserWarning") + def test_realtime_continuous_stream_is_silent(self, da): + source = self.Source(list(xd.split(da, 4, "time"))) + atom = Partial(np.square) + atom.process(source) + + def test_chunked_source_announces_its_splits_upfront(self, da, pipeline): + with pytest.warns(UserWarning, match="1 discontinuity along 'time'"): + pipeline.process(gappy(da), chunks={"time": 30}) + + @pytest.mark.filterwarnings("error::UserWarning") + def test_upfront_scan_skips_non_axis_coordinates(self): + # A dense coordinate has no free discontinuity scan: the source is + # processed without any upfront announcement. + dense = xd.testing.dummy(shape=(52, 5), ctype="dense") + atom = Partial(np.square) + result = atom.process(dense, chunks={"time": 20}) + assert np.allclose(result.values, np.square(dense).values) + + @pytest.mark.filterwarnings("error::UserWarning") + def test_realtime_chunks_without_the_dim_are_not_judged(self, da): + # A realtime chunk that does not carry the chunked dimension leaves + # the seam information untouched rather than resetting it. + aside = xd.testing.dummy(dims=("distance",), shape=(5,), step=(10.0,)) + left, right = xd.split(da, 2, "time") + source = self.Source([left, aside, right]) + atom = Partial(np.square) + atom.process(source) + + @pytest.mark.filterwarnings("error::UserWarning") + def test_realtime_one_sample_chunk_adopts_the_stream_rate(self): + # A one-sample chunk of a sampled coordinate declares no rate of its + # own: continuous with the stream, it inherits the previous chunk's + # delta so the seam after it is still judged correctly. + sampled = xd.testing.dummy(shape=(52, 5), ctype="sampled") + chunks = list(xd.split(sampled, [50, 51], "time")) + source = self.Source(chunks) + atom = Partial(np.square) + atom.process(source) + + def test_watch_is_a_realtime_loader(self, da, tmp_path): + loader = xd.watch(tmp_path) + try: + assert isinstance(loader, xp.RealTimeLoader) + assert loader.unbounded + finally: + loader.stop() + + def test_watch_source_end_to_end(self, da, pipeline, tmp_path): + # Feed the queue directly (the watchdog handler is tested elsewhere) + # and close the stream with the None sentinel. + loader = xd.watch(tmp_path) + for chunk in xd.split(da, 4, "time"): + loader.queue.put(chunk) + loader.queue.put(None) + result = pipeline.process(loader) + expected = pipeline(da) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + +class TestZMQRoundTrip: + def test_publish_process_subscribe(self, da, opened): + address = f"tcp://localhost:{xd.io.get_free_port()}" + packets = list(xd.split(da, 10, "time")) + # Bind before connecting so the subscription is live for packet one. + publisher = opened(xp.ZMQPublisher(address)) + source = opened(xp.get_source(address)) + + def publish(): + publisher.wait_for_subscribers(timeout=60.0) + for packet in packets: + publisher.submit(packet) + + thread = threading.Thread(target=publish) + thread.start() + atom = Partial(np.square) + until = da["time"][-1].values + result = atom.process(source, until=until) + thread.join() + expected = np.square(da) + assert result.coords.equals(expected.coords) + assert np.allclose(result.values, expected.values) + + +def cft(t0=0.0, quiet=False, open_ended=False): + """ + A characteristic function peaking on P and S at ``t0 + 1``. + + ``open_ended`` leaves the trigger on at the last sample, so the pick is + only emitted when the atom is flushed. + """ + lane = [0.0, 0.8, 0.8] if open_ended else [0.0, 0.8, 0.0] + if quiet: + lane = [0.0, 0.0, 0.0] + return xd.DataArray( + data=[[0.0, 0.0, 0.0], lane, lane], + coords={ + "phase": ["N", "P", "S"], + "time": { + "tie_indices": [0, 2], + "tie_values": [t0, t0 + 2.0], + "sampling_interval": 1.0, + }, + }, + ) + + +def seed_tree(): + """A ``network / station / location`` tree of pick-able leaves.""" + return xd.DataCollection( + { + "IA": xd.DataCollection( + { + "DBNFM": xd.DataCollection({"--": cft()}, "location"), + "LBFI": xd.DataCollection({"00": cft(10.0)}, "location"), + }, + "station", + ) + }, + "network", + ) + + +def das_tree(): + """A ``node / cable / acquisition`` tree, whose sequence level folds.""" + return xd.DataCollection( + { + "N1": xd.DataCollection( + {"C1": xd.DataCollection([cft(0.0), cft(10.0)], "acquisition")}, "cable" + ) + }, + "node", + ) + + +def picker(): + """A fresh table-valued atom, whose tables carry their tree path.""" + return xd.trigger(..., thresh={"P": 0.5, "S": 0.5}) + + +class Resets(xd.atoms.Atom): + """An identity atom logging the chunk count it held at every reset.""" + + def __init__(self, log): + super().__init__() + self.log = log + self.count = State(...) + + def initialize(self, x, **flags): + self.count = State(0) + + def call(self, x, **flags): + self.count = State(self.count + 1) + return x + + def reset(self): + self.log.append(self.count) + super().reset() + + +class TestCollectionWalk: + """ + W11: `process` walks a collection exactly as `atom(dc)` walks it. + + The defining invariant is `atom.process(dc, out=None) == atom(dc)`: it is + what stops the streaming and the eager walk drifting apart, and it matters + most where a leaf is too large to call eagerly at all. + """ + + @pytest.mark.parametrize("chunks", [None, {"time": 2}]) + def test_the_seed_tree_streams_to_the_eager_result(self, chunks): + dc = seed_tree() + expected = picker()(dc) + result = picker().process(dc, out=None, chunks=chunks) + assert isinstance(result, pd.DataFrame) + assert result.equals(expected) + + @pytest.mark.parametrize("chunks", [None, {"time": 2}]) + def test_the_das_tree_streams_to_the_eager_result(self, chunks): + dc = das_tree() + expected = picker()(dc) + result = picker().process(dc, out=None, chunks=chunks) + assert list(result["acquisition"]) == [0, 0, 1, 1] + assert result.equals(expected) + + def test_an_array_atom_rebuilds_the_walked_tree(self, da, pipeline): + dc = xd.DataCollection({"a": da, "b": da}, "node") + expected = pipeline(dc) + result = pipeline.process(dc, chunks={"time": 30}) + assert isinstance(result, xd.DataMapping) + assert result.name == "node" and list(result) == ["a", "b"] + for key in expected: + assert result[key].coords.equals(expected[key].coords) + assert np.allclose(result[key].values, expected[key].values) + + def test_a_bare_sequence_folds_like_the_eager_call(self): + dc = xd.DataCollection([cft(0.0), cft(10.0)], "acquisition") + expected = picker()(dc) + result = picker().process(dc, out=None) + assert result.equals(expected) + + def test_the_flushed_tail_goes_to_the_last_element(self): + dc = xd.DataCollection([cft(0.0), cft(10.0, open_ended=True)], "acquisition") + expected = picker()(dc) + result = picker().process(dc, out=None) + assert list(result["acquisition"]) == [0, 0, 1, 1] + assert result.equals(expected) + + def test_a_sequence_with_nothing_to_fold_along_walks_element_by_element(self, da): + atom = Partial(np.square) + dc = xd.DataCollection([da, da]) + result = atom.process(dc) + assert isinstance(result, xd.DataSequence) and len(result) == 2 + assert np.allclose(result[0].values, np.square(da).values) + + def test_a_bare_callable_walks_a_collection_too(self, da): + dc = xd.DataCollection({"a": xd.DataCollection([da], "acquisition")}, "node") + result = xp.process(np.square, dc) + assert np.allclose(result["a"][0].values, np.square(da).values) + + def test_a_mapping_inside_a_folded_sequence_raises(self): + dc = xd.DataCollection([cft(0.0), xd.DataCollection({"a": cft(10.0)})]) + with pytest.raises(NotImplementedError, match="mapping collections"): + picker().process(dc) + + def test_merge_false_keeps_the_labelled_tree(self): + tree = picker().process(seed_tree(), out=None, merge=False) + assert isinstance(tree, xd.DataCollection) + assert tree.name == "network" + leaf = tree["IA"]["DBNFM"]["--"] + assert list(leaf.columns)[:3] == ["network", "station", "location"] + + def test_the_atom_is_reset_between_leaves(self, da): + log = [] + dc = xd.DataCollection({key: da for key in "abc"}, "node") + Resets(log).process(dc, chunks={"time": 30}) + # the walk resets before every leaf, so each of the first two leaves + # left its own chunk count behind and none of them inherited it + assert [count for count in log if isinstance(count, int)] == [4, 4] + + def test_state_does_not_leak_from_one_leaf_to_the_next(self, da, pipeline): + dc = xd.DataCollection({"a": da, "b": da}, "node") + result = pipeline.process(dc, chunks={"time": 30}) + assert np.allclose(result["a"].values, result["b"].values) + + +class TestCollectionSinks: + """W11: one rule per kind of destination, applied leaf by leaf.""" + + def test_one_csv_holds_every_leaf(self, tmp_path): + path = tmp_path / "picks.csv" + result = picker().process(seed_tree(), out=str(path)) + assert path.exists() + written = pd.read_csv(path) + assert list(written["station"]) == ["DBNFM", "DBNFM", "LBFI", "LBFI"] + assert result.equals(written) + + def test_a_csv_no_leaf_wrote_to_stays_absent(self, tmp_path): + path = tmp_path / "picks.csv" + dc = xd.DataCollection({"A": cft(quiet=True)}, "station") + assert picker().process(dc, out=str(path)).empty + assert not path.exists() + + def test_a_shared_csv_no_leaf_emitted_to_closes_to_nothing(self, da, tmp_path): + # Not one leaf emits, so the shared table is never even resolved to a + # writer: closing it has nothing to report and nothing to write. + atom = Partial(lambda x: None) + dc = xd.DataCollection({"a": da, "b": da}, "node") + path = tmp_path / "picks.csv" + assert atom.process(dc, out=str(path)) is None + assert not path.exists() + + def test_a_leaf_emitting_nothing_answers_with_an_empty_collection( + self, da, tmp_path + ): + atom = Partial(lambda x: None) + dc = xd.DataCollection({"a": da, "b": da}, "node") + out = tmp_path / "results" + result = atom.process(dc, out=str(out), chunks={"time": 30}) + assert all(len(leaf) == 0 for leaf in result.values()) + assert not out.exists() # no directory is created for nothing + + def test_a_directory_fans_out_to_the_tree_path(self, da, pipeline, tmp_path): + dc = xd.DataCollection( + {"N1": xd.DataCollection({"C1": da, "C2": da}, "cable")}, "node" + ) + out = tmp_path / "results" + result = pipeline.process(dc, out=str(out), chunks={"time": 30}) + assert sorted(path.name for path in (out / "N1").iterdir()) == ["C1", "C2"] + expected = pipeline(dc) + assert np.allclose(result["N1"]["C1"].values, expected["N1"]["C1"].values) + + def test_a_folded_sequence_writes_to_one_directory(self, da, pipeline, tmp_path): + dc = xd.DataCollection({"C1": xd.DataCollection([da], "acquisition")}, "cable") + out = tmp_path / "results" + result = pipeline.process(dc, out=str(out), chunks={"time": 30}) + assert [path.name for path in out.iterdir()] == ["C1"] + assert np.allclose(result["C1"].values, pipeline(da).values) + + def test_a_writer_instance_is_shared_by_every_leaf(self, tmp_path): + writer = xp.DataFrameWriter(str(tmp_path / "picks.csv")) + result = picker().process(seed_tree(), out=writer) + assert len(result) == 4 + + def test_an_uninferable_out_raises(self): + with pytest.raises(TypeError, match="cannot infer a writer"): + picker().process(seed_tree(), out=42) + + +class TestCollectionGather: + """W11: the gather is consulted before anything is chunked.""" + + def test_a_channel_level_becomes_a_dimension_before_chunking(self): + from tests.test_atoms_ml import component_collection, picker_model + + dc = component_collection(["SHZ", "SHN", "SHE"]) + model = picker_model("original") + expected = xd.pick(dc, model, device="cpu") + result = xd.pick(..., model, device="cpu").process( + dc, out=None, chunks={"time": 8} + ) + assert "channel" not in result.columns # the level became an axis + assert result.equals(expected) + + def test_the_gathered_array_stays_virtual(self, monkeypatch, tmp_path): + from tests.test_atoms_ml import component_trace, picker_model + + keys = ["SHZ", "SHN", "SHE"] + for index, key in enumerate(keys): + (tmp_path / key).mkdir() + trace = component_trace(index, n=64) + for part, chunk in enumerate(xd.split(trace, 4, "time")): + chunk.to_netcdf(tmp_path / key / f"{part:03d}.nc") + dc = xd.DataCollection( + { + key: xd.open_mfdataarray(str(tmp_path / key / "*.nc"), vtype="tiles") + for key in keys + }, + "channel", + ) + atom = xd.pick(..., picker_model("original"), device="cpu") + assert isinstance(atom.gather(dc).data, TileArray) # nothing was read + # a limit below the stacked array (1536 B) but above the pick table + monkeypatch.setitem(Config.config, "memory_limit", 1024) + with pytest.raises(ValueError, match="process"): + atom(dc) # the eager walk refuses to load the stacked array + assert len(atom.process(dc, chunks={"time": 16})) > 0 diff --git a/tests/test_processing.py b/tests/test_processing.py index 96b8e124..58081735 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1,6 +1,6 @@ +import gc import os import threading -import time from pathlib import Path import hdf5plugin @@ -15,6 +15,9 @@ from xdas.atoms import Partial, Sequential from xdas.signal import sosfilt +TIMEOUT = 60.0 +"""Seconds any ZMQ wait is given before failing, so that no test can hang.""" + class TestDataArrayLoader: def test_init(self): @@ -54,6 +57,95 @@ def test_error_handling(self): with pytest.raises(ValueError): xp.DataArrayLoader(da, {"time": 2000}) + def test_unknown_pool_raises(self): + da = xd.testing.dummy(shape=(1000, 100)) + with pytest.raises(ValueError, match="no worker pool"): + list(xp.DataArrayLoader(da, {"time": 100}, pool="fork")) + + +class TestReadOnlyIngress: + """Chunks may arrive immutable (mapped from an arena slot). + + Atoms honor this by allocating their outputs rather than writing into + their input; this pins the contract down without spawning a process pool. + """ + + def test_pipeline_accepts_readonly_chunks(self): + da = xd.testing.dummy(shape=(1000, 100)) + sos = sp.iirfilter(4, 0.1, btype="lowpass", output="sos") + sequence = Sequential([Partial(sosfilt, sos, ..., dim="time", zi=...)]) + expected = xd.concat( + [sequence(chunk, chunk_dim="time") for chunk in xd.split(da, 10, "time")] + ) + chunks = [chunk.copy() for chunk in xd.split(da, 10, "time")] + for chunk in chunks: + chunk.data.setflags(write=False) + sequence.reset() + result = xd.concat([sequence(chunk, chunk_dim="time") for chunk in chunks]) + assert result.equals(expected) + for chunk, original in zip(chunks, xd.split(da, 10, "time")): + np.testing.assert_array_equal(chunk.data, original.data) + + +@pytest.mark.slow +class TestProcessPool: + """Worker processes get past the HDF5 lock, shared memory past the pipe.""" + + def test_loader_chunks_integrity(self, tmp_path): + # A virtual array: the manifest of each chunk is what crosses to the + # worker, which then reads its own files; the loaded chunk comes back + # through an arena slot. + expected = xd.testing.dummy(shape=(1000, 100)) + expected.to_netcdf(tmp_path / "data.nc") + da = xd.open_dataarray(tmp_path / "data.nc") + dl = xp.DataArrayLoader(da, {"time": 100}, 4, 2, pool="processes") + assert xd.concat(list(dl)).equals(expected) + + def test_loader_chunks_are_zero_copy(self): + # Read-only data is the signature of an arena-backed chunk: the parent + # mapped the worker's pages, nothing was pickled back. The arena is + # sized for streaming, so the chunks are read one at a time and + # dropped -- a caller keeping every chunk alive would run it out of + # slots and get the (correct, slower) pickle path instead. + da = xd.testing.dummy(shape=(1000, 100)) + loader = xp.DataArrayLoader(da, {"time": 100}, 2, 2, pool="processes") + assert all(not chunk.data.flags.writeable for chunk in loader) + + def test_loader_equals_threads(self): + da = xd.testing.dummy(shape=(1000, 100)) + threads = list(xp.DataArrayLoader(da, {"time": 100}, 2, 2)) + processes = list(xp.DataArrayLoader(da, {"time": 100}, 2, 2, "processes")) + assert xd.concat(processes).equals(xd.concat(threads)) + + def test_writer_equals_threads(self, tmp_path): + da = xd.testing.dummy(shape=(1000, 100)) + chunks = list(xd.split(da, 10, "time")) + results = [] + for pool in ["threads", "processes"]: + dirpath = tmp_path / pool + dirpath.mkdir() + dw = xp.DataArrayWriter(dirpath, max_buffers=2, max_workers=2, pool=pool) + for chunk in chunks: + dw.write(chunk) + results.append(dw.result()) + assert results[1].load().equals(results[0].load()) + assert results[0].load().equals(da) + + def test_end_to_end(self, tmp_path): + # which pool the chunks travel through cannot change what comes out. + xd.testing.dummy(shape=(1000, 100)).to_netcdf(tmp_path / "data.nc") + da = xd.open_dataarray(tmp_path / "data.nc") + sos = sp.iirfilter(4, 0.1, btype="lowpass", output="sos") + results = [] + for pool in ["threads", "processes"]: + sequence = Sequential([Partial(sosfilt, sos, ..., dim="time", zi=...)]) + loader = xp.DataArrayLoader(da, {"time": 100}, 2, 2, pool=pool) + dirpath = tmp_path / pool + dirpath.mkdir() + writer = xp.DataArrayWriter(dirpath, dim="time") + results.append(xp.process(sequence, loader, writer)) + assert results[1].load().equals(results[0].load()) + class TestDataArrayWriter: def test_init(self, tmp_path): @@ -89,6 +181,23 @@ def test_passing_wrong_input(self, tmp_path): with pytest.raises(TypeError): dw.submit(None) + def test_empty_chunks_are_dropped(self, tmp_path): + expected = xd.testing.dummy(shape=(1000, 100)) + dw = xp.DataArrayWriter(tmp_path, dim="time") + dw.submit(expected.isel(time=slice(0, 0))) + for chunk in xd.split(expected, 10, dim="time"): + dw.submit(chunk) + assert dw.result().equals(expected) + + def test_the_chunked_dimension_need_not_lead(self, tmp_path): + # joining on the first dimension stacks the chunks along the wrong + # axis whenever the chunked dimension does not lead the output. + expected = xd.testing.dummy(shape=(1000, 100)).transpose("distance", "time") + dw = xp.DataArrayWriter(tmp_path, dim="time") + for chunk in xd.split(expected, 10, dim="time"): + dw.submit(chunk) + assert dw.result().equals(expected) + class TestProcessing: def test_stateful(self, tmp_path): @@ -165,12 +274,14 @@ def test_multiple_dataframes(self, tmp_path): assert result.equals(expected) def test_write_empty_dataframe(self, tmp_path): + # Empty chunks are accepted and silently dropped (many flushes + # produce nothing): no file is created for them. dw = xp.DataFrameWriter(tmp_path / "output.csv") expected = pd.DataFrame() dw.submit(expected) result = dw.result() assert result.equals(expected) - assert Path(dw.path).exists() + assert not Path(dw.path).exists() def test_with_existing_file(self, tmp_path): dw1 = xp.DataFrameWriter(tmp_path / "output.csv") @@ -207,22 +318,32 @@ def test_passing_wrong_input(self, tmp_path): class TestZMQ: + @pytest.fixture(autouse=True) + def _close_endpoints(self, opened): + """Close whatever this test opens, so that no socket outlives it.""" + self.opened = opened + def _publish_and_subscribe(self, packets, address, encoding=None): - publisher = xp.ZMQPublisher(address, encoding) + publisher = self.opened(xp.ZMQPublisher(address, encoding)) def publish(): + # A recording, wanted whole. Only the publisher can hold a replay + # back until its audience has landed; a subscriber joining one + # already under way has no way to recover its first packets. + publisher.wait_for_subscribers(timeout=TIMEOUT) for packet in packets: - time.sleep(0.001) publisher.submit(packet) - threading.Thread(target=publish).start() + thread = threading.Thread(target=publish) + thread.start() + subscriber = self.opened(xp.ZMQSubscriber(address, timeout=TIMEOUT)) - subscriber = xp.ZMQSubscriber(address) result = [] for n, packet in enumerate(subscriber, start=1): result.append(packet) if n == len(packets): break + thread.join() return xd.concat(result) def test_publish_and_subscribe(self): @@ -233,6 +354,52 @@ def test_publish_and_subscribe(self): result = self._publish_and_subscribe(packets, address) assert result.equals(expected) + def test_subscriber_joins_a_real_time_flux(self): + # A real-time publisher streams whether or not anyone is listening, so + # a subscriber joining it gets the stream from wherever it lands. + packets = xd.split(xd.testing.dummy(), 10) + address = f"tcp://localhost:{xd.io.get_free_port()}" + publisher = self.opened(xp.ZMQPublisher(address)) + stop = threading.Event() + + def flux(): + while not stop.is_set(): + for packet in packets: + if stop.is_set(): + return + publisher.submit(packet) + + thread = threading.Thread(target=flux) + thread.start() + try: + subscriber = self.opened(xp.ZMQSubscriber(address, timeout=TIMEOUT)) + # The flux greets us as it streams, without ever waiting for us. + subscriber.wait_until_subscribed() + received = [next(subscriber) for _ in range(3)] + finally: + stop.set() + thread.join() + + subscriber.wait_until_subscribed() # greeted already, so this returns + for packet in received: + assert any(packet.equals(published) for published in packets) + + def test_subscriber_timeout(self): + address = f"tcp://localhost:{xd.io.get_free_port()}" + self.opened(xp.ZMQPublisher(address)) # binds, but never publishes + subscriber = self.opened(xp.ZMQSubscriber(address, timeout=0.1)) + with pytest.raises(TimeoutError, match="no packet received"): + next(subscriber) + + def test_wait_until_subscribed_needs_a_publisher_that_publishes(self): + # Only a publisher that publishes can acknowledge anybody: a silent one + # leaves a subscriber waiting, which is what the timeout is for. + address = f"tcp://localhost:{xd.io.get_free_port()}" + self.opened(xp.ZMQPublisher(address)) + subscriber = self.opened(xp.ZMQSubscriber(address, timeout=0.1)) + with pytest.raises(TimeoutError, match="no packet received"): + subscriber.wait_until_subscribed() + def test_encoding(self): expected = xd.testing.dummy() packets = xd.split(expected, 10) @@ -479,20 +646,59 @@ def test_submit_wrong_type_raises(self, tmp_path): with pytest.raises(TypeError): sw.submit("not_a_stream") + def test_result_without_any_chunk_is_an_empty_stream(self, tmp_path): + # A pipeline that emitted nothing leaves no temporary file to merge. + result = xp.StreamWriter(tmp_path, "M").result() + assert isinstance(result, obspy.Stream) + assert len(result) == 0 + class TestZMQPublisherAliases: - def test_write_alias(self): + def test_write_alias(self, opened): address = f"tcp://localhost:{xd.io.get_free_port()}" - publisher = xp.ZMQPublisher(address) + publisher = opened(xp.ZMQPublisher(address)) da = xd.testing.dummy() publisher.write(da) # use write() alias - def test_result_returns_none(self): + def test_result_returns_none(self, opened): address = f"tcp://localhost:{xd.io.get_free_port()}" - publisher = xp.ZMQPublisher(address) + publisher = opened(xp.ZMQPublisher(address)) assert publisher.result() is None +class TestZMQEndpointLifecycle: + def address(self): + return f"tcp://localhost:{xd.io.get_free_port()}" + + def test_closing_releases_the_socket_and_the_context(self): + publisher = xp.ZMQPublisher(self.address()) + socket, context = publisher._socket, publisher._context + publisher.close() + assert socket.closed + assert context.closed + + def test_closing_twice_releases_nothing_more(self): + subscriber = xp.ZMQSubscriber(self.address()) + subscriber.close() + subscriber.close() + assert subscriber._socket is None + assert subscriber._context is None + + def test_the_context_manager_closes_on_the_way_out(self): + with xp.ZMQPublisher(self.address()) as publisher: + assert not publisher._socket.closed + assert publisher._socket is None + + def test_a_dropped_endpoint_is_closed_by_the_collector(self): + # The safety net under an endpoint nobody closed: when it runs is not + # for the caller to know, which is why `close` is the way to write it. + publisher = xp.ZMQPublisher(self.address()) + socket = publisher._socket + del publisher + gc.collect() + assert socket.closed + + class TestHandlerDirect: def test_on_closed(self, tmp_path): from queue import Queue diff --git a/tests/test_routines.py b/tests/test_routines.py index 9afca094..4b0a4c6e 100644 --- a/tests/test_routines.py +++ b/tests/test_routines.py @@ -5,6 +5,7 @@ import xdas as xd from xdas.core.routines import Bag, CompatibilityError +from xdas.virtual import VirtualSource class TestBag: @@ -1289,3 +1290,696 @@ def test_permutes_without_declared_sampling_interval(self, tmp_path): result = xd.sortby(shuffled, "time") assert np.array_equal(result["time"].values, expected["time"].values) np.testing.assert_array_equal(np.asarray(result.data), expected.values) + + +def trace(channel="SHZ", station="SX01", start=0.0, npts=5, values=None): + """A one-dimensional trace with the SEED identifiers as scalar coordinates.""" + time = {"tie_indices": [0, npts - 1], "tie_values": [start, start + npts - 1.0]} + if values is None: + values = np.arange(npts, dtype="float64") + return xd.DataArray( + values, + { + "network": (None, "DX"), + "station": (None, station), + "channel": (None, channel), + "time": time, + }, + ) + + +def regular_trace(channel="SHZ", start=0.0, npts=5, step=1.0): + """A trace whose time coordinate declares its nominal sampling interval.""" + return xd.DataArray( + np.arange(npts, dtype="float64"), + { + "channel": (None, channel), + "time": { + "tie_indices": [0, npts - 1], + "tie_values": [start, start + step * (npts - 1)], + "sampling_interval": step, + }, + }, + ) + + +def instrument(station="SX01", channels=("SHZ", "SHN", "SHE"), **kwargs): + """One station as a `channel` level of traces.""" + return xd.DataCollection( + {code: trace(code, station, **kwargs) for code in channels}, "channel" + ) + + +class TestStack: + def test_channel_level_becomes_a_dimension(self): + dc = xd.DataCollection( + {"SX01": instrument("SX01"), "SX02": instrument("SX02")}, "station" + ) + result = xd.stack(dc, "channel") + assert result.name == "station" + assert list(result) == ["SX01", "SX02"] + da = result["SX01"] + assert da.dims == ("channel", "time") + assert da.shape == (3, 5) + assert da["channel"].values.tolist() == ["SHE", "SHN", "SHZ"] + assert da["station"].values == "SX01" + assert da["network"].values == "DX" + + def test_the_new_dimension_is_named_after_the_level(self): + dc = instrument() + assert xd.stack(dc, "channel").dims == ("channel", "time") + + def test_dim_renames_the_new_dimension_and_keeps_the_keys(self): + dc = instrument() + da = xd.stack(dc, "channel", dim="component") + assert da.dims == ("component", "time") + assert da["component"].values.tolist() == ["SHE", "SHN", "SHZ"] + # the leaves' own `channel` scalar is promoted alongside it + assert da["channel"].dim == "component" + + def test_a_sequence_level_is_keyed_by_position(self): + dc = xd.DataCollection([trace(), trace(), trace()], "acquisition") + da = xd.stack(dc, "acquisition") + assert da.dims == ("acquisition", "time") + assert da["acquisition"].values.tolist() == [0, 1, 2] + + def test_a_single_member_level_gives_a_length_one_dimension(self): + dc = instrument(channels=("SHZ",)) + da = xd.stack(dc, "channel") + assert da.shape == (1, 5) + assert da["channel"].values.tolist() == ["SHZ"] + + def test_levels_below_the_collapsed_one_are_merged_in_lockstep(self): + dc = xd.DataCollection( + { + code: xd.DataCollection( + [trace(code, npts=5), trace(code, start=10.0, npts=5)], + "acquisition", + ) + for code in ("SHZ", "SHN") + }, + "channel", + ) + result = xd.stack(dc, "channel") + assert result.name == "acquisition" + assert len(result) == 2 + for da in result: + assert da.dims == ("channel", "time") + assert da["channel"].values.tolist() == ["SHN", "SHZ"] + + def test_a_mapping_below_the_collapsed_level_is_merged_in_lockstep(self): + dc = xd.DataCollection( + { + code: xd.DataCollection( + {loc: trace(code) for loc in ("00", "10")}, "location" + ) + for code in ("SHZ", "SHN") + }, + "channel", + ) + result = xd.stack(dc, "channel") + assert result.name == "location" + assert list(result) == ["00", "10"] + assert result["10"].dims == ("channel", "time") + + def test_levels_above_the_collapsed_one_are_walked(self): + dc = xd.DataCollection( + {"DX": xd.DataCollection([instrument()], "acquisition")}, "network" + ) + result = xd.stack(dc, "channel") + assert result.name == "network" + assert result["DX"][0].dims == ("channel", "time") + + def test_a_leaf_beside_the_level_is_left_alone(self): + dc = xd.DataCollection( + {"SX01": instrument("SX01"), "SX02": trace(station="SX02")}, "station" + ) + result = xd.stack(dc, "channel") + assert result["SX01"].dims == ("channel", "time") + assert result["SX02"].dims == ("time",) + + def test_varying_scalar_coordinates_are_promoted(self): + dc = xd.DataCollection( + {name: trace(station=name) for name in ("SX01", "SX02")}, "station" + ) + da = xd.stack(dc, "station") + assert da["station"].values.tolist() == ["SX01", "SX02"] + assert da["network"].dim is None + + # --- input validation --- + + def test_unknown_join_raises(self): + with pytest.raises(ValueError, match="unknown join method 'left'"): + xd.stack(instrument(), "channel", join="left") + + def test_a_data_array_is_not_a_collection(self): + with pytest.raises(TypeError, match="can only stack a level"): + xd.stack(trace(), "channel") + + def test_unknown_level_raises(self): + with pytest.raises(KeyError, match="'component' does not name any level"): + xd.stack(instrument(), "component") + + def test_an_empty_level_raises(self): + dc = xd.DataCollection({}, "channel") + with pytest.raises(ValueError, match="level 'channel' is empty"): + xd.stack(dc, "channel") + + def test_dim_colliding_with_a_leaf_dimension_raises(self): + dc = xd.DataCollection( + {code: trace(code) for code in ("SHZ", "SHN")}, "channel" + ) + with pytest.raises(ValueError, match="already has a 'time' dimension"): + xd.stack(dc, "channel", dim="time") + + # --- structural disagreement --- + + def test_a_leaf_facing_a_sub_tree_raises(self): + dc = xd.DataCollection( + { + "SHZ": trace("SHZ"), + "SHN": xd.DataCollection([trace("SHN")], "acquisition"), + }, + "channel", + ) + with pytest.raises(ValueError, match="sub-trees .*do not agree"): + xd.stack(dc, "channel") + + def test_differently_named_sub_levels_raise(self): + dc = xd.DataCollection( + { + "SHZ": xd.DataCollection([trace("SHZ")], "acquisition"), + "SHN": xd.DataCollection([trace("SHN")], "epoch"), + }, + "channel", + ) + with pytest.raises( + ValueError, match="'acquisition' sequence.*'epoch' sequence" + ): + xd.stack(dc, "channel") + + def test_differing_sub_keys_raise(self): + dc = xd.DataCollection( + { + "SHZ": xd.DataCollection({"00": trace("SHZ")}, "location"), + "SHN": xd.DataCollection({"10": trace("SHN")}, "location"), + }, + "channel", + ) + with pytest.raises(ValueError, match="does not hold the same keys"): + xd.stack(dc, "channel") + + def test_differing_sub_lengths_raise(self): + dc = xd.DataCollection( + { + "SHZ": xd.DataCollection([trace("SHZ")], "acquisition"), + "SHN": xd.DataCollection([trace("SHN"), trace("SHN")], "acquisition"), + }, + "channel", + ) + with pytest.raises(ValueError, match="does not hold the same number"): + xd.stack(dc, "channel") + + def test_a_level_nested_under_itself_raises(self): + dc = xd.DataCollection( + { + code: xd.DataCollection({code: trace(code)}, "channel") + for code in ("SHZ", "SHN") + }, + "channel", + ) + with pytest.raises(ValueError, match="nested under itself"): + xd.stack(dc, "channel") + + # --- coordinate disagreement --- + + def test_differing_leaf_dimensions_raise(self): + dc = xd.DataCollection( + {"SHZ": trace("SHZ"), "SHN": trace("SHN").expand_dims("space")}, + "channel", + ) + with pytest.raises(ValueError, match="has dimensions"): + xd.stack(dc, "channel") + + def test_a_missing_coordinate_raises(self): + dc = xd.DataCollection( + {"SHZ": trace("SHZ"), "SHN": trace("SHN").drop_coords("network")}, + "channel", + ) + with pytest.raises(ValueError, match=r"lacks the coordinates \['network'\]"): + xd.stack(dc, "channel") + + def test_a_dimension_without_coordinate_cannot_be_aligned(self): + dc = xd.DataCollection( + { + "SHZ": xd.DataArray(np.arange(4.0)), + "SHN": xd.DataArray(np.arange(5.0)), + }, + "channel", + ) + with pytest.raises(ValueError, match="no coordinate to align on"): + xd.stack(dc, "channel", join="inner") + + def test_a_time_mismatch_raises_naming_the_coordinate(self): + dc = xd.DataCollection( + {"SHZ": trace("SHZ"), "SHN": trace("SHN", start=2.0)}, "channel" + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + + def test_the_mismatch_error_points_at_the_join(self): + dc = xd.DataCollection( + {"SHZ": trace("SHZ"), "SHN": trace("SHN", start=2.0)}, "channel" + ) + with pytest.raises(ValueError, match="pass join='inner' or join='outer'"): + xd.stack(dc, "channel") + + def test_the_mismatch_error_names_where_it_happened(self): + dc = xd.DataCollection( + { + "SHZ": xd.DataCollection([trace("SHZ")], "acquisition"), + "SHN": xd.DataCollection([trace("SHN", start=2.0)], "acquisition"), + }, + "channel", + ) + with pytest.raises(ValueError, match="at acquisition=0"): + xd.stack(dc, "channel") + + def test_an_unjoinable_mismatch_gets_no_join_hint(self): + dc = xd.DataCollection( + {"SHZ": trace("SHZ"), "SHN": trace("SHN").drop_coords("network")}, + "channel", + ) + with pytest.raises(ValueError) as excinfo: + xd.stack(dc, "channel") + assert "join=" not in str(excinfo.value) + + # --- joining --- + + def test_inner_join_keeps_the_shared_span(self): + dc = xd.DataCollection( + { + "SX01": trace(station="SX01", start=0.0, npts=5), + "SX02": trace(station="SX02", start=2.0, npts=5), + }, + "station", + ) + da = xd.stack(dc, "station", join="inner") + assert da.dims == ("station", "time") + assert da["time"].values.tolist() == [2.0, 3.0, 4.0] + npt.assert_array_equal(da.values, [[2.0, 3.0, 4.0], [0.0, 1.0, 2.0]]) + + def test_outer_join_pads_with_nan(self): + dc = xd.DataCollection( + { + "SX01": trace(station="SX01", start=0.0, npts=5), + "SX02": trace(station="SX02", start=2.0, npts=5), + }, + "station", + ) + da = xd.stack(dc, "station", join="outer") + assert da["time"].values.tolist() == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + npt.assert_array_equal( + da.values, + [ + [0.0, 1.0, 2.0, 3.0, 4.0, np.nan, np.nan], + [np.nan, np.nan, 0.0, 1.0, 2.0, 3.0, 4.0], + ], + ) + + def test_outer_join_promotes_an_integer_dtype(self): + dc = xd.DataCollection( + { + "SX01": trace( + station="SX01", npts=3, values=np.arange(3, dtype="int16") + ), + "SX02": trace( + station="SX02", + start=2.0, + npts=3, + values=np.arange(3, dtype="int16"), + ), + }, + "station", + ) + da = xd.stack(dc, "station", join="outer") + assert np.issubdtype(da.dtype, np.floating) + assert np.isnan(da.values).any() + + def test_inner_join_indexes_when_the_overlap_is_not_contiguous(self): + common = xd.DataArray( + np.arange(3.0), {"time": [0.0, 2.0, 3.0], "station": (None, "SX02")} + ) + dc = xd.DataCollection( + { + "SX01": xd.DataArray( + np.arange(4.0), + {"time": [0.0, 1.0, 2.0, 3.0], "station": (None, "SX01")}, + ), + "SX02": common, + }, + "station", + ) + da = xd.stack(dc, "station", join="inner") + assert da["time"].values.tolist() == [0.0, 2.0, 3.0] + npt.assert_array_equal(da.values, [[0.0, 2.0, 3.0], [0.0, 1.0, 2.0]]) + + def test_join_reconciles_equivalent_coordinate_representations(self): + from xdas.coordinates import InterpCoordinate + + redundant = trace("SHN") + redundant["time"] = InterpCoordinate( + {"tie_indices": [0, 2, 4], "tie_values": [0.0, 2.0, 4.0]}, "time" + ) + dc = xd.DataCollection({"SHZ": trace("SHZ"), "SHN": redundant}, "channel") + # structurally different tie points describe the same grid: strict + # equality says no, the join says yes + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + da = xd.stack(dc, "channel", join="inner") + assert da.shape == (2, 5) + assert da["time"].values.tolist() == [0.0, 1.0, 2.0, 3.0, 4.0] + + def test_leaves_sharing_no_coordinate_value_raise(self): + dc = xd.DataCollection( + { + "SX01": trace(station="SX01", start=0.0, npts=3), + "SX02": trace(station="SX02", start=10.0, npts=3), + }, + "station", + ) + with pytest.raises(ValueError, match="share no coordinate value"): + xd.stack(dc, "station", join="inner") + + def test_repeated_coordinate_values_refuse_to_align(self): + dc = xd.DataCollection( + { + "SX01": xd.DataArray( + np.arange(3.0), {"time": [0.0, 0.0, 1.0], "station": (None, "A")} + ), + "SX02": xd.DataArray( + np.arange(2.0), {"time": [0.0, 1.0], "station": (None, "B")} + ), + }, + "station", + ) + with pytest.raises(ValueError, match="repeats coordinate values"): + xd.stack(dc, "station", join="inner") + + def test_padding_refuses_a_second_coordinate_along_the_dimension(self): + dc = xd.DataCollection( + { + "SX01": xd.DataArray( + np.arange(3.0), + {"time": [0.0, 1.0, 2.0], "quality": ("time", [1, 1, 1])}, + ), + "SX02": xd.DataArray( + np.arange(3.0), + {"time": [2.0, 3.0, 4.0], "quality": ("time", [1, 1, 1])}, + ), + }, + "station", + ) + with pytest.raises(ValueError, match="cannot pad along 'time'"): + xd.stack(dc, "station", join="outer") + + def test_a_join_that_does_not_reconcile_still_raises(self): + # aligning `time` cannot make the second coordinate agree + dc = xd.DataCollection( + { + "SX01": xd.DataArray( + np.arange(3.0), + {"time": [0.0, 1.0, 2.0], "quality": ("time", [1, 1, 1])}, + ), + "SX02": xd.DataArray( + np.arange(4.0), + {"time": [0.0, 1.0, 2.0, 3.0], "quality": ("time", [2, 2, 2, 2])}, + ), + }, + "station", + ) + with pytest.raises(ValueError, match="coordinate 'quality' differs"): + xd.stack(dc, "station", join="inner") + + # --- grid snapping --- + + def test_a_subsample_offset_is_the_same_coordinate(self): + # 1e-3 of a sample apart: one grid, two roundings of it + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=1e-3)}, + "channel", + ) + da = xd.stack(dc, "channel") + assert da.shape == (2, 5) + assert da["time"].values.tolist() == [0.0, 1.0, 2.0, 3.0, 4.0] + npt.assert_array_equal(da.values, np.tile(np.arange(5.0), (2, 1))) + + def test_the_first_leaf_of_the_level_wins_the_grid(self): + # the offset leaf comes first here, so its rounding is the one kept — + # the level's own key order decides, not the sorted output order + dc = xd.DataCollection( + {"SHN": regular_trace("SHN", start=1e-3), "SHZ": regular_trace("SHZ")}, + "channel", + ) + da = xd.stack(dc, "channel") + npt.assert_allclose(da["time"].values, np.arange(5.0) + 1e-3) + + def test_a_full_sample_offset_still_raises(self): + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=1.0)}, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + + def test_a_half_sample_offset_still_raises(self): + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=0.5)}, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + + def test_tolerance_false_restores_strict_equality(self): + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=1e-3)}, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel", tolerance=False) + + def test_an_undeclared_grid_is_not_snapped_by_default(self): + # `trace` ties values to indices without declaring a sampling interval: + # there is no grid to snap to, and a fraction of a sample means nothing + dc = xd.DataCollection( + {"SHZ": trace("SHZ"), "SHN": trace("SHN", start=1e-3)}, "channel" + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + da = xd.stack(dc, "channel", tolerance=1e-2) + assert da["time"].values.tolist() == [0.0, 1.0, 2.0, 3.0, 4.0] + + def test_an_explicit_tolerance_is_absolute(self): + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=0.4)}, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel", tolerance=0.1) + assert xd.stack(dc, "channel", tolerance=0.5).shape == (2, 5) + + def test_datetime_grids_snap_within_a_fraction_of_a_sample(self): + def datetime_trace(channel, start): + return xd.DataArray( + np.arange(4.0), + { + "channel": (None, channel), + "time": { + "tie_indices": [0, 3], + "tie_values": [ + np.datetime64(start, "ns"), + np.datetime64(start, "ns") + np.timedelta64(75, "ms"), + ], + "sampling_interval": np.timedelta64(25, "ms"), + }, + }, + ) + + # the reference dataset's own mismatch: one nanosecond at 40 Hz + dc = xd.DataCollection( + { + "SHE": datetime_trace("SHE", "2026-05-20T00:00:00.000000000"), + "SHZ": datetime_trace("SHZ", "2026-05-19T23:59:59.999999999"), + }, + "channel", + ) + da = xd.stack(dc, "channel") + assert da.shape == (2, 4) + assert da["time"][0].values == np.datetime64("2026-05-20T00:00:00.000000000") + # an explicit tolerance is in seconds, as everywhere else + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel", tolerance=1e-12) + + def test_a_dense_coordinate_snaps_onto_a_declared_grid(self): + dense = xd.DataArray( + np.arange(5.0), + {"channel": (None, "SHN"), "time": (np.arange(5.0) + 1e-3).tolist()}, + ) + dc = xd.DataCollection({"SHZ": regular_trace("SHZ"), "SHN": dense}, "channel") + da = xd.stack(dc, "channel") + assert da["time"].values.tolist() == [0.0, 1.0, 2.0, 3.0, 4.0] + assert da["time"].isregular() + + def test_the_budget_can_come_from_the_snapped_leaf(self): + # the reference declares no spacing; the other one does, and that is + # the grid the tolerance is a fraction of + dense = xd.DataArray( + np.arange(5.0), + {"channel": (None, "SHE"), "time": (np.arange(5.0) + 1e-3).tolist()}, + ) + dc = xd.DataCollection({"SHE": dense, "SHZ": regular_trace("SHZ")}, "channel") + da = xd.stack(dc, "channel") + npt.assert_allclose(da["time"].values, np.arange(5.0) + 1e-3) + assert not da["time"].isregular() + + def test_a_sampled_coordinate_snaps_segment_by_segment(self): + from xdas.coordinates import SampledCoordinate + + def segmented(channel, start): + return xd.DataArray( + np.arange(10.0), + { + "channel": (None, channel), + "time": SampledCoordinate( + { + "tie_values": [start, start + 10.0], + "tie_lengths": [5, 5], + "sampling_interval": 1.0, + }, + "time", + ), + }, + ) + + dc = xd.DataCollection( + {"SHZ": segmented("SHZ", 0.0), "SHN": segmented("SHN", 1e-3)}, "channel" + ) + da = xd.stack(dc, "channel") + assert da.shape == (2, 10) + assert da["time"].values.tolist() == [0, 1, 2, 3, 4, 10, 11, 12, 13, 14] + # the gap between the two segments is part of the grid: move it and the + # coordinates are no longer the same one + moved = xd.DataCollection( + {"SHZ": segmented("SHZ", 0.0), "SHN": segmented("SHN", 0.0)}, "channel" + ) + moved["SHN"]["time"] = SampledCoordinate( + { + "tie_values": [0.0, 11.0], + "tie_lengths": [5, 5], + "sampling_interval": 1.0, + }, + "time", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(moved, "channel") + + def test_coordinates_of_different_length_are_not_snapped(self): + # snapping never changes a length; a ragged pair stays a join's business + dc = xd.DataCollection( + { + "SHZ": regular_trace("SHZ", npts=5), + "SHN": regular_trace("SHN", npts=4, start=1e-3), + }, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + + def test_a_string_coordinate_is_never_snapped(self): + def labelled(channel, labels): + return xd.DataArray( + np.arange(3.0), + { + "channel": (None, channel), + "time": [0.0, 1.0, 2.0], + "label": ("time", labels), + }, + ) + + dc = xd.DataCollection( + { + "SHZ": labelled("SHZ", ["a", "b", "c"]), + "SHN": labelled("SHN", ["a", "b", "d"]), + }, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'label' differs"): + xd.stack(dc, "channel") + + def test_a_datetime_and_a_numeric_coordinate_are_never_snapped(self): + dc = xd.DataCollection( + { + "SHZ": xd.DataArray(np.arange(3.0), {"time": [0.0, 1.0, 2.0]}), + "SHN": xd.DataArray( + np.arange(3.0), + {"time": np.arange("2026-01-01", 3, dtype="datetime64[s]")}, + ), + }, + "channel", + ) + with pytest.raises(ValueError, match="coordinate 'time' differs"): + xd.stack(dc, "channel") + + def test_outer_join_refuses_to_interleave_grids(self): + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=0.5)}, + "channel", + ) + with pytest.raises(ValueError, match="would interleave 10 samples"): + xd.stack(dc, "channel", join="outer") + + def test_outer_join_still_extends_a_shared_grid(self): + dc = xd.DataCollection( + {"SHZ": regular_trace("SHZ"), "SHN": regular_trace("SHN", start=2.0)}, + "channel", + ) + da = xd.stack(dc, "channel", join="outer") + assert da.shape == (2, 7) + assert da["time"].values.tolist() == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + + # --- laziness --- + + def test_tile_backed_leaves_stay_virtual(self, tmp_path): + from xdas.virtual import TileArray + + expected = xd.testing.dummy(dims=("time", "space"), shape=(10, 5)) + expected.to_netcdf(tmp_path / "chunk.nc") + dc = xd.DataCollection( + { + code: xd.open_dataarray( + tmp_path / "chunk.nc", engine="xdas", vtype="tiles" + ) + for code in ("SHZ", "SHN") + }, + "channel", + ) + da = xd.stack(dc, "channel") + assert isinstance(da.data, TileArray) + assert da.dims == ("channel", "time", "space") + assert da["channel"].values.tolist() == ["SHN", "SHZ"] + npt.assert_array_equal(np.asarray(da.data)[0], expected.values) + + +class TestConcatNewDimVirtual: + def test_a_new_dimension_over_virtual_sources_loads_instead_of_raising( + self, tmp_path + ): + # `expand_dims` cannot follow a `VirtualSource` — a stack of sources is + # a longer axis, never an extra one — so the result is dense + expected = xd.testing.dummy(dims=("time", "space"), shape=(10, 5)) + expected.to_netcdf(tmp_path / "chunk.nc") + objs = [xd.open_dataarray(tmp_path / "chunk.nc") for _ in range(2)] + assert all(isinstance(da.data, VirtualSource) for da in objs) + da = xd.concat(objs, "channel") + assert da.dims == ("channel", "time", "space") + assert isinstance(da.data, np.ndarray) diff --git a/tests/test_testing.py b/tests/test_testing.py index 42a1ec11..7b9c1757 100644 --- a/tests/test_testing.py +++ b/tests/test_testing.py @@ -1,7 +1,10 @@ import numpy as np +import pandas as pd import pytest import xdas as xd +from xdas.atoms import Partial +from xdas.testing import _assert_same class TestDummy: @@ -23,3 +26,160 @@ def test_mismatched_step(self): def test_datetime_step_passthrough(self): da = xd.testing.dummy(step=(np.timedelta64(10, "ms"), 10.0)) assert da["time"].get_sampling_interval() == 0.01 + + +def _shift_after_the_first_chunk(da): + """Move the time axis by a nanosecond, but not on the first call.""" + seen = getattr(_shift_after_the_first_chunk, "seen", 0) + 1 + _shift_after_the_first_chunk.seen = seen + if seen == 1: + return da + coords = dict(da.coords) + coords["time"] = da["time"].values + np.timedelta64(1, "ns") + return xd.DataArray(da.values, coords, da.dims) + + +class TestAssertChunkInvariant: + def test_elementwise_pipeline_passes(self): + da = xd.testing.dummy() + xd.testing.assert_chunk_invariant(Partial(np.square), da, {"time": 25}) + + def test_stateful_pipeline_passes(self): + da = xd.testing.dummy() + pipeline = xd.filter(..., (None, 10.0), dim="time") + xd.testing.assert_chunk_invariant(pipeline, da, {"time": 25}) + + def test_chunk_hostile_pipeline_fails(self): + def normalize(da): + return da / np.std(da.values) + + da = xd.testing.dummy() + with pytest.raises(AssertionError, match="Not equal to tolerance"): + xd.testing.assert_chunk_invariant(Partial(normalize), da, {"time": 25}) + + def test_reports_a_shape_difference(self): + def drop_last(da): + return da.isel(time=slice(None, -1)) + + da = xd.testing.dummy() + with pytest.raises(AssertionError, match="shape differs"): + xd.testing.assert_chunk_invariant(Partial(drop_last), da, {"time": 25}) + + def test_reports_a_dims_difference(self): + da = xd.testing.dummy() + atom = Partial(lambda x: x.mean("time")) + with pytest.raises(AssertionError, match="dims differ"): + xd.testing.assert_chunk_invariant(atom, da, {"time": 25}) + + def test_reports_a_coordinate_difference(self): + da = xd.testing.dummy() + _shift_after_the_first_chunk.seen = 0 + atom = Partial(_shift_after_the_first_chunk) + with pytest.raises(AssertionError, match="coordinate differs"): + xd.testing.assert_chunk_invariant(atom, da, {"time": 25}) + + def test_coord_atol_admits_a_sub_sample_drift(self): + da = xd.testing.dummy() + _shift_after_the_first_chunk.seen = 0 + atom = Partial(_shift_after_the_first_chunk) + xd.testing.assert_chunk_invariant(atom, da, {"time": 25}, coord_atol=2) + + def test_compares_sequences_of_chunks(self): + left = [xd.testing.dummy(shape=(4, 2)), xd.testing.dummy(shape=(4, 2))] + _assert_same(left, list(left), 1e-7, 0.0, 0, "result") + + def test_reports_a_chunk_count_difference(self): + left = [xd.testing.dummy(shape=(4, 2)), xd.testing.dummy(shape=(4, 2))] + with pytest.raises(AssertionError, match="2 chunks eager vs 1 chunked"): + _assert_same(left, left[:1], 1e-7, 0.0, 0, "result") + + def test_reports_a_result_that_did_not_join(self): + eager = xd.testing.dummy(shape=(4, 2)) + with pytest.raises(AssertionError, match="did not join into"): + _assert_same(eager, [eager], 1e-7, 0.0, 0, "result") + + def test_compares_bare_values(self): + _assert_same(1.0, 1.0, 1e-7, 0.0, 0, "result") + + def test_reports_a_table_that_is_not_a_table(self): + frame = pd.DataFrame({"a": [1.0]}) + with pytest.raises(AssertionError, match="chunked gave a DataArray"): + _assert_same(frame, xd.testing.dummy(shape=(2, 2)), 1e-7, 0.0, 0, "result") + + def test_reports_tables_whose_columns_differ(self): + frame = pd.DataFrame({"a": [1.0]}) + with pytest.raises(AssertionError, match="columns differ"): + _assert_same(frame, pd.DataFrame({"b": [1.0]}), 1e-7, 0.0, 0, "result") + + def test_compares_pick_tables_as_sets_of_rows(self): + # Eager processing walks lane by lane, chunked walks chunk by chunk, so + # the rows arrive in a different order with the same content. + frame = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]}) + _assert_same(frame, frame.iloc[::-1], 1e-7, 0.0, 0, "result") + + def test_gaps_exercise_the_seams(self): + da = xd.testing.dummy() + pipeline = xd.filter(..., (None, 10.0), dim="time") + xd.testing.assert_chunk_invariant(pipeline, da, {"time": 25}, gaps=2) + + def test_gaps_at_explicit_positions(self): + da = xd.testing.dummy() + xd.testing.assert_chunk_invariant( + Partial(np.square), da, {"time": 25}, gaps=[10, 60] + ) + + def test_gaps_shrink_an_oversized_chunking(self): + da = xd.testing.dummy() + # 100 samples minus two gaps leaves 90: the requested size must clamp. + xd.testing.assert_chunk_invariant(Partial(np.square), da, {"time": 100}, gaps=2) + + def test_cut_invariance_catches_a_cut_sensitive_atom(self): + def anchor_on_chunk_start(da): + return da - da.values[0] + + da = xd.testing.dummy() + atom = Partial(anchor_on_chunk_start) + with pytest.raises(AssertionError): + xd.testing.assert_chunk_invariant(atom, da, {"time": 100}) + + def test_explicit_cuts(self): + da = xd.testing.dummy() + xd.testing.assert_chunk_invariant( + Partial(np.square), da, {"time": 25}, cuts=[{"time": 13}, {"time": 7}] + ) + + def test_cuts_zero_restores_the_single_split(self): + da = xd.testing.dummy() + xd.testing.assert_chunk_invariant(Partial(np.square), da, {"time": 25}, cuts=0) + + def test_inject_gaps_makes_discontinuities(self): + da = xd.testing.dummy() + gappy = xd.testing.inject_gaps(da, "time", 2) + assert gappy.sizes["time"] == 90 + coord = gappy["time"] + indices = coord.get_split_indices("discontinuities", coord.tolerance) + assert len(indices) == 2 + + def test_inject_gaps_refuses_to_eat_the_record(self): + da = xd.testing.dummy(shape=(4, 2)) + with pytest.raises(ValueError, match="less than two pieces"): + xd.testing.inject_gaps(da, "time", [0]) + + def test_cuttings_stop_when_no_new_size_exists(self): + from xdas.testing import _cuttings + + # size 2 derives 2 // 2 + 1 == 2 again: nothing new to cut with. + assert _cuttings(3, "time", 2) == [] + + def test_a_dim_without_coordinate_is_skipped(self): + left = xd.DataArray( + np.zeros((3, 2)), {"time": [0.0, 1.0, 2.0]}, dims=("time", "distance") + ) + _assert_same(left, left.copy(), 1e-7, 0.0, 0, "result") + + def test_a_bare_callable_is_not_a_pipeline(self): + # It cannot be streamed at all — the chunked path hands each chunk a + # `chunk_dim` keyword, which a plain function does not accept. + da = xd.testing.dummy() + with pytest.raises(AttributeError, match="reset"): + xd.testing.assert_chunk_invariant(np.square, da, {"time": 25}) diff --git a/tests/test_trigger.py b/tests/test_trigger.py index ff6b8ed0..468da68a 100644 --- a/tests/test_trigger.py +++ b/tests/test_trigger.py @@ -2,35 +2,7 @@ import pandas as pd import xdas as xd -from xdas.trigger import Trigger, _find_picks_numeric, find_picks - - -def test_trigger(): - # test case - cft = xd.DataArray( - data=[[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2]], - coords={ - "space": [0.0], - "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0]}, - }, - ) - - # test monolithic processing - picks = Trigger(thresh=0.5, dim="time")(cft) - expected = pd.DataFrame( - {"space": [0.0, 0.0], "time": [2.0, 7.0], "value": [0.9, 0.7]} - ) - assert picks.equals(expected) - - # test chunked processing - trigger = Trigger(thresh=0.5, dim="time") - chunks = xd.split(cft, 3, dim="time") - result = [] - for chunk in chunks: - picks = trigger(chunk, chunk_dim="time") - result.append(picks) - result = pd.concat(result, ignore_index=True) - assert result.equals(expected) +from xdas.trigger import _find_picks_numeric, find_picks def test_find_picks_numeric(): @@ -134,7 +106,11 @@ def test_find_picks(): data=[[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2]], coords={ "space": [0.0], - "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0]}, + "time": { + "tie_indices": [0, 9], + "tie_values": [0.0, 9.0], + "sampling_interval": 1.0, + }, }, ) @@ -164,16 +140,3 @@ def test_find_picks(): result.append(atom(chunk, chunk_dim="time")) result = pd.concat(result, ignore_index=True) assert result.equals(expected) - - -def test_trigger_1d(): - """1D input (no spatial dimension) covers the coords=() branch in _call_numeric.""" - cft = xd.DataArray( - data=[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2], - coords={ - "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0]}, - }, - ) - picks = Trigger(thresh=0.5, dim="time")(cft) - assert len(picks) == 2 - assert list(picks["time"]) == [2.0, 7.0] diff --git a/tests/test_xdas.py b/tests/test_xdas.py index 775815f8..a99487ce 100644 --- a/tests/test_xdas.py +++ b/tests/test_xdas.py @@ -1,5 +1,7 @@ import re +import numpy as np + import xdas as xd # Release segment, plus the optional PEP 440 pre/post/dev markers (e.g. 0.2.9.dev0). @@ -10,3 +12,16 @@ def test_version(): version = xd.__version__ assert isinstance(version, str) assert VERSION_PATTERN.match(version) + + +class TestSynthetics: + def test_randn_wavefronts_contract(self): + from xdas.synthetics import randn_wavefronts + + da = randn_wavefronts() + assert da.dims == ("time", "distance") + assert da.sizes == {"time": 20000, "distance": 1001} + assert da["time"][0].values == np.datetime64("2024-01-01T00:00:00", "ns") + assert float(da["distance"][-1].values) == 100000.0 + # seeded: two calls give the same wavefronts + assert np.array_equal(da.values[:100], randn_wavefronts().values[:100]) diff --git a/xdas/__init__.py b/xdas/__init__.py index 0145ba61..7a324b84 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -57,7 +57,27 @@ "plot_availability", "sortby", "split", + "stack", "trim_overlaps", + # task atoms (function forms) + "annotate", + "decimate", + "detrend", + "differentiate", + "filter", + "hilbert", + "integrate", + "medfilt", + "mlpicker", + "pick", + "rechunk", + "resample", + "sliding_mean_removal", + "stft", + "taper", + "trigger", + # streaming + "watch", ] from . import ( @@ -73,6 +93,27 @@ testing, virtual, ) + +# The compat module first, so that a later `import xdas.trigger` finds it in +# `sys.modules` and does not rebind the attribute: the lowercase twin below +# stays `xdas.trigger` for everyone. +from . import trigger as _trigger_module # noqa: F401 isort: skip +from .atoms.detect import trigger +from .atoms.kernel import rechunk +from .atoms.ml import annotate, mlpicker, pick +from .atoms.tasks import ( + decimate, + detrend, + differentiate, + filter, + hilbert, + integrate, + medfilt, + resample, + sliding_mean_removal, + stft, + taper, +) from .coordinates import ( Coordinate, Coordinates, @@ -110,6 +151,8 @@ routines, sortby, split, + stack, trim_overlaps, ) from .core.methods import * +from .processing.core import watch diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index 7a251e2b..89be6e6f 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -1,36 +1,55 @@ """ Stateful processing units (atoms) for building chunked data pipelines. -Exports :class:`Atom`, :class:`State`, :class:`Sequential`, :class:`Partial`, -:func:`atomized`, signal-processing atoms, and the ML-based :class:`MLPicker`. +Three layers: + +- :mod:`xdas.atoms.core`: the machinery — :class:`Atom`, :class:`State`, + :class:`Sequential`, :class:`Partial`, :func:`atomized`, + :func:`as_function`, :func:`compose`. +- :mod:`xdas.atoms.kernel`: the expert layer — exact stateful primitives with + machine parameters (:class:`LFilter`, :class:`SOSFilter`, + :class:`DownSample`, :class:`UpSample`, :class:`Polyphase`). +- :mod:`xdas.atoms.tasks`: the public layer — task atoms with physical + parameters only (:class:`Filter`, :class:`Decimate`, :class:`Resample`, + ...), each with a function form exported at the top level of :mod:`xdas`. + +Plus the signal-processing atoms of :mod:`xdas.atoms.signal` and the ML-based +:class:`Annotate`. """ __all__ = [ + "STFT", + "Annotate", "Atom", + "Decimate", + "Differentiate", "DownSample", "FIRFilter", + "Filter", "IIRFilter", + "Integrate", "LFilter", "MLPicker", "Partial", + "Picker", + "Polyphase", + "Rechunk", + "Resample", "ResamplePoly", "SOSFilter", "Sequential", "State", "Trigger", "UpSample", + "as_function", "atomized", + "compose", + "trigger", ] -from ..trigger import Trigger -from .core import Atom, Partial, Sequential, State, atomized -from .ml import MLPicker -from .signal import ( - DownSample, - FIRFilter, - IIRFilter, - LFilter, - ResamplePoly, - SOSFilter, - UpSample, -) +from .core import Atom, Partial, Sequential, State, as_function, atomized, compose +from .detect import Trigger, trigger +from .kernel import DownSample, LFilter, Polyphase, Rechunk, SOSFilter, UpSample +from .ml import Annotate, MLPicker, Picker +from .signal import FIRFilter, IIRFilter, ResamplePoly +from .tasks import STFT, Decimate, Differentiate, Filter, Integrate, Resample diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index 6bf0c88d..252127eb 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -2,20 +2,174 @@ Base classes for stateful processing atoms. Includes :class:`Atom`, :class:`State`, :class:`Sequential`, :class:`Partial`, -and the :func:`atomized` decorator. +the :func:`atomized` decorator, the :func:`as_function` generator of the +function form of an atom class, and the :func:`compose` primitive behind the +``>>`` operator. """ import importlib +import inspect +import re +import warnings from collections.abc import Callable from functools import wraps from typing import Any -from ..core import DataArray, DataCollection, open_datacollection +import numpy as np +import pandas as pd + +from .. import config +from ..coordinates import AxisCoordinate +from ..coordinates.core import parse_scalar_delta +from ..core import ( + DataArray, + DataCollection, + DataMapping, + DataSequence, + concat, + open_datacollection, + split, +) +from ..virtual import VirtualBackend + + +def _announce_splits(x, dim, count): + """ + Warn that a source will be processed as several runs. + + Splitting on discontinuities must not be silent: each reset restarts the + warm-up of every stateful stage, so the user should know how often it + happens. The message names the source by its start so a collection walk + reports every leaf rather than being deduplicated to the first. + """ + start = x.coords[dim].start if dim in getattr(x, "coords", {}) else "?" + plural = "discontinuities" if count > 1 else "discontinuity" + warnings.warn( + f"source starting at {start} has {count} {plural} along {dim!r}; " + "state is flushed and reset at each", + UserWarning, + stacklevel=3, + ) + + +def _aschunks(value): + """ + Normalize an atom output into a list of chunks. + + Atoms follow the transducer contract: ``call()`` maps one input chunk to + zero or more output chunks. A bare object is one chunk, ``None`` is zero + chunks, and a list or :class:`DataSequence` is taken chunk by chunk. Empty + chunks are dropped: they carry no information and their degenerate + coordinates would poison the initialization of downstream atoms. + """ + if value is None: + return [] + if not isinstance(value, (list, DataSequence)): + value = [value] + return [ + chunk for chunk in value if not (isinstance(chunk, DataArray) and chunk.empty) + ] + + +def _join_chunks(chunks, dim=None): + """ + Join heterogeneous output chunks into a single result. + + DataArrays are concatenated along *dim* with the gaps kept in the + coordinates, falling back to a :class:`DataSequence` when concatenation + cannot represent the result in one array; DataFrames are concatenated + with a fresh index. Anything else is returned as a plain list. Zero + chunks give ``None`` and a single chunk is returned bare. + """ + if not chunks: + return None + if len(chunks) == 1: + return chunks[0] + if all(isinstance(chunk, DataArray) for chunk in chunks): + if dim is not None: + try: + return concat(chunks, dim) + except (TypeError, ValueError, KeyError): + pass + return DataCollection(chunks) + if all(isinstance(chunk, pd.DataFrame) for chunk in chunks): + return pd.concat(chunks, ignore_index=True) + return list(chunks) + + +def _extend_path(path, name, key): + """Extend the tree path with the *key* a *name*d collection level was entered by.""" + return path if name is None else path | {name: key} + + +def _annotate_path(result, path): + """ + Tag the tables in *result* with the tree path they were produced under. + + A leaf reached through ``IA / DBNFM / --`` gets a ``network``, a + ``station`` and a ``location`` column, filled with those keys, *as its + result is produced* rather than reconstructed afterwards — which is what + lets a streaming walk hand a leaf straight to a sink and still carry its + identity. Only tables are annotated: an atom returning arrays sees + nothing change, and its tree is rebuilt as before. + + The columns lead, in tree order, so the identity of a pick comes first + whatever the atom put in the table. A column the table already carries — + a `network` scalar coordinate on the leaf, say, which the ObsPy engine + attaches — is *the same column*, not a second one: it is moved into its + leading position and the tree path wins, warning if the two disagree. + + *result* is what one leaf produced, so it is a table, an array, or one of + the containers a leaf can answer with: a :class:`DataSequence` of chunks + the atom could not join, or a plain list of chunks. It is never a mapping + level — those the walk recurses into itself. + """ + if not path: + return result + if isinstance(result, pd.DataFrame): + return _annotate_frame(result, path) + if isinstance(result, DataSequence): + return DataCollection( + [_annotate_path(value, path) for value in result], result.name + ) + if isinstance(result, list): + return [_annotate_path(value, path) for value in result] + return result + + +def _annotate_frame(frame, path): + """Prepend the *path* identity columns to the table *frame*.""" + frame = frame.copy() + for name, key in path.items(): + if name in frame.columns and not (frame[name] == key).all(): + warnings.warn( + f"the {name!r} column of a table disagrees with the tree path " + f"the leaf was reached by ({key!r}): the tree path wins. The " + "column comes from a coordinate of the leaf; rename it or " + "drop it from `coords` to keep both.", + UserWarning, + stacklevel=2, + ) + frame[name] = key + rest = [name for name in frame.columns if name not in path] + return frame[list(path) + rest] + + +def _iter_results(tree): + """Yield the leaf results of a walked collection, in walk order.""" + if isinstance(tree, DataMapping): + for value in tree.values(): + yield from _iter_results(value) + elif isinstance(tree, DataSequence): + for value in tree: + yield from _iter_results(value) + else: + yield tree class State: """ - A class to declare a new state or to update a preexising one into an Atom object. + A class to declare a new state or to update a preexisting one into an Atom object. Parameters ---------- @@ -24,7 +178,7 @@ class State: Examples -------- - In practice the State object is used when implementing new Atom objects. Bellow a + In practice the State object is used when implementing new Atom objects. Below a dummy example without any class declaration. >>> from xdas.atoms import Atom, State @@ -53,7 +207,7 @@ def __init__(self, state): self.state = state -class Atom: +class Atom(np.lib.mixins.NDArrayOperatorsMixin): """ The base class for atoms. Used to implement new Atom objects. @@ -66,48 +220,85 @@ class Atom: that is updated during each execution of the Atom. The Atom object is initialized with an initial state at the first call. In subsequent calls, the state is updated **if and only if the `chunk_dim` flag is provided** along with the dimension along - wich chunking was performed. If this flag is not provided, the state is reset + which chunking was performed. If this flag is not provided, the state is reset between calls. The Atom can be reset manually to its initial state by calling the `reset` method. - When implementing a new Atom object, the user must sublcass the Atom class, and + When implementing a new Atom object, the user must subclass the Atom class, and at minima define the `initialize` and the `call` methods. The `initialize` method is called at the first call to the Atom and is used to initialize the Atom with the input data. The `call` method is called at each subsequent call to the Atom and is used to perform the main processing logic. The Atom class handles when an how those two methods are called. - To reduce the size of the state that need to be stored, a good practive is to also + To reduce the size of the state that need to be stored, a good practice is to also define the `initialize_from_state` method. This method is called in the `initialize` as soon as the minimal set of states is initialized. The other states - that are usefull for the processing but that can be recomputed from the minimal set + that are useful for the processing but that can be recomputed from the minimal set are initialized in the `initialize_from_state` method. + Atoms compose into pipelines with the ``>>`` operator (see :func:`compose`) + and trace ordinary numpy expressions: applying a ufunc to an atom appends + the operation to the pipeline instead of computing it. + Attributes ---------- state: dict Returns the current state of the atom recursively including the state of nested atoms. initialized: bool - Wether the atom has been initialized or not. + Whether the atom has been initialized or not. + on_discontinuity: str + Seam policy for chunked processing: ``"reset"`` (default) flushes + and starts a new run at every gap or rate change, ``"raise"`` + refuses discontinuous input. Overlaps always raise. Methods ------- + gather(mapping) + Optionally collapse a mapping level into an array before the + walk descends into it. See :meth:`gather`. initialize(x, **flags) Initializes the atom with the given input. initialize_from_state() Initializes the atom from its minimal state. call(x, **flags) - Performs the main processing logic of the atom. + Performs the main processing logic of the atom. May return zero + or more output chunks (the transducer contract). + flush() + Drains buffered samples at the end of a run. reset() Resets the atom to its initial state. + merge(results) + Optional. Folds the leaf results of a collection walk into one + object. ``None`` by default, meaning the atom declares none and + the walked tree is rebuilt as it always was; see + :meth:`~xdas.atoms.Trigger.merge` for an implementation. + fresh() + Returns a stateless clone sharing the configuration. + iter_chunks(source) + Streams a chunk source through the atom, seams and flush included. """ + on_discontinuity = "reset" + #: Undefined by default: an atom that does not fold its collection + #: results leaves the walked tree as it is. + merge = None + def __init__(self): object.__setattr__(self, "_config", {}) object.__setattr__(self, "_state", {}) object.__setattr__(self, "_atoms", {}) + object.__setattr__(self, "_seam", None) + + def __eq__(self, other): + return self is other + + def __ne__(self, other): + return self is not other + + __hash__ = object.__hash__ def __repr__(self): name = self.__class__.__name__ @@ -140,11 +331,13 @@ def state(self): @property def initialized(self): - """``True`` if every state key has been initialised (no ``...`` sentinels remain).""" - return all(value is not ... for value in self._state.values()) + """``True`` if every state key, nested atoms included, is initialised.""" + return all(value is not ... for value in self._state.values()) and all( + atom.initialized for atom in self._atoms.values() + ) def initialize(self, x, **flags): - """Initialise the atom from a first chunks of data.""" + """Initialise the atom from a first chunk of data.""" return NotImplemented def initialize_from_state(self): @@ -155,23 +348,568 @@ def call(self, x, **flags): """Process a chunk of data.""" return NotImplemented + def gather(self, mapping): + """ + Return *mapping* collapsed to an array, or ``None`` to map over it. + + Consulted on every mapping level of a collection *before* the walk + descends into it, so an atom that knows a level is really an axis of + its input can take it as one. :class:`~xdas.atoms.Annotate` is the + implementation: it knows the model's ``component_order``, so it knows + that a ``channel`` level keyed ``SHZ``/``SHN``/``SHE`` is the + component dimension of one instrument rather than three + independent leaves, and collapses it with :func:`xdas.stack`. + + Returning ``None`` — the default, and what every other atom does — + leaves the walk exactly as it was. + + Parameters + ---------- + mapping : DataMapping + The collection level about to be walked. + + Returns + ------- + DataArray, DataCollection or None + The collapsed input, which the walk continues on in place of the + level, or ``None`` to map over the level's leaves. + """ + return + def __call__(self, x, **flags): - """Process input data, initializing state if needed and resetting after final chunk.""" + """ + Process input data, returning zero or more output chunks. + + Eager calls (no ``chunk_dim`` flag) auto-split gappy input into runs, + process each run with a fresh state, flush the tails and re-join the + outputs into a single object with gap-aware coordinates when possible + (falling back to a :class:`DataSequence`). Chunked calls (``chunk_dim`` + given) carry state across continuous chunks and handle seams: on a gap + or a rate change the atom flushes, resets and starts a new run (see + `on_discontinuity`); on an overlap it raises. Sequence collections are + folded element by element through the same seam-aware machinery, so + resets emerge from the coordinates; mapping collections map over their + leaves. + + Walking a collection, each leaf's result is annotated with the tree + path it was produced under — one column per level, leading, filled + with the key the level was entered by — and the annotated results are + then folded by the atom's `merge` hook if it declares one, so a + table-valued atom answers a whole collection with one table rather + than with a tree of them. ``merge=False`` opts out and returns the + walked tree, annotations included. Mapping levels are first offered + to `gather`, which may collapse a level into an axis of the input + (see :meth:`gather`). + + A single output chunk is returned bare; several are joined into one + object when their kind allows it (arrays concatenated along the + working dimension, tables into one table) and returned as a + :class:`DataSequence` of chunks when it does not. + """ + merge = flags.pop("merge", True) chunk_dim = flags.get("chunk_dim", None) - if not self.initialized or chunk_dim is None: + self._check_chunk_dim(x, chunk_dim) + if ( + chunk_dim is None + and isinstance(x, DataArray) + and isinstance(x.data, VirtualBackend) + and x.nbytes > config.get("memory_limit") + ): + limit = config.get("memory_limit") + raise ValueError( + f"this eager call would load the full virtual array " + f"(~{x.nbytes / 2**30:.1f} GiB) in memory, above the " + f"'memory_limit' configuration entry ({limit / 2**30:.1f} GiB): " + "stream it chunk by chunk with `.process(da, out=...)` instead, " + "or raise the limit with `xdas.config.set('memory_limit', ...)`" + ) + if isinstance(x, (DataMapping, DataSequence)): + if isinstance(x, DataMapping): + result = self._walk(x, flags, {}) + else: + result = self._fold(x, flags, {}) + if merge and self.merge is not None: + return self.merge(list(_iter_results(result))) + return result + if chunk_dim is None: + dim = self._resolve_dim(x) + runs = self._split_runs(x, dim) + if len(runs) > 1: + _announce_splits(x, dim, len(runs) - 1) + chunks = [] + for run in runs: + self.initialize(run, **flags) + chunks += _aschunks(self.call(run, **flags)) + chunks += self.flush() + self.reset() + return self._join(chunks, dim) + else: + chunks = [] + for run in self._split_runs(x, chunk_dim): + chunks += self._call_run(run, flags) + return self._join(chunks, None) + + def _call_run(self, x, flags): + """Seam-aware chunked call on one internally-regular chunk.""" + if isinstance(x, DataArray) and x.empty: + return [] + chunk_dim = flags["chunk_dim"] + chunks = [] + stateful = self._live_state() + if stateful: + info = self._seam_info(x, chunk_dim) + verdict = self._judge_seam(info) + if verdict in ("gap", "rate"): + match self.on_discontinuity: + case "reset": + chunks += self.flush() + self.reset() + case "raise": + raise ValueError( + f"the incoming chunk is discontinuous with the " + f"stream processed so far ({verdict} detected " + f"along {chunk_dim!r}) and `on_discontinuity` is " + "set to 'raise'" + ) + case other: + raise ValueError( + "`on_discontinuity` must be 'reset' or 'raise', " + f"got {other!r}" + ) + elif verdict == "overlap": + raise ValueError( + f"the incoming chunk overlaps the stream processed so far " + f"(the {chunk_dim!r} coordinate goes backward across the " + "seam); sort or deduplicate the input, or call `reset()` " + "to explicitly start a new run" + ) + if not self.initialized: self.initialize(x, **flags) - y = self.call(x, **flags) - if not chunk_dim: + chunks += _aschunks(self.call(x, **flags)) + if stateful and info is not None: + if info["delta"] is None and verdict == "continuous": + info["delta"] = self._seam["delta"] + object.__setattr__(self, "_seam", info) + return chunks + + def _live_state(self): + """Return ``True`` if any state entry (own or nested) holds a live value.""" + + def live(state): + return any( + live(value) if isinstance(value, dict) else value is not None + for value in state.values() + ) + + return live(self.state) + + def _dim_name(self, x): + """ + Resolve the ``"first"``/``"last"`` alias of this atom against *x*. + + The alias never equals a real dimension name, so it has to be + resolved before it names an axis, a coordinate, or is compared with + the chunked dimension. Resolution is against the *dimensions* of *x*, + which a dimension without a coordinate still has. + """ + dim = getattr(self, "dim", None) + if not isinstance(x, DataArray) or not isinstance(dim, str): + return dim + if dim == "first": + return x.dims[0] + if dim == "last": + return x.dims[-1] + return dim + + def _resolve_dim(self, x): + """Resolve the dimension this atom operates along on *x*, or ``None``.""" + dim = self._dim_name(x) + if not isinstance(dim, str): + return None + return dim if dim in getattr(x, "coords", {}) else None + + def _split_runs(self, x, dim): + """Split *x* at the discontinuities of its *dim* coordinate.""" + if not isinstance(x, DataArray) or dim not in getattr(x, "coords", {}): + return [x] + coord = x.coords[dim] + if not isinstance(coord, AxisCoordinate) or not coord.isregular(): + return [x] + indices = coord.get_split_indices( + "discontinuities", getattr(coord, "tolerance", None) + ) + if not indices.size: + return [x] + return list(split(x, indices, dim)) + + def _seam_info(self, x, chunk_dim): + """Extract the seam-judgment metadata of a chunk, or ``None``.""" + if not isinstance(x, DataArray) or chunk_dim not in x.coords: + return None + coord = x.coords[chunk_dim] + if not isinstance(coord, AxisCoordinate) or coord.empty: + return None + return { + "chunk_dim": chunk_dim, + "start": coord.start, + "end": coord.end, + "delta": coord.get_sampling_interval(cast=False), + "tolerance": parse_scalar_delta( + getattr(coord, "tolerance", None), coord.dtype, default_zero=True + ), + "size": len(coord), + } + + def _judge_seam(self, info): + """ + Compare an incoming chunk with the expected continuation of the stream. + + Returns ``None`` when there is nothing to judge against (first chunk, + non-array chunk, or neither side of the seam declaring a rate), else + one of ``"continuous"``, ``"gap"``, ``"rate"`` or ``"overlap"``. Both + O(1) checks of the regularity contract happen here: the sampling + interval must match within tolerance, and the chunk must start one + interval after the previous end within the jitter budget. + """ + seam = self._seam + if info is None or seam is None or seam["chunk_dim"] != info["chunk_dim"]: + return None + for entry in (seam, info): + if entry["delta"] is None and entry["size"] > 1: + dim = info["chunk_dim"] + raise ValueError( + f"chunked processing along {dim!r} requires a regular " + "coordinate (one that declares its `sampling_interval`); " + "regularize it first, e.g. `da[dim] = da[dim].to_regular()` " + "or open the files with a tolerance" + ) + # a one-sample chunk declares no interval of its own: the seam is + # judged on the rate of whichever side knows one + delta = seam["delta"] if seam["delta"] is not None else info["delta"] + if delta is None: + return None + tolerance = max(seam["tolerance"], info["tolerance"]) + if info["delta"] is not None and np.abs(info["delta"] - delta) > tolerance: + return "rate" + jump = info["start"] - (seam["end"] + delta) + if np.abs(jump) <= tolerance: + return "continuous" + return "gap" if jump > 0 else "overlap" + + def _walk(self, x, flags, path): + """ + Walk a collection leaf by leaf, carrying the tree path down. + + Mapping levels are first offered to `gather`, which may take the whole + level as an axis of the input rather than as leaves to map over; the + level is then consumed and contributes no path column. Otherwise + mapping levels recurse under their key, sequence levels fold (see + `_fold`), and every leaf result is annotated with the path it was + produced under before it goes anywhere else. Carrying the path *down* + rather than rebuilding it on the way up is what a streaming walk + needs: a leaf is complete the moment it is produced. + + One atom instance walks the leaves sequentially — the eager path + already resets it at the end of each run — because an atom holding a + model either saturates the CPU or holds a lot of device memory, so + only one should be live per node. + """ + if isinstance(x, DataMapping): + gathered = self.gather(x) + if gathered is not None: + return self._walk(gathered, flags, path) + if flags.get("chunk_dim", None) is not None: + raise NotImplementedError( + "chunked processing of mapping collections is not supported: " + "process each leaf with its own atom instance" + ) + name = getattr(x, "name", None) + return DataCollection( + { + key: self._walk(value, flags, _extend_path(path, name, key)) + for key, value in x.items() + }, + name, + ) + if isinstance(x, DataSequence): + return self._fold(x, flags, path) + return _annotate_path(self(x, **flags), path) + + def _fold(self, x, flags, path=None): + """ + Fold a sequence collection through the same seam-aware call. + + A collection is multiple chunks delivered at once: each element goes + through the chunked path along the atom's dimension, so state carries + across continuous elements and resets emerge from the coordinates. + + The level contributes its positional keys as a column, each output + chunk taking the index of the element that produced it. The flushed + tail is attributed to the last element, which is where it came out; + no finer answer exists, since a folded element's buffered samples are + released by the element that follows it. + """ + name = getattr(x, "name", None) + path = {} if path is None else path + chunk_dim = flags.get("chunk_dim", None) + if chunk_dim is None: + first = next((el for el in x if isinstance(el, DataArray)), None) + dim = self._resolve_dim(first) + if dim is None: + return DataCollection( + [ + self._walk(el, flags, _extend_path(path, name, index)) + for index, el in enumerate(x) + ], + name, + ) + flags = flags | {"chunk_dim": dim} + chunks = [] + for index, el in enumerate(x): + chunks += _annotate_path( + _aschunks(self(el, **flags)), _extend_path(path, name, index) + ) + chunks += _annotate_path( + self.flush(), _extend_path(path, name, max(len(x) - 1, 0)) + ) self.reset() - return y + return DataCollection(chunks, name) + chunks = [] + for index, el in enumerate(x): + chunks += _annotate_path( + _aschunks(self(el, **flags)), _extend_path(path, name, index) + ) + return DataCollection(chunks, name) + + def _join(self, chunks, dim): + """Re-join output chunks: one chunk bare, else gap-aware concat or sequence.""" + # a call answers with an object, so zero chunks and chunk types with no + # join of their own become an empty and a plain sequence respectively + if not chunks: + return DataCollection([]) + result = _join_chunks(chunks, dim) + return DataCollection(result) if isinstance(result, list) else result + + def flush(self): + """ + Drain buffered samples, returning zero or more output chunks. + + Stateful atoms that hold samples back waiting for the next chunk + override this to emit what remains computable at the end of a run. + Called at the end of the stream, at every seam, and at the end of + every eager call. Default is a no-op. + """ + return [] + + def iter_chunks(self, source, chunk_dim=None): + """ + Iterate over the output chunks of this atom applied to a chunk source. + + The manual chunk-loop surface: wraps seam handling, buffering and the + final flush into a plain generator, so + ``for out in atom.iter_chunks(source): ...`` is a complete streaming + loop. The serial executor is literally this generator plus a writer. + + Parameters + ---------- + source : iterable of DataArray + The chunks to process. Any iterable works; a loader exposing a + ``chunk_dim`` attribute provides the chunked dimension. + chunk_dim : str, optional + The dimension along which chunks follow each other. Defaults to + the source's ``chunk_dim`` attribute, else ``"time"``. + + Yields + ------ + chunk : DataArray or pandas.DataFrame + Zero or more output chunks per input chunk, then the flushed tail. + """ + if chunk_dim is None: + chunk_dim = getattr(source, "chunk_dim", "time") + for chunk in source: + yield from _aschunks(self(chunk, chunk_dim=chunk_dim)) + yield from self.flush() + self.reset() + + def process(self, source, out=None, chunks=None, until=None, merge=True): + """ + Process any chunk source through this atom, writing to any sink. + + The one-call form of chunked execution: the input is resolved into a + chunk source and the output into a writer automatically (see + :func:`xdas.processing.process`, which this method binds to the + atom). The same pipeline that runs eagerly with ``pipeline(da)`` + streams a massive archive with ``pipeline.process(da, out=...)``. + + A :class:`~xdas.DataCollection` is walked exactly as ``atom(dc)`` + walks it — `gather` first, then mapping levels recursed and sequence + levels folded, each leaf streamed in turn — so that + ``atom.process(dc, out=None) == atom(dc)``. The streaming form is + not second-class, which matters most where a leaf is too large to + call eagerly at all. + + Parameters + ---------- + source : DataArray, DataCollection, str, Path, iterable or loader + What to process: an in-memory or virtual :class:`DataArray`, a + :class:`~xdas.DataCollection` to walk leaf by leaf, a file path, + directory or glob pattern, a ``"tcp://..."`` address, + :func:`xdas.watch` for realtime, or any iterable of chunks. + out : str, Path, writer or None, optional + Where to write the output: ``None`` accumulates in memory and + returns the joined result (size-guarded); a path is matched with + the first output chunk (directory for DataArray or Stream + chunks, ``*.csv`` for DataFrames, ``"tcp://..."`` to publish); a + writer instance passes through. Walking a collection, a file, a + URL or a writer instance is shared by every leaf while a + directory fans out into one subdirectory per leaf. + chunks : dict or "auto", optional + Chunk sizes for DataArray sources, e.g. ``{"time": 1000}``. + Virtual arrays default to ``"auto"``: chunk boundaries aligned + to the storage tiling. + until : str, datetime64 or float, optional + Stop at this coordinate value along the chunked dimension; the + clean way to bound an unbounded source. + merge : bool, optional + Whether to fold the per-leaf results of a collection walk + through the `merge` hook, as ``atom(dc)`` does. Walk-level, not + stage-level, and only meaningful for ``out=None``. + + Returns + ------- + result : object + The writer result: the joined output for ``out=None``, whatever + the resolved writer returns otherwise, or ``None`` when the + pipeline emitted no output. + + Examples + -------- + >>> import numpy as np + >>> import xdas as xd + >>> pipeline = xd.decimate(..., target=50.0) >> np.square + >>> pipeline.process(da_virtual, out="results/") # doctest: +SKIP + >>> pipeline.process("archive/*.h5", out="results/") # doctest: +SKIP + >>> pipeline.process(xd.watch("/incoming"), out="sds/") # doctest: +SKIP + """ + from ..processing.core import process + + return process(self, source, out=out, chunks=chunks, until=until, merge=merge) + + def _check_chunk_dim(self, x, chunk_dim): + """Raise if this atom cannot process *x* chunked along *chunk_dim*.""" + + def _refuse_chunked_along(self, dim, chunk_dim, x=None): + """ + Raise if a whole-record operation is being chunked along its own dim. + + The guard for atoms that need the whole record along the dimension + they work on: call it from :meth:`initialize` (or a + :meth:`_check_chunk_dim` override) with the dimension the atom works + along and the dimension the stream is chunked along. *dim* may also + be a mapping whose keys are the working dimensions (a kernel dict). + ``"first"`` and ``"last"`` aliases are resolved against *x* when + given, so the comparison is never made on an unresolved alias. + """ + if chunk_dim is None: + return + dims = list(dim.keys()) if isinstance(dim, dict) else [dim] + if x is not None and hasattr(x, "dims"): + dims = [ + x.dims[0] if d == "first" else x.dims[-1] if d == "last" else d + for d in dims + ] + if any(d is None or d in ("first", "last") or d == chunk_dim for d in dims): + name = ( + getattr(self, "name", None) + or getattr(getattr(self, "func", None), "__name__", None) + or type(self).__name__ + ) + named = ", ".join(repr(d) for d in dims) + raise ValueError( + f"{name} needs the whole record along {named} and cannot " + f"process data chunked along {chunk_dim!r}: process the " + f"stream unchunked, or chunk along another dimension" + ) + + def __rshift__(self, other): + """Compose with *other* into a new pipeline: ``atom >> atom``.""" + if isinstance(other, Atom) or callable(other): + return compose(self, other) + return NotImplemented + + def __rrshift__(self, other): + """Prepend a bare callable, or apply the pipeline: ``da >> atom``.""" + if callable(other): + return compose(other, self) + return self(other) + + __irshift__ = __rshift__ + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + """Append the ufunc to the pipeline instead of computing it (tracing).""" + if ufunc is np.right_shift and method == "__call__" and len(inputs) == 2: + # ``da >> atom``: the DataArray operand dispatched through numpy. + left, right = inputs + if right is self and not isinstance(left, Atom): + return self.__rrshift__(left) + if method != "__call__": + return NotImplemented + if "out" in kwargs: + # In-place operators rebind the target name with their return + # value, so tracing them keeps value semantics: drop the `out` + # and append the out-of-place operation. + if len(kwargs["out"]) != 1 or kwargs["out"][0] is not self: + return NotImplemented + kwargs = {key: value for key, value in kwargs.items() if key != "out"} + if sum(input is self for input in inputs) != 1 or any( + isinstance(input, Atom) and input is not self for input in inputs + ): + # Never bail silently into computation: fan-in is a design + # boundary, and the error must say so at the line that traced it. + raise TypeError( + "traced fan-in is not supported: a pipeline is a single " + "chain, so a ufunc may involve one atom exactly once " + "(combine the branches eagerly, or wrap the whole " + "expression in a function and Partial it)" + ) + args = tuple(... if input is self else input for input in inputs) + return compose(self, Partial(ufunc, *args, **kwargs)) def reset(self): """Reset all state entries to ``...`` (uninitialised sentinel).""" + object.__setattr__(self, "_seam", None) for key in self._state: setattr(self, key, State(...)) for filter in self._atoms.values(): filter.reset() + def fresh(self): + """ + Return a stateless clone of this atom: same config, no state. + + Config is shared *by reference* (a model is never deep-copied), + nested atoms are recursed, and every state entry comes back + uninitialised — where :meth:`reset` wipes this instance, ``fresh`` + leaves it untouched, so one configured atom can serve several + independent runs. + """ + clone = type(self).__new__(type(self)) + Atom.__init__(clone) + for name, value in vars(self).items(): + if name in ("_config", "_state", "_atoms", "_seam"): + continue + if name in self._atoms: + setattr(clone, name, self._atoms[name].fresh()) + elif name in self._state: + setattr(clone, name, State(...)) + elif name in self._config: + setattr(clone, name, value) + else: + # Attributes opted out of the registries with + # `object.__setattr__` (pure helpers) travel as they are. + object.__setattr__(clone, name, value) + return clone + def save_state(self, path): """Serialise the current state to a NetCDF4 file at *path*.""" DataCollection(self.state).to_netcdf(path) @@ -291,10 +1029,102 @@ def __init__(self, atoms: Any, name: str | None = None) -> None: self.name = name def call(self, x: Any, **flags) -> Any: - """Pass *x* through each atom in order and return the final result.""" + """ + Pass *x* through each atom in order and return the output chunks. + + Each stage may emit zero or more chunks per input chunk (seam tails, + rechunking, reductions); the streams are folded stage by stage, so + cadence mismatches are absorbed inside the pipeline. + """ + chunks = [x] + for atom in self: + chunks = [ + chunk + for out in (atom(x, **flags) for x in chunks) + for chunk in _aschunks(out) + ] + return chunks + + def flush(self): + """ + Cascade-flush the pipeline, codec-drain style. + + Flush the first stage and push its tail through the following stages, + then flush the second stage, and so on. Tails flow downstream as + ordinary data: each downstream stage folds them into its own state + before being flushed itself. Returns the drained chunks. + """ + flags = {"chunk_dim": self._seam["chunk_dim"]} if self._seam else {} + chunks = [] + for index, atom in enumerate(self): + tail = atom.flush() + for downstream in self[index + 1 :]: + tail = [ + chunk + for out in (downstream(x, **flags) for x in tail) + for chunk in _aschunks(out) + ] + chunks.extend(tail) + return chunks + + def _resolve_dim(self, x): + """Resolve the operating dimension from the first stage that has one.""" for atom in self: - x = atom(x, **flags) - return x + dim = atom._resolve_dim(x) + if dim is not None: + return dim + return None + + def gather(self, mapping): + """ + Collapse *mapping* through the first stage that claims it. + + The gather happens once, before the *first* stage runs, whichever + stage claimed it — a pipeline is one transformation of one input, and + a level a later stage needs as an axis has to be an axis by the time + the input enters the pipeline. That is what a picking pipeline relies + on: the component level is recognised by the stage that knows the + model, and the earlier filter and resampling stages see the + components as a dimension, which is the only form in which a + per-channel filter can select one of them. + + First claim wins. Claiming is a structural statement about the level + — that it is one axis of the input — so two stages that both claim it + agree on what it is, and can differ only in how they would collapse + it; the first stage in the pipeline is then as good an arbiter as any, + and the only one that keeps the answer independent of the stages + downstream. + """ + for atom in self: + gathered = atom.gather(mapping) + if gathered is not None: + return gathered + return None + + @property + def merge(self): + """ + The `merge` hook of the last stage declaring one, else ``None``. + + The last stage is the one whose output leaves the pipeline, so it is + the one that knows what folding its results means — a pipeline + ending on a :class:`~xdas.atoms.Trigger` merges pick tables without + having to say so. + """ + for atom in reversed(self): + if atom.merge is not None: + return atom.merge + return None + + def fresh(self): + """Return a stateless clone: each stage cloned, config shared.""" + # the stages are not in `vars`, and a subclass that assembles its own + # (a `Picker`, say) cannot be rebuilt by calling its constructor + clone = super().fresh() + for key, atom in enumerate(self): + clone.append(atom.fresh()) + clone._atoms[key] = clone[key] + return clone def __repr__(self) -> str: width = len(str(len(self))) @@ -324,8 +1154,8 @@ class Partial(Atom): keyword arguments, to properly initialize the corresponding states and to return as many additional outputs as there are stateful arguments. - Partial uses several reserved keyword arguments that cannot by passed to `func`: - 'func', 'name' and 'state'. + Partial uses two reserved keyword arguments that cannot be passed to `func`: + 'func' and 'name'. Parameters ---------- @@ -391,6 +1221,34 @@ def __init__( setattr(self, key, value) else: self.kwargs[key] = value + # The operating dimension is resolved from the call arguments so the + # whole-record guard can compare it with the chunked one and so eager + # calls split gappy input into runs along it. A `_whole_record`-marked + # function may name a different argument (a kernel dict, say) as the + # one carrying its working dimensions. + dim_arg = getattr(func, "_whole_record_dim_arg", None) + try: + bound = inspect.signature(func).bind_partial(*self.args, **self.kwargs) + bound.apply_defaults() + default = getattr(func, "_whole_record_default", None) + dim = bound.arguments.get("dim") + if isinstance(dim, dict) and len(dim) == 1: + # {input_dim: output_dim} mapping (e.g. the fft functions): + # the operating dimension is the input one. + ((dim, _),) = dim.items() + self.dim = default if dim is None else dim + refuse_dim = bound.arguments.get(dim_arg) if dim_arg else None + if refuse_dim is None: + refuse_dim = default + except (TypeError, ValueError): + self.dim = None + refuse_dim = None + object.__setattr__(self, "_refuse_dim", refuse_dim) + + def _check_chunk_dim(self, x, chunk_dim): + """Refuse chunking along the working dim of a whole-record function.""" + if getattr(self.func, "_whole_record_dim_arg", None) is not None: + self._refuse_chunked_along(self._refuse_dim, chunk_dim, x) @property def stateful(self): @@ -448,29 +1306,74 @@ def get_state(self): } +def compose(input, output): + """ + Chain *input* then *output* into a new Sequential with value semantics. + + Composition never mutates its operands: each call returns a fresh + Sequential, so intermediate pipelines stay usable on their own. Bare + callables are wrapped into Partial atoms. Unnamed Sequentials are + flattened; a named input Sequential keeps its name, a named output + Sequential stays nested. + + This is the primitive behind the ``>>`` operator and operator tracing. + + Parameters + ---------- + input : Atom or callable + The upstream atom or pipeline. + output : Atom or callable + The atom or pipeline to append. + + Returns + ------- + Sequential + A new pipeline running *input* then *output*. + """ + if not isinstance(input, Atom): + input = Partial(input) + if not isinstance(output, Atom): + output = Partial(output) + head = list(input) if isinstance(input, Sequential) else [input] + tail = ( + list(output) + if isinstance(output, Sequential) and output.name is None + else [output] + ) + name = input.name if isinstance(input, Sequential) else None + return Sequential(head + tail, name=name) + + def atomized(func): """ Make the function return an Atom if `...` or an atom is passed as argument. In case `...` is passed as a positional argument, the function is wrapped into a - Partial object. If an Atom object is passed as a positional argument, the function - is wrapped into a Sequential object. Otherwise, the function is called as is. + Partial object. If an Atom object is passed as a positional argument, a new + Sequential is returned that chains that atom with the atomized function (the + input atom is never mutated). Otherwise, the function is called as is. + + Applied to an Atom subclass, `atomized` instead generates its function + form (see :func:`as_function`): a function taking the data as first + argument followed by the class parameters, with the same `...`/atom + dispatch as above. Parameters ---------- - func: callable - The function to wrap as a Partial atom if any `...` or input atom is a passed. + func: callable or type + The function to wrap as a Partial atom if any `...` or input atom is passed. It must handle the `...` argument as a placeholder for the input data and for the passing states. It must return a unique output except if the function is stateful. In that case, the function must return the processed data as first - output and the updated state as additional outputs. + output and the updated state as additional outputs. If an Atom subclass is + given, its function form is returned instead. Returns ------- output or atom: Any or (Partial or Sequential) if no `...` or Atom object is passed as a positional argument, returns the output of the function. If an Atom object is passed as a positional argument, - returns a Sequential object containing the Atom object and the atomized function. + returns a new Sequential chaining the Atom object and the atomized function. If `...` is passed as a positional argument, returns a Partial object containing the atomized function. This latter has the same documentation and names than the original function. @@ -514,6 +1417,8 @@ def atomized(func): cumsum(...) [stateful] """ + if isinstance(func, type) and issubclass(func, Atom): + return as_function(func) @wraps(func) def wrapper(*args, **kwargs): @@ -526,12 +1431,83 @@ def wrapper(*args, **kwargs): else: raise ValueError("Only one Atom object can be passed as function input") args = tuple(... if isinstance(arg, Atom) else arg for arg in args) - output = Partial(func, *args, **kwargs) - if isinstance(input, Sequential): - return input.append(output) - else: - return Sequential([input, output]) + return compose(input, Partial(func, *args, **kwargs)) else: return func(*args, **kwargs) return wrapper + + +def _whole_record(dim_arg="dim", default=None): + """ + Mark a function as needing the whole record along its working dimension. + + The decorator to apply at the definition site of a whole-record function, + under :func:`atomized`: the resulting atoms refuse chunked execution + along the dimension named by the *dim_arg* argument (resolved from the + call arguments, aliases included), via + :meth:`Atom._refuse_chunked_along`. + + *default* names the dimension the function works along when that argument + is left unset. A function whose ``dim=None`` means "the last one" has to + say so here: an unknown dimension is refused whatever the stream is + chunked along, which would reject the transform of *another* dimension + than the chunked one. + """ + + def decorator(func): + func._whole_record_dim_arg = dim_arg + func._whole_record_default = default + return func + + return decorator + + +def as_function(cls): + """ + Generate the function form of an Atom subclass. + + A lowercase function taking the data as first argument, then the class + parameters: called with data it builds the atom and applies it eagerly, + called with ``...`` in the data slot it returns the configured atom, and + called with an atom or a pipeline it returns a new pipeline extended + with this atom. + + Parameters + ---------- + cls : type + The :class:`Atom` subclass to generate the function form of. + + Returns + ------- + callable + The function form, named after the class in snake case. + """ + parameters = [ + parameter + for key, parameter in inspect.signature(cls.__init__).parameters.items() + if key != "self" + ] + + def wrapper(da, *args, **kwargs): + atom = cls(*args, **kwargs) + if da is ...: + return atom + if isinstance(da, Atom): + return compose(da, atom) + return atom(da) + + name = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", cls.__name__).lower() + wrapper.__name__ = name + wrapper.__qualname__ = name + wrapper.__module__ = cls.__module__ + wrapper.__signature__ = inspect.Signature( + [inspect.Parameter("da", inspect.Parameter.POSITIONAL_OR_KEYWORD), *parameters] + ) + wrapper.__doc__ = ( + f"Apply a :class:`{cls.__name__}` atom to `da`.\n\n" + "Passing ``...`` as `da` returns the atom itself; passing an atom or a\n" + "pipeline returns a new pipeline extended with this atom. The other\n" + f"parameters are those of :class:`{cls.__name__}`, documented below.\n\n" + ) + (inspect.getdoc(cls) or "") + return wrapper diff --git a/xdas/atoms/detect.py b/xdas/atoms/detect.py new file mode 100644 index 00000000..99dfc2d9 --- /dev/null +++ b/xdas/atoms/detect.py @@ -0,0 +1,594 @@ +""" +Detection atoms: turn a characteristic function into picks. + +These atoms emit :class:`pandas.DataFrame` objects rather than arrays — the +transducer contract makes that unremarkable, a pick table is just another +chunk type flowing downstream to a CSV sink. :class:`Trigger` is the +threshold detector; its lowercase twin is :func:`xdas.trigger`. +""" + +from collections.abc import Mapping + +import numpy as np +import pandas as pd +from numba import njit + +from ..coordinates import Coordinate +from ..core import concat_coords +from .core import Atom, State, atomized + +__all__ = ["Trigger", "trigger"] + +#: Name of the label dimension a mapping of thresholds keys on. It is the +#: dimension :class:`~xdas.atoms.Annotate` appends to its characteristic +#: function, so the two agree by construction. +PHASE_DIM = "phase" + + +class Trigger(Atom): + """ + Find picks in a data array along a given axis based on a given threshold. + + The pick findings use a triggering mechanism where triggers are turned on and off + based on the threshold crossings. The trigger off threshold is half of the trigger + on threshold. Picks are determined by finding the maximum value on each triggered + region. + + Parameters + ---------- + thresh : float or mapping + The threshold value for picking. A scalar applies to every lane. A + mapping keyed on the ``phase`` coordinate gives one threshold per + label; labels the mapping does not list are **never** triggered, which + is how a characteristic function keeps carrying its noise class + without that class ever producing a pick. Keying on the label rather + than on its position is a correctness requirement: the label order of + a model is a property of its weight set and flips between them. + dim : str, optional + The dimension along which to find picks. Defaults to "time". + coords : sequence of str, "auto" or None, optional + The coordinates used to annotate the picks, one column per name. Any + coordinate of the input can be named, including non-dimensional ones + (a station identifier attached to the distance dimension, a latitude, + ...): each is indexed along the dimension it varies on. Scalar (0-d) + coordinates are emitted as constant columns, which is what lets a pick + table carry a ``network``/``station``/``location`` identity when + picking a single array. Defaults to ``"auto"``: the scalar + coordinates, then the other dimension coordinates, then the picked + dimension — identity first, measurement last, so the columns do not + depend on the input's dimension order (see :meth:`_names`); ``None`` + restricts it to the dimension coordinates, in the input's own order. + + Notes + ----- + For more details see the documentation of the `initialize`, `call` and `flush` + methods. + + Examples + -------- + >>> import numpy as np + >>> import pandas as pd + >>> import xdas as xd + + Use case: + + >>> cft = xd.DataArray( + ... data=[[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2]], + ... coords={ + ... "space": [0.0], + ... "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0], "sampling_interval": 1.0}, + ... }, + ... ) + + Eager processing with the functional twin: + + >>> xd.trigger(cft, thresh=0.5) + space time value + 0 0.0 2.0 0.9 + 1 0.0 7.0 0.7 + + Chunked processing using atomic processing. `flush` closes whatever trigger + is still open once the stream ends: + + >>> atom = xd.trigger(..., thresh=0.5) + >>> chunks = xd.split(cft, 3, dim="time") + >>> result = [] + >>> for chunk in chunks: + ... picks = atom(chunk, chunk_dim="time") + ... result.append(picks) + >>> result += atom.flush() + >>> result = pd.concat(result, ignore_index=True) + >>> result + space time value + 0 0.0 2.0 0.9 + 1 0.0 7.0 0.7 + + A trigger that never turns off used to lose its pick; it is now closed at + the end of the record, as `obspy.trigger_onset` does: + + >>> tail = cft.isel(time=slice(0, 4)) + >>> xd.trigger(tail, thresh=0.5) + space time value + 0 0.0 2.0 0.9 + + Annotating the picks with a non-dimensional coordinate instead of the + dimension it is attached to: + + >>> cft = cft.assign_coords(station=("space", ["ST01"])) + >>> xd.trigger(cft, thresh=0.5, coords=["time", "station"]) + time station value + 0 2.0 ST01 0.9 + 1 7.0 ST01 0.7 + + A characteristic function labelled by phase takes one threshold per label: + + >>> cft = xd.DataArray( + ... data=[[0.0, 0.9, 0.0], [0.0, 0.8, 0.0], [0.0, 0.9, 0.0]], + ... coords={ + ... "phase": ["N", "P", "S"], + ... "time": {"tie_indices": [0, 2], "tie_values": [0.0, 2.0], "sampling_interval": 1.0}, + ... }, + ... ) + >>> xd.trigger(cft, thresh={"P": 0.5, "S": 0.5}) + phase time value + 0 P 1.0 0.8 + 1 S 1.0 0.9 + + The noise class is the loudest lane here and carries no entry, so it never + triggers. Scalar coordinates come along as constant columns, leading: + + >>> cft = cft.assign_coords(station="ST01") + >>> xd.trigger(cft, thresh={"P": 0.5, "S": 0.5}) + station phase time value + 0 ST01 P 1.0 0.8 + 1 ST01 S 1.0 0.9 + + Picking a whole collection gives one table for the whole network. Each + leaf's picks are labelled with the tree path they were found under, as + they are found, and `merge` folds the tables on the way back up. Here the + ``station`` column has two sources — the tree key and the leaf's own + scalar coordinate — which agree, so it stays one column, leading: + + >>> dc = xd.DataCollection( + ... { + ... "DBNFM": cft.assign_coords(station="DBNFM"), + ... "LBFI": cft.assign_coords(station="LBFI"), + ... }, + ... "station", + ... ) + >>> xd.trigger(dc, thresh={"P": 0.5, "S": 0.5}) + station phase time value + 0 DBNFM P 1.0 0.8 + 1 DBNFM S 1.0 0.9 + 2 LBFI P 1.0 0.8 + 3 LBFI S 1.0 0.9 + + `merge=False` keeps the tree, each leaf already annotated: + + >>> tree = xd.trigger(..., thresh={"P": 0.5, "S": 0.5})(dc, merge=False) + >>> tree["DBNFM"] + station phase time value + 0 DBNFM P 1.0 0.8 + 1 DBNFM S 1.0 0.9 + + """ + + def __init__(self, thresh, dim="time", coords="auto"): + super().__init__() + + # parameters + if isinstance(thresh, Mapping): + self.thresh = {str(key): float(value) for key, value in thresh.items()} + else: + self.thresh = float(thresh) + self.dim = str(dim) + if coords is None or isinstance(coords, str): + if isinstance(coords, str) and coords != "auto": + raise ValueError( + f"`coords` must be 'auto', None or a sequence of " + f"coordinate names, got {coords!r}" + ) + self.coords = coords + else: + self.coords = tuple(coords) + + # states + self.axis = State(...) + self.shape = State(...) + self.thresh_on = State(...) + self.thresh_off = State(...) + self.status = State(...) + self.index = State(...) + self.value = State(...) + self.offset = State(...) + self.coord = State(...) + self.annotations = State(...) + + def initialize(self, cft, **flags): + """ + Initialize the trigger with the following states. + + - "axis": An integer indicating the axis number of the dimension along which to + find picks. + - "shape": A tuple indicating the unravel shape of the lanes along which + the picks will be found. + - "thresh_on"/"thresh_off": Float arrays holding the trigger on and off + thresholds of each lane, raveled like the lanes. + - "status": A boolean array indicating the trigger status for each lane. + - "index": An integer array indicating the index of the last triggered value + for each lane. + - "value": A float array indicating the value of the last triggered value for + each lane. + - "offset": An integer indicating the offset of the chunk. + - "coord": An InterpCoordinate containing coordinate information along 'dim' up + to the last processed chunk. + - "annotations": The resolved pick columns, one ``(name, axis, source)`` + triple each. + + + Parameters + ---------- + cft : DataArray + The characteristic function where picks must be found. + **flags + Optional flags. + + """ + self.axis = State(cft.get_axis_num(self.dim)) + self.shape = State(cft.shape[: self.axis] + cft.shape[self.axis + 1 :]) + thresh_on = self._thresholds(cft) + self.thresh_on = State(thresh_on) + self.thresh_off = State(thresh_on / 2.0) + self.status = State(np.zeros(self.shape, dtype=bool)) + self.index = State(np.zeros(self.shape, dtype=int)) + self.value = State(np.zeros(self.shape, dtype=float)) + self.offset = State(0) + self.coord = State(Coordinate({"tie_indices": [], "tie_values": []}, self.dim)) + self.annotations = State(self._annotations(cft)) + self.labels = State({}) + + def call(self, cft, **flags): + """ + Call the trigger. + + Parameters + ---------- + cft : DataArray + The characteristic function where picks must be found. + **flags + Optional flags. + + Returns + ------- + picks: DataFrame + A DataFrame containing the pick coordinates and their corresponding values. + + Notes + ----- + A trigger that has not turned off by the end of the chunk stays open: its + pick is emitted by the chunk that closes it, or by `flush` at the end of + the run. + + Chunked along a dimension other than `dim`, none of that state carries: + the next chunk holds *other* lanes, so its open triggers, its sample + offset, its accumulated coordinate and the lane values annotating its + picks all belong to the chunk that produced them. Such a chunk is a + whole record on its own, run from a fresh state and closed here, which + is what makes the cross-dimension exemption of the chunk-semantics gate + true of this atom. + + """ + chunk_dim = flags.get("chunk_dim") + independent = chunk_dim is not None and chunk_dim != self.dim + if independent: + self.initialize(cft, **flags) + data = np.asarray(cft.values, dtype=float) + values, indices = self._call_numeric(data) + self.coord = concat_coords([self.coord, cft.coords[self.dim]], tolerance=None) + self.labels = self._accumulate(cft) + picks = self._picks(indices, values) + if independent: + return [picks] + self.flush() + return picks + + def merge(self, results): + """ + Fold the pick tables of a collection walk into one. + + A plain concatenation is all this takes: the walk already gave each + table the columns of the tree path its leaf was reached by, so the + rows carry their identity and nothing has to be reconstructed here. + + Parameters + ---------- + results : sequence of DataFrame + The per-leaf pick tables, in walk order. A leaf that produced no + pick contributes an empty table, which concatenates away. + + Returns + ------- + DataFrame + The concatenated table, reindexed from zero. Leaves disagreeing + on their columns — one carrying a scalar coordinate another does + not — union them, the missing cells left empty. A collection + without a single pick gives an empty table. + + """ + if not results: + return pd.DataFrame() + return pd.concat(results, ignore_index=True) + + def flush(self): + """ + Close the triggers still open at the end of a run. + + `obspy.trigger_onset` closes whatever is on when the array ends, so the + last pick of a record is not lost. Doing it here rather than at the end + of each chunk keeps the result chunk-invariant: a run is closed once, + whether it arrived in one piece or in twenty. + + Returns + ------- + list of DataFrame + One frame of the closed picks, or nothing if no trigger was open. + + """ + if not self.initialized: + return [] + lanes = np.flatnonzero(np.reshape(self.status, (-1,))) + if lanes.size == 0: + return [] + values = np.reshape(self.value, (-1,))[lanes] + indices = np.reshape(self.index, (-1,))[lanes] + self.status = State(np.zeros(self.shape, dtype=bool)) + return [self._picks(self._unravel(lanes, indices), values)] + + def _thresholds(self, cft): + """Resolve `thresh` into one trigger-on threshold per lane, raveled.""" + if not isinstance(self.thresh, dict): + return np.full(self.shape, self.thresh, dtype=float).reshape(-1) + if self.dim == PHASE_DIM: + raise ValueError( + f"cannot key thresholds on {PHASE_DIM!r}: it is the dimension " + "the picks are found along" + ) + if PHASE_DIM not in cft.dims or PHASE_DIM not in cft.coords: + raise ValueError( + f"a mapping of thresholds is keyed on the {PHASE_DIM!r} " + f"coordinate, which the data to pick on does not have " + f"(dimensions: {cft.dims})" + ) + labels = [ + value.decode() if isinstance(value, bytes) else str(value) + for value in cft.coords[PHASE_DIM].values + ] + unknown = [key for key in self.thresh if key not in labels] + if unknown: + raise KeyError( + f"the threshold mapping names labels that are not in the " + f"{PHASE_DIM!r} coordinate ({sorted(unknown)}); it holds {labels}" + ) + # An unlisted label gets an infinite threshold: it can never trigger. + values = np.array([self.thresh.get(label, np.inf) for label in labels]) + axis = cft.get_axis_num(PHASE_DIM) + shape = [1] * len(self.shape) + shape[axis - 1 if axis > self.axis else axis] = values.size + return np.broadcast_to(values.reshape(shape), self.shape).reshape(-1).copy() + + def _annotations(self, cft): + """ + Resolve the requested columns into ``(name, axis, source)`` triples. + + *source* is the coordinate the column is read from, or ``None`` when + it is attached to the picked dimension: those are indexed by absolute + sample number, so they are read off the coordinates accumulated over + the run rather than off the chunk in hand. + """ + annotations = [] + for name in self._names(cft): + if name == self.dim: + annotations.append((name, self.axis, None)) + continue + coord = self._annotation(cft, name) + if coord.dim is None: + annotations.append((name, None, coord)) + elif coord.dim == self.dim: + annotations.append((name, self.axis, None)) + else: + annotations.append((name, cft.get_axis_num(coord.dim), coord)) + return tuple(annotations) + + def _names(self, cft): + """ + Return the names of the columns annotating the picks, in order. + + Identity first, measurement last: the scalar coordinates lead, then + the dimension coordinates that say *which* lane the pick was found in, + then the picked dimension itself, then the value. So a pick table + reads the same — ``network station location phase time value`` — + whether its identity came from scalar coordinates on one array or + from the tree path of a collection (which the walk puts in the same + leading position), and whatever order the input's dimensions came in: + a characteristic function laid out ``(distance, phase, time)`` — as + :class:`~xdas.atoms.Annotate` emits it — and one laid out + ``(time, distance, phase)`` give the same columns. + """ + if self.coords is None: + # A dimension without a coordinate has no labels to annotate with. + return tuple(dim for dim in cft.dims if dim in cft.coords) + if self.coords == "auto": + scalars = tuple( + name for name, coord in cft.coords.items() if coord.dim is None + ) + dims = tuple( + dim for dim in cft.dims if dim in cft.coords and dim != self.dim + ) + picked = (self.dim,) if self.dim in cft.coords else () + return scalars + dims + picked + return self.coords + + def _annotation(self, cft, name): + """Return the coordinate *name* of *cft*, checked as a pick annotation.""" + if name not in cft.coords: + raise KeyError( + f"cannot annotate picks with {name!r}: it is not a coordinate " + f"of the data to pick on (available: {sorted(cft.coords)})" + ) + return cft.coords[name] + + def _accumulate(self, cft): + """ + Extend the coordinates of the picked dimension with those of *cft*. + + A trigger reports the sample its onset was found at, which may lie in + a chunk already gone: the labels of the picked dimension are kept for + the whole run so that an absolute index still names something. The + dimension coordinate itself is kept as a coordinate (`coord`), which + stays compact; the others are plain arrays. + """ + labels = dict(self.labels) + for name, axis, source in self.annotations: + if source is not None or name == self.dim: + continue + values = np.asarray(cft.coords[name].values) + previous = labels.get(name) + labels[name] = ( + values if previous is None else np.concatenate([previous, values]) + ) + labels[self.dim] = self.coord.values + return labels + + def _picks(self, indices, values): + """Build the pick table of the *values* found at *indices*.""" + picks = {} + for name, axis, source in self.annotations: + if source is None: + picks[name] = self.labels[name][indices[axis]] + elif axis is None: + picks[name] = np.full(len(values), source.values, dtype=source.dtype) + else: + picks[name] = source[indices[axis]].values + picks["value"] = values + return pd.DataFrame(picks) + + def _unravel(self, lanes, indices): + """Turn lane numbers and sample indices into one index array per axis.""" + coords = np.unravel_index(lanes, self.shape) if self.shape else () + return coords[: self.axis] + (indices,) + coords[self.axis :] + + def _call_numeric(self, data): + """ + Find picks in a N-dimensional array along a given axis based on a given threshold. + + The pick findings use a triggering mechanism where triggers are turned on and off + based on the threshold crossings. The trigger off threshold is half of the trigger + on threshold. Picks are determined by finding the maximum value on each triggered + region. + + Parameters + ---------- + data : ndarray + The characteristic function where picks must be found. + + Returns + ------- + values : 1d ndarray + The values of the picks. + indices : tuple of 1d ndarray + One index array per axis, locating each pick. + + Notes + ----- + A trigger that has not turned off by the end of the array stays open; `flush` + closes it at the end of the run. + + """ + data = np.moveaxis(data, self.axis, -1) + length = data.shape[-1] + + # ravel additional axes into a unique lanes axis + data = np.reshape(data, (-1, data.shape[-1])) + status_view = np.reshape(self.status, (-1,)) + index_view = np.reshape(self.index, (-1,)) + value_view = np.reshape(self.value, (-1,)) + + lanes, indices, values = _trigger( + data, + self.thresh_on, + self.thresh_off, + status_view, + index_view, + value_view, + self.offset, + ) + self.offset += length + + return values, self._unravel(lanes, indices) + + +@njit( + "Tuple((i8[:], i8[:], f8[:]))(f8[:, :], f8[:], f8[:], b1[:], i8[:], f8[:], i8)", + cache=True, +) +def _trigger( # pragma: no cover + cft, thresh_on, thresh_off, buffer_status, buffer_index, buffer_value, offset +): + """ + Perform trigger detection on the input data. + + Parameters + ---------- + cft : ndarray + 2D array of shape (n, m) representing the input data. Each row is a lane. Each + column is the signal onto perform trigger detection. + thresh_on : ndarray + Float array of shape (n,) holding the threshold value for turning on the + trigger of each lane. An infinite threshold never triggers. + thresh_off : ndarray + Float array of shape (n,) holding the threshold value for turning off the + trigger of each lane. + buffer_status : ndarray + Boolean buffer of shape (n,) holding the trigger status for each lane. + buffer_index : ndarray + Integer buffer of shape (n,) holding the index of the last found pick for each + lane. + buffer_value : ndarray + Float buffer of shape (n,) holding the value of the last found pick for each + lane. + offset : int + The offset to add to the found indices. + + Returns + ------- + tuple of ndarray + A tuple containing three arrays of shape (k,) where k is the number of picks + found. The arrays are: + + - lanes : lanes indices (along first axis) of the picks. + - indices : signal indices (along last axis) of the picks. + - values : values of the picks. + + """ + lanes = [] + indices = [] + values = [] + for (lane, index), value in np.ndenumerate(cft): + index += offset + if buffer_status[lane]: + if value > buffer_value[lane]: + buffer_index[lane] = index + buffer_value[lane] = value + if value < thresh_off[lane]: + buffer_status[lane] = False + lanes.append(lane) + indices.append(buffer_index[lane]) + values.append(buffer_value[lane]) + else: + if value > thresh_on[lane]: + buffer_status[lane] = True + buffer_index[lane] = index + buffer_value[lane] = value + return np.array(lanes), np.array(indices), np.array(values) + + +trigger = atomized(Trigger) diff --git a/xdas/atoms/kernel.py b/xdas/atoms/kernel.py new file mode 100644 index 00000000..7ccef722 --- /dev/null +++ b/xdas/atoms/kernel.py @@ -0,0 +1,590 @@ +""" +Kernel atoms: exact stateful chunked primitives with machine parameters. + +This is the expert layer. Kernel atoms take machine parameters (filter +coefficients, integer factors) whose meaning depends on the sampling rate; +the public task atoms (:mod:`xdas.atoms.tasks`) design them from physical +parameters at the first call. They are the units used to prove that chunked +processing equals unchunked processing. + +Includes :class:`LFilter`, :class:`SOSFilter`, :class:`DownSample`, +:class:`UpSample`, :class:`Polyphase`, :class:`Rechunk`. +""" + +import math + +import numpy as np +import scipy.signal as sp + +from ..coordinates import Coordinate, get_sampling_interval +from ..coordinates.core import parse_scalar_delta +from ..core import DataArray, concat, split +from ..parallel import parallelize +from .core import Atom, State, atomized + + +def _along(axis, ndim, slc): + """Index tuple selecting *slc* along *axis* and everything else elsewhere.""" + return tuple(slc if index == axis else slice(None) for index in range(ndim)) + + +def _carry_labels(coords, name, positions): + """ + Re-index the non-dimensional coordinates attached to *name*. + + An atom that changes the number of samples along a dimension has to say + what became of the *other* coordinates attached to it — the station code + of a channel, the latitude of a sensor — or they keep their input length + and silently label the output with the wrong lanes. They are carried by + taking, for each output sample, the input sample it is drawn from + (*positions*): a label names a source sample rather than a position, so + unlike the dimension coordinate it is not shifted by a filter's group + delay. Depending on the sampling grid alone, that mapping is the same + however the stream was chunked. + + Parameters + ---------- + coords : Coordinates + The output coordinates, modified in place. + name : str + The resampled dimension. + positions : ndarray of int + One input index per output sample, into this chunk. + + Returns + ------- + Coordinates + The same mapping, for chaining. + """ + for key, coord in list(coords.items()): + if key != name and coord.dim == name: + coords[key] = coord[positions] + return coords + + +class LFilter(Atom): + """ + Stateful direct-form IIR/FIR filter using :func:`scipy.signal.lfilter`. + + Parameters + ---------- + b : array-like + Numerator polynomial coefficients. + a : array-like + Denominator polynomial coefficients. + dim : str or int, optional + Dimension to filter along. Defaults to ``"last"``. + parallel : int, bool, or None, optional + Worker count for parallelisation. + """ + + def __init__(self, b, a, dim="last", parallel=None): + super().__init__() + self.b = b + self.a = a + self.dim = dim + self.parallel = parallel + self.axis = State(...) + self.zi = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Set the filter axis and allocate the initial conditions buffer.""" + self.axis = State(da.get_axis_num(self.dim)) + # `dim` may be the "first"/"last" alias, which never equals a real + # dimension name: resolve it against the data before comparing, else + # the seam state is silently never allocated. + dim = self._dim_name(da) + if dim == chunk_dim: + n_sections = max(len(self.a), len(self.b)) - 1 + shape = tuple( + n_sections if name == dim else size for name, size in da.sizes.items() + ) + self.zi = State(np.zeros(shape)) + else: + self.zi = State(None) + + def call(self, da, **flags): + """Apply the filter to *da*, updating the state if chunked.""" + across = int(self.axis == 0) + if self.zi is None: + func = parallelize((None, None, across), across, self.parallel)(sp.lfilter) + data = func(self.b, self.a, da.values, self.axis) + else: + func = parallelize( + (None, None, across, None, across), (across, across), self.parallel + )(sp.lfilter) + data, zf = func(self.b, self.a, da.values, self.axis, self.zi) + self.zi = State(zf) + return da.copy(data=data) + + +class SOSFilter(Atom): + """ + Stateful second-order-sections IIR filter using :func:`scipy.signal.sosfilt`. + + Parameters + ---------- + sos : array-like + SOS filter coefficients, of shape ``(n_sections, 6)``, as returned by + e.g. :func:`scipy.signal.iirfilter`. + dim : str or int, optional + Dimension to filter along. Defaults to ``"last"``. + parallel : int, bool, or None, optional + Worker count for parallelisation. + """ + + def __init__(self, sos, dim="last", parallel=None): + super().__init__() + self.sos = sos + self.dim = dim + self.parallel = parallel + self.axis = State(...) + self.zi = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Set the filter axis and allocate the SOS initial-conditions buffer.""" + self.axis = State(da.get_axis_num(self.dim)) + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._dim_name(da) + if dim == chunk_dim: + n_sections = self.sos.shape[0] + shape = (n_sections,) + tuple( + 2 if index == self.axis else element + for index, element in enumerate(da.shape) + ) + self.zi = State(np.zeros(shape)) + else: + self.zi = State(None) + + def call(self, da, **flags): + """Apply the SOS filter to *da*, updating the state if chunked.""" + across = int(self.axis == 0) + if self.zi is None: + func = parallelize((None, across), across, self.parallel)(sp.sosfilt) + data = func(self.sos, da.values, self.axis) + else: + func = parallelize( + (None, across, None, across + 1), (across, across + 1), self.parallel + )(sp.sosfilt) + data, zf = func(self.sos, da.values, self.axis, self.zi) + self.zi = State(zf) + return da.copy(data=data) + + +class DownSample(Atom): + """ + Stateful integer downsampling by selecting every *factor*-th sample. + + Parameters + ---------- + factor : int + Downsampling factor. + dim : str or int, optional + Dimension to downsample along. Defaults to ``"last"``. + """ + + def __init__(self, factor, dim="last"): + super().__init__() + self.factor = factor + self.dim = dim + self.buffer = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Initialise the carry-over buffer for chunked operation.""" + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._dim_name(da) + if chunk_dim == dim: + self.buffer = State(da.isel({self.dim: slice(0, 0)})) + else: + self.buffer = State(None) + + def call(self, da, **flags): + """Downsample *da*, buffering the trailing partial stride when chunked.""" + if self.factor == 1: + return da + if self.buffer is not None: + da = concat([self.buffer, da], self.dim) + divpoint = da.sizes[self.dim] - da.sizes[self.dim] % self.factor + da, buffer = split(da, [divpoint], self.dim) + self.buffer = State(buffer) + return da.isel({self.dim: slice(None, None, self.factor)}) + + def flush(self): + """ + Emit the buffered samples that fall on the output grid. + + The buffer always starts on an output sample (every emission consumes + a whole number of strides), so its strided selection is the exact + remainder of the downsampled stream. + """ + if not isinstance(self.buffer, DataArray) or self.buffer.sizes[self.dim] == 0: + return [] + out = self.buffer.isel({self.dim: slice(None, None, self.factor)}) + self.buffer = State(self.buffer.isel({self.dim: slice(0, 0)})) + return [out] + + +class UpSample(Atom): + """ + Integer upsampling by zero-insertion (and optional energy scaling). + + Parameters + ---------- + factor : int + Upsampling factor. + scale : bool, optional + If ``True``, scale inserted samples so energy is preserved. + dim : str or int, optional + Dimension to upsample along. Defaults to ``"last"``. + """ + + def __init__(self, factor, scale=True, dim="last"): + super().__init__() + self.factor = factor + self.scale = scale + self.dim = dim + + def call(self, da, **flags): + """Upsample *da* by inserting zeros between every original sample.""" + if self.factor == 1: + return da + # The "first"/"last" alias never matches a real dimension name, so it + # has to be resolved before it names an axis or a coordinate. + name = self._dim_name(da) + shape = tuple( + self.factor * size if dim == name else size + for dim, size in da.sizes.items() + ) + slc = tuple( + slice(None, None, self.factor) if dim == name else slice(None) + for dim in da.dims + ) + data = np.zeros(shape, dtype=da.dtype) + if self.scale: + data[slc] = da.values * self.factor + else: + data[slc] = da.values + coords = da.coords.copy() + delta = get_sampling_interval(da, name, cast=False) + new_delta = delta / self.factor + coord = coords[name] + # Copies: the tie arrays are the input coordinate's own storage. + tie_indices = np.asarray(coord.tie_indices) * self.factor + tie_values = np.asarray(coord.tie_values).copy() + if tie_indices.size == 1: + # A one-sample chunk has a single tie, and the upsampled block + # still spans `factor` samples: it takes a second tie to say so. + tie_indices = np.append(tie_indices, self.factor - 1) + tie_values = np.append( + tie_values, tie_values[-1] + (self.factor - 1) * new_delta + ) + else: + tie_indices[-1] += self.factor - 1 + tie_values[-1] += (self.factor - 1) * new_delta + data_coord = {"tie_indices": tie_indices, "tie_values": tie_values} + if coord.isregular(): + # The derived rate may not be exactly representable (integer datetime + # resolutions truncate), so declare the representation error as jitter + # on top of the inherited one; chunk seams then stay within tolerance. + data_coord["sampling_interval"] = new_delta + data_coord["tolerance"] = coord.tolerance + np.abs( + delta - new_delta * self.factor + ) + # An irregular input gives no rate to inherit and no jitter bound to + # derive one from, so the result stays irregular rather than claiming a + # precision the source never declared. + coords[name] = Coordinate(data_coord, name) + # Each inserted sample is drawn from the input sample it follows. + positions = np.arange(shape[da.get_axis_num(name)]) // self.factor + _carry_labels(coords, name, positions) + return DataArray(data, coords, da.dims, da.name, da.attrs) + + +class Polyphase(Atom): + """ + Stateful polyphase resampler: upsample, FIR filter and downsample in one pass. + + Computes what :class:`UpSample` → FIR :class:`LFilter` → :class:`DownSample` + computes, but only the output samples that survive the decimation, and + without ever materialising the zero-stuffed signal + (:func:`scipy.signal.upfirdn`). The linear-phase group delay of `taps` is + removed from the coordinate, as :class:`~xdas.atoms.FIRFilter` does. + + The taps are cast down to the data dtype when the data is less precise, so + float32 input stays float32 instead of being promoted by the filter. + + Chunked calls carry the filter memory and the output-grid phase across + chunks, and every call emits every output sample the stream can support so + far — nothing is held back, so no flush is needed. + + Parameters + ---------- + taps : array-like + FIR coefficients, designed at the *upsampled* rate ``up * fs``. At + least `up` of them are needed, one per polyphase branch. + up : int, optional + Upsampling factor. Default is 1. + down : int, optional + Downsampling factor. Default is 1. + dim : str or int, optional + Dimension to resample along. Defaults to ``"last"``. + parallel : int, bool, or None, optional + Worker count for parallelisation. + + Examples + -------- + >>> import numpy as np + >>> import scipy.signal as sp + >>> import xdas as xd + >>> from xdas.atoms import Polyphase + + >>> da = xd.testing.dummy(shape=(100, 3)) + >>> taps = sp.firwin(21, 0.4) + >>> Polyphase(taps, down=2, dim="time")(da).sizes["time"] + 50 + + Splitting the input does not change the result: + + >>> eager = Polyphase(taps, up=2, down=5, dim="time")(da) + >>> atom = Polyphase(taps, up=2, down=5, dim="time") + >>> chunked = xd.concat(list(atom.iter_chunks(xd.split(da, 7, "time"))), "time") + >>> bool(np.allclose(chunked.values, eager.values)) + True + + """ + + def __init__(self, taps, up=1, down=1, dim="last", parallel=None): + super().__init__() + self.taps = taps + self.up = up + self.down = down + self.dim = dim + self.parallel = parallel + self.axis = State(...) + self.buffer = State(...) + self.consumed = State(...) + + @property + def lag(self): + """Group delay of the taps, in upsampled samples.""" + return (np.asarray(self.taps).size - 1) // 2 + + @property + def phase(self): + """Period, in input samples, of the output-grid phase.""" + return self.down // math.gcd(self.up, self.down) + + def _history_size(self): + """Input samples to keep: the filter memory plus one phase period.""" + memory = -(-(np.asarray(self.taps).size - 1) // self.up) + return memory + self.phase - 1 + + def initialize(self, da, chunk_dim=None, **flags): + """Set the axis and allocate the history buffer for chunked operation.""" + if np.asarray(self.taps).size < self.up: + raise ValueError( + f"at least {self.up} taps are needed to resample by {self.up}/" + f"{self.down} (one per polyphase branch), got " + f"{np.asarray(self.taps).size}" + ) + self.axis = State(da.get_axis_num(self.dim)) + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._dim_name(da) + if dim == chunk_dim: + shape = tuple( + self._history_size() if name == dim else size + for name, size in da.sizes.items() + ) + self.buffer = State(np.zeros(shape, dtype=da.dtype)) + self.consumed = State(0) + else: + self.buffer = State(None) + self.consumed = State(None) + + def call(self, da, **flags): + """Resample *da*, carrying the filter memory and grid phase if chunked.""" + size = da.sizes[self.dim] + if size == 0: + return [] + axis = self.axis + up, down = self.up, self.down + chunked = self.buffer is not None + start = self.consumed if chunked else 0 + # The surviving outputs are those whose upsampled index falls in this + # chunk; both bounds are ceils since the grid starts on sample zero. + first = -(-start * up // down) + stop = -(-(start + size) * up // down) + # Prepend enough past input to warm up the filter, choosing an amount + # that puts the chunk start on the output grid so `upfirdn`, which + # always emits from its own sample zero, lands on the global phase. + memory = -(-(np.asarray(self.taps).size - 1) // self.up) + nhist = memory + (start - memory) % self.phase + values = da.values + if chunked: + history = self.buffer[ + _along(axis, values.ndim, slice(self.buffer.shape[axis] - nhist, None)) + ] + else: + shape = tuple( + nhist if index == axis else length + for index, length in enumerate(values.shape) + ) + history = np.zeros(shape, dtype=values.dtype) + data = self._resample(np.concatenate([history, values], axis), axis) + if chunked: + self.buffer = State(self._keep_history(values, axis)) + self.consumed = State(start + size) + if stop <= first: + return [] + offset = (start - nhist) * up // down + data = data[_along(axis, data.ndim, slice(first - offset, stop - offset))] + return DataArray( + data, self._coords(da, first, stop, start), da.dims, da.name, da.attrs + ) + + def _resample(self, values, axis): + """Run the polyphase filter over *values*, keeping the data precision.""" + taps = np.asarray(self.taps) + if np.issubdtype(values.dtype, np.floating) and ( + values.dtype.itemsize < taps.dtype.itemsize + ): + taps = taps.astype(values.dtype) + across = int(axis == 0) + func = parallelize((None, across, None, None, None), across, self.parallel)( + sp.upfirdn + ) + return func(taps, values, self.up, self.down, axis) + + def _keep_history(self, values, axis): + """Return the last `_history_size` input samples of the stream so far.""" + size = values.shape[axis] + length = self.buffer.shape[axis] + if size >= length: + return values[_along(axis, values.ndim, slice(size - length, None))] + kept = self.buffer[_along(axis, values.ndim, slice(size, None))] + return np.concatenate([kept, values], axis) + + def _coords(self, da, first, stop, start): + """Build the output coordinates on the resampled, delay-corrected grid.""" + # The "first"/"last" alias would name a new dimension if it reached the + # `Coordinate` built below, so resolve it here. + name = self._dim_name(da) + coord = da.coords[name] + delta = get_sampling_interval(da, name, cast=False) + size = stop - first + # Output `index` sits `index * down - lag` upsampled samples after the + # start of the run, hence that many minus `start * up` after this chunk. + shifts = np.array([first, stop - 1]) * self.down - self.lag - start * self.up + grid = coord.start + self._upsampled(shifts, delta) + origin, last = grid[0], grid[1] + step = self._upsampled(self.down, delta) + if size > 1: + tie_indices, tie_values = [0, size - 1], [origin, last] + drift = abs((last - origin) - (size - 1) * step) + else: + tie_indices, tie_values = [0], [origin] + drift = 0 * step + data = {"tie_indices": tie_indices, "tie_values": tie_values} + if coord.isregular(): + # A rate that the coordinate resolution cannot represent exactly + # makes the tie values drift from the nominal step; declare that + # drift as jitter rather than refusing to call the output regular. + tolerance = getattr(coord, "tolerance", None) + base = parse_scalar_delta(tolerance, coord.dtype, default_zero=True) + data["sampling_interval"] = step + data["tolerance"] = base + drift + coords = da.coords.copy() + coords[name] = Coordinate(data, name) + # Output `index` is drawn from input sample `index * down // up`. + # Flooring — never rounding — is what keeps the position inside this + # chunk by construction: rounding a half-integer up can walk past the + # last input sample the chunk holds, and the clip that would be needed + # to pull it back lands on a different sample than the eager call. + positions = (np.arange(first, stop) * self.down) // self.up - start + return _carry_labels(coords, name, positions) + + def _upsampled(self, count, delta): + """Return the span of *count* upsampled samples, at coordinate resolution.""" + if np.issubdtype(np.asarray(delta).dtype, np.timedelta64): + return (count * delta) // self.up + return count * delta / self.up + + +class Rechunk(Atom): + """ + Merge and split incoming chunks to a fixed size along a dimension. + + Chunk sizes are a performance knob, not science, so they are given in + samples (as in ``process(chunks=...)``): the canonical use is restoring a + workable cadence after a decimation shrank the chunks. Rechunking never + merges across a discontinuity — the seam handling of the base class + flushes the partial buffer at every gap — so chunks stay internally + regular through this stage. Each call returns zero or more chunks of + exactly the target size; :meth:`flush` drains the remainder. + + Eager calls (whole records) pass through unchanged: re-joining the emitted + chunks would reproduce the input. + + Parameters + ---------- + chunks : dict + Mapping of a unique dimension name to the target chunk size in + samples, e.g. ``{"time": 1000}``. + + Examples + -------- + >>> import xdas as xd + >>> from xdas.atoms import Rechunk + >>> da = xd.testing.dummy(shape=(100, 10)) + >>> atom = Rechunk({"time": 30}) + >>> [out.sizes["time"] for out in atom.iter_chunks(xd.split(da, 4, "time"))] + [30, 30, 30, 10] + + """ + + def __init__(self, chunks): + super().__init__() + if not (isinstance(chunks, dict) and len(chunks) == 1): + raise TypeError( + "`chunks` must be a dict that maps a unique " + "dimension to a unique size: {'dim': int}" + ) + ((dim, size),) = chunks.items() + if not (isinstance(size, int) and size > 0): + raise ValueError("the chunk size must be a strictly positive integer") + self.dim = dim + self.size = size + self.buffer = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Initialise the carry-over buffer for chunked operation.""" + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._dim_name(da) + if chunk_dim == dim: + self.buffer = State(da.isel({self.dim: slice(0, 0)})) + else: + self.buffer = State(None) + + def call(self, da, **flags): + """Emit full-size chunks from the buffered stream, keep the remainder.""" + if self.buffer is None: + return da + da = concat([self.buffer, da], self.dim) + divpoint = da.sizes[self.dim] - da.sizes[self.dim] % self.size + out = [ + da.isel({self.dim: slice(index, index + self.size)}) + for index in range(0, divpoint, self.size) + ] + self.buffer = State(da.isel({self.dim: slice(divpoint, None)})) + return out + + def flush(self): + """Emit the remaining partial chunk.""" + if not isinstance(self.buffer, DataArray) or self.buffer.sizes[self.dim] == 0: + return [] + out = self.buffer + self.buffer = State(out.isel({self.dim: slice(0, 0)})) + return [out] + + +rechunk = atomized(Rechunk) diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index 8d892b1f..1a7e425a 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -1,15 +1,20 @@ """ -Machine-learning atom: :class:`MLPicker` wraps SeisBench models as pipeline atoms. +Machine-learning atom: :class:`Annotate` wraps SeisBench models as pipeline atoms. Torch and SeisBench are loaded lazily so they remain optional dependencies. """ import importlib +import warnings +from collections import deque +from fnmatch import fnmatch import numpy as np -from ..core import DataArray, concat -from .core import Atom, State +from ..core import DataArray, concat, stack +from .core import Atom, Sequential, State, atomized +from .detect import Trigger +from .tasks import Filter, Resample class LazyModule: @@ -26,54 +31,746 @@ def __getattr__(self, name): except ImportError: raise ImportError( f"{self._name} is not installed by default, " - f"please install is manually" + f"please install it manually" ) return getattr(self._module, name) torch = LazyModule("torch") +#: Horizontal component letters SeisBench treats as interchangeable when its +#: ``flexible_horizontal_components`` argument is set (which is its default), +#: so that ``Z12H`` weights accept ``ZNE`` data and vice versa. +FLEXIBLE_HORIZONTALS = {"1": "N", "N": "1", "2": "E", "E": "2"} -class MLPicker(Atom): +#: Fallbacks for the annotate arguments these atoms read. SeisBench's own +#: ``_argdict_get_with_default`` falls back to the model's ``_annotate_args``, +#: but ``blinding`` is not in the ``WaveformModel`` base and a model need not +#: declare one at all, so the accessor needs a default of its own rather than a +#: subscript of ``None``. +ANNOTATE_DEFAULTS = { + "overlap": 0, + "stacking": "avg", + "flexible_horizontal_components": True, +} + +#: Number of corners ``obspy``'s ``Stream.filter`` uses when the weight set +#: does not say, and the order :class:`~xdas.atoms.Filter` defaults to. The two +#: agree, which is what makes a model-declared filter reproducible exactly. +OBSPY_CORNERS = 4 + +#: Collection level names :meth:`Annotate.gather` will consider collapsing into +#: a component dimension. A *default*, not a rule: ``components="whatever"`` +#: replaces the list, exactly as SeisBench's ``guess_channel_coord_name`` walks +#: a candidate list on the input side. Our engines name the level ``channel``, +#: but nothing in the implementation may assume the literal. +COMPONENT_LEVELS = ("channel", "component") + +#: Lengths a collection key may have and still be read as a channel code: three +#: for a SEED code (band, instrument, orientation) and one for a bare +#: orientation letter. This is what keeps the gather off a station level, since +#: the last-letter matcher alone resolves ``STA1``/``STA2`` cleanly against +#: ``ZNE`` weights — see :func:`component_keys`. +CHANNEL_CODE_LENGTHS = (1, 3) + +#: Label of the noise class, which is never picked. SeisBench identifies it by +#: this exact name — ``classify_aggregate`` skips ``phase == "N"`` — and so do +#: we, rather than by position: the label *order* is a property of the weight +#: set and flips between them (``NPS`` for PhaseNet's ``original``, ``PSN`` for +#: most others), so a positional rule would silently pick noise. +NOISE_LABEL = "N" + +#: Detection threshold used when neither the weight set's ``default_args`` nor +#: the model's ``_annotate_args`` declares one. SeisBench's ``PhaseNet`` +#: documents the same value under the ``*_threshold`` key. +DEFAULT_THRESHOLD = 0.3 + + +def annotate_arg(model, argdict, key): """ - Wraps a SeisBench phase-picking model as a streaming :class:`Atom`. + Read one annotate argument, SeisBench-style but crash-free. - Uses an overlapping sliding-window strategy to apply the model to - arbitrarily long data and to stitch the per-segment probability outputs - back into a continuous DataArray. + Mirrors ``WaveformModel._argdict_get_with_default``: what the call site + passes wins, else what the weight set declares, else the model's + ``_annotate_args`` documentation default. The last step is where SeisBench + raises on a key its base class does not declare, so this falls back to + :data:`ANNOTATE_DEFAULTS` instead. The two SeisBench internals this module + relies on are isolated here and nowhere else. Parameters ---------- model : seisbench.models.WaveformModel - A SeisBench model in evaluation mode (will be moved to *device*). + The model whose ``_annotate_args`` documents the fallbacks. + argdict : dict + The weight set's ``default_args`` merged with the call's kwargs. + key : str + Name of the annotate argument. + + Returns + ------- + value : object + The value of *key*, from the argdict, the model or the fallbacks. + """ + if key in argdict: + return argdict[key] + entry = getattr(model, "_annotate_args", {}).get(key) + if entry is None: + return ANNOTATE_DEFAULTS[key] + return entry[1] + + +def model_phases(model): + """ + Return the labels a model's output classes carry, in order. + + ``labels`` is per-weight-set and SeisBench lets a weight set leave it unset + or make it a callable, resolving both only when it builds its output + stream. ``None`` falls back to positional labels ``0, 1, ... classes - 1``, + exactly as ``WaveformModel._predictions_to_stream`` does. A callable is + refused: SeisBench feeds it the identity of the trace being annotated, + which these atoms do not carry, and a silently wrong ``phase`` coordinate + would key :class:`~xdas.atoms.Trigger`'s thresholds onto the wrong class. + + Parameters + ---------- + model : seisbench.models.WaveformModel + The model whose weight set is read. + + Returns + ------- + list + One label per output class. + """ + labels = model.labels + if labels is None: + classes = getattr(model, "classes", None) + if classes is None: + raise ValueError( + "the model declares neither `labels` nor `classes`, so " + "neither the number nor the names of its outputs can be " + "known: set `model.labels` on the weight set" + ) + return list(range(classes)) + if callable(labels): + raise TypeError( + "the model's `labels` is a callable, which SeisBench resolves " + "from the identity of the trace it annotates; this atom " + "annotates arrays, so set `model.labels` to the list of " + "output names instead" + ) + return list(labels) + + +def model_pick_labels(model): + """ + List the labels of *model* that are picked, as SeisBench picks them. + + A model may declare which of its outputs are phases in a ``phases`` + attribute — beware the name, SeisBench's ``model.phases`` is this *subset* + while :attr:`Annotate.phases` is the full label list. ``classify_aggregate`` + walks exactly that attribute, so ``EQTransformer``'s detection trace and + ``EQTP``'s polarities are left out of the picking the way noise is: still + emitted in the characteristic function, never triggered on. A model + declaring no subset gets all its labels but :data:`NOISE_LABEL`. + + Parameters + ---------- + model : seisbench.models.WaveformModel + The model whose weight set is read. + + Returns + ------- + list + One label per picked class. + + Examples + -------- + >>> from xdas.atoms.ml import model_pick_labels + + >>> class PhaseNet: + ... labels = "NPS" + >>> model_pick_labels(PhaseNet()) + ['P', 'S'] + + >>> class EQTransformer(PhaseNet): + ... labels, phases = ("Detection", "P", "S"), "PS" + >>> model_pick_labels(EQTransformer()) + ['P', 'S'] + """ + phases = getattr(model, "phases", None) + if phases is None: + return [label for label in model_phases(model) if label != NOISE_LABEL] + return list(phases) + + +def _model_thresholds(model, **annotate_kwargs): + """ + Build the per-phase detection thresholds a weight set declares. + + One entry per picked label, as :func:`model_pick_labels` resolves them — + SeisBench's ``classify_aggregate`` picks no others, and leaving them out of + the mapping is what stops :class:`~xdas.atoms.Trigger` triggering on them + while :class:`Annotate` keeps emitting them. Each threshold is looked up as + SeisBench does: what the call passes wins, else what the weight set + declares in ``default_args[f"{label}_threshold"]``, else the model's own + documented default for that key, else the ``*_threshold`` catch-all, else + :data:`DEFAULT_THRESHOLD`. + + Values are passed through faithfully, including thresholds above one — + PhaseNet's ``iquique`` declares ``P_threshold = 1.12``, which simply never + fires. That is the weight set's own metadata, not ours to clamp. + + Parameters + ---------- + model : seisbench.models.WaveformModel + The model whose weight set is read. + **annotate_kwargs + SeisBench annotate arguments overriding the weight set's, including + the ``f"{label}_threshold"`` keys themselves. + + Returns + ------- + dict + One threshold per picked label, keyed on the label. + + Examples + -------- + >>> from xdas.atoms.ml import _model_thresholds + + A weight set declaring nothing falls back to 0.3 per phase, and the noise + class gets no entry however the labels are ordered: + + >>> class Plain: + ... labels, default_args = "NPS", {} + >>> _model_thresholds(Plain()) + {'P': 0.3, 'S': 0.3} + + What the weight set declares wins, and what the call passes wins over that: + + >>> class Geofon(Plain): + ... labels = "PSN" + ... default_args = {"P_threshold": 0.57, "S_threshold": 0.073} + >>> _model_thresholds(Geofon()) + {'P': 0.57, 'S': 0.073} + >>> _model_thresholds(Geofon(), S_threshold=0.5) + {'P': 0.57, 'S': 0.5} + + A detection trace is no more picked than noise is: + + >>> class EQTransformer(Plain): + ... labels, phases = ("Detection", "P", "S"), "PS" + >>> _model_thresholds(EQTransformer()) + {'P': 0.3, 'S': 0.3} + """ + argdict = dict(model.default_args) | annotate_kwargs + annotate_args = getattr(model, "_annotate_args", {}) + fallback = annotate_args.get("*_threshold", (None, DEFAULT_THRESHOLD))[1] + thresholds = {} + for label in model_pick_labels(model): + key = f"{label}_threshold" + if key in argdict: + value = argdict[key] + elif key in annotate_args: + value = annotate_args[key][1] + else: + value = fallback + thresholds[label] = float(value) + return thresholds + + +def resolve_sample_dim(da, dim): + """ + Resolve *dim* against *da*, accepting the ``first``/``last`` aliases. + + Parameters + ---------- + da : DataArray + The input whose dimensions *dim* is resolved against. dim : str - Dimension name along which the model slides (usually ``"time"``). - device : str or torch.device, optional - Torch device. Defaults to CUDA if available, else CPU. + A dimension name, or ``"first"`` or ``"last"``. + + Returns + ------- + str + The name of the resolved dimension. + """ + if dim == "first": + return da.dims[0] + if dim == "last": + return da.dims[-1] + if dim not in da.dims: + raise ValueError(f"{dim!r} is not a dimension of the input (got {da.dims})") + return dim + + +def component_labels(da, dim): + """ + Return the string labels of *dim*, or ``None`` if it carries none. + + Labels are read as text only: a numeric coordinate is excluded by dtype + rather than by failing to match, without which an integer identifier axis + of ``[1, 2]`` would resolve cleanly against ``Z12H`` weights. + + Parameters + ---------- + da : DataArray + The input carrying the coordinate. + dim : str + Name of the dimension to read. + + Returns + ------- + list of str or None + The labels, or ``None`` when *dim* has no string coordinate. + """ + if dim not in da.coords: + return None + coord = da.coords[dim] + if coord.dtype.kind not in "SU": + return None + return [ + value.decode() if isinstance(value, bytes) else str(value) + for value in coord.values + ] + + +def match_components(labels, order, flexible=True): + """ + Map *labels* onto model input slots, or ``None`` if they are not components. + + Parameters + ---------- + labels : list of str + The labels to match, each ending with a component letter. + order : str + The model's ``component_order``, e.g. ``"ENZ"`` or ``"Z12H"``. + flexible : bool, optional + Whether ``1``/``N`` and ``2``/``E`` are interchangeable, as SeisBench's + ``flexible_horizontal_components`` makes them by default. + + Returns + ------- + list of int or None + One slot index per label, or ``None`` if any label names no component. + """ + order = list(order) + slots = [] + for label in labels: + letter = label[-1:] + if letter not in order and flexible: + letter = FLEXIBLE_HORIZONTALS.get(letter, letter) + if letter not in order: + return None + slots.append(order.index(letter)) + return slots + + +def component_slots(da, dim, order, flexible=True): + """ + Resolve the labels of *dim* into distinct model slots, or ``None``. + + Parameters + ---------- + da : DataArray + The input carrying the coordinate. + dim : str + Name of the candidate component dimension. + order : str + The model's ``component_order``. + flexible : bool, optional + Whether horizontal components are matched flexibly. + + Returns + ------- + list of int or None + One slot index per label, or ``None`` when *dim* is not a component + dimension. + """ + labels = component_labels(da, dim) + if labels is None: + return None + slots = match_components(labels, order, flexible) + if slots is None: + return None + if len(set(slots)) != len(slots): + raise ValueError( + f"the {dim!r} dimension repeats component orientations " + f"({labels}): it mixes several instruments, split them first" + ) + return slots + + +def is_channel_code(key): + """ + Whether *key* has the shape of a channel code. + + True for a three-character SEED code (band, instrument, orientation) and + for a bare orientation letter, false for anything longer, for a purely + numeric key, and for a non-string one. + + Parameters + ---------- + key : object + A key of a collection level. + + Returns + ------- + bool + Whether *key* may be read as a channel code. + + Examples + -------- + >>> from xdas.atoms.ml import is_channel_code + >>> [is_channel_code(key) for key in ("SHZ", "Z", "STA1", "S-Z", "2", 2)] + [True, True, False, False, False, False] + """ + return ( + isinstance(key, str) + and len(key) in CHANNEL_CODE_LENGTHS + and key.isalnum() + and not key.isdigit() + ) + + +def component_keys(keys, order, flexible=True, level=None): + """ + Resolve the keys of a collection level into distinct model input slots. + + The key rule of :meth:`Annotate.gather`, and deliberately stricter than + :func:`match_components`, which is what recognises a component + *dimension*. The flexible horizontal matching makes any label ending in + ``1`` or ``2`` a horizontal, so a station level keyed ``STA1``/``STA2`` + resolves cleanly against ``ZNE`` weights. On a dimension that is harmless, + since the duplicate and count checks catch it; on a tree *level* it would + silently fold two stations into one instrument. So a key must look like a + whole channel code — :func:`is_channel_code` — and the keys must agree on + their length and on their band code, which is the character before the + instrument code. Only the instrument code may vary, because it does: an + OBS station is ``BHZ``/``BH1``/``BH2``/``BDH``, whose hydrophone is a + pressure instrument and whose stem is therefore *not* common. + + Three outcomes, and the distinction between the last two is the point: + + - every key resolves to a distinct slot: the level is the component + dimension, and the slots are returned; + - no key resolves: the level is not a component level despite its name, + and ``None`` is returned so the caller walks its leaves — several DAS + formats call their spatial axis ``channel``; + - some keys resolve: someone meant components and the data disagrees, so + this raises, naming the conflict. + + Parameters + ---------- + keys : iterable + The keys of the collection level. + order : str + The model's ``component_order``, e.g. ``"ENZ"`` or ``"Z12H"``. + flexible : bool, optional + Whether ``1``/``N`` and ``2``/``E`` are interchangeable, as SeisBench's + ``flexible_horizontal_components`` makes them by default. + level : str, optional + Name of the level, used in the error messages only. + + Returns + ------- + list of int or None + One slot index per key, or ``None`` when no key names a component. + + Raises + ------ + ValueError + If only some of the keys name a component, if two of them name the + same one, or if they do not agree on their length and band code. + + Examples + -------- + >>> from xdas.atoms.ml import component_keys + + A station's three components resolve, in the model's own order: + + >>> component_keys(["SHZ", "SHN", "SHE"], "ENZ") + [2, 1, 0] + + A DAS cable whose spatial level happens to be called ``channel`` does not, + which is what sends the caller back to walking the leaves: + + >>> component_keys(["0", "1", "2"], "ZNE") is None + True + """ + keys = list(keys) + resolved = {} + unresolved = [] + for key in keys: + slots = ( + match_components([key], order, flexible) if is_channel_code(key) else None + ) + if slots is None: + unresolved.append(key) + else: + resolved[key] = slots[0] + where = "the level" if level is None else f"the {level!r} level" + if not resolved: + return None + if unresolved: + raise ValueError( + f"{where} names components ({sorted(resolved)}) and other things " + f"({unresolved}) at once: it cannot be collapsed into a component " + f"dimension of {order!r}, split it or rename it" + ) + slots = list(resolved.values()) + if len(set(slots)) != len(slots): + raise ValueError( + f"{where} repeats component orientations ({keys}): it holds " + "several instruments, split them first" + ) + if len({len(key) for key in keys}) > 1 or len({key[:-1][:1] for key in keys}) > 1: + raise ValueError( + f"{where} does not name one instrument ({keys}): its keys differ " + "by more than their orientation, split them first" + ) + return slots + + +def resolve_component_dim(da, sample_dim, order, components=None, flexible=True): + """ + Find the component dimension of *da* and the slots its labels name. + + Detection is by labels, never by name: the component dimension is the one + whose labels each end with a distinct letter of *order*. + + Parameters + ---------- + da : DataArray + The input to inspect. + sample_dim : str + The already-resolved sample dimension, never a component candidate. + order : str + The model's ``component_order``. + components : str, False or None, optional + Name of the component dimension, ``False`` to disable detection, or + ``None`` (default) to detect it. + flexible : bool, optional + Whether horizontal components are matched flexibly. + + Returns + ------- + dim : str or None + The component dimension, or ``None`` when the data has none. + slots : list of int or None + The model input slots its labels name, or ``None`` with *dim*. + """ + if components is False: + return None, None + if components is not None: + if components not in da.dims: + raise ValueError( + f"{components!r} is not a dimension of the input (got {da.dims})" + ) + slots = component_slots(da, components, order, flexible) + if slots is None: + raise ValueError( + f"the {components!r} dimension is not labelled by components: " + f"expected labels ending with distinct letters of {order!r}" + ) + return components, slots + found = {} + for dim in da.dims: + if dim == sample_dim: + continue + slots = component_slots(da, dim, order, flexible) + if slots is not None: + found[dim] = slots + if not found: + return None, None + if len(found) > 1: + raise ValueError( + f"several dimensions could be the component dimension " + f"({sorted(found)}): name it explicitly with `components=`" + ) + return found.popitem() + + +class Annotate(Atom): + """ + Wraps a SeisBench model as a streaming :class:`Atom`. + + Slides the model over the data with the overlap the weight set declares, + stitches the per-window outputs back into one continuous characteristic + function and appends the model's labels as a ``phase`` dimension. Every + parameter the model can decide — the window overlap, the stacking rule, the + normalisation, the blinding — is read off the model instance rather than + assumed, since in SeisBench all of them belong to the *weight set* and not + to the architecture. + + Input dimension names are never assumed: *dim* names the sample dimension + (the repository's ``"first"`` and ``"last"`` aliases resolve too) and the + component dimension is found by its labels, not by its name. The output + keeps the input's order among the remaining dimensions but is laid out + sample-last, ``(..., "phase", dim)``: the characteristic function of one + phase of one channel is then contiguous, which is the layout its consumers + reduce along. + + Parameters + ---------- + model : seisbench.models.WaveformModel + A SeisBench model in evaluation mode (will be moved to *device*). + dim : str, optional + Dimension the model slides along. Defaults to ``"time"``; ``"first"`` + and ``"last"`` resolve against the input. + components : str or False, optional + Name of the component dimension. ``None`` (default) detects it by its + labels: the dimension whose labels each end with a distinct letter of + the model's ``component_order``. ``False`` disables detection, which is + the escape hatch when an identifier axis carries labels that could + collide with component letters. + + On a collection, this also names the *level* to gather (see + :meth:`gather`), overriding the :data:`COMPONENT_LEVELS` candidates; + ``False`` disables the gather along with the detection. component_strategy : str, optional - How to fill the channel dimension: ``"clone"`` replicates the - single-component signal, or pass a component letter to select it. + How the model's input slots are filled from the data: + + ``"auto"`` (default) + Reproduces SeisBench: ``"clone"`` when the data has no component + dimension (its DAS wrapper's default), ``"pad"`` when it has a + partial one (its ``strict=False`` default). + ``"clone"`` + Replicate the single available signal into every slot. + ``"pad"`` + Single signal in the first slot, zeros in the rest. + ``"E"``, ``"N"``, ``"Z"``, ... + Single signal in that named slot of ``component_order``, zeros in + the rest. + ``"strict"`` + Any missing component is an error (SeisBench's ``strict=True``). + + With a component dimension, present components always go to the slot + their label names and the remaining slots are zeroed. + device : str or torch.device, optional + Torch device. Defaults to CUDA if available, else CPU. + tolerance : scalar, None or False, optional + Grid-snapping budget forwarded to :func:`xdas.stack` when a component + level is gathered (see :meth:`gather`). ``None`` (default) spends a + hundredth of a sample, which is what lets three components whose start + times were rounded a nanosecond apart stack; ``False`` restores strict + equality. + max_buffers : int, optional + Depth of the in-flight output queue on a CUDA device, the same + bounded-staging pattern :class:`~xdas.processing.DataArrayLoader` + uses. Each completed window's device-to-host transfer is issued + asynchronously (pinned, non-blocking) and :meth:`call` emits only + the outputs whose transfer has completed — the 0..n contract already + allows late emission — so the CPU keeps feeding the model while + results cross back; :meth:`flush` drains whatever is still in + flight. At most `max_buffers` transfers are left pending (default + 2); ``0`` restores fully synchronous emission. On the CPU the + transfers complete immediately and the queue changes nothing. + **annotate_kwargs + SeisBench annotate arguments (``overlap``, ``stacking``, ``blinding``, + ...) overriding what the weight set declares in ``default_args``. + + Warnings + -------- + ``component_strategy="pad"`` is *positional*: like SeisBench it fills slot + 0 of ``component_order``, which is ``Z`` for the many ``ZNE`` weight sets + but ``E`` for ``ENZ`` ones such as PhaseNet's ``original``. Pass the letter + itself (``"Z"``) to name a slot rather than count to it. + + Examples + -------- + >>> import torch + >>> import xdas as xd + >>> from xdas.atoms import Annotate + + Any SeisBench model works; here is a stand-in small enough to inline, with + a four-sample window, two labels and a 50 % overlap declared by its weights: + + >>> class Model(torch.nn.Module): + ... in_samples, in_channels, classes = 4, 3, 2 + ... labels, component_order = "PS", "ZNE" + ... default_args = {"overlap": 0.5} + ... def annotate_batch_pre(self, batch, argdict): + ... return batch + ... def annotate_batch_post(self, batch, piggyback, argdict): + ... return torch.transpose(batch, -1, -2) + ... def forward(self, batch): + ... return batch[:, : self.classes] + + >>> da = xd.testing.dummy(dims=("time", "distance"), shape=(16, 3)) + >>> atom = Annotate(Model(), dim="time", device="cpu") + >>> result = atom(da) + + The labels become a ``phase`` dimension and the samples end up last: + + >>> result.dims + ('distance', 'phase', 'time') + >>> result.coords["phase"].values + array(['P', 'S'], dtype='>> result.sizes["time"] == da.sizes["time"] + True + + Given a collection whose ``channel`` level holds one component per key, + the atom takes that level as its component dimension rather than + annotating each component on its own (see :meth:`gather`): + + >>> traces = {code: da.isel(distance=0) for code in ("SHZ", "SHN", "SHE")} + >>> atom(xd.DataCollection(traces, "channel")).dims + ('phase', 'time') """ - def __init__(self, model, dim, device=None, component_strategy="clone"): + def __init__( + self, + model, + dim="time", + components=None, + component_strategy="auto", + device=None, + tolerance=None, + max_buffers=2, + **annotate_kwargs, + ): super().__init__() if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") else: device = torch.device(device) - valid = {"clone", *model.component_order} + valid = {"auto", "clone", "pad", "strict", *model.component_order} if component_strategy not in valid: - raise ValueError(f"component_strategy must be one of {valid}") + raise ValueError( + f"component_strategy must be one of {sorted(valid)}, " + f"got {component_strategy!r}" + ) self.device = device self.model = model.eval().to(self.device) self.dim = dim + self.components = components self.component_strategy = component_strategy + self.tolerance = tolerance + self.max_buffers = max_buffers + self.argdict = dict(model.default_args) | annotate_kwargs + if self.stacking not in ("avg", "max"): + raise ValueError(f"stacking must be 'avg' or 'max', got {self.stacking!r}") + if not 0 <= self.noverlap < self.nperseg: + raise ValueError( + f"overlap must be shorter than one model window " + f"({self.noverlap} of {self.nperseg} samples)" + ) + self.sample_dim = State(...) + self.component_dim = State(...) + self.batch_dims = State(...) + self.out_dims = State(...) + self.slots = State(...) + self.started = State(...) self.buffer = State(...) - self.batch_size = State(...) self.circular_input = State(...) - self.model_input = State(...) # TODO: should not be a state + self.model_input = State(...) self.circular_output = State(...) self.circular_counts = State(...) + self.inflight = State(...) + + def _annotate_arg(self, key): + """Read one annotate argument through :func:`annotate_arg`.""" + return annotate_arg(self.model, self.argdict, key) @property def nperseg(self): @@ -82,139 +779,976 @@ def nperseg(self): @property def noverlap(self): - """Number of overlapping samples between consecutive segments.""" - return self.nperseg // 2 + """Overlap between consecutive segments, in samples, as the model declares it.""" + overlap = self._annotate_arg("overlap") + return int(overlap * self.nperseg) if overlap < 1 else int(overlap) @property def step(self): """Stride between the start of consecutive segments.""" return self.nperseg - self.noverlap + @property + def stacking(self): + """How overlapping windows are combined: ``"avg"`` or ``"max"``.""" + return self._annotate_arg("stacking") + @property def phases(self): - """List of phase label strings produced by the model.""" - return list(self.model.labels) + """ + List of the phase labels produced by the model. + + Read off the weight set through :func:`model_phases`, which is also + what :func:`_model_thresholds` keys on, so the ``phase`` coordinate and + the thresholds of a pipeline agree by construction. + """ + return model_phases(self.model) @property def in_channels(self): """Number of input channels the model expects.""" - return self.model.in_channels + return getattr(self.model, "in_channels", len(self.model.component_order)) @property def classes(self): - """Number of output classes (phases) the model produces.""" - return self.model.classes + """ + Number of output classes the model produces. + + Read off the labels rather than ``model.classes``, which counts only + what the architecture's *picking* head emits: ``EQTransformer`` sets it + to 2 while labelling three outputs, the third being its detection + trace. + """ + return len(self.phases) @property - def blinding(self): - """``(left, right)`` blinding samples from the model's default args.""" - return self.model.default_args["blinding"] + def fill(self): + """Neutral element of the stacking rule, used to reset freed samples.""" + return 0.0 if self.stacking == "avg" else -np.inf - def initialize(self, da, chunk_dim=None, **flags): - """Allocate circular buffers sized to *da*'s batch and segment dimensions.""" - self.batch_size = State( - np.prod([size for dim, size in da.sizes.items() if dim != self.dim]) + def gather(self, mapping): + """ + Collapse a component level of a collection into a component dimension. + + So that annotating a collection needs no prior :func:`xdas.stack`: an + ObsPy-style tree keeps one leaf per channel, but three channels of one + instrument are one input to the model, not three. What counts as a + component is a property of the *model* — ``Z12H`` and ``ENZ`` disagree + about which channels group together — which is why the reader cannot + do this and this atom can. + + The gather is deliberately conservative, because folding a station + level into a component axis would silently destroy the distinction + between stations. It fires only when **both** the level's name is one + that *components* accepts (:data:`COMPONENT_LEVELS` by default) and + its keys resolve to distinct components under + :func:`component_keys`. A level whose keys resolve but whose name is + not a candidate is left alone; a level whose name is a candidate but + whose keys resolve to nothing is left alone too, which is what lets a + DAS collection whose spatial level is called ``channel`` walk + straight past. Only a level that is clearly component-ish and + malformed raises. + + This is the one place a *name* is consulted: detection on the array + side is purely label-based. The asymmetry is deliberate — a + mis-detected dimension is caught at once by the count and letter + checks, a mis-collapsed level is not. + + Parameters + ---------- + mapping : DataMapping + The collection level about to be walked. + + Returns + ------- + DataArray, DataCollection or None + The level collapsed onto a dimension named after it, or ``None`` + to walk its leaves. + """ + if self.components is False: + return None + explicit = isinstance(self.components, str) + candidates = (self.components,) if explicit else COMPONENT_LEVELS + level = getattr(mapping, "name", None) + if level not in candidates: + return None + keys = list(mapping) + slots = component_keys( + keys, + self.model.component_order, + self._annotate_arg("flexible_horizontal_components"), + level, ) - self.circular_input = State( - torch.zeros( - self.batch_size, self.nperseg, dtype=torch.float32, device=self.device - ) + if slots is None: + if explicit: + raise ValueError( + f"could not identify the component level of the collection: " + f"tried {list(candidates)}, and the keys of {level!r} " + f"({keys}) name no component of " + f"{self.model.component_order!r}. Please provide it " + "explicitly with `components=`, or pass `components=False` " + "to walk the collection leaf by leaf" + ) + return None + return stack(mapping, level, tolerance=self.tolerance) + + def initialize(self, da, chunk_dim=None, **flags): + """Resolve the dimensions and allocate the sliding-window buffers.""" + dim = resolve_sample_dim(da, self.dim) + component_dim, slots = resolve_component_dim( + da, + dim, + self.model.component_order, + self.components, + self._annotate_arg("flexible_horizontal_components"), ) - self.model_input = State( - torch.zeros( - self.batch_size, - self.in_channels, - self.nperseg, - dtype=torch.float32, - device=self.device, - ) + self.sample_dim = State(dim) + self.component_dim = State(component_dim) + self.slots = State(self._resolve_slots(slots)) + self.batch_dims = State( + tuple(other for other in da.dims if other not in (dim, component_dim)) ) + self.out_dims = State(self.batch_dims + ("phase", dim)) + batch = int(np.prod([da.sizes[other] for other in self.batch_dims], dtype=int)) + ncomp = 1 if component_dim is None else da.sizes[component_dim] + self.circular_input = State(self._zeros(batch, ncomp, self.nperseg)) + self.model_input = State(self._zeros(batch, self.in_channels, self.nperseg)) self.circular_output = State( - torch.zeros( - self.batch_size, - self.classes, - self.nperseg, + torch.full( + (batch, self.classes, self.nperseg + self.step), + self.fill, dtype=torch.float32, device=self.device, ) ) self.circular_counts = State( torch.zeros( - self.batch_size, self.nperseg, dtype=torch.int32, device=self.device + batch, + self.classes, + self.nperseg + self.step, + dtype=torch.int32, + device=self.device, ) ) - if chunk_dim == self.dim: - self.buffer = State(da.isel({self.dim: slice(0, 0)})) - else: - self.buffer = State(None) + self.started = State(False) + self.buffer = State(da.isel({dim: slice(0, 0)}) if chunk_dim == dim else None) + self.inflight = State(deque()) + + def _zeros(self, *shape): + """Allocate a zeroed float32 tensor on the atom's device.""" + return torch.zeros(*shape, dtype=torch.float32, device=self.device) def call(self, da, **flags): - """Run the model over *da*, managing a carry-over buffer for chunked input.""" + """ + Run the model over *da*, managing a carry-over buffer for chunked input. + + A window is annotated as soon as it is complete, but its samples are + only emitted once the *following* window has been annotated: the + end-aligned last window of a record may reach back into them, and + :meth:`flush` is where that is settled. A chunk that completes no new + window therefore produces no output at all. On a CUDA device the + emission is one step later still: each output's device-to-host + transfer is issued asynchronously and only the completed ones are + returned (see `max_buffers`), the rest following on later calls or + at :meth:`flush`. + + Chunked along a dimension other than `dim`, that carry-over would be + a leak: the next chunk holds *other* lanes, so nothing of this one — + neither the tail buffer nor the circular window — applies to it. Such + a chunk is a whole record on its own and is run from a fresh state + and settled here. + """ + chunk_dim = flags.get("chunk_dim") + if chunk_dim is not None and chunk_dim != self.sample_dim: + return self._call_independent(da, **flags) + dim = self.sample_dim if self.buffer is None: - out = self._process(da) + if da.sizes[dim] < self.nperseg: + raise ValueError( + f"the record is shorter along {dim!r} " + f"({da.sizes[dim]} samples) than one model window " + f"({self.nperseg} samples)" + ) + self._process(da) else: - da = concat([self.buffer, da], self.dim) - out = self._process(da) - divpoint = out.sizes[self.dim] - self.buffer = State(da.isel({self.dim: slice(divpoint, None)})) - return out + da = concat([self.buffer, da], dim) + if da.sizes[dim] < self.nperseg: + self.buffer = State(da) + else: + self._process(da) + return self._harvest() or None + + def _call_independent(self, da, **flags): + """Annotate one whole record, settled on the spot, leaving no state behind.""" + dim = self.sample_dim + if da.sizes[dim] < self.nperseg: + raise ValueError( + f"the record is shorter along {dim!r} " + f"({da.sizes[dim]} samples) than one model window " + f"({self.nperseg} samples)" + ) + # Reallocate: the lanes of this chunk are not the lanes the buffers + # were sized and filled for, and the last chunk may hold fewer. + self.initialize(da, **flags) + if flags.get("chunk_dim") == self.component_dim: + # Not a lane axis: the model reads every component of a window at + # once, so a chunk holding a subset of them is not a record. + raise ValueError( + f"{self.component_dim!r} is the component dimension of the " + "model and cannot be chunked: the components of one window " + "are read together" + ) + self._process(da) + chunks = self._harvest() + self.flush() + # One record in, one record out: the pieces are consecutive along + # `dim`, and downstream is joining along the *chunked* dimension. + return concat(chunks, dim) if len(chunks) > 1 else chunks[0] + + def flush(self): + """ + Emit the end-aligned final window and everything it completes. + + SeisBench appends one last window ending on the record's last sample + whenever the stride leaves a remainder, so the output spans the input. + Firing once per run, this stays chunk-invariant. The in-flight + output queue is drained here, so nothing survives the end of a run. + """ + dim = self.sample_dim + if self.started is not True: + buffered = self.buffer + if isinstance(buffered, DataArray) and buffered.sizes[dim] > 0: + # The whole stream was shorter than one window: say so, as + # the eager call does, rather than answering with nothing. + raise ValueError( + f"the record is shorter along {dim!r} " + f"({buffered.sizes[dim]} samples) than one model window " + f"({self.nperseg} samples)" + ) + return self._harvest(block=True) + buffer = self.buffer + remainder = buffer.sizes[dim] - self.nperseg + if remainder > 0: + self._advance(buffer.isel({dim: slice(-remainder, None)}), remainder) + self._emit(buffer, 0, self.step - remainder, self.nperseg + remainder) + self.buffer = State(buffer.isel({dim: slice(0, 0)})) + self.started = State(False) + self.circular_output.fill_(self.fill) + self.circular_counts.fill_(0) + return self._harvest(block=True) def _process(self, da): - self._initialize(da) - chunks = [] - for idx in range(0, da.sizes[self.dim] - self.nperseg, self.step): - data = self._push_chunk(da, idx) - self._roll(self.circular_input, data) - self._roll(self.circular_counts, 0) - self._roll(self.circular_output, 0.0) - normalized = self.model.annotate_batch_pre(self.circular_input, {}) - if self.component_strategy == "clone": - self.model_input[:] = normalized.unsqueeze(1) - else: - ch = list(self.model.component_order).index(self.component_strategy) - self.model_input[:, ch, :] = normalized - self._run_model() - data = self._pull_completed() - chunk = self._attach_metadata(data, da, idx) - chunk = chunk.transpose(self.dim, ...) # TODO: does it make sense? - chunks.append(chunk) - return concat(chunks, self.dim) - - def _initialize(self, da): - chunk = da.isel({self.dim: slice(0, self.noverlap)}) - chunk = chunk.transpose(..., self.dim) - chunk = torch.tensor(chunk.values, dtype=torch.float32, device=self.device) - self.circular_input[:, self.step :] = chunk - - def _push_chunk(self, da, idx): - chunk = da.isel({self.dim: slice(idx + self.noverlap, idx + self.nperseg)}) - chunk = chunk.transpose(..., self.dim) - return torch.tensor(chunk.values, dtype=torch.float32, device=self.device) - - def _roll(self, buffer, values): - buffer[..., : self.noverlap] = buffer[..., self.step :] - buffer[..., self.noverlap :] = values - - def _run_model(self): - with torch.no_grad(): - out = self.model(self.model_input) - slc = slice(self.blinding[0], -self.blinding[1]) - self.circular_counts[:, slc] += 1 - self.circular_output[:, :, slc] += out[:, :, slc] - - def _pull_completed(self): - data = ( - self.circular_output[:, :, : self.step] - / self.circular_counts[:, None, : self.step] - ) - return data.cpu().numpy() - - def _attach_metadata(self, data, da, idx): - coords = da.coords.copy() - coords[self.dim] = coords[self.dim][idx : idx + self.step] + """Annotate every window of *da* that the buffered state can complete.""" + dim = self.sample_dim + nperseg, step = self.nperseg, self.step + if self.started: + first = step # the window at 0 was annotated by the previous call + else: + self._prime(da) + first = 0 + last = None + for idx in range(first, da.sizes[dim] - nperseg + 1, step): + tail = da.isel({dim: slice(idx + nperseg - step, idx + nperseg)}) + self._advance(tail, step) + if idx > 0: # the first window of a run completes nothing yet + self._emit(da, idx - step, 0, step) + last = idx + if last is None: + self.buffer = State(da) + return + self.started = State(True) + self.buffer = State(da.isel({dim: slice(last, None)})) + + def _prime(self, da): + """Stage the head of the first window so that the first slide completes it.""" + head = da.isel({self.sample_dim: slice(0, self.noverlap)}) + self.circular_input.narrow(-1, self.step, self.noverlap).copy_( + self._to_device(head) + ) + + def _advance(self, tail, shift): + """Slide the window by *shift* samples, run the model and stack its output.""" + self._slide(self.circular_input, shift, -1).copy_(self._to_device(tail)) + self._slide(self.circular_output, shift, -1).fill_(self.fill) + self._slide(self.circular_counts, shift, -1).fill_(0) + self._fill_model_input() + with torch.inference_mode(): + # `annotate_batch_post` writes into its input, which is an inference + # tensor: outside this block that raises. + batch = self.model.annotate_batch_pre(self.model_input, self.argdict) + piggyback = None + if isinstance(batch, tuple): + batch, piggyback = batch + out = self.model(batch) + out = self.model.annotate_batch_post( + out, piggyback=piggyback, argdict=self.argdict + ) + self._accumulate(out) + + def _slide(self, buffer, shift, axis): + """Slide *buffer* left by *shift* samples, returning the freed tail view.""" + size = buffer.shape[axis] + kept = buffer.narrow(axis, shift, size - shift).clone() + buffer.narrow(axis, 0, size - shift).copy_(kept) + return buffer.narrow(axis, size - shift, shift) + + def _fill_model_input(self): + """Permute the staged components into the model's input slots.""" + if self.slots is None: + self.model_input[:] = self.circular_input + else: + self.model_input[:, self.slots] = self.circular_input + + def _accumulate(self, out): + """ + Stack one window of ``(batch, samples, classes)`` predictions. + + Transposed on arrival into the ``(batch, classes, samples)`` the buffers + hold, which costs nothing — it is a view, and the reduction runs along + the axis the model's own output is contiguous in either way. + + The model blinds its own output by writing NaN into it, so the sum + ignores NaN and the count tracks what was finite: dividing the two + reproduces SeisBench's ``nanmean`` over covering windows exactly, and + ``stacking="max"`` its ``nanmax``. + + Raises + ------ + ValueError + If the post-processed batch is not ``(..., in_samples, classes)``. + Shipped SeisBench models do land here — ``CRED`` on both counts, + since it keeps the ``WaveformModel`` default + ``annotate_batch_post``, which leaves the batch as + ``(batch, classes, samples)`` rather than transposing it as + ``PhaseNet`` does, *and* predicts 19 samples for a 3000-sample + window. Without this check either mistake surfaces as a bare + ``RuntimeError`` from the broadcast. + """ + expected = (self.nperseg, self.classes) + if tuple(out.shape[-2:]) != expected: + raise ValueError( + f"{type(self.model).__name__}.annotate_batch_post returned a " + f"batch ending in {tuple(out.shape[-2:])}, not " + f"{expected} = (in_samples, classes): this atom adopts " + "SeisBench's stacking contract, `(batch, samples, classes)`. " + "A model leaving the batch as `(batch, classes, samples)` — " + "the `WaveformModel` default, which `PhaseNet` overrides — " + "must transpose it; a model whose window prediction is not " + "one sample per input sample cannot be stacked this way." + ) + out = torch.transpose(out, -1, -2) + window = self.circular_output.narrow(-1, self.step, self.nperseg) + counts = self.circular_counts.narrow(-1, self.step, self.nperseg) + finite = torch.isfinite(out) + values = torch.nan_to_num(out, nan=self.fill) + if self.stacking == "max": + torch.maximum(window, values, out=window) + else: + window += values + counts += finite.to(counts.dtype) + + def _pull(self, start, length): + """ + Reduce *length* stacked samples from *start* into a fresh tensor. + + Fresh so that the device-to-host transfer of the result can stay in + flight while the circular buffers slide on: the reduction allocates, + it never views the stack. + """ + values = self.circular_output.narrow(-1, start, length) + counts = self.circular_counts.narrow(-1, start, length) + if self.stacking == "max": + return torch.where(counts > 0, values, float("nan")) + return values / counts # samples no window covered come out NaN + + def _to_device(self, chunk): + """Stage *chunk* as a float32 tensor on the device, async on CUDA.""" + dims = self.batch_dims + if self.component_dim is not None: + dims += (self.component_dim,) + chunk = chunk.transpose(*dims, self.sample_dim) + values = np.ascontiguousarray(chunk.values, dtype=np.float32) + batch, ncomp = self.circular_input.shape[:2] + values = values.reshape(batch, ncomp, chunk.sizes[self.sample_dim]) + data = torch.from_numpy(values) + if self.device.type == "cuda": # pragma: no cover + # A pinned staging copy lets the transfer overlap with compute. + return data.pin_memory().to(self.device, non_blocking=True) + return data + + def _emit(self, da, offset, start, length): + """Queue the output chunk of *length* samples found at *start* in the stack.""" + dim = self.sample_dim + data = self._pull(start, length) + # Slice the whole chunk, not just the dimension coordinate: a label + # attached to the samples (a tag, a pick id) has to follow them, or + # it keeps the input length and is silently dropped on assembly. + coords = da.isel({dim: slice(offset, offset + length)}).coords.copy() + if self.component_dim is not None: + coords = coords.drop_dims(self.component_dim) coords["phase"] = self.phases - dims = tuple(dim for dim in da.dims if dim != self.dim) + ( - "phase", - self.dim, + shape = tuple(da.sizes[other] for other in self.batch_dims) + shape = (*shape, self.classes, length) + self._submit(data, (shape, coords, self.out_dims, da.name, da.attrs)) + + def _submit(self, data, meta): + """ + Queue one emitted output, its device-to-host transfer in flight. + + On CUDA the transfer is issued into a pinned staging buffer, + non-blocking, with an event marking its completion; the source + tensor rides along in the queue so it outlives the copy. At most + `max_buffers` transfers are left pending — the older ones are + synchronized — which is what bounds the staging memory. On the CPU + there is nothing to wait for and the item is complete on arrival. + """ + if self.device.type == "cuda": # pragma: no cover + values = torch.empty(data.shape, dtype=data.dtype, pin_memory=True) + values.copy_(data, non_blocking=True) + event = torch.cuda.Event() + event.record() + else: + values, event = data, None + self.inflight.append((event, values, data, meta)) + excess = len(self.inflight) - self.max_buffers + if excess > 0: + event = self.inflight[excess - 1][0] + if event is not None: # pragma: no cover + event.synchronize() + + def _harvest(self, block=False): + """ + Emit the queued outputs whose transfer has completed, in order. + + The queue is strictly ordered, so the harvest stops at the first + transfer still in flight; *block* waits every transfer out instead, + which is what :meth:`flush` does to drain the run. + """ + if self.inflight is ...: + return [] + chunks = [] + while self.inflight: + event = self.inflight[0][0] + if event is not None and not block and not event.query(): + break # pragma: no cover + chunks.append(self._realize(self.inflight.popleft())) + return chunks + + def _realize(self, item): + """Build the output chunk of one completed transfer.""" + event, values, _, (shape, coords, dims, name, attrs) = item + if event is not None: # pragma: no cover + event.synchronize() + return DataArray(values.numpy().reshape(shape), coords, dims, name, attrs) + + def _resolve_slots(self, slots): + """Turn ``component_strategy`` into the input slots the data fills.""" + order = list(self.model.component_order) + strategy = self.component_strategy + if slots is None: + if strategy in ("auto", "clone"): + return None # SeisBench's DAS default: clone into every slot + if strategy == "strict": + raise ValueError( + "component_strategy='strict' needs a component dimension, " + "and none was found: name it with `components=`" + ) + return [0 if strategy == "pad" else order.index(strategy)] + if strategy == "strict" and len(slots) < self.in_channels: + raise ValueError( + f"the data fills {len(slots)} of the model's {self.in_channels} " + "input slots and component_strategy is 'strict'" + ) + if strategy == "clone" or strategy in order: + if len(slots) > 1: + raise ValueError( + f"component_strategy={strategy!r} needs a single component " + f"per lane, but the data has {len(slots)}" + ) + return None if strategy == "clone" else [order.index(strategy)] + return slots + + +class _ChannelFilter(Filter): + """ + Filter only the channels whose labels match a glob pattern. + + The subset form of :class:`~xdas.atoms.Filter`, whose parameters it takes, + for the per-channel preprocessing filters some SeisBench weight sets ship: + ``obs`` declares a 0.5 Hz highpass on ``"??H"``, i.e. on its hydrophone + alone. The channels the pattern does not match pass through untouched, and a + pattern matching none of them makes the whole atom a no-op — which is what + SeisBench's ``stream.select(channel=...)`` does with a pattern no trace + answers. + + The channel dimension is found exactly as :class:`Annotate` finds it: by + its labels, each ending with a distinct letter of the model's + ``component_order``, never by its name. + + Parameters + ---------- + pattern : str + Channel glob, matched against the whole label with :mod:`fnmatch`. + freq : tuple of float or None + Corner frequencies ``(low, high)`` in Hz, as :class:`~xdas.atoms.Filter` + takes them: ``(0.5, None)`` is a highpass. + components : str, False or None, optional + Name of the channel dimension. ``None`` (default) detects it by its + labels, ``False`` disables detection, which makes the atom a no-op. + component_order : str, optional + The model's ``component_order``, against which the channel dimension is + recognised. Defaults to ``"ZNE"``. + flexible : bool, optional + Whether ``1``/``N`` and ``2``/``E`` are interchangeable while + recognising the channel dimension, as SeisBench's + ``flexible_horizontal_components`` makes them by default. + **kwargs + Passed to :class:`~xdas.atoms.Filter`: ``order``, ``zerophase`` and + ``dim``. The default order of 4 is also ``obspy``'s ``corners`` + default, which is what makes a model-declared filter reproducible + exactly. ``ftype="fir"`` is refused: the FIR form compensates its group + delay by shifting the coordinate, which would leave the channels this + atom does *not* filter on somebody else's samples. + + Examples + -------- + >>> import numpy as np + >>> import xdas as xd + >>> from xdas.atoms.ml import _ChannelFilter + + A four-channel OBS station, with the hydrophone last: + + >>> da = xd.testing.dummy(dims=("time", "channel"), shape=(1000, 4)) + >>> da["channel"] = np.array(["BHZ", "BH1", "BH2", "BDH"]) + >>> atom = _ChannelFilter("??H", (0.5, None), component_order="Z12H") + >>> result = atom(da) + + The three seismometer channels are untouched, the hydrophone is not: + + >>> np.allclose(result.isel(channel=slice(0, 3)).values, da.values[:, :3]) + True + >>> np.allclose(result.isel(channel=3).values, da.values[:, 3]) + False + """ + + def __init__( + self, + pattern, + freq, + components=None, + component_order="ZNE", + flexible=True, + **kwargs, + ): + super().__init__(freq, **kwargs) + if self.ftype == "fir": + raise ValueError( + "`ftype='fir'` cannot filter a subset of the channels: it " + "compensates its group delay by shifting the coordinate, and " + "the untouched channels would then sit on the wrong samples" + ) + self.pattern = pattern + self.components = components + self.component_order = component_order + self.flexible = flexible + self.mask = State(...) + + def initialize(self, da, **flags): + """Find the channel dimension and build the mask the pattern selects.""" + super().initialize(da, **flags) + channel_dim, _ = resolve_component_dim( + da, + resolve_sample_dim(da, self.dim), + self.component_order, + self.components, + self.flexible, + ) + labels = () if channel_dim is None else component_labels(da, channel_dim) + selected = [fnmatch(label, self.pattern) for label in labels] + shape = [-1 if dim == channel_dim else 1 for dim in da.dims] + self.mask = State(np.reshape(selected, shape) if any(selected) else None) + + def call(self, da, **flags): + """Filter every channel, then restore the ones the pattern left out.""" + if self.mask is None: + return da + # Filtering all of them and dropping what is not wanted keeps the + # inherited filter state one shape and costs a handful of channels. + filtered = super().call(da, **flags) + values = np.where(self.mask, filtered.values, da.values) + return DataArray(values, da.coords, da.dims, da.name, da.attrs) + + +def _model_filter(model, dim="time", components=None, **annotate_kwargs): + """ + Build the preprocessing filter stage a weight set declares, if any. + + SeisBench's ``annotate_stream_pre`` filters the waveforms *before* + resampling them, with whatever the weight set puts in ``filter_args`` and + ``filter_kwargs``; skipping it feeds the network something it was not + trained on. This is not a user-facing bandpass — compose + :class:`~xdas.atoms.Filter` yourself for that — and most weight sets + declare none, which is why the stage is optional. + + Two declarations are understood, both spelled as ``obspy``'s + ``Stream.filter`` arguments: + + - *flat*, e.g. ``filter_args=("highpass",)`` with + ``filter_kwargs={"freq": 1}``, applied to everything; + - *per channel*, a dict from channel glob to arguments with + ``filter_kwargs`` keyed identically, e.g. ``{"??H": ("highpass",)}``, + which becomes one :class:`_ChannelFilter` per pattern, applied in + declaration order. SeisBench's own DAS wrapper refuses this form; here it + works, through the label matching :class:`Annotate` already does. + + The translation is exact rather than approximate: ``obspy`` filters with + ``corners=4``, ``zerophase=False`` and a Butterworth in second-order + sections, which is what :class:`~xdas.atoms.Filter` defaults to. Two + concessions are copied from SeisBench's ``_get_filter_args``: a zero-phase + declaration warns and doubles the order instead, since exact zero-phase IIR + filtering has no causal streaming form, and no corner is allowed above half + the Nyquist of the model's own sampling rate. + + Parameters + ---------- + model : seisbench.models.WaveformModel + The model whose weight set is read. + dim : str, optional + Dimension the filter runs along. Defaults to ``"time"``. + components : str, False or None, optional + Name of the channel dimension, for the per-channel form. ``None`` + (default) detects it by its labels. + **annotate_kwargs + SeisBench annotate arguments overriding the weight set's, of which only + ``flexible_horizontal_components`` is read here. + + Returns + ------- + Atom or None + The stage to run before resampling, or ``None`` when the weight set + declares no filter — in which case the pipeline is one stage shorter. + + Examples + -------- + >>> from xdas.atoms.ml import _model_filter + + A weight set declaring nothing gets no stage: + + >>> class Plain: + ... component_order, sampling_rate = "ZNE", 100 + ... default_args, filter_args, filter_kwargs = {}, None, None + >>> _model_filter(Plain()) is None + True + + The per-channel form of the ``obs`` weight set: + + >>> class OBS(Plain): + ... component_order = "Z12H" + ... filter_args = {"??H": ("highpass",)} + ... filter_kwargs = {"??H": {"freq": 0.5}} + >>> stage = _model_filter(OBS()) + >>> stage.pattern, stage.freq, stage.order + ('??H', (0.5, None), 4) + """ + filter_args = getattr(model, "filter_args", None) + if filter_args is None: + return None + filter_kwargs = getattr(model, "filter_kwargs", None) + sampling_rate = getattr(model, "sampling_rate", None) + if not isinstance(filter_args, dict): + freq, order = _translate_filter(filter_args, filter_kwargs, sampling_rate) + return Filter(freq, order=order, dim=dim) + argdict = dict(model.default_args) | annotate_kwargs + flexible = annotate_arg(model, argdict, "flexible_horizontal_components") + stages = [] + for pattern, args in filter_args.items(): + if not isinstance(filter_kwargs, dict) or pattern not in filter_kwargs: + raise ValueError( + f"the weight set declares a filter for the channels matching " + f"{pattern!r} in `filter_args` but not in `filter_kwargs`" + ) + freq, order = _translate_filter(args, filter_kwargs[pattern], sampling_rate) + stages.append( + _ChannelFilter( + pattern, + freq, + order=order, + dim=dim, + components=components, + component_order=model.component_order, + flexible=flexible, + ) + ) + if not stages: + return None + return stages[0] if len(stages) == 1 else Sequential(stages) + + +def _translate_filter(args, kwargs, sampling_rate=None): + """ + Translate one ``obspy`` filter declaration into `Filter` parameters. + + Returns the ``(low, high)`` corner pair and the filter order, applying + SeisBench's two concessions: zero-phase becomes a doubled order, and no + corner sits above half the Nyquist of *sampling_rate*. + """ + args = tuple(args) + if len(args) != 1: + raise ValueError( + "a weight set's filter declaration must name exactly one obspy " + f"filter type, got {args!r}" + ) + name = args[0] + kwargs = dict(kwargs or {}) + order = kwargs.get("corners", OBSPY_CORNERS) + if kwargs.get("zerophase", False): + warnings.warn( + f"the weight set declares a zero-phase {name} filter, which has no " + f"causal streaming form: filtering forward only with the order " + f"doubled ({order} -> {2 * order}), as SeisBench does", + UserWarning, + stacklevel=3, + ) + order *= 2 + # As SeisBench does: no corner frequency may sit above half the Nyquist, + # so that the filter stays valid on data at a legal but lower rate. + top = np.inf if not sampling_rate else 0.999999 * 0.25 * sampling_rate + match name: + case "highpass": + return (min(kwargs["freq"], top), None), order + case "lowpass": + return (None, min(kwargs["freq"], top)), order + case "bandpass": + return (kwargs["freqmin"], min(kwargs["freqmax"], top)), order + case _: + raise ValueError( + f"the weight set declares a {name!r} filter, which has no " + "`Filter` equivalent: pass the stage yourself, or drop it" + ) + + +class Picker(Sequential): + """ + Pick phases with a SeisBench model: waveforms in, one pick table out. + + The whole pipeline SeisBench's ``model.classify(stream)`` runs, assembled + from the weight set and nothing else:: + + model-declared filter -> Resample -> Annotate -> Trigger + + Every stage is configured from the *weight set*, so two pickers built on + one model class can differ in stage count, sampling rate and thresholds:: + + Picker(PhaseNet("original")) Picker(PhaseNet("obs")) + Resample(100.0) _ChannelFilter('??H', 0.5 Hz) + Annotate Resample(100.0) + Trigger({'P': 0.3, 'S': 0.3}) Annotate + Trigger({'P': 0.2, 'S': 0.1}) + + Being a :class:`Sequential` rather than a factory function, a picker keeps + everything a pipeline can do: ``>>`` composes it, ``repr`` shows its + stages, it pickles, and ``picker.process(source, out=...)`` streams it. It + also inherits the two collection hooks from the stages that define them — + :meth:`Annotate.gather` collapses a ``channel`` level into the component + dimension before the first stage runs, and :meth:`Trigger.merge` folds the + per-leaf tables — so picking a whole network is one call answering with + one table. + + Parameters + ---------- + model : seisbench.models.WaveformModel + The model to pick with, weights loaded. + thresh : float, mapping or None, optional + Trigger-on thresholds, as :class:`~xdas.atoms.Trigger` takes them. + ``None`` (default) reads them off the weight set: one per non-noise + label, so the noise class is never picked. + resample : bool, optional + Whether to resample to the model's own ``sampling_rate`` — which is + not always 100 Hz, PhaseNet's ``diting`` running at 50. ``True`` by + default; the stage is a no-op on data already at that rate, so it + costs nothing to leave in. ``False`` drops it, which is what to pass + when the data is already there or when the polyphase resampling is + not wanted. It is the one stage that does not match SeisBench, and by + a difference of passband rather than of accuracy: SeisBench resamples + with obspy's ``Trace.resample``, whose default ``window="hann"`` is + applied in the frequency domain and halves the amplitude at half the + input Nyquist, where the polyphase filter used here is flat. + filter : bool, optional + Whether to apply the preprocessing filter the weight set ships, if it + ships one. ``True`` by default. Of the 17 cached PhaseNet weight sets + only ``obs`` declares one, which is why most pipelines are three + stages. + dim : str, optional + The sample dimension. Defaults to ``"time"``. Unlike + :class:`Annotate`, the ``"first"``/``"last"`` aliases are refused: the + pick table names its columns after coordinates, so the dimension has + to be nameable before any data is seen. + components : str, False or None, optional + Name of the component dimension, and of the collection level to + gather. ``None`` (default) detects both, ``False`` disables both. See + :class:`Annotate`. + component_strategy : str, optional + How the model's input slots are filled, see :class:`Annotate`. + device : str or torch.device, optional + Torch device. Defaults to CUDA if available, else CPU. + tolerance : scalar, None or False, optional + Grid-snapping budget forwarded to :func:`xdas.stack` when a component + level is gathered, see :meth:`Annotate.gather`. + coords : sequence of str, "auto" or None, optional + The coordinates annotating the picks, as + :class:`~xdas.atoms.Trigger` takes them. Defaults to ``"auto"``: the + scalar coordinates lead, so a pick carries the identity its array + carries. + **annotate_kwargs + SeisBench annotate arguments (``overlap``, ``stacking``, ``blinding``, + ``P_threshold``, ...) overriding what the weight set declares. + + Warnings + -------- + A weight set whose filter is declared *per channel* — ``obs``'s 0.5 Hz + highpass on ``??H`` — can only select the channels it names once they are + labelled. On unlabelled data (a DAS section, a bare trace) the glob + matches nothing and the stage is a silent no-op, exactly as + ``stream.select`` is, while ``component_strategy="clone"`` still clones + the signal into the hydrophone slot. Label the channels, or do not run OBS + weights on unlabelled data. + + Examples + -------- + >>> import numpy as np + >>> import torch + >>> import xdas as xd + >>> from xdas.atoms import Picker + + A stand-in for a real weight set, small enough to inline: an eight-sample + window, three components, three classes and one declared threshold. + + >>> class Model(torch.nn.Module): + ... in_samples, in_channels, classes = 8, 3, 3 + ... labels, component_order = "PSN", "ZNE" + ... sampling_rate = 100.0 + ... default_args = {"overlap": 0, "P_threshold": 0.5} + ... def annotate_batch_pre(self, batch, argdict): + ... return batch + ... def annotate_batch_post(self, batch, piggyback, argdict): + ... return torch.transpose(batch, -1, -2) + ... def forward(self, batch): + ... return batch + + The stages come from the weight set. This one declares no filter, so the + pipeline is three stages, and only the phases it can pick get a threshold + — the noise class never does: + + >>> picker = Picker(Model(), device="cpu") + >>> [type(stage).__name__ for stage in picker] + ['Resample', 'Annotate', 'Trigger'] + >>> picker[-1].thresh + {'P': 0.5, 'S': 0.3} + + Picking a three-component record. The model here returns its input, so + the vertical channel is the ``P`` class: + + >>> values = np.zeros((64, 3)) + >>> values[20, 0] = 0.9 + >>> da = xd.DataArray( + ... values, + ... { + ... "time": { + ... "tie_indices": [0, 63], + ... "tie_values": [0.0, 0.63], + ... "sampling_interval": 0.01, + ... }, + ... "channel": ["SHZ", "SHN", "SHE"], + ... }, + ... ("time", "channel"), + ... ) + >>> da = da.assign_coords(network="IA", station="DBNFM") + >>> xd.pick(da, Model(), device="cpu") + network station phase time value + 0 IA DBNFM P 0.2 0.9 + + """ + + def __init__( + self, + model, + thresh=None, + resample=True, + filter=True, + dim="time", + components=None, + component_strategy="auto", + device=None, + tolerance=None, + coords="auto", + **annotate_kwargs, + ): + if dim in ("first", "last"): + raise ValueError( + f"a picker needs its sample dimension by name, not as {dim!r}: " + "the picks are annotated with coordinates, which are named" + ) + stages = [] + if filter: + stage = _model_filter( + model, dim=dim, components=components, **annotate_kwargs + ) + if stage is not None: + stages.append(stage) + if resample: + rate = getattr(model, "sampling_rate", None) + if not rate: + raise ValueError( + "the weight set declares no `sampling_rate`, so there is " + "nothing to resample to: pass `resample=False` to pick at " + "the data's own rate" + ) + stages.append(Resample(rate, dim=dim)) + stages.append( + Annotate( + model, + dim=dim, + components=components, + component_strategy=component_strategy, + device=device, + tolerance=tolerance, + **annotate_kwargs, + ) ) - return DataArray(data, coords, dims, da.name, da.attrs) + if thresh is None: + thresh = _model_thresholds(model, **annotate_kwargs) + stages.append(Trigger(thresh, dim=dim, coords=coords)) + super().__init__(stages, name="picker") + + +class MLPicker(Annotate): + """ + Deprecated alias of :class:`Annotate`, removed in 0.4. + + The old signature is kept, positional arguments included, but the results + move: the output is laid out sample-last, ``(..., "phase", dim)``, rather + than leading with the sample dimension; the component dimension is found + by its labels rather than assumed; the window overlap is read off the + weight set rather than fixed at half a window; and the end-aligned final + window SeisBench appends is emitted at :meth:`flush`. + """ + + def __init__( + self, model, dim="time", device=None, component_strategy="clone", **kwargs + ): + warnings.warn( + "MLPicker is deprecated and will be removed in 0.4, use Annotate instead", + DeprecationWarning, + stacklevel=3, + ) + super().__init__( + model, + dim=dim, + device=device, + component_strategy=component_strategy, + **kwargs, + ) + + +annotate = atomized(Annotate) +pick = atomized(Picker) +mlpicker = atomized(MLPicker) diff --git a/xdas/atoms/signal.py b/xdas/atoms/signal.py index 36cdc4d2..a5220422 100644 --- a/xdas/atoms/signal.py +++ b/xdas/atoms/signal.py @@ -1,19 +1,18 @@ """ -Signal-processing atoms: stateful wrappers around filtering and resampling. +Composite signal-processing atoms with physical (Hz) parameters. -Includes :class:`ResamplePoly`, :class:`IIRFilter`, :class:`FIRFilter`, -:class:`LFilter`, :class:`SOSFilter`, :class:`DownSample`, :class:`UpSample`. +Includes :class:`ResamplePoly`, :class:`IIRFilter`, :class:`FIRFilter`. The +stateful machine-parameter primitives they orchestrate live in +:mod:`xdas.atoms.kernel`. """ from fractions import Fraction -import numpy as np import scipy.signal as sp -from ..coordinates import Coordinate, get_sampling_interval -from ..core import DataArray, concat, split -from ..parallel import parallelize +from ..coordinates import get_sampling_interval from .core import Atom, State +from .kernel import LFilter, Polyphase, SOSFilter class ResamplePoly(Atom): @@ -91,9 +90,7 @@ def __init__(self, target, maxfactor=100, window=("kaiser", 5.0), dim="last"): self.maxfactor = maxfactor self.window = window self.dim = dim - self.upsampling = UpSample(..., dim=self.dim) self.firfilter = FIRFilter(..., ..., "lowpass", self.window, dim=self.dim) - self.downsampling = DownSample(..., self.dim) self.fs = State(...) def initialize(self, da, **flags): @@ -111,19 +108,16 @@ def initialize_from_state(self): cutoff = min(self.target / 2, self.fs / 2) max_rate = max(up, down) numtaps = 20 * max_rate + 1 - self.upsampling.factor = up self.firfilter.numtaps = numtaps self.firfilter.cutoff = cutoff - self.downsampling.factor = down + self.firfilter.up = up + self.firfilter.down = down def call(self, da, **flags): - """Apply polyphase resampling (upsample → FIR filter → downsample) to *da*.""" - if self.upsampling.factor == 1 and self.downsampling.factor == 1: + """Apply polyphase resampling to *da*.""" + if self.firfilter.up == 1 and self.firfilter.down == 1: return da - da = self.upsampling(da, **flags) - da = self.firfilter(da, **flags) - da = self.downsampling(da, **flags) - return da + return self.firfilter(da, **flags) class IIRFilter(Atom): @@ -286,6 +280,11 @@ class FIRFilter(Atom): Default: ``None`` scale : bool Default: ``True`` + up, down : int + Machine parameters of the polyphase form: the taps are designed at the + upsampled rate ``up * fs`` and applied by :class:`Polyphase`, which + keeps one output sample in ``down``. Both default to 1, i.e. plain + filtering. dim : str or int The dimension along which the downsampling is applied. This is either an index, ``time`` or ``distance``, or ``last``. @@ -350,6 +349,8 @@ def __init__( window="hamming", width=None, scale=True, + up=1, + down=1, dim="last", ): super().__init__() @@ -359,8 +360,10 @@ def __init__( self.window = window self.width = width self.scale = scale + self.up = up + self.down = down self.dim = dim - self.lfilter = LFilter(..., [1.0], self.dim) + self.polyphase = Polyphase(..., self.up, self.down, self.dim) self.fs = State(...) def initialize(self, da, **flags): @@ -369,7 +372,7 @@ def initialize(self, da, **flags): self.initialize_from_state() def initialize_from_state(self): - """Recompute the FIR taps and lag from the current design parameters.""" + """Recompute the FIR taps from the current design parameters.""" taps = sp.firwin( self.numtaps, self.cutoff, @@ -377,214 +380,14 @@ def initialize_from_state(self): window=self.window, pass_zero=self.btype, scale=self.scale, - fs=self.fs, + fs=self.fs * self.up, ) - self.lag = (len(taps) - 1) // 2 - self.lfilter.b = taps + # Interpolation spreads the energy of one input sample over `up` + # upsampled ones, which the taps must compensate. + self.polyphase.taps = self.up * taps + self.polyphase.up = self.up + self.polyphase.down = self.down def call(self, da, **flags): - """Apply the FIR taps to *da* and correct the time coordinate for filter lag.""" - da = self.lfilter(da, **flags) - da[self.dim] -= get_sampling_interval(da, self.dim, cast=False) * self.lag - return da - - -class LFilter(Atom): - """ - Stateful direct-form IIR/FIR filter using :func:`scipy.signal.lfilter`. - - Parameters - ---------- - b : array-like - Numerator polynomial coefficients. - a : array-like - Denominator polynomial coefficients. - dim : str or int, optional - Dimension to filter along. Defaults to ``"last"``. - parallel : int, bool, or None, optional - Worker count for parallelisation. - """ - - def __init__(self, b, a, dim="last", parallel=None): - super().__init__() - self.b = b - self.a = a - self.dim = dim - self.parallel = parallel - self.axis = State(...) - self.zi = State(...) - - def initialize(self, da, chunk_dim=None, **flags): - """Set the filter axis and allocate the initial conditions buffer.""" - self.axis = State(da.get_axis_num(self.dim)) - if self.dim == chunk_dim: - n_sections = max(len(self.a), len(self.b)) - 1 - shape = tuple( - n_sections if name == self.dim else size - for name, size in da.sizes.items() - ) - self.zi = State(np.zeros(shape)) - else: - self.zi = State(None) - - def call(self, da, **flags): - """Apply the filter to *da*, updating the state if chunked.""" - across = int(self.axis == 0) - if self.zi is None: - func = parallelize((None, None, across), across, self.parallel)(sp.lfilter) - data = func(self.b, self.a, da.values, self.axis) - else: - func = parallelize( - (None, None, across, None, across), (across, across), self.parallel - )(sp.lfilter) - data, zf = func(self.b, self.a, da.values, self.axis, self.zi) - self.zi = State(zf) - return da.copy(data=data) - - -class SOSFilter(Atom): - """ - Stateful second-order-sections IIR filter using :func:`scipy.signal.sosfilt`. - - Parameters - ---------- - sos : array-like, shape (n_sections, 6) - SOS filter coefficients as returned by e.g. :func:`scipy.signal.iirfilter`. - dim : str or int, optional - Dimension to filter along. Defaults to ``"last"``. - parallel : int, bool, or None, optional - Worker count for parallelisation. - """ - - def __init__(self, sos, dim="last", parallel=None): - super().__init__() - self.sos = sos - self.dim = dim - self.parallel = parallel - self.axis = State(...) - self.zi = State(...) - - def initialize(self, da, chunk_dim=None, **flags): - """Set the filter axis and allocate the SOS initial-conditions buffer.""" - self.axis = State(da.get_axis_num(self.dim)) - if self.dim == chunk_dim: - n_sections = self.sos.shape[0] - shape = (n_sections,) + tuple( - 2 if index == self.axis else element - for index, element in enumerate(da.shape) - ) - self.zi = State(np.zeros(shape)) - else: - self.zi = State(None) - - def call(self, da, **flags): - """Apply the SOS filter to *da*, updating the state if chunked.""" - across = int(self.axis == 0) - if self.zi is None: - func = parallelize((None, across), across, self.parallel)(sp.sosfilt) - data = func(self.sos, da.values, self.axis) - else: - func = parallelize( - (None, across, None, across + 1), (across, across + 1), self.parallel - )(sp.sosfilt) - data, zf = func(self.sos, da.values, self.axis, self.zi) - self.zi = State(zf) - return da.copy(data=data) - - -class DownSample(Atom): - """ - Stateful integer downsampling by selecting every *factor*-th sample. - - Parameters - ---------- - factor : int - Downsampling factor. - dim : str or int, optional - Dimension to downsample along. Defaults to ``"last"``. - """ - - def __init__(self, factor, dim="last"): - super().__init__() - self.factor = factor - self.dim = dim - self.buffer = State(...) - - def initialize(self, da, chunk_dim=None, **flags): - """Initialise the carry-over buffer for chunked operation.""" - if chunk_dim == self.dim: - self.buffer = State(da.isel({self.dim: slice(0, 0)})) - else: - self.buffer = State(None) - - def call(self, da, **flags): - """Downsample *da*, buffering the trailing partial stride when chunked.""" - if self.factor == 1: - return da - if self.buffer is not None: - da = concat([self.buffer, da], self.dim) - divpoint = da.sizes[self.dim] - da.sizes[self.dim] % self.factor - da, buffer = split(da, [divpoint], self.dim) - self.buffer = State(buffer) - return da.isel({self.dim: slice(None, None, self.factor)}) - - -class UpSample(Atom): - """ - Integer upsampling by zero-insertion (and optional energy scaling). - - Parameters - ---------- - factor : int - Upsampling factor. - scale : bool, optional - If ``True``, scale inserted samples so energy is preserved. - dim : str or int, optional - Dimension to upsample along. Defaults to ``"last"``. - """ - - def __init__(self, factor, scale=True, dim="last"): - super().__init__() - self.factor = factor - self.scale = scale - self.dim = dim - - def call(self, da, **flags): - """Upsample *da* by inserting zeros between every original sample.""" - if self.factor == 1: - return da - shape = tuple( - self.factor * size if dim == self.dim else size - for dim, size in da.sizes.items() - ) - slc = tuple( - slice(None, None, self.factor) if dim == self.dim else slice(None) - for dim in da.dims - ) - data = np.zeros(shape, dtype=da.dtype) - if self.scale: - data[slc] = da.values * self.factor - else: - data[slc] = da.values - coords = da.coords.copy() - delta = get_sampling_interval(da, self.dim, cast=False) - new_delta = delta / self.factor - coord = coords[self.dim] - tie_indices = coord.tie_indices * self.factor - tie_values = coord.tie_values - tie_indices[-1] += self.factor - 1 - tie_values[-1] += (self.factor - 1) * new_delta - data_coord = {"tie_indices": tie_indices, "tie_values": tie_values} - if coord.isregular(): - # The derived rate may not be exactly representable (integer datetime - # resolutions truncate), so declare the representation error as jitter - # on top of the inherited one; chunk seams then stay within tolerance. - data_coord["sampling_interval"] = new_delta - data_coord["tolerance"] = coord.tolerance + np.abs( - delta - new_delta * self.factor - ) - # An irregular input gives no rate to inherit and no jitter bound to - # derive one from, so the result stays irregular rather than claiming a - # precision the source never declared. - coords[self.dim] = Coordinate(data_coord, self.dim) - return DataArray(data, coords, da.dims, da.name, da.attrs) + """Apply the FIR taps to *da*, delay-corrected and resampled.""" + return self.polyphase(da, **flags) diff --git a/xdas/atoms/tasks.py b/xdas/atoms/tasks.py new file mode 100644 index 00000000..62e776cd --- /dev/null +++ b/xdas/atoms/tasks.py @@ -0,0 +1,705 @@ +""" +Task atoms: the public processing vocabulary with physical parameters only. + +Every public parameter keeps its meaning when the sampling rate changes: +frequencies are in Hz, window lengths in seconds (or meters along distance). +Machine parameters (coefficients, factors, taps) live in the kernel layer +(:mod:`xdas.atoms.kernel`) and are designed here from the data at the first +call. + +Each task atom has a function form generated with +:func:`~xdas.atoms.core.as_function`: ``decimate(da, 50.0)`` applies eagerly, +``decimate(..., 50.0)`` returns the atom, and passing an atom extends a +pipeline. Stateless operations are plain ``@atomized`` functions, which behave +identically; the split between functions and classes is invisible to users. +""" + +import numpy as np +import scipy.fft +from scipy.signal import ShortTimeFFT, get_window + +from ..coordinates import get_sampling_interval +from ..core import DataArray, concat +from ..parallel import parallelize +from .core import Atom, State, _whole_record, atomized +from .signal import FIRFilter, IIRFilter, ResamplePoly + +__all__ = [ + "STFT", + "Decimate", + "Differentiate", + "Filter", + "Integrate", + "Resample", + "decimate", + "detrend", + "differentiate", + "filter", + "hilbert", + "integrate", + "medfilt", + "resample", + "sliding_mean_removal", + "stft", + "taper", +] + + +class Filter(Atom): + """ + Bandpass, lowpass or highpass filter with corner frequencies in Hz. + + The band is given as a pair of corner frequencies ``(low, high)`` in Hz, + with ``None`` opening one end: ``(1.0, 10.0)`` is a bandpass, ``(1.0, + None)`` a highpass and ``(None, 10.0)`` a lowpass. + + Parameters + ---------- + freq : tuple of float or None + The pair of corner frequencies (low, high) in Hz. Use None to open one + end of the band. + ftype : {"iir", "fir"} + The filter implementation. "iir" designs a Butterworth filter applied + in second-order sections; it is causal and streams chunk by chunk. + "fir" designs a windowed-sinc linear-phase filter whose group delay is + compensated on the coordinate, making it effectively zero-phase while + remaining streamable. + order : int + The order of the IIR filter. Ignored for FIR filters, whose length is + set by `transition`. Default is 4. + transition : float, optional + The FIR transition bandwidth in Hz. Default is 10% of the lowest given + corner frequency. Ignored for IIR filters. + zerophase : bool + If True with an IIR filter, the filter is applied forwards and + backwards, doubling the effective order and cancelling the phase + shift. Exact zero-phase IIR filtering has no causal streaming form, so + such an atom refuses chunked execution along its dimension; use + ``ftype="fir"`` for a streamable (effectively) zero-phase filter. FIR + filters ignore this parameter as they are always compensated. + dim : str + The dimension along which to filter. Default is "time". + + Examples + -------- + >>> import xdas as xd + >>> from xdas.synthetics import wavelet_wavefronts + >>> da = wavelet_wavefronts() + >>> filtered = xd.filter(da, (1.0, 10.0)) + >>> atom = xd.filter(..., (None, 10.0), ftype="fir") + >>> atom + Filter(freq=(None, 10.0), ftype=fir, order=4, zerophase=False, dim=time, btype=lowpass, cutoff=10.0) + FIRFilter(numtaps=Ellipsis, cutoff=10.0, btype=lowpass, window=hamming, scale=True, up=1, down=1, dim=time) + Polyphase(taps=Ellipsis, up=1, down=1, dim=time) + + """ + + def __init__( + self, freq, ftype="iir", order=4, transition=None, zerophase=False, dim="time" + ): + super().__init__() + try: + low, high = freq + except (TypeError, ValueError): + raise TypeError( + "`freq` must be a pair of corner frequencies, using None to " + "open one end, e.g. (1.0, None) for a highpass" + ) from None + if low is None and high is None: + raise ValueError("at least one corner frequency must be given") + if ftype not in ("iir", "fir"): + raise ValueError("`ftype` must be either 'iir' or 'fir'") + self.freq = (low, high) + self.ftype = ftype + self.order = order + self.transition = transition + self.zerophase = zerophase + self.dim = dim + if low is None: + self.btype = "lowpass" + self.cutoff = high + elif high is None: + self.btype = "highpass" + self.cutoff = low + else: + self.btype = "bandpass" + self.cutoff = (low, high) + if ftype == "fir": + self.filter = FIRFilter(..., self.cutoff, self.btype, dim=self.dim) + self.fs = State(...) + elif not zerophase: + self.filter = IIRFilter(self.order, self.cutoff, self.btype, dim=self.dim) + + def _check_chunk_dim(self, x, chunk_dim): + """Zero-phase IIR has no causal streaming form: whole-record only.""" + if self.ftype == "iir" and self.zerophase: + self._refuse_chunked_along(self.dim, chunk_dim, x) + + def initialize(self, da, **flags): + """Measure the sampling rate to size the FIR filter from `transition`.""" + if self.ftype == "fir": + self.fs = State(1.0 / get_sampling_interval(da, self.dim)) + self.initialize_from_state() + + def initialize_from_state(self): + """Derive the FIR length from the transition bandwidth.""" + if self.ftype == "fir": + if self.transition is None: + transition = 0.1 * min(f for f in self.freq if f is not None) + else: + transition = self.transition + numtaps = int(np.ceil(3.3 * self.fs / transition)) + self.filter.numtaps = numtaps // 2 * 2 + 1 + + def call(self, da, **flags): + """Apply the filter, delegating to the designed child atom.""" + if self.ftype == "iir" and self.zerophase: + from ..signal import filter + + return filter( + da, + self.cutoff, + self.btype, + corners=self.order, + zerophase=True, + dim=self.dim, + ) + return self.filter(da, **flags) + + +class Decimate(Atom): + """ + Decimate to a target sampling rate by an integer factor. + + Composite atom: a lowpass anti-alias FIR filter (group delay compensated + on the coordinate) followed by integer downsampling. The current sampling + rate must be an integer multiple of `target`; for rational ratios use + :class:`Resample`. + + Parameters + ---------- + target : float + The target sampling rate in Hz (or in 1/m along distance). + window : str or tuple + The window used to design the anti-alias filter, compatible with + ``scipy.signal.get_window``. Default is ``("kaiser", 5.0)``. + dim : str + The dimension along which to decimate. Default is "time". + + Examples + -------- + >>> import xdas as xd + >>> from xdas.synthetics import wavelet_wavefronts + >>> da = wavelet_wavefronts() # 50 Hz + >>> xd.decimate(da, 25.0).sizes["time"] + 150 + + """ + + def __init__(self, target, window=("kaiser", 5.0), dim="time"): + super().__init__() + self.target = target + self.window = window + self.dim = dim + self.antialias = FIRFilter(..., ..., "lowpass", self.window, dim=self.dim) + self.fs = State(...) + + def initialize(self, da, **flags): + """Measure the sampling rate and design the anti-alias filter.""" + self.fs = State(1.0 / get_sampling_interval(da, self.dim)) + self.initialize_from_state() + + def initialize_from_state(self): + """Derive the integer factor and anti-alias design from the rate.""" + factor = round(self.fs / self.target) + if factor < 1 or abs(self.fs / self.target - factor) > 1e-6 * factor: + raise ValueError( + f"the sampling rate ({self.fs:g}) is not an integer multiple " + f"of the target ({self.target:g}); use Resample for rational " + "ratios" + ) + self.antialias.numtaps = 20 * factor + 1 + self.antialias.cutoff = self.target / 2 + self.antialias.down = factor + + def call(self, da, **flags): + """Anti-alias filter and downsample in one polyphase pass.""" + if self.antialias.down == 1: + return da + return self.antialias(da, **flags) + + +class Resample(ResamplePoly): + """ + Resample to any target sampling rate by polyphase filtering. + + Task-layer name for :class:`~xdas.atoms.ResamplePoly`: the data is + upsampled, lowpass FIR filtered and downsampled so that the ratio of the + factors matches `target` over the current sampling rate. + + Parameters + ---------- + target : float + The target sampling rate in Hz (or in 1/m along distance). + maxfactor : int + Limit on the intermediate upsampling factor, to avoid accidental + memory overflow. Default is 100. + window : str or tuple + The window used to design the FIR filter, compatible with + ``scipy.signal.get_window``. Default is ``("kaiser", 5.0)``. + dim : str + The dimension along which to resample. Default is "time". + + Examples + -------- + >>> import xdas as xd + >>> from xdas.synthetics import wavelet_wavefronts + >>> da = wavelet_wavefronts() # 50 Hz + >>> xd.resample(da, 20.0).sizes["time"] + 120 + + """ + + def __init__(self, target, maxfactor=100, window=("kaiser", 5.0), dim="time"): + super().__init__(target, maxfactor=maxfactor, window=window, dim=dim) + + +class Integrate(Atom): + """ + Integrate cumulatively along a dimension. + + Stateful: when processing chunk by chunk, the cumulative sum continues + across chunks. + + Parameters + ---------- + midpoints : bool + Whether to move the coordinates by half a step. Default is False. + dim : str + The dimension along which to integrate. Default is "time". + + """ + + def __init__(self, midpoints=False, dim="time"): + super().__init__() + self.midpoints = midpoints + self.dim = dim + self.carry = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Allocate the carried cumulative offset for chunked operation.""" + if chunk_dim == self.dim: + axis = da.get_axis_num(self.dim) + shape = tuple( + 1 if index == axis else size for index, size in enumerate(da.shape) + ) + self.carry = State(np.zeros(shape)) + else: + self.carry = State(None) + + def call(self, da, **flags): + """Integrate the chunk and offset it by the carried cumulative sum.""" + from ..signal import integrate + + out = integrate(da, midpoints=self.midpoints, dim=self.dim) + if self.carry is not None: + axis = out.get_axis_num(self.dim) + out = out.copy(data=out.values + self.carry) + index = tuple( + slice(-1, None) if a == axis else slice(None) for a in range(out.ndim) + ) + self.carry = State(out.values[index]) + return out + + +class Differentiate(Atom): + """ + Differentiate along a dimension. + + Stateful: when processing chunk by chunk, the last sample of each chunk is + carried over so the difference across the seam is not lost. The output has + one sample less than the input in total. + + Parameters + ---------- + midpoints : bool + Whether to move the coordinates by half a step. Default is False. + dim : str + The dimension along which to differentiate. Default is "time". + + """ + + def __init__(self, midpoints=False, dim="time"): + super().__init__() + self.midpoints = midpoints + self.dim = dim + self.buffer = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Initialise the one-sample carry-over buffer for chunked operation.""" + if chunk_dim == self.dim: + self.buffer = State(da.isel({self.dim: slice(0, 0)})) + else: + self.buffer = State(None) + + def call(self, da, **flags): + """Differentiate the chunk, prepending the buffered last sample.""" + from ..signal import differentiate + + if self.buffer is not None: + x = concat([self.buffer, da], self.dim) + self.buffer = State(da.isel({self.dim: slice(-1, None)})) + else: + x = da + return differentiate(x, midpoints=self.midpoints, dim=self.dim) + + +@atomized +@_whole_record() +def detrend(da, type="linear", dim="time", parallel=None): + """ + Remove a trend along the given dimension. + + Whole-record operation: the trend is fitted on the full record, so this + atom refuses chunked execution along its dimension. + + Parameters + ---------- + da : DataArray + The data to detrend. + type : str + Either "linear" or "constant". Default is "linear". + dim : str + The dimension along which to detrend. Default is "time". + parallel : bool or int, optional + Number of threads to use. + + Returns + ------- + DataArray + The detrended data. + + """ + from ..signal import detrend + + return detrend(da, type, dim=dim, parallel=parallel) + + +@atomized +@_whole_record() +def taper(da, window="hann", fftbins=False, dim="time", parallel=None): + """ + Apply a tapering window along the given dimension. + + Whole-record operation: the window spans the full record, so this atom + refuses chunked execution along its dimension. + + Parameters + ---------- + da : DataArray + The data to taper. + window : str or tuple, optional + The window to use, by default "hann". + fftbins : bool, optional + Whether to use a periodic windowing, by default False. + dim : str, optional + The dimension along which to taper. Default is "time". + parallel : bool or int, optional + Number of threads to use. + + Returns + ------- + DataArray + The tapered data. + + """ + from ..signal import taper + + return taper(da, window=window, fftbins=fftbins, dim=dim, parallel=parallel) + + +@atomized +@_whole_record() +def hilbert(da, dim="time", parallel=None): + """ + Compute the analytic signal, using the Hilbert transform. + + Whole-record operation: the transform is acausal, so this atom refuses + chunked execution along its dimension. + + Parameters + ---------- + da : DataArray + Signal data. Must be real. + dim : str, optional + The dimension along which to transform. Default is "time". + parallel : bool or int, optional + Number of threads to use. + + Returns + ------- + DataArray + Analytic signal of `da` along `dim`. + + """ + from ..signal import hilbert + + return hilbert(da, dim=dim, parallel=parallel) + + +@atomized +@_whole_record() +def sliding_mean_removal( + da, wlen, window="hann", pad_mode="reflect", dim="time", parallel=None +): + """ + Remove a sliding mean. + + The window length is physical: seconds along time, meters along distance. + Pending overlap-aware execution, this atom refuses chunked execution along + its dimension. + + Parameters + ---------- + da : DataArray + The data that the sliding mean should be removed from. + wlen : float + Length of the sliding mean, in the units of the `dim` coordinate. + window : str, optional + Tapering window used, by default "hann". + pad_mode : str, optional + Padding mode used, by default "reflect". + dim : str, optional + The dimension along which to remove the sliding mean. Default is + "time". + parallel : bool or int, optional + Number of threads to use. + + Returns + ------- + DataArray + The data with the sliding mean removed. + + """ + from ..signal import sliding_mean_removal + + return sliding_mean_removal( + da, wlen, window=window, pad_mode=pad_mode, dim=dim, parallel=parallel + ) + + +@atomized +@_whole_record(dim_arg="kernel") +def medfilt(da, kernel): + """ + Apply a median filter with kernel lengths in physical units. + + Parameters + ---------- + da : DataArray + The data to filter. + kernel : dict + Mapping of dimension name to kernel length in the units of that + dimension's coordinate (seconds along time, meters along distance). + Each length is converted to the nearest odd number of samples. + Dimensions not listed are not filtered. + + Returns + ------- + DataArray + The median filtered data. + + """ + from ..signal import medfilt + + kernel_dim = {} + for dim, length in kernel.items(): + size = max(1, round(length / get_sampling_interval(da, dim))) + kernel_dim[dim] = size if size % 2 else size + 1 + return medfilt(da, kernel_dim) + + +class STFT(Atom): + """ + Short-Time Fourier Transform with window length and hop in physical units. + + The window length is a target in the units of the `dim` coordinate + (seconds along time): the actual length is the next fast FFT size of the + corresponding number of samples, so transforms stay efficient whatever the + sampling rate. Frames start at the first sample and advance by `hop`; only + fully computable frames are ever emitted, so when processing chunk by + chunk the unconsumed tail is buffered across chunks (and dropped at gaps + and at the end of the stream), and chunked processing emits exactly the + frames of the eager transform. The output gains a "frequency" dimension + (one-sided for real data, centered two-sided for complex data) and the + `dim` coordinate moves to the frame centers. + + Parameters + ---------- + wlen : float + Target window length, in the units of the `dim` coordinate (seconds + along time). The actual length is ``scipy.fft.next_fast_len`` of the + equivalent number of samples. + hop : float, optional + Step between frame starts, in the same units. Must be positive and at + most `wlen`. Like the window length, it is snapped: to the nearest + whole number of samples (at least one, at most the actual window + length), so the frame grid can differ from the request. Default is + half the actual window length. + window : str or tuple + The tapering window, compatible with ``scipy.signal.get_window``. + Default is "hann". + scaling : {"spectrum", "psd"} + The scaling of the complex frames: "spectrum" preserves peak + amplitudes ("magnitude" scaling of `scipy.signal.ShortTimeFFT`), + "psd" makes the squared modulus a power spectral density. Default is + "spectrum". + nfft : int, optional + Expert mode: the FFT length in samples, to zero pad the windowed + frames. Must be at least the actual window length; a common choice is + twice that length, and fast FFT sizes matter. Default is the actual + window length (no padding). + dim : str + The dimension along which to transform. Default is "time". + parallel : bool or int, optional + Number of threads to use. + + Examples + -------- + >>> import xdas as xd + >>> from xdas.synthetics import wavelet_wavefronts + >>> da = wavelet_wavefronts() # 50 Hz + >>> xd.stft(da, 2.0, hop=1.0).sizes + {'time': 5, 'distance': 401, 'frequency': 51} + + """ + + def __init__( + self, + wlen, + hop=None, + window="hann", + scaling="spectrum", + nfft=None, + dim="time", + parallel=None, + ): + super().__init__() + if not wlen > 0: + raise ValueError("`wlen` must be positive") + if hop is not None and not 0 < hop <= wlen: + raise ValueError("`hop` must be positive and at most `wlen`") + if scaling not in ("spectrum", "psd"): + raise ValueError("`scaling` must be 'spectrum' or 'psd'") + self.wlen = wlen + self.hop = hop + self.window = window + self.scaling = scaling + self.nfft = nfft + self.dim = dim + self.parallel = parallel + self.sft = State(...) + self.buffer = State(...) + + def initialize(self, da, chunk_dim=None, **flags): + """Design the transform from the measured sampling rate.""" + fs = 1.0 / get_sampling_interval(da, self.dim) + nperseg = scipy.fft.next_fast_len(max(round(self.wlen * fs), 1)) + if chunk_dim != self.dim and da.sizes[self.dim] < nperseg: + raise ValueError( + f"the record is shorter along {self.dim!r} " + f"({da.sizes[self.dim]} samples) than the window " + f"({nperseg} samples)" + ) + if self.hop is None: + hop = max(nperseg // 2, 1) + else: + hop = min(max(round(self.hop * fs), 1), nperseg) + nfft = nperseg if self.nfft is None else self.nfft + if nfft < nperseg: + raise ValueError( + f"`nfft` ({nfft}) must be at least the window length in " + f"samples ({nperseg})" + ) + self.sft = State( + ShortTimeFFT( + get_window(self.window, nperseg), + hop=hop, + fs=fs, + fft_mode="onesided" if np.isrealobj(da.values) else "centered", + mfft=nfft, + scale_to="magnitude" if self.scaling == "spectrum" else "psd", + phase_shift=None, + ) + ) + if chunk_dim == self.dim: + self.buffer = State(da.isel({self.dim: slice(0, 0)})) + else: + self.buffer = State(None) + + def call(self, da, **flags): + """Emit every fully computable frame, buffering the unconsumed tail.""" + nperseg = self.sft.m_num + hop = self.sft.hop + if self.buffer is None: + return self._transform(da) + da = concat([self.buffer, da], self.dim) + n = da.sizes[self.dim] + if n < nperseg: + self.buffer = State(da) + return None + nframes = (n - nperseg) // hop + 1 + consumed = nframes * hop + out = da.isel({self.dim: slice(0, consumed - hop + nperseg)}) + self.buffer = State(da.isel({self.dim: slice(consumed, None)})) + return self._transform(out) + + def flush(self): + """Discard the buffered tail: only fully computable frames are emitted.""" + if isinstance(self.buffer, DataArray): + self.buffer = State(self.buffer.isel({self.dim: slice(0, 0)})) + return [] + + def _transform(self, da): + """Compute the windowed FFT of every full frame in *da*.""" + sft = self.sft + axis = da.get_axis_num(self.dim) + + def func(x): + frames = np.lib.stride_tricks.sliding_window_view(x, sft.m_num, axis=axis) + slc = [slice(None)] * frames.ndim + slc[axis] = slice(None, None, sft.hop) + frames = sft.win * frames[tuple(slc)] + if sft.onesided_fft: + return scipy.fft.rfft(frames, n=sft.mfft, axis=-1) + return scipy.fft.fftshift( + scipy.fft.fft(frames, n=sft.mfft, axis=-1), axes=-1 + ) + + across = int(axis == 0) + func = parallelize(across, across, self.parallel)(func) + data = func(da.values) + + coord_cls = type(da.coords[self.dim]) + dt = get_sampling_interval(da, self.dim, cast=False) + t0 = da.coords[self.dim].values[0] + time = coord_cls.from_block( + t0 + (sft.m_num // 2) * dt, data.shape[axis], sft.hop * dt + ) + freqs = coord_cls.from_block(sft.f[0], len(sft.f), sft.delta_f) + coords = {} + for name in da.coords: + if name == self.dim: + coords[self.dim] = time + elif da[name].dim != self.dim: # TODO: keep non-dimensional coordinates + coords[name] = da.coords[name] + coords["frequency"] = freqs + return DataArray(data, coords, da.dims + ("frequency",), da.name, da.attrs) + + +filter = atomized(Filter) +decimate = atomized(Decimate) +resample = atomized(Resample) +integrate = atomized(Integrate) +differentiate = atomized(Differentiate) +stft = atomized(STFT) diff --git a/xdas/config.py b/xdas/config.py index c91ea3e3..8b5c35d7 100644 --- a/xdas/config.py +++ b/xdas/config.py @@ -7,11 +7,62 @@ import os from typing import ClassVar +MEMORY_FRACTION = 0.25 +"""Share of the machine's memory the default ``"memory_limit"`` allows.""" + +FALLBACK_MEMORY = 32 * 2**30 +"""Memory assumed when the machine's cannot be determined.""" + +CGROUP_LIMITS = ( + "/sys/fs/cgroup/memory.max", # cgroup v2 + "/sys/fs/cgroup/memory/memory.limit_in_bytes", # cgroup v1 +) +"""Where a container or a batch scheduler declares the memory of a process.""" + + +def total_memory(): + """ + Return the memory this process can use, in bytes. + + The smallest of what the machine has and what a cgroup allows it: a + container or a scheduler allocation is what the process actually gets, + whatever the machine holds. The unlimited sentinel a cgroup writes when + there is no limit is larger than the physical memory, so it loses on its + own. Falls back to `FALLBACK_MEMORY` where neither can be read (Windows, + a sandboxed filesystem). + + Returns + ------- + int + Usable memory in bytes. + """ + limits = [] + for path in CGROUP_LIMITS: + try: + with open(path) as file: + value = file.read().strip() + except OSError: + continue + if value.isdigit(): # "max" spells out an absent limit + limits.append(int(value)) + try: + limits.append(os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")) + except (AttributeError, ValueError, OSError): # pragma: no cover + pass + return min(limits) if limits else FALLBACK_MEMORY + class Config: """Global configuration store backed by a plain dict.""" - config: ClassVar[dict] = {"n_workers": os.cpu_count()} + config: ClassVar[dict] = { + "n_workers": os.cpu_count(), + # A guard against footguns, not a budget: it must sit far enough above + # what one legitimately loads at once that it only ever fires on a + # mistake, which a fixed number cannot do across a laptop and a node + # with a terabyte. + "memory_limit": int(MEMORY_FRACTION * total_memory()), + } def get(key): diff --git a/xdas/coordinates/dense.py b/xdas/coordinates/dense.py index 3001f08c..c8ac3716 100644 --- a/xdas/coordinates/dense.py +++ b/xdas/coordinates/dense.py @@ -68,11 +68,11 @@ def _isvalid(data): @override def _is_monotonic_increasing(self): - if np.issubdtype(self.dtype, np.datetime64): - zero = np.timedelta64(0) - else: - zero = 0 - return np.all(np.diff(self.values) > zero) + # `pandas` is used rather than differencing the values: `np.diff` has no + # `subtract` loop for string dtypes, which would make any label-based + # selection fail. `is_unique` restores the strict increase, since + # `pandas` considers repeated values monotonic increasing. + return self.index.is_monotonic_increasing and self.index.is_unique @override def _get_value(self, index): diff --git a/xdas/core/__init__.py b/xdas/core/__init__.py index 79e23617..a7355506 100644 --- a/xdas/core/__init__.py +++ b/xdas/core/__init__.py @@ -28,6 +28,7 @@ "plot_availability", "sortby", "split", + "stack", "trim_overlaps", ] @@ -52,5 +53,6 @@ plot_availability, sortby, split, + stack, trim_overlaps, ) diff --git a/xdas/core/dataarray.py b/xdas/core/dataarray.py index 717ba22f..f011a347 100644 --- a/xdas/core/dataarray.py +++ b/xdas/core/dataarray.py @@ -137,6 +137,14 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): if method != "__call__": return NotImplemented + if any( + hasattr(input, "__array_ufunc__") + and not isinstance(input, (self.__class__, np.ndarray, np.generic)) + for input in inputs + ): + # Defer to foreign implementations (e.g. atoms tracing pipelines). + return NotImplemented + coords = broadcast_coords( *tuple(input for input in inputs if isinstance(input, self.__class__)) ) @@ -379,15 +387,51 @@ def sel( ------- DataArray The selected part of the original data array. + + Notes + ----- + Naming labels — a scalar, or a list of them — works whatever the order + of the coordinate, which makes categorical axes such as a phase axis + ``["P", "S", "N"]`` selectable; a list returns the labels in the order + it asks for them. Ordered look-ups need an axis whose values increase: + a slice on an axis that goes backwards somewhere is resolved by cutting + it on its overlaps and concatenating, which is slow and warns, and an + inexact look-up (*method*) is refused. + + Examples + -------- + >>> import numpy as np + >>> import xdas as xd + >>> da = xd.DataArray( + ... np.arange(6).reshape(2, 3), + ... {"time": [0.0, 1.0], "phase": ["P", "S", "N"]}, + ... ) + >>> da.sel(phase=["S", "P"]) + + [[1 0] + [4 3]] + Coordinates: + * time (time): [0. 1.] + * phase (phase): ['S' 'P'] """ if indexers is None: indexers = {} indexers.update(indexers_kwargs) - # handle not monotonic increasing coordinates + # Only *ordered* look-ups need a sorted axis. Resolving a slice searches + # for its bounds and `method` searches for a neighbour, so both are + # guarded below: an axis that goes backwards somewhere is cut into + # monotonic chunks for a slice, and refused for a neighbour search. An + # exact label look-up is not ordered — it is a hash look-up, well + # defined whatever the order — so it is left alone, which is what makes + # a categorical axis such as ``["P", "S", "N"]`` selectable by label. + # Labels that designate more than one position are the coordinate's own + # business; `to_index` raises on them. for dim in indexers: - if not self[dim]._is_monotonic_increasing(): - if isinstance(indexers[dim], slice): + is_slice = isinstance(indexers[dim], slice) + is_ordered = is_slice or method is not None + if is_ordered and not self[dim]._is_monotonic_increasing(): + if is_slice: warnings.warn( f"dimension {dim} is not monotonic increasing, " f"spliting on overlaps, slicing and concatenating can be slow..." diff --git a/xdas/core/datacollection.py b/xdas/core/datacollection.py index c8bfe798..6471156b 100644 --- a/xdas/core/datacollection.py +++ b/xdas/core/datacollection.py @@ -9,6 +9,7 @@ from pathlib import Path import h5py +import pandas as pd from .dataarray import DataArray @@ -67,6 +68,11 @@ def __new__(cls, data, name=None): if name is not None: data = data.rename(name) return data + elif isinstance(data, pd.DataFrame): + # A table is a leaf of its own kind: atoms emitting pick tables + # walk collections like any other, and coercing their result into + # a `DataArray` would silently destroy it. + return data else: return DataArray(data, name=name) @@ -262,11 +268,11 @@ def __repr__(self): label = f" {key:{width}}: " else: label = f" {key + ':':{width + 1}} " - if isinstance(value, DataArray): - s += label + repr(value).split("\n")[0] + "\n" - else: + if isinstance(value, DataCollection): s += label + "\n" s += "\n".join(f" {e}" for e in repr(value).split("\n")[:-1]) + "\n" + else: + s += label + repr(value).split("\n")[0] + "\n" return s def __reduce__(self): diff --git a/xdas/core/numpy.py b/xdas/core/numpy.py index 0eaf2cd6..0dcff4a6 100644 --- a/xdas/core/numpy.py +++ b/xdas/core/numpy.py @@ -45,8 +45,12 @@ def decorator(func): def wrapper(*args, **kwargs): """Forward *func* call while preserving or reducing DataArray coordinates.""" ba = sig.bind(*args, **kwargs) + # `defaults` overrides numpy's own defaults, never the caller's + # arguments: bind first, fill what was not given, only then apply + # numpy's remaining defaults. + for name, value in defaults.items(): + ba.arguments.setdefault(name, value) ba.apply_defaults() - ba.arguments.update(defaults) key = next(iter(ba.arguments)) da = ba.arguments.get(key) axis = ba.arguments.get("axis") diff --git a/xdas/core/routines.py b/xdas/core/routines.py index a24ba149..f47f2985 100644 --- a/xdas/core/routines.py +++ b/xdas/core/routines.py @@ -22,6 +22,7 @@ from tqdm import tqdm from ..coordinates import AxisCoordinate, Coordinates +from ..coordinates.core import parse_scalar_delta from ..parallel import get_workers_count from ..virtual import TileArray, VirtualBackend, VirtualSource, VirtualStack from .dataarray import DataArray @@ -1310,9 +1311,6 @@ def concat( return objs[0] if objs else DataArray([]) objs = non_empty - if virtual is None: - virtual = all(isinstance(da.data, (VirtualSource, VirtualStack)) for da in objs) - if dim in objs[0].dims + ("first", "last"): axis = objs[0].get_axis_num(dim) dim = objs[0].dims[axis] # ensure not "first" or "last" @@ -1324,6 +1322,13 @@ def concat( promoted = _get_promoted_coords(objs, dim) objs = [da.expand_dims(dim) for da in objs] + # inferred on what will actually be concatenated: opening a new dimension + # goes through `expand_dims`, which a `VirtualSource` cannot follow (a + # stack of sources is a longer axis, never an extra one) and so loads. + # Inferring beforehand would promise a `VirtualStack` of dense arrays. + if virtual is None: + virtual = all(isinstance(da.data, (VirtualSource, VirtualStack)) for da in objs) + coords = objs[0].coords.drop_dims(dim) name = objs[0].name attrs = objs[0].attrs @@ -1366,6 +1371,575 @@ def concat( concatenate = concat # TODO: deprecate it +#: The alignment strategies :func:`stack` accepts. Deliberately an open +#: enumeration rather than a boolean: SeisBench answers the same mismatch by +#: *splitting the record* into maximal stretches of constant member coverage +#: (``GroupingHelper._get_intervals``), which is the right shape for ragged +#: station deployments and would join this tuple as a further mode rather than +#: replace it. +JOIN_METHODS = (None, "inner", "outer") + +#: Default agreement tolerance of :func:`stack`, as a fraction of the nominal +#: sampling interval. One percent of a sample is far above the sub-nanosecond +#: rounding real acquisitions differ by (a reference three-component station +#: was measured 1 ns apart at 40 Hz, i.e. 4e-8 of a sample) and far below any +#: misalignment worth reporting: it takes a hundred times the budget to hide a +#: one-sample shift, and fifty to hide the half-sample one that would already +#: change which sample a value lands on. +SNAP_FRACTION = 1e-2 + + +def stack(dc, level, dim=None, join=None, tolerance=None): + """ + Collapse a level of a data collection into an array dimension. + + The inverse of :func:`combine_by_coords`, which concatenates *along* an + existing dimension: here the keys of one collection level become the + coordinate of a *new* dimension, and everything below that level is merged + in lock-step. Stacking the ``channel`` level of a seismological collection + turns each station's three traces into one ``(channel, time)`` array. + + The new dimension is named after the level it collapsed, so nothing is + renamed behind your back; pass *dim* to choose another name. + + Parameters + ---------- + dc : DataCollection + The collection to stack. + level : str + The name of the level to collapse. Must name one of ``dc.fields``. + Only the outermost occurrence of that name on each branch is + collapsed. + dim : str, optional + The name of the new dimension. Defaults to *level*. It must not + already name a dimension of the leaves. + join : None or str, optional + How to reconcile leaves that do not share their coordinates. ``None`` + (default) raises, naming what disagreed. ``"inner"`` keeps the + coordinate values every leaf has, ``"outer"`` keeps the values any + leaf has and fills the missing samples with NaN. Aligning materialises + the joined coordinates and, for ``"outer"``, the data. + tolerance : scalar, None, or ``False``, optional + How far apart two leaves may describe the same sampling grid and still + count as agreeing (see the notes). ``None`` (default) spends + :data:`SNAP_FRACTION` of the nominal sampling interval, and only on + coordinates that declare one. ``False`` disables snapping, restoring + strict equality. A scalar is an absolute budget in coordinate units + (seconds for a datetime axis) and applies to every axis coordinate, + declared spacing or not. + + Returns + ------- + DataCollection or DataArray + The collection with the level collapsed. When *level* is the outermost + level and the leaves sit directly below it, the result is a single + data array. + + Raises + ------ + KeyError + If *level* names no level of the collection. + ValueError + If the sub-trees below the collapsed level do not agree structurally, + or if their leaves do not agree on their other coordinates and *join* + does not resolve it. + + Notes + ----- + Stacking is a :func:`concat` over the leaves, so it inherits its + behaviour: the stacked coordinate is sorted, and scalar coordinates that + vary from leaf to leaf are promoted onto the new dimension. + + **Agreement is judged on the sampling grid, not on tie points.** Two + acquisitions of one instrument routinely round their start time + differently by a fraction of a sample; those are the same coordinate, and + comparing them exactly would raise on data that is perfectly aligned — or, + worse, send it to ``join="outer"``, which would interleave the two grids + into an array twice as long. So before any mismatch is reported, leaves + whose axis coordinates have the same length and stay within *tolerance* + of each other everywhere are snapped onto **the first leaf's coordinate**, + in the collection's own key order; the sub-sample offset of the others is + dropped. Only leaves that then still disagree are reported or joined. + + Snapping is deliberately narrow. It never changes a length, it never moves + a value by as much as a sample, and by default it only applies to + coordinates that declare a nominal sampling interval — without one there + is no grid to snap to, and structurally different descriptions of the same + values stay a *join*'s business rather than an equality's. + + Tile-backed leaves stay tile-backed, so stacking a collection of virtual + arrays reads nothing. Alignment is where that can stop: ``"inner"`` slices + the leaves to a shared span, which stays virtual only while the resulting + tile geometries still agree, and ``"outer"`` has to write the NaNs and so + always materialises. + + See Also + -------- + concat : the primitive this is built on, over one list of arrays. + combine_by_coords : concatenate along an existing dimension instead. + + Examples + -------- + >>> import numpy as np + >>> import xdas as xd + + >>> def trace(channel): + ... return xd.DataArray( + ... np.arange(4.0), + ... {"channel": (None, channel), "time": [0.0, 1.0, 2.0, 3.0]}, + ... ) + + >>> dc = xd.DataCollection( + ... { + ... "SX01": ("channel", {code: trace(code) for code in ["SHZ", "SHN"]}), + ... "SX02": ("channel", {code: trace(code) for code in ["SHZ", "SHN"]}), + ... }, + ... "station", + ... ) + >>> dc + Station: + SX01: + Channel: + SHZ: + SHN: + SX02: + Channel: + SHZ: + SHN: + + >>> stacked = xd.stack(dc, "channel") + >>> stacked + Station: + SX01: + SX02: + + >>> stacked["SX01"]["channel"].values + array(['SHN', 'SHZ'], dtype=' 1: + raise ValueError( + f"the sub-trees{_at(path)} of level {level!r} do not agree: " + + ", ".join(f"{key!r} is {kind}" for key, kind in zip(keys, kinds)) + ) + if isinstance(objs[0], DataArray): + return _stack_arrays(objs, keys, level, dim, join, tolerance, path) + name = objs[0].name + if name == level: + raise ValueError( + f"level {level!r} is nested under itself{_at(path)}; stacking two " + "levels sharing a name is not supported" + ) + if objs[0].ismapping(): + subkeys = list(objs[0]) + for key, obj in zip(keys[1:], objs[1:]): + if set(obj) != set(subkeys): + raise ValueError( + f"the {name!r} level{_at(path)} does not hold the same keys " + f"under every {level!r}: {keys[0]!r} has {sorted(subkeys)} " + f"and {key!r} has {sorted(obj)}" + ) + data = { + subkey: _stack_entries( + [obj[subkey] for obj in objs], + keys, + level, + dim, + join, + tolerance, + (*path, f"{name}={subkey}"), + ) + for subkey in subkeys + } + else: + length = len(objs[0]) + for key, obj in zip(keys[1:], objs[1:]): + if len(obj) != length: + raise ValueError( + f"the {name!r} level{_at(path)} does not hold the same number " + f"of elements under every {level!r}: {keys[0]!r} has {length} " + f"and {key!r} has {len(obj)}" + ) + data = [ + _stack_entries( + [obj[index] for obj in objs], + keys, + level, + dim, + join, + tolerance, + (*path, f"{name}={index}"), + ) + for index in range(length) + ] + return DataCollection(data, name) + + +def _stack_arrays(objs, keys, level, dim, join, tolerance, path): + """Concatenate one lock-step group of leaves onto the new dimension *dim*.""" + for key, obj in zip(keys, objs): + if dim in obj.dims: + raise ValueError( + f"cannot stack level {level!r} onto {dim!r}: the leaf {key!r}" + f"{_at(path)} already has a {dim!r} dimension; pass `dim=` to " + "name the new dimension otherwise" + ) + objs = _snap_leaves(objs, tolerance) + messages, joinable = _leaf_mismatches(objs, keys) + if messages and join is not None and joinable: + objs = _join_leaves(objs, keys, joinable, join) + messages, _ = _leaf_mismatches(objs, keys) + joinable = [] # already spent: what remains is not a join away + if messages: + hint = ( + " (pass join='inner' or join='outer' to align them first)" + if joinable and join is None + else "" + ) + raise ValueError( + f"the leaves{_at(path)} of level {level!r} do not agree: " + + "; ".join(messages) + + hint + ) + objs = [_with_key(obj, dim, key) for obj, key in zip(objs, keys)] + return concat(objs, dim) + + +def _snap_leaves(objs, tolerance): + """Put the leaves on one representation of every grid they share within *tolerance*. + + The first leaf is the reference; any other whose axis coordinate has the + same length and stays within *tolerance* of the reference everywhere + adopts it verbatim, so the strict equality that follows sees one + coordinate instead of two roundings of it. ``tolerance=False`` disables + the whole pass. + """ + if tolerance is False or len(objs) < 2: + return objs + reference = objs[0] + out = [reference] + for obj in objs[1:]: + snapped = { + name: reference.coords[name] + for name, coord in obj.coords.items() + if name in reference.coords + and _same_grid(reference.coords[name], coord, tolerance) + } + if snapped: + obj = obj.copy(deep=False) + for name, coord in snapped.items(): + obj.coords[name] = coord.copy() + out.append(obj) + return out + + +def _same_grid(reference, coord, tolerance): + """Whether *coord* describes *reference*'s grid to within *tolerance*. + + Both must be axis coordinates of the same dimension, same length and same + kind of values; the deviation is then measured exactly. Coordinates are + piecewise linear in their index, so comparing them at the union of their + breakpoints bounds their distance everywhere in between. + """ + if not ( + isinstance(reference, AxisCoordinate) + and isinstance(coord, AxisCoordinate) + and reference.dim == coord.dim + and len(reference) == len(coord) + and len(coord) > 0 + ): + return False + if any( + not (np.issubdtype(dtype, np.number) or np.issubdtype(dtype, np.datetime64)) + for dtype in (reference.dtype, coord.dtype) + ): + return False + if np.issubdtype(reference.dtype, np.datetime64) != np.issubdtype( + coord.dtype, np.datetime64 + ): + return False + tolerance = _snap_tolerance(reference, coord, tolerance) + if tolerance is None: + return False + indices = np.union1d(_breakpoints(reference), _breakpoints(coord)) + deviation = np.abs(reference._get_value(indices) - coord._get_value(indices)) + return bool(np.all(deviation <= tolerance)) + + +def _snap_tolerance(reference, coord, tolerance): + """Return the absolute budget to compare two coordinates with, or ``None``. + + An explicit *tolerance* is taken as given, in the coordinate's own units. + The default is :data:`SNAP_FRACTION` of the nominal sampling interval, and + ``None`` — do not snap at all — when neither coordinate declares one: a + fraction of a sample means nothing on an axis that has no samples. + """ + if tolerance is not None: + return parse_scalar_delta(tolerance, reference.dtype) + for candidate in (reference, coord): + sampling_interval = candidate.get_sampling_interval(cast=False) + if sampling_interval is not None: + return np.abs(sampling_interval) * SNAP_FRACTION + return None + + +def _breakpoints(coord): + """Return the indices at which *coord*'s value curve may bend. + + Between two of them the values are affine in the index, which is what lets + a comparison sampled there bound the deviation everywhere. A coordinate + that ties values to indices bends at its tie points (plus the end of each + segment when they carry lengths); one that stores every value is its own + worst case and bends anywhere. + """ + indices = getattr(coord, "tie_indices", None) + if indices is None: + return coord.indices + lengths = getattr(coord, "tie_lengths", None) + if lengths is None: + return np.asarray(indices) + return np.union1d(indices, np.asarray(indices) + np.asarray(lengths) - 1) + + +def _leaf_mismatches(objs, keys): + """Report what the leaves disagree on, and which dimensions a join could fix. + + Returns a list of human-readable messages — empty when the leaves are + stackable as they are — and the names of the dimension coordinates that + differ but could be aligned. + """ + messages = [] + dims = objs[0].dims + for key, obj in zip(keys[1:], objs[1:]): + if obj.dims != dims: + messages.append( + f"{keys[0]!r} has dimensions {dims} and {key!r} has {obj.dims}" + ) + if messages: + return messages, [] + names = list(objs[0].coords) + for key, obj in zip(keys[1:], objs[1:]): + missing = [name for name in names if name not in obj.coords] + extra = [name for name in obj.coords if name not in names] + if missing or extra: + messages.append( + f"{key!r} lacks the coordinates {missing} of {keys[0]!r} and " + f"carries {extra} it does not" + ) + if messages: + return messages, [] + joinable = [] + for dim in dims: + sizes = sorted({obj.sizes[dim] for obj in objs}) + if len(sizes) > 1 and dim not in objs[0].coords: + messages.append( + f"dimension {dim!r} has sizes {sizes} and no coordinate to align on" + ) + for name in names: + coords = [obj.coords[name] for obj in objs] + if all(coord.equals(coords[0]) for coord in coords[1:]): + continue + if all(coord.dim is None for coord in coords): + continue # varying scalars are promoted onto the new dimension + if name in dims and all(isinstance(coord, AxisCoordinate) for coord in coords): + joinable.append(name) + messages.append(f"coordinate {name!r} differs from one leaf to another") + return messages, joinable + + +def _join_leaves(objs, keys, dims, join): + """Reindex the leaves onto a shared index along each of *dims*.""" + for dim in dims: + indices = [pd.Index(obj.coords[dim].values) for obj in objs] + for key, index in zip(keys, indices): + if not index.is_unique: + raise ValueError( + f"cannot align on {dim!r}: the leaf {key!r} repeats coordinate " + "values; resolve its overlaps first (see `trim_overlaps`)" + ) + target = indices[0] + for index in indices[1:]: + if join == "inner": + target = target.intersection(index, sort=False) + else: + target = target.union(index, sort=None) + if len(target) == 0: + raise ValueError( + f"cannot align on {dim!r}: the leaves share no coordinate value" + ) + _refuse_interleaving(objs, dim, target) + objs = [_reindex(obj, dim, target, index) for obj, index in zip(objs, indices)] + objs = _unify_coord(objs, dim) + return objs + + +def _refuse_interleaving(objs, dim, target): + """Raise when the joined index holds more samples than its span can carry. + + An outer join over leaves that are on the same grid spans it once. Over + leaves that are a fraction of a sample apart it spans it as many times as + there are offsets, interleaving grids into an array that looks plausible + and is mostly missing samples. The finest declared sampling interval says + how many samples the joined span may hold; anything beyond that is + interleaving, and silence would be the worst answer. + """ + intervals = [ + obj.coords[dim].get_sampling_interval(cast=False) + for obj in objs + if obj.coords[dim].isregular() + ] + if len(intervals) < len(objs): + return # an irregular leaf declares no grid to violate + sampling_interval = min(np.abs(interval) for interval in intervals) + # rounded, not truncated: a float span is worth a sample either way, while + # interleaving overshoots by a factor, never by one. + expected = round((target.max() - target.min()) / sampling_interval) + 1 + if len(target) > expected: + raise ValueError( + f"cannot align on {dim!r}: the leaves are not on a common sampling " + f"grid, and joining them would interleave {len(target)} samples " + f"where the span holds {expected}; snap them together first with a " + "larger `tolerance`" + ) + + +def _reindex(obj, dim, target, index): + """Return *obj* with its *dim* axis put on *target*, staying lazy when it can.""" + positions = index.get_indexer(target) + if len(index) == len(target) and np.array_equal(positions, np.arange(len(index))): + return obj + if (positions >= 0).all(): + start = int(positions[0]) + # a contiguous run is a slice, and slicing a virtual array reads nothing + if np.array_equal(positions, np.arange(start, start + len(positions))): + return obj.isel({dim: slice(start, start + len(positions))}) + return obj.isel({dim: positions}) + return _pad(obj, dim, target, positions) + + +def _pad(obj, dim, target, positions): + """Return *obj* on *target*, filling the samples it does not have with NaN.""" + others = [ + name for name, coord in obj.coords.items() if coord.dim == dim and name != dim + ] + if others: + raise ValueError( + f"cannot pad along {dim!r}: the leaves carry the coordinates {others} " + "along it, which have no value where the data is missing" + ) + axis = obj.get_axis_num(dim) + present = positions >= 0 + dtype = ( + obj.dtype + if np.issubdtype(obj.dtype, np.inexact) + else np.result_type(obj.dtype, np.float32) + ) + shape = list(obj.shape) + shape[axis] = len(target) + data = np.full(tuple(shape), np.nan, dtype) + key = [slice(None)] * obj.ndim + key[axis] = np.nonzero(present)[0] + data[tuple(key)] = np.asarray(obj.isel({dim: positions[present]}).data) + coords = obj.coords.copy() + coords[dim] = target.values + return DataArray(data, coords, obj.dims, obj.name, obj.attrs) + + +def _unify_coord(objs, dim): + """Give every leaf the same *dim* coordinate object once they agree on its values. + + Reindexing puts every leaf on the same values, but not necessarily on the + same *representation*: two interpolated coordinates sliced out of + differently tied inputs describe one grid with different tie points, and + :meth:`Coordinate.equals` is structural. Normalising here is what lets the + equality check that follows stay strict. + """ + reference = objs[0].coords[dim] + out = [objs[0]] + for obj in objs[1:]: + coord = obj.coords[dim] + if not coord.equals(reference) and np.array_equal( + coord.values, reference.values + ): + obj = obj.copy(deep=False) + obj.coords[dim] = reference.copy() + out.append(obj) + return out + + +def _with_key(obj, dim, key): + """Return *obj* carrying its level key as a scalar coordinate named *dim*. + + That is all it takes for :func:`concat` to open the new dimension with the + keys as its coordinate: ``expand_dims`` promotes the scalar and + ``concat_coords`` concatenates the promoted length-one coordinates. + """ + obj = obj.copy(deep=False) + obj.coords[dim] = key + return obj + + def sortby(da, dim="first", tolerance=None): """ Sort a blocked virtual data array along *dim* by coordinate value, lazily. diff --git a/xdas/fft.py b/xdas/fft.py index 47893c46..73a7bc2a 100644 --- a/xdas/fft.py +++ b/xdas/fft.py @@ -8,12 +8,14 @@ import numpy as np from .atoms import atomized +from .atoms.core import _whole_record from .coordinates import get_sampling_interval from .core import DataArray from .parallel import parallelize @atomized +@_whole_record(default="last") def fft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the discrete Fourier Transform along a given dimension. @@ -88,6 +90,7 @@ def func(x): @atomized +@_whole_record(default="last") def rfft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the discrete Fourier Transform for real inputs along a given dimension. @@ -158,6 +161,7 @@ def rfft(da, n=None, dim=None, norm=None, parallel=None): @atomized +@_whole_record(default="last") def ifft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the inverse of `fft`. @@ -228,6 +232,7 @@ def func(x): @atomized +@_whole_record(default="last") def irfft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the inverse of `rfft`. diff --git a/xdas/io/asn.py b/xdas/io/asn.py index 97175d6c..56a63a4e 100644 --- a/xdas/io/asn.py +++ b/xdas/io/asn.py @@ -15,6 +15,7 @@ from ..coordinates import Coordinate, get_sampling_interval from ..core import DataArray, concat_coords +from ..processing.core import SubscriptionTracker, ZMQEndpoint from ..virtual import TileArray, VirtualSource from .core import Engine @@ -122,7 +123,7 @@ def _get_roi_bound_indices(self, all_dists, n_start, n_end, dx): } -class ZMQSubscriber: +class ZMQSubscriber(ZMQEndpoint): """ Iterator that pulls :class:`DataArray` chunks from a live ASN ZMQ publisher. @@ -130,9 +131,20 @@ class ZMQSubscriber: ---------- address : str ZMQ address of the publisher (e.g. ``"tcp://localhost:5555"``). + timeout : float or None, optional + How many seconds to wait at most for each message. None, the default, + waits forever. + + Methods + ------- + wait_until_subscribed() + Block until the publisher has registered this subscription. Building + the subscriber already does it. + close() + Release the socket and its context. """ - def __init__(self, address): + def __init__(self, address, timeout=None): """ Initialize a ZMQStream object. @@ -140,10 +152,11 @@ def __init__(self, address): ---------- address : str The address to connect to. + timeout : float or None, optional + How many seconds to wait at most for each message. Examples -------- - >>> import time >>> import threading >>> import xdas as xd @@ -157,21 +170,30 @@ def __init__(self, address): >>> chunks = xd.split(da, 10) >>> def publish(): + ... publisher.wait_for_subscribers() # a replay, so no chunk is lost ... for chunk in chunks: - ... time.sleep(0.001) # so that the subscriber can connect in time ... publisher.submit(chunk) - >>> threading.Thread(target=publish).start() + >>> thread = threading.Thread(target=publish) + >>> thread.start() >>> subscriber = ZMQSubscriber(address) >>> for nchunk in range(10): ... chunk = next(subscriber) ... # do something with the chunk + Both ends hold a socket until they are closed, by hand as here or by + using them as context managers. + + >>> thread.join() + >>> subscriber.close() + >>> publisher.close() + """ self.address = address + self.timeout = timeout + self._subscribed = False self._connect(self.address) - message = self._get_message() - self._update_header(message) + self.wait_until_subscribed() def __iter__(self): return self @@ -189,11 +211,53 @@ def _connect(self, address): socket = context.socket(zmq.SUB) socket.connect(address) socket.setsockopt_string(zmq.SUBSCRIBE, "") + self._context = context self._socket = socket + def wait_until_subscribed(self): + """ + Block until the publisher has registered this subscription. + + A publisher drops what it sends to a peer it does not know about yet, + and a subscriber cannot tell from its own side whether its + subscription has arrived — being connected is not being subscribed. + Here the proof comes for free: the header describing the stream is the + greeting an ASN publisher answers a new subscription with, in passing + as it streams. It never waits for anyone, and receiving its header is + proof that nothing it publishes from then on will be missed. + + This is done when the subscriber is built — a packet cannot be decoded + before it — so calling it again returns immediately. The one stream + that keeps a subscriber waiting is one that is not streaming: a + publisher that has gone quiet, or has yet to send its first packet, + acknowledges nobody. + """ + # A subscriber that beat the first publication to the socket gets no + # welcome message, and can be handed data before the header is sent. + while not self._subscribed: + message = self._get_message() + if self._is_header(message): + self._update_header(message) + self._subscribed = True + def _get_message(self): + if self.timeout is not None and not self._socket.poll( + round(1000 * self.timeout) + ): + raise TimeoutError( + f"no message received from {self.address} after {self.timeout} seconds" + ) return self._socket.recv() + @staticmethod + def _is_header(message): + """Whether *message* is a header rather than a data packet.""" + try: + header = json.loads(message.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + return False + return isinstance(header, dict) and "bytesPerPackage" in header + def _is_packet(self, message): return len(message) == self.packet_size @@ -223,7 +287,7 @@ def _unpack(self, message): return DataArray(data, {"time": time, "distance": self.distance}) -class ZMQPublisher: +class ZMQPublisher(SubscriptionTracker): """ A class to stream data using ZeroMQ. @@ -236,11 +300,18 @@ class ZMQPublisher: ---------- address : str The address where the ZeroMQ is bound to. + nsubscribers : int + The number of currently subscribed peers. Methods ------- submit(da) Submits the data array for publishing. + wait_for_subscribers(count, timeout) + Blocks until *count* peers are subscribed, so that nothing published + afterwards is dropped. + close() + Release the socket and its context. Examples -------- @@ -251,15 +322,16 @@ class ZMQPublisher: >>> port = xd.io.get_free_port() >>> address = f"tcp://localhost:{port}" - >>> publisher = ZMQPublisher(address) >>> chunks = xd.split(da, 10) - >>> for chunk in chunks: - ... publisher.submit(chunk) + >>> with ZMQPublisher(address) as publisher: + ... for chunk in chunks: + ... publisher.submit(chunk) """ def __init__(self, address): self.address = address + self._nsubscribers = 0 self._connect(address) self._header = None @@ -272,7 +344,9 @@ def header(self): def header(self, header): """Set the welcome-message header and push it to the ZMQ socket option.""" self._header = header - self.socket.setsockopt(zmq.XPUB_WELCOME_MSG, json.dumps(header).encode("utf-8")) + self._socket.setsockopt( + zmq.XPUB_WELCOME_MSG, json.dumps(header).encode("utf-8") + ) def submit(self, da): """Publish *da* over ZMQ.""" @@ -287,7 +361,8 @@ def _connect(self, address): socket = context.socket(zmq.XPUB) socket.setsockopt(zmq.XPUB_VERBOSE, True) socket.bind(address) - self.socket = socket + self._context = context + self._socket = socket @staticmethod def _get_header(da): @@ -306,11 +381,17 @@ def _get_header(da): return header def _send(self, da): + # Taking the subscriptions the socket has queued is what greets the + # peers behind them with the header — ZeroMQ holds a welcome message + # back until the application reads the subscription it answers — and + # what keeps the subscriber count current. Neither costs any waiting. + self._read_subscriptions(0.0) da = da.transpose("time", "distance") header = self._get_header(da) - if self.header is None: - self.header = header if header != self.header: + # Peers that subscribed before there was a header to welcome them + # with — including the very first one — only learn the layout if it + # is sent down the stream, so the first submit publishes it too. self.header = header self._send_header() self._send_data(da) @@ -327,7 +408,7 @@ def _send_data(self, da): self._send_message(message) def _send_message(self, message): - self.socket.send(message) + self._socket.send(message) def float_to_timedelta(value, unit): diff --git a/xdas/io/tdms.py b/xdas/io/tdms.py index 082130dc..b6ef360e 100644 --- a/xdas/io/tdms.py +++ b/xdas/io/tdms.py @@ -156,29 +156,36 @@ def __init__(self, filename): # TODO: Error if file not big enough to hold header # handle outlives __init__ on purpose: reads are lazy, closed in __exit__ self._tdms_file = open(filename, "rb") # noqa: SIM115 - # Read lead in (28 bytes): - lead_in = self._tdms_file.read(LEAD_IN_LENGTH) - # lead_in is 28 bytes: - # [string of length 4][int32][int32][int64][int64] - fields = struct.unpack("<4siiQQ", lead_in) - - # TODO: validate file - if fields[0].decode() not in "TDSm": - msg = "Not a TDMS file (TDSm tag not found)" - raise (TypeError, msg) - - self.fileinfo = dict(zip(FILEINFO_NAMES, fields)) - self.fileinfo["decimated"] = not bool(self.fileinfo["toc"] & DECIMATE_MASK) - # Make offsets relative to beginning of file: - self.fileinfo["next_segment_offset"] += LEAD_IN_LENGTH - self.fileinfo["raw_data_offset"] += LEAD_IN_LENGTH - self.fileinfo["file_size"] = os.path.getsize(self._tdms_file.name) - - # TODO: Validate lead in: - self.fileinfo["next_segment_offset"] = min( - self.fileinfo["next_segment_offset"], self.file_size - ) - # raise(ValueError, "Next Segment Offset too large in TDMS header") + try: + # Read lead in (28 bytes): + lead_in = self._tdms_file.read(LEAD_IN_LENGTH) + # lead_in is 28 bytes: + # [string of length 4][int32][int32][int64][int64] + fields = struct.unpack("<4siiQQ", lead_in) + + # TODO: validate file + if fields[0].decode() not in "TDSm": + msg = "Not a TDMS file (TDSm tag not found)" + raise (TypeError, msg) + + self.fileinfo = dict(zip(FILEINFO_NAMES, fields)) + self.fileinfo["decimated"] = not bool(self.fileinfo["toc"] & DECIMATE_MASK) + # Make offsets relative to beginning of file: + self.fileinfo["next_segment_offset"] += LEAD_IN_LENGTH + self.fileinfo["raw_data_offset"] += LEAD_IN_LENGTH + self.fileinfo["file_size"] = os.path.getsize(self._tdms_file.name) + + # TODO: Validate lead in: + self.fileinfo["next_segment_offset"] = min( + self.fileinfo["next_segment_offset"], self.file_size + ) + # raise(ValueError, "Next Segment Offset too large in TDMS header") + except Exception: + # A reader that fails to build is never handed to ``with``, so its + # ``__exit__`` never runs and the handle would be left open. Every + # auto-detection sniff of a file that is not TDMS lands here. + self._tdms_file.close() + raise def __enter__(self): return self diff --git a/xdas/processing/__init__.py b/xdas/processing/__init__.py index 9e24e055..dbe71d84 100644 --- a/xdas/processing/__init__.py +++ b/xdas/processing/__init__.py @@ -1,8 +1,8 @@ """ Chunked processing pipeline for larger-than-memory datasets. -Provides loaders, writers, real-time streaming, and the :func:`process` -orchestrator. +Provides loaders, writers, real-time streaming, the :func:`process` +orchestrator and its :func:`get_source` / :func:`get_writer` dispatch. """ __all__ = [ @@ -10,10 +10,14 @@ "DataArrayWriter", "DataFrameWriter", "RealTimeLoader", + "ResultWriter", "StreamWriter", "ZMQPublisher", "ZMQSubscriber", + "get_source", + "get_writer", "process", + "watch", ] from .core import ( @@ -21,8 +25,12 @@ DataArrayWriter, DataFrameWriter, RealTimeLoader, + ResultWriter, StreamWriter, ZMQPublisher, ZMQSubscriber, + get_source, + get_writer, process, + watch, ) diff --git a/xdas/processing/core.py b/xdas/processing/core.py index 5ff15832..fd5f2cac 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -3,10 +3,16 @@ Includes :class:`DataArrayLoader`, :class:`DataArrayWriter`, :class:`DataFrameWriter`, :class:`StreamWriter`, :class:`ZMQPublisher`, -:class:`ZMQSubscriber`, :class:`RealTimeLoader`, and :func:`process`. +:class:`ZMQSubscriber` with the :class:`SubscriptionTracker` shared with the +ASN publisher, :class:`RealTimeLoader`, :func:`watch`, and the :func:`process` +dispatch boundary with its :func:`get_source` / :func:`get_writer` resolution +machinery. """ import os +import re +import time +import warnings from concurrent.futures import ThreadPoolExecutor from glob import glob from pathlib import Path @@ -20,54 +26,909 @@ from watchdog.events import FileSystemEventHandler from watchdog.observers import Observer -from ..core import DataArray, concat, open_dataarray +from .. import config +from ..atoms.core import ( + _annotate_path, + _announce_splits, + _aschunks, + _extend_path, + _iter_results, + _join_chunks, +) +from ..coordinates import AxisCoordinate +from ..coordinates.core import parse_scalar_delta +from ..core import ( + DataArray, + DataCollection, + DataMapping, + DataSequence, + concat, + open_dataarray, + open_mfdataarray, +) +from ..virtual import TileArray, VirtualBackend, VirtualStack from .monitor import Monitor +from .pools import ProcessPool +AUTO_CHUNK_NBYTES = 256 * 2**20 +"""Target in-memory chunk size (in bytes) for ``chunks="auto"``.""" -def process(atom, data_loader, data_writer): +WELCOME = b"xdas" +""" +Greeting a :class:`ZMQPublisher` sends to each subscriber as it registers it. + +Never a packet — those are netCDF binaries — so subscribers skip it. Receiving +it is how a subscriber knows the publisher has applied its subscription. +""" + + +POOLS = { + "threads": lambda max_workers, max_buffers, slot_nbytes: ThreadPoolExecutor( + max_workers + ), + "processes": ProcessPool, +} +"""Worker pools available for chunk ingress and egress, name → factory.""" + + +def get_pool(pool, max_workers, max_buffers=1, slot_nbytes=None): """ - Execute a chunked processing pipeline. + Build the worker pool used to load or write chunks. + + Threads are enough when the work releases the GIL, but compressed HDF5 + does not: reading a chunk goes through one virtual layout, hence one h5py + call holding the global HDF5 lock, so decompression of several chunks + cannot overlap in-process and extra threads only contend. Worker + processes each hold their own lock. What crosses to a worker is the + *manifest* of the chunk (a sliced virtual array, kilobytes), not data, + and the loaded chunk crosses *back* through shared memory: the worker + writes it once into an arena slot, the parent maps the same pages + (read-only). See :class:`~xdas.processing.pools.ProcessPool`. Parameters ---------- - atom : callable - The atomic operation to execute on each chunk of data. - data_loader : DataArrayLoader - The data loader object that provides the chunks of data. - data_writer : DataArrayWriter - The data writer object that writes the processed data. + pool : str + Pool kind, ``"threads"`` (default everywhere) or ``"processes"``. + max_workers : int + Number of workers. + max_buffers : int, optional + Chunks the caller keeps in flight, which sizes the arena. Default 1. + slot_nbytes : int, optional + Largest chunk the arena can carry. Defaults to the ``chunks="auto"`` + budget; pass the real chunk size when it is known, so that no chunk + has to fall back to the pickle path. + + Returns + ------- + executor + A :class:`~concurrent.futures.Executor`-like pool. + + Examples + -------- + >>> from xdas.processing.core import get_pool + >>> with get_pool("threads", 2) as pool: + ... pool.submit(abs, -1).result() + 1 + """ + if pool not in POOLS: + raise ValueError(f"no worker pool named {pool!r}; available: {sorted(POOLS)}") + if slot_nbytes is None: + slot_nbytes = AUTO_CHUNK_NBYTES + return POOLS[pool](max_workers, max_buffers, slot_nbytes) + + +def _load(da): + """Load a (virtual) DataArray. The unit of work shipped to ingress workers.""" + return da.load() + + +def _dump(chunk, path, encoding): + """Write *chunk* to *path* and return it virtually. Egress unit of work.""" + chunk.to_netcdf(path, encoding=encoding) + return open_dataarray(path) + + +def process(atom, source, out=None, chunks=None, until=None, merge=True): + """ + Execute a processing pipeline over any chunk source, into any sink. + + The dispatch boundary of chunked processing: the input is resolved into a + chunk source with :func:`get_source` and the output into a writer with + :func:`get_writer` (deferred to the first output chunk, so the writer can + match what the pipeline actually emits). Explicit loader and writer + instances pass through untouched, so the historical + ``process(atom, data_loader, data_writer)`` form keeps working. + + Parameters + ---------- + atom : Atom or callable + The operation to execute on each chunk of data. It may emit zero or + more output chunks per input chunk (seam tails, rechunking, + reductions); at the end of the stream it is flushed so buffering + atoms emit their remainder. :meth:`Atom.process` is this function + with the atom bound. + source : DataArray, DataCollection, str, Path, iterable or loader + What to process. An in-memory :class:`DataArray` is processed in one + eager call (or chunk by chunk if `chunks` is given); a virtual + DataArray streams through a :class:`DataArrayLoader`; a + :class:`~xdas.DataCollection` is walked leaf by leaf; a path, + directory or glob pattern is opened with :func:`open_mfdataarray` + first; a ``"tcp://..."`` address subscribes to a + :class:`ZMQSubscriber`; any iterable of chunks (including + :func:`watch` and generators) is consumed as is. + out : str, Path, writer or None, optional + Where to write the output. ``None`` (default) accumulates the output + chunks in memory and returns the joined result, guarded by the + ``"memory_limit"`` configuration entry. A path is matched with the + first output chunk: a directory for DataArray (netcdf chunks) or + Stream (SDS) chunks, a ``*.csv`` file for DataFrame chunks, a + ``"tcp://..."`` address publishes DataArrays. Non-inferable + configuration (miniseed data quality, encodings) is passed as a + ready writer instance. Walking a collection, a single file, a URL + or a ready writer is *shared* by every leaf — one table for the + whole collection, the tree-path columns keeping the rows apart — + while a directory fans out, one subdirectory per leaf mirroring the + tree path (see :class:`_CollectionSink`). + chunks : dict or "auto", optional + Chunk sizes for DataArray sources, e.g. ``{"time": 1000}``. + ``"auto"`` (the default for virtual sources) aligns chunk boundaries + to the storage tiling, merged up to ``AUTO_CHUNK_NBYTES``. + until : str, datetime64 or float, optional + Stop processing at this coordinate value along the chunked dimension. + The chunk containing it is truncated; the pipeline is then flushed + normally. This is the clean way to bound an unbounded source. + merge : bool, optional + Whether to fold the per-leaf results of a collection walk through + the atom's :attr:`~xdas.atoms.Atom.merge` hook, as ``atom(dc)`` + does. Walk-level, not stage-level, and only meaningful for + ``out=None``: with any other sink the leaves were written as they + were produced and there is nothing left to fold. Default ``True``. Returns ------- result : object - The result of the processing pipeline. + The writer result: the joined output for ``out=None``, whatever the + resolved writer's ``result()`` returns otherwise, or ``None`` when + the pipeline emitted no output (no empty outputs are created). + Walking a collection, the result of a shared sink is its single + result and everything else answers with the tree, so + ``atom.process(dc, out=None)`` equals ``atom(dc)``. Notes ----- - This function executes a chunked processing pipeline by ingesting the data from - the `data_loader` and flushing the processed data through the `data_writer`. - It iterates over the chunks of data provided by the `data_loader`, applies the - `atom` function to each chunk, and writes the processed data using the `data_writer`. - The progress of the processing is monitored using a `Monitor` object. + Unbounded sources (:func:`watch`, ZMQ subscriptions — anything exposing + ``unbounded = True``) get streaming semantics: no byte total on the + progress monitor, and a clean :exc:`KeyboardInterrupt` stops the loop, + flushes the pipeline and returns the writer result instead of raising. + + A :class:`~xdas.DataCollection` source is *walked*, exactly as + ``atom(dc)`` walks it (see :meth:`~xdas.atoms.Atom.__call__`): each + mapping level is first offered to the atom's + :meth:`~xdas.atoms.Atom.gather` hook — before anything is chunked, so a + channel level becomes a component dimension and the stacked array + streams as one thing — then recursed into, sequence levels are folded + element by element, and every leaf is streamed through the single-source + path with `chunks` and `until` applying per leaf. One atom instance + walks the leaves sequentially, reset between them: an atom holding a + model either saturates the CPU or holds a lot of device memory, so only + one should be live per node. Output chunks carry the tree path of the + leaf that produced them as they are produced, so a streaming walk hands + a leaf straight to a sink and still knows whose it was. + + A :class:`~xdas.DataSequence` handed over as a collection is walked like + any other, so it folds rather than streaming as one concatenated result. + A glob or a directory that *opens* to a sequence is still a single + source, chained by :func:`get_source`; passing ``get_source(sequence)`` + explicitly asks for the same of a collection in hand. + """ + if isinstance(source, (DataMapping, DataSequence)): + return _process_collection(atom, source, out, chunks, until, merge) + return _process_source(atom, source, out, chunks, until) + + +def _process_source(atom, source, out, chunks, until, path=None): + """ + Stream one source through *atom* into one sink (see :func:`process`). + The single-source path, and the only place chunks are actually driven. + *path* is the tree path of the leaf being streamed, when this runs as + one step of a collection walk: every output chunk is labelled with it + on its way to the sink, as it is produced. """ + path = {} if path is None else path + source = get_source(source, chunks) + if isinstance(source, DataArray): + # In-memory, unchunked: direct eager call, then sink dispatch on the + # result so `process(da, out=...)` and `pipeline(da)` stay twins. + if until is not None: + # `until` truncates the source rather than the stream here, which + # is the same cut: inclusive, as `sel` slices. + dim = getattr(atom, "_resolve_dim", lambda _: None)(source) or "time" + source = source.sel({dim: slice(None, until)}) + result = _annotate_path(atom(source), path) + if out is None: + return result + outputs = _aschunks(result) + if not outputs: + return None + # Nothing was chunked, so there is no chunk dimension to name: the + # chunks one eager call returns are joined the way `concat` defaults to. + writer = get_writer(out, outputs[0], "first") + for chunk in outputs: + writer.write(chunk) + result = writer.result() + _close_if_owned(writer, out) + return result if hasattr(atom, "reset"): atom.reset() - if hasattr(data_loader, "nbytes"): - total = data_loader.nbytes - else: - total = None + chunk_dim = getattr(source, "chunk_dim", "time") + unbounded = bool(getattr(source, "unbounded", False)) + total = None if unbounded else getattr(source, "nbytes", None) + if isinstance(until, str): + until = np.datetime64(until) + # Writer instances pass through upfront; inferred writers are deferred to + # the first output chunk (the correct writer depends on what the pipeline + # emits, and deferral avoids creating empty outputs). + writer = out if hasattr(out, "write") and hasattr(out, "result") else None + + def write(chunk): + nonlocal writer + chunk = _annotate_path(chunk, path) + if writer is None: + writer = get_writer(out, chunk, chunk_dim) + writer.write(chunk) + + if not unbounded and isinstance(getattr(source, "da", None), DataArray): + # Free to know upfront: the source coordinate says where the runs + # split, before any data is read. + coord = source.da.coords.get(chunk_dim, None) + if isinstance(coord, AxisCoordinate) and coord.isregular(): + indices = coord.get_split_indices( + "discontinuities", getattr(coord, "tolerance", None) + ) + if indices.size: + _announce_splits(source.da, chunk_dim, int(indices.size)) + previous = None monitor = Monitor(total=total) monitor.tic("read") - for chunk in data_loader: - monitor.tic("proc") - result = atom(chunk, chunk_dim=data_loader.chunk_dim) - monitor.tic("write") - data_writer.write(result) - monitor.toc(chunk.nbytes) - monitor.tic("read") + try: + for chunk in source: + last = False + if until is not None and isinstance(chunk, DataArray): + coord = chunk.coords.get(chunk_dim, None) + if isinstance(coord, AxisCoordinate) and not coord.empty: + if coord.start > until: + break + if coord.end >= until: + # `until` is inclusive, like `sel(slice(None, until))`. + chunk = chunk.sel({chunk_dim: slice(None, until)}) + last = True + if unbounded and isinstance(chunk, DataArray): + # A realtime source cannot be inspected upfront: announce + # each seam as it arrives instead. + previous = _announce_realtime_seam(previous, chunk, chunk_dim) + monitor.tic("proc") + result = atom(chunk, chunk_dim=chunk_dim) + monitor.tic("write") + for chunk_out in _aschunks(result): + write(chunk_out) + monitor.toc(getattr(chunk, "nbytes", 0)) + monitor.tic("read") + if last: + break + except KeyboardInterrupt: + if not unbounded: + raise + finally: + if unbounded and hasattr(source, "stop"): + source.stop() + if hasattr(atom, "flush"): + for chunk_out in atom.flush(): + write(chunk_out) monitor.close() - return data_writer.result() + if writer is None: + return None + result = writer.result() + _close_if_owned(writer, out) + return result + + +def _close_if_owned(writer, out): + """ + Release a sink :func:`process` opened itself. + + A publisher built from a ``"tcp://..."`` spec holds a socket that nothing + else will ever close, whereas one the caller passed in stays theirs to + reuse and to close. + """ + if writer is not out and isinstance(writer, ZMQEndpoint): + writer.close() + + +def _process_collection(atom, dc, out, chunks, until, merge): + """ + Walk a collection leaf by leaf, streaming each leaf into the sink. + + The streaming twin of :meth:`~xdas.atoms.Atom._walk`, and it mirrors it + step for step so the two cannot drift apart: ``atom.process(dc, + out=None)`` returns what ``atom(dc)`` returns. + """ + walk = _Walk(atom, _CollectionSink(out), chunks, until) + result = walk.sink.result(walk.level(dc, {}, ())) + if out is None and merge and getattr(atom, "merge", None) is not None: + return atom.merge(list(_iter_results(result))) + return result + + +class _Walk: + """ + One walk of a collection: what runs, where it writes, how it streams. + + Holds everything a walk keeps constant — the atom, the resolved sink, + and the `chunks` and `until` that apply per leaf — so the recursion + carries only what changes: the level, and the tree path it was reached + by. + + Parameters + ---------- + atom : Atom or callable + The operation, one instance for the whole walk. + sink : _CollectionSink + The out spec, resolved per leaf. + chunks, until + As :func:`process`, applied per leaf. + """ + + def __init__(self, atom, sink, chunks, until): + self.atom = atom + self.sink = sink + self.chunks = chunks + self.until = until + + def level(self, x, path, where): + """Walk one collection level, or stream *x* if it is a leaf.""" + if isinstance(x, DataMapping): + gathered = self.atom.gather(x) if hasattr(self.atom, "gather") else None + if gathered is not None: + # Consulted before anything is chunked: the level becomes an + # axis of the input and the stacked array streams as one + # thing. Being consumed, it contributes no path column. + return self.level(gathered, path, where) + name = getattr(x, "name", None) + return DataCollection( + { + key: self.level(value, _extend_path(path, name, key), (*where, key)) + for key, value in x.items() + }, + name, + ) + if isinstance(x, DataSequence): + return self.fold(x, path, where) + if hasattr(self.atom, "reset"): + self.atom.reset() # one atom instance, the leaves taken one by one + return _asleaf(self.stream(self.atom, x, self.sink.spec(where), path)) + + def fold(self, x, path, where): + """ + Fold a sequence level: one stream delivered in pieces. + + State carries from element to element — the seams are judged from + the coordinates, as everywhere else — so the atom is flushed once, + after the last element, and its tail is attributed to that last + element. The whole sequence shares one sink: it is one stream, so + its chunks belong in one directory and concatenate into one result. + Only the tree-path column tells the elements apart, and only + approximately, since a buffering atom releases element *i-1*'s + samples while element *i* is being fed. + """ + name = getattr(x, "name", None) + if hasattr(self.atom, "reset"): + self.atom.reset() + first = next((el for el in x if isinstance(el, DataArray)), None) + resolve = getattr(self.atom, "_resolve_dim", None) + dim = resolve(first) if resolve is not None else None + if dim is None: + # Nothing to fold along: each element is a leaf of its own. + return DataCollection( + [ + self.level( + element, _extend_path(path, name, index), (*where, index) + ) + for index, element in enumerate(x) + ], + name, + ) + spec = self.sink.spec(where) + writer = _SharedWriter(ResultWriter(None) if spec is None else spec) + for index, element in enumerate(x): + if not isinstance(element, DataArray): + raise NotImplementedError( + "chunked processing of mapping collections is not supported: " + "process each leaf with its own atom instance" + ) + source = get_source(element, self.chunks) + if isinstance(source, DataArray): + source = _Element(source, dim) + self.stream( + _Held(self.atom), source, writer, _extend_path(path, name, index) + ) + # Past the `dim` resolution the operation is an atom, so it flushes + # and resets; the tail is what its last element left buffered. + for chunk in self.atom.flush(): + writer.write( + _annotate_path(chunk, _extend_path(path, name, max(len(x) - 1, 0))) + ) + self.atom.reset() + result = writer.close() + if spec is None: + # As `Atom._fold`: the level answers with the chunks of its stream. + return DataCollection(_aschunks(result), name) + return _asleaf(result) + + def stream(self, atom, source, out, path): + """Stream one leaf, or one element of a folded sequence, into *out*.""" + return _process_source(atom, source, out, self.chunks, self.until, path) + + +def _asleaf(result): + """Normalize a leaf result: nothing written is an empty collection.""" + return DataCollection([]) if result is None else result + + +class _Held: + """ + An atom facade whose flush is held back. + + A sequence level is one stream delivered in pieces, so its elements must + not each end the stream: the state carries across and the tail is + flushed once, by the walk, after the last element. Holding the flush is + also what keeps :func:`_process_source` from resetting the atom between + elements — the facade declares no ``reset``. + """ + + def __init__(self, atom): + self.atom = atom + + def __call__(self, chunk, **flags): + """Process one chunk through the held atom.""" + return self.atom(chunk, **flags) + + def flush(self): + """Emit nothing: the stream is not over yet.""" + return [] + + +class _Element: + """One in-memory element of a folded sequence, as a single-chunk source.""" + + def __init__(self, da, chunk_dim): + self.da = da + self.chunk_dim = chunk_dim + + @property + def nbytes(self): + """Size of the element, in bytes.""" + return self.da.nbytes + + def __iter__(self): + yield self.da + + +class _SharedWriter: + """ + One writer shared by several streams of a collection walk. + + A ``*.csv`` sink is one table for the whole collection — the tree-path + columns keep the leaves apart — so the writer must outlive the leaf that + created it and must not be closed by it. Resolution stays deferred to + the first output chunk (:func:`get_writer`), and :meth:`result` answers + nothing until the walk closes the sink itself. + + Parameters + ---------- + out : str, Path or writer + The out spec to resolve on the first chunk written. + """ + + def __init__(self, out): + self.out = out + self.writer = None + + def write(self, chunk): + """Write one chunk, resolving the underlying writer on the first.""" + if self.writer is None: + self.writer = get_writer(self.out, chunk) + self.writer.write(chunk) + + def result(self): + """Answer nothing: a shared sink is closed by the walk, not by a leaf.""" + return + + def close(self): + """Close the underlying writer and return its result.""" + if self.writer is None: + return None + result = self.writer.result() + _close_if_owned(self.writer, self.out) + return result + + +class _CollectionSink: + """ + The `out` spec of a collection walk, resolved per leaf. + + Three rules, one per kind of destination: + + - ``None`` accumulates each leaf in memory, so the walk answers with the + tree the eager call returns; + - a single file, a URL or a ready writer instance is **shared** by every + leaf — one table, one archive, one socket for the whole collection, + the tree-path columns each chunk already carries keeping the rows + apart — and the walk answers with that one result; + - a directory **fans out**: one subdirectory per leaf, mirroring the + tree path, since a directory of netcdf chunks describes one stream. + + Parameters + ---------- + out : str, Path, writer or None + The out spec given to :func:`process`. + """ + + def __init__(self, out): + self.out = out + self.shared = None + if out is None: + pass + elif hasattr(out, "write") and hasattr(out, "result"): + self.shared = _SharedWriter(out) + elif isinstance(out, (str, Path)): + spec = str(out) + if re.match(r"[a-z0-9+.-]+://", spec) or Path(spec).suffix: + self.shared = _SharedWriter(out) + else: + raise TypeError(f"cannot infer a writer from `out` of type {type(out)}") + + def spec(self, where): + """Return the out spec of the leaf reached by the keys *where*.""" + if self.shared is not None: + return self.shared + if self.out is None: + return None + # keyed on the keys themselves, not on the annotation path: an + # unnamed level contributes no column but still needs its own + # directory, or its leaves would overwrite one another + return os.path.join(str(self.out), *(str(key) for key in where)) + + def result(self, tree): + """Return the walk's answer, given its walked *tree* of leaf results.""" + if self.shared is not None: + return self.shared.close() + return tree + + +def _announce_realtime_seam(previous, chunk, chunk_dim): + """ + Warn when a realtime chunk arrives discontinuous with the previous one. + + Returns the seam information of *chunk*, to pass back on the next call. + """ + coord = chunk.coords.get(chunk_dim, None) + if not isinstance(coord, AxisCoordinate) or coord.empty: + return previous + info = { + "start": coord.start, + "end": coord.end, + "delta": coord.get_sampling_interval(cast=False), + "tolerance": parse_scalar_delta( + getattr(coord, "tolerance", None), coord.dtype, default_zero=True + ), + } + if previous is not None and previous["delta"] is not None: + tolerance = max(previous["tolerance"], info["tolerance"]) + jump = info["start"] - (previous["end"] + previous["delta"]) + if jump > tolerance: + warnings.warn( + f"realtime source has a discontinuity along {chunk_dim!r} at " + f"{info['start']}; state is flushed and reset", + UserWarning, + stacklevel=2, + ) + if info["delta"] is None and previous is not None: + info["delta"] = previous["delta"] + return info + + +def watch(path, engine="xdas"): + """ + Watch a directory for new files, as an unbounded chunk source. + + Sugar over :class:`RealTimeLoader`: every file closed under `path` is + opened with `engine`, loaded, and yielded as a chunk. Realtime is always + *named* — a bare directory path passed to :func:`process` means "process + what is there", never "block forever". + + Parameters + ---------- + path : str or Path + Directory to watch. + engine : str or Engine, optional + Engine used to open arriving files. Defaults to ``"xdas"``. + + Returns + ------- + RealTimeLoader + An unbounded source for :func:`process` / :meth:`Atom.process`. + + Examples + -------- + >>> pipeline.process(xd.watch("/incoming"), out="sds/") # doctest: +SKIP + """ + return RealTimeLoader(path, engine) + + +def get_source(source, chunks=None): + """ + Resolve a :func:`process` input into a chunk source. + + The source contract is a duck-typed protocol: an iterable yielding + chunks, optionally exposing ``chunk_dim``, ``nbytes`` and ``unbounded``. + Anything already satisfying it (loaders, generators, collections) passes + through. An in-memory :class:`DataArray` with no `chunks` is returned as + is, meaning "process eagerly in one call". + + Parameters + ---------- + source : DataArray, str, Path or iterable + See :func:`process`. + chunks : dict or "auto", optional + Chunk sizes for DataArray sources. Defaults to ``"auto"`` for + virtual DataArrays. + + Returns + ------- + DataArray or iterable + The chunk source, or a bare in-memory DataArray for the eager path. + """ + if isinstance(source, (str, Path)): + spec = str(source) + match = re.match(r"(?P[a-z0-9+.-]+)://", spec) + if match: + scheme = match["scheme"] + if scheme not in SOURCE_SCHEMES: + raise ValueError( + f"no source registered for the {scheme!r} URL scheme; " + f"available: {sorted(SOURCE_SCHEMES)}" + ) + return SOURCE_SCHEMES[scheme](spec) + if os.path.isdir(spec): + spec = os.path.join(spec, "*") + source = open_mfdataarray(spec) + if isinstance(source, DataSequence) and all( + isinstance(element, DataArray) and isinstance(element.data, VirtualBackend) + for element in source + ): + # A virtual multi-acquisition collection: stream each run through its + # own loader instead of materializing whole runs as single chunks. + return _ChainSource(source, chunks) + if isinstance(source, DataArray): + if isinstance(chunks, dict): + # A chunk cannot be larger than what it cuts: the last acquisition + # of a sequence is routinely shorter than the size asked for, and + # one chunk is what that means. + chunks = {dim: min(size, source.sizes[dim]) for dim, size in chunks.items()} + if isinstance(source.data, VirtualBackend): + return DataArrayLoader(source, "auto" if chunks is None else chunks) + if chunks is None: + return source + return DataArrayLoader(source, chunks) + if hasattr(source, "__iter__") or hasattr(source, "__next__"): + return source + raise TypeError( + f"cannot use a {type(source).__name__} object as a source: expected a " + "DataArray, a path, directory or glob, a URL, or an iterable of chunks" + ) + + +def get_writer(out, chunk, chunk_dim="time"): + """ + Resolve a :func:`process` output spec into a writer. + + Dispatch happens on *(out spec × first-chunk type)*: the same directory + path means a netcdf chunk store for DataArray chunks and an SDS archive + for Stream chunks. Writer instances (anything with ``write`` and + ``result``) pass through. + + Parameters + ---------- + out : str, Path, dict, writer or None + See :func:`process`. + chunk : object + The first output chunk of the pipeline. + chunk_dim : str, optional + Dimension along which DataArray chunks follow each other; used to + join them, in memory when `out` is ``None`` and virtually when it is + a directory. It is the dimension the *source* was chunked along, + which need not lead the output. + + Returns + ------- + writer + An object with ``write(chunk)`` and ``result()``. + """ + if hasattr(out, "write") and hasattr(out, "result"): + return out + if out is None: + return ResultWriter(chunk_dim) + if not isinstance(out, (str, Path)): + raise TypeError(f"cannot infer a writer from `out` of type {type(out)}") + spec = str(out) + match = re.match(r"(?P[a-z0-9+.-]+)://", spec) + if match: + scheme = match["scheme"] + if scheme not in SINK_SCHEMES: + raise ValueError( + f"no writer registered for the {scheme!r} URL scheme; " + f"available: {sorted(SINK_SCHEMES)}" + ) + return SINK_SCHEMES[scheme](spec) + if isinstance(chunk, DataArray): + if Path(spec).suffix: + raise ValueError( + f"cannot write DataArray chunks to {spec!r}: pass a directory " + "(chunks are stored as netcdf files and virtually " + "concatenated), or a configured writer instance" + ) + return DataArrayWriter(spec, create_dirs=True, dim=chunk_dim) + if isinstance(chunk, pd.DataFrame): + if not spec.endswith(".csv"): + raise ValueError( + f"cannot write DataFrame chunks to {spec!r}: pass a `*.csv` " + "path or a configured writer instance" + ) + return DataFrameWriter(spec, create_dirs=True) + if isinstance(chunk, obspy.Stream): + return StreamWriter(spec, "D") + raise TypeError( + f"no writer known for output chunks of type {type(chunk).__name__}; " + "pass a configured writer instance as `out`" + ) + + +class ResultWriter: + """ + Accumulate output chunks in memory and join them at the end. + + The writer behind ``out=None``: chunks are collected as they arrive and + ``result()`` returns them joined (gap-aware concatenation for DataArrays, + :func:`pandas.concat` for DataFrames, merged :class:`obspy.Stream`). + Accumulation is guarded by the ``"memory_limit"`` configuration entry. + + Parameters + ---------- + chunk_dim : str, optional + Dimension along which DataArray chunks are concatenated. + """ + + def __init__(self, chunk_dim="time"): + self.chunk_dim = chunk_dim + self.chunks = [] + self.nbytes = 0 + + def write(self, chunk): + """Accumulate one chunk, enforcing the in-memory size guard.""" + self.chunks.append(chunk) + self.nbytes += _sizeof(chunk) + limit = config.get("memory_limit") + if self.nbytes > limit: + raise ValueError( + f"the accumulated in-memory result exceeds {_to_human(limit)} " + "(the 'memory_limit' configuration entry): write the output " + "to disk with `out=...`, or raise the limit with " + "`xdas.config.set('memory_limit', ...)`" + ) + + def result(self): + """Return the joined result, or ``None`` if nothing was written.""" + if self.chunks and all(isinstance(c, obspy.Stream) for c in self.chunks): + out = obspy.Stream() + for st in self.chunks: + out += st + return out + return _join_chunks(self.chunks, self.chunk_dim) + + +class _ChainSource: + """Chain per-run loaders over a virtual multi-acquisition collection.""" + + def __init__(self, collection, chunks): + self.loaders = [] + for element in collection: + if isinstance(chunks, dict): + ((dim, size),) = chunks.items() + element_chunks = {dim: min(size, element.sizes[dim])} + else: + element_chunks = "auto" if chunks is None else chunks + self.loaders.append(DataArrayLoader(element, element_chunks)) + + @property + def chunk_dim(self): + """Chunked dimension, taken from the first per-run loader.""" + return self.loaders[0].chunk_dim + + @property + def nbytes(self): + """Total bytes over all runs.""" + return sum(loader.nbytes for loader in self.loaders) + + def __iter__(self): + for loader in self.loaders: + yield from loader + + +def _sizeof(chunk): + """ + Return the in-memory size of an output chunk, in bytes. + + A pipeline emits arrays, tables and streams, and only the first knows + `nbytes`: taking the others as weightless would silence the accumulation + guard exactly where it matters, on the unbounded walk of a collection + into one pick table. + """ + if isinstance(chunk, pd.DataFrame): + return int(chunk.memory_usage(deep=True).sum()) + if hasattr(chunk, "traces"): + return sum(trace.data.nbytes for trace in chunk.traces) + return getattr(chunk, "nbytes", 0) + + +def _to_human(nbytes): + """Format a byte count as a human-readable string.""" + for unit in ("B", "KB", "MB", "GB"): + if nbytes < 1024: + return f"{nbytes:.1f} {unit}" if unit != "B" else f"{nbytes} B" + nbytes /= 1024 + return f"{nbytes:.1f} TB" + + +def _auto_chunks(da): + """ + Derive tile-aligned chunk boundaries from the storage blocking of *da*. + + Returns the chunked dimension and the list of chunk boundaries along it. + The storage blocking is authoritative when there is one (tile extents for + the tiles vtype, per-source extents for stacked HDF5 virtual datasets); + consecutive blocks are merged up to the ``AUTO_CHUNK_NBYTES`` budget. + Sources with no blocking (single files, in-memory arrays) fall back to + fixed-size chunks from the same byte budget. + """ + data = da.data + if isinstance(data, TileArray): + tiling = data.chunks + axis = int(np.argmax([len(extents) for extents in tiling])) + extents = list(tiling[axis]) + elif isinstance(data, VirtualStack): + axis = data.axis + extents = [source.shape[axis] for source in data.sources] + else: + axis, extents = 0, None + dim = da.dims[axis] + size = da.sizes[dim] + nbytes_per_slice = max(da.nbytes // max(size, 1), 1) + target = max(AUTO_CHUNK_NBYTES // nbytes_per_slice, 1) + if extents is None: + step = int(min(target, size)) + divs = list(range(0, size, step)) + [size] + return dim, divs + divs = [0] + accumulated = 0 + for extent in extents: + if accumulated and accumulated + extent > target: + divs.append(divs[-1] + accumulated) + accumulated = 0 + accumulated += extent + # The loop always leaves the last block(s) in the accumulator. + divs.append(divs[-1] + accumulated) + return dim, divs class DataArrayLoader: @@ -81,15 +942,26 @@ class DataArrayLoader: ---------- da : ``DataArray`` The (virtual) DataArray that contains the data to be chunked - chunks : dict + chunks : dict or "auto" The sizes of the chunks along each dimension. Needs to be of the form: ``{"dim": int}``. The key correspond with the dimension (usually "time"), and the value is an integer indicating the size of the chunk (in samples) - along that dimension. + along that dimension. ``"auto"`` aligns chunk boundaries to the storage + blocking of the array (tile extents, per-file extents), merged up to + the ``AUTO_CHUNK_NBYTES`` byte budget. max_buffers : int, default=1 The maximum number of chunks to load into memory at the same time. max_workers : int, default=1 - The maximum number of thread used to load the chunks. + The maximum number of workers used to load the chunks. + pool : {"threads", "processes"}, default="threads" + The kind of workers to load with. Compressed HDF5 decodes under the + global HDF5 lock, so several threads do not decode concurrently and + ``max_workers`` above one only pays off with ``"processes"``: each + worker receives the manifest of its chunk (kilobytes), reads its own + files, and returns the loaded chunk through a shared-memory object + store — zero-copy for the parent, with chunk data arriving + read-only. Requires the optional ``ray`` dependency. See + :func:`get_pool`. Examples -------- @@ -107,55 +979,78 @@ class DataArrayLoader: >>> for chunk in dl: ... process(chunk) # doctest: +SKIP + Decode four chunks at a time, one worker process each + + >>> dl = DataArrayLoader(da, chunks, 4, 4, pool="processes") # doctest: +SKIP + """ - def __init__(self, da, chunks, max_buffers=1, max_workers=1): + def __init__(self, da, chunks, max_buffers=1, max_workers=1, pool="threads"): if not isinstance(da, DataArray): raise TypeError(f"`da` must by a DataArray object, not a {type(da)}") - if not (isinstance(chunks, dict) and len(chunks) == 1): + if isinstance(chunks, str) and chunks == "auto": + chunk_dim, divs = _auto_chunks(da) + chunk_size = None + elif isinstance(chunks, dict) and len(chunks) == 1: + ((chunk_dim, chunk_size),) = chunks.items() + chunk_dim = str(chunk_dim) + chunk_size = int(chunk_size) + if chunk_dim not in da.dims: + raise ValueError( + f"chunking dimension {chunk_dim} not found in `da` " + f"dimensions {da.dims}" + ) + if chunk_size > da.sizes[chunk_dim]: + raise ValueError( + f"chunking size {chunk_size} is greater than `da` " + f"size {da.sizes[chunk_dim]} along dim {chunk_dim}" + ) + size = da.sizes[chunk_dim] + divs = list(range(0, size, chunk_size)) + [size] + else: raise TypeError( "`chunks` must be a dict that maps a unique " - "dimension to a unique size: {'dim': int}" - ) - ((chunk_dim, chunk_size),) = chunks.items() - chunk_dim = str(chunk_dim) - chunk_size = int(chunk_size) - if chunk_dim not in da.dims: - raise ValueError( - f"chunking dimension {chunk_dim} not found in `da` dimensions {da.dims}" - ) - if chunk_size > da.sizes[chunk_dim]: - raise ValueError( - f"chunking size {chunk_size} is greater than `da` " - f"size {da.sizes[chunk_dim]} along dim {chunk_dim}" + "dimension to a unique size ({'dim': int}) or 'auto'" ) self.da = da self.chunk_dim = chunk_dim self.chunk_size = chunk_size + self._divs = divs self.max_buffers = max_buffers self.max_workers = max_workers + self.pool = pool def __len__(self): - div, mod = divmod(self.da.sizes[self.chunk_dim], self.chunk_size) - return div if mod == 0 else div + 1 + return len(self._divs) - 1 - def _get_chunk(self, idx): - start = idx * self.chunk_size - end = (idx + 1) * self.chunk_size + def _select(self, idx): + """Return chunk *idx* as a lazy selection: the manifest, not the data.""" + start = self._divs[idx] + end = self._divs[idx + 1] query = { dim: slice(start, end) if dim == self.chunk_dim else slice(None) for dim in self.da.dims } - return self.da[query].load() + return self.da[query] def __iter__(self): - with ThreadPoolExecutor(self.max_workers) as executor: + # The chunk sizes are known here, so the arena can be cut to fit the + # biggest of them and no chunk ever has to take the pickle path. + largest = int(np.argmax(np.diff(self._divs))) + with get_pool( + self.pool, self.max_workers, self.max_buffers, self._select(largest).nbytes + ) as executor: it = iter(range(len(self))) + def submit(idx): + # The task is the sliced virtual array, so a process worker + # receives kilobytes and reads its own files. + return executor.submit(_load, self._select(idx)) + futures = [] try: for _ in range(self.max_buffers): - futures.append(executor.submit(self._get_chunk, next(it))) + futures.append(submit(next(it))) except StopIteration: pass @@ -164,7 +1059,7 @@ def __iter__(self): result = future.result() try: - futures.append(executor.submit(self._get_chunk, next(it))) + futures.append(submit(next(it))) except StopIteration: pass @@ -189,6 +1084,9 @@ class RealTimeLoader(Observer): :class:`~xdas.io.Engine` instance. Defaults to ``"xdas"``. """ + chunk_dim = "time" + unbounded = True + def __init__(self, path, engine="xdas"): super().__init__() self.path = str(path) if isinstance(path, Path) else path @@ -238,6 +1136,18 @@ class DataArrayWriter: The maximum number of thread used to load the chunks. create_dirs : bool, optional Whether to create parent directories if they do not exist. Default is False. + dim : str, optional + Dimension the chunks follow each other along, used to join them at + :meth:`result`. Defaults to ``"first"``, which is only right when the + chunked dimension leads the output: a pipeline emitting it elsewhere + must name it. + pool : {"threads", "processes"}, default="threads" + The kind of workers to write with. Compression happens under the + global HDF5 lock, so as on the read side several threads do not + compress concurrently; ``"processes"`` sends each chunk to a worker + through the shared-memory object store — one memcpy at memory + bandwidth — in exchange for parallel compression. Requires the + optional ``ray`` dependency. See :func:`get_pool`. Examples -------- @@ -256,7 +1166,14 @@ class DataArrayWriter: """ def __init__( - self, dirpath, encoding=None, max_buffers=1, max_workers=1, create_dirs=False + self, + dirpath, + encoding=None, + max_buffers=1, + max_workers=1, + create_dirs=False, + dim="first", + pool="threads", ): dirpath = str(dirpath) if isinstance(dirpath, Path) else dirpath if create_dirs: @@ -264,10 +1181,12 @@ def __init__( if not os.path.exists(dirpath): raise OSError(f"no directory {dirpath}") self.dirpath = dirpath + self.dim = dim self.encoding = encoding self.max_buffers = max_buffers self.max_workers = max_workers - self._executor = ThreadPoolExecutor(self.max_workers) + self.pool = pool + self._executor = get_pool(pool, self.max_workers, self.max_buffers) self._futures = [] self._results = [] self._count = 0 @@ -279,26 +1198,25 @@ def submit(self, chunk): Parameters ---------- chunk : DataArray - Processed data chunk to persist. + Processed data chunk to persist. Empty chunks are silently + dropped (many flushes produce nothing). """ if not isinstance(chunk, DataArray): raise TypeError(f"`chunk` must by a DataArray object, not a {type(chunk)}") + if chunk.empty: + return if not len(self._futures) < self.max_buffers: future = self._futures.pop(0) result = future.result() self._results.append(result) - self._futures.append(self._executor.submit(self._write, chunk, self._count)) + path = os.path.join(self.dirpath, f"{self._count:09d}") + self._futures.append(self._executor.submit(_dump, chunk, path, self.encoding)) self._count += 1 def write(self, chunk): """Alias for :meth:`submit`.""" return self.submit(chunk) - def _write(self, chunk, count): - path = os.path.join(self.dirpath, f"{count:09d}") - chunk.to_netcdf(path, encoding=self.encoding) - return open_dataarray(path) - def shutdown(self): """Shut down the internal thread pool.""" self._executor.shutdown() @@ -310,7 +1228,7 @@ def result(self): result = future.result() self._results.append(result) self.shutdown() - return concat(self._results) + return concat(self._results, self.dim) class DataFrameWriter: @@ -359,10 +1277,12 @@ def submit(self, df): Parameters ---------- df : pandas.DataFrame - DataFrame chunk to write. + DataFrame chunk to write. Empty frames are silently dropped. """ if not isinstance(df, pd.DataFrame): raise TypeError(f"`df` must by a DataFrame object, not a {type(df)}") + if df.empty: + return if self._future is not None: self._future.result() self._future = self._executor.submit(self._write, df) @@ -372,11 +1292,12 @@ def write(self, df): return self.submit(df) def _write(self, df): - if df is not None: # pragma: no branch - if not os.path.exists(self.path): - df.to_csv(self.path, mode="w", header=True, index=False) - else: - df.to_csv(self.path, mode="a", header=False, index=False) + # A run appends to the table it finds, which is what lets a restarted + # acquisition keep filling the day's file. + if os.path.exists(self.path): + df.to_csv(self.path, mode="a", header=False, index=False) + else: + df.to_csv(self.path, mode="w", header=True, index=False) def shutdown(self): """Shut down the internal thread pool.""" @@ -384,11 +1305,12 @@ def shutdown(self): def result(self): """Flush pending writes and return the full CSV as a :class:`pandas.DataFrame`.""" - self._future.result() + if self._future is not None: + self._future.result() self.shutdown() try: return pd.read_csv(self.path, parse_dates=self.parse_dates) - except pd.errors.EmptyDataError: + except (FileNotFoundError, pd.errors.EmptyDataError): return pd.DataFrame() @@ -546,7 +1468,7 @@ def submit(self, st): Stream chunk to persist. """ if not isinstance(st, obspy.Stream): - raise TypeError(f"`st` must by a DataFrame object, not a {type(st)}") + raise TypeError(f"`st` must be a Stream object, not a {type(st)}") if self._future is not None: self._future.result() self._future = self._executor.submit(self._write, st) @@ -564,6 +1486,11 @@ def shutdown(self): def result(self): """Merge all temporary MiniSEED files and write the final output.""" + if self._future is None: + # A pipeline that emitted nothing leaves nothing to merge; a + # writer passed in by the caller still gets asked for its result. + self.shutdown() + return obspy.Stream() self._future.result() self.shutdown() pattern = f"{self.dirpath}/*_tmp.mseed" @@ -579,7 +1506,140 @@ def result(self): return out -class ZMQPublisher: +class ZMQEndpoint: + """ + Ownership of the ZeroMQ context and socket an endpoint speaks through. + + A socket nobody closes keeps its file descriptor, and its context keeps an + I/O thread, until the process ends. Releasing them is therefore part of the + interface rather than an afterthought: as a context manager where the + endpoint has a scope, with :meth:`close` where it does not, and on garbage + collection for one that is merely dropped — the last of which is a safety + net, not the way to write it, since when it runs is not for the caller + to know. + + Endpoints inheriting this must expose their socket as ``_socket`` and the + context it came from as ``_context``. + """ + + _socket = None + _context = None + + def close(self): + """Close the socket and terminate the context. Closing twice is a no-op.""" + if self._socket is not None: + self._socket.close() + self._socket = None + if self._context is not None: + self._context.term() + self._context = None + + def __del__(self): + self.close() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + +class SubscriptionTracker(ZMQEndpoint): + """ + Subscriber bookkeeping for publishers that own a ``zmq.XPUB`` socket. + + A ZeroMQ publisher silently drops whatever it sends before a subscriber's + subscription has travelled to it — the "slow joiner" problem. Unlike + ``PUB``, an ``XPUB`` socket hands each subscription to the application, and + only once it has been applied to the socket's routing table. Reading one is + therefore proof that the peer is connected and that everything sent from + then on reaches it. + + A real-time publisher never waits: an instrument streams what it measures + whether or not anyone listens, and a subscriber that joins late is meant to + pick the stream up from wherever it lands. So the count is there to be + *observed* — and to be waited on by the one publisher that is not + real-time, the one replaying a recording, where the head of a finite stream + would otherwise be lost to whoever was not connected yet. + + A mixin rather than part of :class:`ZMQPublisher` because the ASN publisher + (:class:`xdas.io.asn.ZMQPublisher`), which speaks the protocol of the + instrument rather than this one, needs exactly the same bookkeeping. + + Publishers mixing this in must expose their bound ``zmq.XPUB`` socket as + ``_socket`` with ``zmq.XPUB_VERBOSE`` set — so that peers subscribing to an + already-subscribed topic are announced too — the context it came from as + ``_context``, their address as ``address``, and initialize + ``_nsubscribers`` to zero. Closing comes with :class:`ZMQEndpoint`. + """ + + @property + def nsubscribers(self): + """The number of currently subscribed peers.""" + self._read_subscriptions(0.0) + return self._nsubscribers + + def wait_for_subscribers(self, count=1, timeout=60.0): + """ + Block until at least *count* peers are subscribed. + + Whatever is published once this returns is guaranteed to reach those + peers. This is for replaying a finite recording, where losing the first + packets to a subscriber that had not connected yet would lose them for + good; a real-time publisher has nothing to wait for and never calls it. + + Parameters + ---------- + count : int, optional + The number of subscribers to wait for. Defaults to one. + timeout : float or None, optional + How many seconds to wait at most. None waits forever. + + Returns + ------- + int + The number of subscribers connected once the wait is over, which + can exceed *count* if several peers joined at once. + + Raises + ------ + TimeoutError + If fewer than *count* subscribers showed up in time. + + """ + deadline = None if timeout is None else time.monotonic() + timeout + while self.nsubscribers < count: + remaining = None if deadline is None else deadline - time.monotonic() + if not self._read_subscriptions(remaining): + raise TimeoutError( + f"got {self._nsubscribers} of the {count} subscriber(s) " + f"expected on {self.address} after {timeout} seconds" + ) + return self._nsubscribers + + def _read_subscriptions(self, timeout): + """ + Fold pending subscription events into the subscriber count. + + Waits up to *timeout* seconds (None waits forever) for a first event, + then folds in whatever else is already queued. Returns whether any + event was read. + """ + wait = None if timeout is None else max(0, round(1000 * timeout)) + received = False + while self._socket.poll(wait, zmq.POLLIN): + # XPUB only ever delivers subscriptions (\x01) and their + # cancellations (\x00), both followed by the topic. + if self._socket.recv().startswith(b"\x01"): + self._nsubscribers += 1 + else: + self._nsubscribers -= 1 + received = True + wait = 0 + return received + + +class ZMQPublisher(SubscriptionTracker): """ A class for publishing DataArray chunks over ZeroMQ. @@ -590,6 +1650,21 @@ class ZMQPublisher: encoding : dict The encoding to use when dumping the DataArrays to bytes. + Attributes + ---------- + nsubscribers : int + The number of currently subscribed peers. + + Methods + ------- + submit(da) + Send a DataArray over ZeroMQ. + wait_for_subscribers(count, timeout) + Blocks until *count* peers are subscribed, so that nothing published + afterwards is dropped. + close() + Release the socket and its context. + Examples -------- >>> import xdas as xd @@ -599,15 +1674,16 @@ class ZMQPublisher: >>> packets = xd.split(xd.testing.dummy(), 10) - We initialize the publisher at a given address + We initialize the publisher at a given address. Used as a context manager, + it releases its socket on the way out. >>> address = f"tcp://localhost:{xd.io.get_free_port()}" - >>> publisher = ZMQPublisher(address) We can then publish the packets - >>> for da in packets: - ... publisher.submit(da) + >>> with ZMQPublisher(address) as publisher: + ... for da in packets: + ... publisher.submit(da) To reduce the size of the packets, we can also specify an encoding @@ -615,17 +1691,24 @@ class ZMQPublisher: >>> address = f"tcp://localhost:{xd.io.get_free_port()}" >>> encoding = {"chunks": (10, 10), **hdf5plugin.Zfp(accuracy=1e-6)} - >>> publisher = ZMQPublisher(address, encoding) - >>> for da in packets: - ... publisher.submit(da) + >>> with ZMQPublisher(address, encoding) as publisher: + ... for da in packets: + ... publisher.submit(da) """ def __init__(self, address, encoding=None): self.address = address self.encoding = encoding + self._nsubscribers = 0 self._context = zmq.Context() - self._socket = self._context.socket(zmq.PUB) + # XPUB publishes exactly like PUB but also reports who subscribes. + self._socket = self._context.socket(zmq.XPUB) + self._socket.setsockopt(zmq.XPUB_VERBOSE, True) + # The greeting is a socket option, not a rendez-vous: the publisher + # waits for nobody, and hands it to each new peer in passing, on the + # next `submit`. A subscriber that receives it knows it is registered. + self._socket.setsockopt(zmq.XPUB_WELCOME_MSG, WELCOME) self._socket.bind(self.address) def submit(self, da): @@ -638,6 +1721,11 @@ def submit(self, da): The DataArray to be sent. """ + # Taking the subscriptions the socket has queued is what greets the + # peers behind them — ZeroMQ holds a welcome message back until the + # application reads the subscription it answers — and what keeps the + # subscriber count current. Neither costs the publisher any waiting. + self._read_subscriptions(0.0) self._socket.send(tobytes(da, self.encoding)) def write(self, da): @@ -649,7 +1737,7 @@ def result(self): return -class ZMQSubscriber: +class ZMQSubscriber(ZMQEndpoint): """ A class for subscribing to DataArray chunks over ZeroMQ. @@ -657,11 +1745,16 @@ class ZMQSubscriber: ---------- address : str The address to connect the subscriber to. + timeout : float or None, optional + How many seconds to wait at most for each packet. None, the default, + waits forever. Methods ------- - submit(da) - Send a DataArray over ZeroMQ. + wait_until_subscribed() + Block until the publisher has registered this subscription. + close() + Release the socket and its context. Examples -------- @@ -675,31 +1768,49 @@ class ZMQSubscriber: >>> da = xd.testing.dummy() >>> packets = xd.split(da, 10) - We then publish the packets asynchronously - >>> address = f"tcp://localhost:{xd.io.get_free_port()}" >>> publisher = ZMQPublisher(address) + A publisher drops what it sends to subscribers it does not know about yet. + Here the packets come from a recording and we want every one of them, so + the replay waits for its audience — a real-time publisher does not, and its + subscribers pick the stream up wherever they land, waiting instead with + :meth:`ZMQSubscriber.wait_until_subscribed`. + >>> def publish(): + ... publisher.wait_for_subscribers() ... for packet in packets: ... publisher.submit(packet) - >>> threading.Thread(target=publish).start() + >>> thread = threading.Thread(target=publish) + >>> thread.start() - Now let's receive the packets + Now let's receive the packets. The subscriber is an infinite iterator, so + we stop it once the whole stream has been received. >>> subscriber = ZMQSubscriber(address) - >>> packets = [] - >>> for n, da in enumerate(subscriber, start=1): - ... packets.append(da) - ... if n == 10: + >>> received = [] + >>> for packet in subscriber: + ... received.append(packet) + ... if len(received) == len(packets): ... break - >>> da = xd.concat(packets) - >>> assert da.equals(da) + >>> assert xd.concat(received).equals(da) + + Both ends hold a socket until they are closed. Where the two cannot be + written as one ``with`` block, as here, closing them by hand does the same. + + >>> thread.join() + >>> subscriber.close() + >>> publisher.close() """ - def __init__(self, address): + chunk_dim = "time" + unbounded = True + + def __init__(self, address, timeout=None): self.address = address + self.timeout = timeout + self._subscribed = False self._context = zmq.Context() self._socket = self._context.socket(zmq.SUB) self._socket.connect(address) @@ -709,8 +1820,48 @@ def __iter__(self): return self def __next__(self): - message = self._socket.recv() - return frombuffer(message) + while True: + message = self._recv() + if message == WELCOME: + # Sent again whenever the socket reconnects, e.g. to a + # publisher that restarted. + self._subscribed = True + else: + return frombuffer(message) + + def wait_until_subscribed(self): + """ + Block until the publisher has registered this subscription. + + A publisher drops what it sends to a peer it does not know about yet, + and a subscriber cannot tell from its own side whether its + subscription has arrived — being connected is not being subscribed. + The publisher answers it with a greeting, in passing as it streams, and + receiving that greeting is proof that nothing published from then on + will be missed. Nothing is asked of the publisher in return, and a + real-time stream is never delayed by anyone. + + It follows that only a publisher that publishes can acknowledge + anybody: waiting on one that has gone quiet — or on an address nothing + is bound to — raises :exc:`TimeoutError` once the subscriber's + ``timeout`` has passed, and waits forever without one. To be sure of + receiving a *recording* whole, it is the publisher that must wait, with + :meth:`ZMQPublisher.wait_for_subscribers`, since nothing a subscriber + does can hold back a replay that has already started. + + Returns immediately once the publisher has greeted this subscriber. + """ + while not self._subscribed: + self._subscribed = self._recv() == WELCOME + + def _recv(self): + if self.timeout is not None and not self._socket.poll( + round(1000 * self.timeout) + ): + raise TimeoutError( + f"no packet received from {self.address} after {self.timeout} seconds" + ) + return self._socket.recv() def tobytes(da, encoding=None): @@ -753,3 +1904,10 @@ def frombuffer(da): with open(path, "wb") as file: file.write(da) return open_dataarray(path).load() + + +SOURCE_SCHEMES = {"tcp": ZMQSubscriber} +"""Registry of URL schemes accepted as sources, scheme → source factory.""" + +SINK_SCHEMES = {"tcp": ZMQPublisher} +"""Registry of URL schemes accepted as sinks, scheme → writer factory.""" diff --git a/xdas/processing/pools.py b/xdas/processing/pools.py new file mode 100644 index 00000000..8b3949ab --- /dev/null +++ b/xdas/processing/pools.py @@ -0,0 +1,591 @@ +""" +The worker pool behind ``pool="processes"``. + +A process pool is the right answer for compressed HDF5 — decompression happens +under the global HDF5 lock, so only separate processes overlap it — but the +loaded chunk crossing back would be pickled through a pipe, which caps out +well below memory bandwidth and spends parent CPU doing it. This module +replaces that pipe with an arena of shared memory: the worker writes the chunk +once into a slot, the parent maps the same pages, and no bytes are serialized +in either direction. + +The arena is a single file in ``/dev/shm``, sliced into fixed-size slots and +reused for the whole run. Reuse is the point: allocating a fresh block per +chunk makes the kernel zero and fault in every page again, which costs more +than it saves. Slots are handed out by the parent alone, so there is no +cross-process allocator and no lock in the data path. + +What crosses the boundary is a :class:`ShmRef` — a shape/dtype placeholder +standing in for the array, exactly the duck type +:class:`~xdas.virtual.VirtualArray` already presents to +:class:`~xdas.core.DataArray`. Anything that is not a plain numeric array (a +virtual array, a DataFrame, an oversized chunk, an arena with no free slot) +simply travels the ordinary pickle path: the transport is an optimization, and +falling back to it is never wrong, only slower. + +Releasing the memory is the other half of the design. Pages die with the last +mapping that holds them, so anything that gets to run Python on the way out — +an interrupt included — frees them at once, and the arena's *name* is unlinked +when the pool shuts down, again by a finalizer at interpreter exit, and swept +at the start of the next run if its owner is gone. Nothing in ``/dev/shm`` +survives two runs. + +A run killed outright is the exception, and not because of the arena: loky's +workers do not notice that their parent has died, so they stay, and what they +still hold mapped stays with them until they are killed too. That is worth +knowing when a scheduler or an out-of-memory killer ends a job, and it is the +same with or without shared memory in the picture. + +All of that rests on POSIX semantics — unlinking a file that is still mapped, +and signal zero as a liveness probe — so the arena is built on POSIX only. +Elsewhere the pool is a plain process pool that pickles its chunks, slower but +identical in what it computes. +""" + +import mmap +import os +import re +import shutil +import tempfile +import threading +import weakref +from uuid import uuid4 + +import numpy as np +from loky import ProcessPoolExecutor, process_executor + +from ..core import DataArray + +SLOT_NBYTES = 256 * 2**20 +"""Default slot size, matching the target chunk size of ``chunks="auto"``.""" + +SHM_FRACTION = 0.25 +"""Share of the free shared-memory filesystem a single arena may claim.""" + +PREFIX = "xdas-shm-" +"""Arena filename prefix, carrying the owner pid so stale ones can be swept.""" + +_ARENAS = {} +"""Mappings held by this process, path → mmap. Populated by :func:`attach`.""" + + +def _directory(): + """Where arenas live: RAM-backed if the platform offers it, else temp.""" + return "/dev/shm" if os.path.isdir("/dev/shm") else tempfile.gettempdir() + + +def _unlink(path): + """Drop an arena's name. The pages go when the last mapping does.""" + try: + os.unlink(path) + except OSError: + pass + + +def _discard(path): + """ + Let go of an arena entirely: its name, and this process's handle on it. + + Dropping the handle is what actually frees the pages, and it has to be + done for a pool that was never shut down as much as for one that was, or a + process that opens pools in a loop would hold every arena it ever made. + The mapping itself is not closed: chunks still in the caller's hands are + views on it, and they keep it alive on their own until they are collected. + """ + _ARENAS.pop(path, None) + _unlink(path) + + +def sweep(): + """ + Unlink arenas whose owner process is gone. + + A pool unlinks its own arena on shutdown and at interpreter exit, so this + only ever finds arenas orphaned by a kill that ran no Python at all. + """ + if os.name != "posix": # pragma: no cover - signal zero kills on Windows + return + directory = _directory() + try: + names = os.listdir(directory) + except OSError: # pragma: no cover - unreadable shm directory + return + for name in names: + match = re.fullmatch(rf"{PREFIX}(\d+)-[0-9a-f]+", name) + if match is None: + continue + try: + os.kill(int(match.group(1)), 0) + except ProcessLookupError: + _unlink(os.path.join(directory, name)) + except PermissionError: # pragma: no cover - alive, another user + pass + + +def attach(path, size): + """ + Map an existing arena into this process. + + Runs in the parent when the arena is created, and in every worker through + :func:`_init_worker`, so a worker is mapped before it is handed any task. + + Parameters + ---------- + path : str + Filesystem path of the arena. + size : int + Its size in bytes, as passed to :meth:`Arena.__init__`. + """ + if path not in _ARENAS: + fd = os.open(path, os.O_RDWR) + try: + _ARENAS[path] = mmap.mmap(fd, size) + finally: + os.close(fd) + + +def _init_worker(path, size): + """ + Prepare a worker: map the arena, and stop loky reading it as a leak. + + A worker filling a slot sees its own resident memory grow by the size of + the chunk it wrote, because shared pages are counted like any other. Loky + watches that figure and recycles a worker that grows by more than 300 MB, + so chunks of any real size had workers quitting and respawning mid-run. + The arena is bounded and reused, so it is not what that check is for; it + is discounted here, leaving the original allowance on top of it to catch + leaks that are real. + + Parameters + ---------- + path : str + Filesystem path of the arena. + size : int + Its size in bytes. + """ + if hasattr(process_executor, "_MAX_MEMORY_LEAK_SIZE"): # pragma: no branch + process_executor._MAX_MEMORY_LEAK_SIZE += size + attach(path, size) + + +class ShmRef: + """ + A placeholder for an array parked in an arena slot. + + Presents the ``shape``/``dtype``/``ndim`` surface + :class:`~xdas.core.DataArray` asks of its data, so a chunk can carry one + across the process boundary in place of its array and be rebuilt on the + far side by whoever has the arena mapped. + + Parameters + ---------- + path : str + Arena the slot belongs to. + offset : int + Byte offset of the slot within the arena. + shape : tuple of int + Shape of the parked array. + dtype : str + Its dtype, as a ``numpy`` type string. + """ + + def __init__(self, path, offset, shape, dtype): + self.path = path + self.offset = offset + self.shape = tuple(shape) + self.dtype = dtype + + @property + def ndim(self): + """Number of dimensions of the parked array.""" + return len(self.shape) + + @property + def nbytes(self): + """Size of the parked array in bytes.""" + return int(np.prod(self.shape)) * np.dtype(self.dtype).itemsize + + def __array__(self, dtype=None, copy=None): + """Resolve to the shared pages. Requires this process to be mapped.""" + out = view(self) + return out.astype(dtype) if dtype is not None else out + + def __repr__(self): + return f"ShmRef(shape={self.shape}, dtype={self.dtype})" + + +def view(ref): + """ + Return the array a :class:`ShmRef` points at, without copying. + + Parameters + ---------- + ref : ShmRef + Reference to resolve. + + Returns + ------- + numpy.ndarray + A view on the shared pages. + """ + arena = _ARENAS.get(ref.path) + if arena is None: + raise RuntimeError(f"arena {ref.path} is not mapped in process {os.getpid()}") + return np.ndarray(ref.shape, np.dtype(ref.dtype), arena, ref.offset) + + +def _offloadable(value, capacity): + """Whether *value* is a chunk whose array can be parked in a slot.""" + if not isinstance(value, DataArray): + return False + data = value.data + return ( + isinstance(data, np.ndarray) + and not data.dtype.hasobject + and 0 < data.nbytes <= capacity + ) + + +def _park(value, path, offset): + """Copy *value*'s array into a slot and return it carrying a reference.""" + ref = ShmRef(path, offset, value.shape, value.data.dtype.str) + np.copyto(view(ref), value.data) + return value.copy(data=ref) + + +def _resolve(value): + """Rebuild a chunk whose array was parked by the other side.""" + if isinstance(value, DataArray) and isinstance(value.data, ShmRef): + return value.copy(data=view(value.data)) + return value + + +def _run(fn, outbox, args, kwargs): + """ + Run one task in a worker, translating chunks in and out of the arena. + + Arguments arrive as references and are resolved to views; the result is + parked in *outbox* when there is a slot for it and it is worth parking. + """ + args = tuple(_resolve(arg) for arg in args) + result = fn(*args, **kwargs) + if outbox is not None: + path, offset, capacity = outbox + if _offloadable(result, capacity): + result = _park(result, path, offset) + return result + + +class Arena: + """ + A file of shared memory, sliced into reusable fixed-size slots. + + Slots are reserved and released by the parent only, so the free list needs + no cross-process synchronization — just a lock, because releases can come + from a finalizer running on whichever thread the garbage collector was on. + + Parameters + ---------- + nslots : int + Number of slots. + slot_nbytes : int + Size of each slot, and so the largest array that can be parked. + """ + + def __init__(self, nslots, slot_nbytes): + sweep() + self.nslots = nslots + self.slot_nbytes = slot_nbytes + self.size = nslots * slot_nbytes + self.path = os.path.join( + _directory(), f"{PREFIX}{os.getpid()}-{uuid4().hex[:8]}" + ) + fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600) + try: + os.ftruncate(fd, self.size) + finally: + os.close(fd) + # Unlink on garbage collection and at interpreter exit, whichever + # comes first; the sweep covers the case where neither runs. + self._finalizer = weakref.finalize(self, _discard, self.path) + try: + attach(self.path, self.size) + except OSError: # pragma: no cover - mapping refused + self._finalizer() + raise + self._free = list(range(nslots)) + self._lock = threading.Lock() + + @classmethod + def create(cls, nslots, slot_nbytes): + """ + Build an arena, or return ``None`` if shared memory cannot host it. + + Sizing is capped at :data:`SHM_FRACTION` of what the filesystem has + free: the file is sparse, so an oversized arena costs nothing until it + is used, but writing into pages the filesystem cannot back would take + the worker down with a bus error rather than an exception. + + Parameters + ---------- + nslots : int + Desired number of slots; reduced to fit if need be. + slot_nbytes : int + Size of each slot. + + Returns + ------- + Arena or None + ``None`` when fewer than two slots fit, when the platform is not + POSIX, or when the filesystem cannot be read. + """ + if os.name != "posix": # pragma: no cover - no unlink-while-mapped + return None + try: + free = shutil.disk_usage(_directory()).free + except OSError: # pragma: no cover - unreadable directory + return None + nslots = min(nslots, int(free * SHM_FRACTION // slot_nbytes)) + if nslots < 2: + return None + try: + return cls(nslots, slot_nbytes) + except OSError: # pragma: no cover - creation refused + return None + + def reserve(self): + """ + Take a free slot, or ``None`` if they are all out on loan. + + Returns + ------- + tuple or None + The ``(path, offset, capacity)`` an outbox is described by. + """ + with self._lock: + if not self._free: + return None + index = self._free.pop() + return self.path, index * self.slot_nbytes, self.slot_nbytes + + def release(self, offset): + """ + Hand a slot back to the free list. + + Parameters + ---------- + offset : int + Byte offset of the slot being released. + """ + with self._lock: + self._free.append(offset // self.slot_nbytes) + + def close(self): + """ + Unlink the arena. Mappings already handed out stay valid. + + The pages are freed by the kernel once the last mapping goes, which is + when the last chunk still pointing into them is collected. + """ + self._finalizer() + + +class ProcessFuture: + """ + The handle :class:`ProcessPool` hands out for a submitted task. + + Carries the same ``result()`` surface the loader and writer use, adding + the step where a parked result is turned back into an array — a view on + the shared pages, not a copy. + + Parameters + ---------- + pool : ProcessPool + Pool that submitted the task. + inner : concurrent.futures.Future + The underlying worker future. + outbox : tuple or None + Slot reserved for the result, if any. + """ + + def __init__(self, pool, inner, outbox): + self._pool = pool + self._inner = inner + self._outbox = outbox + self._value = None + self._collected = False + + def result(self, timeout=None): + """ + Wait for the task and return its result. + + Returns + ------- + object + Whatever the task returned. A chunk that travelled through the + arena comes back read-only: its data is the shared page range + itself, which the pool recycles once nothing refers to it any + more. + """ + if self._collected: + return self._value + try: + value = self._inner.result(timeout) + except BaseException: + # Only a task that is over has stopped writing into its slot; one + # that merely outran `timeout` still owns those pages. + if self._inner.done(): + self._release() + raise + self._value = self._pool._collect(value, self._outbox) + self._collected = True + self._outbox = None + return self._value + + def cancel(self): + """Try to cancel the task, releasing its slot if that worked.""" + cancelled = self._inner.cancel() + if cancelled: + self._release() + return cancelled + + def done(self): + """Whether the task has finished, one way or another.""" + return self._inner.done() + + def _release(self): + """Give back the result slot of a task that produced nothing.""" + if self._outbox is not None and self._pool._arena is not None: + self._pool._arena.release(self._outbox[1]) + self._outbox = None + + +class ProcessPool: + """ + A process pool whose chunks cross through shared memory. + + Quacks like a :class:`~concurrent.futures.Executor` as far as the loader + and the writer need — ``submit``/``shutdown``/context manager — and keeps + the per-process HDF5 locks a process pool is chosen for. What it adds is + the data path: a chunk returned by a worker is written once into an arena + slot and mapped by the parent, and a chunk sent to a worker is staged the + same way. Neither direction serializes the array. + + The price of mapping rather than copying is immutability: chunk data + arrives read-only, since the pages are the pool's to recycle. Atoms honor + this already by allocating their outputs. + + Whatever cannot be parked — a chunk larger than a slot, an arena with no + free slot, a result that is not an array — travels the ordinary pickle + path instead, which is slower but always correct. Should shared memory be + unavailable altogether, the pool degrades to exactly that. + + Parameters + ---------- + max_workers : int + Number of worker processes. + max_buffers : int, optional + Chunks the caller keeps in flight, which is what the arena is sized + from. Default is 1. + slot_nbytes : int, optional + Largest chunk the arena can carry. Defaults to :data:`SLOT_NBYTES`. + """ + + def __init__(self, max_workers, max_buffers=1, slot_nbytes=None): + if slot_nbytes is None: + slot_nbytes = SLOT_NBYTES + # Sized well past what streaming needs -- one slot per chunk in + # flight, plus the one being consumed and the one whose release is + # still pending. Spare slots are sparse pages that cost nothing until + # they are written to, whereas running out quietly costs a pickle. + nslots = 2 * max_buffers + max_workers + 4 + self._arena = Arena.create(nslots, slot_nbytes) + if self._arena is None: + self._executor = ProcessPoolExecutor(max_workers) + else: + self._executor = ProcessPoolExecutor( + max_workers, + initializer=_init_worker, + initargs=(self._arena.path, self._arena.size), + ) + + def submit(self, fn, /, *args, **kwargs): + """ + Schedule ``fn(*args, **kwargs)`` on a worker. + + Parameters + ---------- + fn : callable + The unit of work. Chunk arguments are staged into the arena and + resolved back to arrays inside the worker. + + Returns + ------- + ProcessFuture + A ``result()``-able handle. + """ + outbox = None + staged = [] + if self._arena is not None: + args = tuple(self._stage(arg, staged) for arg in args) + outbox = self._arena.reserve() + inner = self._executor.submit(_run, fn, outbox, args, kwargs) + if staged: + # The worker reads a staged argument while it runs, so the slot + # comes back only once the task is over -- whatever its outcome. + inner.add_done_callback(lambda _: self._release(staged)) + return ProcessFuture(self, inner, outbox) + + def shutdown(self, wait=True): + """ + Shut the workers down and unlink the arena. + + Chunks still held by the caller keep their pages: the mapping outlives + the name, and the kernel frees it when the last one is dropped. + + Parameters + ---------- + wait : bool, optional + Whether to block until running tasks complete. + """ + self._executor.shutdown(wait=wait) + if self._arena is not None: + self._arena.close() + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self.shutdown() + + def _stage(self, value, staged): + """Park an outgoing chunk in a slot, if it is one and one is free.""" + if not _offloadable(value, self._arena.slot_nbytes): + return value + slot = self._arena.reserve() + if slot is None: + return value + path, offset, _ = slot + staged.append(offset) + return _park(value, path, offset) + + def _release(self, offsets): + """Return staged argument slots to the free list.""" + for offset in offsets: + self._arena.release(offset) + + def _collect(self, value, outbox): + """Turn a parked result into a read-only view that recycles itself.""" + if outbox is None: + return value + offset = outbox[1] + if not (isinstance(value, DataArray) and isinstance(value.data, ShmRef)): + self._arena.release(offset) + return value + data = view(value.data) + data.flags.writeable = False + # The slot belongs to the chunk now: it goes back to the free list + # when the array is collected, which is exactly when nothing can read + # those pages any more. + weakref.finalize(data, self._arena.release, offset) + return value.copy(data=data) diff --git a/xdas/testing.py b/xdas/testing.py index f03f13db..fe671b7f 100644 --- a/xdas/testing.py +++ b/xdas/testing.py @@ -1,9 +1,12 @@ """Test utilities for xdas.""" +import warnings + import numpy as np +import pandas as pd from .coordinates import Coordinate -from .core import DataArray +from .core import DataArray, DataCollection, concat, split def dummy( @@ -78,3 +81,264 @@ def dummy( coords[dim] = Coordinate[ctype].from_block(start, size, s, dim=dim) return DataArray(data=data, coords=coords) + + +def assert_chunk_invariant( + pipeline, da, chunks, rtol=1e-7, atol=0.0, coord_atol=0, gaps=None, cuts=1 +): + """ + Assert that *pipeline* answers the same however the stream is cut. + + Chunk-safety is a claim an atom makes about itself; this is the evidence. + The pipeline is run once eagerly on the whole array and once streamed + chunk by chunk with the given `chunks`, and the two results are required + to match — values, coordinates and all. On top of that single split, the + invariant is quantified over *cuts*: the same stream is re-chunked at + other sizes (whose boundaries fall elsewhere, including across any gap) + and every cutting must answer the same. With `gaps`, discontinuities are + injected into the input first, so seam resets are exercised at chunk + boundaries that do not line up with them. + + Parameters + ---------- + pipeline : Atom + The pipeline to check. It has to be an atom: a bare callable cannot + be streamed, since the chunked path hands each chunk a ``chunk_dim``. + Wrap one with :class:`xdas.atoms.Partial`. + da : DataArray + The input to run it on. + chunks : dict + Chunk sizes for the streamed run, e.g. ``{"time": 100}``. + rtol, atol : float, optional + Tolerances for the value comparison, as in + :func:`numpy.testing.assert_allclose`. + coord_atol : int or float, optional + Tolerance on the dimension coordinates, in their own units + (nanoseconds for datetime axes). Zero — the default — demands an + exact match. Rational resampling to a rate that is not an exact + number of nanoseconds reconstructs its output grid segment by + segment, so eager and chunked coordinates may differ by a nanosecond + with bit-identical values; that is what this admits, explicitly. + gaps : int or sequence of int, optional + Inject gaps into `da` along the chunked dimension before comparing: + an int places that many evenly spaced gaps, a sequence gives the + sample indices where each gap starts. Each gap drops a twentieth of + the record (at least one sample), which is well beyond any jitter + tolerance, so the seams judge them as real discontinuities. + cuts : int or sequence of dict, optional + How many extra cuttings to check beyond `chunks` (cut-invariance): + the result must be a function of *which samples were processed*, + never of where the stream was cut. An int derives that many + alternative chunk sizes from `chunks` (each smaller and coprime-ish, + so the boundaries land elsewhere); a sequence gives explicit + ``chunks``-style dicts. ``0`` restores the single-split check. + + Raises + ------ + AssertionError + If any two runs disagree, in shape, in coordinates or in values. + + Notes + ----- + Pick tables are compared as *sets* of rows: eager processing walks the + whole record lane by lane while chunked processing walks chunk by chunk, + so the rows come out in a different order with the same content, and it + is the content that the invariant is about. + + Examples + -------- + >>> import numpy as np + >>> import xdas as xd + >>> from xdas.atoms import Partial + >>> da = xd.testing.dummy() + >>> xd.testing.assert_chunk_invariant(Partial(np.square), da, {"time": 25}) + + Injecting gaps exercises the seam resets as well: + + >>> xd.testing.assert_chunk_invariant( + ... xd.filter(..., (None, 10.0), dim="time"), da, {"time": 25}, gaps=2 + ... ) + + A pipeline that is *not* chunk-invariant says so rather than passing: + + >>> def normalize(da): + ... return da / np.std(da.values) + >>> xd.testing.assert_chunk_invariant(Partial(normalize), da, {"time": 25}) + ... # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + AssertionError + + """ + ((dim, size),) = chunks.items() + with warnings.catch_warnings(): + if gaps is not None: + # The gaps are this function's own doing: the split announcement + # would only tell the test suite what it already decided. + warnings.filterwarnings("ignore", message="source starting at") + da = inject_gaps(da, dim, gaps) + size = min(size, da.sizes[dim]) + chunks = {dim: size} + pipeline.reset() + eager = pipeline(da) + pipeline.reset() + chunked = _stream(pipeline, da, chunks) + _assert_same(eager, chunked, rtol, atol, coord_atol, "result") + for cut in _cuttings(cuts, dim, size): + pipeline.reset() + recut = _stream(pipeline, da, cut) + _assert_same(eager, recut, rtol, atol, coord_atol, f"result cut as {cut}") + + +def _stream(pipeline, da, chunks): + """Run *pipeline* chunk by chunk over *da* and join the outputs.""" + ((dim, size),) = chunks.items() + indices = list(range(size, da.sizes[dim], size)) + pieces = split(da, indices, dim) if indices else [da] + outs = list(pipeline.iter_chunks(pieces, chunk_dim=dim)) + return pipeline._join(outs, dim) + + +def _cuttings(cuts, dim, size): + """Derive the extra chunk sizes the cut-invariance pass runs with.""" + if not isinstance(cuts, int): + return list(cuts) + sizes = [] + current = size + for _ in range(cuts): + # Not a divisor of the previous size, so the boundaries move. + current = current // 2 + 1 if current > 1 else current + 1 + if current == size or current < 1: + break + sizes.append(current) + return [{dim: value} for value in sizes] + + +def inject_gaps(da, dim, gaps): + """ + Return *da* with gaps injected along *dim*, for seam testing. + + Each gap drops a twentieth of the record (at least one sample), which is + well beyond any jitter tolerance, so downstream seam judgment sees a real + discontinuity in the coordinate. + + Parameters + ---------- + da : DataArray + The array to make gappy. + dim : str + The dimension along which to drop samples. + gaps : int or sequence of int + An int places that many evenly spaced gaps; a sequence gives the + sample indices where each gap starts. + + Returns + ------- + DataArray + The gappy array: same values minus the dropped spans, with the gaps + kept in the coordinates. + + Examples + -------- + >>> import xdas as xd + >>> da = xd.testing.dummy() + >>> gappy = xd.testing.inject_gaps(da, "time", 2) + >>> gappy.sizes["time"] + 90 + """ + size = da.sizes[dim] + if isinstance(gaps, int): + starts = [round((index + 1) * size / (gaps + 1)) for index in range(gaps)] + else: + starts = sorted(int(start) for start in gaps) + width = max(1, size // 20) + pieces = [] + previous = 0 + for start in starts: + pieces.append(da.isel({dim: slice(previous, start)})) + previous = min(start + width, size) + pieces.append(da.isel({dim: slice(previous, None)})) + pieces = [piece for piece in pieces if piece.sizes[dim]] + if len(pieces) < 2: + raise ValueError( + f"cannot inject {gaps!r} gaps into a record of {size} samples " + f"along {dim!r}: the gaps leave less than two pieces" + ) + return concat(pieces, dim) + + +def _assert_same(eager, chunked, rtol, atol, coord_atol, path): + """Compare two pipeline outputs of any supported chunk type.""" + if isinstance(eager, DataArray): + if not isinstance(chunked, DataArray): + raise AssertionError( + f"{path}: eager gave a DataArray, chunked gave a " + f"{type(chunked).__name__} — the chunked run did not join into " + "one array" + ) + if eager.dims != chunked.dims: + raise AssertionError( + f"{path}: dims differ, {eager.dims} eager vs {chunked.dims} chunked" + ) + if eager.shape != chunked.shape: + raise AssertionError( + f"{path}: shape differs, {eager.shape} eager vs {chunked.shape} chunked" + ) + np.testing.assert_allclose( + chunked.values, eager.values, rtol=rtol, atol=atol, err_msg=path + ) + for dim in eager.dims: + if dim in eager.coords: + _assert_same_coord( + chunked.coords[dim], eager.coords[dim], coord_atol, path, dim + ) + elif isinstance(eager, pd.DataFrame): + _assert_same_frame(eager, chunked, rtol, atol, path) + elif isinstance(eager, (list, DataCollection)): + if len(eager) != len(chunked): + raise AssertionError( + f"{path}: {len(eager)} chunks eager vs {len(chunked)} chunked" + ) + for index, (left, right) in enumerate(zip(eager, chunked)): + _assert_same(left, right, rtol, atol, coord_atol, f"{path}[{index}]") + else: + np.testing.assert_allclose(chunked, eager, rtol=rtol, atol=atol, err_msg=path) + + +def _assert_same_coord(chunked, eager, coord_atol, path, dim): + """Compare one dimension coordinate, within *coord_atol* of its own units.""" + left = np.asarray(chunked.values) + right = np.asarray(eager.values) + message = f"{path}: the {dim!r} coordinate differs" + if not coord_atol: + np.testing.assert_array_equal(left, right, err_msg=message) + elif np.issubdtype(right.dtype, np.datetime64): + # Stay in integers: an epoch nanosecond does not survive float64, + # whose ULP up there is a few hundred nanoseconds — wider than any + # tolerance worth expressing. + difference = np.abs((left - right).astype("int64")) + np.testing.assert_array_less(difference, coord_atol + 1, err_msg=message) + else: + np.testing.assert_allclose( + left, right, rtol=0, atol=coord_atol, err_msg=message + ) + + +def _assert_same_frame(eager, chunked, rtol, atol, path): + """Compare two pick tables as sets of rows (see the Notes of the caller).""" + if not isinstance(chunked, pd.DataFrame): + raise AssertionError( + f"{path}: eager gave a DataFrame, chunked gave a {type(chunked).__name__}" + ) + if list(eager.columns) != list(chunked.columns): + raise AssertionError( + f"{path}: columns differ, {list(eager.columns)} eager vs " + f"{list(chunked.columns)} chunked" + ) + columns = list(eager.columns) + pd.testing.assert_frame_equal( + chunked.sort_values(columns).reset_index(drop=True), + eager.sort_values(columns).reset_index(drop=True), + rtol=rtol, + atol=atol, + ) diff --git a/xdas/trigger.py b/xdas/trigger.py index 4f07fbac..1ac66f74 100644 --- a/xdas/trigger.py +++ b/xdas/trigger.py @@ -1,213 +1,21 @@ """ -Threshold-based triggering atom :class:`Trigger`. +Compatibility home of the threshold trigger. -Detects phase arrivals in :class:`DataArray` objects using an on/off -mechanism. +:class:`Trigger` now lives in :mod:`xdas.atoms.detect` and is re-exported +here unchanged. :func:`find_picks` is the historical functional form and +stays as is. """ import numpy as np import pandas as pd from numba import njit -from .atoms.core import Atom, State, atomized +from .atoms.core import atomized +from .atoms.detect import Trigger from .coordinates import Coordinate from .core import concat_coords - -class Trigger(Atom): - """ - Find picks in a data array along a given axis based on a given threshold. - - The pick findings use a triggering mechanism where triggers are turned on and off - based on the threshold crossings. The trigger off threshold is half of the trigger - on threshold. Picks are determined by finding the maximum value on each triggered - region. - - Parameters - ---------- - thresh : float - The threshold value for picking. - dim : str, optional - The dimension along which to find picks. Defaults to "last". - - Notes - ----- - For more details see the documentation of the `initialize` and `call` methods. - - Examples - -------- - >>> import numpy as np - >>> import xdas as xd - >>> from xdas.atoms import Trigger - - Use case: - - >>> cft = xd.DataArray( - ... data=[[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2]], - ... coords={ - ... "space": [0.0], - ... "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0]}, - ... }, - ... ) - - Chunked processing using atomic processing: - - >>> atom = Trigger(thresh=0.5, dim="time") - >>> chunks = xd.split(cft, 3, dim="time") - >>> result = [] - >>> for chunk in chunks: - ... picks = atom(chunk, chunk_dim="time") - ... result.append(picks) - >>> result = pd.concat(result, ignore_index=True) - >>> result - space time value - 0 0.0 2.0 0.9 - 1 0.0 7.0 0.7 - - """ - - def __init__(self, thresh, dim="last"): - super().__init__() - - # parameters - self.thresh_on = float(thresh) - self.thresh_off = float(thresh) / 2.0 - self.dim = str(dim) - - # states - self.axis = State(...) - self.shape = State(...) - self.status = State(...) - self.index = State(...) - self.value = State(...) - self.offset = State(...) - self.coord = State(...) - - def initialize(self, cft, **flags): - """ - Initialize the trigger with the following states. - - - "axis": An integer indicating the axis number of the dimension along which to - find picks. - - "shape": A tuple indicating the unravel shape of the lanes along wigh the - the picks will be found. - - "status": A boolean array indicating the trigger status for each lane. - - "index": An integer array indicating the index of the last triggered value - for each lane. - - "value": A float array indicating the value of the last triggered value for - each lane. - - "offset": An integer indicating the offset of the chunk. - - "coord": An InterpCoordinate containing coordinate information along 'dim' up - to the last processed chunk. - - - Parameters - ---------- - cft : DataArray - The characteristic function where picks must be found. - **flags - Optional flags. - - """ - self.axis = State(cft.get_axis_num(self.dim)) - self.shape = State(cft.shape[: self.axis] + cft.shape[self.axis + 1 :]) - self.status = State(np.zeros(self.shape, dtype=bool)) - self.index = State(np.zeros(self.shape, dtype=int)) - self.value = State(np.zeros(self.shape, dtype=float)) - self.offset = State(0) - self.coord = State(Coordinate({"tie_indices": [], "tie_values": []}, self.dim)) - - def call(self, cft, **flags): - """ - Call the trigger. - - Parameters - ---------- - cft : DataArray - The characteristic function where picks must be found. - **flags - Optional flags. - - Returns - ------- - picks: DataFrame - A DataFrame containing the pick coordinates and their corresponding values. - - Notes - ----- - In the trigger does not turn off at the end of the array, the last pick will - not be found. This can be fixed by appending a zero to the end of the array. - - """ - data = np.asarray(cft.values, dtype=float) - values, coords = self._call_numeric(data) - self.coord = concat_coords([self.coord, cft.coords[self.dim]], tolerance=None) - - picks = {} - for axis, dim in enumerate(cft.dims): - if dim == self.dim: - picks[dim] = self.coord[coords[axis]].values - else: - picks[dim] = cft.coords[dim][coords[axis]].values - picks["value"] = values - return pd.DataFrame(picks) - - def _call_numeric(self, data): - """ - Find picks in a N-dimensional array along a given axis based on a given threshold. - - The pick findings use a triggering mechanism where triggers are turned on and off - based on the threshold crossings. The trigger off threshold is half of the trigger - on threshold. Picks are determined by finding the maximum value on each triggered - region. - - Parameters - ---------- - data : DataArray - The characteristic function where picks must be found. - - Returns - ------- - coords : tuple of 1d ndarray - A tuple containing the coordinates of the picks. - values : 1d ndarray - The values of the picks. - - Notes - ----- - If the trigger does not turn off at the end of the array, the last pick will \ - not be found. This can be fixed by appending a zero to the end of the array. - - """ - data = np.moveaxis(data, self.axis, -1) - length = data.shape[-1] - - # ravel additional axes into a unique lanes axis - data = np.reshape(data, (-1, data.shape[-1])) - status_view = np.reshape(self.status, (-1,)) - index_view = np.reshape(self.index, (-1,)) - value_view = np.reshape(self.value, (-1,)) - - lanes, indices, values = _trigger( - data, - self.thresh_on, - self.thresh_off, - status_view, - index_view, - value_view, - self.offset, - ) - self.offset += length - - # unravel lanes indices - if self.shape: - coords = np.unravel_index(lanes, self.shape) - else: - coords = () - - # insert found indices into the original axis position - coords = coords[: self.axis] + (indices,) + coords[self.axis :] - return values, coords +__all__ = ["Trigger", "find_picks"] @atomized @@ -265,7 +73,7 @@ def find_picks(cft, thresh, dim="last", state_dict=None): # TODO: state_dict => ... data=[[0.0, 0.1, 0.9, 0.8, 0.2, 0.1, 0.6, 0.7, 0.3, 0.2]], ... coords={ ... "space": [0.0], - ... "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0]}, + ... "time": {"tie_indices": [0, 9], "tie_values": [0.0, 9.0], "sampling_interval": 1.0}, ... }, ... ) diff --git a/xdas/virtual/hdf5.py b/xdas/virtual/hdf5.py index 5993b38d..515b97ea 100644 --- a/xdas/virtual/hdf5.py +++ b/xdas/virtual/hdf5.py @@ -365,7 +365,7 @@ class VirtualSource(VirtualArray): sliced to indicate which regions should be used. Sliced VirtualSource eventually can be assigned to a VirtualLayout to - Best practive is to pass it a `h5py.Dataset` obtain destructuring a `h5py.File`. + Best practice is to pass it a `h5py.Dataset` obtain destructuring a `h5py.File`. Otherwise the exact filename, dataset name, shape and dtype must be passed. The data can be accessed using `numpy.asarray` or the `__array__` special method. @@ -476,7 +476,7 @@ class Selection: """ Used to perform lazy selection. - It is usefull when dealing with lazy array to avoid loading unneccessary data. + It is useful when dealing with lazy array to avoid loading unnecessary data. It must be initialized with the shape of the underlying array. It allows to track the succesive slice or single element selections made along the different dimensions of the array. Once the overall selection must be aaplied, the diff --git a/xdas/virtual/tiles.py b/xdas/virtual/tiles.py index 25012c90..ed74dbe6 100644 --- a/xdas/virtual/tiles.py +++ b/xdas/virtual/tiles.py @@ -669,6 +669,11 @@ class TileArray(VirtualBackend, np.lib.mixins.NDArrayOperatorsMixin, vtype="tile # so multi-file scans can drain batches (see VirtualBackend) consolidates = True + #: Common directory of the tile sources, the stored per-tile paths being + #: relative to it. Read-only: an archive is relocated by rewriting the + #: ``root`` entry of the stored header, not this attribute. + root: str + def __init__(self, dataset, dtype=None, engine=None): # canonical string dtype is fixed-width bytes: str-valued # variables (hand-built manifests) recode here