1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import os
7import logging
8from pyrocko import util, trace
10from . import (mseed, sac, kan, segy, yaff, seisan_waveform, gse1, gcf,
11 datacube, suds, css, gse2, tdms_idas)
12from .io_common import FileLoadError, FileSaveError
14import numpy as num
17__doc__ = util.parse_md(__file__)
19logger = logging.getLogger('pyrocko.io')
22def allowed_formats(operation, use=None, default=None):
23 if operation == 'load':
24 lst = ['detect', 'from_extension', 'mseed', 'sac', 'segy', 'seisan',
25 'seisan.l', 'seisan.b', 'kan', 'yaff', 'gse1', 'gse2', 'gcf',
26 'datacube', 'suds', 'css', 'tdms_idas']
28 elif operation == 'save':
29 lst = ['mseed', 'sac', 'text', 'yaff', 'gse2']
31 if use == 'doc':
32 return ', '.join("``'%s'``" % fmt for fmt in lst)
34 elif use == 'cli_help':
35 return ', '.join(fmt + ['', ' [default]'][fmt == default]
36 for fmt in lst)
38 else:
39 return lst
42def load(filename, format='mseed', getdata=True, substitutions=None):
43 '''
44 Load traces from file.
46 :param format: format of the file (%s)
47 :param getdata: if ``True`` (the default), read data, otherwise only read
48 traces metadata
49 :param substitutions: dict with substitutions to be applied to the traces
50 metadata
52 :returns: list of loaded traces
54 When *format* is set to ``'detect'``, the file type is guessed from the
55 first 512 bytes of the file. Only Mini-SEED, SAC, GSE1, and YAFF format are
56 detected. When *format* is set to ``'from_extension'``, the filename
57 extension is used to decide what format should be assumed. The filename
58 extensions considered are (matching is case insensitive): ``'.sac'``,
59 ``'.kan'``, ``'.sgy'``, ``'.segy'``, ``'.yaff'``, everything else is
60 assumed to be in Mini-SEED format.
62 This function calls :py:func:`iload` and aggregates the loaded traces in a
63 list.
64 '''
66 return list(iload(
67 filename, format=format, getdata=getdata, substitutions=substitutions))
70load.__doc__ %= allowed_formats('load', 'doc')
73def detect_format(filename):
74 try:
75 f = open(filename, 'rb')
76 data = f.read(512)
77 except OSError as e:
78 raise FileLoadError(e)
79 finally:
80 f.close()
82 formats = [
83 (yaff, 'yaff'),
84 (mseed, 'mseed'),
85 (sac, 'sac'),
86 (gse1, 'gse1'),
87 (gse2, 'gse2'),
88 (datacube, 'datacube'),
89 (suds, 'suds'),
90 (tdms_idas, 'tdms_idas')]
92 for mod, fmt in formats:
93 if mod.detect(data):
94 return fmt
96 raise FileLoadError(UnknownFormat(filename))
99def iload(filename, format='mseed', getdata=True, substitutions=None):
100 '''
101 Load traces from file (iterator version).
103 This function works like :py:func:`load`, but returns an iterator which
104 yields the loaded traces.
105 '''
106 load_data = getdata
108 toks = format.split('.', 1)
109 if len(toks) == 2:
110 format, subformat = toks
111 else:
112 subformat = None
114 try:
115 mtime = os.stat(filename)[8]
116 except OSError as e:
117 raise FileLoadError(e)
119 def subs(tr):
120 make_substitutions(tr, substitutions)
121 tr.set_mtime(mtime)
122 return tr
124 extension_to_format = {
125 '.yaff': 'yaff',
126 '.sac': 'sac',
127 '.kan': 'kan',
128 '.segy': 'segy',
129 '.sgy': 'segy',
130 '.gse': 'gse2',
131 '.wfdisc': 'css',
132 '.tdms': 'tdms_idas'
133 }
135 if format == 'from_extension':
136 format = 'mseed'
137 extension = os.path.splitext(filename)[1]
138 format = extension_to_format.get(extension.lower(), 'mseed')
140 if format == 'detect':
141 format = detect_format(filename)
143 format_to_module = {
144 'kan': kan,
145 'segy': segy,
146 'yaff': yaff,
147 'sac': sac,
148 'mseed': mseed,
149 'seisan': seisan_waveform,
150 'gse1': gse1,
151 'gse2': gse2,
152 'gcf': gcf,
153 'datacube': datacube,
154 'suds': suds,
155 'css': css,
156 'tdms_idas': tdms_idas
157 }
159 add_args = {
160 'seisan': {'subformat': subformat},
161 }
163 if format not in format_to_module:
164 raise UnsupportedFormat(format)
166 mod = format_to_module[format]
168 for tr in mod.iload(
169 filename, load_data=load_data, **add_args.get(format, {})):
171 yield subs(tr)
174def save(traces, filename_template, format='mseed', additional={},
175 stations=None, overwrite=True, **kwargs):
176 '''
177 Save traces to file(s).
179 :param traces: a trace or an iterable of traces to store
180 :param filename_template: filename template with placeholders for trace
181 metadata. Uses normal python '%%(placeholder)s' string templates.
182 The following placeholders are considered: ``network``,
183 ``station``, ``location``, ``channel``, ``tmin``
184 (time of first sample), ``tmax`` (time of last sample),
185 ``tmin_ms``, ``tmax_ms``, ``tmin_us``, ``tmax_us``. The versions
186 with '_ms' include milliseconds, the versions with '_us' include
187 microseconds.
188 :param format: %s
189 :param additional: dict with custom template placeholder fillins.
190 :param overwrite': if ``False``, raise an exception if file exists
191 :returns: list of generated filenames
193 .. note::
194 Network, station, location, and channel codes may be silently truncated
195 to file format specific maximum lengthes.
196 '''
198 if isinstance(traces, trace.Trace):
199 traces = [traces]
201 if format == 'from_extension':
202 format = os.path.splitext(filename_template)[1][1:]
204 if format == 'mseed':
205 return mseed.save(traces, filename_template, additional,
206 overwrite=overwrite, **kwargs)
208 elif format == 'gse2':
209 return gse2.save(traces, filename_template, additional,
210 overwrite=overwrite, **kwargs)
212 elif format == 'sac':
213 fns = []
214 for tr in traces:
215 fn = tr.fill_template(filename_template, **additional)
216 if not overwrite and os.path.exists(fn):
217 raise FileSaveError('file exists: %s' % fn)
219 if fn in fns:
220 raise FileSaveError('file just created would be overwritten: '
221 '%s (multiple traces map to same filename)'
222 % fn)
224 util.ensuredirs(fn)
226 f = sac.SacFile(from_trace=tr)
227 if stations:
228 s = stations[tr.network, tr.station, tr.location]
229 f.stla = s.lat
230 f.stlo = s.lon
231 f.stel = s.elevation
232 f.stdp = s.depth
233 f.cmpinc = s.get_channel(tr.channel).dip + 90.
234 f.cmpaz = s.get_channel(tr.channel).azimuth
236 f.write(fn)
237 fns.append(fn)
239 return fns
241 elif format == 'text':
242 fns = []
243 for tr in traces:
244 fn = tr.fill_template(filename_template, **additional)
245 if not overwrite and os.path.exists(fn):
246 raise FileSaveError('file exists: %s' % fn)
248 if fn in fns:
249 raise FileSaveError('file just created would be overwritten: '
250 '%s (multiple traces map to same filename)'
251 % fn)
253 util.ensuredirs(fn)
254 x, y = tr.get_xdata(), tr.get_ydata()
255 num.savetxt(fn, num.transpose((x, y)))
256 fns.append(fn)
257 return fns
259 elif format == 'yaff':
260 return yaff.save(traces, filename_template, additional,
261 overwrite=overwrite, **kwargs)
262 else:
263 raise UnsupportedFormat(format)
266save.__doc__ %= allowed_formats('save', 'doc')
269class UnknownFormat(Exception):
270 def __init__(self, filename):
271 Exception.__init__(self, 'Unknown file format: %s' % filename)
274class UnsupportedFormat(Exception):
275 def __init__(self, format):
276 Exception.__init__(self, 'Unsupported file format: %s' % format)
279def make_substitutions(tr, substitutions):
280 if substitutions:
281 tr.set_codes(**substitutions)