"""
Testing utilities for ObsPlus.
"""
from __future__ import annotations
import os
from collections import Counter
from contextlib import contextmanager
from os.path import join
from pathlib import Path
from typing import ClassVar
import numpy as np
import obspy
import pandas as pd
from obsplus.bank.core import _Bank
from obsplus.constants import NSLC, utc_able_type
from obsplus.utils.bank import _natify_paths
from obsplus.utils.misc import iter_files
from obsplus.utils.time import make_time_chunks, to_utc
[docs]
@contextmanager
def instrument_methods(obj):
"""
Temporarily instrument an object's methods.
This allows the calls to each of the objects methods to be counted.
"""
old_methods = {}
counter = Counter()
setattr(obj, "_counter", counter)
for attr in dir(obj):
# skip dunders
if attr.startswith("__"):
continue
method = getattr(obj, attr, None)
# skip anything that isnt callable
if not callable(method):
continue
# append old method to old_methods and create new method
old_methods[attr] = method
def func(*args, __method_name=attr, **kwargs):
counter.update({__method_name: 1})
return old_methods[__method_name](*args, **kwargs)
setattr(obj, attr, func)
# yield monkey patched object
yield obj
# reset methods
for attr, method in old_methods.items():
setattr(obj, attr, method)
# delete counter
delattr(obj, "_counter")
[docs]
class ArchiveDirectory:
"""class for creating a simple archive"""
def __init__(
self,
path,
starttime=None,
endtime=None,
sampling_rate=1,
duration=3600,
overlap=0,
gaps=None,
seed_ids=("TA.M17A..VHZ", "TA.BOB..VHZ"),
):
self.path = path
if not os.path.exists(path):
os.makedirs(path)
self.starttime = starttime
self.endtime = endtime
self.sampling_rate = sampling_rate
self.duration = duration
self.overlap = overlap
self.seed_ids = seed_ids
self.gaps = gaps
[docs]
def create_stream(
self,
starttime: utc_able_type,
endtime: utc_able_type,
seed_ids: list[str] | None = None,
sampling_rate: float | int | None = None,
) -> obspy.Stream:
"""Create a waveforms from random data"""
t1 = to_utc(starttime)
t2 = to_utc(endtime)
sr = sampling_rate or self.sampling_rate
ar_len = int((t2.timestamp - t1.timestamp) * sr)
st = obspy.Stream()
rand_gen = np.random.default_rng(seed=20)
for seed in seed_ids or self.seed_ids:
net, sta, loc, cha = seed.split(".")
meta = {
"sampling_rate": sr,
"starttime": t1,
"network": net,
"station": sta,
"location": loc,
"channel": cha,
}
data = rand_gen.standard_normal(ar_len)
tr = obspy.Trace(data=data, header=meta)
st.append(tr)
return st
[docs]
def get_gap_stream(self, t1, t2, gaps):
"""Return streams with gaps in it"""
assert len(gaps) == 1
gap = gaps.iloc[0]
ts1, ts2 = t1.timestamp, t2.timestamp
# if gap covers time completely
if gap.start <= ts1 and gap.end >= ts2:
raise ValueError("gapped out")
# if gap is contained by time frame
elif gap.start > ts1 and gap.end < ts2:
st1 = self.create_stream(ts1, gap.start)
st2 = self.create_stream(gap.end, ts2)
return st1 + st2
# if gap only effects endtime
elif ts1 < gap.start < ts2 <= gap.end:
return self.create_stream(ts1, gap.start)
# if gap only effects starttime
elif gap.start <= ts1 < gap.end < ts2:
return self.create_stream(gap.end, ts2)
else: # should not reach here
raise ValueError("something went wrong!") # pragma: no cover
[docs]
def create_directory(self):
"""Create the directory with gaps in it"""
# get a dataframe of the gaps
if self.gaps is not None:
df = pd.DataFrame(self.gaps, columns=["start", "end"])
df["start"] = df["start"].apply(lambda x: x.timestamp)
df["end"] = df["end"].apply(lambda x: x.timestamp)
else:
df = pd.DataFrame(columns=["start", "end"])
assert self.starttime and self.endtime, "needs defined times"
for t1, t2 in make_time_chunks(
self.starttime, self.endtime, self.duration, self.overlap
):
# figure out of this time lies in a gap
gap = df[~((df.start >= t2) | (df.end <= t1))]
if not gap.empty:
try:
st = self.get_gap_stream(t1, t2, gap)
except ValueError:
continue
else:
st = self.create_stream(t1, t2)
finame = str(t1).split(".")[0].replace(":", "-") + ".mseed"
path = join(self.path, finame)
st.write(path, "mseed")
[docs]
def create_directory_from_bulk_args(self, bulk_args):
"""Create a directory from bulk waveform arguments"""
# ensure directory exists
path = Path(self.path)
path.mkdir(exist_ok=True, parents=True)
for net, sta, loc, chan, t1, t2 in bulk_args:
nslc = ".".join([net, sta, loc, chan])
st = self.create_stream(t1, t2, (nslc,))
time_name = str(t1).split(".")[0].replace(":", "-") + ".mseed"
save_name = path / f"{net}_{sta}_{time_name}"
st.write(str(save_name), "mseed")
[docs]
def assert_streams_almost_equal(
st1: obspy.Stream,
st2: obspy.Stream,
basic_stats: bool = True,
atol: float = 1e-05,
rtol: float = 1e-08,
equal_nan: bool = True,
allow_off_by_one: bool = False,
) -> None:
"""
Assert that two streams are almost equal else raise helpful exceptions.
Parameters
----------
st1
The first stream
st2
The second stream
basic_stats
If True, only compare basic stats of the streams including:
network, station, location, channel, starttime, endtime
atol
The absolute tolerance parameter
rtol
The relative tolerance parameter
equal_nan
If True evaluate NaNs as equal
allow_off_by_one
If True, allow the arrays and time alignments to be off by one sample.
Notes
-----
See numpy.allclose for description of atol and rtol paramter.
Raises
------
AssertionError if streams are not about equal.
"""
kwargs = dict(
basic_stats=basic_stats,
atol=atol,
rtol=rtol,
equal_nan=equal_nan,
allow_off_by_one=allow_off_by_one,
)
_StreamEqualTester(**kwargs)(st1, st2)
class _StreamEqualTester:
"""
Simple class for testing if streams are (almost) equal.
This class is not intended to be used directly, instead use
:func:`obsplus.utils.testing.assert_streams_almost_equal`.
"""
skeys: ClassVar = [*list(NSLC), "sampling_rate", "starttime", "endtime"]
def __init__(
self,
basic_stats: bool = True,
atol: float = 1e-05,
rtol: float = 1e-08,
equal_nan: bool = True,
allow_off_by_one: bool = False,
):
self.basic_stats = basic_stats
self.atol = atol
self.rtol = rtol
self.equal_nan = equal_nan
self.allow_off_by_one = allow_off_by_one
def _assert_stats_equal(self, tr1, tr2):
"""Assert that the stats dicts are almost equal."""
skeys, basic_stats = self.skeys, self.basic_stats
sta1 = {x: tr1.stats[x] for x in skeys} if basic_stats else tr1.stats
sta2 = {x: tr2.stats[x] for x in skeys} if basic_stats else tr2.stats
if not sta1 == sta2:
stats_equal = False
# see if the start and end times are within one sample rate
if self.allow_off_by_one:
times = ("starttime", "endtime")
sta1_new = {i: v for i, v in sta1.items() if i not in times}
sta2_new = {i: v for i, v in sta2.items() if i not in times}
if sta1["sampling_rate"] == sta2["sampling_rate"]:
sr = sta1["sampling_rate"]
t1_diff = abs(sta1["starttime"] - sta2["starttime"])
t2_diff = abs(sta1["endtime"] - sta2["starttime"])
if t1_diff < sr and t2_diff < sr and sta1_new == sta2_new:
stats_equal = True
if not stats_equal:
msg = f"Stats are not the same for the traces: \n{sta1}\n{sta2}"
assert 0, msg
def _assert_arrays_almost_equal(self, tr1, tr2):
"""Assert that the data arrays of the traces are almost equal."""
ars = sorted([tr1.data, tr2.data], key=lambda x: len(x))
len1, len2 = len(ars[0]), len(ars[1])
len_diff = len2 - len1
kwargs = dict(atol=self.atol, rtol=self.rtol, equal_nan=self.equal_nan)
# check for off by one error
if len_diff != 0: # if they aren't equal in len
close = False
else:
close = np.allclose(ars[0], ars[1], **kwargs)
if not close and self.allow_off_by_one:
# If the arrays are within 2 samples of each other in length
if abs(len(ars[0]) - len(ars[1])) <= 2:
sub_ar = ars[0][1:-1]
# slide the smaller array over the larger, return true if found
for i in range(len(ars[1]) - len(ars[0]) + 3):
if np.allclose(sub_ar, ars[1][i : i + len(sub_ar)], **kwargs):
close = True
break
assert close, "Data of traces are not nearly equal"
def __call__(self, st1, st2):
"""
Assert that two streams are almost equal else raise AssertionError.
Parameters
----------
st1
The first stream
st2
The second stream
"""
st1.sort(), st2.sort()
if len(st1) != len(st2):
assert 0, "streams do not have the same number of traces"
# iterate each trace and raise if stats and arrays are not almost equal.
for tr1, tr2 in zip(st1, st2):
self._assert_stats_equal(tr1, tr2)
self._assert_arrays_almost_equal(tr1, tr2)
[docs]
def check_index_paths(bank: _Bank):
"""
Make sure the paths in a bank's index can be resolved correctly
Parameters
----------
bank:
A Bank (either WaveBank or EventBank) to verify
"""
bank_path = bank.bank_path
index = bank.read_index()
index_paths = _natify_paths(index["path"])
file_paths = set([bank.bank_path / pth for pth in index_paths])
for file_path in iter_files(str(bank_path), ext="mseed"):
# go up two levels to match path reference
file_path = Path(file_path)
assert file_path in file_paths
![Logo of ObsPlus [ 0.3.3 ]](../../../_static/obsplus_panda.png)