use python style for subarray with ghost cells

Author: wkliao

examples/ghost_cell.py

@@ -66,62 +66,38 @@
 import pnetcdf
 
 
-def parse_help():
-    help_flag = "-h" in sys.argv or "--help" in sys.argv
-    if help_flag and rank == 0:
-        help_text = (
-            "Usage: {} [-h] | [-q] [file_name]\n"
-            "       [-h] Print help\n"
-            "       [-q] Quiet mode (reports when fail)\n"
-            "       [-k format] file format: 1 for CDF-1, 2 for CDF-2, 5 for CDF-5\n"
-            "       [-l len] size of each dimension of the local array\n"
-            "       [filename] (Optional) output netCDF file name\n"
-        ).format(sys.argv[0])
-        print(help_text)
-    return help_flag
-
-
 def pnetcdf_io(filename, file_format, length):
 
     count = [length, length + 1]
     psizes = MPI.Compute_dims(nprocs, 2)
 
+    nghosts = 2
+
     if verbose and rank == 0:
-        print("psizes=", psizes[0], psizes[1])
+        print("number of MPI processes =", nprocs)
+        print("number of ghost cells =", nghosts)
+        print("psizes=", psizes)
 
     # set global array sizes
     gsizes = np.zeros(2, dtype=np.int64)
     gsizes[0] = length * psizes[0]  # global array size
     gsizes[1] = (length + 1) * psizes[1]
     if verbose and rank == 0:
-        print("global variable shape:", gsizes[0], gsizes[1])
+        print("global variable shape:", gsizes)
 
     # find its local rank IDs along each dimension
     local_rank = np.zeros(2, dtype=np.int32)
     local_rank[0] = rank // psizes[1]
     local_rank[1] = rank % psizes[1]
     if verbose:
-        print("rank {}: dim rank= {} {}".format(rank, local_rank[0], local_rank[1]))
+        print("rank ",rank,": local_rank = ", local_rank)
 
     # set subarray access pattern
     count = np.array([length, length + 1], dtype=np.int64)
     start = np.array([local_rank[0] * count[0], local_rank[1] * count[1]],
                       dtype=np.int64)
     if verbose:
-        print("start= {} {} count= {} {}".format(start[0], start[1], count[0], count[1]))
-
-   # allocate and initialize buffer with ghost cells on both ends of each dim
-    nghosts = 2
-    bufsize = (count[0] + 2 * nghosts) * (count[1] + 2 * nghosts)
-    buf = np.empty(bufsize, dtype=np.int32)
-    for i in range(count[0] + 2 * nghosts):
-        for j in range(count[1] + 2 * nghosts):
-            if nghosts <= i < count[0] + nghosts and \
-               nghosts <= j < count[1] + nghosts:
-                buf[i * (count[1] + 2 * nghosts) + j] = rank
-            else:
-                # set values of all ghost cells to -8
-                buf[i * (count[1] + 2 * nghosts) + j] = -8
+        print("rank ",rank,": start = ",start," count =", count)
 
     # Create the file
     f = pnetcdf.File(filename = filename,
@@ -140,19 +116,39 @@ def pnetcdf_io(filename, file_format, length):
     # Exit the define mode
     f.enddef()
 
-    # set imap pattern for local buffer
-    imap = np.zeros(2, dtype=np.int64)
-    imap[1] = 1
-    imap[0] = count[1] + 2 * nghosts
-    buf_ptr = buf[nghosts * (count[1] + 2 * nghosts + 1):]
+    # allocate and initialize buffer with ghost cells on both ends of each dim
+    buf = np.empty([2*nghosts+count[0], 2*nghosts+count[1]], dtype=np.int32)
+    buf.fill(-8)
+
+    # keep contents of ghost cells to -8, all others 'rank'
+    buf[nghosts:nghosts+count[0], nghosts:nghosts+count[1]].fill(rank)
 
     # Write data to the variable
-    var.put_var_all(buf_ptr, start = start, count = count, imap = imap)
+    end = np.add(start, count)
+    var[start[0]:end[0], start[1]:end[1]] = buf[nghosts:nghosts+count[0], nghosts:nghosts+count[1]]
+
+    # Equivalently, below uses function call
+    var.put_var_all(buf[nghosts:nghosts+count[0], nghosts:nghosts+count[1]], start = start, count = count)
 
     # Close the file
     f.close()
 
 
+def parse_help():
+    help_flag = "-h" in sys.argv or "--help" in sys.argv
+    if help_flag and rank == 0:
+        help_text = (
+            "Usage: {} [-h] | [-q] [file_name]\n"
+            "       [-h] Print help\n"
+            "       [-q] Quiet mode (reports when fail)\n"
+            "       [-k format] file format: 1 for CDF-1, 2 for CDF-2, 5 for CDF-5\n"
+            "       [-l len] size of each dimension of the local array\n"
+            "       [filename] (Optional) output netCDF file name\n"
+        ).format(sys.argv[0])
+        print(help_text)
+    return help_flag
+
+
 if __name__ == "__main__":
     comm = MPI.COMM_WORLD
     rank = comm.Get_rank()