1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
5from __future__ import division, absolute_import
7import os
8import logging
9from pyrocko import util, trace
11from . import (mseed, sac, kan, segy, yaff, seisan_waveform, gse1, gcf,
12 datacube, suds, css, gse2, tdms_idas)
13from .io_common import FileLoadError, FileSaveError
15import numpy as num
18__doc__ = util.parse_md(__file__)
20logger = logging.getLogger('pyrocko.io')
23def allowed_formats(operation, use=None, default=None):
24 if operation == 'load':
25 lst = ['detect', 'from_extension', 'mseed', 'sac', 'segy', 'seisan',
26 'seisan.l', 'seisan.b', 'kan', 'yaff', 'gse1', 'gse2', 'gcf',
27 'datacube', 'suds', 'css', 'tdms_idas']
29 elif operation == 'save':
30 lst = ['mseed', 'sac', 'text', 'yaff', 'gse2']
32 if use == 'doc':
33 return ', '.join("``'%s'``" % fmt for fmt in lst)
35 elif use == 'cli_help':
36 return ', '.join(fmt + ['', ' [default]'][fmt == default]
37 for fmt in lst)
39 else:
40 return lst
43def load(filename, format='mseed', getdata=True, substitutions=None):
44 '''
45 Load traces from file.
47 :param format: format of the file (%s)
48 :param getdata: if ``True`` (the default), read data, otherwise only read
49 traces metadata
50 :param substitutions: dict with substitutions to be applied to the traces
51 metadata
53 :returns: list of loaded traces
55 When *format* is set to ``'detect'``, the file type is guessed from the
56 first 512 bytes of the file. Only Mini-SEED, SAC, GSE1, and YAFF format are
57 detected. When *format* is set to ``'from_extension'``, the filename
58 extension is used to decide what format should be assumed. The filename
59 extensions considered are (matching is case insensitive): ``'.sac'``,
60 ``'.kan'``, ``'.sgy'``, ``'.segy'``, ``'.yaff'``, everything else is
61 assumed to be in Mini-SEED format.
63 This function calls :py:func:`iload` and aggregates the loaded traces in a
64 list.
65 '''
67 return list(iload(
68 filename, format=format, getdata=getdata, substitutions=substitutions))
71load.__doc__ %= allowed_formats('load', 'doc')
74def detect_format(filename):
75 try:
76 f = open(filename, 'rb')
77 data = f.read(512)
78 except OSError as e:
79 raise FileLoadError(e)
80 finally:
81 f.close()
83 formats = [
84 (yaff, 'yaff'),
85 (mseed, 'mseed'),
86 (sac, 'sac'),
87 (gse1, 'gse1'),
88 (gse2, 'gse2'),
89 (datacube, 'datacube'),
90 (suds, 'suds'),
91 (tdms_idas, 'tdms_idas')]
93 for mod, fmt in formats:
94 if mod.detect(data):
95 return fmt
97 raise FileLoadError(UnknownFormat(filename))
100def iload(filename, format='mseed', getdata=True, substitutions=None):
101 '''
102 Load traces from file (iterator version).
104 This function works like :py:func:`load`, but returns an iterator which
105 yields the loaded traces.
106 '''
107 load_data = getdata
109 toks = format.split('.', 1)
110 if len(toks) == 2:
111 format, subformat = toks
112 else:
113 subformat = None
115 try:
116 mtime = os.stat(filename)[8]
117 except OSError as e:
118 raise FileLoadError(e)
120 def subs(tr):
121 make_substitutions(tr, substitutions)
122 tr.set_mtime(mtime)
123 return tr
125 extension_to_format = {
126 '.yaff': 'yaff',
127 '.sac': 'sac',
128 '.kan': 'kan',
129 '.segy': 'segy',
130 '.sgy': 'segy',
131 '.gse': 'gse2',
132 '.wfdisc': 'css',
133 '.tdms': 'tdms_idas'
134 }
136 if format == 'from_extension':
137 format = 'mseed'
138 extension = os.path.splitext(filename)[1]
139 format = extension_to_format.get(extension.lower(), 'mseed')
141 if format == 'detect':
142 format = detect_format(filename)
144 format_to_module = {
145 'kan': kan,
146 'segy': segy,
147 'yaff': yaff,
148 'sac': sac,
149 'mseed': mseed,
150 'seisan': seisan_waveform,
151 'gse1': gse1,
152 'gse2': gse2,
153 'gcf': gcf,
154 'datacube': datacube,
155 'suds': suds,
156 'css': css,
157 'tdms_idas': tdms_idas
158 }
160 add_args = {
161 'seisan': {'subformat': subformat},
162 }
164 if format not in format_to_module:
165 raise UnsupportedFormat(format)
167 mod = format_to_module[format]
169 for tr in mod.iload(
170 filename, load_data=load_data, **add_args.get(format, {})):
172 yield subs(tr)
175def save(traces, filename_template, format='mseed', additional={},
176 stations=None, overwrite=True, **kwargs):
177 '''
178 Save traces to file(s).
180 :param traces: a trace or an iterable of traces to store
181 :param filename_template: filename template with placeholders for trace
182 metadata. Uses normal python '%%(placeholder)s' string templates.
183 The following placeholders are considered: ``network``,
184 ``station``, ``location``, ``channel``, ``tmin``
185 (time of first sample), ``tmax`` (time of last sample),
186 ``tmin_ms``, ``tmax_ms``, ``tmin_us``, ``tmax_us``. The versions
187 with '_ms' include milliseconds, the versions with '_us' include
188 microseconds.
189 :param format: %s
190 :param additional: dict with custom template placeholder fillins.
191 :param overwrite': if ``False``, raise an exception if file exists
192 :returns: list of generated filenames
194 .. note::
195 Network, station, location, and channel codes may be silently truncated
196 to file format specific maximum lengthes.
197 '''
199 if isinstance(traces, trace.Trace):
200 traces = [traces]
202 if format == 'from_extension':
203 format = os.path.splitext(filename_template)[1][1:]
205 if format == 'mseed':
206 return mseed.save(traces, filename_template, additional,
207 overwrite=overwrite, **kwargs)
209 elif format == 'gse2':
210 return gse2.save(traces, filename_template, additional,
211 overwrite=overwrite, **kwargs)
213 elif format == 'sac':
214 fns = []
215 for tr in traces:
216 fn = tr.fill_template(filename_template, **additional)
217 if not overwrite and os.path.exists(fn):
218 raise FileSaveError('file exists: %s' % fn)
220 if fn in fns:
221 raise FileSaveError('file just created would be overwritten: '
222 '%s (multiple traces map to same filename)'
223 % fn)
225 util.ensuredirs(fn)
227 f = sac.SacFile(from_trace=tr)
228 if stations:
229 s = stations[tr.network, tr.station, tr.location]
230 f.stla = s.lat
231 f.stlo = s.lon
232 f.stel = s.elevation
233 f.stdp = s.depth
234 f.cmpinc = s.get_channel(tr.channel).dip + 90.
235 f.cmpaz = s.get_channel(tr.channel).azimuth
237 f.write(fn)
238 fns.append(fn)
240 return fns
242 elif format == 'text':
243 fns = []
244 for tr in traces:
245 fn = tr.fill_template(filename_template, **additional)
246 if not overwrite and os.path.exists(fn):
247 raise FileSaveError('file exists: %s' % fn)
249 if fn in fns:
250 raise FileSaveError('file just created would be overwritten: '
251 '%s (multiple traces map to same filename)'
252 % fn)
254 util.ensuredirs(fn)
255 x, y = tr.get_xdata(), tr.get_ydata()
256 num.savetxt(fn, num.transpose((x, y)))
257 fns.append(fn)
258 return fns
260 elif format == 'yaff':
261 return yaff.save(traces, filename_template, additional,
262 overwrite=overwrite, **kwargs)
263 else:
264 raise UnsupportedFormat(format)
267save.__doc__ %= allowed_formats('save', 'doc')
270class UnknownFormat(Exception):
271 def __init__(self, filename):
272 Exception.__init__(self, 'Unknown file format: %s' % filename)
275class UnsupportedFormat(Exception):
276 def __init__(self, format):
277 Exception.__init__(self, 'Unsupported file format: %s' % format)
280def make_substitutions(tr, substitutions):
281 if substitutions:
282 tr.set_codes(**substitutions)