Implementation of a more balanced partition

Partition.py

@@ -6,41 +6,65 @@ log_master, log_world = getargs.getLogger(__name__)
 INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
 INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug
 #
-def halfpartition(partin) :
-    partout=[]
-    if partin[0]["nbi"] > partin[0]["nbj"] :
-        for dom in partin :
-            h=int(np.trunc(dom["nbi"]/2))
-            partout.append({"nbi":h,           "nbj":dom["nbj"],"istart":dom["istart"],  "jstart":dom["jstart"]})
-            partout.append({"nbi":dom["nbi"]-h,"nbj":dom["nbj"],"istart":dom["istart"]+h,"jstart":dom["jstart"]})
+def halfpartition(partin, land) :
+    """
+    First partition of the domain
+    The domain with more land points is halved
+    """
+    A = [p["nbland"] for p in partin]
+    i = np.argmax(A)
+    partout = []
+    partout = [partin[k] for k in range(len(partin)) if k != i]
+
+    dom = partin[i]
+
+    if dom["nbi"] > dom["nbj"] :
+        hh = [h for h in range(dom["nbi"])]
+        nb = np.array([np.ma.sum(np.ma.filled(land[dom["jstart"]:dom["jstart"]+dom["nbj"],dom["istart"]:dom["istart"]+h],0)) - dom['nbland']/2. for h in hh])
+        i = np.nanargmin(np.abs(nb))
+        new_dom = {"nbi":dom["nbi"]-hh[i],"nbj":dom["nbj"],  \
+                   "istart":dom["istart"]+hh[i],"jstart":dom["jstart"]}
+        dom["nbi"] = hh[i]
+        
     else :
-        for dom in partin :
-            h=int(np.trunc(dom["nbj"]/2))
-            partout.append({"nbi":dom["nbi"],"nbj":h,           "istart":dom["istart"],"jstart":dom["jstart"]})
-            partout.append({"nbi":dom["nbi"],"nbj":dom["nbj"]-h,"istart":dom["istart"],"jstart":dom["jstart"]+h})
+        hh = [h for h in range(dom["nbj"])]
+        nb = np.array([np.ma.sum(np.ma.filled(land[dom["jstart"]:dom["jstart"]+h,dom["istart"]:dom["istart"]+dom["nbi"]],0)) - int(dom['nbland']/2.) for h in hh])
+        i = np.nanargmin(np.abs(nb))
+        new_dom = {"nbi":dom["nbi"],"nbj":dom["nbj"]-hh[i], \
+                 "istart":dom["istart"],"jstart":dom["jstart"]+hh[i]}
+        dom["nbj"] = hh[i]
+
+    new_dom['nbland'] = int(np.nansum(land[new_dom["jstart"]:new_dom["jstart"]+new_dom["nbj"],new_dom["istart"]:new_dom["istart"]+new_dom["nbi"]]))
+    dom['nbland'] = dom['nbland']-new_dom['nbland']
+
+    partout.append(dom)
+    partout.append(new_dom)
+
     return partout
 #
-# Add number of land points to partitions
-#
-def addnbland(partin, land) :
-    for dom in partin :
-        dom['nbland']=int(np.sum(land[dom["jstart"]:dom["jstart"]+dom["nbj"],dom["istart"]:dom["istart"]+dom["nbi"]]))
-#
 # Divide in two the domain with the most number of land points
 #
-def halflargest(part) :
-    nbland=np.array([dom['nbland'] for dom in part])
-    imax = np.argmax(nbland)
-    if part[imax]["nbi"] > part[imax]["nbj"] :
-        h=int(np.trunc(part[imax]["nbi"]/2))
-        part.append({"nbi":part[imax]["nbi"]-h,"nbj":part[imax]["nbj"],"istart":part[imax]["istart"]+h,"jstart":part[imax]["jstart"]})
-        part[imax]["nbi"]=h
-    else :
-        h=int(np.trunc(part[imax]["nbj"]/2))
-        part.append({"nbi":part[imax]["nbi"],"nbj":part[imax]["nbj"]-h,"istart":part[imax]["istart"],"jstart":part[imax]["jstart"]+h})
-        part[imax]["nbj"]=h
+def fit_partition(partin, land):
+    partout = []
+    for i, dom in enumerate(partin):
+        l = land[dom["jstart"]:dom["jstart"]+dom["nbj"],dom["istart"]:dom["istart"]+dom["nbi"]]
+        l1 = np.sum(l, axis = 0)
+        
+        i0 = np.where(l1>0)[0][0]
+        i1 = np.where(l1>0)[0][-1]
 
