util: now handle non-arrays

bh107 · May 31, 2018 · 824539a · 824539a
1 parent ff87663
commit 824539a
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 17 deletions.
diff --git a/benchpress/benchmarks/lattice_boltzmann_D2Q9/python_numpy/empty b/benchpress/benchmarks/lattice_boltzmann_D2Q9/python_numpy/empty
diff --git a/benchpress/benchmarks/lattice_boltzmann_D2Q9/python_numpy/issues.rst b/benchpress/benchmarks/lattice_boltzmann_D2Q9/python_numpy/issues.rst
diff --git a/benchpress/benchmarks/lattice_boltzmann_D2Q9/python_numpy/lattice_boltzmann_D2Q9.py b/benchpress/benchmarks/lattice_boltzmann_D2Q9/python_numpy/lattice_boltzmann_D2Q9.py
@@ -37,11 +37,6 @@ def cylinder(height, width, obstacle=True):
 
     col = np.arange(1, ly - 1)
 
-    # We need some of the constants as 3D numpy arrays
-    t_3d = np.asarray(t)[:, np.newaxis, np.newaxis]
-    cx_3d = np.asarray(cx)[:, np.newaxis, np.newaxis]
-    cy_3d = np.asarray(cy)[:, np.newaxis, np.newaxis]
-
     # Place obstacle
     if obstacle:
         y, x = np.meshgrid(np.arange(ly), np.arange(lx))
@@ -51,8 +46,13 @@ def cylinder(height, width, obstacle=True):
     else:
         bbRegion = None
 
+    # We need some of the constants as 3D numpy arrays
+    t_3d = np.asarray(t)[:, np.newaxis, np.newaxis]
+    cx_3d = np.asarray(cx)[:, np.newaxis, np.newaxis]
+    cy_3d = np.asarray(cy)[:, np.newaxis, np.newaxis]
+
     # Initial condition: (rho=0, u=0) ==> fIn[i] = t[i]
-    fIn = t_3d.repeat(lx, axis=1).repeat(ly, axis=2)
+    fIn = np.array(t_3d.repeat(lx, axis=1).repeat(ly, axis=2))
 
     state['fIn'] = fIn
     state['cx_3d'] = cx_3d
@@ -65,8 +65,8 @@ def cylinder(height, width, obstacle=True):
 
 def solve(state, iterations, viz=None):
     # load the state
-    ly = state['ly']
-    lx = state['lx']
+    ly = int(state['ly'])
+    lx = int(state['lx'])
     fIn = state['fIn']
     cx_3d = state['cx_3d']
     cy_3d = state['cy_3d']

diff --git a/benchpress/util.py b/benchpress/util.py
@@ -304,14 +304,15 @@ def stop(self):
         self.__elapsed = time.time() - self.__elapsed
 
     def save_data(self, data_dict):
-        """Save `data_dict` as a npz archive when --outputfn is used """
+        """Save `data_dict` as a npz archive when --outputfn is used"""
         assert(isinstance(data_dict, dict))
         if self.outputfn is not None:
             # Clean `data_dict` for Bohrium arrays
-            nobh_data = {}
+            nobh_data = {"_bhary_keys": []}
             for k in data_dict.keys():
                 if hasattr(data_dict[k], "copy2numpy"):
                     nobh_data[k] = data_dict[k].copy2numpy()
+                    nobh_data["_bhary_keys"].append(k)
                 else:
                     nobh_data[k] = data_dict[k]
             np.savez_compressed(self.outputfn, **nobh_data)
@@ -322,11 +323,15 @@ def load_data(self):
             return None
         else:
             nobh_data = np.load(self.inputfn)
+            bhary_keys = nobh_data["_bhary_keys"]
             ret = {}
-            # Convert numpy arrays into bohrium arrays
             for k in nobh_data.keys():
-                if bh_is_loaded_as_np and hasattr(nobh_data[k], "shape") and hasattr(nobh_data[k], "dtype"):
-                    ret[k] = np.array(nobh_data[k])
+                if k == "_bhary_keys":
+                    continue
+                # Convert numpy arrays into bohrium arrays
+                if bh_is_loaded_as_np and k in bhary_keys:
+                    a = nobh_data[k]
+                    ret[k] = bh.array(a, bohrium=True)
                 else:
                     ret[k] = nobh_data[k]
             return ret