diff --git a/python/mspasspy/util/decorators.py b/python/mspasspy/util/decorators.py index d9806300a9..74c9497e94 100644 --- a/python/mspasspy/util/decorators.py +++ b/python/mspasspy/util/decorators.py @@ -1,3 +1,5 @@ +import copy + from decorator import decorator from mspasspy.util.converter import Stream2Seismogram, Trace2TimeSeries @@ -725,30 +727,90 @@ def is_input_dead(*args, **kwargs): return False -def timeseries_copy_helper(ts1, ts2): - """ """ - # keys in this list are ignored - for now only one but made a list - # to simplify additions later - metadata_to_ignore = ["CONVERTER_ENSEMBLE_KEYS"] - ts1.npts = ts2.npts - ts1.dt = ts2.dt - ts1.tref = ts2.tref - ts1.live = ts2.live - ts1.t0 = ts2.t0 - for k in ts2.keys(): # other metadata copy - if k not in metadata_to_ignore: - # because this is used only in decorator we can assure - # overrides are consistent - ts1[k] = ts2[k] - # We need to clear this/these too as if they were already there the - # above logic preserves them - for k in metadata_to_ignore: - if ts1.is_defined(k): - ts1.erase(k) - ts1.data = ts2.data - - -def timeseries_ensemble_copy_helper(es1, es2): +_CONVERTER_METADATA_KEYS = { + "CONVERTER_ENSEMBLE_KEYS", + "_mspass_time_reference", + "_mspass_t0_shift", + "_mspass_relative_t0", + "_mspass_live", + "_mspass_member_index", + "_mspass_component_index", +} + + +def _capture_time_state(data): + """Capture state that ObsPy cannot represent without loss.""" + shift_key = "starttime_shift" + return { + "tref": data.tref, + "t0": data.t0, + "shifted": data.shifted(), + "t0_shift": data.get_t0shift(), + "shift_key_defined": data.is_defined(shift_key), + "shift_key_value": data[shift_key] if data.is_defined(shift_key) else None, + "live": data.live, + "relative_t0_anchor": False, + } + + +def _post_time_state(trace, state, member_index=None, component_index=None): + """Post MsPASS-only state and ordering information to an ObsPy Trace.""" + trace.stats["_mspass_time_reference"] = ( + "Relative" if state["tref"].name == "Relative" else "UTC" + ) + trace.stats["_mspass_t0_shift"] = state["t0_shift"] + trace.stats["_mspass_live"] = state["live"] + if state["tref"].name == "Relative": + trace.stats["_mspass_relative_t0"] = state["t0"] + if member_index is not None: + trace.stats["_mspass_member_index"] = member_index + if component_index is not None: + trace.stats["_mspass_component_index"] = component_index + + +def _restore_time_state(destination, state): + """Restore state that was saved before conversion to ObsPy.""" + destination.tref = state["tref"] + if state["shifted"]: + destination.force_t0_shift(state["t0_shift"]) + shift_key = "starttime_shift" + if state["shift_key_defined"]: + destination[shift_key] = state["shift_key_value"] + elif destination.is_defined(shift_key): + destination.erase(shift_key) + destination.live = state["live"] + + +def _copy_atomic_state(destination, source, state): + """Copy a processed atomic datum while keeping its core state coherent.""" + for key in source.keys(): + if key not in _CONVERTER_METADATA_KEYS: + destination[key] = source[key] + for key in _CONVERTER_METADATA_KEYS: + if destination.is_defined(key): + destination.erase(key) + + # set_npts allocates and clears the sample buffer, so data must follow it. + destination.npts = source.npts + destination.data = source.data + destination.dt = source.dt + if state["relative_t0_anchor"]: + destination.t0 = state["t0"] + source.t0 + else: + destination.t0 = source.t0 + if destination.is_defined("endtime"): + destination["endtime"] = destination.endtime() + _restore_time_state(destination, state) + + +def timeseries_copy_helper(ts1, ts2, _time_state=None): + """Copy processed samples and sampling state into a TimeSeries.""" + if _time_state is None: + _time_state = _capture_time_state(ts1) + _copy_atomic_state(ts1, ts2, _time_state) + + +def timeseries_ensemble_copy_helper(es1, es2, _time_states=None): """ Copy converted time-series ensemble state back to the original object. @@ -756,13 +818,64 @@ def timeseries_ensemble_copy_helper(es1, es2): :param es2: source ensemble reconstructed after conversion. :return: ``None``; ``es1`` is modified in place. """ + if len(es1.member) != len(es2.member): + raise ValueError( + "Processed Stream member count does not match the input " + "TimeSeriesEnsemble" + ) + if _time_states is None: + _time_states = [_capture_time_state(member) for member in es1.member] for i in range(len(es1.member)): - timeseries_copy_helper(es1.member[i], es2.member[i]) - # fixme: not sure in what algorithm the length of es1 and es2 would be different - # also, the following does not address uninitiated history issues in the extended elements - # same problem applies to the seismogram_ensemble_copy_helper - if len(es2.member) > len(es1.member): - es1.member.extend(es2.member[len(es1.member) :]) + timeseries_copy_helper(es1.member[i], es2.member[i], _time_states[i]) + + +def _timeseries_to_trace(data, state, member_index=None): + """Convert through an isolated copy and post MsPASS-only state.""" + source = TimeSeries(data) + source.set_live() + if state["tref"].name == "Relative": + state["relative_t0_anchor"] = True + source.t0 = 0.0 + trace = source.toTrace() + trace.stats = copy.deepcopy(trace.stats) + _post_time_state(trace, state, member_index=member_index) + return trace + + +def _timeseries_ensemble_to_stream(ensemble, states): + """Convert an ensemble without mutating its members during conversion.""" + source = TimeSeriesEnsemble(ensemble) + for index, member in enumerate(source.member): + member.set_live() + if states[index]["tref"].name == "Relative": + states[index]["relative_t0_anchor"] = True + member.t0 = 0.0 + stream = source.toStream() + for index, trace in enumerate(stream): + trace.stats = copy.deepcopy(trace.stats) + _post_time_state(trace, states[index], member_index=index) + return stream + + +def _validate_timeseries_ensemble_stream(stream, member_count): + if len(stream) != member_count: + raise ValueError( + "Processed Stream member count does not match the input " + "TimeSeriesEnsemble" + ) + for index, trace in enumerate(stream): + if trace.stats.get("_mspass_member_index") != index: + raise ValueError( + "Processed Stream member order does not match the input " + "TimeSeriesEnsemble" + ) + + +def _return_ensemble_identity(result, conversions): + for original, converted in conversions: + if result is None or result is converted: + return original + return result @decorator @@ -779,35 +892,48 @@ def timeseries_as_trace(func, *args, **kwargs): :param kwargs: extra kv arguments :return: the output of func """ + if is_input_dead(*args, **kwargs): return converted_args = [] converted_args_ids = [] + time_states = {} converted_kwargs = {} converted_kwargs_keys = [] + kw_time_states = {} for i in range(len(args)): if isinstance(args[i], TimeSeries): - converted_args.append(args[i].toTrace()) + state = _capture_time_state(args[i]) + converted = _timeseries_to_trace(args[i], state) + converted_args.append(converted) converted_args_ids.append(i) + time_states[i] = state else: converted_args.append(args[i]) for k in kwargs: if isinstance(kwargs[k], TimeSeries): - converted_kwargs[k] = kwargs[k].toTrace() + state = _capture_time_state(kwargs[k]) + converted = _timeseries_to_trace(kwargs[k], state) + converted_kwargs[k] = converted converted_kwargs_keys.append(k) + kw_time_states[k] = state else: converted_kwargs[k] = kwargs[k] res = func(*converted_args, **converted_kwargs) + converted_results = {} + for i in converted_args_ids: + converted_results[i] = Trace2TimeSeries(converted_args[i]) + converted_kw_results = {} + for k in converted_kwargs_keys: + converted_kw_results[k] = Trace2TimeSeries(converted_kwargs[k]) for i in converted_args_ids: - ts = Trace2TimeSeries(converted_args[i]) - timeseries_copy_helper(args[i], ts) + timeseries_copy_helper(args[i], converted_results[i], time_states[i]) for k in converted_kwargs_keys: - ts = Trace2TimeSeries(converted_kwargs[k]) - timeseries_copy_helper(kwargs[k], ts) + timeseries_copy_helper(kwargs[k], converted_kw_results[k], kw_time_states[k]) return res -def seismogram_copy_helper(seis1, seis2): +def seismogram_copy_helper(seis1, seis2, _time_state=None): """ Copy converted seismogram state back to the original object. @@ -815,28 +941,12 @@ def seismogram_copy_helper(seis1, seis2): :param seis2: source seismogram reconstructed after conversion. :return: ``None``; ``seis1`` is modified in place. """ - # keys in this list are ignored - for now only one but made a list - # to simplify additions later - metadata_to_ignore = ["CONVERTER_ENSEMBLE_KEYS"] - seis1.npts = seis2.npts - seis1.dt = seis2.dt - seis1.tref = seis2.tref - seis1.live = seis2.live - seis1.t0 = seis2.t0 - for k in seis2.keys(): # other metadata copy - if k not in metadata_to_ignore: - # because this is used only in decorator we can assure - # overrides are consistent - seis1[k] = seis2[k] - # We need to clear this/these too as if they were already there the - # above logic preserves them - for k in metadata_to_ignore: - if seis1.is_defined(k): - seis1.erase(k) - seis1.data = seis2.data - - -def seismogram_ensemble_copy_helper(es1, es2): + if _time_state is None: + _time_state = _capture_time_state(seis1) + _copy_atomic_state(seis1, seis2, _time_state) + + +def seismogram_ensemble_copy_helper(es1, es2, _time_states=None): """ Copy converted seismogram ensemble state back to the original object. @@ -844,10 +954,74 @@ def seismogram_ensemble_copy_helper(es1, es2): :param es2: source ensemble reconstructed after conversion. :return: ``None``; ``es1`` is modified in place. """ + if len(es1.member) != len(es2.member): + raise ValueError( + "Processed Stream member count does not match the input " + "SeismogramEnsemble" + ) + if _time_states is None: + _time_states = [_capture_time_state(member) for member in es1.member] for i in range(len(es1.member)): - seismogram_copy_helper(es1.member[i], es2.member[i]) - if len(es2.member) > len(es1.member): - es1.member.extend(es2.member[len(es1.member) :]) + seismogram_copy_helper(es1.member[i], es2.member[i], _time_states[i]) + + +def _seismogram_to_stream(data, state, member_index=None): + """Convert through an isolated copy and post MsPASS-only state.""" + source = Seismogram(data) + source.set_live() + if state["tref"].name == "Relative": + state["relative_t0_anchor"] = True + source.t0 = 0.0 + stream = source.toStream() + for component_index, trace in enumerate(stream): + trace.stats = copy.deepcopy(trace.stats) + _post_time_state( + trace, + state, + member_index=member_index, + component_index=component_index, + ) + return stream + + +def _seismogram_ensemble_to_stream(ensemble, states): + """Convert an ensemble without mutating its members during conversion.""" + source = SeismogramEnsemble(ensemble) + for index, member in enumerate(source.member): + member.set_live() + if states[index]["tref"].name == "Relative": + states[index]["relative_t0_anchor"] = True + member.t0 = 0.0 + stream = source.toStream() + for trace_index, trace in enumerate(stream): + trace.stats = copy.deepcopy(trace.stats) + member_index, component_index = divmod(trace_index, 3) + _post_time_state( + trace, + states[member_index], + member_index=member_index, + component_index=component_index, + ) + return stream + + +def _validate_seismogram_stream(stream, member_count, ensemble_name): + if len(stream) != 3 * member_count: + raise ValueError( + "Processed Stream member count does not match the input " + ensemble_name + ) + for trace_index, trace in enumerate(stream): + member_index, component_index = divmod(trace_index, 3) + if trace.stats.get("_mspass_component_index") != component_index: + raise ValueError( + "Processed Stream component order does not match the input " + + ensemble_name + ) + if member_count > 1 and trace.stats.get("_mspass_member_index") != member_index: + raise ValueError( + "Processed Stream member order does not match the input " + + ensemble_name + ) @decorator @@ -869,37 +1043,43 @@ def seismogram_as_stream(func, *args, **kwargs): converted_kwargs = {} converted_args_ids = [] converted_kwargs_keys = [] - converted = False + time_states = {} + kw_time_states = {} for i in range(len(args)): if isinstance(args[i], Seismogram): - converted_args.append(args[i].toStream()) + state = _capture_time_state(args[i]) + converted = _seismogram_to_stream(args[i], state) + converted_args.append(converted) converted_args_ids.append(i) - converted = True + time_states[i] = state else: converted_args.append(args[i]) for k in kwargs: if isinstance(kwargs[k], Seismogram): - converted_kwargs[k] = kwargs[k].toStream() + state = _capture_time_state(kwargs[k]) + converted = _seismogram_to_stream(kwargs[k], state) + converted_kwargs[k] = converted converted_kwargs_keys.append(k) - converted = True + kw_time_states[k] = state else: converted_kwargs[k] = kwargs[k] res = func(*converted_args, **converted_kwargs) - if converted: - # todo save relative time attribute - # fixme cardinal here - for i in converted_args_ids: - seis = Stream2Seismogram(converted_args[i], cardinal=True) - args[i].data = seis.data - # metadata copy - for k in seis.keys(): - args[i][k] = seis[k] - for k in converted_kwargs_keys: - seis = Stream2Seismogram(converted_kwargs[k], cardinal=True) - kwargs[k].data = seis.data - # metadata copy - for key in seis.keys(): - kwargs[k][key] = seis[key] + for i in converted_args_ids: + _validate_seismogram_stream(converted_args[i], 1, "Seismogram") + for k in converted_kwargs_keys: + _validate_seismogram_stream(converted_kwargs[k], 1, "Seismogram") + converted_results = { + i: Stream2Seismogram(converted_args[i], cardinal=True) + for i in converted_args_ids + } + converted_kw_results = { + k: Stream2Seismogram(converted_kwargs[k], cardinal=True) + for k in converted_kwargs_keys + } + for i in converted_args_ids: + seismogram_copy_helper(args[i], converted_results[i], time_states[i]) + for k in converted_kwargs_keys: + seismogram_copy_helper(kwargs[k], converted_kw_results[k], kw_time_states[k]) return res @@ -922,30 +1102,47 @@ def timeseries_ensemble_as_stream(func, *args, **kwargs): converted_kwargs = {} converted_args_ids = [] converted_kwargs_keys = [] - converted = False + time_states = {} + kw_time_states = {} + conversions = [] for i in range(len(args)): if isinstance(args[i], TimeSeriesEnsemble): - converted_args.append(args[i].toStream()) + states = [_capture_time_state(member) for member in args[i].member] + converted = _timeseries_ensemble_to_stream(args[i], states) + converted_args.append(converted) converted_args_ids.append(i) - converted = True + time_states[i] = states + conversions.append((args[i], converted)) else: converted_args.append(args[i]) for k in kwargs: if isinstance(kwargs[k], TimeSeriesEnsemble): - converted_kwargs[k] = kwargs[k].toStream() + states = [_capture_time_state(member) for member in kwargs[k].member] + converted = _timeseries_ensemble_to_stream(kwargs[k], states) + converted_kwargs[k] = converted converted_kwargs_keys.append(k) - converted = True + kw_time_states[k] = states + conversions.append((kwargs[k], converted)) else: converted_kwargs[k] = kwargs[k] res = func(*converted_args, **converted_kwargs) - if converted: - for i in converted_args_ids: - tse = converted_args[i].toTimeSeriesEnsemble() - timeseries_ensemble_copy_helper(args[i], tse) - for k in converted_kwargs_keys: - tse = converted_kwargs[k].toTimeSeriesEnsemble() - timeseries_ensemble_copy_helper(kwargs[k], tse) - return res + for i in converted_args_ids: + _validate_timeseries_ensemble_stream(converted_args[i], len(args[i].member)) + for k in converted_kwargs_keys: + _validate_timeseries_ensemble_stream(converted_kwargs[k], len(kwargs[k].member)) + converted_results = { + i: converted_args[i].toTimeSeriesEnsemble() for i in converted_args_ids + } + converted_kw_results = { + k: converted_kwargs[k].toTimeSeriesEnsemble() for k in converted_kwargs_keys + } + for i in converted_args_ids: + timeseries_ensemble_copy_helper(args[i], converted_results[i], time_states[i]) + for k in converted_kwargs_keys: + timeseries_ensemble_copy_helper( + kwargs[k], converted_kw_results[k], kw_time_states[k] + ) + return _return_ensemble_identity(res, conversions) @decorator @@ -967,30 +1164,51 @@ def seismogram_ensemble_as_stream(func, *args, **kwargs): converted_kwargs = {} converted_args_ids = [] converted_kwargs_keys = [] - converted = False + time_states = {} + kw_time_states = {} + conversions = [] for i in range(len(args)): if isinstance(args[i], SeismogramEnsemble): - converted_args.append(args[i].toStream()) + states = [_capture_time_state(member) for member in args[i].member] + converted = _seismogram_ensemble_to_stream(args[i], states) + converted_args.append(converted) converted_args_ids.append(i) - converted = True + time_states[i] = states + conversions.append((args[i], converted)) else: converted_args.append(args[i]) for k in kwargs: if isinstance(kwargs[k], SeismogramEnsemble): - converted_kwargs[k] = kwargs[k].toStream() + states = [_capture_time_state(member) for member in kwargs[k].member] + converted = _seismogram_ensemble_to_stream(kwargs[k], states) + converted_kwargs[k] = converted converted_kwargs_keys.append(k) - converted = True + kw_time_states[k] = states + conversions.append((kwargs[k], converted)) else: converted_kwargs[k] = kwargs[k] res = func(*converted_args, **converted_kwargs) - if converted: - for i in converted_args_ids: - seis_e = converted_args[i].toSeismogramEnsemble() - seismogram_ensemble_copy_helper(args[i], seis_e) - for k in converted_kwargs_keys: - seis_e = converted_kwargs[k].toSeismogramEnsemble() - seismogram_ensemble_copy_helper(kwargs[k], seis_e) - return res + for i in converted_args_ids: + _validate_seismogram_stream( + converted_args[i], len(args[i].member), "SeismogramEnsemble" + ) + for k in converted_kwargs_keys: + _validate_seismogram_stream( + converted_kwargs[k], len(kwargs[k].member), "SeismogramEnsemble" + ) + converted_results = { + i: converted_args[i].toSeismogramEnsemble() for i in converted_args_ids + } + converted_kw_results = { + k: converted_kwargs[k].toSeismogramEnsemble() for k in converted_kwargs_keys + } + for i in converted_args_ids: + seismogram_ensemble_copy_helper(args[i], converted_results[i], time_states[i]) + for k in converted_kwargs_keys: + seismogram_ensemble_copy_helper( + kwargs[k], converted_kw_results[k], kw_time_states[k] + ) + return _return_ensemble_identity(res, conversions) @decorator diff --git a/python/tests/util/test_decorators.py b/python/tests/util/test_decorators.py index 68ad3c49c8..5a2ac3e07e 100644 --- a/python/tests/util/test_decorators.py +++ b/python/tests/util/test_decorators.py @@ -384,18 +384,20 @@ def test_timeseries_ensemble_as_stream(): tse = get_live_timeseries_ensemble(2) assert len(tse.member) == 2 cp = TimeSeriesEnsemble(tse) - dummy_func_timeseries_ensemble_as_stream(tse) - assert len(tse.member) == 5 - np.isclose(cp.member[0].data, tse.member[0].data).all() - np.isclose(cp.member[0].data, tse.member[1].data).all() + with pytest.raises(ValueError, match="member count"): + dummy_func_timeseries_ensemble_as_stream(tse) + assert len(tse.member) == 2 + for original, unchanged in zip(cp.member, tse.member): + assert np.isclose(original.data, unchanged.data).all() tse = get_live_timeseries_ensemble(2) assert len(tse.member) == 2 cp = TimeSeriesEnsemble(tse) - dummy_func_timeseries_ensemble_as_stream_2(data=tse) - assert len(tse.member) == 5 - np.isclose(cp.member[0].data, tse.member[0].data).all() - np.isclose(cp.member[0].data, tse.member[1].data).all() + with pytest.raises(ValueError, match="member count"): + dummy_func_timeseries_ensemble_as_stream_2(data=tse) + assert len(tse.member) == 2 + for original, unchanged in zip(cp.member, tse.member): + assert np.isclose(original.data, unchanged.data).all() @seismogram_ensemble_as_stream @@ -418,26 +420,24 @@ def test_seismogram_ensemble_as_stream(): seis_e = get_live_seismogram_ensemble(2) assert len(seis_e.member) == 2 cp = SeismogramEnsemble(seis_e) - dummy_func_seismogram_ensemble_as_stream(seis_e) - assert len(seis_e.member) == 3 - assert all( - np.isclose(a, b).all() for a, b in zip(cp.member[0].data, seis_e.member[0].data) - ) - assert all( - np.isclose(a, b).all() for a, b in zip(cp.member[1].data, seis_e.member[1].data) - ) + with pytest.raises(ValueError, match="member count"): + dummy_func_seismogram_ensemble_as_stream(seis_e) + assert len(seis_e.member) == 2 + for original, unchanged in zip(cp.member, seis_e.member): + assert all( + np.isclose(a, b).all() for a, b in zip(original.data, unchanged.data) + ) seis_e = get_live_seismogram_ensemble(2) assert len(seis_e.member) == 2 cp = SeismogramEnsemble(seis_e) - dummy_func_seismogram_ensemble_as_stream_2(data=seis_e) - assert len(seis_e.member) == 3 - assert all( - np.isclose(a, b).all() for a, b in zip(cp.member[0].data, seis_e.member[0].data) - ) - assert all( - np.isclose(a, b).all() for a, b in zip(cp.member[1].data, seis_e.member[1].data) - ) + with pytest.raises(ValueError, match="member count"): + dummy_func_seismogram_ensemble_as_stream_2(data=seis_e) + assert len(seis_e.member) == 2 + for original, unchanged in zip(cp.member, seis_e.member): + assert all( + np.isclose(a, b).all() for a, b in zip(original.data, unchanged.data) + ) class dummy_class_method_wrapper: diff --git a/python/tests/util/test_decorators_obspy_state.py b/python/tests/util/test_decorators_obspy_state.py new file mode 100644 index 0000000000..ab81b20718 --- /dev/null +++ b/python/tests/util/test_decorators_obspy_state.py @@ -0,0 +1,360 @@ +import copy + +import numpy as np +import pytest + +from mspasspy.ccore.seismic import ( + Seismogram, + SeismogramEnsemble, + TimeReferenceType, + TimeSeries, + TimeSeriesEnsemble, +) +from mspasspy.util.decorators import ( + seismogram_as_stream, + seismogram_ensemble_as_stream, + timeseries_as_trace, + timeseries_ensemble_as_stream, +) + +NSAMPLES = 128 +DT = 0.05 +UTC_START = 1_700_000_000.0 +T0_SHIFT = UTC_START - 2.5 +INTERPOLATE_TIME_SHIFT = 0.375 +SIDE_CHANNEL_KEYS = { + "_mspass_time_reference", + "_mspass_t0_shift", + "_mspass_relative_t0", + "_mspass_live", + "_mspass_member_index", + "_mspass_component_index", +} + + +def _make_timeseries(relative=False, live=True, phase=0.0): + data = TimeSeries(NSAMPLES) + data.dt = DT + data.t0 = UTC_START + data.tref = TimeReferenceType.UTC + data.set_live() + for index in range(NSAMPLES): + data.data[index] = np.sin(0.2 * index + phase) + 0.01 * index + data["member_label"] = phase + data["nested"] = {"values": [phase]} + data["processing"] = [] + if relative: + data.ator(T0_SHIFT) + data["starttime_shift"] = T0_SHIFT + if not live: + data.kill() + return data + + +def _make_seismogram(relative=False, live=True, phase=0.0): + data = Seismogram(NSAMPLES) + data.dt = DT + data.t0 = UTC_START + data.tref = TimeReferenceType.UTC + data.set_live() + for component in range(3): + for index in range(NSAMPLES): + data.data[component, index] = ( + np.sin(0.2 * index + phase + component) + 0.01 * index + ) + data["member_label"] = phase + data["nested"] = {"values": [phase]} + data["processing"] = [] + if relative: + data.ator(T0_SHIFT) + data["starttime_shift"] = T0_SHIFT + if not live: + data.kill() + return data + + +def _make_timeseries_ensemble(): + ensemble = TimeSeriesEnsemble() + ensemble["ensemble_label"] = "timeseries" + ensemble.member.append(_make_timeseries(phase=0.0)) + ensemble.member.append(_make_timeseries(relative=True, phase=0.5)) + ensemble.member.append(_make_timeseries(relative=True, live=False, phase=1.0)) + ensemble.set_live() + return ensemble + + +def _make_seismogram_ensemble(): + ensemble = SeismogramEnsemble() + ensemble["ensemble_label"] = "seismogram" + ensemble.member.append(_make_seismogram(phase=0.0)) + ensemble.member.append(_make_seismogram(relative=True, phase=0.5)) + ensemble.member.append(_make_seismogram(relative=True, live=False, phase=1.0)) + ensemble.set_live() + return ensemble + + +@timeseries_as_trace +@seismogram_as_stream +@timeseries_ensemble_as_stream +@seismogram_ensemble_as_stream +def _run_obspy_operation(data, operation): + if operation == "detrend": + return data.detrend("linear") + if operation == "filter": + return data.filter("lowpass", freq=2.0) + if operation == "interpolate": + return data.interpolate( + 40.0, method="linear", time_shift=INTERPOLATE_TIME_SHIFT + ) + if operation == "resample": + return data.resample(40.0) + raise ValueError("unknown operation") + + +def _time_state(data): + return { + "tref": data.tref, + "shifted": data.shifted(), + "t0_shift": data.get_t0shift(), + "shift_defined": data.is_defined("starttime_shift"), + "shift_value": ( + data["starttime_shift"] if data.is_defined("starttime_shift") else None + ), + "live": data.live, + "t0": data.t0, + } + + +def _assert_time_state(data, state, t0_offset=0.0): + assert data.tref == state["tref"] + assert data.shifted() == state["shifted"] + assert data.get_t0shift() == pytest.approx(state["t0_shift"]) + assert data.is_defined("starttime_shift") == state["shift_defined"] + if state["shift_defined"]: + assert data["starttime_shift"] == state["shift_value"] + assert data.live == state["live"] + assert data.t0 == pytest.approx(state["t0"] + t0_offset) + + if data.time_is_relative() and data.shifted(): + restored = type(data)(data) + restored.set_live() + restored.rtoa() + assert restored.time_is_UTC() + assert restored.t0 == pytest.approx(data.t0 + state["t0_shift"]) + + +def _assert_core_state_is_coherent(data): + samples = np.array(data.data) + assert data.npts == samples.shape[-1] + assert data["npts"] == data.npts + assert data["delta"] == pytest.approx(data.dt) + assert data["starttime"] == pytest.approx(data.t0) + assert data["endtime"] == pytest.approx(data.endtime()) + if data.is_defined("sampling_rate"): + assert data["sampling_rate"] == pytest.approx(1.0 / data.dt) + assert SIDE_CHANNEL_KEYS.isdisjoint(data.keys()) + + +@pytest.mark.parametrize("factory", [_make_timeseries, _make_seismogram]) +@pytest.mark.parametrize("relative", [False, True]) +@pytest.mark.parametrize("operation", ["detrend", "filter", "interpolate", "resample"]) +def test_atomic_obspy_copyback_preserves_time_and_life(factory, relative, operation): + data = factory(relative=relative) + original_state = _time_state(data) + original_samples = np.array(data.data) + + _run_obspy_operation(data, operation) + + t0_offset = INTERPOLATE_TIME_SHIFT if operation == "interpolate" else 0.0 + _assert_time_state(data, original_state, t0_offset=t0_offset) + _assert_core_state_is_coherent(data) + if operation in ("interpolate", "resample"): + assert data.dt == pytest.approx(0.025) + assert data.npts != NSAMPLES + else: + assert data.dt == pytest.approx(DT) + assert data.npts == NSAMPLES + assert not np.array_equal(np.array(data.data), original_samples) + + +@pytest.mark.parametrize("factory", [_make_timeseries, _make_seismogram]) +def test_relative_t0_is_not_interpreted_as_an_epoch(factory): + data = factory(relative=True) + data.t0 = 1.0e12 # outside the datetime range accepted by UTCDateTime + original_state = _time_state(data) + + _run_obspy_operation(data, "detrend") + + _assert_time_state(data, original_state) + _assert_core_state_is_coherent(data) + + +@pytest.mark.parametrize( + "factory", [_make_timeseries_ensemble, _make_seismogram_ensemble] +) +@pytest.mark.parametrize("operation", ["detrend", "filter", "interpolate", "resample"]) +def test_ensemble_obspy_copyback_preserves_each_member(factory, operation): + ensemble = factory() + original_members = list(ensemble.member) + original_states = [_time_state(member) for member in ensemble.member] + original_labels = [member["member_label"] for member in ensemble.member] + + result = _run_obspy_operation(ensemble, operation) + + assert result is ensemble + assert len(ensemble.member) == len(original_members) + for index, member in enumerate(ensemble.member): + assert member is original_members[index] + t0_offset = INTERPOLATE_TIME_SHIFT if operation == "interpolate" else 0.0 + _assert_time_state(member, original_states[index], t0_offset=t0_offset) + _assert_core_state_is_coherent(member) + assert member["member_label"] == original_labels[index] + if operation in ("interpolate", "resample"): + assert member.dt == pytest.approx(0.025) + assert member.npts != NSAMPLES + + +def _metadata_snapshot(data): + return {key: copy.deepcopy(data[key]) for key in data.keys()} + + +def _atomic_snapshot(data): + return { + "metadata": _metadata_snapshot(data), + "samples": np.array(data.data), + "time": _time_state(data), + "npts": data.npts, + "dt": data.dt, + "elog_size": data.elog.size(), + "history_stages": data.number_of_stages(), + } + + +def _ensemble_snapshot(ensemble): + return { + "metadata": _metadata_snapshot(ensemble), + "live": ensemble.live, + "members": [_atomic_snapshot(member) for member in ensemble.member], + } + + +def _assert_atomic_unchanged(data, snapshot): + assert _metadata_snapshot(data) == snapshot["metadata"] + np.testing.assert_array_equal(np.array(data.data), snapshot["samples"]) + assert _time_state(data) == snapshot["time"] + assert data.npts == snapshot["npts"] + assert data.dt == snapshot["dt"] + assert data.elog.size() == snapshot["elog_size"] + assert data.number_of_stages() == snapshot["history_stages"] + + +def _assert_ensemble_unchanged(ensemble, snapshot): + assert _metadata_snapshot(ensemble) == snapshot["metadata"] + assert ensemble.live == snapshot["live"] + assert len(ensemble.member) == len(snapshot["members"]) + for member, member_snapshot in zip(ensemble.member, snapshot["members"]): + _assert_atomic_unchanged(member, member_snapshot) + + +@pytest.mark.parametrize( + "wrapper,factory,is_ensemble", + [ + (timeseries_as_trace, _make_timeseries, False), + (seismogram_as_stream, _make_seismogram, False), + (timeseries_ensemble_as_stream, _make_timeseries_ensemble, True), + (seismogram_ensemble_as_stream, _make_seismogram_ensemble, True), + ], +) +@pytest.mark.parametrize("as_keyword", [False, True], ids=["positional", "keyword"]) +def test_dead_atomic_and_all_dead_ensemble_short_circuit_without_calling_obspy( + wrapper, factory, is_ensemble, as_keyword +): + if is_ensemble: + data = factory() + for member in data.member: + member.kill() + data.set_live() + snapshot = _ensemble_snapshot(data) + else: + data = factory(live=False) + snapshot = _atomic_snapshot(data) + calls = [] + + @wrapper + def operation(data): + calls.append(data) + raise AssertionError("dead input must not reach the wrapped callable") + + result = operation(data=data) if as_keyword else operation(data) + + assert result is None + assert calls == [] + if is_ensemble: + _assert_ensemble_unchanged(data, snapshot) + else: + _assert_atomic_unchanged(data, snapshot) + + +@timeseries_ensemble_as_stream +def _change_timeseries_stream_size(stream, change): + stream[0].data[0] += 100.0 + stream[0].stats.nested["values"].append("changed") + if change == "shorter": + stream.pop() + elif change == "longer": + stream.append(stream[-1].copy()) + elif change == "reordered": + stream.traces.reverse() + return stream + + +@seismogram_ensemble_as_stream +def _change_seismogram_stream_size(stream, change): + stream[0].data[0] += 100.0 + stream[0].stats.nested["values"].append("changed") + if change == "shorter": + stream.pop() + elif change == "longer": + stream.append(stream[-1].copy()) + elif change == "reordered": + stream.traces[0], stream.traces[1] = stream.traces[1], stream.traces[0] + return stream + + +@timeseries_ensemble_as_stream +def _mismatch_second_stream(first, second): + first[0].data[0] += 100.0 + first[0].stats.nested["values"].append("changed") + second.pop() + return first + + +@pytest.mark.parametrize( + "factory,wrapped", + [ + (_make_timeseries_ensemble, _change_timeseries_stream_size), + (_make_seismogram_ensemble, _change_seismogram_stream_size), + ], +) +@pytest.mark.parametrize("change", ["shorter", "longer", "reordered"]) +def test_ensemble_stream_mismatch_is_atomic(factory, wrapped, change): + ensemble = factory() + snapshot = _ensemble_snapshot(ensemble) + + with pytest.raises(ValueError, match="Processed Stream"): + wrapped(ensemble, change) + + _assert_ensemble_unchanged(ensemble, snapshot) + + +def test_all_streams_are_validated_before_any_copyback(): + first = _make_timeseries_ensemble() + second = _make_timeseries_ensemble() + first_snapshot = _ensemble_snapshot(first) + second_snapshot = _ensemble_snapshot(second) + + with pytest.raises(ValueError, match="member count"): + _mismatch_second_stream(first, second) + + _assert_ensemble_unchanged(first, first_snapshot) + _assert_ensemble_unchanged(second, second_snapshot)