diff --git a/autotest/test_usg.py b/autotest/test_usg.py index 4fb00a654..d5ba72296 100644 --- a/autotest/test_usg.py +++ b/autotest/test_usg.py @@ -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] diff --git a/flopy/mfusg/mfusgwel.py b/flopy/mfusg/mfusgwel.py index d9e39314f..bf800317c 100644 --- a/flopy/mfusg/mfusgwel.py +++ b/flopy/mfusg/mfusgwel.py @@ -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 @@ -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 ------- @@ -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): @@ -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(): diff --git a/flopy/modflow/mfwel.py b/flopy/modflow/mfwel.py index 86c815c37..d4ffbfad6 100644 --- a/flopy/modflow/mfwel.py +++ b/flopy/modflow/mfwel.py @@ -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(): diff --git a/flopy/pakbase.py b/flopy/pakbase.py index b9bb0d9fa..2754d7b35 100644 --- a/flopy/pakbase.py +++ b/flopy/pakbase.py @@ -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 @@ -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}" @@ -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 @@ -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: @@ -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() @@ -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, diff --git a/flopy/utils/flopy_io.py b/flopy/utils/flopy_io.py index 2b74ac634..18c604da1 100644 --- a/flopy/utils/flopy_io.py +++ b/flopy/utils/flopy_io.py @@ -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. @@ -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. """ @@ -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"): @@ -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: @@ -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