-    return
+        l2 = np.ma.sum(l, axis = 1)
+        j0 = np.where(l2>0)[0][0]
+        j1 = np.where(l2>0)[0][-1]
+
+        dom["jstart"] = j0 + dom["jstart"] 
+        dom["nbj"] = j1-j0+1
+        dom["istart"] = i0 + dom["istart"] 
+        dom["nbi"] = i1-i0+1
+        dom["nbland"] = int(np.nansum(land[dom["jstart"]:dom["jstart"]+dom["nbj"],dom["istart"]:dom["istart"]+dom["nbi"]]))
+
+        partout.append(dom)        
+    return partout
 #
 #
 #
@@ -56,30 +80,37 @@ def adjustpart(partin, land, nbcore) :
     #
     # Divide domains with the largest number of land points
     #
+    nbptproc=np.array([dom['nbi']*dom['nbj'] for dom in partin])
     freecores = nbcore-len(partin)
-    while nbcore-len(partin) > 0 :
-        halflargest(partin)
-        addnbland(partin, land)
-        
+
+    while nbcore-len(partin) > 0:
+        partin = halfpartition(partin, land)
+
+    partin = fit_partition(partin, land)
+
     nbptproc=np.array([dom['nbi']*dom['nbj'] for dom in partin])
+    nbland = [dom["nbland"] for dom in partin]
+
     if np.min(nbptproc) < 3 :
         ERROR("The smallest domain in the partition is of a size less than 5. This will probably not work.")
         sys.exit()        
-    return
+
+    return partin
 #
 # Partition domain in two steps : 1) halving, 2) adjusting by nb land points
 #
 def partitiondom(land, nbcore, rank) :
     njg,nig = land.shape
     partout=[{"nbi":nig,"nbj":njg,"istart":0,"jstart":0, "nbland":int(np.sum(land))}]
-    while len(partout) <= nbcore/2 :
-        partout = halfpartition(partout)
-    addnbland(partout, land)
+    while len(partout) < nbcore :
+        partout = halfpartition(partout, land)
     #
-    adjustpart(partout, land, nbcore)
+    partout = adjustpart(partout, land, nbcore)
     #
     gindland=np.zeros((njg,nig), dtype=np.int32)
     gindland[:,:]=-1
+    test = np.zeros((njg,nig), dtype=np.int32)
+    test[:,:] = land[:,:]
     n=0
     for i in range(nig) :
         for j in range(njg) :
@@ -95,10 +126,17 @@ def partitiondom(land, nbcore, rank) :
         locgind=np.copy(gindland[jstart:jstart+nbj,istart:istart+nbi])
         loccore=np.copy(gindland[jstart:jstart+nbj,istart:istart+nbi])
         loccore[loccore >= 0] = ic
+        test[jstart:jstart+nbj,istart:istart+nbi] = test[jstart:jstart+nbj,istart:istart+nbi] - land[jstart:jstart+nbj,istart:istart+nbi]
         partout[ic]["glandindex"] = locgind
         partout[ic]["corenb"] = loccore
         partout[ic]["nblocland"] = np.count_nonzero(locgind >= 0)
     #
+    nbland = [part["nbland"] for part in partout]
+    if rank == 0:
+        INFO("Distribution of the domain : {}".format(nbland))
+    INFO("Test number of error on domain distribution: {0}".format(len(np.where(test == 1)[0])))
+    INFO("Test domain partition (only value must be 0): {0}".format(np.unique(test)))
+    #
     return partout
 #
 # Build global map of processors