Merge branch 'master' into remove-slow-cuda-tests

pftools/python/parflow/tools/core.py

@@ -497,7 +497,7 @@ def dist(self, pfb_file, **kwargs):
         q = kwargs.get('Q', self.Process.Topology.Q)
         r = kwargs.get('R', self.Process.Topology.R)
 
-        with ParflowBinaryReader(pfb_file_full_path) as pfb:
+        with ParflowBinaryReader(pfb_file_full_path, read_sg_info=True) as pfb:
             array = pfb.read_all_subgrids()
             header = pfb.header
 

pftools/python/parflow/tools/io.py

@@ -401,22 +401,7 @@ def __init__(
             self.header['q'] = q
             self.header['r'] = r
 
-        if not ('p' in self.header
-            and 'q' in self.header
-            and 'r' in self.header):
-            # If p, q, and r aren't given we can precompute them
-            # NOTE: This is a bit of a fallback and may not always work
-            eps = 1 - 1e-6
-            first_sg_head = self.read_subgrid_header()
-            self.header['p'] = int((self.header['nx'] / first_sg_head['nx']) + eps)
-            self.header['q'] = int((self.header['ny'] / first_sg_head['ny']) + eps)
-            self.header['r'] = int((self.header['nz'] / first_sg_head['nz']) + eps)
-
-        if precompute_subgrid_info:
-            self.compute_subgrid_info()
-
-        if read_sg_info:
-            self.read_subgrid_info()
+        self.read_subgrid_info()
 
     def close(self):
         self.f.close()
@@ -465,15 +450,43 @@ def read_subgrid_info(self):
             sg_starts.append([sg_head['ix'], sg_head['iy'], sg_head['iz']])
             sg_shapes.append([sg_head['nx'], sg_head['ny'], sg_head['nz']])
             sg_offs.append(off)
+
+            # Finally, move past the current subgrid before next iteration
+            off += sg_head['sg_size']*8
 
+        # Calculate p, q, r from subgrid shapes
+        p, q, r = 0, 0, 0
+        x, y, z = 0, 0, 0
+
+        for shape in sg_shapes:
+            if x == self.header['nx']:
+                break
+            p = p + 1
+            x = x + shape[0]
+
+        for shape in sg_shapes[::p]:
+            if y == self.header['ny']:
+                break
+            q = q + 1
+            y = y + shape[1]
+
+        for shape in sg_shapes[::p * q]:
+            if z == self.header['nz']:
+                break
+            r = r + 1
+            z = z + shape[2]
+
+        self.header['p'] = p
+        self.header['q'] = q
+        self.header['r'] = r
+
+        for sg_num in range(self.header['n_subgrids']):            
             # Calculate subgrid locs instead of reading from file
             sg_p, sg_q, sg_r = get_subgrid_loc(sg_num, self.header['p'], 
                                                        self.header['q'], 
                                                        self.header['r'])
             sg_locs.append((sg_p, sg_q, sg_r))
-
-            # Finally, move past the current subgrid before next iteration
-            off += sg_head['sg_size']*8
+
 
         self.subgrid_offsets = np.array(sg_offs)
         self.subgrid_locations = np.array(sg_locs)

test/python/CMakeLists.txt

@@ -17,6 +17,7 @@ set(TESTS
   bc_flux_file
   Dirichlet
   LW_surface_press
+  overland_tiltedV_MGSemiTest
 )
 
 if(${PARFLOW_HAVE_HYPRE})

test/python/new_features/CMakeLists.txt

@@ -17,6 +17,7 @@ set(TESTS
   normalize_location
   os_function
   # pfb_mask
+  pfb_reader
   pfidb_pfset
   pfset_test
   selection_methods

test/python/new_features/pfb_reader.py

@@ -0,0 +1,27 @@
+import os
+import sys
+import numpy as np
+from parflow import Run
+from parflow.tools.io import read_pfb, write_pfb, ParflowBinaryReader
+from parflow.tools.fs import mkdir, get_absolute_path
+
+# Check that the ParflowBinaryReader reads the correct values of p, q, r even
+# when they don't exactly divide nx, ny, nz.
+
+data = np.ones((89, 91))
+run = Run("test_dist", __file__)
+working_dir = get_absolute_path(os.path.join("test_output", "dist"))
+mkdir(working_dir)
+
+path = os.path.join(working_dir, 'data.pfb')
+write_pfb(path, data)
+
+run.dist(path, P=15, Q=15)
+rd = ParflowBinaryReader(path, read_sg_info=True)
+if rd.header['p'] != 15 or rd.header['q'] != 15:
+    sys.exit(1)
+
+run.dist(path, P=14, Q=15)
+rd = ParflowBinaryReader(path, read_sg_info=True)
+if rd.header['p'] != 14 or rd.header['q'] != 15:
+    sys.exit(1)