Skip to content

Commit

Permalink
feat(modflow): support dataframe for pkg data
Browse files Browse the repository at this point in the history
  • Loading branch information
wpbonelli committed Nov 20, 2023
1 parent 7cc5e6f commit 9637bc6
Show file tree
Hide file tree
Showing 17 changed files with 76 additions and 33 deletions.
22 changes: 17 additions & 5 deletions autotest/test_mnw.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ def mnw1_path(example_data_path):
return example_data_path / "mf2005_test"


def test_load(function_tmpdir, example_data_path, mnw2_examples_path):
def test_load(function_tmpdir, mnw2_examples_path):
"""t027 test load of MNW2 Package"""
# load in the test problem (1 well, 3 stress periods)
m = Modflow.load(
Expand Down Expand Up @@ -94,7 +94,8 @@ def test_mnw1_load_write(function_tmpdir, mnw1_path):
assert np.array_equal(v, m2.mnw1.stress_period_data[k])


def test_make_package(function_tmpdir):
@pytest.mark.parametrize("dataframe", [True, False])
def test_make_package(function_tmpdir, dataframe):
"""t027 test make MNW2 Package"""
ws = function_tmpdir
m4 = Modflow("mnw2example", model_ws=ws)
Expand Down Expand Up @@ -195,6 +196,9 @@ def test_make_package(function_tmpdir):
).view(np.recarray),
}

if dataframe:
node_data = pd.DataFrame(node_data)

mnw2_4 = ModflowMnw2(
model=m4,
mnwmax=2,
Expand Down Expand Up @@ -257,6 +261,9 @@ def test_make_package(function_tmpdir):
).view(np.recarray),
}

if dataframe:
node_data = pd.DataFrame(node_data)

mnw2_4 = ModflowMnw2(
model=m4,
mnwmax=2,
Expand Down Expand Up @@ -294,7 +301,8 @@ def test_make_package(function_tmpdir):
)


def test_mnw2_create_file(function_tmpdir):
@pytest.mark.parametrize("dataframe", [True, False])
def test_mnw2_create_file(function_tmpdir, dataframe):
"""
Test for issue #556, Mnw2 crashed if wells have
multiple node lengths
Expand Down Expand Up @@ -341,8 +349,12 @@ def test_mnw2_create_file(function_tmpdir):
wellids[i],
nnodes=nlayers[i],
nper=len(stress_period_data.index),
node_data=node_data.to_records(index=False),
stress_period_data=stress_period_data.to_records(index=False),
node_data=node_data.to_records(index=False)
if not dataframe
else node_data,
stress_period_data=stress_period_data.to_records(index=False)
if not dataframe
else stress_period_data,
)

wells.append(wl)
Expand Down
1 change: 1 addition & 0 deletions autotest/test_modflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from pathlib import Path

import numpy as np
import pandas as pd
import pytest
from autotest.conftest import get_example_data_path
from modflow_devtools.markers import excludes_platform, requires_exe
Expand Down
17 changes: 13 additions & 4 deletions flopy/modflow/mfag.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import os

import numpy as np
import pandas as pd

from ..pakbase import Package
from ..utils.flopy_io import multi_line_strip
Expand All @@ -29,9 +30,9 @@ class ModflowAg(Package):
model object
options : flopy.utils.OptionBlock object
option block object
time_series : np.recarray
time_series : np.recarray or pd.DataFrame
numpy recarray for the time series block
well_list : np.recarray
well_list : np.recarray or pd.DataFrame
recarray of the well_list block
irrdiversion : dict {per: np.recarray}
dictionary of the irrdiversion block
Expand Down Expand Up @@ -269,8 +270,16 @@ def __init__(
else:
self.options = OptionBlock("", ModflowAg)

self.time_series = time_series
self.well_list = well_list
self.time_series = (
time_series.to_records(index=False)
if isinstance(time_series, pd.DataFrame)
else time_series
)
self.well_list = (
well_list.to_records(index=False)
if isinstance(well_list, pd.DataFrame)
else well_list
)
self.irrdiversion = irrdiversion
self.irrwell = irrwell
self.supwell = supwell
Expand Down
3 changes: 1 addition & 2 deletions flopy/modflow/mfchd.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,7 @@ class ModflowChd(Package):
model : model object
The model object (of type :class:`flopy.modflow.mf.Modflow`) to which
this package will be added.
stress_period_data : list of boundaries, recarrays, or dictionary of
boundaries.
stress_period_data : list, recarray, dataframe, or dictionary of boundaries.
Each chd cell is defined through definition of
layer (int), row (int), column (int), shead (float), ehead (float)
Expand Down
3 changes: 1 addition & 2 deletions flopy/modflow/mfdrn.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@ class ModflowDrn(Package):
A flag that is used to determine if cell-by-cell budget data should be
saved. If ipakcb is non-zero cell-by-cell budget data will be saved.
(default is None).
stress_period_data : list of boundaries, recarrays, or dictionary of
boundaries.
stress_period_data : list, recarray, dataframe or dictionary of boundaries.
Each drain cell is defined through definition of
layer(int), row(int), column(int), elevation(float),
conductance(float).
Expand Down
3 changes: 1 addition & 2 deletions flopy/modflow/mfdrt.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@ class ModflowDrt(Package):
A flag that is used to determine if cell-by-cell budget data should be
saved. If ipakcb is non-zero cell-by-cell budget data will be saved.
(default is None).
stress_period_data : list of boundaries, recarrays, or dictionary of
boundaries.
stress_period_data : list, recarray, dataframe or dictionary of boundaries.
Each drain return cell is defined through definition of
layer(int), row(int), column(int), elevation(float),
conductance(float), layerR (int) , rowR (int), colR (int) and rfprop (float).
Expand Down
9 changes: 7 additions & 2 deletions flopy/modflow/mffhb.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
"""
import numpy as np
import pandas as pd

