code cleanup and commenting

This commit is contained in:
Marcel Paffrath 2016-05-19 11:20:37 +02:00
parent dbaead4754
commit 3138bbfa93

View File

@ -6,23 +6,15 @@ import datetime
import numpy as np import numpy as np
class Tomo3d(object): class Tomo3d(object):
def __init__(self, nproc, iterations, citer = 0, overwrite = False): def __init__(self, citer = 0, overwrite = False):
''' '''
Class build from FMTOMO script tomo3d. Can be used to run several instances of FMM code in parallel. Class build from FMTOMO script tomo3d. Can be used to run several instances of FMM code in parallel.
:param: nproc, number of parallel processes
:type: integer
:param: iterations, number of iterations
:type: integer
:param: citer, current iteration (default = 0: start new model) :param: citer, current iteration (default = 0: start new model)
:type: integer :type: integer
''' '''
self.setCWD() self.setCWD()
self.defParas() self.defParas()
self.nproc = nproc
self.iter = iterations # number of iterations
self.citer = citer # current iteration self.citer = citer # current iteration
self.sources = self.readSrcFile() self.sources = self.readSrcFile()
self.traces = self.readTraces() self.traces = self.readTraces()
@ -34,21 +26,37 @@ class Tomo3d(object):
self.defInvParas() self.defInvParas()
def defFMMParas(self): def defFMMParas(self):
'''
Initiates parameters for the forward calculation.
'''
# Name of fast marching program
self.fmm = '{0}/fm3d'.format(self.cwd) self.fmm = '{0}/fm3d'.format(self.cwd)
# Name of program calculating Frechet derivatives
self.frechgen = '{0}/frechgen'.format(self.cwd) self.frechgen = '{0}/frechgen'.format(self.cwd)
# Name of current velocity/inversion grid
self.cvg = 'vgrids.in' self.cvg = 'vgrids.in'
# Name of current interfaces grid
self.cig = 'interfaces.in' self.cig = 'interfaces.in'
# Name of file containing current source locations
self.csl = 'sources.in' self.csl = 'sources.in'
# Name of file containing propagation grid
self.pg = 'propgrid.in' self.pg = 'propgrid.in'
# Name of file containing receiver coordinates
self.rec = 'receivers.in' self.rec = 'receivers.in'
self.frech = 'frechet.in' self.frech = 'frechet.in'
self.frechout = 'frechet.dat' self.frechout = 'frechet.dat'
# Name of file containing measured data
self.ot = 'otimes.dat' self.ot = 'otimes.dat'
# Name of file containing output velocity information
self.ttim = 'arrivals.dat' self.ttim = 'arrivals.dat'
self.mode = 'mode_set.in' self.mode = 'mode_set.in'
# Name of temporary folders created for each process
self.folder = '.proc_' self.folder = '.proc_'
def defInvParas(self): def defInvParas(self):
'''
Initiates inversion parameters for FMTOMO.
'''
# Name of program for performing inversion # Name of program for performing inversion
self.inv = '{0}/invert3d'.format(self.cwd) self.inv = '{0}/invert3d'.format(self.cwd)
# Name of file containing current model traveltimes # Name of file containing current model traveltimes
@ -67,19 +75,41 @@ class Tomo3d(object):
self.resout = 'residuals.dat' self.resout = 'residuals.dat'
def copyRef(self): def copyRef(self):
# Copies reference grids to used grids (e.g. sourcesref.in to sources.in) '''
Copies reference grids to used grids (e.g. sourcesref.in to sources.in)
'''
os.system('cp %s %s'%(self.ivg, self.cvg)) os.system('cp %s %s'%(self.ivg, self.cvg))
os.system('cp %s %s'%(self.iig, self.cig)) os.system('cp %s %s'%(self.iig, self.cig))
os.system('cp %s %s'%(self.isl, self.csl)) os.system('cp %s %s'%(self.isl, self.csl))
def setCWD(self): def setCWD(self, directory = None):
self.cwd = subprocess.check_output(['pwd'])[0:-1] '''
print('Working directory is pwd: %s'%self.cwd) Set working directory containing all necessary files.
Default: pwd
'''
if directory == None:
directory = subprocess.check_output(['pwd'])[0:-1]
self.cwd = directory
print('Working directory is: %s'%self.cwd)
def runFrech(self): def runFrech(self):
os.system(self.frechgen) os.system(self.frechgen)
def runTOMO3D(self): def runTOMO3D(self, nproc, iterations):
'''
Starts up the FMTOMO code for the set number of iterations on nproc parallel processes.
:param: nproc, number of parallel processes
:type: integer
:param: iterations, number of iterations
:type: integer
'''
self.nproc = nproc
self.iter = iterations # number of iterations
starttime = datetime.datetime.now() starttime = datetime.datetime.now()
print('Starting TOMO3D on %s parallel processes for %s iteration(s).' print('Starting TOMO3D on %s parallel processes for %s iteration(s).'
%(self.nproc, self.iter)) %(self.nproc, self.iter))
@ -104,6 +134,10 @@ class Tomo3d(object):
print('runTOMO3D: Finished %s iterations after %s.'%(self.iter, tdelta)) print('runTOMO3D: Finished %s iterations after %s.'%(self.iter, tdelta))
def runFmm(self, directory, logfile, processes): def runFmm(self, directory, logfile, processes):
'''
Calls an instance of the FMM code in the process directory.
Requires a list of all active processes and returns an updated list.
'''
os.chdir(directory) os.chdir(directory)
fout = open(logfile, 'w') fout = open(logfile, 'w')
processes.append(subprocess.Popen(self.fmm, stdout = fout)) processes.append(subprocess.Popen(self.fmm, stdout = fout))
@ -112,6 +146,9 @@ class Tomo3d(object):
return processes return processes
def startForward(self, logdir): def startForward(self, logdir):
'''
Runs an instance of the FMM code in the process directory.
'''
self._printLine() self._printLine()
print('Starting forward simulation for iteration %s.'%(self.citer)) print('Starting forward simulation for iteration %s.'%(self.citer))
@ -128,8 +165,8 @@ class Tomo3d(object):
logfn = 'fm3dlog_' + str(procID) + '.out' logfn = 'fm3dlog_' + str(procID) + '.out'
log_out = os.path.join(logdir, logfn) log_out = os.path.join(logdir, logfn)
self.writeSrcFile(procID, directory) self.writeSrcFile(procID)
self.writeTracesFile(procID, directory) self.writeTracesFile(procID)
os.system('cp {cvg} {cig} {mode} {pg} {frechin} {dest}' os.system('cp {cvg} {cig} {mode} {pg} {frechin} {dest}'
.format(cvg=self.cvg, cig=self.cig, frechin=self.frech, .format(cvg=self.cvg, cig=self.cig, frechin=self.frech,
mode=self.mode, pg=self.pg, dest=directory)) mode=self.mode, pg=self.pg, dest=directory))
@ -149,10 +186,16 @@ class Tomo3d(object):
print('Finished Forward calculation after %s'%tdelta) print('Finished Forward calculation after %s'%tdelta)
def startInversion(self): def startInversion(self):
'''
Simply calls the inversion program.
'''
print('Calling %s...'%self.inv) print('Calling %s...'%self.inv)
os.system(self.inv) os.system(self.inv)
def calcRes(self): def calcRes(self):
'''
Calls residual calculation program.
'''
resout = os.path.join(self.cwd, self.resout) resout = os.path.join(self.cwd, self.resout)
if self.citer == 0: if self.citer == 0:
os.system('%s > %s'%(self.resid, resout)) os.system('%s > %s'%(self.resid, resout))
@ -185,11 +228,17 @@ class Tomo3d(object):
raise RuntimeError('Could not create directory: %s'%directory) raise RuntimeError('Could not create directory: %s'%directory)
def makeDirectories(self): def makeDirectories(self):
'''
Makes temporary directories for all processes.
'''
for procID in range(1, self.nproc + 1): for procID in range(1, self.nproc + 1):
directory = self.getProcDir(procID) directory = self.getProcDir(procID)
self.makeDir(directory) self.makeDir(directory)
def makeInvIterDir(self): def makeInvIterDir(self):
'''
Makes directories for each iteration step for the output.
'''
invIterDir = self.cwd + '/it_%s'%(self.citer) invIterDir = self.cwd + '/it_%s'%(self.citer)
err = os.system('mkdir %s'%invIterDir) err = os.system('mkdir %s'%invIterDir)
if err is 256: if err is 256:
@ -200,12 +249,18 @@ class Tomo3d(object):
self.cInvIterDir = invIterDir self.cInvIterDir = invIterDir
def clearDir(self, directory): def clearDir(self, directory):
'''
Wipes a certain directory.
'''
print('Wiping directory %s...'%directory) print('Wiping directory %s...'%directory)
for filename in os.listdir(directory): for filename in os.listdir(directory):
filenp = os.path.join(directory, filename) filenp = os.path.join(directory, filename)
os.remove(filenp) os.remove(filenp)
def clearDirectories(self): def clearDirectories(self):
'''
Wipes all generated temporary directories.
'''
for directory in self.directories: for directory in self.directories:
self.clearDir(directory) self.clearDir(directory)
@ -214,14 +269,24 @@ class Tomo3d(object):
return os.rmdir(directory) return os.rmdir(directory)
def removeDirectories(self): def removeDirectories(self):
'''
Removes all generated temporary directories.
'''
for directory in self.directories: for directory in self.directories:
self.rmDir(directory) self.rmDir(directory)
self.directories = [] self.directories = []
def getProcDir(self, procID): def getProcDir(self, procID):
'''
Returns the temporary directory for a certain process
with procID = process number.
'''
return os.path.join(self.cwd, self.folder) + str(procID) return os.path.join(self.cwd, self.folder) + str(procID)
def getTraceIDs4Sources(self, sourceIDs): def getTraceIDs4Sources(self, sourceIDs):
'''
Returns corresponding trace IDs for a set of given source IDs.
'''
traceIDs = [] traceIDs = []
for traceID in self.traces.keys(): for traceID in self.traces.keys():
if self.traces[traceID]['source'] in sourceIDs: if self.traces[traceID]['source'] in sourceIDs:
@ -229,6 +294,9 @@ class Tomo3d(object):
return traceIDs return traceIDs
def getTraceIDs4Source(self, sourceID): def getTraceIDs4Source(self, sourceID):
'''
Returns corresponding trace IDs for a source ID.
'''
traceIDs = [] traceIDs = []
for traceID in self.traces.keys(): for traceID in self.traces.keys():
if self.traces[traceID]['source'] == sourceID: if self.traces[traceID]['source'] == sourceID:
@ -236,22 +304,38 @@ class Tomo3d(object):
return traceIDs return traceIDs
def copyArrivals(self, target = None): def copyArrivals(self, target = None):
'''
Copies the FMM output file (self.ttim) to a specific target file.
Default target is self.mtrav (model travel times).
'''
if target == None: if target == None:
target = os.path.join(self.cwd, self.mtrav) target = os.path.join(self.cwd, self.mtrav)
os.system('cp %s %s'%(os.path.join( os.system('cp %s %s'%(os.path.join(
self.cInvIterDir, self.ttim), target)) self.cInvIterDir, self.ttim), target))
def saveVgrid(self): def saveVgrid(self):
vgpath = os.path.join(self.cwd, 'vgrids.in') '''
Saves the current velocity grid for the current iteration step.
'''
vgpath = os.path.join(self.cwd, self.cvg)
os.system('cp %s %s'%(vgpath, self.cInvIterDir)) os.system('cp %s %s'%(vgpath, self.cInvIterDir))
def calcSrcPerKernel(self): def calcSrcPerKernel(self):
'''
(Equally) distributes all sources depending on the number of processes (kernels).
Returns two integer values.
First: minimum number of sources for each process
Second: remaining sources (always less than number of processes)
'''
nsrc = self.readNsrc() nsrc = self.readNsrc()
if self.nproc > nsrc: if self.nproc > nsrc:
raise ValueError('Number of spawned processes higher than number of sources') raise ValueError('Number of spawned processes higher than number of sources')
return nsrc/self.nproc, nsrc%self.nproc return nsrc/self.nproc, nsrc%self.nproc
def srcIDs4Kernel(self, procID): def srcIDs4Kernel(self, procID):
'''
Calculates and returns all source IDs for a given process ID.
'''
proc = procID - 1 proc = procID - 1
nsrc = self.readNsrc() nsrc = self.readNsrc()
srcPK, remain = self.calcSrcPerKernel() srcPK, remain = self.calcSrcPerKernel()
@ -274,12 +358,18 @@ class Tomo3d(object):
return nsrc return nsrc
def readNtraces(self): def readNtraces(self):
'''
Reads the total number of traces from self.rec header.
'''
recfile = open(self.rec, 'r') recfile = open(self.rec, 'r')
nrec = int(recfile.readline()) nrec = int(recfile.readline())
recfile.close() recfile.close()
return nrec return nrec
def readSrcFile(self): def readSrcFile(self):
'''
Reads the whole sourcefile and returns structured information in a dictionary.
'''
nsrc = self.readNsrc() nsrc = self.readNsrc()
srcfile = open(self.csl, 'r') srcfile = open(self.csl, 'r')
@ -309,6 +399,10 @@ class Tomo3d(object):
return sources return sources
def readTraces(self): def readTraces(self):
'''
Reads the receiver input file and returns the information
in a structured dictionary.
'''
recfile = open(self.rec, 'r') recfile = open(self.rec, 'r')
ntraces = self.readNtraces() ntraces = self.readNtraces()
@ -330,8 +424,13 @@ class Tomo3d(object):
return traces return traces
def readArrivals(self, procID): def readArrivals(self, procID):
'''
Reads the arrival times from a temporary process directory,
changes local to global sourceIDs and traceIDs and returns
a list of arrival times.
'''
directory = self.getProcDir(procID) directory = self.getProcDir(procID)
arrfile = open(directory + '/arrivals.dat', 'r') arrfile = open(os.path.join(directory, self.ttime), 'r')
sourceIDs = self.srcIDs4Kernel(procID) sourceIDs = self.srcIDs4Kernel(procID)
arrivals = [] arrivals = []
@ -347,6 +446,10 @@ class Tomo3d(object):
return arrivals return arrivals
def readRays(self, procID): def readRays(self, procID):
'''
Reads rays output from a temporary process directory and returns
the information in a structured dictionary.
'''
directory = self.getProcDir(procID) directory = self.getProcDir(procID)
raysfile = open(directory + '/rays.dat', 'r') raysfile = open(directory + '/rays.dat', 'r')
sourceIDs = self.srcIDs4Kernel(procID) sourceIDs = self.srcIDs4Kernel(procID)
@ -385,8 +488,12 @@ class Tomo3d(object):
} }
return rays return rays
def writeSrcFile(self, procID, directory): def writeSrcFile(self, procID):
srcfile = open('%s/sources.in'%directory, 'w') '''
Writes a source input file for a process with ID = procID.
'''
directory = self.getProcDir(procID)
srcfile = open(os.path,join(directory, self.csl), 'w')
sourceIDs = self.srcIDs4Kernel(procID) sourceIDs = self.srcIDs4Kernel(procID)
srcfile.write('%s\n'%len(sourceIDs)) srcfile.write('%s\n'%len(sourceIDs))
@ -401,7 +508,11 @@ class Tomo3d(object):
srcfile.write('%s %s\n'%(int(interactions[0]), int(interactions[1]))) srcfile.write('%s %s\n'%(int(interactions[0]), int(interactions[1])))
srcfile.write('%s\n'%source['veltype']) srcfile.write('%s\n'%source['veltype'])
def writeTracesFile(self, procID, directory): def writeTracesFile(self, procID):
'''
Writes a receiver input file for a process with ID = procID.
'''
directory = self.getProcDir(procID)
recfile = open('%s/receivers.in'%directory, 'w') recfile = open('%s/receivers.in'%directory, 'w')
sourceIDs = self.srcIDs4Kernel(procID) sourceIDs = self.srcIDs4Kernel(procID)
traceIDs = self.getTraceIDs4Sources(sourceIDs) traceIDs = self.getTraceIDs4Sources(sourceIDs)
@ -417,9 +528,12 @@ class Tomo3d(object):
recfile.write('%s\n'%trace['path']) recfile.write('%s\n'%trace['path'])
def mergeArrivals(self, directory): def mergeArrivals(self, directory):
'''
Merges the arrival times for all processes to self.cInvIterDir.
'''
arrfn = os.path.join(directory, self.ttim) arrfn = os.path.join(directory, self.ttim)
arrivalsOut = open(arrfn, 'w') arrivalsOut = open(arrfn, 'w')
print('Merging arrivals.dat...') print('Merging %s...'%self.ttim)
for procID in range(1, self.nproc + 1): for procID in range(1, self.nproc + 1):
arrivals = self.readArrivals(procID) arrivals = self.readArrivals(procID)
for line in arrivals: for line in arrivals:
@ -428,6 +542,9 @@ class Tomo3d(object):
os.system('ln -fs %s %s'%(arrfn, os.path.join(self.cwd, self.ttim))) os.system('ln -fs %s %s'%(arrfn, os.path.join(self.cwd, self.ttim)))
def mergeRays(self, directory): def mergeRays(self, directory):
'''
Merges the ray paths for all processes to self.cInvIterDir.
'''
print('Merging rays.dat...') print('Merging rays.dat...')
with open(directory + '/rays.dat', 'w') as outfile: with open(directory + '/rays.dat', 'w') as outfile:
for procID in range(1, self.nproc + 1): for procID in range(1, self.nproc + 1):
@ -448,9 +565,11 @@ class Tomo3d(object):
outfile.writelines(raysec['raypoints']) outfile.writelines(raysec['raypoints'])
def mergeFrechet(self, directory): def mergeFrechet(self, directory):
print('Merging frechet.dat...') '''
Merges the frechet derivatives for all processes to self.cInvIterDir.
'''
frechfnout = os.path.join(directory, self.frechout) frechfnout = os.path.join(directory, self.frechout)
print('Merging %s...'%self.frechout)
with open(frechfnout, 'w') as outfile: with open(frechfnout, 'w') as outfile:
for procID in range(1, self.nproc + 1): for procID in range(1, self.nproc + 1):
filename = os.path.join(self.getProcDir(procID), self.frechout) filename = os.path.join(self.getProcDir(procID), self.frechout)
@ -465,6 +584,9 @@ class Tomo3d(object):
os.system('ln -fs %s %s'%(frechfnout, os.path.join(self.cwd, self.frechout))) os.system('ln -fs %s %s'%(frechfnout, os.path.join(self.cwd, self.frechout)))
def mergeOutput(self, directory): def mergeOutput(self, directory):
'''
Calls self.mergeArrivals, self.mergeFrechet and self.mergeRays.
'''
self.mergeArrivals(directory) self.mergeArrivals(directory)
self.mergeFrechet(directory) self.mergeFrechet(directory)
self.mergeRays(directory) self.mergeRays(directory)