slimmed down the code

pylot/core/active/fmtomoUtils.py

@@ -6,83 +6,27 @@ def vgrids2VTK(inputfile = 'vgrids.in', outputfile = 'vgrids.vtk', absOrRel = 'a
     '''
     Generate a vtk-file readable by e.g. paraview from FMTOMO output vgrids.in
     '''
-    def getDistance(angle):
-        PI = np.pi
-        R = 6371.
-        distance = angle / 180 * (PI * R)
-        return distance
-
-    def readNumberOfPoints(filename):
-        fin = open(filename, 'r')
-        vglines = fin.readlines()
-
-        nR = int(vglines[1].split()[0])
-        nTheta = int(vglines[1].split()[1])
-        nPhi = int(vglines[1].split()[2])
-
-        print('readNumberOf Points: Awaiting %d grid points in %s'
-              %(nR*nTheta*nPhi, filename))
-        fin.close()
-        return nR, nTheta, nPhi
-
-    def readDelta(filename):
-        fin = open(filename, 'r')
-        vglines = fin.readlines()
-
-        dR = float(vglines[2].split()[0])
-        dTheta = float(vglines[2].split()[1])
-        dPhi = float(vglines[2].split()[2])
-
-        fin.close()
-        return dR, dTheta, dPhi
-
-    def readStartpoints(filename):
-        fin = open(filename, 'r')
-        vglines = fin.readlines()
-
-        sR = float(vglines[3].split()[0])
-        sTheta = float(vglines[3].split()[1])
-        sPhi = float(vglines[3].split()[2])
-
-        fin.close()
-        return sR, sTheta, sPhi
-
-    def readVelocity(filename):
-        '''
-        Reads in velocity from vgrids file and returns a list containing all values in the same order
-        '''
-        vel = []; count = 0
-        fin = open(filename, 'r')
-        vglines = fin.readlines()
-
-        for line in vglines:
-            count += 1
-            if count > 4:
-                vel.append(float(line.split()[0]))
-
-        print("Read %d points out of file: %s" %(count - 4, filename))
-        return vel
-
     R = 6371. # earth radius
     outfile = open(outputfile, 'w')
 
-    # Theta, Phi in radians, R in km
+    number, delta, start, vel = _readVgrid(inputfile)
-    nR, nTheta, nPhi = readNumberOfPoints(inputfile)
+
-    dR, dTheta, dPhi = readDelta(inputfile)
+    nR, nTheta, nPhi = number
-    sR, sTheta, sPhi = readStartpoints(inputfile)
+    dR, dTheta, dPhi = delta
-    vel = readVelocity(inputfile)
+    sR, sTheta, sPhi = start
+    
+    thetaGrid, phiGrid, rGrid = _generateGrids(number, delta, start)
+
+    nPoints = nR * nTheta * nPhi
 
     nX = nPhi; nY = nTheta; nZ = nR
 
     sZ = sR - R
-    sX = getDistance(np.rad2deg(sPhi))
+    sX = _getDistance(sPhi)
-    sY = getDistance(np.rad2deg(sTheta))
+    sY = _getDistance(sTheta)
-
-    dX = getDistance(np.rad2deg(dPhi))
-    dY = getDistance(np.rad2deg(dTheta))
-
-    nPoints = nX * nY * nZ
 
+    dX = _getDistance(dPhi)
+    dY = _getDistance(dTheta)
     dZ = dR
 
     # write header
@@ -109,7 +53,8 @@ def vgrids2VTK(inputfile = 'vgrids.in', outputfile = 'vgrids.vtk', absOrRel = 'a
         for velocity in vel:
             outfile.writelines('%10f\n' %velocity)
     elif absOrRel == 'rel':
-        velref = readVelocity(inputfileref)
+        nref, dref, sref, velref = _readVgrid(inputfileref)
+        nR_ref, nTheta_ref, nPhi_ref = nref
         if not len(velref) == len(vel):
             print('ERROR: Number of gridpoints mismatch for %s and %s'%(inputfile, inputfileref))
             return
@@ -122,7 +67,6 @@ def vgrids2VTK(inputfile = 'vgrids.in', outputfile = 'vgrids.vtk', absOrRel = 'a
             else:
                 velrel.append(0)            
 
-        nR_ref, nTheta_ref, nPhi_ref = readNumberOfPoints(inputfileref)
         if not nR_ref == nR and nTheta_ref == nTheta and nPhi_ref == nPhi:
             print('ERROR: Dimension mismatch of grids %s and %s'%(inputfile, inputfileref))
             return
@@ -142,12 +86,6 @@ def rays2VTK(fnin, fdirout = './vtk_files/', nthPoint = 50):
     :param: nthPoint, plot every nth point of the ray
     :type: integer
     '''
-    def getDistance(angle):
-        PI = np.pi
-        R = 6371.
-        distance = angle / 180 * (PI * R)
-        return distance
-
     infile = open(fnin, 'r')
     R = 6371
     rays = {}
@@ -155,7 +93,6 @@ def rays2VTK(fnin, fdirout = './vtk_files/', nthPoint = 50):
     nPoints = 0
 
     ### NOTE: rays.dat seems to be in km and radians
-
     while True:
         raynumber += 1
         firstline = infile.readline()
@@ -173,7 +110,7 @@ def rays2VTK(fnin, fdirout = './vtk_files/', nthPoint = 50):
         for index in range(nRayPoints):
             if index % nthPoint is 0 or index == (nRayPoints - 1):
                 rad, lat, lon = infile.readline().split()
-                rays[shotnumber][raynumber].append([getDistance(np.rad2deg(float(lon))), getDistance(np.rad2deg(float(lat))), float(rad) - R])
+                rays[shotnumber][raynumber].append([_getDistance(np.rad2deg(float(lon))), _getDistance(np.rad2deg(float(lat))), float(rad) - R])
             else:
                 dummy = infile.readline()
 
@@ -199,7 +136,6 @@ def rays2VTK(fnin, fdirout = './vtk_files/', nthPoint = 50):
 
         # write coordinates
         #print("Writing coordinates to VTK file...")
-
         for raynumber in rays[shotnumber].keys():
             for raypoint in rays[shotnumber][raynumber]:
                 outfile.writelines('%10f %10f %10f \n' %(raypoint[0], raypoint[1], raypoint[2]))
@@ -208,7 +144,6 @@ def rays2VTK(fnin, fdirout = './vtk_files/', nthPoint = 50):
 
         # write indices
         #print("Writing indices to VTK file...")
-
         count = 0
         for raynumber in rays[shotnumber].keys():
             outfile.writelines('%d ' %(len(rays[shotnumber][raynumber])))
@@ -217,29 +152,7 @@ def rays2VTK(fnin, fdirout = './vtk_files/', nthPoint = 50):
                 count += 1
             outfile.writelines('\n')
 
-    # outfile.writelines('POINT_DATA %15d\n' %(nPoints))
+def _readVgrid(filename):
-    # outfile.writelines('SCALARS rays float %d\n' %(1))
-    # outfile.writelines('LOOKUP_TABLE default\n')
-
-    # # write velocity
-    # print("Writing velocity values to VTK file...")
-    # for velocity in vel:
-    #     outfile.writelines('%10f\n' %velocity)
-
-    # outfile.close()
-    # print("Wrote velocity grid for %d points to file: %s" %(nPoints, outputfile))
-
-def addCheckerboard(spacing = 10., pertubation = 0.1, inputfile = 'vgrids.in',
-                    outputfile = 'vgrids_cb.in', ampmethod = 'linear', rect = (None, None)):
-    '''
-    Add a checkerboard to an existing vgrids.in velocity model.
-
-    :param: spacing, size of the tiles
-    type: float
-
-    :param: pertubation, pertubation (default: 0.1 = 10%)
-    type: float
-    '''
     def readNumberOfPoints(filename):
         fin = open(filename, 'r')
         vglines = fin.readlines()
@@ -291,6 +204,46 @@ def addCheckerboard(spacing = 10., pertubation = 0.1, inputfile = 'vgrids.in',
         print("Read %d points out of file: %s" %(count - 4, filename))
         return vel
 
+    # Theta, Phi in radians, R in km
+    nR, nTheta, nPhi = readNumberOfPoints(filename)
+    dR, dThetaRad, dPhiRad = readDelta(filename)
+    sR, sThetaRad, sPhiRad = readStartpoints(filename)
+    vel = readVelocity(filename)
+
+    dTheta, dPhi = np.rad2deg((dThetaRad, dPhiRad))
+    sTheta, sPhi = np.rad2deg((sThetaRad, sPhiRad))
+
+    number = (nR, nTheta, nPhi)
+    delta = (dR, dTheta, dPhi)
+    start = (sR, sTheta, sPhi)
+    return number, delta, start, vel
+
+def _generateGrids(number, delta, start):
+    nR, nTheta, nPhi = number
+    dR, dTheta, dPhi = delta
+    sR, sTheta, sPhi = start
+    
+    eR = sR + (nR - 1) * dR
+    ePhi = sPhi + (nPhi - 1) * dPhi
+    eTheta = sTheta + (nTheta - 1) * dTheta
+
+    thetaGrid = np.linspace(sTheta, eTheta, num = nTheta)
+    phiGrid = np.linspace(sPhi, ePhi, num = nPhi)
+    rGrid = np.linspace(sR, eR, num = nR)
+
+    return (thetaGrid, phiGrid, rGrid)
+
+def addCheckerboard(spacing = 10., pertubation = 0.1, inputfile = 'vgrids.in',
+                    outputfile = 'vgrids_cb.in', ampmethod = 'linear', rect = (None, None)):
+    '''
+    Add a checkerboard to an existing vgrids.in velocity model.
+
+    :param: spacing, size of the tiles
+    type: float
+
+    :param: pertubation, pertubation (default: 0.1 = 10%)
+    type: float
+    '''
     def correctSpacing(spacing, delta, disttype = None):
         if spacing > delta:
             spacing_corr = round(spacing / delta) * delta
@@ -320,35 +273,24 @@ def addCheckerboard(spacing = 10., pertubation = 0.1, inputfile = 'vgrids.in',
         else:
             print('ampFunc: Could not amplify cb pattern')
 
-
-    R = 6371. # earth radius
     decm = 0.3 # diagonal elements of the covariance matrix (grid3dg's default value is 0.3)
     outfile = open(outputfile, 'w')
 
-    # Theta, Phi in radians, R in km
+    number, delta, start, vel = _readVgrid(inputfile)
-    nR, nTheta, nPhi = readNumberOfPoints(inputfile)
-    dR, dThetaRad, dPhiRad = readDelta(inputfile)
-    sR, sThetaRad, sPhiRad = readStartpoints(inputfile)
-    vel = readVelocity(inputfile)
 
-    dTheta, dPhi = np.rad2deg((dThetaRad, dPhiRad))
+    nR, nTheta, nPhi = number
-    sTheta, sPhi = np.rad2deg((dThetaRad, dPhiRad))
+    dR, dTheta, dPhi = delta
+    sR, sTheta, sPhi = start
     
-    eR = sR + (nR - 1) * dR
+    thetaGrid, phiGrid, rGrid = _generateGrids(number, delta, start)
-    ePhi = sPhi + (nPhi - 1) * dPhi
-    eTheta = sTheta + (nTheta - 1) * dTheta
 
     nPoints = nR * nTheta * nPhi
 
-    thetaGrid = np.linspace(sTheta, eTheta, num = nTheta)
-    phiGrid = np.linspace(sPhi, ePhi, num = nPhi)
-    rGrid = np.linspace(sR, eR, num = nR)
-
     # write header for velocity grid file (in RADIANS)
     outfile.writelines('%10s %10s \n' %(1, 1))
     outfile.writelines('%10s %10s %10s\n' %(nR, nTheta, nPhi))
-    outfile.writelines('%10s %10s %10s\n' %(dR, dThetaRad, dPhiRad))
+    outfile.writelines('%10s %10s %10s\n' %(dR, np.deg2rad(dTheta), np.deg2rad(dPhi)))
-    outfile.writelines('%10s %10s %10s\n' %(sR, sThetaRad, sPhiRad))
+    outfile.writelines('%10s %10s %10s\n' %(sR, np.deg2rad(sTheta), np.deg2rad(sPhi)))
 
     spacR = correctSpacing(spacing, dR, '[meter], R')
     spacTheta = correctSpacing(_getAngle(spacing), dTheta, '[degree], Theta')
@@ -402,6 +344,85 @@ def addCheckerboard(spacing = 10., pertubation = 0.1, inputfile = 'vgrids.in',
           'Outputfile: %s.'%(inputfile, spacing, pertubation, outputfile))
     outfile.close()
 
+def addBox(x = (None, None), y = (None, None), z = (None, None),
+           boxvelocity = 1.0, inputfile = 'vgrids.in',
+           outputfile = 'vgrids_box.in'):
+    '''
+    Add a box with constant velocity to an existing vgrids.in velocity model.
+
+    :param: x, borders of the box (xleft, xright)
+    type: tuple
+
+    :param: y, borders of the box (yleft, yright)
+    type: tuple
+
+    :param: z, borders of the box (bot, top)
+    type: tuple
+
+    :param: boxvelocity, default: 1.0 km/s
+    type: float
+    '''
+    R = 6371.
+    decm = 0.3 # diagonal elements of the covariance matrix (grid3dg's default value is 0.3)
+    outfile = open(outputfile, 'w')
+
+    theta1 = _getAngle(y[0])
+    theta2 = _getAngle(y[1])
+    phi1 = _getAngle(x[0])
+    phi2 = _getAngle(x[1])
+    r1 = R + z[0]
+    r2 = R + z[1]
+
+    print('Adding box to grid with theta = (%s, %s), phi = (%s, %s), '
+          'r = (%s, %s), velocity = %s [km/s]'
+          %(theta1, theta2, phi1, phi2, r1, r2, boxvelocity))
+    
+    number, delta, start, vel = _readVgrid(inputfile)
+
+    nR, nTheta, nPhi = number
+    dR, dTheta, dPhi = delta
+    sR, sTheta, sPhi = start
+    
+    thetaGrid, phiGrid, rGrid = _generateGrids(number, delta, start)
+
+    nPoints = nR * nTheta * nPhi
+
+    # write header for velocity grid file (in RADIANS)
+    outfile.writelines('%10s %10s \n' %(1, 1))
+    outfile.writelines('%10s %10s %10s\n' %(nR, nTheta, nPhi))
+    outfile.writelines('%10s %10s %10s\n' %(dR, np.deg2rad(dTheta), np.deg2rad(dPhi)))
+    outfile.writelines('%10s %10s %10s\n' %(sR, np.deg2rad(sTheta), np.deg2rad(sPhi)))
+
+    count = 0
+    for radius in rGrid:
+        if r1 <= radius <= r2:
+            rFlag = 1
+        else:
+            rFlag = 0
+        for theta in thetaGrid:
+            if theta1 <= theta <= theta2:
+                thetaFlag = 1
+            else:
+                thetaFlag = 0
+            for phi in phiGrid:
+                if phi1 <= phi <= phi2:
+                    phiFlag = 1
+                else:
+                    phiFlag = 0
+                velocity = vel[count]
+                if rFlag * thetaFlag * phiFlag is not 0:
+                    velocity = boxvelocity
+
+                outfile.writelines('%10s %10s\n'%(velocity, decm))
+                count += 1
+
+                progress = float(count) / float(nPoints) * 100
+                _update_progress(progress)
+
+    print('Added box to the grid in file %s. '
+          'Outputfile: %s.'%(inputfile, outputfile))
+    outfile.close()
+
 def _update_progress(progress):
     sys.stdout.write("%d%% done   \r" % (progress) )
     sys.stdout.flush()
@@ -415,3 +436,9 @@ def _getAngle(distance):
     angle = distance * 180. / (PI * R)
     return angle
 
+def _getDistance(angle):
+    PI = np.pi
+    R = 6371.
+    distance = angle / 180 * (PI * R)
+    return distance
+