Coverage for /usr/local/lib/python3.7/dist-packages/pyrocko/scenario/collection.py : 27%

# http://pyrocko.org - GPLv3 

# 

# The Pyrocko Developers, 21st Century 

# ---|P------/S----------~Lg---------- 

from __future__ import absolute_import, division, print_function 

import os 

import tarfile 

import errno 

import numpy as num 

import time 

from pyrocko.guts import Object, Timestamp 

from pyrocko import gf, guts, util, pile, gmtpy 

from .scenario import draw_scenario_gmt 

from .error import ScenarioError 

op = os.path 

guts_prefix = 'pf.scenario' 

def mtime(p): 

return os.stat(p).st_mtime 

class ScenarioCollectionItem(Object): 

scenario_id = gf.StringID.T() 

time_created = Timestamp.T() 

def __init__(self, **kwargs): 

Object.__init__(self, **kwargs) 

self._path = None 

self._pile = None 

self._engine = None 

self._scenes = None 

def set_base_path(self, path): 

self._path = path 

def get_base_path(self): 

if self._path is None: 

raise EnvironmentError('Base path not set!') 

return self._path 

def init_modelling(self, engine): 

self._engine = engine 

def get_path(self, *entry): 

return op.join(*((self._path,) + entry)) 

def get_generator(self): 

generator = guts.load(filename=self.get_path('generator.yaml')) 

generator.init_modelling(self._engine) 

return generator 

def get_time_range(self): 

return self.get_generator().get_time_range() 

def have_waveforms(self, tmin, tmax): 

p = self.get_waveform_pile() 

trs_have = p.all( 

tmin=tmin, tmax=tmax, load_data=False, degap=False) 

return any(tr.data_len() > 0 for tr in trs_have) 

def get_waveform_pile(self): 

self.ensure_data() 

if self._pile is None: 

path_waveforms = self.get_path('waveforms') 

util.ensuredir(path_waveforms) 

fns = util.select_files( 

[path_waveforms], show_progress=False) 

self._pile = pile.Pile() 

if fns: 

self._pile.load_files( 

fns, fileformat='mseed', show_progress=False) 

return self._pile 

def get_insar_scenes(self): 

from kite import Scene 

if self._scenes is None: 

self._scenes = [] 

path_insar = self.get_path('insar') 

util.ensuredir(path_insar) 

fns = util.select_files([path_insar], regex='\\.(npz)$', 

show_progress=False) 

for f in fns: 

self._scenes.append(Scene.load(f)) 

return self._scenes 

def get_gnss_campaigns(self): 

return self.get_generator().get_gnss_campaigns() 

def make_map(self, path_pdf): 

draw_scenario_gmt(self.get_generator(), path_pdf) 

def get_map(self, format='pdf'): 

path_pdf = self.get_path('map.pdf') 

if not op.exists(path_pdf): 

self.make_map(path_pdf) 

path = self.get_path('map.%s' % format) 

outdated = op.exists(path) and mtime(path) < mtime(path_pdf) 

if not op.exists(path) or outdated: 

gmtpy.convert_graph(path_pdf, path) 

return path 

def ensure_data(self, tmin=None, tmax=None): 

return self.get_generator().ensure_data( 

self.get_path(), tmin, tmax) 

def get_archive(self): 

self.ensure_data() 

path_tar = self.get_path('archive.tar') 

if not op.exists(path_tar): 

path_base = self.get_path() 

path_waveforms = self.get_path('waveforms') 

self.ensure_data() 

fns = util.select_files( 

[path_waveforms], show_progress=False) 

f = tarfile.TarFile(path_tar, 'w') 

for fn in fns: 

fna = fn[len(path_base)+1:] 

f.add(fn, fna) 

f.close() 

return path_tar 

class ScenarioCollection(object): 

def __init__(self, path, engine): 

self._scenario_suffix = 'scenario' 

self._path = path 

util.ensuredir(self._path) 

self._engine = engine 

self._load_scenarios() 

def _load_scenarios(self): 

scenarios = [] 

base_path = self.get_path() 

for path_entry in os.listdir(base_path): 

scenario_id, suffix = op.splitext(path_entry) 

if suffix == '.' + self._scenario_suffix: 

path = op.join(base_path, path_entry, 'scenario.yaml') 

scenario = guts.load(filename=path) 

assert scenario.scenario_id == scenario_id 

scenario.set_base_path(op.join(base_path, path_entry)) 

scenario.init_modelling(self._engine) 

scenarios.append(scenario) 

self._scenarios = scenarios 

self._scenarios.sort(key=lambda s: s.time_created) 

def get_path(self, scenario_id=None, *entry): 

if scenario_id is not None: 

return op.join(self._path, '%s.%s' % ( 

scenario_id, self._scenario_suffix), *entry) 

else: 

return self._path 

def add_scenario(self, scenario_id, scenario_generator): 

if scenario_generator.seed is None: 

scenario_generator = guts.clone(scenario_generator) 

scenario_generator.seed = num.random.randint(1, 2**32-1) 

path = self.get_path(scenario_id) 

try: 

os.mkdir(path) 

except OSError as e: 

if e.errno == errno.EEXIST: 

raise ScenarioError( 

'Scenario id is already in use: %s' % scenario_id) 

else: 

raise 

scenario = ScenarioCollectionItem( 

scenario_id=scenario_id, 

time_created=time.time()) 

scenario_path = self.get_path(scenario_id, 'scenario.yaml') 

guts.dump(scenario, filename=scenario_path) 

generator_path = self.get_path(scenario_id, 'generator.yaml') 

guts.dump(scenario_generator, filename=generator_path) 

scenario.set_base_path(self.get_path(scenario_id)) 

scenario.init_modelling(self._engine) 

self._scenarios.append(scenario) 

def list_scenarios(self, ilo=None, ihi=None): 

return self._scenarios[ilo:ihi] 

def get_scenario(self, scenario_id): 

for scenario in self._scenarios: 

if scenario.scenario_id == scenario_id: 

return scenario 

raise KeyError(scenario_id)