fix(mflist): fixed record shape in get_empty; test add_record

modflowpy · Dec 12, 2019 · e0d6544 · e0d6544
1 parent 14b0ddf
commit e0d6544
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 9 deletions.
diff --git a/autotest/t021_test.py b/autotest/t021_test.py
@@ -48,6 +48,32 @@ def test_single_mflist_entry_load():
     assert mm.wel.stress_period_data
     mm.write_input()
 
+
+def test_mflist_add_record():
+    ml = flopy.modflow.Modflow()
+    _ = flopy.modflow.ModflowDis(ml, nper=2)
+    wel = flopy.modflow.ModflowWel(ml)
+    assert len(wel.stress_period_data.data) == 0
+
+    wel.stress_period_data.add_record(0, [0, 1, 2], [1.0])
+    assert len(wel.stress_period_data.data) == 1
+    wel_dtype = [('k', int), ('i', int), ('j', int), ('flux', np.float32)]
+    check0 = np.array([(0, 1, 2, 1.0)], dtype=wel_dtype)
+    np.testing.assert_array_equal(wel.stress_period_data[0], check0)
+
+    wel.stress_period_data.add_record(0, [0, 1, 1], [8.0])
+    assert len(wel.stress_period_data.data) == 1
+    check0 = np.array([(0, 1, 2, 1.0), (0, 1, 1, 8.0)], dtype=wel_dtype)
+    np.testing.assert_array_equal(wel.stress_period_data[0], check0)
+
+    wel.stress_period_data.add_record(1, [0, 1, 1], [5.0])
+    assert len(wel.stress_period_data.data) == 2
+    check1 = np.array([(0, 1, 1, 5.0)], dtype=wel_dtype)
+    np.testing.assert_array_equal(wel.stress_period_data[0], check0)
+    np.testing.assert_array_equal(wel.stress_period_data[1], check1)
+
+
 if __name__ == '__main__':
     test_mflist_external()
     test_single_mflist_entry_load()
+    test_mflist_add_record()
diff --git a/flopy/utils/util_list.py b/flopy/utils/util_list.py
@@ -122,8 +122,8 @@ def plotable(self):
         return True
 
     def get_empty(self, ncell=0):
-        d = np.zeros((ncell, len(self.dtype)), dtype=self.dtype)
-        d[:, :] = -1.0E+10
+        d = np.zeros(ncell, dtype=self.dtype)
+        d[:] = -1.0E+10
         return d
 
     def export(self, f, **kwargs):
@@ -499,20 +499,20 @@ def add_record(self, kper, index, values):
             "length of value arg != length of self dtype"
         # If we already have something for this kper, then add to it
         if kper in list(self.__data.keys()):
-            # If a 0 or -1, reset
             if self.vtype[kper] == int:
+                # If a 0 or -1, reset
                 self.__data[kper] = self.get_empty(1)
                 self.__vtype[kper] = np.recarray
-            # If filename, load into recarray
-            if self.vtype[kper] == str:
+            elif self.vtype[kper] == str:
+                # If filename, load into recarray
                 d = self.__fromfile(self.data[kper])
                 d.resize(d.shape[0], d.shape[1])
                 self.__data[kper] = d
                 self.__vtype[kper] = np.recarray
-            # Extend the recarray
-            if self.vtype[kper] == np.recarray:
-                shape = self.__data[kper].shape
-                self.__data[kper].resize(shape[0] + 1, shape[1])
+            elif self.vtype[kper] == np.recarray:
+                # Extend the recarray
+                self.__data[kper] = np.append(
+                    self.__data[kper], self.get_empty(1))
         else:
             self.__data[kper] = self.get_empty(1)
             self.__vtype[kper] = np.recarray