fix: auto-set max_columns_of_data for external writes

autotest/test_mf6_max_columns_setting.py

@@ -164,6 +164,73 @@ def check_columns(expected):
     check_columns(expected=1)
 
 
+def test_auto_set_ncol_before_set_all_data_external(function_tmpdir):
+    """
+    Auto-set to ncol before set_all_data_external(), not just write_simulation().
+    Repro https://github.com/modflowpy/flopy/issues/2572#issuecomment-3164881948
+    """
+    ncol = 25
+    sim_ws = function_tmpdir / "test_auto_set_ncol"
+    sim_ws.mkdir(exist_ok=True)
+
+    sim = flopy.mf6.MFSimulation(sim_name="test_sim", sim_ws=sim_ws)
+    flopy.mf6.ModflowTdis(sim, time_units="DAYS", nper=1, perioddata=[(1.0, 1, 1.0)])
+    gwf = flopy.mf6.ModflowGwf(sim, modelname="model", model_nam_file="model.nam")
+    flopy.mf6.ModflowGwfdis(
+        gwf, nlay=1, nrow=2, ncol=ncol, delr=100.0, delc=100.0, top=100.0, botm=0.0
+    )
+    flopy.mf6.ModflowGwfic(
+        gwf, strt=np.arange(2 * ncol).reshape(1, 2, ncol).astype(float)
+    )
+    flopy.mf6.ModflowGwfnpf(gwf, save_flows=True, icelltype=0, k=1.0)
+    ims = flopy.mf6.ModflowIms(sim)
+    sim.register_ims_package(ims, [gwf.name])
+
+    sim.set_all_data_external()
+    assert sim.simulation_data.max_columns_of_data == ncol
+    sim.write_simulation(silent=True)
+
+    strt_file = sim_ws / "model.ic_strt.txt"
+    assert strt_file.exists()
+    data_lines = [
+        line.strip()
+        for line in strt_file.read_text().splitlines()
+        if line.strip() and not line.strip().startswith("#")
+    ]
+    for line in data_lines:
+        assert len(line.split()) == ncol
+
+    # also test the model-level entry point
+    sim2_ws = function_tmpdir / "test_auto_set_ncol_model"
+    sim2_ws.mkdir(exist_ok=True)
+    sim2 = flopy.mf6.MFSimulation(sim_name="test_sim2", sim_ws=sim2_ws)
+    flopy.mf6.ModflowTdis(sim2, time_units="DAYS", nper=1, perioddata=[(1.0, 1, 1.0)])
+    gwf2 = flopy.mf6.ModflowGwf(sim2, modelname="model", model_nam_file="model.nam")
+    flopy.mf6.ModflowGwfdis(
+        gwf2, nlay=1, nrow=2, ncol=ncol, delr=100.0, delc=100.0, top=100.0, botm=0.0
+    )
+    flopy.mf6.ModflowGwfic(
+        gwf2, strt=np.arange(2 * ncol).reshape(1, 2, ncol).astype(float)
+    )
+    flopy.mf6.ModflowGwfnpf(gwf2, save_flows=True, icelltype=0, k=1.0)
+    ims2 = flopy.mf6.ModflowIms(sim2)
+    sim2.register_ims_package(ims2, [gwf2.name])
+
+    gwf2.set_all_data_external()
+    assert sim2.simulation_data.max_columns_of_data == ncol
+    sim2.write_simulation(silent=True)
+
+    strt_file2 = sim2_ws / "model.ic_strt.txt"
+    assert strt_file2.exists()
+    data_lines2 = [
+        line.strip()
+        for line in strt_file2.read_text().splitlines()
+        if line.strip() and not line.strip().startswith("#")
+    ]
+    for line in data_lines2:
+        assert len(line.split()) == ncol
+
+
 reason_ext = (
     "set_all_data_external() writes files immediately. Changing output settings"
     "afterward has no effect unless set_all_data_external() is called again. "

flopy/mf6/mfmodel.py

@@ -1044,13 +1044,9 @@ def load_base(
                     print(f"    loading package {ftype}...")
                 # load package
                 instance.load_package(ftype, fname, pname, strict, None)
-                sim_data = simulation.simulation_data
-                if ftype == "dis" and sim_data._max_columns_set_by != 'user':
-                    # set column wrap to ncol
-                    dis = instance.get_package("dis", type_only=True)
-                    if dis is not None and hasattr(dis, "ncol"):
-                        sim_data._max_columns_of_data = dis.ncol.get_data()
-                        sim_data._max_columns_set_by = 'auto'
+                if ftype == "dis":
+                    # instance is not yet in simulation._models at this point
+                    simulation._auto_set_max_columns(models=[instance])
         # load referenced packages
         if modelname in instance.simulation_data.referenced_files:
             for ref_file in instance.simulation_data.referenced_files[
@@ -1376,11 +1372,7 @@ def write(self, ext_file_action=ExtFileAction.copy_relative_paths):
 
         self.name_file.write(ext_file_action=ext_file_action)
 
-        if self.simulation_data._max_columns_set_by != 'user':
-            grid_type = self.get_grid_type()
-            if grid_type == DiscretizationType.DIS:
-                self.simulation_data._max_columns_of_data = self.dis.ncol.get_data()
-                self.simulation_data._max_columns_set_by = 'auto'
+        self.simulation._auto_set_max_columns()
 
         # write packages
         for pp in self.packagelist:
@@ -1895,6 +1887,8 @@ def set_all_data_external(
                 will not be rewritten. Default is False.
 
         """
+        self.simulation._auto_set_max_columns()
+
         for package in self.packagelist:
             package.set_all_data_external(
                 check_data,

flopy/mf6/mfsimbase.py

@@ -1597,6 +1597,26 @@ def rename_all_packages(self, name):
         for model in self._models.values():
             model.rename_all_packages(name)
 
+    def _auto_set_max_columns(self, models=None):
+        """Set max_columns_of_data to ncol of the first structured DIS model found,
+        unless the value has already been explicitly set by the user.
+
+        Parameters
+        ----------
+        models : iterable, optional
+            Models to search. Defaults to all registered models. Pass an
+            explicit list when a model is not yet registered (e.g. mid-load).
+        """
+        sim_data = self.simulation_data
+        if sim_data._max_columns_set_by == "user":
+            return
+        for model in (models if models is not None else self._models.values()):
+            dis = model.get_package("dis", type_only=True)
+            if dis is not None and hasattr(dis, "ncol"):
+                sim_data._max_columns_of_data = dis.ncol.get_data()
+                sim_data._max_columns_set_by = "auto"
+                return
+
     def set_all_data_external(
         self,
         check_data=True,
@@ -1641,6 +1661,8 @@ def set_all_data_external(
                 will not be rewritten. Default is False.
         """
 
+        self._auto_set_max_columns()
+
         # copy any files whose paths have changed
         self.simulation_data.mfpath.copy_files()
         # set data external for all packages in all models
@@ -1710,14 +1732,7 @@ def write_simulation(
                 Writes out the simulation in silent mode (verbosity_level = 0)
 
         """
-        sim_data = self.simulation_data
-        if sim_data._max_columns_set_by != 'user':
-            # search for dis packages
-            for model in self._models.values():
-                dis = model.get_package("dis", type_only=True)
-                if dis is not None and hasattr(dis, "ncol"):
-                    sim_data._max_columns_of_data = dis.ncol.get_data()
-                    sim_data._max_columns_set_by = 'auto'
+        self._auto_set_max_columns()
 
         saved_verb_lvl = self.simulation_data.verbosity_level
         if silent: