def _unpack_if_needed(possible_tar_file):
if os.path.exists(possible_tar_file):
utils.print_ylw("Unpacking %s..." % (possible_tar_file))
tar = tarfile.open(possible_tar_file)
tar.extractall(path=os.path.dirname(possible_tar_file))
os.remove(possible_tar_file)
def submit_mesher(params, run_type):
"""
Submits the jobarray script in the mesh directory.
"""
_change_job_and_par_file(params, run_type)
forward_run_dir = params['forward_run_dir']
os.chdir(os.path.join(forward_run_dir, 'mesh'))
utils.print_ylw("Cleaning old mesh files...")
for file in os.listdir("DATABASES_MPI"):
os.remove(os.path.join("DATABASES_MPI", file))
subprocess.Popen(['sbatch', 'job_mesher_daint.sbatch']).wait()
def clean_event_kernels(params):
"""
Goes through the simulation directories for an iteration, and cleans out
the individual event kernels (make sure you've summed them already!).
"""
forward_run_dir = params['forward_run_dir']
event_list = params['event_list']
for dir in sorted(os.listdir(forward_run_dir)):
if dir in event_list:
utils.print_ylw("Cleaning " + dir + "...")
databases_mpi = os.path.join(forward_run_dir, dir, 'DATABASES_MPI')
for file in os.listdir(databases_mpi):
if file.endswith('kernel.bin'):
os.remove(os.path.join(databases_mpi, file))
utils.print_blu('Done.')
def clean_attenuation_dumps(params):
"""
Goes through the simulation directories for an iteration, and cleans out
the massive adios attenuation snapshot files.
"""
forward_run_dir = params['forward_run_dir']
event_list = params['event_list']
for dir in sorted(os.listdir(forward_run_dir)):
if dir in event_list:
utils.print_ylw("Cleaning " + dir + "...")
databases_mpi = os.path.join(forward_run_dir, dir, 'DATABASES_MPI')
for file in os.listdir(databases_mpi):
if file.startswith('save_frame_at'):
os.remove(os.path.join(databases_mpi, file))
utils.print_blu('Done.')
def clean_failed(params):
"""
Deletes error files after a failed run.
"""
forward_run_dir = params['forward_run_dir']
for dir in os.listdir(forward_run_dir):
utils.print_ylw("Cleaning " + dir + "...")
if os.path.exists(os.path.join(forward_run_dir, dir, 'OUTPUT_FILES')):
output_files = os.path.join(forward_run_dir, dir, 'OUTPUT_FILES')
for file in os.listdir(output_files):
if 'error' in file:
os.remove(os.path.join(output_files, file))
if 'sac' in file:
os.remove(os.path.join(output_files, file))
if os.path.exists(os.path.join(forward_run_dir, dir, 'core')):
os.remove(os.path.join(forward_run_dir, dir, 'core'))
utils.print_blu('Done.')
def distribute_adjoint_sources(params):
"""
Searches through the OUTPUT directory of the lasif project, and distributes
the adjoint sources to their appropriate simulated directories. Also writes
the STATIONS_ADJOINT file.
"""
lasif_output_dir = os.path.join(params['lasif_path'], 'OUTPUT')
forward_run_dir = params['forward_run_dir']
for dir in sorted(os.listdir(lasif_output_dir))[::-1]:
if 'adjoint' not in dir:
continue
adjoint_source_dir = os.path.join(lasif_output_dir, dir)
event_name = adjoint_source_dir[adjoint_source_dir.find('GCMT'):]
solver_data_dir = os.path.join(forward_run_dir, event_name, 'DATA')
solver_sem_dir = os.path.join(forward_run_dir, event_name, 'SEM')
adjoint_stations = os.listdir(adjoint_source_dir)
utils.print_ylw("Copying adjoint sources for " + event_name + "...")
utils.mkdir_p(solver_sem_dir)
# This is necessary because the adjoint source names expected by
# specfem3d_globe are opposite of what lasif puts out. Could just fix
# this in lasif instead. This line would then reduce to a
# 'copy_directory'.
for old_name in os.listdir(adjoint_source_dir):
fields = old_name.split('.')
new_name = fields[1] + '.' + fields[0] + '.' + fields[2] + '.' + \
fields[3]
utils.safe_copy_file(os.path.join(adjoint_source_dir, old_name),
os.path.join(solver_sem_dir, new_name))
adjoint_stations = sorted(
list(
set([
fields0.split('.')[0] + '.' + fields0.split('.')[1]
for fields0 in adjoint_stations
])))
with open(os.path.join(solver_data_dir, 'STATIONS_ADJOINT'), 'w') as \
sta_adj:
sta = open(os.path.join(solver_data_dir, 'STATIONS'), 'r')
for line in sta:
if line.split()[0] + '.' + line.split()[1] in adjoint_stations:
sta_adj.write(line)
utils.print_blu('Done.')
def _tar_seismograms(dir):
utils.print_ylw("Tarring: " + dir + "...")
tar_files = []
tar_name = os.path.join(dir, "data.tar")
for file in os.listdir(dir):
if file != "data.tar":
tar_files.append(os.path.join(dir, file))
if not tar_files:
return
with tarfile.open(tar_name, "w") as tar:
for file in tar_files:
tar.add(file, arcname=os.path.basename(file))
for file in tar_files:
os.remove(file)
def _copy_files_and_tar((source, dest, tar)):
utils.print_ylw("Moving: " + source.split('/')[-2])
tar_name = os.path.join(dest, 'data.tar')
if os.path.exists(tar_name):
os.remove(tar_name)
if not os.path.exists(dest):
utils.mkdir_p(dest)
source_files = [x for x in os.listdir(source) if x.endswith('.sac')]
if tar:
with tarfile.open(tar_name, "w") as tar:
for file in source_files:
tar.add(os.path.join(source, file), arcname=file)
else:
utils.move_directory(source, dest, ends='.sac')
def _print_help(fcts):
"""
Prints the master help.
"""
utils.print_ylw(
"\n\nWelcome to the oval office. There's no clever acronym "
"here, I was just reading the Hunt for the Red October while writing "
"this.\n\n")
fct_groups = {}
for fct_name, fct in fcts.iteritems():
group_name = fct.group_name if hasattr(fct, "group_name") else "Misc"
fct_groups.setdefault(group_name, {})
fct_groups[group_name][fct_name] = fct
for group_name in sorted(fct_groups.iterkeys()):
utils.print_red(("{0:=>25s} Functions".format(" " + group_name)))
current_functions = fct_groups[group_name]
for name in sorted(current_functions.keys()):
utils.print_cyn(name)
utils.print_gry(_get_cmd_description(fcts[name]))
def prepare_solver(params):
"""
Sets up symbolic link to generated mesh files.
"""
forward_run_dir = params['forward_run_dir']
event_list = params['event_list']
utils.print_ylw("Preparing solver directories...")
databases_mpi = os.path.join(params['forward_run_dir'], 'mesh',
'DATABASES_MPI')
output_files = os.path.join(params['forward_run_dir'], 'mesh',
'OUTPUT_FILES')
for dir in sorted(os.listdir(params['forward_run_dir'])):
if dir not in event_list:
continue
utils.print_ylw('Linking ' + dir + '...')
dest_mpi = os.path.join(forward_run_dir, dir, 'DATABASES_MPI')
dest_out = os.path.join(forward_run_dir, dir, 'OUTPUT_FILES')
utils.sym_link_directory(databases_mpi, dest_mpi)
utils.sym_link_directory(output_files, dest_out)
utils.print_blu('Done.')
def _download_bulk_waveforms(
(event, networks, stations, start_times, end_times, download_path,
event_xml_directory, recording_time, padding_time)):
time.sleep(1)
c = Client("IRIS")
# Make directory.
utils.mkdir_p(os.path.join(download_path, event))
filename = os.path.join(download_path, event, 'data.mseed')
# Don't re-download files.
if os.path.exists(filename):
return
# Find start time and end time.
event_info = os.path.join(event_xml_directory, event + '.xml')
tree = ET.parse(event_info)
root = tree.getroot()
for tag in root.iter():
if tag.tag == XML_STRING + 'time':
time.start = obspy.UTCDateTime(
tag.findall(XML_STRING + 'value')[0].text) - padding_time
time.end = time.start + recording_time + 2 * padding_time
break
# Set up request.
bulk = []
for x, y, s_time, e_time in zip(networks, stations, start_times, end_times):
if time.start < s_time or time.start > e_time:
continue
bulk.append((x, y, '*', '*H*', time.start, time.end))
utils.print_ylw('Downloading %s...' % (event))
c.get_waveforms_bulk(bulk, filename=filename, quality='B')
def prefilter_data(params):
"""
Rotates the downloaded data in the "DOWNLOADED" folder. Will fail if no
StationXML files exist, so make sure you've done this first.
"""
# Local variables.
channel_priority = ['LH*', 'BH*']
lasif_path = params['lasif_path']
lasif_data_path = os.path.join(params['lasif_path'], 'DOWNLOADED_DATA')
event_xml_directory = os.path.join(params['lasif_path'], 'EVENTS')
lasif_stations_path = os.path.join(
params['lasif_path'], 'STATIONS', 'StationXML')
# Get starttime for event.
for event in sorted(os.listdir(event_xml_directory)):
utils.print_ylw("Prefiltering event: " + event[:-4])
download_path = os.path.join(lasif_data_path, event[:-4])
# Make raw data directory in LASIF project normal spot.
lasif_raw_data_path = os.path.join(
lasif_path, 'DATA', event[:-4], 'raw')
utils.mkdir_p(lasif_raw_data_path)
st = obspy.read(os.path.join(download_path, 'data.mseed'))
st_filt = obspy.Stream()
unique_stations = set([x.stats.station for x in st])
for s in unique_stations:
station_components = st.select(station=s)
for c in channel_priority:
if len(station_components.select(channel=c)):
st_filt += station_components.select(channel=c)
break
write_filename = os.path.join(lasif_raw_data_path, 'data.mseed')
st_filt.write(filename=write_filename, format='mseed')
def setup_new_iteration(params, old_iteration, new_iteration):
"""
Sets up a new iteration, and links the mesh files from the old iteration to
the new one.
"""
forward_stage_dir = params['forward_stage_dir']
event_list = params['event_list']
params.update({'iteration_name': new_iteration})
new_forward_run_dir = os.path.join(forward_stage_dir, new_iteration)
params.update({'forward_run_dir': new_forward_run_dir})
setup_solver(params)
old_database_dir = os.path.join(forward_stage_dir, old_iteration, 'mesh',
'DATABASES_MPI')
new_database_dir = os.path.join(forward_stage_dir, new_iteration, 'mesh',
'DATABASES_MPI')
old_optimization_dir = os.path.join(forward_stage_dir, old_iteration,
'OPTIMIZATION', 'PROCESSED_KERNELS')
new_optimization_dir = os.path.join(forward_stage_dir, new_iteration,
'OPTIMIZATION', 'PROCESSED_KERNELS')
utils.print_ylw("Copying mesh information...")
utils.copy_directory(old_database_dir, new_database_dir)
utils.print_ylw("Copying kernels...")
utils.copy_directory(old_optimization_dir,
new_optimization_dir,
ends='smooth.bin')
utils.print_ylw("Copying DATA files...")
for event in event_list:
old_data_dir = os.path.join(forward_stage_dir, old_iteration, event,
'DATA')
new_data_dir = os.path.join(forward_stage_dir, new_iteration, event,
'DATA')
utils.copy_directory(old_data_dir,
new_data_dir,
only=['Par_file', 'CMTSOLUTION', 'STATIONS'])
old_mesh_dir = os.path.join(forward_stage_dir, old_iteration, 'mesh',
'DATA')
new_mesh_dir = os.path.join(forward_stage_dir, new_iteration, 'mesh',
'DATA')
utils.copy_directory(old_mesh_dir,
new_mesh_dir,
only=['Par_file', 'CMTSOLUTION', 'STATIONS'])
utils.print_blu('Done.')
def delete_adjoint_sources_for_iteration(params):
"""
Deletes the directories on both /scratch and /project which contain the
adjoint sources for this iteration. As well, cleans the SEM and
STATIONS_ADJOINT files for the solver.
"""
forward_run_dir = params['forward_run_dir']
lasif_scratch_dir_output = os.path.join(params['lasif_scratch_path'],
'OUTPUT')
lasif_dir_output = os.path.join(params['lasif_path'], 'OUTPUT')
lasif_scratch_dir_adjoint = os.path.join(params['lasif_scratch_path'],
'ADJOINT_SOURCES_AND_WINDOWS')
lasif_project_dir_adjoint = os.path.join(params['lasif_path'],
'ADJOINT_SOURCES_AND_WINDOWS')
iteration_name = params['iteration_name']
utils.print_ylw("Cleaning forward runs...")
for dir in os.listdir(forward_run_dir):
if not os.path.exists(os.path.join(forward_run_dir, dir, 'SEM')):
continue
for file in os.listdir(os.path.join(forward_run_dir, dir, 'SEM')):
os.remove(os.path.join(forward_run_dir, dir, 'SEM', file))
utils.print_ylw("Cleaning LASIF scratch...")
for dir in os.listdir(lasif_scratch_dir_output):
if iteration_name and 'adjoint_sources' in dir:
shutil.rmtree(os.path.join(lasif_scratch_dir_output, dir))
for dir in os.listdir(lasif_scratch_dir_adjoint):
for sub_dir in os.listdir(os.path.join(lasif_scratch_dir_adjoint,
dir)):
shutil.rmtree(os.path.join(lasif_scratch_dir_adjoint, dir,
sub_dir))
utils.print_ylw("Cleaning LASIF project...")
for dir in os.listdir(lasif_dir_output):
if iteration_name and 'adjoint_sources' in dir:
shutil.rmtree(os.path.join(lasif_dir_output, dir))
for dir in os.listdir(lasif_project_dir_adjoint):
for sub_dir in os.listdir(os.path.join(lasif_project_dir_adjoint,
dir)):
shutil.rmtree(os.path.join(lasif_project_dir_adjoint, dir,
sub_dir))
utils.print_blu('Done.')
def _change_job_and_par_file(params, run_type):
"""
Runs and changes the simulation_type and simulation time in the submission
files.
"""
forward_run_dir = params['forward_run_dir']
forward_stage_dir = params['forward_stage_dir']
event_list = params['event_list']
if run_type == 'adjoint_run':
simulation_type = '= 3'
save_forward = '= .false.\n'
sbatch_time = '#SBATCH --time=04:00:00\n'
undo_attenuation = '= .true.\n'
elif run_type == 'first_iteration' or run_type == 'update_mesh' \
or run_type == 'forward_run':
simulation_type = '= 1'
save_forward = '= .true.\n'
sbatch_time = '#SBATCH --time=02:00:00\n'
undo_attenuation = '= .true.\n'
elif run_type == 'line_search':
simulation_type = '= 1'
save_forward = '= .false.\n'
sbatch_time = '#SBATCH --time=02:00:00\n'
undo_attenuation = '= .false.\n'
utils.print_ylw("Modifying Par_files for run type...")
os.chdir(forward_run_dir)
for dir in os.listdir('./'):
if dir not in event_list and dir != 'mesh':
continue
par_path = os.path.join(dir, 'DATA', 'Par_file')
par_path_new = os.path.join(dir, 'DATA', 'Par_file_new')
new_file = open(par_path_new, 'w')
with open(par_path, 'r') as file:
for line in file:
fields = line.split('=')
if len(fields) == 1:
new_file.write(line)
continue
if 'SIMULATION_TYPE' in fields[0]:
new_file.write(fields[0] + simulation_type + fields[1][2:])
elif 'SAVE_FORWARD' in fields[0]:
new_file.write(fields[0] + save_forward)
elif fields[0].startswith('MODEL') and \
run_type == 'update_mesh':
new_file.write(fields[0] + '= CEM_GLL\n')
elif fields[0].startswith('MODEL') and \
run_type == 'first_iteration':
new_file.write(fields[0] + '= CEM_ACCEPT\n')
elif fields[0].startswith('UNDO_ATTENUATION'):
new_file.write(fields[0] + undo_attenuation)
else:
new_file.write(line)
os.rename(par_path_new, par_path)
utils.print_ylw("Modifying .sbatch file for run type...")
os.chdir(forward_stage_dir)
job_array = os.path.join('./', 'jobArray_solver_daint.sbatch')
new_job_array_name = os.path.join('./', 'jobArray_solver_daint.sbatch_new')
new_job_array = open(new_job_array_name, 'w')
with open(job_array, 'r') as file:
for line in file:
if '--time' in line:
new_job_array.write(sbatch_time)
else:
new_job_array.write(line)
new_job_array.close()
os.remove(job_array)
os.rename(new_job_array_name, job_array)
def download_data(params, station_list, with_waveforms, recording_time,
padding_time):
"""
Still a work in progress (perhaps never finished). Sorts a text file
obtained from IRIS (see manual), and parses out the STS and KS instruments
(apparently the best ones). Then passes
"""
# Domain boundaries
min_lat = -65
max_lat = 45
min_lon = -47.5
max_lon = 75
# Set up paths and such.
lasif_data_path = os.path.join(params['lasif_path'], 'DOWNLOADED_DATA')
iteration_xml_path = params['iteration_xml_path']
event_xml_directory = os.path.join(params['lasif_path'], 'EVENTS')
event_list = params['event_list']
lasif_stations_path = os.path.join(params['lasif_path'], 'STATIONS',
'StationXML')
# Set up station tuple and allowable instruments.
station = namedtuple('station', ['network', 'station', 'location',
'sensor', 's_time', 'e_time'])
# Read the data and parse out the important components.
stations_list = pd.read_csv(station_list, delimiter='|')
stations_list.fillna('00', inplace=True)
# Filter based on domain boundaries
stations_list = stations_list[stations_list.Latitude > min_lat]
stations_list = stations_list[stations_list.Latitude < max_lat]
stations_list = stations_list[stations_list.Longitude > min_lon]
stations_list = stations_list[stations_list.Longitude < max_lon]
stations_list = stations_list[stations_list.Location == '00']
stations_list['StartTime'] = \
stations_list['StartTime'].astype(obspy.UTCDateTime)
stations_list['EndTime'] = \
stations_list['EndTime'].astype(obspy.UTCDateTime)
# Number of events.
num_events = len(os.listdir(event_xml_directory))
event_names = sorted([x[:-4] for x in os.listdir(event_xml_directory)])
# Event arrays.
networks = stations_list.Network
stations = stations_list.Station
start_time = stations_list.StartTime
end_time = stations_list.EndTime
# Waveforms.
pool = Pool(processes=NUM_THREADS)
pool.map(_download_bulk_waveforms, zip(
event_names, repeat(networks), repeat(stations), repeat(start_time),
repeat(end_time), repeat(lasif_data_path), repeat(event_xml_directory),
repeat(recording_time), repeat(padding_time)))
if with_waveforms:
return
# Get stations.
c = Client("IRIS")
for x in stations_filt:
station_filename = os.path.join(
lasif_stations_path, 'station.%s_%s.xml' % (x.network, x.station))
if os.path.exists(station_filename):
continue
utils.print_ylw(
"Downloading StationXML for: %s.%s" % (x.network, x.station))
try:
c.get_stations(
network=x.network, station=x.station, location="*", channel="*",
level="response", filename=station_filename)
except:
utils.print_red("No data for %s" % (station_filename))
def setup_solver(params):
"""
It's dylan, you know the drill.
"""
# Setup local parameters.
forward_run_dir = params['forward_run_dir']
forward_stage_dir = params['forward_stage_dir']
event_list = params['event_list']
lasif_path = params['lasif_path']
iteration_name = params['iteration_name']
specfem_root = params['specfem_root']
compiler_suite = params['compiler_suite']
project_name = params['project_name']
# Set up the mesh directory.
_setup_dir_tree('mesh', forward_run_dir)
# Set up the optimization directory.
optimization_base = os.path.join(forward_run_dir, 'OPTIMIZATION')
utils.mkdir_p(optimization_base)
utils.mkdir_p(os.path.join(optimization_base, 'bin'))
utils.mkdir_p(os.path.join(optimization_base, 'PROCESSED_KERNELS'))
utils.mkdir_p(os.path.join(optimization_base, 'GRADIENT_INFO'))
utils.mkdir_p(os.path.join(optimization_base, 'LOGS'))
utils.mkdir_p(os.path.join(optimization_base, 'DATA'))
utils.mkdir_p(os.path.join(optimization_base, 'VTK_FILES'))
# Create the forward modelling directories. Also copy relevant parameter
# files from the LASIF project. _copy_input_files also copies the input
# files to the specfem_root directory if mesh == True.
utils.print_ylw("Creating forward modelling directories...")
mesh = True
for i, event in enumerate(event_list):
_setup_dir_tree(event, forward_run_dir)
_copy_input_files(event,
forward_run_dir,
lasif_path,
iteration_name,
specfem_root,
mesh=mesh)
mesh = False
# Copy the files in SUBMISSION to the specfem root directory.
par_file = os.path.join(lasif_path, 'SUBMISSION', project_name, 'Par_file')
dest = os.path.join(specfem_root, 'DATA')
utils.safe_copy(par_file, dest)
# Change to specfem root directory and compile.
utils.print_ylw("Compiling...")
os.chdir(specfem_root)
with open('compilation_log.txt', 'w') as output:
proc = subprocess.Popen(['./mk_daint.sh', compiler_suite, 'adjoint'],
stdout=output,
stderr=output)
proc.communicate()
proc.wait()
# Distribute binaries and Par_file to directories.
utils.print_ylw('Copying compiled binaries...')
bin_directory = os.path.join('./bin')
opt_bin_directory = os.path.join(optimization_base, 'bin')
opt_dat_directory = os.path.join(optimization_base, 'DATA')
utils.copy_directory(bin_directory, opt_bin_directory)
for event in os.listdir(forward_run_dir):
event_bin = os.path.join(forward_run_dir, event, 'bin')
event_dat = os.path.join(forward_run_dir, event, 'DATA')
compile_par = os.path.join(specfem_root, 'DATA', 'Par_file')
utils.safe_copy(compile_par, event_dat)
utils.copy_directory(bin_directory,
event_bin,
only=['xspecfem3D', 'xmeshfem3D'])
# Also copy to the optimization directory. Recompile with vectorized cray
# compiler.
utils.print_ylw("Recompiling for vectorized smoother CRAY smoother...")
with open('compilation_log_tomo.txt', 'w') as output:
proc = subprocess.Popen(['./mk_daint.sh', 'cray.tomo', 'adjoint'],
stdout=output,
stderr=output)
proc.communicate()
proc.wait()
utils.copy_directory(bin_directory, opt_bin_directory)
compile_par = os.path.join(specfem_root, 'DATA', 'Par_file')
utils.safe_copy(compile_par, opt_dat_directory)
# Copy jobarray script to base directory.
utils.print_ylw('Copying jobarray sbatch script...')
source = os.path.join(lasif_path, 'SUBMISSION', project_name,
'jobArray_solver_daint.sbatch')
utils.safe_copy(source, forward_stage_dir)
utils.mkdir_p(os.path.join(forward_stage_dir, 'logs'))
# Copy mesh submission script.
source = os.path.join(lasif_path, 'SUBMISSION', project_name,
'job_mesher_daint.sbatch')
dest = os.path.join(forward_run_dir, 'mesh')
utils.safe_copy(source, dest)
# Copy topography information to mesh directory.
utils.print_ylw('Copying topography information...')
master_topo_path = os.path.join(specfem_root, 'DATA', 'topo_bathy')
mesh_topo_path = os.path.join(forward_run_dir, 'mesh', 'DATA',
'topo_bathy')
utils.mkdir_p(mesh_topo_path)
utils.copy_directory(master_topo_path, mesh_topo_path)
utils.print_blu('Done.')