From 6c36248497bbde32aa76ffd3436c6372d4a36290 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:26:36 +0200 Subject: [PATCH 01/48] sel works on string and categorical coordinates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Label selection on a string coordinate raised UFuncTypeError before any selection began: the overlap guard differenced the coordinate values, and numpy has no subtract loop for string dtypes. The monotonicity check now goes through the pandas index, which is dtype-generic; is_unique keeps it strict, since pandas considers repeated values monotonic increasing. The guard itself now covers only ordered look-ups. A slice resolves its bounds by searching and `method` searches for a neighbour, so both need an axis whose values increase and stay guarded. Naming a label is a hash look-up, well defined in any order, so it goes straight through — a categorical axis such as ["P", "S", "N"] is selectable by label, and a list returns its labels in the requested order. Ambiguous labels remain the coordinate's business: to_index raises on them. --- docs/release-notes.md | 2 + tests/coordinates/test_dense.py | 178 ++++++++++++++++++++++++++++++++ xdas/coordinates/dense.py | 10 +- xdas/core/dataarray.py | 42 +++++++- 4 files changed, 224 insertions(+), 8 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index eba1ea92..e3cd6591 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -30,6 +30,8 @@ - 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). +- Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@atrabattoni). +- Fix `sel` refusing an exact label look-up on a coordinate whose values are not sorted, such as a categorical axis like `["P", "S", "N"]`. The overlap guard covered every kind of selection, but only ordered look-ups — a slice, or `method="nearest"` and friends — need a sorted axis; naming a label does not. Those stay guarded, `da.sel(phase=["S", "P"])` now works and returns the labels in the requested order (@atrabattoni). ## 0.2.8 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/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/dataarray.py b/xdas/core/dataarray.py index 717ba22f..17de6f07 100644 --- a/xdas/core/dataarray.py +++ b/xdas/core/dataarray.py @@ -379,15 +379,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..." From e0d6c30d638844e73794da0fd73dc84564cb0589 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:27:53 +0200 Subject: [PATCH 02/48] registered numpy defaults no longer override explicit arguments The dispatch wrapper applied its registration defaults after binding the caller's arguments, so a registered `axis=-1` overwrote whatever the caller passed: np.cumsum(da, 0) accumulated along the last axis, positional or keyword alike. Defaults now fill in only the parameters the caller did not bind, before numpy's own signature defaults. --- docs/release-notes.md | 1 + tests/test_numpy.py | 16 ++++++++++++++++ xdas/core/numpy.py | 6 +++++- 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index e3cd6591..bc8283a2 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -31,6 +31,7 @@ - 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). - Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@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. Registered defaults now fill in only when the caller says nothing (@atrabattoni). - Fix `sel` refusing an exact label look-up on a coordinate whose values are not sorted, such as a categorical axis like `["P", "S", "N"]`. The overlap guard covered every kind of selection, but only ordered look-ups — a slice, or `method="nearest"` and friends — need a sorted axis; naming a label does not. Those stay guarded, `da.sel(phase=["S", "P"])` now works and returns the labels in the requested order (@atrabattoni). ## 0.2.8 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/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") From ac305c301e87a92e8d507818bb266a68ec098255 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:29:15 +0200 Subject: [PATCH 03/48] a directory sink joins along the dimension it was chunked along DataArrayWriter concatenated with concat's "first" default, which is only right when the chunked dimension leads the output: (distance, time) chunks from a time-chunked source were stacked along distance. The writer now takes the dimension as a dim argument; left unsaid it keeps the "first" default. --- docs/release-notes.md | 1 + tests/test_processing.py | 9 +++++++++ xdas/processing/core.py | 16 ++++++++++++++-- 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index bc8283a2..b94b3050 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -31,6 +31,7 @@ - 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). - Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@atrabattoni). +- Fix a directory sink joining its chunks along the wrong dimension when the pipeline's output does not lead with the chunked one — `(distance, time)` chunks written from a time-chunked source were stacked along `distance`. `DataArrayWriter` now takes the dimension as a `dim` argument (still `"first"` when left unsaid) (@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. Registered defaults now fill in only when the caller says nothing (@atrabattoni). - Fix `sel` refusing an exact label look-up on a coordinate whose values are not sorted, such as a categorical axis like `["P", "S", "N"]`. The overlap guard covered every kind of selection, but only ordered look-ups — a slice, or `method="nearest"` and friends — need a sorted axis; naming a label does not. Those stay guarded, `da.sel(phase=["S", "P"])` now works and returns the labels in the requested order (@atrabattoni). diff --git a/tests/test_processing.py b/tests/test_processing.py index 96b8e124..41a3b9c9 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -89,6 +89,15 @@ def test_passing_wrong_input(self, tmp_path): with pytest.raises(TypeError): dw.submit(None) + 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): diff --git a/xdas/processing/core.py b/xdas/processing/core.py index 5ff15832..0f593730 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -238,6 +238,11 @@ 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. Examples -------- @@ -256,7 +261,13 @@ 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", ): dirpath = str(dirpath) if isinstance(dirpath, Path) else dirpath if create_dirs: @@ -264,6 +275,7 @@ 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 @@ -310,7 +322,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: From b61e9fe0f5b3a125755ab1801f7e6cc0ebe62d6b Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:30:39 +0200 Subject: [PATCH 04/48] a pandas DataFrame is a collection leaf of its own kind The tree rebuild wrapped every non-DataArray leaf in DataArray(...), so a collection holding tables raised on repr. A DataFrame now passes through untouched, and the mapping repr asks the value what it is rather than assuming everything non-array is a subtree. --- docs/release-notes.md | 1 + tests/test_datacollection.py | 20 ++++++++++++++++++++ xdas/core/datacollection.py | 12 +++++++++--- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index b94b3050..e6353432 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -31,6 +31,7 @@ - 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). - Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@atrabattoni). +- Fix `DataCollection` coercing a `pandas.DataFrame` leaf into a `DataArray`, producing collections whose leaves raised on `repr`. A table is now a leaf of its own kind (@atrabattoni). - Fix a directory sink joining its chunks along the wrong dimension when the pipeline's output does not lead with the chunked one — `(distance, time)` chunks written from a time-chunked source were stacked along `distance`. `DataArrayWriter` now takes the dimension as a `dim` argument (still `"first"` when left unsaid) (@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. Registered defaults now fill in only when the caller says nothing (@atrabattoni). - Fix `sel` refusing an exact label look-up on a coordinate whose values are not sorted, such as a categorical axis like `["P", "S", "N"]`. The overlap guard covered every kind of selection, but only ordered look-ups — a slice, or `method="nearest"` and friends — need a sorted axis; naming a label does not. Those stay guarded, `da.sel(phase=["S", "P"])` now works and returns the labels in the requested order (@atrabattoni). diff --git a/tests/test_datacollection.py b/tests/test_datacollection.py index c34d6bec..7c437612 100644 --- a/tests/test_datacollection.py +++ b/tests/test_datacollection.py @@ -1,4 +1,5 @@ import h5py +import pandas as pd import pytest import xdas as xd @@ -511,3 +512,22 @@ 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 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): From 3d7c9dfa6ff1c309c749159b1c06ef65458cd14e Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:31:45 +0200 Subject: [PATCH 05/48] xd.stack collapses a collection level into a dimension The inverse of combine_by_coords, which concatenates along an existing dimension: xd.stack(dc, "channel") turns each station's traces into one (channel, time) array keyed by the level's own keys. The new dimension is named after the level it collapsed (dim= chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay lazy. Leaves that disagree on their other coordinates raise naming what disagreed; join="inner" trims to the common part and join="outer" pads with NaN. Coordinate agreement is judged on the sampling grid within a tolerance rather than exactly: real networks record a channel a nanosecond off, and exact comparison made outer joins silently interleave two grids into mostly-NaN samples. The default tolerance is 1/100 sample where a nominal interval is declared, tolerance=False keeps the strict comparison, and true interleaving is refused rather than padded. Also fixes concat opening a new dimension over VirtualSource-backed arrays raising TypeError: whether the result can stay virtual is decided before expand_dims, which no virtual source can follow. --- docs/api/xdas.md | 1 + docs/release-notes.md | 2 + tests/test_routines.py | 694 +++++++++++++++++++++++++++++++++++++++++ xdas/__init__.py | 2 + xdas/core/__init__.py | 2 + xdas/core/routines.py | 580 +++++++++++++++++++++++++++++++++- 6 files changed, 1278 insertions(+), 3 deletions(-) diff --git a/docs/api/xdas.md b/docs/api/xdas.md index 8012fb94..66a05a25 100644 --- a/docs/api/xdas.md +++ b/docs/api/xdas.md @@ -34,6 +34,7 @@ get_sampling_interval sortby split + stack plot_availability ``` diff --git a/docs/release-notes.md b/docs/release-notes.md index e6353432..2678ea99 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -7,6 +7,7 @@ - **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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@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). @@ -28,6 +29,7 @@ ### 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). +- `xdas.concat` opening a *new* dimension over `VirtualSource`-backed arrays no longer raises `TypeError: only VirtualSource object can be provided`. Whether the result can stay virtual was decided before `expand_dims`, which no virtual source can follow — a stack of sources is a longer axis, never an extra one — so a `VirtualStack` was promised over arrays that had already been loaded (@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). - Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@atrabattoni). 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/xdas/__init__.py b/xdas/__init__.py index 0145ba61..219af67c 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -57,6 +57,7 @@ "plot_availability", "sortby", "split", + "stack", "trim_overlaps", ] @@ -110,6 +111,7 @@ routines, sortby, split, + stack, trim_overlaps, ) from .core.methods import * 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/routines.py b/xdas/core/routines.py index a24ba149..c32b0af0 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, joinable = _leaf_mismatches(objs, keys) + joinable = [] # already spent + 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. From 1fdc3532c268f6aaa3f214fdb035a73f5a1425b3 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:37:22 +0200 Subject: [PATCH 06/48] chunk ingress and egress can run in worker processes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reading a chunk goes through one virtual layout, hence one h5py call under the global HDF5 lock, so loader threads never decode concurrently and extra threads only contend; compression on the write side contends the same way. DataArrayLoader and DataArrayWriter now accept pool="processes": each worker holds its own HDF5 lock, receives the manifest of its chunk (a sliced virtual array, kilobytes) and reads its own files. The chunks come back through Ray's shared-memory object store rather than a pickle pipe: written once by the worker, mapped zero-copy by the parent. A pipe caps out well below memory bandwidth and spends parent CPU on deserialization the signal chain needs — measured end to end on a compressed ZFP archive at 16 workers, ingest 137 -> 1378 MiB/s and egress 151 -> 1556 where a pickling pool lost end to end. max_workers is enforced by parking submissions beyond it; Ray initializes lazily, an already initialized runtime is left untouched, and shutdown leaves it up because it is a session-wide resource. The price of zero-copy is immutability: chunk data arrives read-only, which atoms honor by allocating their outputs — pinned by a test running a stateful pipeline over read-only chunks, ray installed or not. Ray is the optional extra xdas[ray]; pool="threads" stays the default. --- docs/release-notes.md | 1 + pyproject.toml | 5 +- tests/test_processing.py | 153 +++++++++++++++++++++++++ xdas/processing/core.py | 237 ++++++++++++++++++++++++++++++++++++--- 4 files changed, 380 insertions(+), 16 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index 2678ea99..13f22841 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,6 +12,7 @@ ### 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). +- **Process pools for chunk ingress and egress.** `DataArrayLoader` and `DataArrayWriter` accept `pool="processes"`, which reads and writes chunks in worker processes instead of threads: compressed HDF5 decodes and compresses under the global HDF5 lock, so extra *threads* only contend, while processes each hold their own lock. What crosses to a worker on the read side is the manifest of the chunk — a sliced virtual array, kilobytes — so each worker reads its own files, and the loaded chunk comes back through Ray's shared-memory object store: written once by the worker, mapped zero-copy by the parent, arriving read-only (the immutability convention atoms already follow). End to end on a compressed ZFP archive at 16 workers, ingest goes from 137 to 1378 MiB/s and egress from 151 to 1556. Ray is an optional dependency (`pip install xdas[ray]`); `pool="threads"` remains the default (@atrabattoni). - **`xdas.sortby`.** Sort a tile- or stack-backed data array along a dimension by coordinate value, lazily: the blocks are permuted through the manifest without reading any data (@atrabattoni). - `xdas.concat` opening a *new* dimension now checks that the inputs agree on their other coordinates and promotes the scalar ones that vary. Stacking the components of a station is `xd.concat(traces, "channel")`, lazily, with `channel` becoming a real coordinate (@atrabattoni). - Acquisitions interleaved in time now group by compatibility, one array each, instead of splitting at every alternation (@atrabattoni). diff --git a/pyproject.toml b/pyproject.toml index 9cf5df22..3b1d175b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,9 @@ dependencies = [ "pyzmq", ] +[project.optional-dependencies] +ray = ["ray"] + [dependency-groups] dev = ["ruff", "pytest", "pytest-cov"] docs = [ @@ -40,7 +43,7 @@ docs = [ "sphinx-copybutton", "sphinx", ] -tests = ["dascore", "dask<2025.4.0", "psutil", "seisbench", "torch"] +tests = ["dascore", "dask<2025.4.0", "psutil", "ray", "seisbench", "torch"] # Single source of truth for the version: xdas/__init__.py [tool.setuptools.dynamic] diff --git a/tests/test_processing.py b/tests/test_processing.py index 41a3b9c9..be45aebc 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1,3 +1,4 @@ +import importlib.util import os import threading import time @@ -54,6 +55,158 @@ 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 (zero-copy from an object store). + + Atoms honor this by allocating their outputs rather than writing into + their input; this pins the contract down without needing ray installed. + """ + + 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) + + +def test_missing_ray_points_at_the_extra(monkeypatch): + """Without ray installed, the pool names the optional dependency.""" + import builtins + + from xdas.processing.core import ProcessPool + + real_import = builtins.__import__ + + def no_ray(name, *args, **kwargs): + if name == "ray": + raise ImportError("no module named ray") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_ray) + with pytest.raises(ImportError, match=r"xdas\[ray\]"): + ProcessPool(1) + + +requires_ray = pytest.mark.skipif( + importlib.util.find_spec("ray") is None, reason="ray is not installed" +) + + +@requires_ray +@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 the object store. + 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 a store-backed array: the parent + # mapped shared memory, nothing was pickled back. + da = xd.testing.dummy(shape=(1000, 100)) + chunks = list(xp.DataArrayLoader(da, {"time": 100}, 2, 2, pool="processes")) + assert all(not chunk.data.flags.writeable for chunk in chunks) + + 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()) + + def test_backlog_is_parked_and_ordered(self): + # More submissions than workers: the excess is parked, launched as + # capacity frees, and resolved in submission order. The tasks sleep + # so none can finish during the submit loop and free a slot. + from xdas.processing.core import ProcessPool + + with ProcessPool(2) as pool: + futures = [pool.submit(_slow_double, i) for i in range(8)] + assert sum(future._ref is None for future in futures) == 6 + assert len(pool._running) == 2 + assert [future.result() for future in futures] == [2 * i for i in range(8)] + + def test_errors_propagate(self): + from xdas.processing.core import ProcessPool + + with ProcessPool(1) as pool, pytest.raises(ValueError, match="broken task"): + pool.submit(_raise).result() + + def test_shutdown_cancels_backlog(self): + from concurrent.futures import CancelledError + + from xdas.processing.core import ProcessPool + + pool = ProcessPool(1) + futures = [pool.submit(_slow_double, i) for i in range(4)] + pool.shutdown() + pool.shutdown() # idempotent: nothing left to wait for + assert futures[0].result() == 0 + with pytest.raises(CancelledError): + futures[-1].result() + + +def _slow_double(x): + """Worker task used by the pool tests, slow enough to keep a backlog parked.""" + time.sleep(0.2) + return 2 * x + + +def _raise(): + """Worker task raising, to check error propagation through the store.""" + raise ValueError("broken task") + class TestDataArrayWriter: def test_init(self, tmp_path): diff --git a/xdas/processing/core.py b/xdas/processing/core.py index 0f593730..99ebafe5 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -7,7 +7,8 @@ """ import os -from concurrent.futures import ThreadPoolExecutor +from collections import deque +from concurrent.futures import CancelledError, ThreadPoolExecutor from glob import glob from pathlib import Path from queue import Queue @@ -24,6 +25,187 @@ from .monitor import Monitor +class _RayFuture: + """A future handed out by :class:`ProcessPool`, resolved via the object store.""" + + def __init__(self, pool, task): + self._pool = pool + self._task = task + self._ref = None + self._cancelled = False + + def result(self): + """Block until the task ran and return its output (or raise its error).""" + return self._pool._result(self) + + +class ProcessPool: + """ + A pool of worker processes whose results cross through shared memory. + + Each task runs as a Ray task in its own process. The pool quacks like a + :class:`~concurrent.futures.Executor` as far as the loader and writer + need (``submit``/``shutdown``/context manager), but the data crossing + back is never pickled through a pipe: a task result lands in Ray's + shared-memory object store, written once by the worker, and ``result()`` + maps it zero-copy into the parent. Large task *arguments* — the chunk a + writer sends out — take the same path, one memcpy into the store instead + of a serialize-transfer-deserialize round. The price of zero-copy is + immutability: array data coming out of the store is read-only, which + atoms honor by allocating their outputs. + + ``max_workers`` is enforced by parking submissions beyond it and + launching them as running tasks finish, mirroring how a process pool + queues its backlog. Ray is initialized lazily on first use (an already + initialized Ray, e.g. configured by the user, is left untouched). + + Parameters + ---------- + max_workers : int + Maximum number of tasks running concurrently. + """ + + def __init__(self, max_workers): + try: + import ray + except ImportError: + raise ImportError( + "pool='processes' requires the ray package: pip install xdas[ray]" + ) from None + if not ray.is_initialized(): + ray.init(include_dashboard=False) + self._ray = ray + self._max_workers = max_workers + self._remotes = {} + self._pending = deque() + self._running = [] + + def submit(self, fn, /, *args, **kwargs): + """ + Schedule ``fn(*args, **kwargs)`` as a task and return its future. + + Parameters + ---------- + fn : callable + The unit of work; large array arguments go to the object store. + + Returns + ------- + _RayFuture + A ``result()``-able handle, resolved zero-copy from the store. + """ + future = _RayFuture(self, (fn, args, kwargs)) + self._pending.append(future) + self._launch() + return future + + def shutdown(self, wait=True): + """ + Cancel parked tasks and (by default) wait for the running ones. + + The Ray runtime itself is left up: it is a session-wide resource, + shared with the other pools of the run and with whatever the user + configured before xdas started. + + Parameters + ---------- + wait : bool, optional + Whether to block until in-flight tasks complete. + """ + while self._pending: + self._pending.popleft()._cancelled = True + if wait and self._running: + self._ray.wait(self._running, num_returns=len(self._running)) + self._running.clear() + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self.shutdown() + + def _remote(self, fn): + """Wrap *fn* as a Ray remote function, once per distinct callable.""" + if fn not in self._remotes: + self._remotes[fn] = self._ray.remote(fn) + return self._remotes[fn] + + def _launch(self): + """Drop finished tasks and start parked ones up to ``max_workers``.""" + if self._running: + _, self._running = self._ray.wait( + self._running, num_returns=len(self._running), timeout=0 + ) + while self._pending and len(self._running) < self._max_workers: + future = self._pending.popleft() + fn, args, kwargs = future._task + future._ref = self._remote(fn).remote(*args, **kwargs) + self._running.append(future._ref) + + def _result(self, future): + """Resolve *future*, first waiting out the backlog ahead of it.""" + while future._ref is None: + if future._cancelled: + raise CancelledError() + self._ray.wait(self._running, num_returns=1) + self._launch() + return self._ray.get(future._ref) + + +POOLS = {"threads": ThreadPoolExecutor, "processes": ProcessPool} +"""Worker pools available for chunk ingress and egress, name → factory.""" + + +def get_pool(pool, max_workers): + """ + 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 a shared-memory object + store: the worker writes it once, the parent maps it zero-copy + (read-only). ``"processes"`` requires the optional ``ray`` dependency + (``pip install xdas[ray]``). + + Parameters + ---------- + pool : str + Pool kind, ``"threads"`` (default everywhere) or ``"processes"``. + max_workers : int + Number of workers. + + 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)}") + return POOLS[pool](max_workers) + + +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, data_loader, data_writer): """ Execute a chunked processing pipeline. @@ -89,7 +271,16 @@ class DataArrayLoader: 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,9 +298,13 @@ 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): @@ -134,28 +329,35 @@ def __init__(self, da, chunks, max_buffers=1, max_workers=1): self.chunk_size = chunk_size 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 - def _get_chunk(self, idx): + def _select(self, idx): + """Return chunk *idx* as a lazy selection: the manifest, not the data.""" start = idx * self.chunk_size end = (idx + 1) * self.chunk_size 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: + with get_pool(self.pool, self.max_workers) 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 +366,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 @@ -243,6 +445,13 @@ class DataArrayWriter: :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 -------- @@ -268,6 +477,7 @@ def __init__( max_workers=1, create_dirs=False, dim="first", + pool="threads", ): dirpath = str(dirpath) if isinstance(dirpath, Path) else dirpath if create_dirs: @@ -279,7 +489,8 @@ def __init__( 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._futures = [] self._results = [] self._count = 0 @@ -299,18 +510,14 @@ def submit(self, chunk): 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() From 45b3a1da34468940229fd47852240d240ec30e51 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 12:48:41 +0200 Subject: [PATCH 07/48] atoms compose with >>, trace ufuncs, and clone with fresh() The core grows the vocabulary pipelines are built from: - compose() and >>/>>=/da >> atom, with value semantics everywhere: composing never mutates an operand, so intermediate pipelines stay usable on their own. This also fixes the atomized Sequential-append aliasing bug (it mutated the input and returned None). - Ufunc tracing under the ... seed: applying a numpy ufunc to an atom appends the operation to the pipeline instead of computing it, so 20 * np.log10(np.abs(atom)) is a pipeline. The traced surface is ufuncs exactly; an expression involving two atoms (fan-in) raises at the line that wrote it, never bails silently into computation. Atom equality becomes identity so atoms can live in sets and traces cannot be confused by ==. - as_function(cls) generates the function form of an atom class: a lowercase function taking the data first, eager on data, returning the configured atom on ..., extending a pipeline on an atom. atomized dispatches classes to it and keeps its function behaviour. - fresh(): a stateless clone, config shared by reference (a model is never deep-copied), nested atoms recursed. initialized now recurses into nested atoms, so a pipeline holding an uninitialised filter no longer reports itself ready. - A private whole-record guard: functions that need the whole record along their working dimension are marked at their definition site and refuse chunked execution along that dimension with a pointed error, resolving first/last aliases against the data before comparing. DataArray defers to foreign __array_ufunc__ implementations so da >> atom dispatches to the atom rather than being computed. --- docs/release-notes.md | 2 + tests/test_atoms.py | 208 ++++++++++++++++++++++++++++- xdas/atoms/__init__.py | 7 +- xdas/atoms/core.py | 289 ++++++++++++++++++++++++++++++++++++++--- xdas/core/dataarray.py | 8 ++ 5 files changed, 494 insertions(+), 20 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index 13f22841..a626212e 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,8 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **`>>` composition and operator tracing.** Atoms compose into pipelines with `>>`/`>>=` (bare callables auto-wrap, `da >> atom` applies), and ordinary numpy expressions trace under the `...` seed: `20 * np.log10(np.abs(atom))` appends `absolute → log10 → multiply` to the pipeline instead of computing. Tracing covers ufuncs exactly — a traced expression involving two atoms (fan-in) raises at the line that wrote it rather than silently computing. Composition has value semantics: passing a `Sequential` to an atomized function returns a new extended pipeline instead of mutating (and aliasing) the input — the mutating form also returned `None`, breaking chained composition. `xdas.atoms.as_function` generates the function form of any atom class, and atoms gain `fresh()` (a stateless clone whose config is shared by reference) while `initialized` now recurses into nested atoms (@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). diff --git a/tests/test_atoms.py b/tests/test_atoms.py index 05266c1b..725deb1c 100644 --- a/tests/test_atoms.py +++ b/tests/test_atoms.py @@ -16,7 +16,9 @@ ResamplePoly, Sequential, UpSample, + atomized, ) +from xdas.atoms.core import _whole_record from xdas.signal import lfilter from xdas.synthetics import randn_wavefronts @@ -383,11 +385,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 @@ -521,3 +525,201 @@ def test_mlpicker_invalid_component_strategy(self): 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 + + +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 + + +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/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index 7a251e2b..eec0e534 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -2,7 +2,8 @@ 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`. +:func:`atomized`, :func:`as_function`, :func:`compose`, signal-processing +atoms, and the ML-based :class:`MLPicker`. """ __all__ = [ @@ -19,11 +20,13 @@ "State", "Trigger", "UpSample", + "as_function", "atomized", + "compose", ] from ..trigger import Trigger -from .core import Atom, Partial, Sequential, State, atomized +from .core import Atom, Partial, Sequential, State, as_function, atomized, compose from .ml import MLPicker from .signal import ( DownSample, diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index 6bf0c88d..fe1e342d 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -2,14 +2,20 @@ 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 from collections.abc import Callable from functools import wraps from typing import Any +import numpy as np + from ..core import DataArray, DataCollection, open_datacollection @@ -53,7 +59,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. @@ -83,6 +89,10 @@ class Atom: that are usefull 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 @@ -101,6 +111,8 @@ class Atom: Performs the main processing logic of the atom. reset() Resets the atom to its initial state. + fresh() + Returns a stateless clone sharing the configuration. """ @@ -109,6 +121,14 @@ def __init__(self): object.__setattr__(self, "_state", {}) object.__setattr__(self, "_atoms", {}) + 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__ sig = ", ".join( @@ -140,8 +160,10 @@ 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.""" @@ -158,6 +180,7 @@ def call(self, x, **flags): def __call__(self, x, **flags): """Process input data, initializing state if needed and resetting after final chunk.""" chunk_dim = flags.get("chunk_dim", None) + self._check_chunk_dim(x, chunk_dim) if not self.initialized or chunk_dim is None: self.initialize(x, **flags) y = self.call(x, **flags) @@ -165,6 +188,83 @@ def __call__(self, x, **flags): self.reset() return y + 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. ``"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 + if x is not None and hasattr(x, "dims"): + if dim == "first": + dim = x.dims[0] + elif dim == "last": + dim = x.dims[-1] + if dim is None or dim in ("first", "last") or dim == chunk_dim: + name = ( + getattr(self, "name", None) + or getattr(getattr(self, "func", None), "__name__", None) + or type(self).__name__ + ) + raise ValueError( + f"{name} needs the whole record along {dim!r} 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 kwargs["out"] != (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).""" for key in self._state: @@ -172,6 +272,33 @@ def reset(self): 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"): + 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) @@ -296,6 +423,10 @@ def call(self, x: Any, **flags) -> Any: x = atom(x, **flags) return x + def fresh(self): + """Return a stateless clone: each stage cloned, config shared.""" + return type(self)([atom.fresh() for atom in self], name=self.name) + def __repr__(self) -> str: width = len(str(len(self))) name = self.name if self.name is not None else "sequence" @@ -391,6 +522,23 @@ def __init__( setattr(self, key, value) else: self.kwargs[key] = value + # A whole-record function marked with `_whole_record` refuses chunked + # execution along its working dimension; that dimension is resolved + # from the call arguments so the guard can compare it with the + # chunked one. + dim_arg = getattr(func, "_whole_record_dim_arg", None) + if dim_arg is not None: + try: + bound = inspect.signature(func).bind_partial(*self.args, **self.kwargs) + bound.apply_defaults() + self.dim = bound.arguments.get(dim_arg) + except (TypeError, ValueError): + self.dim = None + + 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.dim, chunk_dim, x) @property def stateful(self): @@ -448,29 +596,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 +707,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 +721,76 @@ 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"): + """ + 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`. + """ + + def decorator(func): + func._whole_record_dim_arg = dim_arg + 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/core/dataarray.py b/xdas/core/dataarray.py index 17de6f07..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__)) ) From 8e373ac5246ee650ff3ed1ba5b591a15fcf81811 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:06:05 +0200 Subject: [PATCH 08/48] kernel atoms get their own layer, resampling goes polyphase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LFilter, SOSFilter, DownSample and UpSample move to xdas.atoms.kernel, the expert layer of exact stateful primitives with machine parameters (taps, integer factors) whose meaning depends on the sampling rate; they stay importable from xdas.atoms. The layer is born with Polyphase: upsample, FIR filter and downsample in one scipy.signal.upfirdn pass, computing only the output samples that survive the decimation and never materialising the zero-stuffed signal. FIRFilter gains up=/down= and applies its taps through Polyphase (designed at the upsampled rate, energy-compensated), and ResamplePoly collapses its upsample/filter/downsample trio to that one child atom — 2.6x on a decimation by two along distance, 8.7x on 62.5 -> 50 Hz on a 254 MiB chunk. Chunked calls carry the filter memory and the output-grid phase across chunks, so splitting the input does not change the result. Taps are cast down to the data precision (float32 in, float32 out instead of an lfilter promotion doubling downstream memory traffic), and a target rate the coordinate resolution cannot represent exactly declares its residual drift as jitter on the output coordinate. The output-phase arithmetic is pinned against scipy's own upfirdn and against the explicit kernel chain across up/down combinations and chunk cuts. --- docs/api/atoms.md | 22 ++- docs/release-notes.md | 1 + tests/test_atoms.py | 89 ++++++++- xdas/atoms/__init__.py | 26 +-- xdas/atoms/kernel.py | 429 +++++++++++++++++++++++++++++++++++++++++ xdas/atoms/signal.py | 257 +++--------------------- 6 files changed, 580 insertions(+), 244 deletions(-) create mode 100644 xdas/atoms/kernel.py diff --git a/docs/api/atoms.md b/docs/api/atoms.md index c0b26cc3..0ec86ab7 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -75,13 +75,15 @@ Methods Partial.get_state ``` -## Decorators +## Decorators and composition ```{eval-rst} .. autosummary:: :toctree: ../_autosummary + as_function atomized + compose ``` ## Signal processing @@ -90,13 +92,25 @@ Methods .. autosummary:: :toctree: ../_autosummary - DownSample FIRFilter IIRFilter - LFilter MLPicker ResamplePoly - SOSFilter Trigger +``` + +## 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 + SOSFilter UpSample ``` \ No newline at end of file diff --git a/docs/release-notes.md b/docs/release-notes.md index a626212e..2a1dbf2c 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,7 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **Polyphase resampling in a kernel layer.** The exact machine-parameter atoms move to the expert layer `xdas.atoms.kernel` (`LFilter`, `SOSFilter`, `DownSample`, `UpSample`, still importable from `xdas.atoms`), joined by the new `Polyphase` kernel: upsample, FIR filter and downsample fused into a single `scipy.signal.upfirdn` pass that computes only the output samples surviving the decimation and never materialises the zero-stuffed signal (which for `up=4` allocated a four times larger, mostly-zero array). `FIRFilter` is born with `up=`/`down=` and `ResamplePoly` rides it, so the upsample/filter/downsample trio collapses to one child atom — on a 254 MiB chunk that is 2.6× on a decimation by two along distance and 8.7× on a 62.5 → 50 Hz resampling. The taps are cast down to the data precision, so float32 stays float32 instead of being promoted by the filter; a target rate the coordinate resolution cannot represent exactly (100 Hz → 30 Hz is 10/3 ns per sample) declares its residual drift as jitter instead of rejecting its own sampling interval (@atrabattoni). - **`>>` composition and operator tracing.** Atoms compose into pipelines with `>>`/`>>=` (bare callables auto-wrap, `da >> atom` applies), and ordinary numpy expressions trace under the `...` seed: `20 * np.log10(np.abs(atom))` appends `absolute → log10 → multiply` to the pipeline instead of computing. Tracing covers ufuncs exactly — a traced expression involving two atoms (fan-in) raises at the line that wrote it rather than silently computing. Composition has value semantics: passing a `Sequential` to an atomized function returns a new extended pipeline instead of mutating (and aliasing) the input — the mutating form also returned `None`, breaking chained composition. `xdas.atoms.as_function` generates the function form of any atom class, and atoms gain `fresh()` (a stateless clone whose config is shared by reference) while `initialized` now recurses into nested atoms (@atrabattoni). ### Improvements diff --git a/tests/test_atoms.py b/tests/test_atoms.py index 725deb1c..40d9f55a 100644 --- a/tests/test_atoms.py +++ b/tests/test_atoms.py @@ -11,8 +11,10 @@ DownSample, FIRFilter, IIRFilter, + LFilter, MLPicker, Partial, + Polyphase, ResamplePoly, Sequential, UpSample, @@ -207,6 +209,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") @@ -215,7 +219,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) @@ -258,6 +265,86 @@ def test_nothing_to_do(self): assert result.equals(da) +class TestPolyphase: + """The fused kernel against the upsample/filter/downsample chain it replaces.""" + + @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.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, full[: result.sizes["time"]], atol=1e-15 + ) + + @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) + + class TestMLPicker: @pytest.mark.slow def test_picker(self): diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index eec0e534..ba393042 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -1,9 +1,17 @@ """ Stateful processing units (atoms) for building chunked data pipelines. -Exports :class:`Atom`, :class:`State`, :class:`Sequential`, :class:`Partial`, -:func:`atomized`, :func:`as_function`, :func:`compose`, signal-processing -atoms, and the ML-based :class:`MLPicker`. +Two layers so far: + +- :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`). + +Plus the signal-processing atoms of :mod:`xdas.atoms.signal` and the ML-based +:class:`MLPicker`. """ __all__ = [ @@ -14,6 +22,7 @@ "LFilter", "MLPicker", "Partial", + "Polyphase", "ResamplePoly", "SOSFilter", "Sequential", @@ -27,13 +36,6 @@ from ..trigger import Trigger from .core import Atom, Partial, Sequential, State, as_function, atomized, compose +from .kernel import DownSample, LFilter, Polyphase, SOSFilter, UpSample from .ml import MLPicker -from .signal import ( - DownSample, - FIRFilter, - IIRFilter, - LFilter, - ResamplePoly, - SOSFilter, - UpSample, -) +from .signal import FIRFilter, IIRFilter, ResamplePoly diff --git a/xdas/atoms/kernel.py b/xdas/atoms/kernel.py new file mode 100644 index 00000000..2947fb90 --- /dev/null +++ b/xdas/atoms/kernel.py @@ -0,0 +1,429 @@ +""" +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`. +""" + +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 + + +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)) + + +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) + + +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") + >>> outs = [atom(chunk, chunk_dim="time") for chunk in xd.split(da, 7, "time")] + >>> chunked = xd.concat(outs, "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)) + if self.dim == chunk_dim: + shape = tuple( + self._history_size() if name == self.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.""" + coord = da.coords[self.dim] + delta = get_sampling_interval(da, self.dim, 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[self.dim] = Coordinate(data, self.dim) + return coords + + 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 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) From 09e28145eab19e50dd6482c2a99e5b0e45e89d48 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:10:36 +0200 Subject: [PATCH 09/48] task atoms speak physical units, with function forms at the top level xdas.atoms.tasks is the public processing vocabulary: Filter (a (low, high) corner pair in Hz with None opening one end, ftype iir/fir, zerophase), Decimate and Resample (target rate in Hz), Integrate and Differentiate (stateful, carrying their seam state across chunks), plus whole-record detrend, taper, hilbert, sliding_mean_removal and medfilt with kernel lengths in physical units. Machine parameters (taps, factors) are designed from the data at the first call and live in the kernel layer, so a pipeline keeps its meaning when the sampling rate changes. Decimate applies its anti-alias taps through the polyphase kernel (antialias.down carries the factor), never filtering at the full rate to throw most of the result away. Zero-phase IIR filtering has no causal streaming form, so that atom refuses chunked execution along its dimension, as do the whole-record functions, each marked at its definition site. Every task atom gets a function form exported at the top level: xd.decimate(da, 50.0) applies eagerly, xd.decimate(..., 50.0) returns the atom, and passing an atom extends a pipeline. --- docs/api/atoms.md | 41 +++ docs/release-notes.md | 2 + tests/test_atoms_tasks.py | 295 +++++++++++++++++++++ xdas/__init__.py | 23 ++ xdas/atoms/__init__.py | 11 +- xdas/atoms/core.py | 18 +- xdas/atoms/tasks.py | 521 ++++++++++++++++++++++++++++++++++++++ 7 files changed, 902 insertions(+), 9 deletions(-) create mode 100644 tests/test_atoms_tasks.py create mode 100644 xdas/atoms/tasks.py diff --git a/docs/api/atoms.md b/docs/api/atoms.md index 0ec86ab7..3f571eb6 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -86,6 +86,47 @@ Methods compose ``` +## Task atoms + +Public processing vocabulary with physical parameters only. Each task atom has +a function form exported at the top level of `xdas` (e.g. `xdas.filter`, +`xdas.decimate`). + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + Decimate + Differentiate + Filter + Integrate + Resample +``` + +```{eval-rst} +.. currentmodule:: xdas +``` + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + decimate + detrend + differentiate + filter + hilbert + integrate + medfilt + resample + sliding_mean_removal + taper +``` + +```{eval-rst} +.. currentmodule:: xdas.atoms +``` + ## Signal processing ```{eval-rst} diff --git a/docs/release-notes.md b/docs/release-notes.md index 2a1dbf2c..558a6eb0 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,8 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **Task atoms with physical units.** A new public processing vocabulary where every parameter keeps its meaning when the sampling rate changes: `Filter` (one atom for all bands — a `(low, high)` corner pair in Hz with `None` opening one end, `ftype="iir"/"fir"`, `zerophase`), `Decimate` and `Resample` (target rate in Hz, both riding the polyphase kernel — the filter-at-full-rate-then-discard chain is never taken), `Integrate` and `Differentiate` (chunk-correct, carrying state across seams), plus whole-record `detrend`, `taper`, `hilbert`, `sliding_mean_removal` and `medfilt` (kernel lengths now in seconds/meters), each refusing chunked execution along its working dimension instead of silently answering wrong. Task atoms default to `dim="time"` and live in `xdas.atoms.tasks` (@atrabattoni). +- **Function forms at the top level.** Every task atom generates a top-level function with a synthesized signature and docstring: `xdas.decimate(da, 50.0)` applies eagerly, `xdas.decimate(..., 50.0)` returns the atom, and passing an atom extends a pipeline — so the same code runs eagerly on a slice and chunked on an archive by seeding it with `...` (@atrabattoni). - **Polyphase resampling in a kernel layer.** The exact machine-parameter atoms move to the expert layer `xdas.atoms.kernel` (`LFilter`, `SOSFilter`, `DownSample`, `UpSample`, still importable from `xdas.atoms`), joined by the new `Polyphase` kernel: upsample, FIR filter and downsample fused into a single `scipy.signal.upfirdn` pass that computes only the output samples surviving the decimation and never materialises the zero-stuffed signal (which for `up=4` allocated a four times larger, mostly-zero array). `FIRFilter` is born with `up=`/`down=` and `ResamplePoly` rides it, so the upsample/filter/downsample trio collapses to one child atom — on a 254 MiB chunk that is 2.6× on a decimation by two along distance and 8.7× on a 62.5 → 50 Hz resampling. The taps are cast down to the data precision, so float32 stays float32 instead of being promoted by the filter; a target rate the coordinate resolution cannot represent exactly (100 Hz → 30 Hz is 10/3 ns per sample) declares its residual drift as jitter instead of rejecting its own sampling interval (@atrabattoni). - **`>>` composition and operator tracing.** Atoms compose into pipelines with `>>`/`>>=` (bare callables auto-wrap, `da >> atom` applies), and ordinary numpy expressions trace under the `...` seed: `20 * np.log10(np.abs(atom))` appends `absolute → log10 → multiply` to the pipeline instead of computing. Tracing covers ufuncs exactly — a traced expression involving two atoms (fan-in) raises at the line that wrote it rather than silently computing. Composition has value semantics: passing a `Sequential` to an atomized function returns a new extended pipeline instead of mutating (and aliasing) the input — the mutating form also returned `None`, breaking chained composition. `xdas.atoms.as_function` generates the function form of any atom class, and atoms gain `fresh()` (a stateless clone whose config is shared by reference) while `initialized` now recurses into nested atoms (@atrabattoni). diff --git a/tests/test_atoms_tasks.py b/tests/test_atoms_tasks.py new file mode 100644 index 00000000..b57530cc --- /dev/null +++ b/tests/test_atoms_tasks.py @@ -0,0 +1,295 @@ +import inspect + +import numpy as np +import pytest + +import xdas as xd +import xdas.signal as xs +from xdas.atoms import ( + Decimate, + Differentiate, + Filter, + Integrate, + Partial, + ResamplePoly, + Sequential, +) +from xdas.synthetics import wavelet_wavefronts + + +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) diff --git a/xdas/__init__.py b/xdas/__init__.py index 219af67c..f6be2e8b 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -59,6 +59,17 @@ "split", "stack", "trim_overlaps", + # task atoms (function forms) + "decimate", + "detrend", + "differentiate", + "filter", + "hilbert", + "integrate", + "medfilt", + "resample", + "sliding_mean_removal", + "taper", ] from . import ( @@ -74,6 +85,18 @@ testing, virtual, ) +from .atoms.tasks import ( + decimate, + detrend, + differentiate, + filter, + hilbert, + integrate, + medfilt, + resample, + sliding_mean_removal, + taper, +) from .coordinates import ( Coordinate, Coordinates, diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index ba393042..05a2b873 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -1,7 +1,7 @@ """ Stateful processing units (atoms) for building chunked data pipelines. -Two layers so far: +Three layers: - :mod:`xdas.atoms.core`: the machinery — :class:`Atom`, :class:`State`, :class:`Sequential`, :class:`Partial`, :func:`atomized`, @@ -9,6 +9,9 @@ - :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:`MLPicker`. @@ -16,13 +19,18 @@ __all__ = [ "Atom", + "Decimate", + "Differentiate", "DownSample", "FIRFilter", + "Filter", "IIRFilter", + "Integrate", "LFilter", "MLPicker", "Partial", "Polyphase", + "Resample", "ResamplePoly", "SOSFilter", "Sequential", @@ -39,3 +47,4 @@ from .kernel import DownSample, LFilter, Polyphase, SOSFilter, UpSample from .ml import MLPicker from .signal import FIRFilter, IIRFilter, ResamplePoly +from .tasks import Decimate, Differentiate, Filter, Integrate, Resample diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index fe1e342d..d7f42412 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -198,18 +198,20 @@ def _refuse_chunked_along(self, dim, chunk_dim, x=None): 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. ``"first"`` and - ``"last"`` aliases are resolved against *x* when given, so the - comparison is never made on an unresolved alias. + 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"): - if dim == "first": - dim = x.dims[0] - elif dim == "last": - dim = x.dims[-1] - if dim is None or dim in ("first", "last") or dim == chunk_dim: + 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) diff --git a/xdas/atoms/tasks.py b/xdas/atoms/tasks.py new file mode 100644 index 00000000..c87f5fb7 --- /dev/null +++ b/xdas/atoms/tasks.py @@ -0,0 +1,521 @@ +""" +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 + +from ..coordinates import get_sampling_interval +from ..core import concat +from .core import Atom, State, _whole_record, atomized +from .signal import FIRFilter, IIRFilter, ResamplePoly + +__all__ = [ + "Decimate", + "Differentiate", + "Filter", + "Integrate", + "Resample", + "decimate", + "detrend", + "differentiate", + "filter", + "hilbert", + "integrate", + "medfilt", + "resample", + "sliding_mean_removal", + "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) + + +filter = atomized(Filter) +decimate = atomized(Decimate) +resample = atomized(Resample) +integrate = atomized(Integrate) +differentiate = atomized(Differentiate) From 3cc1e7901e92a978a2ec16ed845de8f5220b92b6 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:19:45 +0200 Subject: [PATCH 10/48] atoms understand the continuous run: seams, flush, 0..n chunks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every stateful atom now judges the seams of its own input stream from the chunk coordinates: a continuous chunk carries state across, a gap or rate change flushes the previous run and reinitialises (redesigning coefficients for the new rate), a backward overlap raises, and on_discontinuity="reset"|"raise" makes strict runs opt-in. Incoming chunks are split at internal discontinuities so state never crosses a gap, and chunked stateful processing requires a regular coordinate along the chunked dimension, raising with a pointer to to_regular() instead of silently carrying state across unverifiable seams. call() follows the transducer contract: one chunk in, zero or more chunks out. The new flush() lifecycle drains buffered samples at the end of the stream, at every seam, and at the end of every eager call; Sequential.flush cascades codec-drain style, process() drains the pipeline, and writers drop empty chunks. This fixes chunked DownSample dropping its trailing samples when the length is not a multiple of the factor. Reductions fall out of the contract (accumulate in call, emit at flush), Atom.iter_chunks exposes the manual chunk loop as a plain generator, and the new Rechunk kernel atom (function form xdas.rechunk) restores a target chunk cadence without ever merging across a discontinuity. Eager calls auto-split gappy input into runs and re-join the outputs with the gaps kept in the coordinates — announced, not silent: a warning states how many discontinuities the source has and that state is flushed and reset at each, named by the source's start so a collection walk reports every leaf. The first/last dimension aliases are resolved against the data before any comparison with the chunked dimension, so a kernel built with its documented default no longer skips allocating its seam state; UpSample survives one-sample chunks. The commutation invariant — concat(atom(split(da, anywhere))) equals atom(da) for arbitrary split points, tails included — is now in the test suite for the stateful vocabulary. --- docs/api/atoms.md | 4 + docs/release-notes.md | 4 + tests/test_atoms_runs.py | 378 +++++++++++++++++++++++++++++++++++ tests/test_processing.py | 4 +- tests/test_trigger.py | 18 +- xdas/__init__.py | 2 + xdas/atoms/__init__.py | 3 +- xdas/atoms/core.py | 412 ++++++++++++++++++++++++++++++++++++--- xdas/atoms/kernel.py | 163 +++++++++++++--- xdas/processing/core.py | 28 ++- xdas/trigger.py | 4 +- 11 files changed, 960 insertions(+), 60 deletions(-) create mode 100644 tests/test_atoms_runs.py diff --git a/docs/api/atoms.md b/docs/api/atoms.md index 3f571eb6..f051a860 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -32,7 +32,9 @@ Methods Atom.initialize Atom.initialize_from_state Atom.call + Atom.flush Atom.reset + Atom.iter_chunks Atom.save_state Atom.set_state Atom.load_state @@ -118,6 +120,7 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, hilbert integrate medfilt + rechunk resample sliding_mean_removal taper @@ -152,6 +155,7 @@ parameters, designed by the task atoms from the data at the first call. DownSample LFilter Polyphase + Rechunk SOSFilter UpSample ``` \ No newline at end of file diff --git a/docs/release-notes.md b/docs/release-notes.md index 558a6eb0..566c8253 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,10 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **Continuous-run semantics.** Stateful atoms now understand gaps: every atom judges the seams of its own input stream from the chunk coordinates — a continuous chunk carries state across, a gap or rate change flushes the previous run and restarts (redesigning coefficients on rate changes), an overlap raises, and the `on_discontinuity="reset"|"raise"` policy makes strict runs opt-in. Eager calls auto-split gappy input into runs, process each with a fresh state and re-join the outputs with the gaps kept in the coordinates, so filters never cross discontinuities — and the split is announced: a warning states how many discontinuities the source has and that state is flushed and reset at each. Sequence collections fold through the same seam-aware machinery — `concat(atom(split(da, anywhere)))` equals `atom(da)` for arbitrary split points — and mapping collections map over their leaves. Chunked processing along a dimension now requires a regular coordinate (a declared `sampling_interval`) on it, raising with a pointer to `to_regular()` instead of silently carrying state across unverifiable seams (@atrabattoni). +- **`flush()` lifecycle and the transducer contract.** `call()` now maps one input chunk to zero or more output chunks, and the new `Atom.flush()` drains what remains: buffering atoms emit their tail at the end of the stream, at every seam and at the end of every eager call (`Sequential.flush` cascades codec-drain style, and `process()` drains the pipeline at the end of the stream). Reductions fall out for free: a `call()` that accumulates and returns nothing plus a `flush()` that emits the result gives constant-memory streaming statistics. `Atom.iter_chunks(source)` exposes the whole machinery as a plain generator — the manual chunk loop with buffering, seams and flushing handled inside — and writers now silently drop empty chunks (@atrabattoni). +- **`Rechunk` kernel atom.** `Rechunk({"time": n})` (and its function form `xdas.rechunk`) merges and splits streaming chunks to a target size in samples — a performance knob, e.g. to restore a workable cadence after a decimation shrank the chunks — without ever merging across a discontinuity (@atrabattoni). +- Chunked `DownSample` (and thus the stateful decimation path) no longer drops its trailing samples when the stream length is not a multiple of the factor: the buffered remainder is emitted by the new `flush()` lifecycle. The `"first"`/`"last"` dimension aliases are now resolved against the data before being compared with the chunked dimension, so a kernel built with its documented default no longer skips allocating its seam state, and `UpSample` handles one-sample chunks (@atrabattoni). - **Task atoms with physical units.** A new public processing vocabulary where every parameter keeps its meaning when the sampling rate changes: `Filter` (one atom for all bands — a `(low, high)` corner pair in Hz with `None` opening one end, `ftype="iir"/"fir"`, `zerophase`), `Decimate` and `Resample` (target rate in Hz, both riding the polyphase kernel — the filter-at-full-rate-then-discard chain is never taken), `Integrate` and `Differentiate` (chunk-correct, carrying state across seams), plus whole-record `detrend`, `taper`, `hilbert`, `sliding_mean_removal` and `medfilt` (kernel lengths now in seconds/meters), each refusing chunked execution along its working dimension instead of silently answering wrong. Task atoms default to `dim="time"` and live in `xdas.atoms.tasks` (@atrabattoni). - **Function forms at the top level.** Every task atom generates a top-level function with a synthesized signature and docstring: `xdas.decimate(da, 50.0)` applies eagerly, `xdas.decimate(..., 50.0)` returns the atom, and passing an atom extends a pipeline — so the same code runs eagerly on a slice and chunked on an archive by seeding it with `...` (@atrabattoni). - **Polyphase resampling in a kernel layer.** The exact machine-parameter atoms move to the expert layer `xdas.atoms.kernel` (`LFilter`, `SOSFilter`, `DownSample`, `UpSample`, still importable from `xdas.atoms`), joined by the new `Polyphase` kernel: upsample, FIR filter and downsample fused into a single `scipy.signal.upfirdn` pass that computes only the output samples surviving the decimation and never materialises the zero-stuffed signal (which for `up=4` allocated a four times larger, mostly-zero array). `FIRFilter` is born with `up=`/`down=` and `ResamplePoly` rides it, so the upsample/filter/downsample trio collapses to one child atom — on a 254 MiB chunk that is 2.6× on a decimation by two along distance and 8.7× on a 62.5 → 50 Hz resampling. The taps are cast down to the data precision, so float32 stays float32 instead of being promoted by the filter; a target rate the coordinate resolution cannot represent exactly (100 Hz → 30 Hz is 10/3 ns per sample) declares its residual drift as jitter instead of rejecting its own sampling interval (@atrabattoni). diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py new file mode 100644 index 00000000..99a990a8 --- /dev/null +++ b/tests/test_atoms_runs.py @@ -0,0 +1,378 @@ +""" +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 warnings + +import numpy as np +import pytest + +import xdas as xd +from xdas.atoms import ( + 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: 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 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)) + + def test_gapless_input_is_silent(self, da): + with warnings.catch_warnings(): + warnings.simplefilter("error") + 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 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_processing.py b/tests/test_processing.py index be45aebc..855a1e3e 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -327,12 +327,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") diff --git a/tests/test_trigger.py b/tests/test_trigger.py index ff6b8ed0..5af59e6e 100644 --- a/tests/test_trigger.py +++ b/tests/test_trigger.py @@ -11,7 +11,11 @@ def test_trigger(): 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, + }, }, ) @@ -134,7 +138,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, + }, }, ) @@ -171,7 +179,11 @@ def test_trigger_1d(): 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]}, + "time": { + "tie_indices": [0, 9], + "tie_values": [0.0, 9.0], + "sampling_interval": 1.0, + }, }, ) picks = Trigger(thresh=0.5, dim="time")(cft) diff --git a/xdas/__init__.py b/xdas/__init__.py index f6be2e8b..a714ce50 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -67,6 +67,7 @@ "hilbert", "integrate", "medfilt", + "rechunk", "resample", "sliding_mean_removal", "taper", @@ -85,6 +86,7 @@ testing, virtual, ) +from .atoms.kernel import rechunk from .atoms.tasks import ( decimate, detrend, diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index 05a2b873..77f9f8cb 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -30,6 +30,7 @@ "MLPicker", "Partial", "Polyphase", + "Rechunk", "Resample", "ResamplePoly", "SOSFilter", @@ -44,7 +45,7 @@ from ..trigger import Trigger from .core import Atom, Partial, Sequential, State, as_function, atomized, compose -from .kernel import DownSample, LFilter, Polyphase, SOSFilter, UpSample +from .kernel import DownSample, LFilter, Polyphase, Rechunk, SOSFilter, UpSample from .ml import MLPicker from .signal import FIRFilter, IIRFilter, ResamplePoly from .tasks import Decimate, Differentiate, Filter, Integrate, Resample diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index d7f42412..4ac69e0f 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -10,13 +10,85 @@ import importlib import inspect import re +import warnings from collections.abc import Callable from functools import wraps from typing import Any import numpy as np -from ..core import DataArray, DataCollection, open_datacollection +from ..coordinates import AxisCoordinate +from ..coordinates.core import parse_scalar_delta +from ..core import ( + DataArray, + DataCollection, + DataMapping, + DataSequence, + concat, + open_datacollection, + split, +) + + +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 _flush_through(atoms, **flags): + """ + Codec-drain a linear chain of atoms. + + Flush the first atom and push its tail through the remaining atoms, then + flush the second one, and so on. Tails flow downstream as ordinary data: + each downstream atom folds them into its own state before being flushed + itself. + """ + atoms = list(atoms) + chunks = [] + for index, atom in enumerate(atoms): + tail = atom.flush() + for downstream in atoms[index + 1 :]: + tail = [ + chunk + for out in (downstream(x, **flags) for x in tail) + for chunk in _aschunks(out) + ] + chunks.extend(tail) + return chunks class State: @@ -100,6 +172,10 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): the state of nested atoms. initialized: bool Wether 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 ------- @@ -108,18 +184,26 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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. 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" + 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 @@ -178,15 +262,264 @@ def call(self, x, **flags): return NotImplemented 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. + + A single output chunk is returned bare; otherwise a + :class:`DataSequence` of chunks is returned. + """ chunk_dim = flags.get("chunk_dim", None) self._check_chunk_dim(x, chunk_dim) - if not self.initialized or chunk_dim is None: + if isinstance(x, DataMapping): + if chunk_dim is not None: + raise NotImplementedError( + "chunked processing of mapping collections is not supported: " + "process each leaf with its own atom instance" + ) + return DataCollection( + {key: self(value, **flags) for key, value in x.items()}, + getattr(x, "name", None), + ) + if isinstance(x, DataSequence): + return self._fold(x, flags) + 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 _resolve_dim(self, x): + """Resolve the dimension this atom operates along on *x*, or ``None``.""" + dim = getattr(self, "dim", None) + if not isinstance(x, DataArray) or not isinstance(dim, str): + return None + if dim == "first": + dim = x.dims[0] + elif dim == "last": + dim = x.dims[-1] + return dim if dim in 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), 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" + ) + if seam["delta"] is None: + return None + tolerance = max(seam["tolerance"], info["tolerance"]) + if info["delta"] is not None and np.abs(info["delta"] - seam["delta"]) > ( + tolerance + ): + return "rate" + jump = info["start"] - (seam["end"] + seam["delta"]) + if np.abs(jump) <= tolerance: + return "continuous" + return "gap" if jump > 0 else "overlap" + + def _fold(self, x, flags): + """ + 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. + """ + name = getattr(x, "name", None) + 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(el, **flags) for el in x], name) + flags = flags | {"chunk_dim": dim} + chunks = [] + for el in x: + chunks += _aschunks(self(el, **flags)) + chunks += self.flush() self.reset() - return y + return DataCollection(chunks, name) + chunks = [] + for el in x: + chunks += _aschunks(self(el, **flags)) + return DataCollection(chunks, name) + + def _join(self, chunks, dim): + """Re-join output chunks: one chunk bare, else gap-aware concat or sequence.""" + if len(chunks) == 1: + return chunks[0] + if dim is not None and chunks and all(isinstance(c, DataArray) for c in chunks): + try: + return concat(chunks, dim) + except (TypeError, ValueError): + return DataCollection(chunks) + return DataCollection(chunks) + + 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 + ------ + 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 _check_chunk_dim(self, x, chunk_dim): """Raise if this atom cannot process *x* chunked along *chunk_dim*.""" @@ -269,6 +602,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **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(): @@ -287,7 +621,7 @@ def fresh(self): clone = type(self).__new__(type(self)) Atom.__init__(clone) for name, value in vars(self).items(): - if name in ("_config", "_state", "_atoms"): + if name in ("_config", "_state", "_atoms", "_seam"): continue if name in self._atoms: setattr(clone, name, self._atoms[name].fresh()) @@ -420,10 +754,39 @@ 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. Returns the drained chunks. + """ + flags = {"chunk_dim": self._seam["chunk_dim"]} if self._seam else {} + return _flush_through(self, **flags) + + 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 fresh(self): """Return a stateless clone: each stage cloned, config shared.""" @@ -524,23 +887,26 @@ def __init__( setattr(self, key, value) else: self.kwargs[key] = value - # A whole-record function marked with `_whole_record` refuses chunked - # execution along its working dimension; that dimension is resolved - # from the call arguments so the guard can compare it with the - # chunked one. + # 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) - if dim_arg is not None: - try: - bound = inspect.signature(func).bind_partial(*self.args, **self.kwargs) - bound.apply_defaults() - self.dim = bound.arguments.get(dim_arg) - except (TypeError, ValueError): - self.dim = None + try: + bound = inspect.signature(func).bind_partial(*self.args, **self.kwargs) + bound.apply_defaults() + self.dim = bound.arguments.get("dim") + refuse_dim = bound.arguments.get(dim_arg) if dim_arg else None + 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.dim, chunk_dim, x) + self._refuse_chunked_along(self._refuse_dim, chunk_dim, x) @property def stateful(self): diff --git a/xdas/atoms/kernel.py b/xdas/atoms/kernel.py index 2947fb90..cbddcac5 100644 --- a/xdas/atoms/kernel.py +++ b/xdas/atoms/kernel.py @@ -8,7 +8,7 @@ processing equals unchunked processing. Includes :class:`LFilter`, :class:`SOSFilter`, :class:`DownSample`, -:class:`UpSample`, :class:`Polyphase`. +:class:`UpSample`, :class:`Polyphase`, :class:`Rechunk`. """ import math @@ -20,7 +20,7 @@ from ..coordinates.core import parse_scalar_delta from ..core import DataArray, concat, split from ..parallel import parallelize -from .core import Atom, State +from .core import Atom, State, atomized def _along(axis, ndim, slc): @@ -56,11 +56,14 @@ def __init__(self, b, a, dim="last", parallel=None): 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: + # `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._resolve_dim(da) or self.dim + if 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() + n_sections if name == dim else size for name, size in da.sizes.items() ) self.zi = State(np.zeros(shape)) else: @@ -106,7 +109,9 @@ def __init__(self, sos, dim="last", parallel=None): 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: + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._resolve_dim(da) or self.dim + if dim == chunk_dim: n_sections = self.sos.shape[0] shape = (n_sections,) + tuple( 2 if index == self.axis else element @@ -151,7 +156,9 @@ def __init__(self, factor, dim="last"): def initialize(self, da, chunk_dim=None, **flags): """Initialise the carry-over buffer for chunked operation.""" - if chunk_dim == self.dim: + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._resolve_dim(da) or self.dim + if chunk_dim == dim: self.buffer = State(da.isel({self.dim: slice(0, 0)})) else: self.buffer = State(None) @@ -167,6 +174,20 @@ def call(self, da, **flags): 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): """ @@ -192,12 +213,15 @@ 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._resolve_dim(da) or self.dim shape = tuple( - self.factor * size if dim == self.dim else size + self.factor * size if dim == name else size for dim, size in da.sizes.items() ) slc = tuple( - slice(None, None, self.factor) if dim == self.dim else slice(None) + slice(None, None, self.factor) if dim == name else slice(None) for dim in da.dims ) data = np.zeros(shape, dtype=da.dtype) @@ -206,13 +230,22 @@ def call(self, da, **flags): else: data[slc] = da.values coords = da.coords.copy() - delta = get_sampling_interval(da, self.dim, cast=False) + delta = get_sampling_interval(da, name, 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 + 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 @@ -225,7 +258,7 @@ def call(self, da, **flags): # 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) + coords[name] = Coordinate(data_coord, name) return DataArray(data, coords, da.dims, da.name, da.attrs) @@ -276,8 +309,7 @@ class Polyphase(Atom): >>> eager = Polyphase(taps, up=2, down=5, dim="time")(da) >>> atom = Polyphase(taps, up=2, down=5, dim="time") - >>> outs = [atom(chunk, chunk_dim="time") for chunk in xd.split(da, 7, "time")] - >>> chunked = xd.concat(outs, "time") + >>> chunked = xd.concat(list(atom.iter_chunks(xd.split(da, 7, "time"))), "time") >>> bool(np.allclose(chunked.values, eager.values)) True @@ -318,9 +350,11 @@ def initialize(self, da, chunk_dim=None, **flags): f"{np.asarray(self.taps).size}" ) self.axis = State(da.get_axis_num(self.dim)) - if self.dim == chunk_dim: + # Resolve the "first"/"last" alias before comparing (see `LFilter`). + dim = self._resolve_dim(da) or self.dim + if dim == chunk_dim: shape = tuple( - self._history_size() if name == self.dim else size + self._history_size() if name == dim else size for name, size in da.sizes.items() ) self.buffer = State(np.zeros(shape, dtype=da.dtype)) @@ -394,8 +428,11 @@ def _keep_history(self, values, axis): def _coords(self, da, first, stop, start): """Build the output coordinates on the resampled, delay-corrected grid.""" - coord = da.coords[self.dim] - delta = get_sampling_interval(da, self.dim, cast=False) + # The "first"/"last" alias would name a new dimension if it reached the + # `Coordinate` built below, so resolve it here. + name = self._resolve_dim(da) or self.dim + 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. @@ -419,7 +456,7 @@ def _coords(self, da, first, stop, start): data["sampling_interval"] = step data["tolerance"] = base + drift coords = da.coords.copy() - coords[self.dim] = Coordinate(data, self.dim) + coords[name] = Coordinate(data, name) return coords def _upsampled(self, count, delta): @@ -427,3 +464,83 @@ def _upsampled(self, count, delta): 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._resolve_dim(da) or self.dim + 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/processing/core.py b/xdas/processing/core.py index 99ebafe5..f85077ee 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -21,6 +21,7 @@ from watchdog.events import FileSystemEventHandler from watchdog.observers import Observer +from ..atoms.core import _aschunks from ..core import DataArray, concat, open_dataarray from .monitor import Monitor @@ -229,8 +230,11 @@ def process(atom, data_loader, data_writer): 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. + `atom` function to each chunk, and writes the processed data using the + `data_writer`. An atom may emit zero or more output chunks per input chunk + (seam tails, rechunking, reductions); at the end of the stream the atom is + flushed so buffering atoms emit their remainder. The progress of the + processing is monitored using a `Monitor` object. """ if hasattr(atom, "reset"): @@ -245,9 +249,13 @@ def process(atom, data_loader, data_writer): monitor.tic("proc") result = atom(chunk, chunk_dim=data_loader.chunk_dim) monitor.tic("write") - data_writer.write(result) + for out in _aschunks(result): + data_writer.write(out) monitor.toc(chunk.nbytes) monitor.tic("read") + if hasattr(atom, "flush"): + for out in atom.flush(): + data_writer.write(out) monitor.close() return data_writer.result() @@ -502,10 +510,13 @@ 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() @@ -578,10 +589,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) @@ -603,11 +616,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() diff --git a/xdas/trigger.py b/xdas/trigger.py index 4f07fbac..d736d27f 100644 --- a/xdas/trigger.py +++ b/xdas/trigger.py @@ -46,7 +46,7 @@ class Trigger(Atom): ... 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}, ... }, ... ) @@ -265,7 +265,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}, ... }, ... ) From d66424c3535056c649a6bce8a614a77587f6e1e6 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:24:34 +0200 Subject: [PATCH 11/48] assert_chunk_invariant: the evidence for chunk safety MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chunk-safety is a claim an atom makes about itself; xdas.testing.assert_chunk_invariant is the evidence. It runs a pipeline once eagerly and once streamed chunk by chunk and asserts the two agree — shapes, values, coordinates, and pick tables compared as sets of rows since eager and chunked walks order them differently. coord_atol admits an explicitly declared sub-sample coordinate drift where rational resampling reconstructs its grid segment by segment. The invariant quantifies over cuts and gaps: the same stream is re-chunked at derived non-divisor sizes so the boundaries land elsewhere (cuts=, explicit dicts accepted), and inject_gaps places real discontinuities in the input first so seam resets are exercised at chunk boundaries that do not line up with them. A negative control is part of the test suite: a pipeline that is not chunk-invariant fails loudly rather than passing. --- docs/api/testing.md | 2 + docs/release-notes.md | 1 + tests/test_testing.py | 148 +++++++++++++++++++++++ xdas/testing.py | 266 +++++++++++++++++++++++++++++++++++++++++- 4 files changed, 416 insertions(+), 1 deletion(-) 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/release-notes.md b/docs/release-notes.md index 566c8253..c8e59b33 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,7 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **`xdas.testing.assert_chunk_invariant`.** The chunk-safety story in one call: run a pipeline eagerly and streamed and assert the two agree — values, coordinates and all. The invariant is quantified over *cuts* (the same stream re-chunked at derived non-divisor sizes, so boundaries land elsewhere) and over *gaps* (`xdas.testing.inject_gaps` places real discontinuities in the input first, so seam resets are exercised at boundaries that do not line up with them). It is both the CI harness for every stateful atom xdas ships and the tool to run on your own pipelines before trusting them chunked (@atrabattoni). - **Continuous-run semantics.** Stateful atoms now understand gaps: every atom judges the seams of its own input stream from the chunk coordinates — a continuous chunk carries state across, a gap or rate change flushes the previous run and restarts (redesigning coefficients on rate changes), an overlap raises, and the `on_discontinuity="reset"|"raise"` policy makes strict runs opt-in. Eager calls auto-split gappy input into runs, process each with a fresh state and re-join the outputs with the gaps kept in the coordinates, so filters never cross discontinuities — and the split is announced: a warning states how many discontinuities the source has and that state is flushed and reset at each. Sequence collections fold through the same seam-aware machinery — `concat(atom(split(da, anywhere)))` equals `atom(da)` for arbitrary split points — and mapping collections map over their leaves. Chunked processing along a dimension now requires a regular coordinate (a declared `sampling_interval`) on it, raising with a pointer to `to_regular()` instead of silently carrying state across unverifiable seams (@atrabattoni). - **`flush()` lifecycle and the transducer contract.** `call()` now maps one input chunk to zero or more output chunks, and the new `Atom.flush()` drains what remains: buffering atoms emit their tail at the end of the stream, at every seam and at the end of every eager call (`Sequential.flush` cascades codec-drain style, and `process()` drains the pipeline at the end of the stream). Reductions fall out for free: a `call()` that accumulates and returns nothing plus a `flush()` that emits the result gives constant-memory streaming statistics. `Atom.iter_chunks(source)` exposes the whole machinery as a plain generator — the manual chunk loop with buffering, seams and flushing handled inside — and writers now silently drop empty chunks (@atrabattoni). - **`Rechunk` kernel atom.** `Rechunk({"time": n})` (and its function form `xdas.rechunk`) merges and splits streaming chunks to a target size in samples — a performance knob, e.g. to restore a workable cadence after a decimation shrank the chunks — without ever merging across a discontinuity (@atrabattoni). diff --git a/tests/test_testing.py b/tests/test_testing.py index 42a1ec11..320e6260 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,148 @@ 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_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/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, + ) From 3209d2066cda1d8e251310c3bbc4ae3381227825 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:29:19 +0200 Subject: [PATCH 12/48] STFT streams exactly the eager frames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The spectral vocabulary joins the task-atom route: STFT takes its window length and hop in physical units, both snapped — the window to the next fast FFT size of the target so transforms stay efficient whatever the sampling rate, the hop to a whole number of samples — so the actual grid can differ from the request, and the docstring says so. scaling= chooses "spectrum" (peak amplitudes) or "psd", so np.abs(stft)**2 composes to an exact spectrogram; an expert nfft zero-pads the windowed frames. Frames start at the first sample and advance by the hop; only fully computable frames are ever emitted. Chunk by chunk, the unconsumed tail is buffered across chunks and dropped at gaps and at the end of the stream, so chunked processing emits exactly the frames of the eager transform and no frame ever spans a discontinuity — held under assert_chunk_invariant over cuts and gaps. One-sided spectra for real data, centered two-sided for complex; scipy.signal.ShortTimeFFT does the design and scaling internally. xd.stft is the function form. The xdas.fft functions (fft, rfft, ifft, irfft) now declare whole-record semantics: used as atoms in a chunked pipeline they raise along the transformed dimension instead of silently computing one transform per chunk, and Partial resolves the {input_dim: output_dim} mapping form of dim, so transforming along another dimension than the chunked one keeps working. --- docs/api/atoms.md | 2 + docs/release-notes.md | 2 + tests/test_atoms_runs.py | 26 ++++++ tests/test_atoms_tasks.py | 120 ++++++++++++++++++++++++ tests/test_fft.py | 28 ++++++ xdas/__init__.py | 2 + xdas/atoms/__init__.py | 3 +- xdas/atoms/core.py | 7 +- xdas/atoms/tasks.py | 186 +++++++++++++++++++++++++++++++++++++- xdas/fft.py | 5 + 10 files changed, 378 insertions(+), 3 deletions(-) diff --git a/docs/api/atoms.md b/docs/api/atoms.md index f051a860..75feccf4 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -103,6 +103,7 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, Filter Integrate Resample + STFT ``` ```{eval-rst} @@ -123,6 +124,7 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, rechunk resample sliding_mean_removal + stft taper ``` diff --git a/docs/release-notes.md b/docs/release-notes.md index c8e59b33..325aaf0a 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,8 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **`STFT`.** The spectral vocabulary joins the task-atom route: `STFT` streams complex frames with window length and hop in physical units — both are snapped, the window to the next fast FFT size of the target and the hop to a whole sample count — with an expert `nfft` to zero-pad and a `scaling=` of `"spectrum"` or `"psd"`, so `np.abs(stft)**2` composes to an exact spectrogram. Only fully computable frames are ever emitted: the unconsumed tail is buffered across chunks and dropped at gaps, so chunked processing emits exactly the eager frames and no frame ever spans a discontinuity. Built on `scipy.signal.ShortTimeFFT` internally, with the `xdas.stft` function form at the top level (@atrabattoni). +- The `xdas.fft` functions (`fft`, `rfft`, `ifft`, `irfft`) now declare whole-record semantics: used as atoms in a chunked pipeline they raise along the transformed dimension instead of silently computing one transform per chunk. Transforming along another dimension than the chunked one keeps working (@atrabattoni). - **`xdas.testing.assert_chunk_invariant`.** The chunk-safety story in one call: run a pipeline eagerly and streamed and assert the two agree — values, coordinates and all. The invariant is quantified over *cuts* (the same stream re-chunked at derived non-divisor sizes, so boundaries land elsewhere) and over *gaps* (`xdas.testing.inject_gaps` places real discontinuities in the input first, so seam resets are exercised at boundaries that do not line up with them). It is both the CI harness for every stateful atom xdas ships and the tool to run on your own pipelines before trusting them chunked (@atrabattoni). - **Continuous-run semantics.** Stateful atoms now understand gaps: every atom judges the seams of its own input stream from the chunk coordinates — a continuous chunk carries state across, a gap or rate change flushes the previous run and restarts (redesigning coefficients on rate changes), an overlap raises, and the `on_discontinuity="reset"|"raise"` policy makes strict runs opt-in. Eager calls auto-split gappy input into runs, process each with a fresh state and re-join the outputs with the gaps kept in the coordinates, so filters never cross discontinuities — and the split is announced: a warning states how many discontinuities the source has and that state is flushed and reset at each. Sequence collections fold through the same seam-aware machinery — `concat(atom(split(da, anywhere)))` equals `atom(da)` for arbitrary split points — and mapping collections map over their leaves. Chunked processing along a dimension now requires a regular coordinate (a declared `sampling_interval`) on it, raising with a pointer to `to_regular()` instead of silently carrying state across unverifiable seams (@atrabattoni). - **`flush()` lifecycle and the transducer contract.** `call()` now maps one input chunk to zero or more output chunks, and the new `Atom.flush()` drains what remains: buffering atoms emit their tail at the end of the stream, at every seam and at the end of every eager call (`Sequential.flush` cascades codec-drain style, and `process()` drains the pipeline at the end of the stream). Reductions fall out for free: a `call()` that accumulates and returns nothing plus a `flush()` that emits the result gives constant-memory streaming statistics. `Atom.iter_chunks(source)` exposes the whole machinery as a plain generator — the manual chunk loop with buffering, seams and flushing handled inside — and writers now silently drop empty chunks (@atrabattoni). diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py index 99a990a8..13bb4c97 100644 --- a/tests/test_atoms_runs.py +++ b/tests/test_atoms_runs.py @@ -17,6 +17,7 @@ import xdas as xd from xdas.atoms import ( + STFT, Atom, Decimate, DownSample, @@ -63,6 +64,7 @@ class TestCommutation: 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]), ] @@ -322,6 +324,30 @@ def flush(self): 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) diff --git a/tests/test_atoms_tasks.py b/tests/test_atoms_tasks.py index b57530cc..327163bb 100644 --- a/tests/test_atoms_tasks.py +++ b/tests/test_atoms_tasks.py @@ -5,7 +5,9 @@ import xdas as xd import xdas.signal as xs +import xdas.spectral from xdas.atoms import ( + STFT, Decimate, Differentiate, Filter, @@ -15,6 +17,7 @@ Sequential, ) from xdas.synthetics import wavelet_wavefronts +from xdas.testing import dummy def through_chunks(atom, da, nchunk=6, dim="time"): @@ -293,3 +296,120 @@ def test_chunked_along_other_dim_passes(self): 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 + ) diff --git a/tests/test_fft.py b/tests/test_fft.py index 0c9977e5..a30172ee 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,30 @@ 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_is_conservative(self): + da = xd.testing.dummy() + atom = xfft.fft(...) + with pytest.raises(ValueError, match="whole record"): + atom(da.isel(time=slice(0, 50)), chunk_dim="time") + + 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/xdas/__init__.py b/xdas/__init__.py index a714ce50..372564b4 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -70,6 +70,7 @@ "rechunk", "resample", "sliding_mean_removal", + "stft", "taper", ] @@ -97,6 +98,7 @@ medfilt, resample, sliding_mean_removal, + stft, taper, ) from .coordinates import ( diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index 77f9f8cb..c4fb1a4f 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -18,6 +18,7 @@ """ __all__ = [ + "STFT", "Atom", "Decimate", "Differentiate", @@ -48,4 +49,4 @@ from .kernel import DownSample, LFilter, Polyphase, Rechunk, SOSFilter, UpSample from .ml import MLPicker from .signal import FIRFilter, IIRFilter, ResamplePoly -from .tasks import Decimate, Differentiate, Filter, Integrate, Resample +from .tasks import STFT, Decimate, Differentiate, Filter, Integrate, Resample diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index 4ac69e0f..b69ddf73 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -896,7 +896,12 @@ def __init__( try: bound = inspect.signature(func).bind_partial(*self.args, **self.kwargs) bound.apply_defaults() - self.dim = bound.arguments.get("dim") + 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 = dim refuse_dim = bound.arguments.get(dim_arg) if dim_arg else None except (TypeError, ValueError): self.dim = None diff --git a/xdas/atoms/tasks.py b/xdas/atoms/tasks.py index c87f5fb7..62e776cd 100644 --- a/xdas/atoms/tasks.py +++ b/xdas/atoms/tasks.py @@ -15,13 +15,17 @@ """ import numpy as np +import scipy.fft +from scipy.signal import ShortTimeFFT, get_window from ..coordinates import get_sampling_interval -from ..core import concat +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", @@ -36,6 +40,7 @@ "medfilt", "resample", "sliding_mean_removal", + "stft", "taper", ] @@ -514,8 +519,187 @@ def medfilt(da, kernel): 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/fft.py b/xdas/fft.py index 47893c46..44380454 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() 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() 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() 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() def irfft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the inverse of `rfft`. From 717d11404e895a0a8148ddc95fe1ad0d8a7ed80c Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:42:21 +0200 Subject: [PATCH 13/48] pin the edges of tracing, seams and folding Tests for the corners the main suites walked past: in-place ufuncs trace out of place, out= to a foreign atom and non-call ufunc methods raise, right_shift with the atom on the left is an ordinary traced ufunc; fresh() recurses into nested class atoms and the refusal helper is conservative on unresolvable aliases and checks the keys of kernel dicts; empty chunks are skipped, a one-sample chunk of a sampled coordinate inherits the stream's rate at the seam, single-sample streams have nothing to judge, first/last aliases resolve on eager calls, dimensionless atoms map collections eagerly, folds work chunked, unconcatenatable outputs fall back to a sequence; writers drop empty chunks; the cut derivation stops when no new size exists; UpSample survives a one-sample record and Polyphase an empty one. --- tests/test_atoms.py | 65 ++++++++++++++++++++++++++++++++++++ tests/test_atoms_runs.py | 71 ++++++++++++++++++++++++++++++++++++++++ tests/test_processing.py | 8 +++++ tests/test_testing.py | 12 +++++++ 4 files changed, 156 insertions(+) diff --git a/tests/test_atoms.py b/tests/test_atoms.py index 40d9f55a..13b504bd 100644 --- a/tests/test_atoms.py +++ b/tests/test_atoms.py @@ -344,6 +344,18 @@ def test_too_few_taps(self): 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 TestMLPicker: @pytest.mark.slow @@ -723,6 +735,30 @@ def test_equality_is_identity(self): 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_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.""" @@ -803,6 +839,35 @@ def test_fresh_recurses_into_sequences(self): assert clone[0].func is seq[0].func +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") diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py index 13bb4c97..17c06841 100644 --- a/tests/test_atoms_runs.py +++ b/tests/test_atoms_runs.py @@ -215,6 +215,77 @@ def test_chunked_internal_gap_splits(self, da): 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_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_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)) diff --git a/tests/test_processing.py b/tests/test_processing.py index 855a1e3e..403c9dd3 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -242,6 +242,14 @@ 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. diff --git a/tests/test_testing.py b/tests/test_testing.py index 320e6260..7b9c1757 100644 --- a/tests/test_testing.py +++ b/tests/test_testing.py @@ -165,6 +165,18 @@ def test_inject_gaps_refuses_to_eat_the_record(self): 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. From 60f215b93a9026543b1e70a2f461617caed0fe4f Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 13:49:51 +0200 Subject: [PATCH 14/48] process() is the dispatch boundary of chunked processing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit process() becomes a method on every atom: pipeline.process(da, out="results/") infers both ends. Sources dispatch on the input value (get_source): an in-memory DataArray runs eagerly (or chunk by chunk with chunks=), a virtual one streams through a DataArrayLoader whose chunks="auto" aligns boundaries to the storage blocking (tile extents, per-file extents) merged up to a byte budget, a path/directory/glob opens with open_mfdataarray (multi-acquisition collections chain one loader per run), tcp:// subscribes over ZeroMQ through a scheme registry, and any iterable of chunks is consumed as is — the source contract is iteration plus optional chunk_dim/nbytes/unbounded. Sinks dispatch on the out spec crossed with the first output chunk (get_writer), deferring writer creation to what the pipeline actually emits: directories store DataArray chunks joined along the chunked dimension (or SDS archives for Streams), *.csv appends DataFrames, tcp:// publishes, out=None accumulates and returns the joined result, writer instances pass through, and no empty outputs are ever created. The historical process(atom, loader, writer) form keeps working. Realtime is named: xd.watch(path, engine=...) wraps RealTimeLoader, and unbounded sources get streaming semantics — no byte total on the monitor, a clean KeyboardInterrupt that flushes the pipeline and returns the writer result, and until= to stop at a coordinate value (inclusive, truncating the last chunk). Discontinuities are announced at the boundary too: a chunked source warns once upfront with the count read off its coordinate (no data touched), and a realtime source — which cannot be inspected upfront — warns at each seam as it arrives. The new memory_limit configuration entry (default 8 GiB) guards the two footguns: an eager call on a huge virtual array and an out=None accumulation that outgrows the limit both raise with the estimated size and a pointer to .process(out=...). --- docs/api/atoms.md | 1 + docs/api/processing.md | 4 + docs/release-notes.md | 3 + tests/test_process.py | 505 +++++++++++++++++++++++++++++++++ xdas/__init__.py | 3 + xdas/atoms/core.py | 92 ++++++ xdas/config.py | 5 +- xdas/processing/__init__.py | 12 +- xdas/processing/core.py | 540 ++++++++++++++++++++++++++++++++---- 9 files changed, 1106 insertions(+), 59 deletions(-) create mode 100644 tests/test_process.py diff --git a/docs/api/atoms.md b/docs/api/atoms.md index 75feccf4..a1deea09 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -34,6 +34,7 @@ Methods Atom.call Atom.flush Atom.reset + Atom.process Atom.iter_chunks Atom.save_state Atom.set_state diff --git a/docs/api/processing.md b/docs/api/processing.md index e4774ae1..c6c701af 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 diff --git a/docs/release-notes.md b/docs/release-notes.md index 325aaf0a..893fa86e 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,6 +9,9 @@ - **`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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). +- **`process()` with source and sink auto-dispatch.** `process()` is now a method on every atom and a dispatch boundary: `pipeline.process(da, out="results/")` infers both ends. Sources dispatch on the input value — an in-memory `DataArray` runs eagerly (or chunk by chunk with `chunks=`), a virtual one streams through a loader with storage-aligned `chunks="auto"`, a path/directory/glob opens with `open_mfdataarray`, `"tcp://..."` subscribes over ZeroMQ, and any iterable of chunks (a generator, a loader) is consumed as is. Sinks dispatch on the out spec crossed with the first output chunk, so writer creation is deferred to what the pipeline actually emits: a directory stores `DataArray` chunks joined along the chunked dimension (or an SDS archive for `Stream` chunks), `*.csv` appends DataFrames, `"tcp://..."` publishes, `out=None` accumulates and returns the joined result, and a configured writer instance passes through. A chunked source with discontinuities announces them upfront — one warning with the count, read off the source coordinate before any data. The historical `process(atom, loader, writer)` form keeps working unchanged (@atrabattoni). +- **`xdas.watch` and unbounded sources.** Realtime is now *named*: `pipeline.process(xd.watch("/incoming", engine=...), out=...)` watches a directory forever, and a bare directory path always means "process what is there". Unbounded sources (watch, ZMQ subscriptions) get streaming semantics — throughput-style progress, a clean `KeyboardInterrupt` that flushes the pipeline and returns the writer result, `until=` to stop at a coordinate value (inclusive, truncating the last chunk), and a warning at each seam as it arrives, since a realtime source cannot be inspected upfront (@atrabattoni). +- **Memory guards.** The new `"memory_limit"` configuration entry (default 8 GiB) makes footguns loud: an eager call on a huge virtual array and an `out=None` accumulation that outgrows the limit both raise with the estimated size and a pointer to `.process(out=...)` (@atrabattoni). - **`STFT`.** The spectral vocabulary joins the task-atom route: `STFT` streams complex frames with window length and hop in physical units — both are snapped, the window to the next fast FFT size of the target and the hop to a whole sample count — with an expert `nfft` to zero-pad and a `scaling=` of `"spectrum"` or `"psd"`, so `np.abs(stft)**2` composes to an exact spectrogram. Only fully computable frames are ever emitted: the unconsumed tail is buffered across chunks and dropped at gaps, so chunked processing emits exactly the eager frames and no frame ever spans a discontinuity. Built on `scipy.signal.ShortTimeFFT` internally, with the `xdas.stft` function form at the top level (@atrabattoni). - The `xdas.fft` functions (`fft`, `rfft`, `ifft`, `irfft`) now declare whole-record semantics: used as atoms in a chunked pipeline they raise along the transformed dimension instead of silently computing one transform per chunk. Transforming along another dimension than the chunked one keeps working (@atrabattoni). - **`xdas.testing.assert_chunk_invariant`.** The chunk-safety story in one call: run a pipeline eagerly and streamed and assert the two agree — values, coordinates and all. The invariant is quantified over *cuts* (the same stream re-chunked at derived non-divisor sizes, so boundaries land elsewhere) and over *gaps* (`xdas.testing.inject_gaps` places real discontinuities in the input first, so seam resets are exercised at boundaries that do not line up with them). It is both the CI harness for every stateful atom xdas ships and the tool to run on your own pipelines before trusting them chunked (@atrabattoni). diff --git a/tests/test_process.py b/tests/test_process.py new file mode 100644 index 00000000..5a45ed80 --- /dev/null +++ b/tests/test_process.py @@ -0,0 +1,505 @@ +""" +Dispatch tests for the `process()` boundary: sources, sinks, guards. + +The governing invariant (plan §6) is that `pipeline.process(da, out=None)` +equals `pipeline(da)` whatever the source and chunking, and that sinks are +resolved from *(out spec × first-chunk type)* with writer instantiation +deferred to the first output chunk. +""" + +import os +import threading +import time + +import numpy as np +import pandas as pd +import pytest + +import xdas as xd +import xdas.processing as xp +from xdas.atoms import Partial +from xdas.atoms.core import _join_chunks +from xdas.config import Config +from xdas.processing.core import _auto_chunks, _ChainSource, _to_human + + +@pytest.fixture +def da(): + # 101 samples: an awkward length that exercises the flushed tails. + return xd.testing.dummy(shape=(101, 5)) + + +@pytest.fixture +def pipeline(): + return xd.decimate(..., target=25.0) >> 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): + address = f"tcp://localhost:{xd.io.get_free_port()}" + source = 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_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) + + def test_realtime_continuous_stream_is_silent(self, da): + import warnings + + source = self.Source(list(xd.split(da, 4, "time"))) + atom = Partial(np.square) + with warnings.catch_warnings(): + warnings.simplefilter("error") + 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}) + + 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): + 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 = xp.ZMQPublisher(address) + source = xp.get_source(address) + + def publish(): + time.sleep(0.1) + for packet in packets: + time.sleep(0.001) + 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) diff --git a/xdas/__init__.py b/xdas/__init__.py index 372564b4..d3b8ad4e 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -72,6 +72,8 @@ "sliding_mean_removal", "stft", "taper", + # streaming + "watch", ] from . import ( @@ -142,3 +144,4 @@ trim_overlaps, ) from .core.methods import * +from .processing.core import watch diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index b69ddf73..9da436a3 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -16,7 +16,9 @@ from typing import Any 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 ( @@ -28,6 +30,7 @@ open_datacollection, split, ) +from ..virtual import VirtualBackend def _announce_splits(x, dim, count): @@ -68,6 +71,32 @@ def _aschunks(value): ] +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 _flush_through(atoms, **flags): """ Codec-drain a linear chain of atoms. @@ -281,6 +310,19 @@ def __call__(self, x, **flags): """ chunk_dim = flags.get("chunk_dim", 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") + ): + raise ValueError( + f"this eager call would load the full virtual array " + f"(~{x.nbytes / 2**30:.1f} GiB, above the 'memory_limit' " + "configuration entry) in memory: stream it chunk by chunk " + "with `.process(da, out=...)` instead, or raise the limit " + "with `xdas.config.set('memory_limit', ...)`" + ) if isinstance(x, DataMapping): if chunk_dim is not None: raise NotImplementedError( @@ -521,6 +563,56 @@ def iter_chunks(self, source, chunk_dim=None): yield from self.flush() self.reset() + def process(self, source, out=None, chunks=None, until=None): + """ + 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=...)``. + + Parameters + ---------- + source : DataArray, str, Path, iterable or loader + What to process: an in-memory or virtual :class:`DataArray`, 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. + 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. + + 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) + def _check_chunk_dim(self, x, chunk_dim): """Raise if this atom cannot process *x* chunked along *chunk_dim*.""" diff --git a/xdas/config.py b/xdas/config.py index c91ea3e3..707009d6 100644 --- a/xdas/config.py +++ b/xdas/config.py @@ -11,7 +11,10 @@ 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(), + "memory_limit": 8 * 2**30, + } def get(key): 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 f85077ee..a548fd22 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -3,10 +3,14 @@ Includes :class:`DataArrayLoader`, :class:`DataArrayWriter`, :class:`DataFrameWriter`, :class:`StreamWriter`, :class:`ZMQPublisher`, -:class:`ZMQSubscriber`, :class:`RealTimeLoader`, and :func:`process`. +:class:`ZMQSubscriber`, :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 warnings from collections import deque from concurrent.futures import CancelledError, ThreadPoolExecutor from glob import glob @@ -21,10 +25,17 @@ from watchdog.events import FileSystemEventHandler from watchdog.observers import Observer -from ..atoms.core import _aschunks -from ..core import DataArray, concat, open_dataarray +from .. import config +from ..atoms.core import _announce_splits, _aschunks, _join_chunks +from ..coordinates import AxisCoordinate +from ..coordinates.core import parse_scalar_delta +from ..core import DataArray, DataSequence, concat, open_dataarray, open_mfdataarray +from ..virtual import TileArray, VirtualBackend, VirtualStack from .monitor import Monitor +AUTO_CHUNK_NBYTES = 256 * 2**20 +"""Target in-memory chunk size (in bytes) for ``chunks="auto"``.""" + class _RayFuture: """A future handed out by :class:`ProcessPool`, resolved via the object store.""" @@ -207,57 +218,452 @@ def _dump(chunk, path, encoding): return open_dataarray(path) -def process(atom, data_loader, data_writer): +def process(atom, source, out=None, chunks=None, until=None): """ - Execute a chunked processing pipeline. + 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 : 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. + 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, 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 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. + 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. 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). 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`. An atom may emit zero or more output chunks per input chunk - (seam tails, rechunking, reductions); at the end of the stream the atom is - flushed so buffering atoms emit their remainder. 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. """ + 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. + result = atom(source) + 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) + return writer.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 + 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") - for out in _aschunks(result): - data_writer.write(out) - 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 out in atom.flush(): - data_writer.write(out) + for chunk_out in atom.flush(): + write(chunk_out) monitor.close() - return data_writer.result() + return writer.result() if writer is not None else None + + +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(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 += getattr(chunk, "nbytes", 0) + 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 _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: @@ -271,11 +677,13 @@ 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 @@ -315,38 +723,45 @@ class DataArrayLoader: 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 _select(self, idx): """Return chunk *idx* as a lazy selection: the manifest, not the data.""" - start = idx * self.chunk_size - end = (idx + 1) * self.chunk_size + 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 @@ -399,6 +814,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 @@ -931,6 +1349,9 @@ class ZMQSubscriber: >>> assert da.equals(da) """ + chunk_dim = "time" + unbounded = True + def __init__(self, address): self.address = address self._context = zmq.Context() @@ -986,3 +1407,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.""" From d8aedd588250fa219ded878eeacd3103e05c39f8 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 14:11:33 +0200 Subject: [PATCH 15/48] pin the edges of the realtime seam announcements Tests for the three corners of process() the main suites walked past: a source whose chunked coordinate is dense gets no upfront discontinuity scan, a realtime chunk that does not carry the chunked dimension leaves the seam information untouched, and a realtime one-sample chunk of a sampled coordinate inherits the stream's rate so the seam after it is still judged correctly. --- tests/test_process.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_process.py b/tests/test_process.py index 5a45ed80..117eb92f 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -459,6 +459,45 @@ 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}) + 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. + import warnings + + dense = xd.testing.dummy(shape=(52, 5), ctype="dense") + atom = Partial(np.square) + with warnings.catch_warnings(): + warnings.simplefilter("error") + result = atom.process(dense, chunks={"time": 20}) + assert np.allclose(result.values, np.square(dense).values) + + 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. + import warnings + + 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) + with warnings.catch_warnings(): + warnings.simplefilter("error") + atom.process(source) + + 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. + import warnings + + 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) + with warnings.catch_warnings(): + warnings.simplefilter("error") + atom.process(source) + def test_watch_is_a_realtime_loader(self, da, tmp_path): loader = xd.watch(tmp_path) try: From b6f47cd11d48d936db27ba1dbf837c2728093241 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 14:23:19 +0200 Subject: [PATCH 16/48] fixture: a fake SeisBench model spanning the weight-set axes Everything a SeisBench picker does is a property of the weight set rather than the architecture, so tests/fakemodel.py makes that an executable contract instead of a warning in a document. FakeModel is a real WaveformModel subclass with no weights and a closed-form forward pass; WEIGHT_SETS holds five archetypes drawn from the cached PhaseNet metadata, between them spanning ENZ/ZNE/Z12H, 3 and 4 input channels, the NPS/PSN label flip, 50 and 100 Hz, blinding and overlap declared and absent, per-phase thresholds declared and absent, and no filter, a flat filter and a per-channel one. tests/conftest.py exposes the factory as a fake_model fixture; import the module directly to parametrise on it. --- tests/conftest.py | 13 +++ tests/fakemodel.py | 259 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 272 insertions(+) create mode 100644 tests/fakemodel.py diff --git a/tests/conftest.py b/tests/conftest.py index 969d28e8..607283bd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,19 @@ def pytest_configure(config): xdas.config.set("n_workers", 1) +@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/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() From be7de66990e8143740024270c39ca60d82ba4c63 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 14:39:17 +0200 Subject: [PATCH 17/48] rework MLPicker into Annotate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Everything a SeisBench picker reads is a property of the weight set, so the atom now reads it there: the overlap (in SeisBench's fraction-or- samples form), the stacking rule, the blinding and every annotate argument come off the model instance, overridable per call, with the model's own annotate_batch_pre/post driving normalisation and blinding. The component dimension is found by its labels ending with distinct letters of component_order — never by its name — with the flexible horizontal matching, and component_strategy spans SeisBench's range: auto, clone, pad, a named slot, strict. The output keeps the input's order among the batch dimensions but comes out sample-last, (..., 'phase', dim): the characteristic function of one phase of one channel is contiguous, which is the layout its consumers reduce along. The end-aligned final window SeisBench appends is emitted at flush(), so the output spans the input and stays chunk-invariant. A model whose annotate_batch_post breaks the (batch, samples, classes) stacking contract is named instead of surfacing as a broadcast error. Chunked along its own dimension the sliding window carries across chunks; chunked along another dimension the carry-over would leak one chunk's time tail onto the next chunk's other lanes, so such a chunk is a whole record run from a fresh state and settled on the spot — and the component dimension is refused as a chunk axis, since the model reads every component of a window at once. MLPicker and xdas.mlpicker stay as DeprecationWarning aliases until 0.4. The value pin holds: the numbers today's MLPicker produced on the DAS layouts are unchanged, asserted through dimension names so they outlive the layout change. --- docs/release-notes.md | 1 + tests/test_atoms.py | 71 +-- tests/test_atoms_ml.py | 984 +++++++++++++++++++++++++++++++++++++++++ xdas/__init__.py | 3 + xdas/atoms/__init__.py | 5 +- xdas/atoms/ml.py | 665 +++++++++++++++++++++++----- 6 files changed, 1548 insertions(+), 181 deletions(-) create mode 100644 tests/test_atoms_ml.py diff --git a/docs/release-notes.md b/docs/release-notes.md index 893fa86e..d51cb3bc 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,6 +12,7 @@ - **`process()` with source and sink auto-dispatch.** `process()` is now a method on every atom and a dispatch boundary: `pipeline.process(da, out="results/")` infers both ends. Sources dispatch on the input value — an in-memory `DataArray` runs eagerly (or chunk by chunk with `chunks=`), a virtual one streams through a loader with storage-aligned `chunks="auto"`, a path/directory/glob opens with `open_mfdataarray`, `"tcp://..."` subscribes over ZeroMQ, and any iterable of chunks (a generator, a loader) is consumed as is. Sinks dispatch on the out spec crossed with the first output chunk, so writer creation is deferred to what the pipeline actually emits: a directory stores `DataArray` chunks joined along the chunked dimension (or an SDS archive for `Stream` chunks), `*.csv` appends DataFrames, `"tcp://..."` publishes, `out=None` accumulates and returns the joined result, and a configured writer instance passes through. A chunked source with discontinuities announces them upfront — one warning with the count, read off the source coordinate before any data. The historical `process(atom, loader, writer)` form keeps working unchanged (@atrabattoni). - **`xdas.watch` and unbounded sources.** Realtime is now *named*: `pipeline.process(xd.watch("/incoming", engine=...), out=...)` watches a directory forever, and a bare directory path always means "process what is there". Unbounded sources (watch, ZMQ subscriptions) get streaming semantics — throughput-style progress, a clean `KeyboardInterrupt` that flushes the pipeline and returns the writer result, `until=` to stop at a coordinate value (inclusive, truncating the last chunk), and a warning at each seam as it arrives, since a realtime source cannot be inspected upfront (@atrabattoni). - **Memory guards.** The new `"memory_limit"` configuration entry (default 8 GiB) makes footguns loud: an eager call on a huge virtual array and an `out=None` accumulation that outgrows the limit both raise with the estimated size and a pointer to `.process(out=...)` (@atrabattoni). +- **`Annotate`.** The SeisBench wrapper is rebuilt around what a *weight set* declares rather than what the architecture suggests: the window overlap, the stacking rule (`"avg"` or `"max"`, reproducing SeisBench's `nanmean`/`nanmax` over covering windows exactly), the blinding and the preprocessing arguments are all read off the model instance, and any annotate argument can be overridden at the call (`Annotate(model, scale=2.0)`). The component dimension is found by its labels — each ending with a distinct letter of the model's `component_order`, with SeisBench's flexible horizontal matching — never by its name, and `component_strategy` covers SeisBench's whole range (`"auto"`, `"clone"`, `"pad"`, a named slot, `"strict"`). The output is laid out sample-last, `(..., "phase", dim)`, so the characteristic function of one phase of one channel is contiguous; the end-aligned final window SeisBench appends is emitted at `flush()`, so the output spans the input; and a model whose `annotate_batch_post` breaks the `(batch, samples, classes)` stacking contract is named instead of surfacing as a bare broadcast error. Chunked along its own dimension the sliding window carries across chunks exactly; chunked along another dimension each chunk is a whole record settled on the spot. `MLPicker` and `xdas.mlpicker` remain as deprecated aliases until 0.4 (@atrabattoni). - **`STFT`.** The spectral vocabulary joins the task-atom route: `STFT` streams complex frames with window length and hop in physical units — both are snapped, the window to the next fast FFT size of the target and the hop to a whole sample count — with an expert `nfft` to zero-pad and a `scaling=` of `"spectrum"` or `"psd"`, so `np.abs(stft)**2` composes to an exact spectrogram. Only fully computable frames are ever emitted: the unconsumed tail is buffered across chunks and dropped at gaps, so chunked processing emits exactly the eager frames and no frame ever spans a discontinuity. Built on `scipy.signal.ShortTimeFFT` internally, with the `xdas.stft` function form at the top level (@atrabattoni). - The `xdas.fft` functions (`fft`, `rfft`, `ifft`, `irfft`) now declare whole-record semantics: used as atoms in a chunked pipeline they raise along the transformed dimension instead of silently computing one transform per chunk. Transforming along another dimension than the chunked one keeps working (@atrabattoni). - **`xdas.testing.assert_chunk_invariant`.** The chunk-safety story in one call: run a pipeline eagerly and streamed and assert the two agree — values, coordinates and all. The invariant is quantified over *cuts* (the same stream re-chunked at derived non-divisor sizes, so boundaries land elsewhere) and over *gaps* (`xdas.testing.inject_gaps` places real discontinuities in the input first, so seam resets are exercised at boundaries that do not line up with them). It is both the CI harness for every stateful atom xdas ships and the tool to run on your own pipelines before trusting them chunked (@atrabattoni). diff --git a/tests/test_atoms.py b/tests/test_atoms.py index 13b504bd..1b3ca4cf 100644 --- a/tests/test_atoms.py +++ b/tests/test_atoms.py @@ -12,7 +12,6 @@ FIRFilter, IIRFilter, LFilter, - MLPicker, Partial, Polyphase, ResamplePoly, @@ -22,7 +21,6 @@ ) from xdas.atoms.core import _whole_record from xdas.signal import lfilter -from xdas.synthetics import randn_wavefronts class TestAbstractAtom: @@ -357,66 +355,6 @@ def test_upsample_single_sample(self): assert result.sizes["time"] == 3 -class TestMLPicker: - @pytest.mark.slow - def test_picker(self): - from seisbench.models import PhaseNet - - 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) - - @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 - np.testing.assert_allclose( - _result.values, _expected.values, rtol=1e-5, atol=1e-7 - ) - np.testing.assert_array_max_ulp(_result.values, _expected.values, maxulp=300) - - class TestAtomCoreMissingBranches: def test_repr_with_nested_atoms(self): @@ -610,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 @@ -618,13 +556,6 @@ 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): diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py new file mode 100644 index 00000000..691d5060 --- /dev/null +++ b/tests/test_atoms_ml.py @@ -0,0 +1,984 @@ +""" +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 numpy as np +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, MLPicker + +#: 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_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"))) diff --git a/xdas/__init__.py b/xdas/__init__.py index d3b8ad4e..d53732a2 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -60,6 +60,7 @@ "stack", "trim_overlaps", # task atoms (function forms) + "annotate", "decimate", "detrend", "differentiate", @@ -67,6 +68,7 @@ "hilbert", "integrate", "medfilt", + "mlpicker", "rechunk", "resample", "sliding_mean_removal", @@ -90,6 +92,7 @@ virtual, ) from .atoms.kernel import rechunk +from .atoms.ml import annotate, mlpicker from .atoms.tasks import ( decimate, detrend, diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index c4fb1a4f..aab7a407 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -14,11 +14,12 @@ ...), 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:`MLPicker`. +:class:`Annotate`. """ __all__ = [ "STFT", + "Annotate", "Atom", "Decimate", "Differentiate", @@ -47,6 +48,6 @@ from ..trigger import Trigger from .core import Atom, Partial, Sequential, State, as_function, atomized, compose from .kernel import DownSample, LFilter, Polyphase, Rechunk, SOSFilter, UpSample -from .ml import MLPicker +from .ml import Annotate, MLPicker from .signal import FIRFilter, IIRFilter, ResamplePoly from .tasks import STFT, Decimate, Differentiate, Filter, Integrate, Resample diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index 8d892b1f..25de00de 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -1,15 +1,17 @@ """ -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 typing import ClassVar import numpy as np from ..core import DataArray, concat -from .core import Atom, State +from .core import Atom, State, atomized class LazyModule: @@ -33,48 +35,185 @@ def __getattr__(self, 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): + +class Annotate(Atom): """ - Wraps a SeisBench phase-picking model as a streaming :class:`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. - 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. + 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 - Dimension name along which the model slides (usually ``"time"``). - device : str or torch.device, optional - Torch device. Defaults to CUDA if available, else CPU. + 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. 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. + **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 """ - def __init__(self, model, dim, device=None, component_strategy="clone"): + #: Fallbacks for the annotate arguments this atom reads. 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: ClassVar[dict] = { + "overlap": 0, + "stacking": "avg", + "flexible_horizontal_components": True, + } + + def __init__( + self, + model, + dim="time", + components=None, + component_strategy="auto", + device=None, + **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.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(...) + def _annotate_arg(self, key): + """ + Read one annotate argument, SeisBench-style but crash-free. + + 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 :attr:`_annotate_defaults` instead. The two SeisBench + internals this atom relies on are isolated here and nowhere else. + """ + if key in self.argdict: + return self.argdict[key] + entry = getattr(self.model, "_annotate_args", {}).get(key) + if entry is None: + return self._annotate_defaults[key] + return entry[1] + @property def nperseg(self): """Number of samples per segment (= model input length).""" @@ -82,14 +221,20 @@ 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.""" @@ -98,123 +243,425 @@ def phases(self): @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 + return getattr(self.model, "classes", 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]) - ) - self.circular_input = State( - torch.zeros( - self.batch_size, self.nperseg, dtype=torch.float32, device=self.device - ) - ) - self.model_input = State( - torch.zeros( - self.batch_size, - self.in_channels, - self.nperseg, - dtype=torch.float32, - device=self.device, - ) + """Resolve the dimensions and allocate the sliding-window buffers.""" + dim = self._resolve_sample_dim(da) + component_dim, slots = self._resolve_component_dim(da, dim) + 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) + + 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. + + 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)" + ) 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) + return None + return self._process(da) + + 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" + ) + chunk = self._process(da) + chunks = ([] if chunk is None else [chunk]) + 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. + """ + if self.started is not True: + return [] + dim = self.sample_dim + buffer = self.buffer + remainder = buffer.sizes[dim] - self.nperseg + if remainder > 0: + self._advance(buffer.isel({dim: slice(-remainder, None)}), remainder) + chunk = 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 [chunk] def _process(self, da): - self._initialize(da) + """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 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] + 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 + chunks.append(self._emit(da, idx - step, 0, step)) + last = idx + if last is None: + self.buffer = State(da) + return None + self.started = State(True) + self.buffer = State(da.isel({dim: slice(last, None)})) + return concat(chunks, dim) if chunks else 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 numpy array.""" + values = self.circular_output.narrow(-1, start, length) + counts = self.circular_counts.narrow(-1, start, length) + if self.stacking == "max": + data = torch.where(counts > 0, values, float("nan")) + else: + data = values / counts # samples no window covered come out NaN return data.cpu().numpy() - def _attach_metadata(self, data, da, idx): + 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): + """Build the output chunk of *length* samples found at *start* in the stack.""" + dim = self.sample_dim + data = self._pull(start, length) coords = da.coords.copy() - coords[self.dim] = coords[self.dim][idx : idx + self.step] + if self.component_dim is not None: + coords = coords.drop_dims(self.component_dim) + coords[dim] = coords[dim][offset : offset + length] 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) + return DataArray( + data.reshape(*shape, self.classes, length), + coords, + self.out_dims, + da.name, + da.attrs, + ) + + def _resolve_sample_dim(self, da): + """Resolve *dim* against the input, accepting the ``first``/``last`` aliases.""" + if self.dim == "first": + return da.dims[0] + if self.dim == "last": + return da.dims[-1] + if self.dim not in da.dims: + raise ValueError( + f"{self.dim!r} is not a dimension of the input (got {da.dims})" + ) + return self.dim + + def _labels(self, da, dim): + """Return the string labels of *dim*, or ``None`` if it carries none.""" + 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(self, labels): + """Map *labels* onto model input slots, or ``None`` if they are not components.""" + order = list(self.model.component_order) + flexible = self._annotate_arg("flexible_horizontal_components") + 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(self, da, dim): + """Resolve the labels of *dim* into distinct model slots, or ``None``.""" + labels = self._labels(da, dim) + if labels is None: + return None + slots = self._match_components(labels) + 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 _resolve_component_dim(self, da, sample_dim): + """Find the component dimension of *da* and the slots its labels name.""" + if self.components is False: + return None, None + if self.components is not None: + if self.components not in da.dims: + raise ValueError( + f"{self.components!r} is not a dimension of the input " + f"(got {da.dims})" + ) + slots = self._component_slots(da, self.components) + if slots is None: + raise ValueError( + f"the {self.components!r} dimension is not labelled by " + f"components: expected labels ending with distinct letters " + f"of {self.model.component_order!r}" + ) + return self.components, slots + found = {} + for dim in da.dims: + if dim == sample_dim: + continue + slots = self._component_slots(da, dim) + 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() + + 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 MLPicker(Annotate): + """ + Deprecated alias of :class:`Annotate`, removed in 0.4. + + Beyond the name, the output of this atom is now laid out sample-last, + ``(..., "phase", dim)``, rather than leading with the sample dimension. + """ + + def __init__(self, *args, **kwargs): + warnings.warn( + "MLPicker is deprecated and will be removed in 0.4, use Annotate instead", + DeprecationWarning, + stacklevel=2, ) - return DataArray(data, coords, dims, da.name, da.attrs) + super().__init__(*args, **kwargs) + + +annotate = atomized(Annotate) +mlpicker = atomized(MLPicker) From 8e4cb0998f83dc5761013626750de0ca8cc81351 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 14:43:20 +0200 Subject: [PATCH 18/48] apply the preprocessing filter a weight set ships MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SeisBench's annotate_stream_pre filters the waveforms before resampling them with whatever the weight set declares in filter_args/filter_kwargs; skipping it feeds the network something it was not trained on. The _model_filter builder understands both declared forms — flat, applied to everything, and per channel, one glob pattern each, which SeisBench's own DAS wrapper refuses — and translates them exactly: obspy filters with corners=4, zerophase=False and a Butterworth in second-order sections, which is what Filter defaults to. A zero-phase declaration warns and doubles the order (exact zero-phase IIR has no causal streaming form) and no corner may sit above half the Nyquist of the model's own rate, both as SeisBench concedes. The per-channel form is _ChannelFilter, a Filter subclass that filters everything and restores the channels its pattern does not match: the channel dimension is found by its labels exactly as Annotate finds it, a pattern matching nothing is a silent no-op (stream.select semantics), and ftype='fir' is refused since compensating the group delay on the coordinate would leave the untouched channels on the wrong samples. Both names are private: the stage exists for the picker assembly to come, not as a user-facing bandpass — compose Filter yourself for that. The agreement with obspy is pinned bitwise across every band that has a Filter equivalent. --- tests/test_atoms_ml.py | 306 +++++++++++++++++++- xdas/atoms/ml.py | 619 +++++++++++++++++++++++++++++++++-------- 2 files changed, 806 insertions(+), 119 deletions(-) diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py index 691d5060..61cb9503 100644 --- a/tests/test_atoms_ml.py +++ b/tests/test_atoms_ml.py @@ -36,13 +36,15 @@ """ import numpy as np +import obspy 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, MLPicker +from xdas.atoms import Annotate, Filter, MLPicker, Resample, Sequential +from xdas.atoms.ml import _ChannelFilter, _model_filter #: A short, exactly representable signal: two lanes that are neither equal nor #: proportional, so a bug that mixes lanes shows up in the golden. @@ -982,3 +984,305 @@ 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() diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index 25de00de..9d7a2f75 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -6,12 +6,13 @@ import importlib import warnings -from typing import ClassVar +from fnmatch import fnmatch import numpy as np from ..core import DataArray, concat -from .core import Atom, State, atomized +from .core import Atom, Sequential, State, atomized +from .tasks import Filter class LazyModule: @@ -40,6 +41,236 @@ def __getattr__(self, name): #: so that ``Z12H`` weights accept ``ZNE`` data and vice versa. FLEXIBLE_HORIZONTALS = {"1": "N", "N": "1", "2": "E", "E": "2"} +#: 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 + + +def annotate_arg(model, argdict, key): + """ + Read one annotate argument, SeisBench-style but crash-free. + + 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 + 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 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 + 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 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): """ @@ -140,17 +371,6 @@ class Annotate(Atom): True """ - #: Fallbacks for the annotate arguments this atom reads. 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: ClassVar[dict] = { - "overlap": 0, - "stacking": "avg", - "flexible_horizontal_components": True, - } - def __init__( self, model, @@ -197,22 +417,8 @@ def __init__( self.circular_counts = State(...) def _annotate_arg(self, key): - """ - Read one annotate argument, SeisBench-style but crash-free. - - 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 :attr:`_annotate_defaults` instead. The two SeisBench - internals this atom relies on are isolated here and nowhere else. - """ - if key in self.argdict: - return self.argdict[key] - entry = getattr(self.model, "_annotate_args", {}).get(key) - if entry is None: - return self._annotate_defaults[key] - return entry[1] + """Read one annotate argument through :func:`annotate_arg`.""" + return annotate_arg(self.model, self.argdict, key) @property def nperseg(self): @@ -257,8 +463,14 @@ def fill(self): def initialize(self, da, chunk_dim=None, **flags): """Resolve the dimensions and allocate the sliding-window buffers.""" - dim = self._resolve_sample_dim(da) - component_dim, slots = self._resolve_component_dim(da, dim) + 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.sample_dim = State(dim) self.component_dim = State(component_dim) self.slots = State(self._resolve_slots(slots)) @@ -531,93 +743,6 @@ def _emit(self, da, offset, start, length): da.attrs, ) - def _resolve_sample_dim(self, da): - """Resolve *dim* against the input, accepting the ``first``/``last`` aliases.""" - if self.dim == "first": - return da.dims[0] - if self.dim == "last": - return da.dims[-1] - if self.dim not in da.dims: - raise ValueError( - f"{self.dim!r} is not a dimension of the input (got {da.dims})" - ) - return self.dim - - def _labels(self, da, dim): - """Return the string labels of *dim*, or ``None`` if it carries none.""" - 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(self, labels): - """Map *labels* onto model input slots, or ``None`` if they are not components.""" - order = list(self.model.component_order) - flexible = self._annotate_arg("flexible_horizontal_components") - 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(self, da, dim): - """Resolve the labels of *dim* into distinct model slots, or ``None``.""" - labels = self._labels(da, dim) - if labels is None: - return None - slots = self._match_components(labels) - 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 _resolve_component_dim(self, da, sample_dim): - """Find the component dimension of *da* and the slots its labels name.""" - if self.components is False: - return None, None - if self.components is not None: - if self.components not in da.dims: - raise ValueError( - f"{self.components!r} is not a dimension of the input " - f"(got {da.dims})" - ) - slots = self._component_slots(da, self.components) - if slots is None: - raise ValueError( - f"the {self.components!r} dimension is not labelled by " - f"components: expected labels ending with distinct letters " - f"of {self.model.component_order!r}" - ) - return self.components, slots - found = {} - for dim in da.dims: - if dim == sample_dim: - continue - slots = self._component_slots(da, dim) - 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() - def _resolve_slots(self, slots): """Turn ``component_strategy`` into the input slots the data fills.""" order = list(self.model.component_order) @@ -646,6 +771,264 @@ def _resolve_slots(self, slots): 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 MLPicker(Annotate): """ Deprecated alias of :class:`Annotate`, removed in 0.4. From ef316ccfc9beef6ad5981b0465dabcc989c42a03 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 14:48:51 +0200 Subject: [PATCH 19/48] Trigger moves to atoms/detect: per-phase thresholds, scalar annotations, flush MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Trigger gains the three pieces the picker needs from it, in its new home xdas/atoms/detect.py alongside the rest of the detection vocabulary. thresh accepts a mapping keyed on the phase coordinate Annotate produces, as well as the scalar it took before. A label the mapping does not name gets an infinite threshold and therefore never triggers, which is what lets Annotate keep emitting the noise class: nothing downstream has to slice it out of the characteristic function. Keying on the label rather than on its position is a correctness requirement — the label order belongs to the weight set and flips between them — so the numba kernel now takes one threshold per lane instead of two scalars. coords accepts scalar (0-d) coordinates and emits them as constant columns, so a pick table carries its network/station/location identity. Its default becomes 'auto' — identity first, measurement last: scalars, then the other dimension coordinates, then the picked dimension — so the columns do not depend on the input's dimension order. The columns are resolved once at initialize, which is also what lets flush build a table without the chunk it no longer has. flush() closes the triggers still open at the end of a run, as obspy.trigger_onset does at the end of an array; the last pick of a record used to be dropped. Chunked along a dimension other than dim, none of the state carries — such a chunk is a whole record of other lanes, run from a fresh state and closed on the spot. Atom._join learns to concatenate DataFrame chunks, since an eager call can now answer with the picks of its call plus those of its flush and must still return one table. xdas/trigger.py becomes a compatibility module: find_picks stays verbatim, Trigger is re-exported (its dim default is now 'time', not 'last'), and the module is imported eagerly so that the new lowercase twin xd.trigger — which joins the top level with the other function forms — is never shadowed by a later import of the module. --- docs/release-notes.md | 1 + tests/test_atoms_detect.py | 439 ++++++++++++++++++++++++++++++++ tests/test_trigger.py | 51 +--- xdas/__init__.py | 7 + xdas/atoms/__init__.py | 3 +- xdas/atoms/core.py | 4 + xdas/atoms/detect.py | 507 +++++++++++++++++++++++++++++++++++++ xdas/trigger.py | 206 +-------------- 8 files changed, 968 insertions(+), 250 deletions(-) create mode 100644 tests/test_atoms_detect.py create mode 100644 xdas/atoms/detect.py diff --git a/docs/release-notes.md b/docs/release-notes.md index d51cb3bc..1ed5614f 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,6 +12,7 @@ - **`process()` with source and sink auto-dispatch.** `process()` is now a method on every atom and a dispatch boundary: `pipeline.process(da, out="results/")` infers both ends. Sources dispatch on the input value — an in-memory `DataArray` runs eagerly (or chunk by chunk with `chunks=`), a virtual one streams through a loader with storage-aligned `chunks="auto"`, a path/directory/glob opens with `open_mfdataarray`, `"tcp://..."` subscribes over ZeroMQ, and any iterable of chunks (a generator, a loader) is consumed as is. Sinks dispatch on the out spec crossed with the first output chunk, so writer creation is deferred to what the pipeline actually emits: a directory stores `DataArray` chunks joined along the chunked dimension (or an SDS archive for `Stream` chunks), `*.csv` appends DataFrames, `"tcp://..."` publishes, `out=None` accumulates and returns the joined result, and a configured writer instance passes through. A chunked source with discontinuities announces them upfront — one warning with the count, read off the source coordinate before any data. The historical `process(atom, loader, writer)` form keeps working unchanged (@atrabattoni). - **`xdas.watch` and unbounded sources.** Realtime is now *named*: `pipeline.process(xd.watch("/incoming", engine=...), out=...)` watches a directory forever, and a bare directory path always means "process what is there". Unbounded sources (watch, ZMQ subscriptions) get streaming semantics — throughput-style progress, a clean `KeyboardInterrupt` that flushes the pipeline and returns the writer result, `until=` to stop at a coordinate value (inclusive, truncating the last chunk), and a warning at each seam as it arrives, since a realtime source cannot be inspected upfront (@atrabattoni). - **Memory guards.** The new `"memory_limit"` configuration entry (default 8 GiB) makes footguns loud: an eager call on a huge virtual array and an `out=None` accumulation that outgrows the limit both raise with the estimated size and a pointer to `.process(out=...)` (@atrabattoni). +- **`Trigger` joins the task vocabulary, in `xdas.atoms.detect`.** `thresh` now also takes a mapping keyed on the `phase` coordinate — one threshold per label, labels the mapping does not list never trigger, which is how a characteristic function keeps carrying its noise class without that class ever producing a pick (keying on the label rather than its position matters: the label order of a model belongs to its weight set and flips between them). `coords` gains `"auto"` (the default): scalar coordinates lead as constant columns, then the other dimension coordinates, then the picked dimension — identity first, measurement last, whatever the input's dimension order — and non-dimensional coordinates can be named too. `flush()` closes the triggers still open at the end of a run, as `obspy.trigger_onset` does, so the last pick of a record is no longer lost, and chunking along another dimension than the picked one now answers exactly (each such chunk is a whole record of other lanes, run from a fresh state). The lowercase twin `xdas.trigger` joins the top level; `xdas.trigger` the module remains importable as a compatibility home re-exporting `Trigger` and keeping `find_picks` unchanged. Note: the re-exported `Trigger`'s `dim` default is now `"time"`, not `"last"` (@atrabattoni). - **`Annotate`.** The SeisBench wrapper is rebuilt around what a *weight set* declares rather than what the architecture suggests: the window overlap, the stacking rule (`"avg"` or `"max"`, reproducing SeisBench's `nanmean`/`nanmax` over covering windows exactly), the blinding and the preprocessing arguments are all read off the model instance, and any annotate argument can be overridden at the call (`Annotate(model, scale=2.0)`). The component dimension is found by its labels — each ending with a distinct letter of the model's `component_order`, with SeisBench's flexible horizontal matching — never by its name, and `component_strategy` covers SeisBench's whole range (`"auto"`, `"clone"`, `"pad"`, a named slot, `"strict"`). The output is laid out sample-last, `(..., "phase", dim)`, so the characteristic function of one phase of one channel is contiguous; the end-aligned final window SeisBench appends is emitted at `flush()`, so the output spans the input; and a model whose `annotate_batch_post` breaks the `(batch, samples, classes)` stacking contract is named instead of surfacing as a bare broadcast error. Chunked along its own dimension the sliding window carries across chunks exactly; chunked along another dimension each chunk is a whole record settled on the spot. `MLPicker` and `xdas.mlpicker` remain as deprecated aliases until 0.4 (@atrabattoni). - **`STFT`.** The spectral vocabulary joins the task-atom route: `STFT` streams complex frames with window length and hop in physical units — both are snapped, the window to the next fast FFT size of the target and the hop to a whole sample count — with an expert `nfft` to zero-pad and a `scaling=` of `"spectrum"` or `"psd"`, so `np.abs(stft)**2` composes to an exact spectrogram. Only fully computable frames are ever emitted: the unconsumed tail is buffered across chunks and dropped at gaps, so chunked processing emits exactly the eager frames and no frame ever spans a discontinuity. Built on `scipy.signal.ShortTimeFFT` internally, with the `xdas.stft` function form at the top level (@atrabattoni). - The `xdas.fft` functions (`fft`, `rfft`, `ifft`, `irfft`) now declare whole-record semantics: used as atoms in a chunked pipeline they raise along the transformed dimension instead of silently computing one transform per chunk. Transforming along another dimension than the chunked one keeps working (@atrabattoni). diff --git a/tests/test_atoms_detect.py b/tests/test_atoms_detect.py new file mode 100644 index 00000000..ccc30757 --- /dev/null +++ b/tests/test_atoms_detect.py @@ -0,0 +1,439 @@ +""" +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_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_trigger.py b/tests/test_trigger.py index 5af59e6e..468da68a 100644 --- a/tests/test_trigger.py +++ b/tests/test_trigger.py @@ -2,39 +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], - "sampling_interval": 1.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(): @@ -172,20 +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], - "sampling_interval": 1.0, - }, - }, - ) - picks = Trigger(thresh=0.5, dim="time")(cft) - assert len(picks) == 2 - assert list(picks["time"]) == [2.0, 7.0] diff --git a/xdas/__init__.py b/xdas/__init__.py index d53732a2..7d2c3824 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -74,6 +74,7 @@ "sliding_mean_removal", "stft", "taper", + "trigger", # streaming "watch", ] @@ -91,6 +92,12 @@ 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 from .atoms.tasks import ( diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index aab7a407..969b3eed 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -43,10 +43,11 @@ "as_function", "atomized", "compose", + "trigger", ] -from ..trigger import Trigger 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 from .signal import FIRFilter, IIRFilter, ResamplePoly diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index 9da436a3..5bf06beb 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -521,6 +521,10 @@ def _join(self, chunks, dim): return concat(chunks, dim) except (TypeError, ValueError): return DataCollection(chunks) + if chunks and all(isinstance(c, pd.DataFrame) for c in chunks): + # A pick table is a chunk type of its own: an atom emitting one per + # run, plus one at flush, still answers with a single table. + return pd.concat(chunks, ignore_index=True) return DataCollection(chunks) def flush(self): diff --git a/xdas/atoms/detect.py b/xdas/atoms/detect.py new file mode 100644 index 00000000..7221f282 --- /dev/null +++ b/xdas/atoms/detect.py @@ -0,0 +1,507 @@ +""" +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 + + """ + + 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 wigh the + 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)) + + 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) + picks = self._picks(indices, values) + if independent: + return [picks] + self.flush() + return picks + + 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.""" + annotations = [] + for name in self._names(cft): + if name == self.dim: + # The picked dimension is the chunked one: its indices are + # absolute, so they index the coordinate accumulated so far. + annotations.append((name, self.axis, None)) + continue + coord = self._annotation(cft, name) + if coord.dim is None: + annotations.append((name, None, coord)) + 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: + return cft.dims + 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 _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.coord[indices[axis]].values + 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 : 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 + ----- + 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/trigger.py b/xdas/trigger.py index d736d27f..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], "sampling_interval": 1.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 From 05c42e58ee36e5a6c24c8362e9939f29b9a24385 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 14:57:57 +0200 Subject: [PATCH 20/48] label collection results with their tree path, merge and gather them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A collection walk now carries the tree path down to the leaves instead of rebuilding provenance afterwards. Each leaf's result is labelled with the path it was reached by as it is produced — one column per named level, leading the table, filled with that level's key — so a pick found under IA / DBNFM / -- comes out with its network, station and location before its time and its value. Producing the labels at the leaf rather than folding them in on the way back up is what will let a streaming walk hand a leaf straight to a sink and still know whose it was. Atoms may then declare a merge(results) hook folding those labelled results. It is undefined on Atom, so an atom returning arrays sees nothing change and its tree is rebuilt as before; Trigger.merge is a plain concat, which is all it takes once the columns are there, so xd.trigger(dc, ...) answers a whole network with one flat table; and Sequential.merge delegates to the last stage declaring one. merge=False opts out and returns the labelled tree. A column with two sources — the tree key and a scalar coordinate of the same name — stays a single column the tree path fills, warning on a genuine disagreement. Positional levels contribute their index, and the flushed tail of a folded sequence is attributed to the last element it came out of. Annotating a collection no longer needs a prior xd.stack: Atom.gather is a hook the walk consults on every mapping level before descending — return the level collapsed to an array, or None (the default) to map over its leaves. Annotate implements it, because what counts as a component is a property of the model: Z12H and ENZ disagree about which channels group together. It collapses through xd.stack, so the structural checks, the grid snapping and the error messages are shared, and tolerance= reaches them. Sequential.gather delegates to the first stage that claims the level, so the gather happens once and before the first stage runs — the per-channel filter a weight set ships selects ??H out of Z12H and can only do that once the channels are a dimension. Recognition is deliberately conservative, since folding a station level into a component axis would silently destroy the distinction between stations: both the level's name (COMPONENT_LEVELS, or components=) and its keys must resolve, a key must look like a whole channel code, and the keys must agree on their length and band code. Keys that resolve to nothing leave the level alone, so a DAS collection whose spatial axis is called channel walks straight past; keys that resolve only in part raise, naming the conflict. Pre-stacking keeps working: annotate(dc) equals annotate(xd.stack(dc, 'channel')). --- tests/test_atoms_ml.py | 284 +++++++++++++++++++++++++++++++++++ tests/test_datacollection.py | 184 +++++++++++++++++++++++ xdas/atoms/core.py | 256 ++++++++++++++++++++++++++++--- xdas/atoms/detect.py | 55 +++++++ xdas/atoms/ml.py | 233 +++++++++++++++++++++++++++- 5 files changed, 992 insertions(+), 20 deletions(-) diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py index 61cb9503..c3f02771 100644 --- a/tests/test_atoms_ml.py +++ b/tests/test_atoms_ml.py @@ -36,6 +36,7 @@ """ import numpy as np +import numpy.testing as npt import obspy import pytest import torch @@ -1286,3 +1287,286 @@ def test_the_stage_composes_ahead_of_resample_and_annotate(self): 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") diff --git a/tests/test_datacollection.py b/tests/test_datacollection.py index 7c437612..a9c47ebe 100644 --- a/tests/test_datacollection.py +++ b/tests/test_datacollection.py @@ -1,4 +1,7 @@ +import warnings + import h5py +import numpy as np import pandas as pd import pytest @@ -531,3 +534,184 @@ def test_repr_shows_the_table(self): 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 + + 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", + ) + with warnings.catch_warnings(): + warnings.simplefilter("error") + 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/xdas/atoms/core.py b/xdas/atoms/core.py index 5bf06beb..a419cf35 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -97,6 +97,76 @@ def _join_chunks(chunks, dim=None): 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 + + def _flush_through(atoms, **flags): """ Codec-drain a linear chain of atoms. @@ -205,9 +275,17 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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. + merge: callable or None + The optional hook folding the per-leaf results of a collection + walk into one object. ``None`` (the default) means the atom has + none and the tree is rebuilt as it always was. See + :meth:`~xdas.atoms.Trigger.merge` for an implementation. 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() @@ -219,6 +297,9 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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; undefined by default. fresh() Returns a stateless clone sharing the configuration. iter_chunks(source) @@ -227,6 +308,9 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): """ 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", {}) @@ -290,6 +374,34 @@ 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, returning zero or more output chunks. @@ -305,9 +417,20 @@ def __call__(self, x, **flags): 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; otherwise a :class:`DataSequence` of chunks is returned. """ + merge = flags.pop("merge", True) chunk_dim = flags.get("chunk_dim", None) self._check_chunk_dim(x, chunk_dim) if ( @@ -323,18 +446,14 @@ def __call__(self, x, **flags): "with `.process(da, out=...)` instead, or raise the limit " "with `xdas.config.set('memory_limit', ...)`" ) - if isinstance(x, DataMapping): - if chunk_dim is not None: - raise NotImplementedError( - "chunked processing of mapping collections is not supported: " - "process each leaf with its own atom instance" - ) - return DataCollection( - {key: self(value, **flags) for key, value in x.items()}, - getattr(x, "name", None), - ) - if isinstance(x, DataSequence): - return self._fold(x, flags) + 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) @@ -485,31 +604,89 @@ def _judge_seam(self, info): return "continuous" return "gap" if jump > 0 else "overlap" - def _fold(self, x, flags): + 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(el, **flags) for el in x], name) + 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 el in x: - chunks += _aschunks(self(el, **flags)) - chunks += self.flush() + 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 DataCollection(chunks, name) chunks = [] - for el in x: - chunks += _aschunks(self(el, **flags)) + 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): @@ -884,6 +1061,47 @@ def _resolve_dim(self, x): 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.""" return type(self)([atom.fresh() for atom in self], name=self.name) diff --git a/xdas/atoms/detect.py b/xdas/atoms/detect.py index 7221f282..9f5d859d 100644 --- a/xdas/atoms/detect.py +++ b/xdas/atoms/detect.py @@ -143,6 +143,34 @@ class Trigger(Atom): 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"): @@ -261,6 +289,33 @@ def call(self, cft, **flags): 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. Leaves that produced no + pick at all contribute nothing rather than an empty table. + + 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. diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index 9d7a2f75..a87186fe 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -10,7 +10,7 @@ import numpy as np -from ..core import DataArray, concat +from ..core import DataArray, concat, stack from .core import Atom, Sequential, State, atomized from .tasks import Filter @@ -57,6 +57,20 @@ def __getattr__(self, name): #: 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) + def annotate_arg(model, argdict, key): """ @@ -213,6 +227,137 @@ def component_slots(da, dim, order, flexible=True): 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. @@ -305,6 +450,10 @@ class Annotate(Atom): 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 the model's input slots are filled from the data: @@ -326,6 +475,12 @@ class Annotate(Atom): 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. **annotate_kwargs SeisBench annotate arguments (``overlap``, ``stacking``, ``blinding``, ...) overriding what the weight set declares in ``default_args``. @@ -369,6 +524,14 @@ class Annotate(Atom): 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__( @@ -378,6 +541,7 @@ def __init__( components=None, component_strategy="auto", device=None, + tolerance=None, **annotate_kwargs, ): super().__init__() @@ -396,6 +560,7 @@ def __init__( self.dim = dim self.components = components self.component_strategy = component_strategy + self.tolerance = tolerance 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}") @@ -461,6 +626,72 @@ 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 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, + ) + 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) From c54cb10e99e1c21eb0b6bfe22a8415a2f192eb78 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 15:02:28 +0200 Subject: [PATCH 21/48] Picker, the pipeline a weight set describes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Picker(model) assembles the whole pipeline SeisBench's model.classify runs — the weight set's own preprocessing filter when it ships one, Resample to the weight set's own rate (not always 100 Hz: diting runs at 50), Annotate, and Trigger with the weight set's own per-phase thresholds — from the weight set and nothing else. Two pickers built on one model class can differ in stage count, sampling rate and thresholds, which is the point: everything a SeisBench picker does is a property of the weights. The thresholds come from _model_thresholds: one entry per picked label as model_pick_labels resolves them — a model declaring a phases subset (the EQTransformer family) picks exactly that subset, anything else picks every label but the noise class — each looked up as SeisBench does: the call wins, else the weight set's default_args, else the model's documented default for that key, else the *_threshold catch-all, else 0.3. Values pass through faithfully, including iquique's P_threshold of 1.12, which simply never fires. Annotate reads its labels through model_phases — labels=None falls back to positional labels exactly as WaveformModel._predictions_to_stream does, a callable is refused by name — and its class count now comes from the labels rather than model.classes, which counts only the picking head: EQTransformer sets it to 2 while labelling three outputs, and sizing the buffers on it made the whole family unrunnable. Being a Sequential rather than a factory function, a picker keeps everything a pipeline can do — >> composes it, repr shows its stages, it pickles, picker.process streams it — and inherits Annotate.gather and Trigger.merge, so xd.pick(dc, model) answers a whole network tree with one flat table. The dimension aliases first/last are refused: the picks are annotated with coordinates, which are named. --- docs/release-notes.md | 1 + tests/test_atoms_ml.py | 394 ++++++++++++++++++++++++++++++++++++++++- xdas/__init__.py | 3 +- xdas/atoms/__init__.py | 3 +- xdas/atoms/ml.py | 390 +++++++++++++++++++++++++++++++++++++++- 5 files changed, 782 insertions(+), 9 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index 1ed5614f..99a65416 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,6 +12,7 @@ - **`process()` with source and sink auto-dispatch.** `process()` is now a method on every atom and a dispatch boundary: `pipeline.process(da, out="results/")` infers both ends. Sources dispatch on the input value — an in-memory `DataArray` runs eagerly (or chunk by chunk with `chunks=`), a virtual one streams through a loader with storage-aligned `chunks="auto"`, a path/directory/glob opens with `open_mfdataarray`, `"tcp://..."` subscribes over ZeroMQ, and any iterable of chunks (a generator, a loader) is consumed as is. Sinks dispatch on the out spec crossed with the first output chunk, so writer creation is deferred to what the pipeline actually emits: a directory stores `DataArray` chunks joined along the chunked dimension (or an SDS archive for `Stream` chunks), `*.csv` appends DataFrames, `"tcp://..."` publishes, `out=None` accumulates and returns the joined result, and a configured writer instance passes through. A chunked source with discontinuities announces them upfront — one warning with the count, read off the source coordinate before any data. The historical `process(atom, loader, writer)` form keeps working unchanged (@atrabattoni). - **`xdas.watch` and unbounded sources.** Realtime is now *named*: `pipeline.process(xd.watch("/incoming", engine=...), out=...)` watches a directory forever, and a bare directory path always means "process what is there". Unbounded sources (watch, ZMQ subscriptions) get streaming semantics — throughput-style progress, a clean `KeyboardInterrupt` that flushes the pipeline and returns the writer result, `until=` to stop at a coordinate value (inclusive, truncating the last chunk), and a warning at each seam as it arrives, since a realtime source cannot be inspected upfront (@atrabattoni). - **Memory guards.** The new `"memory_limit"` configuration entry (default 8 GiB) makes footguns loud: an eager call on a huge virtual array and an `out=None` accumulation that outgrows the limit both raise with the estimated size and a pointer to `.process(out=...)` (@atrabattoni). +- **`Picker` / `xd.pick`, the headline: waveforms in, one pick table out.** `Picker(model)` assembles the whole pipeline SeisBench's `model.classify(stream)` runs — the preprocessing filter the weight set ships (if any), `Resample` to the weight set's own rate, `Annotate`, `Trigger` with the weight set's own per-phase thresholds — from the weight set and nothing else, so two pickers built on one model class can differ in stage count, sampling rate and thresholds. Being a `Sequential`, a picker composes with `>>`, pickles, shows its stages in `repr` and streams with `picker.process(source, out=...)`; it inherits `Annotate.gather` and `Trigger.merge`, so `xd.pick(dc, model)` answers a whole ObsPy-style network tree with one flat table, each pick labelled `network station location phase time value`. The one deliberate difference from SeisBench is the resampler: obspy's `Trace.resample` halves the amplitude at half the input Nyquist (its frequency-domain hann window), where the polyphase filter used here is flat; `resample=False` drops the stage (@atrabattoni). - **`Trigger` joins the task vocabulary, in `xdas.atoms.detect`.** `thresh` now also takes a mapping keyed on the `phase` coordinate — one threshold per label, labels the mapping does not list never trigger, which is how a characteristic function keeps carrying its noise class without that class ever producing a pick (keying on the label rather than its position matters: the label order of a model belongs to its weight set and flips between them). `coords` gains `"auto"` (the default): scalar coordinates lead as constant columns, then the other dimension coordinates, then the picked dimension — identity first, measurement last, whatever the input's dimension order — and non-dimensional coordinates can be named too. `flush()` closes the triggers still open at the end of a run, as `obspy.trigger_onset` does, so the last pick of a record is no longer lost, and chunking along another dimension than the picked one now answers exactly (each such chunk is a whole record of other lanes, run from a fresh state). The lowercase twin `xdas.trigger` joins the top level; `xdas.trigger` the module remains importable as a compatibility home re-exporting `Trigger` and keeping `find_picks` unchanged. Note: the re-exported `Trigger`'s `dim` default is now `"time"`, not `"last"` (@atrabattoni). - **`Annotate`.** The SeisBench wrapper is rebuilt around what a *weight set* declares rather than what the architecture suggests: the window overlap, the stacking rule (`"avg"` or `"max"`, reproducing SeisBench's `nanmean`/`nanmax` over covering windows exactly), the blinding and the preprocessing arguments are all read off the model instance, and any annotate argument can be overridden at the call (`Annotate(model, scale=2.0)`). The component dimension is found by its labels — each ending with a distinct letter of the model's `component_order`, with SeisBench's flexible horizontal matching — never by its name, and `component_strategy` covers SeisBench's whole range (`"auto"`, `"clone"`, `"pad"`, a named slot, `"strict"`). The output is laid out sample-last, `(..., "phase", dim)`, so the characteristic function of one phase of one channel is contiguous; the end-aligned final window SeisBench appends is emitted at `flush()`, so the output spans the input; and a model whose `annotate_batch_post` breaks the `(batch, samples, classes)` stacking contract is named instead of surfacing as a bare broadcast error. Chunked along its own dimension the sliding window carries across chunks exactly; chunked along another dimension each chunk is a whole record settled on the spot. `MLPicker` and `xdas.mlpicker` remain as deprecated aliases until 0.4 (@atrabattoni). - **`STFT`.** The spectral vocabulary joins the task-atom route: `STFT` streams complex frames with window length and hop in physical units — both are snapped, the window to the next fast FFT size of the target and the hop to a whole sample count — with an expert `nfft` to zero-pad and a `scaling=` of `"spectrum"` or `"psd"`, so `np.abs(stft)**2` composes to an exact spectrogram. Only fully computable frames are ever emitted: the unconsumed tail is buffered across chunks and dropped at gaps, so chunked processing emits exactly the eager frames and no frame ever spans a discontinuity. Built on `scipy.signal.ShortTimeFFT` internally, with the `xdas.stft` function form at the top level (@atrabattoni). diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py index c3f02771..06ddf82a 100644 --- a/tests/test_atoms_ml.py +++ b/tests/test_atoms_ml.py @@ -35,17 +35,20 @@ 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, Resample, Sequential -from xdas.atoms.ml import _ChannelFilter, _model_filter +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. @@ -1570,3 +1573,390 @@ def test_the_gather_is_a_no_op_on_a_collection_of_stacked_arrays(self): 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 diff --git a/xdas/__init__.py b/xdas/__init__.py index 7d2c3824..7a324b84 100644 --- a/xdas/__init__.py +++ b/xdas/__init__.py @@ -69,6 +69,7 @@ "integrate", "medfilt", "mlpicker", + "pick", "rechunk", "resample", "sliding_mean_removal", @@ -99,7 +100,7 @@ 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 +from .atoms.ml import annotate, mlpicker, pick from .atoms.tasks import ( decimate, detrend, diff --git a/xdas/atoms/__init__.py b/xdas/atoms/__init__.py index 969b3eed..89be6e6f 100644 --- a/xdas/atoms/__init__.py +++ b/xdas/atoms/__init__.py @@ -31,6 +31,7 @@ "LFilter", "MLPicker", "Partial", + "Picker", "Polyphase", "Rechunk", "Resample", @@ -49,6 +50,6 @@ 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 +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/ml.py b/xdas/atoms/ml.py index a87186fe..3588c89c 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -12,7 +12,8 @@ from ..core import DataArray, concat, stack from .core import Atom, Sequential, State, atomized -from .tasks import Filter +from .detect import Trigger +from .tasks import Filter, Resample class LazyModule: @@ -71,6 +72,18 @@ def __getattr__(self, name): #: ``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): """ @@ -105,6 +118,165 @@ def annotate_arg(model, argdict, 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. @@ -608,8 +780,14 @@ def stacking(self): @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): @@ -618,8 +796,15 @@ def in_channels(self): @property def classes(self): - """Number of output classes (phases) the model produces.""" - return getattr(self.model, "classes", len(self.phases)) + """ + 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 fill(self): @@ -1260,6 +1445,200 @@ def _translate_filter(args, kwargs, sampling_rate=None): ) +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, + ) + ) + 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. @@ -1278,4 +1657,5 @@ def __init__(self, *args, **kwargs): annotate = atomized(Annotate) +pick = atomized(Picker) mlpicker = atomized(MLPicker) From 17beea2f79406bbf466f026004600c40c05736e5 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 15:08:29 +0200 Subject: [PATCH 22/48] process walks a collection the way the eager call does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit atom.process(dc, ...) now walks a DataCollection exactly as atom(dc) does, so atom.process(dc, out=None) == atom(dc). Anything else makes the streaming form second-class precisely where streaming matters most: a leaf too large to call eagerly at all is what the memory guard already refuses, pointing at .process(da, out=...). The walk mirrors Atom._walk step for step. Each mapping level is offered to gather 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 fold element by element, the state carrying across and the atom flushed once at the end, its tail attributed to the last element; and each leaf streams through the existing single-source path, with chunks= and until= applying per leaf. One atom instance takes the leaves one at a time, reset between them, because an atom holding a model either saturates the CPU or holds a lot of device memory. Sinks gain a rule per destination, and the produce-time path columns are what make them work — every output chunk is labelled on its way to the writer, never reconstructed afterwards: - out=None accumulates per leaf and merges, giving the eager result; - a *.csv, a URL or a ready writer instance is shared: every leaf appends to one table, the path columns 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 — a folded sequence is one stream and so writes to one directory. merge= is accepted on process and popped walk-level, as __call__ does. A DataSequence handed over as a collection now folds rather than streaming as one concatenated result; a glob or directory that opens to a sequence is still a single source, and get_source(sequence) asks for the same of a collection in hand. Verified while porting: xd.stack keeps a tile-backed leaf tile-backed, so the gather of an ObsPy-style collection reads nothing. --- tests/test_process.py | 258 ++++++++++++++++++++++++++++++- xdas/atoms/core.py | 24 ++- xdas/processing/core.py | 331 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 600 insertions(+), 13 deletions(-) diff --git a/tests/test_process.py b/tests/test_process.py index 117eb92f..353d2ea3 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -17,10 +17,11 @@ import xdas as xd import xdas.processing as xp -from xdas.atoms import Partial +from xdas.atoms import Partial, State from xdas.atoms.core import _join_chunks from xdas.config import Config from xdas.processing.core import _auto_chunks, _ChainSource, _to_human +from xdas.virtual import TileArray @pytest.fixture @@ -542,3 +543,258 @@ def publish(): 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_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 + monkeypatch.setitem(Config.config, "memory_limit", 1) + 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/xdas/atoms/core.py b/xdas/atoms/core.py index a419cf35..87891638 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -744,7 +744,7 @@ def iter_chunks(self, source, chunk_dim=None): yield from self.flush() self.reset() - def process(self, source, out=None, chunks=None, until=None): + def process(self, source, out=None, chunks=None, until=None, merge=True): """ Process any chunk source through this atom, writing to any sink. @@ -754,18 +754,28 @@ def process(self, source, out=None, chunks=None, until=None): 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, str, Path, iterable or loader + source : DataArray, DataCollection, str, Path, iterable or loader What to process: an in-memory or virtual :class:`DataArray`, a - file path, directory or glob pattern, a ``"tcp://..."`` address, + :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. + 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 @@ -773,6 +783,10 @@ def process(self, source, out=None, chunks=None, until=None): 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 ------- @@ -792,7 +806,7 @@ def process(self, source, out=None, chunks=None, until=None): """ from ..processing.core import process - return process(self, source, out=out, chunks=chunks, until=until) + 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*.""" diff --git a/xdas/processing/core.py b/xdas/processing/core.py index a548fd22..d6a600ca 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -26,10 +26,25 @@ from watchdog.observers import Observer from .. import config -from ..atoms.core import _announce_splits, _aschunks, _join_chunks +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, DataSequence, concat, open_dataarray, open_mfdataarray +from ..core import ( + DataArray, + DataCollection, + DataMapping, + DataSequence, + concat, + open_dataarray, + open_mfdataarray, +) from ..virtual import TileArray, VirtualBackend, VirtualStack from .monitor import Monitor @@ -218,7 +233,7 @@ def _dump(chunk, path, encoding): return open_dataarray(path) -def process(atom, source, out=None, chunks=None, until=None): +def process(atom, source, out=None, chunks=None, until=None, merge=True): """ Execute a processing pipeline over any chunk source, into any sink. @@ -237,10 +252,11 @@ def process(atom, source, out=None, chunks=None, until=None): 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, str, Path, iterable or loader + 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 path, + 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 @@ -253,7 +269,11 @@ def process(atom, source, out=None, chunks=None, until=None): 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. + 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 @@ -262,6 +282,12 @@ def process(atom, source, out=None, chunks=None, until=None): 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 ------- @@ -269,6 +295,9 @@ def process(atom, source, out=None, chunks=None, until=None): 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 ----- @@ -276,12 +305,47 @@ def process(atom, source, out=None, chunks=None, until=None): ``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. - result = atom(source) + result = _annotate_path(atom(source), path) if out is None: return result outputs = _aschunks(result) @@ -307,6 +371,7 @@ def process(atom, source, out=None, chunks=None, until=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) @@ -362,6 +427,258 @@ def write(chunk): return writer.result() if writer is not None else None +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): + """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) + name = getattr(x, "name", None) + return DataCollection( + { + key: self.level(value, _extend_path(path, name, key)) + for key, value in x.items() + }, + name, + ) + if isinstance(x, DataSequence): + return self.fold(x, path) + 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(path), path)) + + def fold(self, x, path): + """ + 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)) + for index, element in enumerate(x) + ], + name, + ) + spec = self.sink.spec(path) + 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.""" + return None if self.writer is None else self.writer.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, path): + """Return the out spec of the leaf reached by *path*.""" + if self.shared is not None: + return self.shared + if self.out is None: + return None + return os.path.join(str(self.out), *(str(key) for key in path.values())) + + 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. From 8098b09df4c9269a0a384cd9ae77487c7e172456 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 15:15:43 +0200 Subject: [PATCH 23/48] the GPU atom is asynchronous behind a bounded output queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The process loop stays serial and CPU atoms stay internally parallel; only Annotate is asynchronous, behind a small bounded queue — the same max_buffers pattern DataArrayLoader and DataArrayWriter already use. Each completed window's reduction now allocates a fresh tensor (it never views the circular stack), so on CUDA its device-to-host transfer can be issued immediately — pinned staging buffer, non-blocking copy, an event marking completion — while the sliding window moves on. call() emits only the outputs whose transfer has completed, in order (the 0..n contract already allows late emission); flush() drains whatever is still in flight, so nothing survives the end of a run and chunk invariance holds unchanged. At most max_buffers transfers are left pending (default 2, 0 restores fully synchronous emission), which is what bounds the staging memory. On the CPU there is nothing to wait for: the queue completes on arrival and behavior is exactly the synchronous reference. Together with the H2D staging path that was already pinned and non-blocking, the CPU keeps preparing and feeding windows while the GPU computes and while results cross back, with no executor machinery: no thread-per-stage, no stage budget, no second level of parallelism to oversubscribe the first. --- tests/test_atoms_ml.py | 63 +++++++++++++++++++++ xdas/atoms/ml.py | 122 ++++++++++++++++++++++++++++++++--------- 2 files changed, 159 insertions(+), 26 deletions(-) diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py index 06ddf82a..3232ba07 100644 --- a/tests/test_atoms_ml.py +++ b/tests/test_atoms_ml.py @@ -1960,3 +1960,66 @@ def test_a_gap_splits_the_record_into_runs(self): ) 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) diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index 3588c89c..8dd06ed8 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -6,6 +6,7 @@ import importlib import warnings +from collections import deque from fnmatch import fnmatch import numpy as np @@ -653,6 +654,17 @@ class Annotate(Atom): 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``. @@ -714,6 +726,7 @@ def __init__( component_strategy="auto", device=None, tolerance=None, + max_buffers=2, **annotate_kwargs, ): super().__init__() @@ -733,6 +746,7 @@ def __init__( 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}") @@ -752,6 +766,7 @@ def __init__( 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`.""" @@ -917,6 +932,7 @@ def initialize(self, da, chunk_dim=None, **flags): ) 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.""" @@ -930,7 +946,11 @@ def call(self, da, **flags): 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. + 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 — @@ -949,12 +969,14 @@ def call(self, da, **flags): f"({da.sizes[dim]} samples) than one model window " f"({self.nperseg} samples)" ) + self._process(da) else: da = concat([self.buffer, da], dim) if da.sizes[dim] < self.nperseg: self.buffer = State(da) - return None - return self._process(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.""" @@ -976,8 +998,8 @@ def _call_independent(self, da, **flags): "model and cannot be chunked: the components of one window " "are read together" ) - chunk = self._process(da) - chunks = ([] if chunk is None else [chunk]) + self.flush() + 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] @@ -988,21 +1010,22 @@ def flush(self): 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. + Firing once per run, this stays chunk-invariant. The in-flight + output queue is drained here, so nothing survives the end of a run. """ if self.started is not True: - return [] + return self._harvest(block=True) dim = self.sample_dim buffer = self.buffer remainder = buffer.sizes[dim] - self.nperseg if remainder > 0: self._advance(buffer.isel({dim: slice(-remainder, None)}), remainder) - chunk = self._emit(buffer, 0, self.step - remainder, self.nperseg + 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 [chunk] + return self._harvest(block=True) def _process(self, da): """Annotate every window of *da* that the buffered state can complete.""" @@ -1013,20 +1036,18 @@ def _process(self, da): else: self._prime(da) first = 0 - chunks = [] 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 - chunks.append(self._emit(da, idx - step, 0, step)) + self._emit(da, idx - step, 0, step) last = idx if last is None: self.buffer = State(da) - return None + return self.started = State(True) self.buffer = State(da.isel({dim: slice(last, None)})) - return concat(chunks, dim) if chunks else None def _prime(self, da): """Stage the head of the first window so that the first slide completes it.""" @@ -1117,14 +1138,18 @@ def _accumulate(self, out): counts += finite.to(counts.dtype) def _pull(self, start, length): - """Reduce *length* stacked samples from *start* into a numpy array.""" + """ + 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": - data = torch.where(counts > 0, values, float("nan")) - else: - data = values / counts # samples no window covered come out NaN - return data.cpu().numpy() + 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.""" @@ -1142,7 +1167,7 @@ def _to_device(self, chunk): return data def _emit(self, da, offset, start, length): - """Build the output chunk of *length* samples found at *start* in the stack.""" + """Queue the output chunk of *length* samples found at *start* in the stack.""" dim = self.sample_dim data = self._pull(start, length) coords = da.coords.copy() @@ -1151,13 +1176,58 @@ def _emit(self, da, offset, start, length): coords[dim] = coords[dim][offset : offset + length] coords["phase"] = self.phases shape = tuple(da.sizes[other] for other in self.batch_dims) - return DataArray( - data.reshape(*shape, self.classes, length), - coords, - self.out_dims, - da.name, - da.attrs, - ) + 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.""" From 7c6de84abbbcc1f1d8fa7222d3e864c0974e5ded Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 15:17:55 +0200 Subject: [PATCH 24/48] docs: a picking walkthrough, and the phase's release notes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A new user-guide page next to the DAS pipeline material: xd.open then xd.pick over a real 8-station network day, what the picker builds from a weight set (three stages for original, four for obs), why thresholds and classes are keyed by label rather than by position, and the same walk on a DAS collection eagerly and through process(..., out=...). Its code is not executed at build time — SeisBench and its weights are not documentation dependencies — but every output shown was produced by running it on feature/atoms, whose picking pipeline this branch reproduces stage for stage. The page states the one real SeisBench deviation as what it is: obspy's Trace.resample defaults to window='hann' applied in the frequency domain, which halves the amplitude at half the input Nyquist, where the polyphase filter is flat. Fed the same resampled data, the two agree pick for pick. Annotate, Trigger and Picker get their own API section rather than sitting under signal processing, the deprecated MLPicker drops out of the listing, and the function forms annotate/pick/trigger join the top-level list. Release notes gain the walk labelling/merge/gather block, the process() collection walk and the async GPU queue. --- docs/api/atoms.md | 20 +- docs/release-notes.md | 3 + docs/user-guide/index.md | 3 +- docs/user-guide/pipeline/index.md | 6 +- docs/user-guide/pipeline/picking.md | 306 ++++++++++++++++++++++++++++ 5 files changed, 334 insertions(+), 4 deletions(-) create mode 100644 docs/user-guide/pipeline/picking.md diff --git a/docs/api/atoms.md b/docs/api/atoms.md index a1deea09..1c79d485 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -115,6 +115,7 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, .. autosummary:: :toctree: ../_autosummary + annotate decimate detrend differentiate @@ -122,11 +123,13 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, hilbert integrate medfilt + pick rechunk resample sliding_mean_removal stft taper + trigger ``` ```{eval-rst} @@ -141,9 +144,24 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, FIRFilter IIRFilter - MLPicker ResamplePoly +``` + +## 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 diff --git a/docs/release-notes.md b/docs/release-notes.md index 99a65416..e2b7609d 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,6 +12,9 @@ - **`process()` with source and sink auto-dispatch.** `process()` is now a method on every atom and a dispatch boundary: `pipeline.process(da, out="results/")` infers both ends. Sources dispatch on the input value — an in-memory `DataArray` runs eagerly (or chunk by chunk with `chunks=`), a virtual one streams through a loader with storage-aligned `chunks="auto"`, a path/directory/glob opens with `open_mfdataarray`, `"tcp://..."` subscribes over ZeroMQ, and any iterable of chunks (a generator, a loader) is consumed as is. Sinks dispatch on the out spec crossed with the first output chunk, so writer creation is deferred to what the pipeline actually emits: a directory stores `DataArray` chunks joined along the chunked dimension (or an SDS archive for `Stream` chunks), `*.csv` appends DataFrames, `"tcp://..."` publishes, `out=None` accumulates and returns the joined result, and a configured writer instance passes through. A chunked source with discontinuities announces them upfront — one warning with the count, read off the source coordinate before any data. The historical `process(atom, loader, writer)` form keeps working unchanged (@atrabattoni). - **`xdas.watch` and unbounded sources.** Realtime is now *named*: `pipeline.process(xd.watch("/incoming", engine=...), out=...)` watches a directory forever, and a bare directory path always means "process what is there". Unbounded sources (watch, ZMQ subscriptions) get streaming semantics — throughput-style progress, a clean `KeyboardInterrupt` that flushes the pipeline and returns the writer result, `until=` to stop at a coordinate value (inclusive, truncating the last chunk), and a warning at each seam as it arrives, since a realtime source cannot be inspected upfront (@atrabattoni). - **Memory guards.** The new `"memory_limit"` configuration entry (default 8 GiB) makes footguns loud: an eager call on a huge virtual array and an `out=None` accumulation that outgrows the limit both raise with the estimated size and a pointer to `.process(out=...)` (@atrabattoni). +- **Collection walks carry their tree path, and atoms can merge or gather them.** Walking a collection, each leaf's result is now labelled with the path it was reached by *as it is produced* — one leading column per named level, filled with that level's key, so a pick found under `IA / DBNFM / --` comes out with its network, station and location before its time and value (positional levels contribute their index; a column that already exists as a scalar coordinate stays one column, the tree path winning with a warning on a genuine disagreement). Atoms may declare a `merge(results)` hook folding those labelled results — `Trigger.merge` is a plain concat, so `xd.trigger(dc, ...)` answers a whole network with one flat table, and `merge=False` opts out — and a `gather(mapping)` hook offered each mapping level before the walk descends: `Annotate` implements it, collapsing a channel level into the component dimension through `xd.stack` (what counts as a component is a property of the *model*, which is why the reader cannot do this and the atom can), with deliberately conservative recognition so a station level is never silently folded. `Sequential` delegates gather to its first claiming stage and merge to its last declaring one (@atrabattoni). +- **`process()` walks collections.** `atom.process(dc, chunks=..., out=...)` walks a `DataCollection` exactly as `atom(dc)` does — gather first, mapping levels recursed, sequence levels folded with state carried across the elements — with each leaf streamed through the single-source path, so `atom.process(dc, out=None) == atom(dc)` and the streaming form is not second-class precisely where it matters most: a leaf too large to call eagerly at all. Sinks gain one rule per destination: `out=None` accumulates per leaf and merges; a `*.csv`, a URL or a ready writer is *shared*, every leaf appending to one table with the tree-path columns keeping the rows apart; a directory *fans out*, one subdirectory per leaf mirroring the tree path (@atrabattoni). +- **The GPU atom is asynchronous behind a bounded output queue.** The process loop stays serial and CPU atoms stay internally parallel; only `Annotate` overlaps its work with the device: each completed window's device-to-host transfer is issued asynchronously (pinned staging, non-blocking, an event marking completion) and `call()` emits only the outputs whose transfer has completed, `flush()` draining the rest — the 0..n contract already allows late emission, so chunk order, seams and results are untouched. `max_buffers` bounds the in-flight transfers (default 2; 0 restores synchronous emission); on the CPU the queue completes on arrival and nothing changes (@atrabattoni). - **`Picker` / `xd.pick`, the headline: waveforms in, one pick table out.** `Picker(model)` assembles the whole pipeline SeisBench's `model.classify(stream)` runs — the preprocessing filter the weight set ships (if any), `Resample` to the weight set's own rate, `Annotate`, `Trigger` with the weight set's own per-phase thresholds — from the weight set and nothing else, so two pickers built on one model class can differ in stage count, sampling rate and thresholds. Being a `Sequential`, a picker composes with `>>`, pickles, shows its stages in `repr` and streams with `picker.process(source, out=...)`; it inherits `Annotate.gather` and `Trigger.merge`, so `xd.pick(dc, model)` answers a whole ObsPy-style network tree with one flat table, each pick labelled `network station location phase time value`. The one deliberate difference from SeisBench is the resampler: obspy's `Trace.resample` halves the amplitude at half the input Nyquist (its frequency-domain hann window), where the polyphase filter used here is flat; `resample=False` drops the stage (@atrabattoni). - **`Trigger` joins the task vocabulary, in `xdas.atoms.detect`.** `thresh` now also takes a mapping keyed on the `phase` coordinate — one threshold per label, labels the mapping does not list never trigger, which is how a characteristic function keeps carrying its noise class without that class ever producing a pick (keying on the label rather than its position matters: the label order of a model belongs to its weight set and flips between them). `coords` gains `"auto"` (the default): scalar coordinates lead as constant columns, then the other dimension coordinates, then the picked dimension — identity first, measurement last, whatever the input's dimension order — and non-dimensional coordinates can be named too. `flush()` closes the triggers still open at the end of a run, as `obspy.trigger_onset` does, so the last pick of a record is no longer lost, and chunking along another dimension than the picked one now answers exactly (each such chunk is a whole record of other lanes, run from a fresh state). The lowercase twin `xdas.trigger` joins the top level; `xdas.trigger` the module remains importable as a compatibility home re-exporting `Trigger` and keeping `find_picks` unchanged. Note: the re-exported `Trigger`'s `dim` default is now `"time"`, not `"last"` (@atrabattoni). - **`Annotate`.** The SeisBench wrapper is rebuilt around what a *weight set* declares rather than what the architecture suggests: the window overlap, the stacking rule (`"avg"` or `"max"`, reproducing SeisBench's `nanmean`/`nanmax` over covering windows exactly), the blinding and the preprocessing arguments are all read off the model instance, and any annotate argument can be overridden at the call (`Annotate(model, scale=2.0)`). The component dimension is found by its labels — each ending with a distinct letter of the model's `component_order`, with SeisBench's flexible horizontal matching — never by its name, and `component_strategy` covers SeisBench's whole range (`"auto"`, `"clone"`, `"pad"`, a named slot, `"strict"`). The output is laid out sample-last, `(..., "phase", dim)`, so the characteristic function of one phase of one channel is contiguous; the end-aligned final window SeisBench appends is emitted at `flush()`, so the output spans the input; and a model whose `annotate_batch_post` breaks the `(batch, samples, classes)` stacking contract is named instead of surfacing as a bare broadcast error. Chunked along its own dimension the sliding window carries across chunks exactly; chunked along another dimension each chunk is a whole record settled on the spot. `MLPicker` and `xdas.mlpicker` remain as deprecated aliases until 0.4 (@atrabattoni). 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/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..e2aefa7b --- /dev/null +++ b/docs/user-guide/pipeline/picking.md @@ -0,0 +1,306 @@ +# 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, keeping the order of the dimensions it was +given: + +```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` rows the repr shows 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. From bdd1a61347f213d61ff876d5bfd8219dfdb4fa51 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 15:38:50 +0200 Subject: [PATCH 25/48] pin the corners Phase D's test moves uncovered MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The slow real-weight MLPicker tests took two small facts with them when the fake-model suite replaced them: randn_wavefronts had no contract of its own (shape, coordinates, seeded reproducibility — it is the DAS synthetic the picking walkthrough builds on), and no test constructed an Annotate without naming a device, so the CUDA-if-available default went unexercised. --- tests/test_atoms_ml.py | 7 +++++++ tests/test_xdas.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py index 3232ba07..725cd219 100644 --- a/tests/test_atoms_ml.py +++ b/tests/test_atoms_ml.py @@ -2023,3 +2023,10 @@ def test_flush_drains_the_queue(self): 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_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]) From 35b161657682051cd48d70036fb728b5f12fde43 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 18:19:40 +0200 Subject: [PATCH 26/48] resampling carries the labels of the samples it keeps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A stage that changes the number of samples along a dimension has to say what became of the *other* coordinates attached to it. Polyphase and UpSample copied them through untouched, so a decimated array kept its full-length station codes: the dimension had 5 000 samples and the coordinate naming them still had 10 000 values, and every lane was labelled with the code of the lane at its own index — the first half of the fibre, shifted by one channel each. Picking a DAS acquisition answers with those codes, so the picks came out attributed to the wrong channels; found by running the ABYSS pipeline against its archived reference, whose stations sit one every 30.6 m where these sat one every 15.3 m. The labels now follow the samples: output k is drawn from input k * down / up, and that is the input whose label it takes. They carry no group-delay shift — a label names a source sample, where the dimension coordinate names a position, which is what the delay compensation is about — and the mapping depends on the sampling grid alone, so chunking cannot move a label. DownSample was already right: isel subsamples every coordinate attached to the dimension. Note for the record: xd.concat drops non-dimensional coordinates along the concatenated dimension, so re-joining chunked output loses them. That is pre-existing (a plain split/concat round trip loses them too) and left alone here; the streaming path is unaffected, since each chunk reaches the sink already labelled. --- tests/test_atoms_tasks.py | 61 +++++++++++++++++++++++++++++++++++++++ xdas/atoms/core.py | 10 ++----- xdas/atoms/kernel.py | 43 ++++++++++++++++++++++++++- 3 files changed, 106 insertions(+), 8 deletions(-) diff --git a/tests/test_atoms_tasks.py b/tests/test_atoms_tasks.py index 327163bb..becd2f09 100644 --- a/tests/test_atoms_tasks.py +++ b/tests/test_atoms_tasks.py @@ -1,6 +1,7 @@ import inspect import numpy as np +import numpy.testing as npt import pytest import xdas as xd @@ -413,3 +414,63 @@ def test_pipeline_chunk_invariant_over_cuts_and_gaps(self): 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/xdas/atoms/core.py b/xdas/atoms/core.py index 87891638..b5891390 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -275,12 +275,6 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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. - merge: callable or None - The optional hook folding the per-leaf results of a collection - walk into one object. ``None`` (the default) means the atom has - none and the tree is rebuilt as it always was. See - :meth:`~xdas.atoms.Trigger.merge` for an implementation. - Methods ------- gather(mapping) @@ -299,7 +293,9 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): Resets the atom to its initial state. merge(results) Optional. Folds the leaf results of a collection walk into one - object; undefined by default. + 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) diff --git a/xdas/atoms/kernel.py b/xdas/atoms/kernel.py index cbddcac5..1935f8c0 100644 --- a/xdas/atoms/kernel.py +++ b/xdas/atoms/kernel.py @@ -28,6 +28,40 @@ def _along(axis, ndim, slc): 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`. @@ -259,6 +293,9 @@ def call(self, da, **flags): # 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) @@ -457,7 +494,11 @@ def _coords(self, da, first, stop, start): data["tolerance"] = base + drift coords = da.coords.copy() coords[name] = Coordinate(data, name) - return coords + # Output `index` is drawn from input sample `index * down / up`, which + # `first`/`stop` keep inside this chunk by construction. + positions = np.rint(np.arange(first, stop) * self.down / self.up) + positions = np.clip(positions.astype(int) - start, 0, da.sizes[name] - 1) + return _carry_labels(coords, name, positions) def _upsampled(self, count, delta): """Return the span of *count* upsampled samples, at coordinate resolution.""" From c3f689abd085aeb69fef8496553992bcf79bfd9d Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 18:27:27 +0200 Subject: [PATCH 27/48] release notes: the resampled-labels fix --- docs/release-notes.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/release-notes.md b/docs/release-notes.md index e2b7609d..b122c1dd 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -57,6 +57,7 @@ - Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@atrabattoni). - Fix `DataCollection` coercing a `pandas.DataFrame` leaf into a `DataArray`, producing collections whose leaves raised on `repr`. A table is now a leaf of its own kind (@atrabattoni). - Fix a directory sink joining its chunks along the wrong dimension when the pipeline's output does not lead with the chunked one — `(distance, time)` chunks written from a time-chunked source were stacked along `distance`. `DataArrayWriter` now takes the dimension as a `dim` argument (still `"first"` when left unsaid) (@atrabattoni). +- Fix resampling losing track of the *other* coordinates of the dimension it resamples. `Polyphase` and `UpSample` rebuilt the dimension coordinate and copied everything else through untouched, so decimating a DAS acquisition by two left its `station` coordinate at full length and labelled every lane with the code of the lane at its own index — the first half of the fibre. Picking answers with those codes, so the picks came out attributed to the wrong channels. Non-dimensional coordinates now follow the samples: output `k` takes the label of input `k * down / up`, carrying no group-delay shift (a label names a source sample, where the dimension coordinate names a position), and the mapping depends on the sampling grid alone so chunking cannot move a label (@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. Registered defaults now fill in only when the caller says nothing (@atrabattoni). - Fix `sel` refusing an exact label look-up on a coordinate whose values are not sorted, such as a categorical axis like `["P", "S", "N"]`. The overlap guard covered every kind of selection, but only ordered look-ups — a slice, or `method="nearest"` and friends — need a sorted axis; naming a label does not. Those stay guarded, `da.sel(phase=["S", "P"])` now works and returns the labels in the requested order (@atrabattoni). From f1bf88cc70912ebef091747e252e8754d813ae57 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Tue, 11 Aug 2026 18:40:10 +0200 Subject: [PATCH 28/48] restore the blank lines the docstring sections need --- xdas/atoms/core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index b5891390..3150b2b3 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -275,6 +275,7 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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) From ab961cf05b2e10035bd1d260e29042cda8bbd153 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 06:40:54 +0200 Subject: [PATCH 29/48] the seam a one-sample chunk leaves behind is still judged A chunk of one sample declares no sampling interval of its own, so the seam it stored had no rate to compare the next chunk against and the judgement was skipped entirely: a gap right after such a chunk carried state across silently. The seam is now judged on the rate of whichever side knows one. Two smaller repairs on the same layer: `Sequential.fresh` rebuilt the clone by calling the constructor, which a pipeline that assembles its own stages (a `Picker`) does not accept, and the `out=` guard of the ufunc tracer compared tuples, so `out=` an array re-entered tracing instead of refusing. Along the way: the near-duplicate of the chunk joiner folded into the one it duplicated, the one-call flush helper inlined into the cascade that calls it, and the docstrings of the layer corrected. --- tests/test_atoms.py | 22 +++++++ tests/test_atoms_runs.py | 11 ++++ xdas/atoms/core.py | 120 +++++++++++++++++++-------------------- 3 files changed, 90 insertions(+), 63 deletions(-) diff --git a/tests/test_atoms.py b/tests/test_atoms.py index 1b3ca4cf..b54e003c 100644 --- a/tests/test_atoms.py +++ b/tests/test_atoms.py @@ -679,6 +679,13 @@ def test_out_to_another_atom_raises(self): 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): @@ -769,6 +776,21 @@ def test_fresh_recurses_into_sequences(self): 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): diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py index 17c06841..f73ad8b2 100644 --- a/tests/test_atoms_runs.py +++ b/tests/test_atoms_runs.py @@ -238,6 +238,17 @@ def test_single_sample_chunk_adopts_the_stream_rate(self, da): 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") diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index 3150b2b3..14702266 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -167,32 +167,9 @@ def _iter_results(tree): yield tree -def _flush_through(atoms, **flags): - """ - Codec-drain a linear chain of atoms. - - Flush the first atom and push its tail through the remaining atoms, then - flush the second one, and so on. Tails flow downstream as ordinary data: - each downstream atom folds them into its own state before being flushed - itself. - """ - atoms = list(atoms) - chunks = [] - for index, atom in enumerate(atoms): - tail = atom.flush() - for downstream in atoms[index + 1 :]: - tail = [ - chunk - for out in (downstream(x, **flags) for x in tail) - for chunk in _aschunks(out) - ] - chunks.extend(tail) - return chunks - - 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 ---------- @@ -201,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 @@ -243,21 +220,21 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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`) @@ -270,7 +247,7 @@ class Atom(np.lib.mixins.NDArrayOperatorsMixin): 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"`` @@ -360,7 +337,7 @@ def initialized(self): ) 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): @@ -424,8 +401,10 @@ def __call__(self, x, **flags): to `gather`, which may collapse a level into an axis of the input (see :meth:`gather`). - A single output chunk is returned bare; otherwise a - :class:`DataSequence` of chunks is returned. + 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) @@ -571,11 +550,11 @@ 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), 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. + 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"]: @@ -589,14 +568,15 @@ def _judge_seam(self, info): "regularize it first, e.g. `da[dim] = da[dim].to_regular()` " "or open the files with a tolerance" ) - if seam["delta"] is None: + # 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"] - seam["delta"]) > ( - tolerance - ): + if info["delta"] is not None and np.abs(info["delta"] - delta) > tolerance: return "rate" - jump = info["start"] - (seam["end"] + seam["delta"]) + jump = info["start"] - (seam["end"] + delta) if np.abs(jump) <= tolerance: return "continuous" return "gap" if jump > 0 else "overlap" @@ -688,18 +668,12 @@ def _fold(self, x, flags, path=None): def _join(self, chunks, dim): """Re-join output chunks: one chunk bare, else gap-aware concat or sequence.""" - if len(chunks) == 1: - return chunks[0] - if dim is not None and chunks and all(isinstance(c, DataArray) for c in chunks): - try: - return concat(chunks, dim) - except (TypeError, ValueError): - return DataCollection(chunks) - if chunks and all(isinstance(c, pd.DataFrame) for c in chunks): - # A pick table is a chunk type of its own: an atom emitting one per - # run, plus one at flush, still answers with a single table. - return pd.concat(chunks, ignore_index=True) - return DataCollection(chunks) + # 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): """ @@ -732,7 +706,8 @@ def iter_chunks(self, source, chunk_dim=None): Yields ------ - Zero or more output chunks per input chunk, then the flushed tail. + 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") @@ -834,8 +809,9 @@ def _refuse_chunked_along(self, dim, chunk_dim, x=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 {dim!r} and cannot " + 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" ) @@ -867,7 +843,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **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 kwargs["out"] != (self,): + 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( @@ -1059,10 +1035,22 @@ 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. Returns the drained chunks. + 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 {} - return _flush_through(self, **flags) + 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.""" @@ -1115,7 +1103,13 @@ def merge(self): def fresh(self): """Return a stateless clone: each stage cloned, config shared.""" - return type(self)([atom.fresh() for atom in self], name=self.name) + # 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))) @@ -1145,8 +1139,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 ---------- From 2015beb7f381fcc2522c4572104eecb3e181cf66 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 06:41:18 +0200 Subject: [PATCH 30/48] spell the root of a tile array out, and two error messages right `TileArray.root` is the path the whole archive relocates by, so it is part of the class rather than an attribute the constructor happens to set: declared and documented, it is also what the API page can point at. Plus the typos the docs build walked past. --- xdas/atoms/ml.py | 2 +- xdas/virtual/hdf5.py | 4 ++-- xdas/virtual/tiles.py | 4 ++++ 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index 8dd06ed8..cf841e44 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -31,7 +31,7 @@ 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) 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..3f23316b 100644 --- a/xdas/virtual/tiles.py +++ b/xdas/virtual/tiles.py @@ -669,6 +669,10 @@ 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. Rewriting it relocates the whole archive. + 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 From 99249c3427a065b6c67ba4eaa23ef6d78215a233 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 06:41:31 +0200 Subject: [PATCH 31/48] the guides teach the pipeline this branch actually ships MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The atoms and processing pages still taught `Sequential([...])`, `Partial` and hand-built loader/writer pairs — correct, but no longer what one would write. They now open on the function forms and `>>`, on physical parameters, and on `process(source, out=...)` with the source and sink tables, keeping the explicit form as what to reach for when the ends need configuring. The picking page's annotation repr predated the sample-last layout, the obspy page duplicated `xd.stack` by hand and called the legacy engine an alias of the new one, and the API pages were missing `trim_overlaps`, `DataCollection.select` and the collection hooks. --- docs/api/atoms.md | 13 ++- docs/api/xdas.md | 2 + docs/user-guide/io/obspy.md | 26 ++--- docs/user-guide/pipeline/atoms.md | 135 ++++++++++++++++--------- docs/user-guide/pipeline/picking.md | 22 ++-- docs/user-guide/pipeline/processing.md | 124 ++++++++++++++++------- 6 files changed, 208 insertions(+), 114 deletions(-) diff --git a/docs/api/atoms.md b/docs/api/atoms.md index 1c79d485..d2685e61 100644 --- a/docs/api/atoms.md +++ b/docs/api/atoms.md @@ -36,6 +36,8 @@ Methods Atom.reset Atom.process Atom.iter_chunks + Atom.gather + Atom.merge Atom.save_state Atom.set_state Atom.load_state @@ -91,9 +93,7 @@ Methods ## Task atoms -Public processing vocabulary with physical parameters only. Each task atom has -a function form exported at the top level of `xdas` (e.g. `xdas.filter`, -`xdas.decimate`). +Public processing vocabulary with physical parameters only. ```{eval-rst} .. autosummary:: @@ -107,6 +107,11 @@ a function form exported at the top level of `xdas` (e.g. `xdas.filter`, 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 ``` @@ -179,4 +184,4 @@ parameters, designed by the task atoms from the data at the first call. Rechunk SOSFilter UpSample -``` \ No newline at end of file +``` diff --git a/docs/api/xdas.md b/docs/api/xdas.md index 66a05a25..40ed6d73 100644 --- a/docs/api/xdas.md +++ b/docs/api/xdas.md @@ -35,6 +35,7 @@ sortby split stack + trim_overlaps plot_availability ``` @@ -154,6 +155,7 @@ Methods .. autosummary:: :toctree: ../_autosummary + DataCollection.select DataCollection.query DataCollection.issequence DataCollection.ismapping 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..b751bd02 100644 --- a/docs/user-guide/pipeline/atoms.md +++ b/docs/user-guide/pipeline/atoms.md @@ -14,73 +14,118 @@ 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 a processing vocabulary — filtering, resampling, integration, +spectra, machine-learning picking — as *atoms*: elementary operations that +compose into a pipeline. A pipeline built this way runs unchanged on an array +in memory and, chunk by chunk, on an archive that does not fit in one (see +[](processing.md)), which is what makes it worth defining one rather than +calling functions in a row. -## Chaining elementary operations (atoms) +## Applying and composing -There are three "flavours" declaring the atoms that can be used to compose a sequence, illustrated by the following example: +Every atom has a function form at the top level of `xdas`. Called on data, it +applies: -```{code-cell} +```{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.filter(da, (5.0, None), 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.decimate(..., 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. +The parameters are physical: corner frequencies in hertz, target rates in +hertz, window lengths in seconds. They keep their meaning whatever the sampling +rate of the data the pipeline is later given. + +Calling the pipeline applies it: + +```{code-cell} +result = pipeline(da) +result.plot(yincrease=False) +``` -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). +The same pipeline can be reused: it is defined once and carries no data. -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. +Ordinary NumPy expressions compose too. Under the `...` seed they are *traced* +— appended to the pipeline rather than computed — so an expression reads as +mathematics: -## Executing a sequence +```{code-cell} +energy = 20 * np.log10(np.abs(xd.decimate(..., 25.0, dim="time"))) +energy +``` -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: +## Wrapping your own functions -```{code-cell} -from xdas.synthetics import wavelet_wavefronts +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: -For executing a sequence on chunked data (e.g., larger-than-memory data sets), see the next section: [](processing.md). +```{code-cell} +from xdas.atoms import Partial + +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/picking.md b/docs/user-guide/pipeline/picking.md index e2aefa7b..72105f9f 100644 --- a/docs/user-guide/pipeline/picking.md +++ b/docs/user-guide/pipeline/picking.md @@ -160,20 +160,16 @@ Coordinates: ``` {py:class}`~xdas.atoms.Annotate` consumes the component dimension and appends -its classes as a `phase` dimension, keeping the order of the dimensions it was -given: +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]] + +[[nan nan nan ... nan nan nan] + [nan nan nan ... nan nan nan] + [nan nan nan ... nan nan nan]] Coordinates: network: 'IA' station: 'DBNFM' @@ -182,9 +178,9 @@ Coordinates: * time (time): 2026-05-20T00:21:59.771 to 2026-05-20T00:23:59.761 ``` -The `nan` rows the repr shows 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 `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 diff --git a/docs/user-guide/pipeline/processing.md b/docs/user-guide/pipeline/processing.md index 2de8e131..c8b5301a 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,53 @@ 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 (8 GiB by default) it +raises rather than filling the machine. + +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") +``` From de8dc3982faa7268cbda9daab85faba6b458e473 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 06:41:38 +0200 Subject: [PATCH 32/48] release notes: the 0.2.9 story, not its development Rewritten from the development log it had become into what a 0.2.8 user meets: a section each for the tiles backend and the atoms rework, the rest of the new API in one paragraph apiece, and the breaking section reserved for what one actually runs into. Everything else moved to improvements or refactoring, and the details left to the docs. --- docs/release-notes.md | 86 ++++++++++++++++++++----------------------- 1 file changed, 39 insertions(+), 47 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index b122c1dd..9b151d5f 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -2,64 +2,56 @@ ## 0.2.9 (unreleased) +### 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 +- **Composition and function forms.** Atoms compose into pipelines with `>>`, ordinary numpy expressions trace under the `...` seed (`20 * np.log10(np.abs(atom))` extends the pipeline instead of computing), and every task atom has a top-level function form: `xdas.decimate(da, 50.0)` applies eagerly, `xdas.decimate(..., 50.0)` returns the atom (@atrabattoni). +- **Task atoms speak physical units.** The new `xdas.atoms.tasks` vocabulary — `Filter`, `Decimate`, `Resample`, `Integrate`, `Differentiate`, `STFT`, `detrend`, `taper`, `hilbert` and friends — takes corner frequencies and target rates in Hz and window lengths in seconds, so every parameter keeps its meaning when the sampling rate changes (@atrabattoni). +- **Polyphase resampling.** The machine-parameter atoms (`LFilter`, `SOSFilter`, `DownSample`, …) move to the expert layer `xdas.atoms.kernel`, joined by a `Polyphase` kernel that fuses upsampling, FIR filtering and downsampling into a single pass — 2.6–8.7× faster on typical resamplings, without promoting float32 data (@atrabattoni). +- **Gap-aware chunked processing.** Stateful atoms judge the seams of their input stream: state carries across continuous chunks and is flushed and restarted at gaps and rate changes; eager calls split gappy input the same way, so results no longer depend on chunking. The new `flush()` lifecycle drains buffered tails at the end of a stream, atoms that cannot answer correctly chunk by chunk (such as the `fft` functions along the chunked dimension) now raise instead of answering wrong, and `xdas.testing.assert_chunk_invariant` asserts that a pipeline returns identical results eagerly and streamed, gaps included (@atrabattoni). +- **`process()` on every atom, with source and sink auto-dispatch.** `pipeline.process(source, out=...)` infers both ends: a `DataArray` runs eagerly or chunk by chunk, a virtual array streams with storage-aligned chunks, a path opens it, `"tcp://..."` streams over ZeroMQ and `xdas.watch(dir)` follows a growing directory; `out=` takes a directory, a `.csv` file, a URL, a configured writer, or `None` to accumulate. `process()` also walks `DataCollection`s, labelling each leaf's result with its tree path, and memory guards make footguns loud: an eager call or accumulation beyond the `"memory_limit"` configuration entry (default 8 GiB) raises with a pointer to the streaming path (@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, coordinate selection and a `flush` that 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). + ### 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). - **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`.** Collapse a level of a collection into an array dimension — the inverse of `combine_by_coords`, which concatenates *along* an existing one: `xd.stack(dc, "channel")` turns each station's traces into one `(channel, time)` array, keyed by the level's own keys. The new dimension is named after the level it collapsed, so nothing is renamed behind your back (`dim=` chooses otherwise), everything below the level is merged in lock-step, and tile-backed leaves stay tile-backed, so a collection of virtual arrays stacks without reading a byte. Leaves that do not share their other coordinates raise, naming what disagreed; `join="inner"` trims them to what they all have and `join="outer"` pads the rest with NaN. Agreement is judged on the sampling grid, not on tie points: leaves that describe one grid to within `tolerance` — by default a hundredth of a sample, and only where a nominal sampling interval is declared — are snapped onto the first leaf's coordinate, so three components whose start times were rounded a nanosecond apart stack without a join. `tolerance=False` restores strict equality, and a join that would interleave two grids into an array longer than their span can hold now raises instead of returning it (@atrabattoni). - -- **`process()` with source and sink auto-dispatch.** `process()` is now a method on every atom and a dispatch boundary: `pipeline.process(da, out="results/")` infers both ends. Sources dispatch on the input value — an in-memory `DataArray` runs eagerly (or chunk by chunk with `chunks=`), a virtual one streams through a loader with storage-aligned `chunks="auto"`, a path/directory/glob opens with `open_mfdataarray`, `"tcp://..."` subscribes over ZeroMQ, and any iterable of chunks (a generator, a loader) is consumed as is. Sinks dispatch on the out spec crossed with the first output chunk, so writer creation is deferred to what the pipeline actually emits: a directory stores `DataArray` chunks joined along the chunked dimension (or an SDS archive for `Stream` chunks), `*.csv` appends DataFrames, `"tcp://..."` publishes, `out=None` accumulates and returns the joined result, and a configured writer instance passes through. A chunked source with discontinuities announces them upfront — one warning with the count, read off the source coordinate before any data. The historical `process(atom, loader, writer)` form keeps working unchanged (@atrabattoni). -- **`xdas.watch` and unbounded sources.** Realtime is now *named*: `pipeline.process(xd.watch("/incoming", engine=...), out=...)` watches a directory forever, and a bare directory path always means "process what is there". Unbounded sources (watch, ZMQ subscriptions) get streaming semantics — throughput-style progress, a clean `KeyboardInterrupt` that flushes the pipeline and returns the writer result, `until=` to stop at a coordinate value (inclusive, truncating the last chunk), and a warning at each seam as it arrives, since a realtime source cannot be inspected upfront (@atrabattoni). -- **Memory guards.** The new `"memory_limit"` configuration entry (default 8 GiB) makes footguns loud: an eager call on a huge virtual array and an `out=None` accumulation that outgrows the limit both raise with the estimated size and a pointer to `.process(out=...)` (@atrabattoni). -- **Collection walks carry their tree path, and atoms can merge or gather them.** Walking a collection, each leaf's result is now labelled with the path it was reached by *as it is produced* — one leading column per named level, filled with that level's key, so a pick found under `IA / DBNFM / --` comes out with its network, station and location before its time and value (positional levels contribute their index; a column that already exists as a scalar coordinate stays one column, the tree path winning with a warning on a genuine disagreement). Atoms may declare a `merge(results)` hook folding those labelled results — `Trigger.merge` is a plain concat, so `xd.trigger(dc, ...)` answers a whole network with one flat table, and `merge=False` opts out — and a `gather(mapping)` hook offered each mapping level before the walk descends: `Annotate` implements it, collapsing a channel level into the component dimension through `xd.stack` (what counts as a component is a property of the *model*, which is why the reader cannot do this and the atom can), with deliberately conservative recognition so a station level is never silently folded. `Sequential` delegates gather to its first claiming stage and merge to its last declaring one (@atrabattoni). -- **`process()` walks collections.** `atom.process(dc, chunks=..., out=...)` walks a `DataCollection` exactly as `atom(dc)` does — gather first, mapping levels recursed, sequence levels folded with state carried across the elements — with each leaf streamed through the single-source path, so `atom.process(dc, out=None) == atom(dc)` and the streaming form is not second-class precisely where it matters most: a leaf too large to call eagerly at all. Sinks gain one rule per destination: `out=None` accumulates per leaf and merges; a `*.csv`, a URL or a ready writer is *shared*, every leaf appending to one table with the tree-path columns keeping the rows apart; a directory *fans out*, one subdirectory per leaf mirroring the tree path (@atrabattoni). -- **The GPU atom is asynchronous behind a bounded output queue.** The process loop stays serial and CPU atoms stay internally parallel; only `Annotate` overlaps its work with the device: each completed window's device-to-host transfer is issued asynchronously (pinned staging, non-blocking, an event marking completion) and `call()` emits only the outputs whose transfer has completed, `flush()` draining the rest — the 0..n contract already allows late emission, so chunk order, seams and results are untouched. `max_buffers` bounds the in-flight transfers (default 2; 0 restores synchronous emission); on the CPU the queue completes on arrival and nothing changes (@atrabattoni). -- **`Picker` / `xd.pick`, the headline: waveforms in, one pick table out.** `Picker(model)` assembles the whole pipeline SeisBench's `model.classify(stream)` runs — the preprocessing filter the weight set ships (if any), `Resample` to the weight set's own rate, `Annotate`, `Trigger` with the weight set's own per-phase thresholds — from the weight set and nothing else, so two pickers built on one model class can differ in stage count, sampling rate and thresholds. Being a `Sequential`, a picker composes with `>>`, pickles, shows its stages in `repr` and streams with `picker.process(source, out=...)`; it inherits `Annotate.gather` and `Trigger.merge`, so `xd.pick(dc, model)` answers a whole ObsPy-style network tree with one flat table, each pick labelled `network station location phase time value`. The one deliberate difference from SeisBench is the resampler: obspy's `Trace.resample` halves the amplitude at half the input Nyquist (its frequency-domain hann window), where the polyphase filter used here is flat; `resample=False` drops the stage (@atrabattoni). -- **`Trigger` joins the task vocabulary, in `xdas.atoms.detect`.** `thresh` now also takes a mapping keyed on the `phase` coordinate — one threshold per label, labels the mapping does not list never trigger, which is how a characteristic function keeps carrying its noise class without that class ever producing a pick (keying on the label rather than its position matters: the label order of a model belongs to its weight set and flips between them). `coords` gains `"auto"` (the default): scalar coordinates lead as constant columns, then the other dimension coordinates, then the picked dimension — identity first, measurement last, whatever the input's dimension order — and non-dimensional coordinates can be named too. `flush()` closes the triggers still open at the end of a run, as `obspy.trigger_onset` does, so the last pick of a record is no longer lost, and chunking along another dimension than the picked one now answers exactly (each such chunk is a whole record of other lanes, run from a fresh state). The lowercase twin `xdas.trigger` joins the top level; `xdas.trigger` the module remains importable as a compatibility home re-exporting `Trigger` and keeping `find_picks` unchanged. Note: the re-exported `Trigger`'s `dim` default is now `"time"`, not `"last"` (@atrabattoni). -- **`Annotate`.** The SeisBench wrapper is rebuilt around what a *weight set* declares rather than what the architecture suggests: the window overlap, the stacking rule (`"avg"` or `"max"`, reproducing SeisBench's `nanmean`/`nanmax` over covering windows exactly), the blinding and the preprocessing arguments are all read off the model instance, and any annotate argument can be overridden at the call (`Annotate(model, scale=2.0)`). The component dimension is found by its labels — each ending with a distinct letter of the model's `component_order`, with SeisBench's flexible horizontal matching — never by its name, and `component_strategy` covers SeisBench's whole range (`"auto"`, `"clone"`, `"pad"`, a named slot, `"strict"`). The output is laid out sample-last, `(..., "phase", dim)`, so the characteristic function of one phase of one channel is contiguous; the end-aligned final window SeisBench appends is emitted at `flush()`, so the output spans the input; and a model whose `annotate_batch_post` breaks the `(batch, samples, classes)` stacking contract is named instead of surfacing as a bare broadcast error. Chunked along its own dimension the sliding window carries across chunks exactly; chunked along another dimension each chunk is a whole record settled on the spot. `MLPicker` and `xdas.mlpicker` remain as deprecated aliases until 0.4 (@atrabattoni). -- **`STFT`.** The spectral vocabulary joins the task-atom route: `STFT` streams complex frames with window length and hop in physical units — both are snapped, the window to the next fast FFT size of the target and the hop to a whole sample count — with an expert `nfft` to zero-pad and a `scaling=` of `"spectrum"` or `"psd"`, so `np.abs(stft)**2` composes to an exact spectrogram. Only fully computable frames are ever emitted: the unconsumed tail is buffered across chunks and dropped at gaps, so chunked processing emits exactly the eager frames and no frame ever spans a discontinuity. Built on `scipy.signal.ShortTimeFFT` internally, with the `xdas.stft` function form at the top level (@atrabattoni). -- The `xdas.fft` functions (`fft`, `rfft`, `ifft`, `irfft`) now declare whole-record semantics: used as atoms in a chunked pipeline they raise along the transformed dimension instead of silently computing one transform per chunk. Transforming along another dimension than the chunked one keeps working (@atrabattoni). -- **`xdas.testing.assert_chunk_invariant`.** The chunk-safety story in one call: run a pipeline eagerly and streamed and assert the two agree — values, coordinates and all. The invariant is quantified over *cuts* (the same stream re-chunked at derived non-divisor sizes, so boundaries land elsewhere) and over *gaps* (`xdas.testing.inject_gaps` places real discontinuities in the input first, so seam resets are exercised at boundaries that do not line up with them). It is both the CI harness for every stateful atom xdas ships and the tool to run on your own pipelines before trusting them chunked (@atrabattoni). -- **Continuous-run semantics.** Stateful atoms now understand gaps: every atom judges the seams of its own input stream from the chunk coordinates — a continuous chunk carries state across, a gap or rate change flushes the previous run and restarts (redesigning coefficients on rate changes), an overlap raises, and the `on_discontinuity="reset"|"raise"` policy makes strict runs opt-in. Eager calls auto-split gappy input into runs, process each with a fresh state and re-join the outputs with the gaps kept in the coordinates, so filters never cross discontinuities — and the split is announced: a warning states how many discontinuities the source has and that state is flushed and reset at each. Sequence collections fold through the same seam-aware machinery — `concat(atom(split(da, anywhere)))` equals `atom(da)` for arbitrary split points — and mapping collections map over their leaves. Chunked processing along a dimension now requires a regular coordinate (a declared `sampling_interval`) on it, raising with a pointer to `to_regular()` instead of silently carrying state across unverifiable seams (@atrabattoni). -- **`flush()` lifecycle and the transducer contract.** `call()` now maps one input chunk to zero or more output chunks, and the new `Atom.flush()` drains what remains: buffering atoms emit their tail at the end of the stream, at every seam and at the end of every eager call (`Sequential.flush` cascades codec-drain style, and `process()` drains the pipeline at the end of the stream). Reductions fall out for free: a `call()` that accumulates and returns nothing plus a `flush()` that emits the result gives constant-memory streaming statistics. `Atom.iter_chunks(source)` exposes the whole machinery as a plain generator — the manual chunk loop with buffering, seams and flushing handled inside — and writers now silently drop empty chunks (@atrabattoni). -- **`Rechunk` kernel atom.** `Rechunk({"time": n})` (and its function form `xdas.rechunk`) merges and splits streaming chunks to a target size in samples — a performance knob, e.g. to restore a workable cadence after a decimation shrank the chunks — without ever merging across a discontinuity (@atrabattoni). -- Chunked `DownSample` (and thus the stateful decimation path) no longer drops its trailing samples when the stream length is not a multiple of the factor: the buffered remainder is emitted by the new `flush()` lifecycle. The `"first"`/`"last"` dimension aliases are now resolved against the data before being compared with the chunked dimension, so a kernel built with its documented default no longer skips allocating its seam state, and `UpSample` handles one-sample chunks (@atrabattoni). -- **Task atoms with physical units.** A new public processing vocabulary where every parameter keeps its meaning when the sampling rate changes: `Filter` (one atom for all bands — a `(low, high)` corner pair in Hz with `None` opening one end, `ftype="iir"/"fir"`, `zerophase`), `Decimate` and `Resample` (target rate in Hz, both riding the polyphase kernel — the filter-at-full-rate-then-discard chain is never taken), `Integrate` and `Differentiate` (chunk-correct, carrying state across seams), plus whole-record `detrend`, `taper`, `hilbert`, `sliding_mean_removal` and `medfilt` (kernel lengths now in seconds/meters), each refusing chunked execution along its working dimension instead of silently answering wrong. Task atoms default to `dim="time"` and live in `xdas.atoms.tasks` (@atrabattoni). -- **Function forms at the top level.** Every task atom generates a top-level function with a synthesized signature and docstring: `xdas.decimate(da, 50.0)` applies eagerly, `xdas.decimate(..., 50.0)` returns the atom, and passing an atom extends a pipeline — so the same code runs eagerly on a slice and chunked on an archive by seeding it with `...` (@atrabattoni). -- **Polyphase resampling in a kernel layer.** The exact machine-parameter atoms move to the expert layer `xdas.atoms.kernel` (`LFilter`, `SOSFilter`, `DownSample`, `UpSample`, still importable from `xdas.atoms`), joined by the new `Polyphase` kernel: upsample, FIR filter and downsample fused into a single `scipy.signal.upfirdn` pass that computes only the output samples surviving the decimation and never materialises the zero-stuffed signal (which for `up=4` allocated a four times larger, mostly-zero array). `FIRFilter` is born with `up=`/`down=` and `ResamplePoly` rides it, so the upsample/filter/downsample trio collapses to one child atom — on a 254 MiB chunk that is 2.6× on a decimation by two along distance and 8.7× on a 62.5 → 50 Hz resampling. The taps are cast down to the data precision, so float32 stays float32 instead of being promoted by the filter; a target rate the coordinate resolution cannot represent exactly (100 Hz → 30 Hz is 10/3 ns per sample) declares its residual drift as jitter instead of rejecting its own sampling interval (@atrabattoni). -- **`>>` composition and operator tracing.** Atoms compose into pipelines with `>>`/`>>=` (bare callables auto-wrap, `da >> atom` applies), and ordinary numpy expressions trace under the `...` seed: `20 * np.log10(np.abs(atom))` appends `absolute → log10 → multiply` to the pipeline instead of computing. Tracing covers ufuncs exactly — a traced expression involving two atoms (fan-in) raises at the line that wrote it rather than silently computing. Composition has value semantics: passing a `Sequential` to an atomized function returns a new extended pipeline instead of mutating (and aliasing) the input — the mutating form also returned `None`, breaking chained composition. `xdas.atoms.as_function` generates the function form of any atom class, and atoms gain `fresh()` (a stateless clone whose config is shared by reference) while `initialized` now recurses into nested atoms (@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). -- **Process pools for chunk ingress and egress.** `DataArrayLoader` and `DataArrayWriter` accept `pool="processes"`, which reads and writes chunks in worker processes instead of threads: compressed HDF5 decodes and compresses under the global HDF5 lock, so extra *threads* only contend, while processes each hold their own lock. What crosses to a worker on the read side is the manifest of the chunk — a sliced virtual array, kilobytes — so each worker reads its own files, and the loaded chunk comes back through Ray's shared-memory object store: written once by the worker, mapped zero-copy by the parent, arriving read-only (the immutability convention atoms already follow). End to end on a compressed ZFP archive at 16 workers, ingest goes from 137 to 1378 MiB/s and egress from 151 to 1556. Ray is an optional dependency (`pip install xdas[ray]`); `pool="threads"` remains the default (@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. Ray is an optional dependency (`pip install xdas[ray]`); `pool="threads"` remains the default (@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). -- `xdas.concat` opening a *new* dimension over `VirtualSource`-backed arrays no longer raises `TypeError: only VirtualSource object can be provided`. Whether the result can stay virtual was decided before `expand_dims`, which no virtual source can follow — a stack of sources is a longer axis, never an extra one — so a `VirtualStack` was promised over arrays that had already been loaded (@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). -- Fix `sel` raising on a string-labelled coordinate: the overlap guard differenced the coordinate values, which numpy cannot do on strings, so `da.sel(phase="P")` failed before any selection happened. Label selection — scalar, list, reordering list and slice — now works, and `isel` with hard-coded positions is no longer the only option (@atrabattoni). -- Fix `DataCollection` coercing a `pandas.DataFrame` leaf into a `DataArray`, producing collections whose leaves raised on `repr`. A table is now a leaf of its own kind (@atrabattoni). -- Fix a directory sink joining its chunks along the wrong dimension when the pipeline's output does not lead with the chunked one — `(distance, time)` chunks written from a time-chunked source were stacked along `distance`. `DataArrayWriter` now takes the dimension as a `dim` argument (still `"first"` when left unsaid) (@atrabattoni). -- Fix resampling losing track of the *other* coordinates of the dimension it resamples. `Polyphase` and `UpSample` rebuilt the dimension coordinate and copied everything else through untouched, so decimating a DAS acquisition by two left its `station` coordinate at full length and labelled every lane with the code of the lane at its own index — the first half of the fibre. Picking answers with those codes, so the picks came out attributed to the wrong channels. Non-dimensional coordinates now follow the samples: output `k` takes the label of input `k * down / up`, carrying no group-delay shift (a label names a source sample, where the dimension coordinate names a position), and the mapping depends on the sampling grid alone so chunking cannot move a label (@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. Registered defaults now fill in only when the caller says nothing (@atrabattoni). -- Fix `sel` refusing an exact label look-up on a coordinate whose values are not sorted, such as a categorical axis like `["P", "S", "N"]`. The overlap guard covered every kind of selection, but only ordered look-ups — a slice, or `method="nearest"` and friends — need a sorted axis; naming a label does not. Those stay guarded, `da.sel(phase=["S", "P"])` now works and returns the labels in the requested order (@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 `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 From caa8bf1a77c02dbecbb138cd9e7e8c0cfe64c1a2 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 06:46:29 +0200 Subject: [PATCH 33/48] the realtime page says how a stream is processed, and the build is quiet The streaming page described the publisher and the subscriber but not what one does with them now: naming an address on either end of `process`, following a directory with `watch`, stopping on `until` or on a keyboard interrupt. The FAQ still answered the chunked-filter question with the kernel layer. A stray dangling reference and the one ambiguous cross-reference left in the build go with them. --- docs/user-guide/faq.md | 9 +++---- docs/user-guide/pipeline/streaming.md | 34 ++++++++++++++++++++++++--- xdas/atoms/kernel.py | 5 ++-- 3 files changed, 39 insertions(+), 9 deletions(-) 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/pipeline/streaming.md b/docs/user-guide/pipeline/streaming.md index e2f5d16a..764c8e0d 100644 --- a/docs/user-guide/pipeline/streaming.md +++ b/docs/user-guide/pipeline/streaming.md @@ -92,8 +92,36 @@ encoding = {"chunks": (10, 10), **hdf5plugin.Zfp(accuracy=1e-6)} publisher = ZMQPublisher(address, encoding) # Add encoding here, the rest is the same ``` -{py:class}`~xdas.io.asn.ZMQSubscriber` - ```{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/xdas/atoms/kernel.py b/xdas/atoms/kernel.py index 1935f8c0..c8c8b303 100644 --- a/xdas/atoms/kernel.py +++ b/xdas/atoms/kernel.py @@ -124,8 +124,9 @@ class SOSFilter(Atom): Parameters ---------- - sos : array-like, shape (n_sections, 6) - SOS filter coefficients as returned by e.g. :func:`scipy.signal.iirfilter`. + 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 From 0df7ddfdf1e91d322a7abda30b3f07220dff7ad2 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 06:47:13 +0200 Subject: [PATCH 34/48] release notes: tighter still on the three longest entries --- docs/release-notes.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index 9b151d5f..c49e61b4 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,9 +12,9 @@ - **Composition and function forms.** Atoms compose into pipelines with `>>`, ordinary numpy expressions trace under the `...` seed (`20 * np.log10(np.abs(atom))` extends the pipeline instead of computing), and every task atom has a top-level function form: `xdas.decimate(da, 50.0)` applies eagerly, `xdas.decimate(..., 50.0)` returns the atom (@atrabattoni). - **Task atoms speak physical units.** The new `xdas.atoms.tasks` vocabulary — `Filter`, `Decimate`, `Resample`, `Integrate`, `Differentiate`, `STFT`, `detrend`, `taper`, `hilbert` and friends — takes corner frequencies and target rates in Hz and window lengths in seconds, so every parameter keeps its meaning when the sampling rate changes (@atrabattoni). - **Polyphase resampling.** The machine-parameter atoms (`LFilter`, `SOSFilter`, `DownSample`, …) move to the expert layer `xdas.atoms.kernel`, joined by a `Polyphase` kernel that fuses upsampling, FIR filtering and downsampling into a single pass — 2.6–8.7× faster on typical resamplings, without promoting float32 data (@atrabattoni). -- **Gap-aware chunked processing.** Stateful atoms judge the seams of their input stream: state carries across continuous chunks and is flushed and restarted at gaps and rate changes; eager calls split gappy input the same way, so results no longer depend on chunking. The new `flush()` lifecycle drains buffered tails at the end of a stream, atoms that cannot answer correctly chunk by chunk (such as the `fft` functions along the chunked dimension) now raise instead of answering wrong, and `xdas.testing.assert_chunk_invariant` asserts that a pipeline returns identical results eagerly and streamed, gaps included (@atrabattoni). -- **`process()` on every atom, with source and sink auto-dispatch.** `pipeline.process(source, out=...)` infers both ends: a `DataArray` runs eagerly or chunk by chunk, a virtual array streams with storage-aligned chunks, a path opens it, `"tcp://..."` streams over ZeroMQ and `xdas.watch(dir)` follows a growing directory; `out=` takes a directory, a `.csv` file, a URL, a configured writer, or `None` to accumulate. `process()` also walks `DataCollection`s, labelling each leaf's result with its tree path, and memory guards make footguns loud: an eager call or accumulation beyond the `"memory_limit"` configuration entry (default 8 GiB) raises with a pointer to the streaming path (@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, coordinate selection and a `flush` that 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). +- **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 (8 GiB by default) 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). ### New Features - **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). From aea2946da455a363247a0d853b45e2870679fce1 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 07:15:35 +0200 Subject: [PATCH 35/48] a resampled label, and a chunked kernel, no longer depend on the chunking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three ways the same promise was broken. Rounding an output sample's source position could walk past the last input sample of the chunk, and the clip that pulled it back landed on a different sample than the eager call picked: flooring is in range by construction, so the label of an output no longer depends on where the stream was cut. The "first"/"last" alias was resolved against the *coordinates*, so a dimension carrying none kept the literal alias, never matched the chunked dimension, and the seam state was silently never allocated — chunked output quietly differed from eager. Resolution now goes against the dimensions, which a bare axis still has. And a whole-record function whose `dim` defaults to the last dimension refused *every* chunked dimension, since an unknown working dimension is refused against all of them: the fft functions could not be used in a pipeline chunked along another dimension, which is exactly what they document. The default is now declared where the guard can read it. --- tests/test_atoms_runs.py | 19 ++++++++++++++++++ tests/test_fft.py | 8 ++++++-- xdas/atoms/core.py | 42 +++++++++++++++++++++++++++++++--------- xdas/atoms/kernel.py | 24 ++++++++++++----------- xdas/fft.py | 8 ++++---- 5 files changed, 75 insertions(+), 26 deletions(-) diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py index f73ad8b2..d8bfd7e2 100644 --- a/tests/test_atoms_runs.py +++ b/tests/test_atoms_runs.py @@ -256,6 +256,25 @@ def test_stream_of_single_samples_has_nothing_to_judge(self): 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) diff --git a/tests/test_fft.py b/tests/test_fft.py index a30172ee..24071fdc 100644 --- a/tests/test_fft.py +++ b/tests/test_fft.py @@ -97,11 +97,15 @@ def test_chunked_along_transform_dim_raises(self): with pytest.raises(ValueError, match="whole record"): atom(chunk, chunk_dim="time") - def test_default_dim_is_conservative(self): + 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(time=slice(0, 50)), chunk_dim="time") + 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() diff --git a/xdas/atoms/core.py b/xdas/atoms/core.py index 14702266..125d98d0 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -502,16 +502,30 @@ def live(state): return live(self.state) - def _resolve_dim(self, x): - """Resolve the dimension this atom operates along on *x*, or ``None``.""" + 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 None + return dim if dim == "first": - dim = x.dims[0] - elif dim == "last": - dim = x.dims[-1] - return dim if dim in x.coords else None + 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.""" @@ -1215,13 +1229,16 @@ def __init__( 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 = dim + 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 @@ -1420,7 +1437,7 @@ def wrapper(*args, **kwargs): return wrapper -def _whole_record(dim_arg="dim"): +def _whole_record(dim_arg="dim", default=None): """ Mark a function as needing the whole record along its working dimension. @@ -1429,10 +1446,17 @@ def _whole_record(dim_arg="dim"): 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 diff --git a/xdas/atoms/kernel.py b/xdas/atoms/kernel.py index c8c8b303..7ccef722 100644 --- a/xdas/atoms/kernel.py +++ b/xdas/atoms/kernel.py @@ -93,7 +93,7 @@ def initialize(self, da, chunk_dim=None, **flags): # `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._resolve_dim(da) or self.dim + dim = self._dim_name(da) if dim == chunk_dim: n_sections = max(len(self.a), len(self.b)) - 1 shape = tuple( @@ -145,7 +145,7 @@ 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._resolve_dim(da) or self.dim + dim = self._dim_name(da) if dim == chunk_dim: n_sections = self.sos.shape[0] shape = (n_sections,) + tuple( @@ -192,7 +192,7 @@ def __init__(self, factor, dim="last"): 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._resolve_dim(da) or self.dim + dim = self._dim_name(da) if chunk_dim == dim: self.buffer = State(da.isel({self.dim: slice(0, 0)})) else: @@ -250,7 +250,7 @@ def call(self, da, **flags): 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._resolve_dim(da) or self.dim + name = self._dim_name(da) shape = tuple( self.factor * size if dim == name else size for dim, size in da.sizes.items() @@ -389,7 +389,7 @@ def initialize(self, da, chunk_dim=None, **flags): ) self.axis = State(da.get_axis_num(self.dim)) # Resolve the "first"/"last" alias before comparing (see `LFilter`). - dim = self._resolve_dim(da) or self.dim + dim = self._dim_name(da) if dim == chunk_dim: shape = tuple( self._history_size() if name == dim else size @@ -468,7 +468,7 @@ 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._resolve_dim(da) or self.dim + name = self._dim_name(da) coord = da.coords[name] delta = get_sampling_interval(da, name, cast=False) size = stop - first @@ -495,10 +495,12 @@ def _coords(self, da, first, stop, start): data["tolerance"] = base + drift coords = da.coords.copy() coords[name] = Coordinate(data, name) - # Output `index` is drawn from input sample `index * down / up`, which - # `first`/`stop` keep inside this chunk by construction. - positions = np.rint(np.arange(first, stop) * self.down / self.up) - positions = np.clip(positions.astype(int) - start, 0, da.sizes[name] - 1) + # 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): @@ -557,7 +559,7 @@ def __init__(self, chunks): 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._resolve_dim(da) or self.dim + dim = self._dim_name(da) if chunk_dim == dim: self.buffer = State(da.isel({self.dim: slice(0, 0)})) else: diff --git a/xdas/fft.py b/xdas/fft.py index 44380454..73a7bc2a 100644 --- a/xdas/fft.py +++ b/xdas/fft.py @@ -15,7 +15,7 @@ @atomized -@_whole_record() +@_whole_record(default="last") def fft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the discrete Fourier Transform along a given dimension. @@ -90,7 +90,7 @@ def func(x): @atomized -@_whole_record() +@_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. @@ -161,7 +161,7 @@ def rfft(da, n=None, dim=None, norm=None, parallel=None): @atomized -@_whole_record() +@_whole_record(default="last") def ifft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the inverse of `fft`. @@ -232,7 +232,7 @@ def func(x): @atomized -@_whole_record() +@_whole_record(default="last") def irfft(da, n=None, dim=None, norm=None, parallel=None): """ Compute the inverse of `rfft`. From c0a176ab6dbe3c99c2fc5f4f8dd660bc4b215b98 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 07:15:45 +0200 Subject: [PATCH 36/48] a fanned-out leaf writes to its own directory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sink path was taken from the annotation path, which drops the levels that have no name — the default for a plain `DataCollection({...})`. Every leaf then wrote into the same directory, each writer restarting its numbering, and reading one back gave another's data. The keys the walk descended by now address the sink, whether or not they name a column. Three more on the same path: `until=` was silently ignored on the eager source (the truncation only existed in the chunked loop); the accumulation guard sized chunks with `nbytes`, which tables and streams do not have, so the one case that most needs a limit — a collection walk accumulating into one pick table — was unbounded; and folding a sequence asked for chunks larger than its last element, which raises where the sibling path clamps. --- tests/test_process.py | 3 +- xdas/core/routines.py | 4 +-- xdas/processing/core.py | 75 ++++++++++++++++++++++++++++++----------- 3 files changed, 60 insertions(+), 22 deletions(-) diff --git a/tests/test_process.py b/tests/test_process.py index 353d2ea3..9fde6e75 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -794,7 +794,8 @@ def test_the_gathered_array_stays_virtual(self, monkeypatch, tmp_path): ) atom = xd.pick(..., picker_model("original"), device="cpu") assert isinstance(atom.gather(dc).data, TileArray) # nothing was read - monkeypatch.setitem(Config.config, "memory_limit", 1) + # 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/xdas/core/routines.py b/xdas/core/routines.py index c32b0af0..f47f2985 100644 --- a/xdas/core/routines.py +++ b/xdas/core/routines.py @@ -1649,8 +1649,8 @@ def _stack_arrays(objs, keys, level, dim, join, tolerance, path): messages, joinable = _leaf_mismatches(objs, keys) if messages and join is not None and joinable: objs = _join_leaves(objs, keys, joinable, join) - messages, joinable = _leaf_mismatches(objs, keys) - joinable = [] # already spent + 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)" diff --git a/xdas/processing/core.py b/xdas/processing/core.py index d6a600ca..8218186b 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -345,6 +345,11 @@ def _process_source(atom, source, out, chunks, until, path=None): 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 @@ -436,7 +441,7 @@ def _process_collection(atom, dc, out, chunks, until, merge): out=None)`` returns what ``atom(dc)`` returns. """ walk = _Walk(atom, _CollectionSink(out), chunks, until) - result = walk.sink.result(walk.level(dc, {})) + 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 @@ -467,7 +472,7 @@ def __init__(self, atom, sink, chunks, until): self.chunks = chunks self.until = until - def level(self, x, path): + 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 @@ -475,22 +480,22 @@ def level(self, x, path): # 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) + return self.level(gathered, path, where) name = getattr(x, "name", None) return DataCollection( { - key: self.level(value, _extend_path(path, name, key)) + 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) + 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(path), path)) + return _asleaf(self.stream(self.atom, x, self.sink.spec(where), path)) - def fold(self, x, path): + def fold(self, x, path, where): """ Fold a sequence level: one stream delivered in pieces. @@ -513,12 +518,14 @@ def fold(self, x, path): # Nothing to fold along: each element is a leaf of its own. return DataCollection( [ - self.level(element, _extend_path(path, name, index)) + self.level( + element, _extend_path(path, name, index), (*where, index) + ) for index, element in enumerate(x) ], name, ) - spec = self.sink.spec(path) + 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): @@ -664,13 +671,16 @@ def __init__(self, out): else: raise TypeError(f"cannot infer a writer from `out` of type {type(out)}") - def spec(self, path): - """Return the out spec of the leaf reached by *path*.""" + 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 - return os.path.join(str(self.out), *(str(key) for key in path.values())) + # 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.""" @@ -784,6 +794,11 @@ def get_source(source, chunks=None): # 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: @@ -885,7 +900,7 @@ def __init__(self, chunk_dim="time"): def write(self, chunk): """Accumulate one chunk, enforcing the in-memory size guard.""" self.chunks.append(chunk) - self.nbytes += getattr(chunk, "nbytes", 0) + self.nbytes += _sizeof(chunk) limit = config.get("memory_limit") if self.nbytes > limit: raise ValueError( @@ -933,6 +948,22 @@ def __iter__(self): 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"): @@ -1339,11 +1370,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.""" @@ -1514,7 +1546,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) @@ -1532,6 +1564,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" From 7ae9e0cd4210720347c841d199c4f98f7745358e Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 07:15:57 +0200 Subject: [PATCH 37/48] picks keep the labels of the samples they were found at MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A coordinate attached to the picked dimension — a tag, a code, anything naming the samples rather than the lanes — was indexed with the absolute sample number of the run against a single chunk's values, which raised once the run was longer than a chunk. Those coordinates are now accumulated over the run like the dimension coordinate they sit on, so a trigger's onset still names something when the chunk it was found in is gone. `coords=None` no longer raises on a dimension carrying no coordinate: there is nothing to annotate with there. `Annotate` shipped its companion coordinates at the length of the input rather than of the emitted chunk, which the assembly then dropped, and swallowed a stream shorter than one model window in silence where the eager call says so. `MLPicker` gets its old signature back: the new third positional argument is `components`, so the documented `MLPicker(model, dim, device)` was setting the wrong one. --- xdas/atoms/detect.py | 54 ++++++++++++++++++++++++++++++++++--------- xdas/atoms/ml.py | 39 ++++++++++++++++++++++++------- xdas/virtual/tiles.py | 5 ++-- 3 files changed, 77 insertions(+), 21 deletions(-) diff --git a/xdas/atoms/detect.py b/xdas/atoms/detect.py index 9f5d859d..99dfc2d9 100644 --- a/xdas/atoms/detect.py +++ b/xdas/atoms/detect.py @@ -210,7 +210,7 @@ def initialize(self, cft, **flags): - "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 + - "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. @@ -245,6 +245,7 @@ def initialize(self, cft, **flags): 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): """ @@ -284,6 +285,7 @@ def call(self, 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() @@ -300,8 +302,8 @@ def merge(self, results): Parameters ---------- results : sequence of DataFrame - The per-leaf pick tables, in walk order. Leaves that produced no - pick at all contribute nothing rather than an empty table. + The per-leaf pick tables, in walk order. A leaf that produced no + pick contributes an empty table, which concatenates away. Returns ------- @@ -374,17 +376,24 @@ def _thresholds(self, cft): 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.""" + """ + 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: - # The picked dimension is the chunked one: its indices are - # absolute, so they index the coordinate accumulated so far. 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) @@ -405,7 +414,8 @@ def _names(self, cft): ``(time, distance, phase)`` give the same columns. """ if self.coords is None: - return cft.dims + # 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 @@ -426,12 +436,34 @@ def _annotation(self, cft, name): ) 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.coord[indices[axis]].values + picks[name] = self.labels[name][indices[axis]] elif axis is None: picks[name] = np.full(len(values), source.values, dtype=source.dtype) else: @@ -455,15 +487,15 @@ def _call_numeric(self, data): Parameters ---------- - data : DataArray + data : ndarray 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. + indices : tuple of 1d ndarray + One index array per axis, locating each pick. Notes ----- diff --git a/xdas/atoms/ml.py b/xdas/atoms/ml.py index cf841e44..1a7e425a 100644 --- a/xdas/atoms/ml.py +++ b/xdas/atoms/ml.py @@ -1013,9 +1013,18 @@ def flush(self): 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) - dim = self.sample_dim buffer = self.buffer remainder = buffer.sizes[dim] - self.nperseg if remainder > 0: @@ -1170,10 +1179,12 @@ 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) - coords = da.coords.copy() + # 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[dim] = coords[dim][offset : offset + length] coords["phase"] = self.phases shape = tuple(da.sizes[other] for other in self.batch_dims) shape = (*shape, self.classes, length) @@ -1713,17 +1724,29 @@ class MLPicker(Annotate): """ Deprecated alias of :class:`Annotate`, removed in 0.4. - Beyond the name, the output of this atom is now laid out sample-last, - ``(..., "phase", dim)``, rather than leading with the sample dimension. + 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, *args, **kwargs): + 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=2, + stacklevel=3, + ) + super().__init__( + model, + dim=dim, + device=device, + component_strategy=component_strategy, + **kwargs, ) - super().__init__(*args, **kwargs) annotate = atomized(Annotate) diff --git a/xdas/virtual/tiles.py b/xdas/virtual/tiles.py index 3f23316b..ed74dbe6 100644 --- a/xdas/virtual/tiles.py +++ b/xdas/virtual/tiles.py @@ -669,8 +669,9 @@ 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. Rewriting it relocates the whole archive. + #: 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): From 4fcd3438d00362fe6a0c8cdaffdeefbb654d498c Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 07:16:52 +0200 Subject: [PATCH 38/48] release notes: keep the mislabelled-lanes fix, which 0.2.8 could hit --- docs/release-notes.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/release-notes.md b/docs/release-notes.md index c49e61b4..6a0dbcec 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -44,6 +44,7 @@ ### Bug Fixes - 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). From 7c75bcfd61b78e3820509aaaeaab00307b3d5fb6 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 07:53:47 +0200 Subject: [PATCH 39/48] release notes: lead the atoms story with the two faces of an operation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The point is not that a few routines gained physical parameters: it is that the norm is now one operation with two faces — a user-friendly functional form, and the atom behind it that composes and streams — and that the default vocabulary has moved off SciPy's parameters onto the quantities of the measurement. Those two lead, with the whole functional roster named. Which demotes the exact machine-parameter atoms to what they are: an expert layer one should not have to meet, said at the end of the section rather than as a headline about polyphase throughput. And the leftover "New Features" heading, now holding nothing but the obspy engine, stack, trim_overlaps and select, is named for what it holds. --- docs/release-notes.md | 8 ++-- docs/user-guide/pipeline/atoms.md | 64 +++++++++++++++++++++++-------- 2 files changed, 52 insertions(+), 20 deletions(-) diff --git a/docs/release-notes.md b/docs/release-notes.md index 6a0dbcec..27008484 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -9,14 +9,14 @@ - **`xdas.sortby`** sorts a virtual data array along a dimension by coordinate value without reading any data (@atrabattoni). ### Atoms -- **Composition and function forms.** Atoms compose into pipelines with `>>`, ordinary numpy expressions trace under the `...` seed (`20 * np.log10(np.abs(atom))` extends the pipeline instead of computing), and every task atom has a top-level function form: `xdas.decimate(da, 50.0)` applies eagerly, `xdas.decimate(..., 50.0)` returns the atom (@atrabattoni). -- **Task atoms speak physical units.** The new `xdas.atoms.tasks` vocabulary — `Filter`, `Decimate`, `Resample`, `Integrate`, `Differentiate`, `STFT`, `detrend`, `taper`, `hilbert` and friends — takes corner frequencies and target rates in Hz and window lengths in seconds, so every parameter keeps its meaning when the sampling rate changes (@atrabattoni). -- **Polyphase resampling.** The machine-parameter atoms (`LFilter`, `SOSFilter`, `DownSample`, …) move to the expert layer `xdas.atoms.kernel`, joined by a `Polyphase` kernel that fuses upsampling, FIR filtering and downsampling into a single pass — 2.6–8.7× faster on typical resamplings, without promoting float32 data (@atrabattoni). +- **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 (8 GiB by default) 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). -### New Features +### 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.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). diff --git a/docs/user-guide/pipeline/atoms.md b/docs/user-guide/pipeline/atoms.md index b751bd02..0a9f9f9d 100644 --- a/docs/user-guide/pipeline/atoms.md +++ b/docs/user-guide/pipeline/atoms.md @@ -14,24 +14,31 @@ os.chdir("../../_data") # Composing a processing sequence -*Xdas* ships a processing vocabulary — filtering, resampling, integration, -spectra, machine-learning picking — as *atoms*: elementary operations that -compose into a pipeline. A pipeline built this way runs unchanged on an array -in memory and, chunk by chunk, on an archive that does not fit in one (see -[](processing.md)), which is what makes it worth defining one rather than -calling functions in a row. +*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)). + +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`. ## Applying and composing -Every atom has a function form at the top level of `xdas`. Called on data, it -applies: +Called on data, the function applies: ```{code-cell} import numpy as np import xdas as xd da = xd.synthetics.wavelet_wavefronts() -xd.filter(da, (5.0, None), dim="time") +xd.resample(da, 25.0, dim="time") ``` Called on `...` — the placeholder standing for the data to come — the same @@ -41,16 +48,15 @@ function returns the atom instead, and atoms compose with `>>`: pipeline = ( xd.taper(..., dim="time") >> xd.filter(..., (5.0, None), dim="time") - >> xd.decimate(..., 25.0, dim="time") + >> xd.resample(..., 25.0, dim="time") ) pipeline ``` -The parameters are physical: corner frequencies in hertz, target rates in -hertz, window lengths in seconds. They keep their meaning whatever the sampling -rate of the data the pipeline is later given. - -Calling the pipeline applies it: +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: ```{code-cell} result = pipeline(da) @@ -64,10 +70,36 @@ Ordinary NumPy expressions compose too. Under the `...` seed they are *traced* mathematics: ```{code-cell} -energy = 20 * np.log10(np.abs(xd.decimate(..., 25.0, dim="time"))) +energy = 20 * np.log10(np.abs(xd.resample(..., 25.0, dim="time"))) energy ``` +## 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 From d7c97f874278033f766e59c22963e672b0a599a5 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 08:03:46 +0200 Subject: [PATCH 40/48] the memory limit is a share of the machine, not a fixed eight gigabytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A guard against footguns has to sit far enough above what one legitimately loads at once that it only fires on a mistake, and no fixed number does that across a laptop and a node with a terabyte: 8 GiB is a quarter of the one and a hundredth of a percent of the other, refusing ordinary work on the big machine while still being generous on the small. The default is now a quarter of what the *process* can use — the smaller of the machine's physical memory and any cgroup limit, so a container or a batch allocation is respected rather than the hardware behind it — and the error says what the limit is, since it is no longer a number one can recite. --- docs/release-notes.md | 2 +- docs/user-guide/pipeline/processing.md | 12 +++++- tests/test_config.py | 53 ++++++++++++++++++++++++++ xdas/atoms/core.py | 9 +++-- xdas/config.py | 50 +++++++++++++++++++++++- 5 files changed, 118 insertions(+), 8 deletions(-) create mode 100644 tests/test_config.py diff --git a/docs/release-notes.md b/docs/release-notes.md index 27008484..a05e4915 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -12,7 +12,7 @@ - **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 (8 GiB by default) raises rather than filling the machine (@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). diff --git a/docs/user-guide/pipeline/processing.md b/docs/user-guide/pipeline/processing.md index c8b5301a..b5dc7e60 100644 --- a/docs/user-guide/pipeline/processing.md +++ b/docs/user-guide/pipeline/processing.md @@ -113,8 +113,16 @@ and on the destination it is given: | 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 (8 GiB by default) it -raises rather than filling the machine. +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: 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/xdas/atoms/core.py b/xdas/atoms/core.py index 125d98d0..252127eb 100644 --- a/xdas/atoms/core.py +++ b/xdas/atoms/core.py @@ -415,12 +415,13 @@ def __call__(self, x, **flags): 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, above the 'memory_limit' " - "configuration entry) in memory: stream it chunk by chunk " - "with `.process(da, out=...)` instead, or raise the limit " - "with `xdas.config.set('memory_limit', ...)`" + 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): diff --git a/xdas/config.py b/xdas/config.py index 707009d6..8b5c35d7 100644 --- a/xdas/config.py +++ b/xdas/config.py @@ -7,13 +7,61 @@ 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(), - "memory_limit": 8 * 2**30, + # 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()), } From 424f7aefbc7645c4577faef5de11738227d8089f Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 14:02:46 +0200 Subject: [PATCH 41/48] a subscriber knows when the publisher has registered it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ZeroMQ drops what a publisher sends to a peer whose subscription has not reached it yet, and being connected is not being subscribed. There is no way to ask from the receiving end, so the tests slept and hoped — and hung outright whenever the sleep turned out to be too short, which on a loaded machine it eventually is. A publisher now answers each new subscription with a greeting, in passing as it streams. That is what the XPUB welcome message is for, but ZeroMQ withholds it until the application reads the subscription it answers, which nothing here ever did: submitting a packet now does, so the greeting has in fact never been delivered before this. It is the ASN header, which is why a subscriber joining a running interrogator can now be told the shape of the stream — and skip forward to it, rather than choke on the first packet it happens to land on. `wait_until_subscribed` returns on that greeting, and returning is proof that nothing published from then on will be missed. None of this asks anything of a real-time publisher, which streams whether or not anyone listens and is never held up by its audience. Replaying a recording is the one case that needs the other end to wait, since nothing a subscriber does can hold back a stream already under way, and `wait_for_subscribers` is for that alone. Both subscribers take a timeout, so a stream gone quiet raises where it used to hang. The ZMQ tests keep no sleep, cover joining a live flux, and no longer end in a doctest whose publisher thread raced the reader for a rebound global — that one hung the suite for the fifty minutes it took to notice. --- docs/api/processing.md | 10 ++ docs/release-notes.md | 1 + docs/user-guide/pipeline/streaming.md | 54 +++++-- tests/io/test_asn.py | 147 ++++++++++++----- tests/test_process.py | 4 +- tests/test_processing.py | 56 ++++++- xdas/io/asn.py | 89 +++++++++-- xdas/processing/core.py | 217 +++++++++++++++++++++++--- 8 files changed, 496 insertions(+), 82 deletions(-) diff --git a/docs/api/processing.md b/docs/api/processing.md index c6c701af..1c21cdbb 100644 --- a/docs/api/processing.md +++ b/docs/api/processing.md @@ -70,4 +70,14 @@ ZMQPublisher.submit ZMQPublisher.write ZMQPublisher.result + ZMQPublisher.wait_for_subscribers +``` + +### ZMQSubscriber + +```{eval-rst} +.. autosummary:: + :toctree: ../_autosummary + + ZMQSubscriber.wait_until_subscribed ``` \ No newline at end of file diff --git a/docs/release-notes.md b/docs/release-notes.md index a05e4915..c439b16c 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -25,6 +25,7 @@ ### Improvements - **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. Ray is an optional dependency (`pip install xdas[ray]`); `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). diff --git a/docs/user-guide/pipeline/streaming.md b/docs/user-guide/pipeline/streaming.md index 764c8e0d..c779fd39 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,50 @@ 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) ``` ## Using encoding @@ -92,6 +98,34 @@ encoding = {"chunks": (10, 10), **hdf5plugin.Zfp(accuracy=1e-6)} publisher = ZMQPublisher(address, encoding) # Add encoding here, the rest is the same ``` +## 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 +subscriber = ZMQSubscriber("tcp://interrogator:5555") +for packet in subscriber: + ... +``` + +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 to get data in real-time from one instrument of that kind. Note that compression is not available with that protocol yet. ``` diff --git a/tests/io/test_asn.py b/tests/io/test_asn.py index 624c5930..8822ad31 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") @@ -205,33 +208,27 @@ def test_init_conect_set_header(self): address = get_free_local_address() pub = 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.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.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() @@ -242,11 +239,9 @@ def test_send_chunks(self): pub = 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() @@ -258,15 +253,12 @@ def test_several_subscribers(self): pub = 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() @@ -283,15 +275,12 @@ def test_change_header(self): pub = 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 +293,40 @@ 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): + def get_socket(self, pub, nsubscribers=1): + """Subscribe to *pub* and hand back the socket, once *pub* knows of it.""" socket = zmq.Context().socket(zmq.SUB) - socket.connect(address) + 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 = 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 = 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: def test_one_chunk(self): address = get_free_local_address() pub = 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 +339,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 @@ -343,7 +353,7 @@ def test_several_chunks(self): pub = 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 @@ -363,11 +373,11 @@ def test_several_subscribers(self): 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) @@ -382,7 +392,7 @@ def test_change_header(self): 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) @@ -392,7 +402,7 @@ def test_roiDec(self): pub = 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' @@ -453,13 +463,76 @@ def test_iter(self): pub = 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 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 = 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 = 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() + ZMQPublisher(address) # binds, but never publishes anything + with pytest.raises(TimeoutError, match="no message received"): + ZMQSubscriber(address, timeout=0.1) diff --git a/tests/test_process.py b/tests/test_process.py index 9fde6e75..91e6ea2b 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -9,7 +9,6 @@ import os import threading -import time import numpy as np import pandas as pd @@ -529,9 +528,8 @@ def test_publish_process_subscribe(self, da): source = xp.get_source(address) def publish(): - time.sleep(0.1) + publisher.wait_for_subscribers(timeout=60.0) for packet in packets: - time.sleep(0.001) publisher.submit(packet) thread = threading.Thread(target=publish) diff --git a/tests/test_processing.py b/tests/test_processing.py index 403c9dd3..79c1a441 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -16,6 +16,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): @@ -383,13 +386,16 @@ def _publish_and_subscribe(self, packets, address, encoding=None): publisher = 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() + subscriber = xp.ZMQSubscriber(address, timeout=TIMEOUT) - subscriber = xp.ZMQSubscriber(address) result = [] for n, packet in enumerate(subscriber, start=1): result.append(packet) @@ -405,6 +411,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 = 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 = 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()}" + xp.ZMQPublisher(address) # binds, but never publishes + subscriber = 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()}" + xp.ZMQPublisher(address) + subscriber = 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) diff --git a/xdas/io/asn.py b/xdas/io/asn.py index 97175d6c..822969d4 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 from ..virtual import TileArray, VirtualSource from .core import Engine @@ -130,9 +131,18 @@ 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. """ - def __init__(self, address): + def __init__(self, address, timeout=None): """ Initialize a ZMQStream object. @@ -140,10 +150,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,8 +168,8 @@ 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() @@ -169,9 +180,10 @@ def __init__(self, address): """ 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 @@ -191,9 +203,50 @@ def _connect(self, address): socket.setsockopt_string(zmq.SUBSCRIBE, "") 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 +276,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 +289,16 @@ 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. Examples -------- @@ -260,6 +318,7 @@ class ZMQPublisher: def __init__(self, address): self.address = address + self._nsubscribers = 0 self._connect(address) self._header = None @@ -272,7 +331,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 +348,7 @@ def _connect(self, address): socket = context.socket(zmq.XPUB) socket.setsockopt(zmq.XPUB_VERBOSE, True) socket.bind(address) - self.socket = socket + self._socket = socket @staticmethod def _get_header(da): @@ -306,11 +367,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 +394,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/processing/core.py b/xdas/processing/core.py index 8218186b..07a9d6a6 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -3,13 +3,15 @@ Includes :class:`DataArrayLoader`, :class:`DataArrayWriter`, :class:`DataFrameWriter`, :class:`StreamWriter`, :class:`ZMQPublisher`, -:class:`ZMQSubscriber`, :class:`RealTimeLoader`, :func:`watch`, and the -:func:`process` dispatch boundary with its :func:`get_source` / -:func:`get_writer` resolution machinery. +: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 collections import deque from concurrent.futures import CancelledError, ThreadPoolExecutor @@ -51,6 +53,14 @@ AUTO_CHUNK_NBYTES = 256 * 2**20 """Target in-memory chunk size (in bytes) for ``chunks="auto"``.""" +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. +""" + class _RayFuture: """A future handed out by :class:`ProcessPool`, resolved via the object store.""" @@ -1584,7 +1594,101 @@ def result(self): return out -class ZMQPublisher: +class SubscriptionTracker: + """ + 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 — their address as ``address``, + and initialize ``_nsubscribers`` to zero. + """ + + @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. @@ -1595,6 +1699,19 @@ 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. + Examples -------- >>> import xdas as xd @@ -1629,8 +1746,15 @@ class ZMQPublisher: 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): @@ -1643,6 +1767,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): @@ -1662,11 +1791,14 @@ 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. Examples -------- @@ -1680,34 +1812,41 @@ 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() - 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) """ chunk_dim = "time" unbounded = True - def __init__(self, address): + 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) @@ -1717,8 +1856,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): From e429f03a7ad72c9db81267feb2627c62d0dd4de3 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 15:30:35 +0200 Subject: [PATCH 42/48] chunks cross to the workers through shared memory of our own Ray earned its place by getting chunks across a process boundary without pickling them, which is what caps a process pool: a loaded chunk crossing back was serialized, transferred and deserialized, well below memory bandwidth and on the parent's CPU. It brought a scheduler, a raylet, an object store sized against the machine's memory and session state under /tmp for a job whose chunks are bounded in size and whose flow is already back-pressured -- everything that store is built to survive, and none of what it is needed for. An arena of shared memory does the same work in one module. The parent cuts one /dev/shm file into fixed slots and hands them out; a worker writes its chunk once into a slot and the parent maps the same pages, and a chunk on its way to a writer is staged the same way. Reuse is the point: allocating a block per chunk makes the kernel zero and fault every page again, which is why the obvious spelling of this measures slower than the pipe it replaces. What crosses is a ShmRef, the shape/dtype duck type DataArray already accepts from a virtual array, so nothing else in the path changes. Chunks arrive read-only, as they did from the object store. Slots come back when the chunk that owns them is collected, so a streaming consumer turns a handful of them forever. Anything that will not fit -- an oversized chunk, an exhausted arena, a result that is not an array -- takes the ordinary pickle path, which is slower and never wrong. The pages die with the mappings that hold them, so any crash frees them. What could outlive a run is the arena's name, unlinked at shutdown, again by a finalizer at exit, and swept at the start of the next run if its owner is gone. Two things had to be taught. Loky recycles a worker whose resident memory grows 300 MB past its baseline, and shared pages count toward it, so the arena is discounted from that check or workers quit and respawn mid-run. And the whole design rests on unlinking a mapped file and on signal zero as a liveness probe, so it is POSIX-only; elsewhere the pool pickles as it did before, which is what it did everywhere without ray installed. pool="processes" is unchanged as a name, an interface and a contract. What goes is the optional dependency. --- docs/release-notes.md | 2 +- pyproject.toml | 5 +- tests/test_pools.py | 354 ++++++++++++++++++++++++ tests/test_processing.py | 85 +----- xdas/processing/core.py | 166 ++---------- xdas/processing/pools.py | 572 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 966 insertions(+), 218 deletions(-) create mode 100644 tests/test_pools.py create mode 100644 xdas/processing/pools.py diff --git a/docs/release-notes.md b/docs/release-notes.md index c439b16c..83f5d720 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -24,7 +24,7 @@ ### Improvements - **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. Ray is an optional dependency (`pip install xdas[ray]`); `pool="threads"` remains the default (@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). diff --git a/pyproject.toml b/pyproject.toml index 3b1d175b..9cf5df22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,9 +29,6 @@ dependencies = [ "pyzmq", ] -[project.optional-dependencies] -ray = ["ray"] - [dependency-groups] dev = ["ruff", "pytest", "pytest-cov"] docs = [ @@ -43,7 +40,7 @@ docs = [ "sphinx-copybutton", "sphinx", ] -tests = ["dascore", "dask<2025.4.0", "psutil", "ray", "seisbench", "torch"] +tests = ["dascore", "dask<2025.4.0", "psutil", "seisbench", "torch"] # Single source of truth for the version: xdas/__init__.py [tool.setuptools.dynamic] diff --git a/tests/test_pools.py b/tests/test_pools.py new file mode 100644 index 00000000..a0fbc66b --- /dev/null +++ b/tests/test_pools.py @@ -0,0 +1,354 @@ +import os +import shutil +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor + +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, +) + + +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,), " Date: Wed, 12 Aug 2026 16:13:51 +0200 Subject: [PATCH 43/48] a worker does not have to import the test suite to run its tasks The tasks these tests submit are defined beside them, and cloudpickle sends a function defined in an importable module by reference -- so the worker was being asked to import `tests`. It could, but only by accident: `python -m pytest` puts the working directory on the path and loky passes that path on to its workers. Run as `uv run pytest`, as the CI does, nothing puts the project root on the path -- the import mode this suite uses names the module `tests.test_pools` without making it importable -- and the worker failed to unserialize the call. Registering the module for pickle by value sends the tasks themselves instead, which is what already happens for anything defined in a script. Nothing about what is under test changes: the code the pool submits in earnest lives in `xdas`, which a worker can always import. --- tests/test_pools.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_pools.py b/tests/test_pools.py index a0fbc66b..6716ef40 100644 --- a/tests/test_pools.py +++ b/tests/test_pools.py @@ -1,7 +1,9 @@ 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 @@ -27,6 +29,12 @@ 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.""" From db2008133c2bc9553e90a5784f1e3553e7766203 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Wed, 12 Aug 2026 16:14:09 +0200 Subject: [PATCH 44/48] an arena the collector takes is let go of, not just unnamed Shutting a pool down dropped both the arena's name and this process's handle on it, but the finalizer that stands in when nobody shuts it down only dropped the name. The handle stayed in the module's table for good, and with it the pages, so a process that opened pools in a loop and let them fall out of scope accumulated one whole arena per pool -- twelve of them after twelve pools, in the run that turned this up. The mapping is still not closed, only forgotten: chunks the caller is holding are views on it and keep it alive by themselves, which is the same promise shutdown already made. The module also claimed a crash of any kind frees the pages. That is true of anything that runs Python on the way out, an interrupt included, but not of a run killed outright: loky's workers never notice their parent has gone, and they hold what they have mapped until somebody kills them too. It reproduces with a bare loky pool and no shared memory anywhere near it, so the note now says what actually happens rather than what the arena alone would do. --- tests/test_pools.py | 16 ++++++++++++++++ xdas/processing/pools.py | 33 ++++++++++++++++++++++++++------- 2 files changed, 42 insertions(+), 7 deletions(-) diff --git a/tests/test_pools.py b/tests/test_pools.py index 6716ef40..8799840b 100644 --- a/tests/test_pools.py +++ b/tests/test_pools.py @@ -1,3 +1,4 @@ +import gc import os import shutil import sys @@ -85,6 +86,21 @@ def test_close_unlinks_but_keeps_mappings_valid(self): # The pages outlive the name: whoever still holds a chunk can read it. np.testing.assert_array_equal(data, [1.0, 2.0, 3.0, 4.0]) + def test_an_arena_left_to_the_collector_lets_go_of_everything(self): + # A pool that is never shut down still has to release its arena, or a + # process opening pools in a loop would hold every one it ever made. + arena = Arena(2, 1024) + path = arena.path + _, offset, _ = arena.reserve() + data = view(ShmRef(path, offset, (2,), " Date: Wed, 12 Aug 2026 16:30:05 +0200 Subject: [PATCH 45/48] Remove warnings raising. --- tests/test_process.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_process.py b/tests/test_process.py index 91e6ea2b..6906b7b1 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -452,7 +452,6 @@ def test_realtime_continuous_stream_is_silent(self, da): source = self.Source(list(xd.split(da, 4, "time"))) atom = Partial(np.square) with warnings.catch_warnings(): - warnings.simplefilter("error") atom.process(source) def test_chunked_source_announces_its_splits_upfront(self, da, pipeline): @@ -467,7 +466,6 @@ def test_upfront_scan_skips_non_axis_coordinates(self): dense = xd.testing.dummy(shape=(52, 5), ctype="dense") atom = Partial(np.square) with warnings.catch_warnings(): - warnings.simplefilter("error") result = atom.process(dense, chunks={"time": 20}) assert np.allclose(result.values, np.square(dense).values) @@ -481,7 +479,6 @@ def test_realtime_chunks_without_the_dim_are_not_judged(self, da): source = self.Source([left, aside, right]) atom = Partial(np.square) with warnings.catch_warnings(): - warnings.simplefilter("error") atom.process(source) def test_realtime_one_sample_chunk_adopts_the_stream_rate(self): @@ -495,7 +492,6 @@ def test_realtime_one_sample_chunk_adopts_the_stream_rate(self): source = self.Source(chunks) atom = Partial(np.square) with warnings.catch_warnings(): - warnings.simplefilter("error") atom.process(source) def test_watch_is_a_realtime_loader(self, da, tmp_path): From 1604024f5ae1b9e1bdc0808cce85b36906e223bc Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Thu, 13 Aug 2026 10:30:30 +0200 Subject: [PATCH 46/48] a reader that fails to build lets go of its file TdmsReader opens the file in __init__ and closes it in __exit__, but a reader whose __init__ raises is never handed to the "with" that would have closed it. Every auto-detection sniff of a file that is not TDMS went down that path and leaked a descriptor. --- xdas/io/tdms.py | 53 ++++++++++++++++++++++++++++--------------------- 1 file changed, 30 insertions(+), 23 deletions(-) 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 From ea649aab1c5edbdb4559ec5a90878613d073da8a Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Thu, 13 Aug 2026 10:30:58 +0200 Subject: [PATCH 47/48] an endpoint releases the socket it speaks through None of the four ZMQ publishers and subscribers ever closed anything: a socket kept its file descriptor, and its context an I/O thread, for the life of the process, and the two ASN classes dropped their context on the floor at construction. Closing is now part of the interface, shared by all four through ZMQEndpoint: as a context manager where the endpoint has a scope, with close() where it does not, and on garbage collection for one that is merely dropped. A publisher process() opened itself from a "tcp://" spec is closed by process(); one passed in stays the caller's. --- docs/api/processing.md | 2 + docs/release-notes.md | 2 + docs/user-guide/pipeline/streaming.md | 24 +++++- tests/conftest.py | 15 ++++ tests/io/test_asn.py | 56 +++++++++----- tests/test_processing.py | 73 ++++++++++++++---- xdas/io/asn.py | 26 +++++-- xdas/processing/core.py | 104 ++++++++++++++++++++++---- 8 files changed, 245 insertions(+), 57 deletions(-) diff --git a/docs/api/processing.md b/docs/api/processing.md index 1c21cdbb..f0f7b4f8 100644 --- a/docs/api/processing.md +++ b/docs/api/processing.md @@ -71,6 +71,7 @@ ZMQPublisher.write ZMQPublisher.result ZMQPublisher.wait_for_subscribers + ZMQPublisher.close ``` ### ZMQSubscriber @@ -80,4 +81,5 @@ :toctree: ../_autosummary ZMQSubscriber.wait_until_subscribed + ZMQSubscriber.close ``` \ No newline at end of file diff --git a/docs/release-notes.md b/docs/release-notes.md index 83f5d720..5c70ebb7 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -43,6 +43,8 @@ - `DataCollection.query` raises a `KeyError` on an indexer naming no level of the collection, instead of silently returning everything unchanged (@atrabattoni). ### Bug Fixes +- **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). diff --git a/docs/user-guide/pipeline/streaming.md b/docs/user-guide/pipeline/streaming.md index c779fd39..0063551f 100644 --- a/docs/user-guide/pipeline/streaming.md +++ b/docs/user-guide/pipeline/streaming.md @@ -84,6 +84,23 @@ 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 To reduce the volume of the transmitted data, compression is often useful. Xdas enable the use of the ZFP algorithm when storing data but also when streaming it. Encoding is declared the same way. @@ -96,6 +113,7 @@ 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 @@ -106,9 +124,9 @@ 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 -subscriber = ZMQSubscriber("tcp://interrogator:5555") -for packet in subscriber: - ... +with ZMQSubscriber("tcp://interrogator:5555") as subscriber: + for packet in subscriber: + ... ``` The waiting moves to the receiving end, where it costs the stream nothing. A diff --git a/tests/conftest.py b/tests/conftest.py index 607283bd..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,19 @@ 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(): """ diff --git a/tests/io/test_asn.py b/tests/io/test_asn.py index 8822ad31..c468e136 100644 --- a/tests/io/test_asn.py +++ b/tests/io/test_asn.py @@ -44,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) @@ -190,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 @@ -206,13 +216,13 @@ 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) 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) socket = self.get_socket(pub) pub.submit(da_float32) # a packet must be sent once subscriber is connected @@ -220,7 +230,7 @@ def test_send_header(self): def test_send_data(self): address = get_free_local_address() - pub = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) pub.submit(da_float32) socket = self.get_socket(pub) pub.submit(da_float32) # a packet must be sent once subscriber is connected @@ -236,7 +246,7 @@ 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]) socket = self.get_socket(pub) @@ -250,7 +260,7 @@ 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]) socket1 = self.get_socket(pub) @@ -272,7 +282,7 @@ 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]) socket = self.get_socket(pub) @@ -295,7 +305,8 @@ def test_change_header(self): def get_socket(self, pub, nsubscribers=1): """Subscribe to *pub* and hand back the socket, once *pub* knows of it.""" - socket = zmq.Context().socket(zmq.SUB) + 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"") @@ -304,14 +315,14 @@ def get_socket(self, pub, nsubscribers=1): def test_wait_for_subscribers_times_out(self): address = get_free_local_address() - pub = ZMQPublisher(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 = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) socket = self.get_socket(pub) assert pub.nsubscribers == 1 socket.setsockopt(zmq.UNSUBSCRIBE, b"") @@ -320,9 +331,14 @@ def test_unsubscribing_is_accounted_for(self): 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] thread = threading.Thread(target=self.publish, args=(pub, chunks)) thread.start() @@ -350,7 +366,7 @@ 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 = self.get_subscriber(address) @@ -369,7 +385,7 @@ 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() @@ -388,7 +404,7 @@ 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() @@ -399,7 +415,7 @@ def test_change_header(self): 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 = self.get_subscriber(address) @@ -460,7 +476,7 @@ 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 = self.get_subscriber(address) @@ -470,7 +486,7 @@ def test_iter(self): def get_subscriber(self, address): """A subscriber that raises rather than waiting forever on a lost stream.""" - return ZMQSubscriber(address, timeout=TIMEOUT) + return self.opened(ZMQSubscriber(address, timeout=TIMEOUT)) def publish(self, pub, chunks, nsubscribers=1): """ @@ -489,7 +505,7 @@ def test_subscriber_joins_a_real_time_flux(self): # 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 = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) chunks = xd.split(da_float32, 10) stop = threading.Event() @@ -517,7 +533,7 @@ 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 = ZMQPublisher(address) + pub = self.opened(ZMQPublisher(address)) header = json.dumps(ZMQPublisher._get_header(da_float32)).encode("utf-8") def publish(): @@ -533,6 +549,6 @@ def publish(): def test_timeout(self): address = get_free_local_address() - ZMQPublisher(address) # binds, but never publishes anything + 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_processing.py b/tests/test_processing.py index 11d70440..58081735 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1,3 +1,4 @@ +import gc import os import threading from pathlib import Path @@ -317,8 +318,13 @@ 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 @@ -328,14 +334,16 @@ def publish(): for packet in packets: publisher.submit(packet) - threading.Thread(target=publish).start() - subscriber = xp.ZMQSubscriber(address, timeout=TIMEOUT) + thread = threading.Thread(target=publish) + thread.start() + subscriber = self.opened(xp.ZMQSubscriber(address, timeout=TIMEOUT)) 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): @@ -351,7 +359,7 @@ def test_subscriber_joins_a_real_time_flux(self): # 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 = xp.ZMQPublisher(address) + publisher = self.opened(xp.ZMQPublisher(address)) stop = threading.Event() def flux(): @@ -364,7 +372,7 @@ def flux(): thread = threading.Thread(target=flux) thread.start() try: - subscriber = xp.ZMQSubscriber(address, timeout=TIMEOUT) + 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)] @@ -378,8 +386,8 @@ def flux(): def test_subscriber_timeout(self): address = f"tcp://localhost:{xd.io.get_free_port()}" - xp.ZMQPublisher(address) # binds, but never publishes - subscriber = xp.ZMQSubscriber(address, timeout=0.1) + 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) @@ -387,8 +395,8 @@ 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()}" - xp.ZMQPublisher(address) - subscriber = xp.ZMQSubscriber(address, timeout=0.1) + 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() @@ -638,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/xdas/io/asn.py b/xdas/io/asn.py index 822969d4..56a63a4e 100644 --- a/xdas/io/asn.py +++ b/xdas/io/asn.py @@ -15,7 +15,7 @@ from ..coordinates import Coordinate, get_sampling_interval from ..core import DataArray, concat_coords -from ..processing.core import SubscriptionTracker +from ..processing.core import SubscriptionTracker, ZMQEndpoint from ..virtual import TileArray, VirtualSource from .core import Engine @@ -123,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. @@ -140,6 +140,8 @@ class ZMQSubscriber: 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, timeout=None): @@ -171,13 +173,21 @@ def __init__(self, address, timeout=None): ... publisher.wait_for_subscribers() # a replay, so no chunk is lost ... for chunk in chunks: ... 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 @@ -201,6 +211,7 @@ 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): @@ -299,6 +310,8 @@ class ZMQPublisher(SubscriptionTracker): 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 -------- @@ -309,10 +322,10 @@ class ZMQPublisher(SubscriptionTracker): >>> 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) """ @@ -348,6 +361,7 @@ def _connect(self, address): socket = context.socket(zmq.XPUB) socket.setsockopt(zmq.XPUB_VERBOSE, True) socket.bind(address) + self._context = context self._socket = socket @staticmethod diff --git a/xdas/processing/core.py b/xdas/processing/core.py index 82f13f90..fd5f2cac 100644 --- a/xdas/processing/core.py +++ b/xdas/processing/core.py @@ -256,7 +256,9 @@ def _process_source(atom, source, out, chunks, until, path=None): writer = get_writer(out, outputs[0], "first") for chunk in outputs: writer.write(chunk) - return writer.result() + result = writer.result() + _close_if_owned(writer, out) + return result if hasattr(atom, "reset"): atom.reset() chunk_dim = getattr(source, "chunk_dim", "time") @@ -324,7 +326,23 @@ def write(chunk): for chunk_out in atom.flush(): write(chunk_out) monitor.close() - return writer.result() if writer is not None else None + 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): @@ -528,7 +546,11 @@ def result(self): def close(self): """Close the underlying writer and return its result.""" - return None if self.writer is None else self.writer.result() + if self.writer is None: + return None + result = self.writer.result() + _close_if_owned(self.writer, self.out) + return result class _CollectionSink: @@ -1484,7 +1506,45 @@ def result(self): return out -class SubscriptionTracker: +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. @@ -1508,8 +1568,9 @@ class SubscriptionTracker: 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 — their address as ``address``, - and initialize ``_nsubscribers`` to zero. + 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 @@ -1601,6 +1662,8 @@ class ZMQPublisher(SubscriptionTracker): 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 -------- @@ -1611,15 +1674,16 @@ class ZMQPublisher(SubscriptionTracker): >>> 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 @@ -1627,9 +1691,9 @@ class ZMQPublisher(SubscriptionTracker): >>> 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) """ @@ -1673,7 +1737,7 @@ def result(self): return -class ZMQSubscriber: +class ZMQSubscriber(ZMQEndpoint): """ A class for subscribing to DataArray chunks over ZeroMQ. @@ -1689,6 +1753,8 @@ class ZMQSubscriber: ------- wait_until_subscribed() Block until the publisher has registered this subscription. + close() + Release the socket and its context. Examples -------- @@ -1716,7 +1782,8 @@ class ZMQSubscriber: ... 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. The subscriber is an infinite iterator, so we stop it once the whole stream has been received. @@ -1728,6 +1795,13 @@ class ZMQSubscriber: ... if len(received) == len(packets): ... break >>> 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() """ chunk_dim = "time" From 1a9d064b5b39f39287ab8401553e01252988b735 Mon Sep 17 00:00:00 2001 From: Alister Trabattoni Date: Thu, 13 Aug 2026 10:31:11 +0200 Subject: [PATCH 48/48] a test that asserts silence names the silence it asserts simplefilter("error") turned every warning into a failure, including the ones the test has no opinion about, which is why the blocks it guarded were emptied rather than kept. Naming the category instead says what the test means -- no discontinuity was announced -- and leaves unrelated warnings alone, so the four blocks that had been left standing with nothing in them assert something again. Also covers the last four paths the suite never reached: a coordinate riding the picked dimension, a chunked record shorter than one model window, an unchunked source truncated by "until", and a shared table no leaf ever wrote to. --- tests/test_atoms_detect.py | 18 +++++++++++++ tests/test_atoms_ml.py | 9 +++++++ tests/test_atoms_runs.py | 7 ++--- tests/test_datacollection.py | 7 ++--- tests/test_process.py | 52 +++++++++++++++++++++--------------- 5 files changed, 62 insertions(+), 31 deletions(-) diff --git a/tests/test_atoms_detect.py b/tests/test_atoms_detect.py index ccc30757..5ad2464a 100644 --- a/tests/test_atoms_detect.py +++ b/tests/test_atoms_detect.py @@ -131,6 +131,24 @@ def test_chunked_annotation_matches_monolithic(self): 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"): diff --git a/tests/test_atoms_ml.py b/tests/test_atoms_ml.py index 725cd219..bcb21178 100644 --- a/tests/test_atoms_ml.py +++ b/tests/test_atoms_ml.py @@ -624,6 +624,15 @@ def test_a_chunk_completing_no_window_holds_everything_back(self, indices): 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() == [] diff --git a/tests/test_atoms_runs.py b/tests/test_atoms_runs.py index d8bfd7e2..aa67ac88 100644 --- a/tests/test_atoms_runs.py +++ b/tests/test_atoms_runs.py @@ -10,8 +10,6 @@ elements carry state across, discontinuous ones reset, and tails are flushed. """ -import warnings - import numpy as np import pytest @@ -330,10 +328,9 @@ def test_singular_wording(self, 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): - with warnings.catch_warnings(): - warnings.simplefilter("error") - xd.filter(da, (1.0, 10.0)) + 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 diff --git a/tests/test_datacollection.py b/tests/test_datacollection.py index a9c47ebe..7e09d287 100644 --- a/tests/test_datacollection.py +++ b/tests/test_datacollection.py @@ -1,5 +1,3 @@ -import warnings - import h5py import numpy as np import pandas as pd @@ -664,6 +662,7 @@ def test_an_empty_collection_gives_an_empty_table(self): 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 @@ -675,9 +674,7 @@ def test_an_agreeing_scalar_coordinate_dedupes_silently(self): }, "station", ) - with warnings.catch_warnings(): - warnings.simplefilter("error") - result = xd.trigger(dc, thresh=self.thresh) + result = xd.trigger(dc, thresh=self.thresh) assert list(result.columns) == ["station", "phase", "time", "value"] assert list(result["station"]) == ["DBNFM", "DBNFM", "LBFI", "LBFI"] diff --git a/tests/test_process.py b/tests/test_process.py index 6906b7b1..6900e999 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -112,9 +112,9 @@ def test_invalid_source_raises(self): with pytest.raises(TypeError, match="source"): xp.get_source(42) - def test_tcp_scheme(self): + def test_tcp_scheme(self, opened): address = f"tcp://localhost:{xd.io.get_free_port()}" - source = xp.get_source(address) + source = opened(xp.get_source(address)) assert isinstance(source, xp.ZMQSubscriber) assert source.unbounded assert source.chunk_dim == "time" @@ -407,6 +407,15 @@ def test_until_truncates(self, da, pipeline): 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")) @@ -446,53 +455,45 @@ def test_realtime_seam_warns_per_seam(self, da): 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): - import warnings - source = self.Source(list(xd.split(da, 4, "time"))) atom = Partial(np.square) - with warnings.catch_warnings(): - atom.process(source) + 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. - import warnings - dense = xd.testing.dummy(shape=(52, 5), ctype="dense") atom = Partial(np.square) - with warnings.catch_warnings(): - result = atom.process(dense, chunks={"time": 20}) + 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. - import warnings - 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) - with warnings.catch_warnings(): - atom.process(source) + 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. - import warnings - 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) - with warnings.catch_warnings(): - atom.process(source) + atom.process(source) def test_watch_is_a_realtime_loader(self, da, tmp_path): loader = xd.watch(tmp_path) @@ -516,12 +517,12 @@ def test_watch_source_end_to_end(self, da, pipeline, tmp_path): class TestZMQRoundTrip: - def test_publish_process_subscribe(self, da): + 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 = xp.ZMQPublisher(address) - source = xp.get_source(address) + publisher = opened(xp.ZMQPublisher(address)) + source = opened(xp.get_source(address)) def publish(): publisher.wait_for_subscribers(timeout=60.0) @@ -718,6 +719,15 @@ def test_a_csv_no_leaf_wrote_to_stays_absent(self, tmp_path): 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 ):