diff --git a/autotest/t021_test.py b/autotest/t021_test.py index 47696e62f4..0ede027098 100644 --- 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 index 9c73a4b750..1abc0a4292 100644 --- 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