Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 95 additions & 0 deletions autotest/test_usg.py
Original file line number Diff line number Diff line change
Expand Up @@ -471,3 +471,98 @@ def test_free_format_npl_constructor():

m2 = MfUsg()
assert m2.free_format_npl is None


def test_wel_comment_columns_read_write(function_tmpdir, freyberg_usg_model_path):
"""Test that trailing comment text in WEL files is preserved on roundtrip."""
import shutil

# Copy freyberg model to tmpdir so we can modify the WEL file
for f in freyberg_usg_model_path.iterdir():
shutil.copy2(f, function_tmpdir)

# Rewrite the WEL file with trailing comments on each data line
wel_file = function_tmpdir / "freyberg.usg.wel"
with open(wel_file) as f:
lines = f.readlines()

with open(wel_file, "w") as f:
for i, line in enumerate(lines):
if i >= 3 and line.strip() and not line.startswith("#"):
# Append comment text to data lines
f.write(line.rstrip() + f" zone1 well_{i}\n")
else:
f.write(line)

# Load the model
m = MfUsg.load("freyberg.usg.nam", model_ws=str(function_tmpdir))
wel = m.get_package("WEL")

# Verify comment columns are captured
spd = wel.stress_period_data[0]
assert "comment1" in spd.dtype.names, "comment1 not found in dtype"
assert "comment2" in spd.dtype.names, "comment2 not found in dtype"
assert spd["comment1"][0] == "zone1"
assert spd["comment2"][0] == "well_3"

# Write the model back out
out_dir = function_tmpdir / "output"
out_dir.mkdir()
m.model_ws = str(out_dir)
m.write_input()

# Read the written WEL file and verify comments are preserved
written_wel = out_dir / "freyberg.usg.wel"
with open(written_wel) as f:
written_lines = f.readlines()

# Header should NOT contain "aux comment"
header = written_lines[1].lower()
assert "aux comment" not in header, (
f"aux comment found in header: {written_lines[1]}"
)

# Data lines should contain the comment text
data_lines = [l for l in written_lines if "zone1" in l]
assert len(data_lines) > 0, "Comment text not found in written WEL file"


def test_wel_comment_columns_fresh_model(function_tmpdir):
"""Test creating a fresh MfUsgWel with comment columns programmatically."""
m = MfUsg(modelname="test_fresh", structured=False, model_ws=function_tmpdir)

# Create stress period data with comments using get_empty
spd = MfUsgWel.get_empty(ncells=2, structured=False, n_comments=2)
spd["node"] = [0, 1]
spd["flux"] = [-100.0, -200.0]
spd["comment1"] = ["zone1", "zone2"]
spd["comment2"] = ["well_X", "well_Y"]

# Verify dtype has comment columns with object dtype
assert spd.dtype["comment1"] == object
assert spd.dtype["comment2"] == object

# Create the WEL package with matching dtype
wel = MfUsgWel(m, stress_period_data={0: spd}, dtype=spd.dtype, add_package=False)

# Verify comments are not treated as AUX
if wel.options:
for opt in wel.options:
assert "comment" not in opt.lower(), (
f"comment incorrectly declared as AUX: {opt}"
)

# Write the WEL file
wel_file = function_tmpdir / "test_fresh.wel"
wel.write_file(str(wel_file))

# Read and verify
with open(wel_file) as f:
lines = f.readlines()

header = lines[1].lower()
assert "aux comment" not in header

