parallel processing bursts

contrib/stack/topsStack/topo.py

@@ -10,6 +10,8 @@ import sys
 import s1a_isce_utils as ut
 from isceobj.Planet.Planet import Planet
 from zerodop.topozero import createTopozero
+import multiprocessing as mp
+
 
 def createParser():
     parser = argparse.ArgumentParser( description='Generates lat/lon/h and los for each pixel')
@@ -30,43 +32,22 @@ def cmdLineParse(iargs=None):
     return parser.parse_args(args=iargs)
 
 
-def main(iargs=None):
+def call_topo(input):
 
-    inps = cmdLineParse(iargs)
+    (dirname, demImage, master, ind) = input
 
-    swathList = ut.getSwathList(inps.master)
-
-    demImage = isceobj.createDemImage()
-    demImage.load(inps.dem + '.xml')
-
-    boxes = []
-    for swath in swathList:
-        master =  ut.loadProduct(os.path.join(inps.master , 'IW{0}.xml'.format(swath)))
-    
-        ###Check if geometry directory already exists.
-        dirname = os.path.join(inps.geom_masterDir, 'IW{0}'.format(swath))
-
-        if os.path.isdir(dirname):
-            print('Geometry directory {0} already exists.'.format(dirname))
-        else:
-            os.makedirs(dirname)
-
-
-    ###For each burst
-        for ind in range(master.numberOfBursts):
     burst = master.bursts[ind]
-
+    latname = os.path.join(dirname, 'lat_%02d.rdr' % (ind + 1))
-            latname = os.path.join(dirname, 'lat_%02d.rdr'%(ind+1))
+    lonname = os.path.join(dirname, 'lon_%02d.rdr' % (ind + 1))
-            lonname = os.path.join(dirname, 'lon_%02d.rdr'%(ind+1))
+    hgtname = os.path.join(dirname, 'hgt_%02d.rdr' % (ind + 1))
-            hgtname = os.path.join(dirname, 'hgt_%02d.rdr'%(ind+1))
+    losname = os.path.join(dirname, 'los_%02d.rdr' % (ind + 1))
-            losname = os.path.join(dirname, 'los_%02d.rdr'%(ind+1))
+    maskname = os.path.join(dirname, 'shadowMask_%02d.rdr' % (ind + 1))
-            maskname = os.path.join(dirname, 'shadowMask_%02d.rdr'%(ind+1))
+    incname = os.path.join(dirname, 'incLocal_%02d.rdr' % (ind + 1))
-            incname = os.path.join(dirname, 'incLocal_%02d.rdr'%(ind+1))
     #####Run Topo
     planet = Planet(pname='Earth')
     topo = createTopozero()
     topo.slantRangePixelSpacing = burst.rangePixelSize
-            topo.prf = 1.0/burst.azimuthTimeInterval
+    topo.prf = 1.0 / burst.azimuthTimeInterval
     topo.radarWavelength = burst.radarWavelength
     topo.orbit = burst.orbit
     topo.width = burst.numberOfSamples
@@ -78,7 +59,7 @@ def main(iargs=None):
     topo.lookSide = -1
     topo.sensingStart = burst.sensingStart
     topo.rangeFirstSample = burst.startingRange
-            topo.demInterpolationMethod='BIQUINTIC'
+    topo.demInterpolationMethod = 'BIQUINTIC'
     topo.latFilename = latname
     topo.lonFilename = lonname
     topo.heightFilename = hgtname
@@ -88,14 +69,50 @@ def main(iargs=None):
     topo.topo()
 
     bbox = [topo.minimumLatitude, topo.maximumLatitude, topo.minimumLongitude, topo.maximumLongitude]
-            boxes.append(bbox)
 
     topo = None
 
+    return bbox
+
+
+def main(iargs=None):
+
+    inps = cmdLineParse(iargs)
+
+    swathList = ut.getSwathList(inps.master)
+
+    demImage = isceobj.createDemImage()
+    demImage.load(inps.dem + '.xml')
+
+    boxes = []
+    inputs = []
+
+    for swath in swathList:
+        master =  ut.loadProduct(os.path.join(inps.master , 'IW{0}.xml'.format(swath)))
+    
+        ###Check if geometry directory already exists.
+        dirname = os.path.join(inps.geom_masterDir, 'IW{0}'.format(swath))
+
+        if os.path.isdir(dirname):
+            print('Geometry directory {0} already exists.'.format(dirname))
+        else:
+            os.makedirs(dirname)
+
+        for ind in range(master.numberOfBursts):
+            inputs.append((dirname, demImage, master, ind))
+
+    pool = mp.Pool(mp.cpu_count())
+    results = pool.map(call_topo, inputs)
+    pool.close()
+
+    for bbox in results:
+        boxes.append(bbox)
+
     boxes = np.array(boxes)
     bbox = [np.min(boxes[:,0]), np.max(boxes[:,1]), np.min(boxes[:,2]), np.max(boxes[:,3])]
     print ('bbox : ',bbox)
     
+
 if __name__ == '__main__':
 
     main()