from ..pakbase import Package
from ..utils import read1d
Expand Down Expand Up @@ -64,7 +65,7 @@ class ModflowFhb(Package):
(default is 0.0)
cnstm5 : float
A constant multiplier for data list flwrat. (default is 1.0)
ds5 : list or numpy array or recarray
ds5 : list or numpy array or recarray or pandas dataframe
Each FHB flwrat cell (dataset 5) is defined through definition of
layer(int), row(int), column(int), iaux(int), flwrat[nbdtime](float).
There should be nflw entries. (default is None)
Expand All @@ -81,7 +82,7 @@ class ModflowFhb(Package):
cnstm7 : float
A constant multiplier for data list sbhedt. (default is 1.0)
ds7 : list or numpy array or recarray
ds7 : list or numpy array or recarray or pandas dataframe
Each FHB sbhed cell (dataset 7) is defined through definition of
layer(int), row(int), column(int), iaux(int), sbhed[nbdtime](float).
There should be nhed entries. (default is None)
Expand Down Expand Up @@ -211,6 +212,8 @@ def __init__(
raise TypeError(msg)
elif isinstance(ds5, list):
ds5 = np.array(ds5)
elif isinstance(ds5, pd.DataFrame):
ds5 = ds5.to_records(index=False)
# convert numpy array to a recarray
if ds5.dtype != dtype:
ds5 = np.core.records.fromarrays(ds5.transpose(), dtype=dtype)
Expand All @@ -228,6 +231,8 @@ def __init__(
raise TypeError(msg)
elif isinstance(ds7, list):
ds7 = np.array(ds7)
elif isinstance(ds7, pd.DataFrame):
ds7 = ds7.to_records(index=False)
# convert numpy array to a recarray
if ds7.dtype != dtype:
ds7 = np.core.records.fromarrays(ds7.transpose(), dtype=dtype)
Expand Down
5 changes: 4 additions & 1 deletion flopy/modflow/mfgage.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import os

import numpy as np
import pandas as pd

from ..pakbase import Package
from ..utils import read_fixed_var, write_fixed_var
Expand All @@ -27,7 +28,7 @@ class ModflowGage(Package):
this package will be added.
numgage : int
The total number of gages included in the gage file (default is 0).
gage_data : list or numpy array
gage_data : list or numpy array or recarray or pandas dataframe
data for dataset 2a and 2b in the gage package. If a list is provided
then the list includes 2 to 3 entries (LAKE UNIT [OUTTYPE]) for each
LAK Package entry and 4 entries (GAGESEG GAGERCH UNIT OUTTYPE) for
Expand Down Expand Up @@ -132,6 +133,8 @@ def __init__(
gage_data = np.core.records.fromarrays(
gage_data.transpose(), dtype=dtype
)
elif isinstance(gage_data, pd.DataFrame):
gage_data = gage_data.to_records(index=False)
elif isinstance(gage_data, list):
d = ModflowGage.get_empty(ncells=numgage)
for n in range(len(gage_data)):
Expand Down
3 changes: 1 addition & 2 deletions flopy/modflow/mfghb.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@ class ModflowGhb(Package):
A flag that is used to determine if cell-by-cell budget data should be
saved. If ipakcb is non-zero cell-by-cell budget data will be saved.
(default is 0).
stress_period_data : list of boundaries, recarray of boundaries or,
dictionary of boundaries.
stress_period_data : list, recarray, dataframe or dictionary of boundaries.
Each ghb cell is defined through definition of
layer(int), row(int), column(int), stage(float), conductance(float)
Expand Down
5 changes: 4 additions & 1 deletion flopy/modflow/mfhyd.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
"""
import numpy as np
import pandas as pd

from ..pakbase import Package
from ..utils.recarray_utils import create_empty_recarray
Expand All @@ -31,7 +32,7 @@ class ModflowHyd(Package):
is a user-specified value that is output if a value cannot be computed
at a hydrograph location. For example, the cell in which the hydrograph
is located may be a no-flow cell. (default is -999.)
obsdata : list of lists, numpy array, or numpy recarray (nhyd, 7)
obsdata : list of lists, numpy array or recarray, or pandas dataframe (nhyd, 7)
Each row of obsdata includes data defining pckg (3 character string),
arr (2 character string), intyp (1 character string) klay (int),
xl (float), yl (float), hydlbl (14 character string) for each
Expand Down Expand Up @@ -158,6 +159,8 @@ def __init__(

dtype = ModflowHyd.get_default_dtype()
obs = ModflowHyd.get_empty(nhyd)
if isinstance(obsdata, pd.dataFrame):
obsdata = obsdata.to_records(index=False)
if isinstance(obsdata, list):
if len(obsdata) != nhyd:
raise RuntimeError(
Expand Down
7 changes: 7 additions & 0 deletions flopy/modflow/mfmnw2.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import warnings

import numpy as np
import pandas as pd

from ..pakbase import Package
from ..utils import MfList, check
Expand Down Expand Up @@ -451,6 +452,8 @@ def __init__(
# does this need to be Mflist?
self.stress_period_data = self.get_empty_stress_period_data(nper)
if stress_period_data is not None:
if isinstance(stress_period_data, pd.DataFrame):
stress_period_data = stress_period_data.to_records(index=False)
for n in stress_period_data.dtype.names:
self.stress_period_data[n] = stress_period_data[n]

Expand All @@ -459,6 +462,8 @@ def __init__(
np.abs(nnodes), aux_names=self.aux
)
if node_data is not None:
if isinstance(node_data, pd.DataFrame):
node_data = node_data.to_records(index=False)
for n in node_data.dtype.names:
self.node_data[n] = node_data[n]
# convert strings to lower case
Expand Down Expand Up @@ -1054,6 +1059,8 @@ def __init__(
self.node_data = self.get_empty_node_data(0, aux_names=aux)

if node_data is not None:
if isinstance(node_data, pd.DataFrame):
node_data = node_data.to_records(index=False)
self.node_data = self.get_empty_node_data(
len(node_data), aux_names=aux
)
Expand Down
3 changes: 1 addition & 2 deletions flopy/modflow/mfriv.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@ class ModflowRiv(Package):
A flag that is used to determine if cell-by-cell budget data should be
saved. If ipakcb is non-zero cell-by-cell budget data will be saved.
(default is 0).
stress_period_data : list of boundaries, or recarray of boundaries, or
dictionary of boundaries.
stress_period_data : list, recarray, dataframe, or dictionary of boundaries.
Each river cell is defined through definition of
layer (int), row (int), column (int), stage (float), cond (float),
rbot (float).
Expand Down
8 changes: 6 additions & 2 deletions flopy/modflow/mfsfr2.py
Original file line number Diff line number Diff line change
Expand Up @@ -172,13 +172,13 @@ class ModflowSfr2(Package):
simulations (and would need to be converted to whatever units are being
used in the particular simulation). (default is 0.0001; for
MODFLOW-2005 simulations only when irtflg > 0)
reach_data : recarray
reach_data : recarray or dataframe
Numpy record array of length equal to nstrm, with columns for each
variable entered in item 2 (see SFR package input instructions). In
following flopy convention, layer, row, column and node number
(for unstructured grids) are zero-based; segment and reach are
one-based.
segment_data : recarray
segment_data : recarray or dataframe
Numpy record array of length equal to nss, with columns for each
variable entered in items 6a, 6b and 6c (see SFR package input
instructions). Segment numbers are one-based.
Expand Down Expand Up @@ -431,6 +431,8 @@ def __init__(
)
if segment_data is not None:
# segment_data is a zero-d array
if isinstance(segment_data, pd.DataFrame):
segment_data = segment_data.to_records(index=False)
if not isinstance(segment_data, dict):
if len(segment_data.shape) == 0:
segment_data = np.atleast_1d(segment_data)
Expand Down Expand Up @@ -479,6 +481,8 @@ def __init__(
# Dataset 2.
self.reach_data = self.get_empty_reach_data(np.abs(self._nstrm))
if reach_data is not None:
if isinstance(reach_data, pd.DataFrame):
reach_data = reach_data.to_records(index=False)
for n in reach_data.dtype.names:
self.reach_data[n] = reach_data[n]

Expand Down
7 changes: 5 additions & 2 deletions flopy/modflow/mfstr.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
"""
import numpy as np
import pandas as pd

from ..pakbase import Package
from ..utils import MfList, read_fixed_var, write_fixed_var
Expand Down Expand Up @@ -81,8 +82,8 @@ class ModflowStr(Package):
datasets 6 and 8.
The value for stress period data for a stress period can be an integer
(-1 or 0), a list of lists, a numpy array, or a numpy recarray. If
stress period data for a stress period contains an integer, a -1
(-1 or 0), a list of lists, a numpy array or recarray, or a pandas
dataframe. If data for a stress period contains an integer, a -1
denotes data from the previous stress period will be reused and a 0
indicates there are no str reaches for this stress period.
Expand Down Expand Up @@ -367,6 +368,8 @@ def __init__(
for key, d in stress_period_data.items():
if isinstance(d, list):
d = np.array(d)
if isinstance(d, pd.DataFrame):
d = d.to_records(index=False)
if isinstance(d, np.recarray):
e = (
"ModflowStr error: recarray dtype: {} does not match "
Expand Down
3 changes: 1 addition & 2 deletions flopy/modflow/mfwel.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,7 @@ class ModflowWel(Package):
A flag that is used to determine if cell-by-cell budget data should be
saved. If ipakcb is non-zero cell-by-cell budget data will be saved.
(default is 0).
stress_period_data : list of boundaries, or recarray of boundaries, or
dictionary of boundaries
stress_period_data : list, recarray, dataframe or dictionary of boundaries.
Each well is defined through definition of
layer (int), row (int), column (int), flux (float).
The simplest form is a dictionary with a lists of boundaries for each
Expand Down
1 change: 0 additions & 1 deletion flopy/utils/mflistfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
import numpy as np
import pandas as pd

from ..utils import import_optional_dependency
from ..utils.flopy_io import get_ts_sp
from ..utils.utils_def import totim_to_datetime

Expand Down
Loading

0 comments on commit 9637bc6

Please sign in to comment.