1import math
2import numpy as num
3from pyrocko.guts import Object, String, Unicode, List, Int, SObject, Any
4from pyrocko.guts_array import Array
5from pyrocko import geometry, cake
6from pyrocko import orthodrome as od
7from pyrocko.util import num_full
10guts_prefix = 'pf'
13def nextpow2(i):
14 return 2**int(math.ceil(math.log(i)/math.log(2.)))
17def ncols(arr):
18 return 1 if arr.ndim == 1 else arr.shape[1]
21def nrows(arr):
22 return arr.shape[0]
25def resize_shape(shape, n):
26 return (n, ) if len(shape) == 1 else (n, shape[1])
29class DType(SObject):
30 dummy_for = num.dtype
33class SubHeader(Object):
34 name = String.T()
35 unit = Unicode.T(optional=True)
36 default = Any.T(optional=True)
37 label = Unicode.T(optional=True)
39 def __init__(self, name, unit=None, default=None, label=None, **kwargs):
40 Object.__init__(
41 self, name=name, unit=unit, default=default, label=label, **kwargs)
43 def get_caption(self):
44 s = self.label or self.name
45 if self.unit:
46 s += ' [%s]' % self.unit
48 return s
50 def get_ncols(self):
51 return 1
54class Header(SubHeader):
55 sub_headers = List.T(SubHeader.T())
56 dtype = DType.T(default=num.dtype('float64'), optional=True)
58 def __init__(
59 self, name,
60 unit=None,
61 sub_headers=[],
62 dtype=None,
63 default=None,
64 label=None):
66 sub_headers = [anything_to_sub_header(sh) for sh in sub_headers]
68 kwargs = dict(sub_headers=sub_headers, dtype=dtype)
70 SubHeader.__init__(self, name, unit, default, label, **kwargs)
72 def get_ncols(self):
73 return max(1, len(self.sub_headers))
75 def default_array(self, nrows):
76 val = self.dtype(self.default)
77 if not self.sub_headers:
78 return num_full((nrows,), val, dtype=self.dtype)
79 else:
80 return num_full((nrows, self.get_ncols()), val, dtype=self.dtype)
83def anything_to_header(args):
84 if isinstance(args, Header):
85 return args
86 elif isinstance(args, str):
87 return Header(name=args)
88 elif isinstance(args, tuple):
89 return Header(*args)
90 else:
91 raise ValueError('argument of type Header, str or tuple expected')
94def anything_to_sub_header(args):
95 if isinstance(args, SubHeader):
96 return args
97 elif isinstance(args, str):
98 return SubHeader(name=args)
99 elif isinstance(args, tuple):
100 return SubHeader(*args)
101 else:
102 raise ValueError('argument of type SubHeader, str or tuple expected')
105class Description(Object):
106 name = String.T(optional=True)
107 headers = List.T(Header.T())
108 nrows = Int.T()
109 ncols = Int.T()
111 def __init__(self, table=None, **kwargs):
112 if table:
113 Object.__init__(
114 self,
115 name=table._name,
116 headers=table._headers,
117 nrows=table.get_nrows(),
118 ncols=table.get_ncols())
119 else:
120 Object.__init__(self, **kwargs)
123class NoSuchRecipe(Exception):
124 pass
127class NoSuchCol(Exception):
128 pass
131class Recipe(Object):
133 def __init__(self):
134 self._table = None
135 self._table = Table()
137 self._required_headers = []
138 self._headers = []
139 self._col_update_map = {}
140 self._name_to_header = {}
142 def has_col(self, name):
143 return name in self._name_to_header
145 def get_col_names(self):
146 names = []
147 for h in self._headers:
148 names.append(h.name)
149 for sh in h.sub_headers:
150 names.append(sh.name)
152 return names
154 def get_table(self):
155 return self._table
157 def get_header(self, name):
158 try:
159 return self._name_to_header[name]
160 except KeyError:
161 for h in self._required_headers:
162 if h.name == name:
163 return h
165 raise KeyError(name)
167 def _add_required_cols(self, table):
168 for h in self._headers:
169 if not table.has_col(h.name):
170 table.add_col(h)
172 def _update_col(self, table, name):
173 if not self._table.has_col(name):
174 self._col_update_map[name](table)
176 def _add_rows_handler(self, table, nrows_added):
177 pass
179 def _register_required_col(self, header):
180 self._required_headers.append(header)
182 def _register_computed_col(self, header, updater):
183 self._headers.append(header)
184 self._name_to_header[header.name] = header
185 self._col_update_map[header.name] = updater
186 for sh in header.sub_headers:
187 self._col_update_map[sh.name] = updater
188 self._name_to_header[sh.name] = sh
191class Table(Object):
193 description__ = Description.T()
194 arrays__ = List.T(Array.T(serialize_as='base64+meta'))
195 recipes__ = List.T(Recipe.T())
197 def __init__(
198 self,
199 name=None,
200 nrows_capacity=None,
201 nrows_capacity_min=0,
202 description=None,
203 arrays=None,
204 recipes=[]):
206 Object.__init__(self, init_props=False)
208 self._name = name
209 self._buffers = []
210 self._arrays = []
211 self._headers = []
212 self._cols = {}
213 self._recipes = []
214 self.nrows_capacity_min = nrows_capacity_min
215 self._nrows_capacity = 0
216 if nrows_capacity is not None:
217 self.set_nrows_capacity(max(nrows_capacity, nrows_capacity_min))
219 if description and arrays:
220 self.T.get_property('arrays').validate(
221 arrays, regularize=True, depth=0)
222 self._name = description.name
223 self.add_cols(description.headers, arrays)
224 for recipe in recipes:
225 self.add_recipe(recipe)
227 @property
228 def description(self):
229 return self.get_description()
231 @property
232 def arrays(self):
233 return self._arrays
235 @property
236 def recipes(self):
237 return self._recipes
239 def add_recipe(self, recipe):
240 self._recipes.append(recipe)
241 # recipe._add_required_cols(self)
243 def get_nrows(self):
244 if not self._arrays:
245 return 0
246 else:
247 return nrows(self._arrays[0])
249 def get_nrows_capacity(self):
250 return self._nrows_capacity
252 def set_nrows_capacity(self, nrows_capacity_new):
253 if self.get_nrows_capacity() != nrows_capacity_new:
254 if self.get_nrows() > nrows_capacity_new:
255 raise ValueError('new capacity too small to hold current data')
257 new_buffers = []
258 for buf in self._buffers:
259 shape = resize_shape(buf.shape, nrows_capacity_new)
260 new_buffers.append(num.zeros(shape, dtype=buf.dtype))
262 ncopy = min(self.get_nrows(), nrows_capacity_new)
264 new_arrays = []
265 for arr, buf in zip(self._arrays, new_buffers):
266 buf[:ncopy, ...] = arr[:ncopy, ...]
267 new_arrays.append(buf[:ncopy, ...])
269 self._buffers = new_buffers
270 self._arrays = new_arrays
271 self._nrows_capacity = nrows_capacity_new
273 def get_ncols(self):
274 return len(self._arrays)
276 def add_col(self, header, array=None):
277 header = anything_to_header(header)
279 nrows_current = self.get_nrows()
280 if array is None:
281 array = header.default_array(nrows_current)
283 array = num.asarray(array)
285 assert header.get_ncols() == ncols(array)
286 assert array.ndim in (1, 2)
287 if self._arrays:
288 assert nrows(array) == nrows_current
290 if nrows_current == 0:
291 nrows_current = nrows(array)
292 self.set_nrows_capacity(
293 max(nrows_current, self.nrows_capacity_min))
295 iarr = len(self._arrays)
297 shape = resize_shape(array.shape, self.get_nrows_capacity())
298 if shape != array.shape:
299 buf = num.zeros(shape, dtype=array.dtype)
300 buf[:nrows_current, ...] = array[:, ...]
301 else:
302 buf = array
304 self._buffers.append(buf)
305 self._arrays.append(buf[:nrows_current, ...])
306 self._headers.append(header)
308 self._cols[header.name] = iarr, None
310 for icol, sub_header in enumerate(header.sub_headers):
311 self._cols[sub_header.name] = iarr, icol
313 def add_cols(self, headers, arrays=None):
314 if arrays is None:
315 arrays = [None] * len(headers)
317 for header, array in zip(headers, arrays):
318 self.add_col(header, array)
320 def add_rows(self, arrays):
321 assert self.get_ncols() == len(arrays)
322 arrays = [num.asarray(arr) for arr in arrays]
324 nrows_add = nrows(arrays[0])
325 nrows_current = self.get_nrows()
326 nrows_new = nrows_current + nrows_add
327 if self.get_nrows_capacity() < nrows_new:
328 self.set_nrows_capacity(max(
329 self.nrows_capacity_min, nextpow2(nrows_new)))
331 new_arrays = []
332 for buf, arr in zip(self._buffers, arrays):
333 assert ncols(arr) == ncols(buf)
334 assert nrows(arr) == nrows_add
335 buf[nrows_current:nrows_new, ...] = arr[:, ...]
336 new_arrays.append(buf[:nrows_new, ...])
338 self._arrays = new_arrays
340 for recipe in self._recipes:
341 recipe._add_rows_handler(self, nrows_add)
343 def get_col(self, name, mask=slice(None)):
344 if name in self._cols:
345 if isinstance(mask, str):
346 mask = self.get_col(mask)
348 iarr, icol = self._cols[name]
349 if icol is None:
350 return self._arrays[iarr][mask]
351 else:
352 return self._arrays[iarr][mask, icol]
353 else:
354 recipe = self.get_recipe_for_col(name)
355 recipe._update_col(self, name)
357 return recipe.get_table().get_col(name, mask)
359 def get_header(self, name):
360 if name in self._cols:
361 iarr, icol = self._cols[name]
362 if icol is None:
363 return self._headers[iarr]
364 else:
365 return self._headers[iarr].sub_headers[icol]
366 else:
367 recipe = self.get_recipe_for_col(name)
368 return recipe.get_header(name)
370 def has_col(self, name):
371 return name in self._cols or \
372 any(rec.has_col(name) for rec in self._recipes)
374 def get_col_names(self, sub_headers=True):
375 names = []
376 for h in self._headers:
377 names.append(h.name)
378 if sub_headers:
379 for sh in h.sub_headers:
380 names.append(sh.name)
382 for recipe in self._recipes:
383 names.extend(recipe.get_col_names())
385 return names
387 def get_recipe_for_col(self, name):
388 for recipe in self._recipes:
389 if recipe.has_col(name):
390 return recipe
392 raise NoSuchCol(name)
394 def get_description(self):
395 return Description(self)
397 def get_as_text(self):
398 scols = []
399 formats = {
400 num.dtype('float64'): '%e'}
402 for name in self.get_col_names(sub_headers=False):
403 array = self.get_col(name)
404 header = self.get_header(name)
405 fmt = formats.get(array.dtype, '%s')
406 if array.ndim == 1:
407 scol = [header.get_caption(), '']
408 for val in array:
409 scol.append(fmt % val)
411 scols.append(scol)
412 else:
413 for icol in range(ncols(array)):
414 sub_header = header.sub_headers[icol]
415 scol = [header.get_caption(), sub_header.get_caption()]
416 for val in array[:, icol]:
417 scol.append(fmt % val)
419 scols.append(scol)
421 for scol in scols:
422 width = max(len(s) for s in scol)
423 for i in range(len(scol)):
424 scol[i] = scol[i].rjust(width)
426 return '\n'.join(' '.join(s for s in srow) for srow in zip(*scols))
428 def add_computed_col(self, header, func):
429 header = anything_to_header(header)
430 self.add_recipe(SimpleRecipe(header, func))
432 def get_row(self, irow):
433 if irow == -1:
434 return [arr[-1:] for arr in self.arrays]
435 else:
436 return [arr[irow:irow+1] for arr in self.arrays]
439class SimpleRecipe(Recipe):
441 def __init__(self, header, func):
442 Recipe.__init__(self)
443 self._col_name = header.name
445 def call_func(tab):
446 self._table.add_col(header, func(tab))
448 self._register_computed_col(header, call_func)
450 def _add_rows_handler(self, table, nrows_added):
451 Recipe._add_rows_handler(self, table, nrows_added)
452 if self._table.has_col(self._col_name):
453 self._table.remove_col(self._col_name)
456class LocationRecipe(Recipe):
458 def __init__(self):
459 Recipe.__init__(self)
460 self.c5_header = Header(name='c5', sub_headers=[
461 SubHeader(name='ref_lat', unit='degrees'),
462 SubHeader(name='ref_lon', unit='degrees'),
463 SubHeader(name='north_shift', unit='m'),
464 SubHeader(name='east_shift', unit='m'),
465 SubHeader(name='depth', unit='m')])
467 self._register_required_col(self.c5_header)
469 self._latlon_header = Header(name='latlon', sub_headers=[
470 SubHeader(name='lat', unit='degrees'),
471 SubHeader(name='lon', unit='degrees')])
473 self._register_computed_col(self._latlon_header, self._update_latlon)
475 self._xyz_header = Header(name='xyz', sub_headers=[
476 SubHeader(name='x', unit='m'),
477 SubHeader(name='y', unit='m'),
478 SubHeader(name='z', unit='m')])
480 self._register_computed_col(self._xyz_header, self._update_xyz)
482 def set_depth_offset(self, depth_offset):
483 self.depth_offset = depth_offset
485 def _add_rows_handler(self, table, nrows_added):
486 Recipe._add_rows_handler(self, table, nrows_added)
488 for colname in ['latlon', 'xyz']:
489 if self._table.has_col(colname):
490 self._table.remove_col(colname)
492 def _update_latlon(self, table):
493 lats, lons = od.ne_to_latlon(
494 table.get_col('ref_lat'),
495 table.get_col('ref_lon'),
496 table.get_col('north_shift'),
497 table.get_col('east_shift'))
499 latlons = num.zeros((lats.size, 2))
500 latlons[:, 0] = lats
501 latlons[:, 1] = lons
503 self._table.add_col(self._latlon_header, latlons)
505 def _update_xyz(self, table):
506 self._update_latlon(table)
508 xyzs = geometry.latlondepth2xyz(
509 num.concatenate((
510 table.get_col('lat').reshape(-1, 1),
511 table.get_col('lon').reshape(-1, 1),
512 table.get_col('depth').reshape(-1, 1)),
513 axis=1),
514 planetradius=cake.earthradius)
516 self._table.add_col(self._xyz_header, xyzs)
519class EventRecipe(LocationRecipe):
521 def __init__(self):
522 LocationRecipe.__init__(self)
523 self._register_required_col(Header(name='time', unit='s'))
524 self._register_required_col(Header(name='magnitude'))
526 def iter_events(self, table):
527 from pyrocko import model
528 for vec in zip(table.get_col(x) for x in [
529 'time', 'lat', 'lon',
530 'north_shift', 'east_shift', 'depth',
531 'magnitude']):
533 yield model.Event(
534 time=vec[0],
535 lat=vec[1],
536 lon=vec[2],
537 north_shift=vec[3],
538 east_shift=vec[4],
539 depth=vec[5],
540 magnitude=vec[6])
542 def get_events(self, table):
543 return list(self.iter_events(table))
546class MomentTensorRecipe(Recipe):
548 def __init__(self):
549 Recipe.__init__(self)
550 self.m6_header = Header(name='m6', sub_headers=[
551 SubHeader(name='Mnn', unit='Nm'),
552 SubHeader(name='Mee', unit='Nm'),
553 SubHeader(name='Mdd', unit='Nm'),
554 SubHeader(name='Mne', unit='Nm'),
555 SubHeader(name='Mnd', unit='Nm'),
556 SubHeader(name='Med', unit='Nm')])
558 self._register_required_col(self.m6_header)