Skip to content

Commit

Permalink
fix MfList dict of df support, raise ValueError instead of Exception …
Browse files Browse the repository at this point in the history 
…where appropriate, add tests
  • Loading branch information
@wpbonelli
wpbonelli committed Nov 21, 2023
1 parent 15a2189 commit 7b70240
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 27 deletions.
88 changes: 84 additions & 4 deletions autotest/test_modflow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1288,7 +1288,87 @@ def test_load_with_list_reader(function_tmpdir):
assert np.array_equal(originalwelra, m2.wel.stress_period_data[0])


def get_basic_modflow_model(ws, name):
@pytest.mark.parametrize(
"container",
[
str(np.recarray),
str(pd.DataFrame),
str(dict[int, np.recarray]),
str(dict[int, pd.DataFrame]),
],
)
def test_pkg_data_containers(function_tmpdir, container):
"""Test various containers for package data (list, ndarray, recarray, dataframe, dict of such)"""

name = "pkg_data"
ws = function_tmpdir
m = Modflow(name, model_ws=ws)
size = 100
nlay = 10
nper = 1
nsfr = int((size**2) / 5)

# grid discretization
dis = ModflowDis(
m,
nper=nper,
nlay=nlay,
nrow=size,
ncol=size,
top=nlay,
botm=list(range(nlay)),
)

# recharge pkg
rch = ModflowRch(
m, rech={k: 0.001 - np.cos(k) * 0.001 for k in range(nper)}
)

# well pkg, setup data per 'container' parameter
# to test support for various container types
ra = ModflowWel.get_empty(size**2)
ra_per = ra.copy()
ra_per["k"] = 1
ra_per["i"] = (
(np.ones((size, size)) * np.arange(size))
.transpose()
.ravel()
.astype(int)
)
ra_per["j"] = list(range(size)) * size
wel_dtype = np.dtype(
[
("k", int),
("i", int),
("j", int),
("flux", np.float32),
]
)
df_per = pd.DataFrame(ra_per)
if "'numpy.recarray'" in container:
well_spd = ra_per
elif "'pandas.core.frame.DataFrame'" in container:
well_spd = df_per
elif "dict[int, numpy.recarray]" in container:
well_spd = {}
well_spd[0] = ra_per
elif "dict[int, pandas.core.frame.DataFrame]" in container:
well_spd = {}
well_spd[0] = df_per
wel = ModflowWel(m, stress_period_data=well_spd)

# streamflow routing pkg
rd = ModflowSfr2.get_empty_reach_data(nsfr)
rd["iseg"] = range(len(rd))
rd["ireach"] = 1
sd = ModflowSfr2.get_empty_segment_data(nsfr)
sd["nseg"] = range(len(sd))
sfr = ModflowSfr2(reach_data=rd, segment_data=sd, model=m)

return m


def get_perftest_model(ws, name):
m = Modflow(name, model_ws=ws)

size = 100
Expand Down Expand Up @@ -1339,20 +1419,20 @@ def get_basic_modflow_model(ws, name):
@pytest.mark.slow
def test_model_init_time(function_tmpdir, benchmark):
name = inspect.getframeinfo(inspect.currentframe()).function
benchmark(lambda: get_basic_modflow_model(ws=function_tmpdir, name=name))
benchmark(lambda: get_perftest_model(ws=function_tmpdir, name=name))


@pytest.mark.slow
def test_model_write_time(function_tmpdir, benchmark):
name = inspect.getframeinfo(inspect.currentframe()).function
model = get_basic_modflow_model(ws=function_tmpdir, name=name)
model = get_perftest_model(ws=function_tmpdir, name=name)
benchmark(lambda: model.write_input())


@pytest.mark.slow
def test_model_load_time(function_tmpdir, benchmark):
name = inspect.getframeinfo(inspect.currentframe()).function
model = get_basic_modflow_model(ws=function_tmpdir, name=name)
model = get_perftest_model(ws=function_tmpdir, name=name)
model.write_input()
benchmark(
lambda: Modflow.load(
Expand Down
42 changes: 19 additions & 23 deletions flopy/utils/util_list.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,6 @@ def __init__(
if package.parent.version == "mf2k":
list_free_format = False
self.list_free_format = list_free_format
return

@property
def name(self):
Expand Down Expand Up @@ -306,41 +305,37 @@ def __cast_data(self, data):
try:
data = np.array(data)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList error: casting list to ndarray: {e!s}"
)

# If data is a dict, the we have to assume it is keyed on kper
if isinstance(data, dict):
if not list(data.keys()):
raise Exception("MfList error: data dict is empty")
raise ValueError("MfList error: data dict is empty")
for kper, d in data.items():
try:
kper = int(kper)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList error: data dict key {kper} not integer: "
f"{type(kper)}\n{e!s}"
)
# Same as before, just try...
if isinstance(d, list):
# warnings.warn("MfList: casting list to array at " +\
# "kper {0:d}".format(kper))
try:
d = np.array(d)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList error: casting list to ndarray: {e}"
)