data_lines = [l for l in lines[3:] if l.strip() and not l.startswith("#")]
assert "well_X" in data_lines[0]
assert "well_Y" in data_lines[1]
15 changes: 14 additions & 1 deletion flopy/mfusg/mfusgwel.py
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,9 @@ def __init__(
self.parent.add_package(self)

@staticmethod
def get_empty(ncells=0, aux_names=None, structured=True, wellbot=False):
def get_empty(
ncells=0, aux_names=None, structured=True, wellbot=False, n_comments=0
):
"""Get empty recarray for MFUSG wells.

Parameters
Expand All @@ -252,6 +254,8 @@ def get_empty(ncells=0, aux_names=None, structured=True, wellbot=False):
Whether grid is structured
wellbot : bool
Whether WELLBOT option is used
n_comments : int
Number of comment columns to add (default is 0)

Returns
-------
Expand All @@ -275,6 +279,12 @@ def get_empty(ncells=0, aux_names=None, structured=True, wellbot=False):
if aux_names is not None:
dtype = Package.add_to_dtype(dtype, aux_names, np.float32)

# Add comment columns
if n_comments > 0:
comment_names = [f"comment{i + 1}" for i in range(n_comments)]
for name in comment_names:
dtype = Package.add_to_dtype(dtype, name, object)

return create_empty_recarray(ncells, dtype, default_value=-1.0e10)

def _check_for_aux(self, options, cln=False):
Expand All @@ -297,6 +307,9 @@ def _check_for_aux(self, options, cln=False):
dt = self.get_default_dtype(structured=self.parent.structured)
if len(self.dtype.names) > len(dt.names):
for name in self.dtype.names[len(dt.names) :]:
# Skip comment columns (object dtype) — not AUX variables
if self.dtype.fields[name][0] == object:
continue
ladd = True
for option in options:
if name.lower() in option.lower():
Expand Down
3 changes: 3 additions & 0 deletions flopy/modflow/mfwel.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,9 @@ def __init__(
dt = self.get_default_dtype(structured=self.parent.structured)
if len(self.dtype.names) > len(dt.names):
for name in self.dtype.names[len(dt.names) :]:
# Skip comment columns (object dtype) — not AUX variables
if self.dtype.fields[name][0] == object:
continue
ladd = True
for option in options:
if name.lower() in option.lower():
Expand Down
137 changes: 132 additions & 5 deletions flopy/pakbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,15 @@
import numpy as np
from numpy.lib.recfunctions import stack_arrays

from .utils import MfList, OptionBlock, Transient2d, Util2d, Util3d, check
from .utils import (
MfList,
OptionBlock,
Transient2d,
Util2d,
Util3d,
check,
create_empty_recarray,
)
from .utils.flopy_io import ulstrd


Expand Down Expand Up @@ -1085,6 +1093,9 @@ def load(
bnd_output_cln = None
stress_period_data_cln = {}
current_cln = None
is_mfusgwel = "mfusgwel" in pak_type_str
max_comment_cols = 0
max_comment_cols_cln = 0
for iper in range(nper):
if model.verbose:
msg = f" loading {pak_type} for kper {iper + 1:5d}"
Expand Down Expand Up @@ -1126,7 +1137,41 @@ def load(
current = pak_type.get_empty(
itmp, aux_names=aux_names, structured=model.structured, **usg_args
)
current = ulstrd(f, itmp, current, model, sfac_columns, ext_unit_dict)
if is_mfusgwel:
current, extra_tokens = ulstrd(
f,
itmp,
current,
model,
sfac_columns,
ext_unit_dict,
capture_comments=True,
)
n_comments = max((len(ct) for ct in extra_tokens), default=0)
if n_comments > 0:
max_comment_cols = max(max_comment_cols, n_comments)
new_dtype = current.dtype
for ci in range(n_comments):
new_dtype = Package.add_to_dtype(
new_dtype, f"comment{ci + 1}", object
)
new_ra = create_empty_recarray(
len(current), new_dtype, default_value=-1.0e10
)
for name in current.dtype.names:
new_ra[name] = current[name]
for ii in range(len(current)):
for ci in range(n_comments):
cname = f"comment{ci + 1}"
if ci < len(extra_tokens[ii]):
new_ra[cname][ii] = extra_tokens[ii][ci]
else:
new_ra[cname][ii] = ""
current = new_ra
else:
current = ulstrd(
f, itmp, current, model, sfac_columns, ext_unit_dict
)
if model.structured:
current["k"] -= 1
current["i"] -= 1
Expand All @@ -1149,9 +1194,48 @@ def load(
current_cln = pak_type.get_empty(
itmp_cln, aux_names=aux_names, structured=False, **usg_args
)
current_cln = ulstrd(
f, itmp_cln, current_cln, model, sfac_columns, ext_unit_dict
)
if is_mfusgwel:
current_cln, extra_tokens_cln = ulstrd(
f,
itmp_cln,
current_cln,
model,
sfac_columns,
ext_unit_dict,
capture_comments=True,
)
n_comments_cln = max(
(len(ct) for ct in extra_tokens_cln), default=0
)
if n_comments_cln > 0:
max_comment_cols_cln = max(max_comment_cols_cln, n_comments_cln)
new_dtype = current_cln.dtype
for ci in range(n_comments_cln):
new_dtype = Package.add_to_dtype(
new_dtype, f"comment{ci + 1}", object
)
new_ra = create_empty_recarray(
len(current_cln), new_dtype, default_value=-1.0e10
)
for name in current_cln.dtype.names:
new_ra[name] = current_cln[name]
for ii in range(len(current_cln)):
for ci in range(n_comments_cln):
cname = f"comment{ci + 1}"
if ci < len(extra_tokens_cln[ii]):
new_ra[cname][ii] = extra_tokens_cln[ii][ci]
else:
new_ra[cname][ii] = ""
current_cln = new_ra
else:
current_cln = ulstrd(
f,
itmp_cln,
current_cln,
model,
sfac_columns,
ext_unit_dict,
)
current_cln["node"] -= 1
bnd_output_cln = np.recarray.copy(current_cln)
else:
Expand Down Expand Up @@ -1231,6 +1315,43 @@ def load(
0, aux_names=aux_names, structured=model.structured, **usg_args
).dtype

# Normalize comment columns across stress periods for MfUsgWel
if is_mfusgwel and (max_comment_cols > 0 or max_comment_cols_cln > 0):
for spd_dict, max_cc in [
(stress_period_data, max_comment_cols),
(stress_period_data_cln, max_comment_cols_cln),
]:
if max_cc == 0:
continue
for iper in spd_dict:
spd = spd_dict[iper]
if not isinstance(spd, np.recarray):
continue
existing = sum(
1 for n in spd.dtype.names if n.startswith("comment")
)
if existing < max_cc:
new_dtype = spd.dtype
for ci in range(existing, max_cc):
new_dtype = Package.add_to_dtype(
new_dtype, f"comment{ci + 1}", object
)
new_ra = create_empty_recarray(
len(spd), new_dtype, default_value=-1.0e10
)
for name in spd.dtype.names:
new_ra[name] = spd[name]
for ci in range(existing, max_cc):
cname = f"comment{ci + 1}"
for ii in range(len(new_ra)):
new_ra[cname][ii] = ""
spd_dict[iper] = new_ra

# Update dtype to include comment columns
if max_comment_cols > 0:
for ci in range(max_comment_cols):
dtype = Package.add_to_dtype(dtype, f"comment{ci + 1}", object)

if openfile:
f.close()

Expand All @@ -1248,6 +1369,12 @@ def load(
cln_dtype = pak_type.get_empty(
0, aux_names=aux_names, structured=False, **usg_args
).dtype
# Update cln_dtype to include comment columns
if max_comment_cols_cln > 0:
for ci in range(max_comment_cols_cln):
cln_dtype = Package.add_to_dtype(
cln_dtype, f"comment{ci + 1}", object
)
pak = pak_type(
model,
ipakcb=ipakcb,
Expand Down
15 changes: 14 additions & 1 deletion flopy/utils/flopy_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ def get_url_text(url, error_msg=None):
return


def ulstrd(f, nlist, ra, model, sfac_columns, ext_unit_dict):
def ulstrd(f, nlist, ra, model, sfac_columns, ext_unit_dict, capture_comments=False):
"""
Read a list and allow for open/close, binary, external, sfac, etc.

Expand All @@ -404,9 +404,17 @@ def ulstrd(f, nlist, ra, model, sfac_columns, ext_unit_dict):
then in this case ext_unit_dict is required, which can be
constructed using the function
:class:`flopy.utils.mfreadnam.parsenamefile`.
capture_comments : bool, optional
If True, capture extra tokens beyond the expected number of columns
and return them as a list of lists alongside the recarray.
(default is False)

Returns
-------
ra : np.recarray
The filled record array. If capture_comments is True, returns a
tuple of (ra, comment_tokens) where comment_tokens is a list of
lists containing extra tokens for each row.

"""

Expand All @@ -419,6 +427,7 @@ def ulstrd(f, nlist, ra, model, sfac_columns, ext_unit_dict):
close_the_file = False
file_handle = f
mode = "r"
comment_tokens = [[] for _ in range(nlist)] if capture_comments else None

# check for external
if line.strip().lower().startswith("external"):
Expand Down Expand Up @@ -488,6 +497,8 @@ def ulstrd(f, nlist, ra, model, sfac_columns, ext_unit_dict):
if model.free_format_input:
# whitespace separated
t = line_parse(line)
if capture_comments and len(t) > ncol:
comment_tokens[ii] = t[ncol:]
if len(t) < ncol:
t = t + (ncol - len(t)) * [0.0]
else:
Expand All @@ -509,6 +520,8 @@ def ulstrd(f, nlist, ra, model, sfac_columns, ext_unit_dict):
if close_the_file:
file_handle.close()

if capture_comments:
return ra, comment_tokens
return ra


Expand Down
Loading