Calculate p, q, r explicitly from subgrid shapes

gartavanis · Jul 15, 2024 · 8ad8b94 · 8ad8b94
1 parent 40c1348
commit 8ad8b94
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/pftools/python/parflow/tools/io.py b/pftools/python/parflow/tools/io.py
@@ -465,15 +465,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)

diff --git a/test/python/new_features/pfb_reader.py b/test/python/new_features/pfb_reader.py
@@ -1,4 +1,5 @@
 import os
+import sys
 import numpy as np
 from parflow import Run
 from parflow.tools.io import read_pfb, write_pfb, ParflowBinaryReader
@@ -16,11 +17,11 @@
 write_pfb(path, data)
 
 run.dist(path, P=15, Q=15)
-rd = ParflowBinaryReader(path)
+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)
+rd = ParflowBinaryReader(path, read_sg_info=True)
 if rd.header['p'] != 14 or rd.header['q'] != 15:
     sys.exit(1)