# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------
'''
Reading and writing of carpets in the RUG file format.
RUG is Pyrocko's native file format for :py:class:`~pyrocko.carpet.Carpet`
objects (experimental). It uses the same record-based container as the
:py:mod:`~pyrocko.io.yaff` trace format (see :py:mod:`pyrocko.file`), with
its own record types.
A rug file is a sequence of carpets. Each carpet is stored as a sequence of
records: a ``carpet_header`` record, optional ``carpet_comp_codes`` records
(one per component, omitted if all component codes are empty), optional
``carpet_comp_axis`` records (one per component axis), an optional
``carpet_mask`` record (bit-packed mask of invalid samples, omitted if no
sample is masked) and a ``carpet_data`` record.
Usually, rug files are accessed through Squirrel, which can index and read
them like any other waveform file format. For direct access, use
:py:func:`iload` and :py:func:`save`.
'''
import os
from collections import defaultdict
import numpy as num
from struct import unpack
from os import SEEK_SET
from pyrocko.file import (File, numtype2type, NoDataAvailable,
size_record_header, FileError)
from pyrocko.util import ensuredirs
from pyrocko.carpet import (
Carpet, CarpetOverlapError, check_overlaps, deoverlap, join)
from pyrocko.squirrel.model import CodesNSLCE
from .io_common import FileLoadError, FileSaveError
record_formats = {
'carpet_header': {
'network': 'string',
'station': 'string',
'location': 'string',
'channel': 'string',
'extra': 'string',
'tmin': 'time_string',
'deltat': 'f8',
'nsamples': 'i8',
'ncomponents': 'i8',
},
'carpet_comp_codes': {
'network': 'string',
'station': 'string',
'location': 'string',
'channel': 'string',
'extra': 'string',
},
'carpet_comp_axis': {
'name': 'string',
'data': '@f8',
},
'carpet_mask': {
'mask': '@u1',
},
'carpet_data': {
'data': ('@i2', '@i4', '@i8', '@i2', '@i4', '@i8', '@f4', '@f8'),
},
}
[docs]
class RugLoadError(FileLoadError):
'''
Raised when a rug file is invalid.
'''
pass
[docs]
def pack_mask(mask):
'''
Pack a boolean mask into bytes, 8 samples per byte.
'''
return num.packbits(num.asarray(mask, dtype=bool).ravel())
[docs]
def unpack_mask(packed, shape):
'''
Unpack a mask packed with :py:func:`pack_mask`.
'''
n = int(num.prod(shape))
packed = num.asarray(packed, dtype=num.uint8)
if packed.size != (n + 7) // 8:
raise RugLoadError('Size of mask does not match shape of data.')
return num.unpackbits(packed, count=n).astype(bool).reshape(shape)
EMPTY_CODES = CodesNSLCE()
[docs]
class CarpetFileIO(File):
'''
Low-level access to the carpet records in a rug file.
Most users should use :py:func:`iload` and :py:func:`save` instead.
:param file:
Open file object, in binary mode.
'''
def __init__(self, file):
File.__init__(
self, file,
type_label='YAFF',
version='0000',
record_formats=record_formats)
[docs]
def get_type(self, key, value):
'''
Get the record value type to use for an array.
'''
return numtype2type[value.dtype.type]
[docs]
def codes_from_dict(self, d):
'''
Make codes from the code fields of an unpacked record.
'''
return CodesNSLCE(
d['network'],
d['station'],
d['location'],
d['channel'],
d['extra'])
[docs]
def carpet_from_dict(self, d):
'''
Make a carpet from the collected record contents.
'''
data = d['data']
mask = None
if data is not None:
shape = (d['ncomponents'], d['nsamples'])
data = data.reshape(shape)
if d['mask'] is not None:
mask = unpack_mask(d['mask'], shape)
component_codes = d['component_codes']
if not component_codes:
component_codes = None
if data is None and component_codes is None:
component_codes = [CodesNSLCE()] * d['ncomponents']
return Carpet(
nsamples=d['nsamples'],
codes=self.codes_from_dict(d),
component_codes=component_codes,
component_axes=d['component_axes'],
deltat=d['deltat'],
tmin=d['tmin'],
data=data,
mask=mask,
)
[docs]
def to_component_codes_dict(self, codes):
'''
Get the contents of a ``carpet_comp_codes`` record.
'''
return {
'network': codes.network,
'station': codes.station,
'location': codes.location,
'channel': codes.channel,
'extra': codes.extra,
}
[docs]
def to_component_axis_dict(self, pair):
'''
Get the contents of a ``carpet_comp_axis`` record from a
``(name, values)`` pair.
'''
return {
'name': pair[0],
'data': pair[1],
}
[docs]
def to_data_dict(self, data):
'''
Get the contents of the ``carpet_data`` record.
'''
return {
'data': data,
}
[docs]
def load(self, load_data=True):
'''
Read carpets from the current position of the file.
:param load_data:
If ``False``, only the headers and component codes are read; the
component axes, mask and data are skipped.
:type load_data:
bool
:yields:
``(offset, carpet)`` pairs, ``offset`` being the position [bytes]
of the carpet's first record in the file.
'''
state = 0
d = None
position = self._f.tell()
offset = position
while True:
try:
r = self.next_record()
if state == 0:
if d is not None:
yield offset, self.carpet_from_dict(d)
if not r.type == 'carpet_header':
raise RugLoadError(
'Expected record of type "%s" but got "%s".' % (
'carpet_header', r.type))
d = r.unpack()
d['component_codes'] = []
d['component_axes'] = {}
d['mask'] = None
d['data'] = None
state = 1
offset = position
elif state == 1:
if r.type == 'carpet_comp_codes':
d['component_codes'].append(
self.codes_from_dict(r.unpack()))
elif r.type == 'carpet_comp_axis':
if load_data:
d_ca = r.unpack()
d['component_axes'][d_ca['name']] = d_ca['data']
elif r.type == 'carpet_mask':
if load_data:
d['mask'] = r.unpack()['mask']
elif r.type == 'carpet_data':
if load_data:
d['data'] = r.unpack()['data']
state = 0
else:
raise RugLoadError(
'Expected record of type "%s", "%s", "%s", or '
'"%s" but got "%s".' % (
'carpet_comp_codes',
'carpet_comp_axis',
'carpet_mask',
'carpet_data',
r.type))
position += r.size_record
except NoDataAvailable:
break
if d is not None:
yield offset, self.carpet_from_dict(d)
[docs]
def save(self, carpets):
'''
Write carpets at the current position of the file.
:param carpets:
Carpets to write.
:type carpets:
iterable of :py:class:`~pyrocko.carpet.Carpet`
'''
for carpet in carpets:
r = self.add_record('carpet_header', make_hash=True)
r.pack(self.to_header_dict(carpet))
r.close()
if not all(
codes == EMPTY_CODES
for codes in carpet.component_codes):
for codes in carpet.component_codes:
r = self.add_record('carpet_comp_codes', make_hash=True)
r.pack(self.to_component_codes_dict(codes))
r.close()
for pair in carpet.component_axes.items():
r = self.add_record('carpet_comp_axis', make_hash=True)
r.pack(self.to_component_axis_dict(pair))
r.close()
if carpet.mask is not None:
r = self.add_record('carpet_mask', make_hash=True)
r.pack({'mask': pack_mask(carpet.mask)})
r.close()
r = self.add_record('carpet_data', make_hash=True)
r.pack(self.to_data_dict(carpet.data))
r.close()
[docs]
def iload(filename, load_data=True, offset=0, nsegments=0):
'''
Read carpets from a rug file.
:param filename:
Path to the file.
:type filename:
str
:param load_data:
If ``False``, only the headers and component codes are read; the
component axes, mask and data are skipped. The resulting carpets have
no data.
:type load_data:
bool
:param offset:
Position [bytes] in the file to start reading at, as yielded by an
earlier call.
:type offset:
int
:param nsegments:
Number of carpets to read at most, ``0`` for all.
:type nsegments:
int
:yields:
``(offset, carpet)`` pairs, ``offset`` being the position [bytes] of
the carpet in the file.
:raises:
:py:exc:`~pyrocko.io.io_common.FileLoadError` if the file cannot be
read or is invalid.
'''
f = None
tf = None
try:
f = open(filename, 'rb')
f.seek(offset, SEEK_SET)
tf = CarpetFileIO(f)
isegment = 0
for offset, carpet in tf.load(load_data=load_data):
yield offset, carpet
isegment += 1
if isegment == nsegments:
break
except (OSError, FileError) as e:
raise FileLoadError(e)
finally:
if tf is not None:
tf.close()
if f is not None:
f.close()
[docs]
def save(
carpets,
filename_template,
additional={},
max_open_files=10,
overwrite=True,
append=False,
check_append=False,
check_append_merge=False,
check_append_hook=None):
'''
Save carpets to rug files.
The carpets are distributed to files by filling the placeholders in
``filename_template`` (see :py:meth:`~pyrocko.carpet.Carpet.fill_template`)
and the files are written in turn. Missing directories are created.
If a file exists, it is overwritten or appended to, depending on
``overwrite``, ``append`` and the result of ``check_append_hook``:
========== ======== ===================== ============================
overwrite append check_append_hook file
========== ======== ===================== ============================
False False error
False True ``None`` or ``True`` append
False True ``False`` error
True False truncate
True True ``None`` or ``True`` append
True True ``False`` truncate
========== ======== ===================== ============================
The error is a :py:exc:`~pyrocko.io.io_common.FileSaveError`.
:param carpets:
Carpets to save.
:type carpets:
iterable of :py:class:`~pyrocko.carpet.Carpet`
:param filename_template:
Template for the file names.
:type filename_template:
str
:param additional:
Values for additional placeholders in ``filename_template``.
:type additional:
dict
:param max_open_files:
Unused, for compatibility with the other save functions.
:type max_open_files:
int
:param overwrite:
Whether existing files may be overwritten, see above.
:type overwrite:
bool
:param append:
Whether to append to existing files, see above.
:type append:
bool
:param check_append:
When appending, check that the new carpets do not overlap with the
carpets already in the file (same codes and overlapping time spans).
:type check_append:
bool
:param check_append_merge:
With ``check_append``, merge overlapping carpets with the ones in the
file instead of raising an error. The new carpets take precedence
where they overlap, and the file is rewritten with the merged result.
:type check_append_merge:
bool
:param check_append_hook:
Function called with the path of an existing file, deciding whether
it may be appended to (``True``) or not (``False``), see above.
:type check_append_hook:
callable
:returns:
Paths of the files written.
:rtype:
:py:class:`list` of :py:class:`str`
:raises:
:py:exc:`~pyrocko.io.io_common.FileSaveError` if a file exists and
may not be overwritten or appended to, or if ``check_append`` finds
overlaps and ``check_append_merge`` is not set.
'''
# check_append_hook only makes a decision about whether appending
# is permitted, so it is meaningless without append.
if not append:
check_append = False
check_append_hook = None
paths = set()
by_path = defaultdict(list)
for carpet in carpets:
path = carpet.fill_template(filename_template, **additional)
by_path[path].append(carpet)
for path in sorted(by_path.keys()):
carpets_thisfile = by_path[path]
force_overwrite = False
if path not in paths:
if os.path.exists(path):
# Same overwrite/append/check_append_hook(fn) logic as
# pyrocko.io.mseed.save(): a hook may veto appending to
# a given file; if it does, fall back to truncating it
# (if permitted) or fail outright.
if not overwrite:
if not append or (
append and check_append_hook
and not check_append_hook(path)):
raise FileSaveError('file exists: %s' % path)
else:
if not append or (
append and check_append_hook
and not check_append_hook(path)):
os.unlink(path)
if check_append and os.path.exists(path):
carpets_infile = list(
cp for (_, cp) in iload(path, load_data=False))
try:
check_overlaps(
carpets_thisfile,
carpets_infile,
message='Carpets to be stored would overlap '
'with carpet already stored in file.'
'\n File: %s' % path)
except CarpetOverlapError as e:
if not check_append_merge:
# Wrapped in FileSaveError, like every other
# save failure in this module (and like
# pyrocko.io.mseed.save() does for
# overlapping traces), so callers can catch
# one exception type regardless of format.
raise FileSaveError(str(e)) from e
# Carpets already in the file need their data
# loaded this time, to actually merge them with
# the new ones rather than just checking codes
# and time spans.
carpets_infile = list(
cp for (_, cp)
in iload(path, load_data=True))
# New carpets are listed first so they take
# precedence over already-stored data where
# they overlap.
carpets_thisfile = join(deoverlap(
list(carpets_thisfile) + carpets_infile,
precedence='first'))
# The merged result replaces the file's
# contents entirely, so this has to be a
# truncating write, not an append, even though
# the caller asked to append.
force_overwrite = True
ensuredirs(path)
if force_overwrite:
mode = 'wb'
else:
mode = ['wb', 'ab'][append or path in paths]
with open(path, mode) as f:
paths.add(path)
tf = CarpetFileIO(f)
tf.save(carpets_thisfile)
tf.close()
return list(paths)
[docs]
def detect(first512):
'''
Check whether data is the start of a rug file.
:param first512:
First bytes of a file (at least the size of a record header).
:type first512:
bytes
:returns:
``True`` if the data starts with a carpet header record.
:rtype:
bool
'''
if len(first512) < size_record_header:
return False
label, version, size_record, size_payload, hash, type = unpack(
'>4s4sQQ20s20s', first512[:size_record_header])
if label == b'YAFF' and version == b'0000' \
and type.strip() == b'carpet_header':
return True
return False