# super hack - sick of recarrays already
# if (isinstance(d,np.ndarray) and len(d.dtype.fields) > 1):
# d = d.view(np.recarray)

if isinstance(d, np.recarray):
self.__cast_recarray(kper, d)
elif isinstance(d, np.ndarray):
self.__cast_ndarray(kper, d)
elif isinstance(d, pd.DataFrame):
self.__cast_dataframe(kper, d)
elif isinstance(d, int):
self.__cast_int(kper, d)
elif isinstance(d, str):
Expand All @@ -349,7 +344,7 @@ def __cast_data(self, data):
self.__data[kper] = -1
self.__vtype[kper] = None
else:
raise Exception(
raise ValueError(
"MfList error: unsupported data type: "
f"{type(d)} at kper {kper}"
)
Expand All @@ -367,7 +362,7 @@ def __cast_data(self, data):
elif isinstance(data, str):
self.__cast_str(0, data)
else:
raise Exception(
raise ValueError(
f"MfList error: unsupported data type: {type(data)}"
)

Expand All @@ -383,7 +378,7 @@ def __cast_str(self, kper, d):
def __cast_int(self, kper, d):
# If d is an integer, then it must be 0 or -1
if d > 0:
raise Exception(
raise ValueError(
"MfList error: dict integer value for "
"kper {:10d} must be 0 or -1, "
"not {:10d}".format(kper, d)
Expand Down Expand Up @@ -415,7 +410,7 @@ def __cast_ndarray(self, kper, d):
d.transpose(), dtype=self.dtype
)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList error: casting ndarray to recarray: {e!s}"
)
self.__vtype[kper] = np.recarray
Expand Down Expand Up @@ -537,7 +532,7 @@ def add_record(self, kper, index, values):
try:
self.__data[kper][-1] = tuple(rec)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList.add_record() error: adding record to recarray: {e}"
)

Expand All @@ -551,7 +546,7 @@ def __getitem__(self, kper):
try:
kper = int(kper)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList error: _getitem__() passed invalid kper index: {kper}"
)
if kper not in list(self.data.keys()):
Expand Down Expand Up @@ -581,7 +576,7 @@ def __setitem__(self, kper, data):
try:
data = np.array(data)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList error: casting list to ndarray: {e!s}"
)
# cast data
Expand All @@ -596,7 +591,7 @@ def __setitem__(self, kper, data):
elif isinstance(data, str):
self.__cast_str(kper, data)
else:
raise Exception(
raise ValueError(
f"MfList error: unsupported data type: {type(data)}"
)

Expand All @@ -607,7 +602,7 @@ def __fromfile(self, f):
try:
d = np.genfromtxt(f, dtype=self.dtype)
except Exception as e:
raise Exception(
raise ValueError(
f"MfList.__fromfile() error reading recarray from file {e!s}"
)
return d
Expand Down Expand Up @@ -1100,8 +1095,9 @@ def to_array(self, kper=0, mask=False):
for name, arr in arrays.items():
arrays[name][:] = np.NaN
return arrays
else:
raise Exception("MfList: something bad happened")
raise ValueError(
f"MfList: expected no entries for period {kper} but found {sarr}"
)

for name, arr in arrays.items():
if unstructured:
Expand Down Expand Up @@ -1228,7 +1224,7 @@ def masked4D_arrays_to_stress_period_data(dtype, m4ds):
for i2, key2 in enumerate(keys[i1:]):
a2 = np.isnan(m4ds[key2])
if not np.array_equal(a1, a2):
raise Exception(
raise ValueError(
f"Transient2d error: masking not equal for {key1} and {key2}"
)

Expand Down

0 comments on commit 7b70240

Please sign in to comment.