Coverage for C:\src\imod-python\imod\mf6\simulation.py: 96%
484 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1from __future__ import annotations
3import collections
4import pathlib
5import subprocess
6import warnings
7from copy import deepcopy
8from pathlib import Path
9from typing import Any, Callable, DefaultDict, Iterable, Optional, Union, cast
11import cftime
12import dask
13import jinja2
14import numpy as np
15import tomli
16import tomli_w
17import xarray as xr
18import xugrid as xu
20import imod
21import imod.logging
22import imod.mf6.exchangebase
23from imod.logging import standard_log_decorator
24from imod.mf6.gwfgwf import GWFGWF
25from imod.mf6.gwfgwt import GWFGWT
26from imod.mf6.gwtgwt import GWTGWT
27from imod.mf6.ims import Solution
28from imod.mf6.interfaces.isimulation import ISimulation
29from imod.mf6.model import Modflow6Model
30from imod.mf6.model_gwf import GroundwaterFlowModel
31from imod.mf6.model_gwt import GroundwaterTransportModel
32from imod.mf6.multimodel.exchange_creator_structured import ExchangeCreator_Structured
33from imod.mf6.multimodel.exchange_creator_unstructured import (
34 ExchangeCreator_Unstructured,
35)
36from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
37from imod.mf6.out import open_cbc, open_conc, open_hds
38from imod.mf6.package import Package
39from imod.mf6.ssm import SourceSinkMixing
40from imod.mf6.statusinfo import NestedStatusInfo
41from imod.mf6.utilities.mask import _mask_all_models
42from imod.mf6.utilities.regrid import _regrid_like
43from imod.mf6.write_context import WriteContext
44from imod.schemata import ValidationError
45from imod.typing import GridDataArray, GridDataset
46from imod.typing.grid import (
47 concat,
48 is_equal,
49 is_unstructured,
50 merge_partitions,
51)
53OUTPUT_FUNC_MAPPING: dict[str, Callable] = {
54 "head": open_hds,
55 "concentration": open_conc,
56 "budget-flow": open_cbc,
57 "budget-transport": open_cbc,
58}
60OUTPUT_MODEL_MAPPING: dict[
61 str, type[GroundwaterFlowModel] | type[GroundwaterTransportModel]
62] = {
63 "head": GroundwaterFlowModel,
64 "concentration": GroundwaterTransportModel,
65 "budget-flow": GroundwaterFlowModel,
66 "budget-transport": GroundwaterTransportModel,
67}
70def get_models(simulation: Modflow6Simulation) -> dict[str, Modflow6Model]:
71 return {k: v for k, v in simulation.items() if isinstance(v, Modflow6Model)}
74def get_packages(simulation: Modflow6Simulation) -> dict[str, Package]:
75 return {
76 pkg_name: pkg
77 for pkg_name, pkg in simulation.items()
78 if isinstance(pkg, Package)
79 }
82class Modflow6Simulation(collections.UserDict, ISimulation):
83 def _initialize_template(self):
84 loader = jinja2.PackageLoader("imod", "templates/mf6")
85 env = jinja2.Environment(loader=loader, keep_trailing_newline=True)
86 self._template = env.get_template("sim-nam.j2")
88 def __init__(self, name):
89 super().__init__()
90 self.name = name
91 self.directory = None
92 self._initialize_template()
94 def __setitem__(self, key, value):
95 super().__setitem__(key, value)
97 def update(self, *args, **kwargs):
98 for k, v in dict(*args, **kwargs).items():
99 self[k] = v
101 def time_discretization(self, times):
102 warnings.warn(
103 f"{self.__class__.__name__}.time_discretization() is deprecated. "
104 f"In the future call {self.__class__.__name__}.create_time_discretization().",
105 DeprecationWarning,
106 )
107 self.create_time_discretization(additional_times=times)
109 def create_time_discretization(self, additional_times, validate: bool = True):
110 """
111 Collect all unique times from model packages and additional given
112 `times`. These unique times are used as stress periods in the model. All
113 stress packages must have the same starting time. Function creates
114 TimeDiscretization object which is set to self["time_discretization"]
116 The time discretization in imod-python works as follows:
118 - The datetimes of all packages you send in are always respected
119 - Subsequently, the input data you use is always included fully as well
120 - All times are treated as starting times for the stress: a stress is
121 always applied until the next specified date
122 - For this reason, a final time is required to determine the length of
123 the last stress period
124 - Additional times can be provided to force shorter stress periods &
125 more detailed output
126 - Every stress has to be defined on the first stress period (this is a
127 modflow requirement)
129 Or visually (every letter a date in the time axes):
131 >>> recharge a - b - c - d - e - f
132 >>> river g - - - - h - - - - j
133 >>> times - - - - - - - - - - - i
134 >>> model a - b - c h d - e - f i
136 with the stress periods defined between these dates. I.e. the model
137 times are the set of all times you include in the model.
139 Parameters
140 ----------
141 additional_times : str, datetime; or iterable of str, datetimes.
142 Times to add to the time discretization. At least one single time
143 should be given, which will be used as the ending time of the
144 simulation.
146 Note
147 ----
148 To set the other parameters of the TimeDiscretization object, you have
149 to set these to the object after calling this function.
151 Example
152 -------
153 >>> simulation = imod.mf6.Modflow6Simulation("example")
154 >>> simulation.create_time_discretization(times=["2000-01-01", "2000-01-02"])
155 >>> # Set number of timesteps
156 >>> simulation["time_discretization"]["n_timesteps"] = 5
157 """
158 self.use_cftime = any(
159 model._use_cftime()
160 for model in self.values()
161 if isinstance(model, Modflow6Model)
162 )
164 times = [
165 imod.util.time.to_datetime_internal(time, self.use_cftime)
166 for time in additional_times
167 ]
168 for model in self.values():
169 if isinstance(model, Modflow6Model):
170 times.extend(model._yield_times())
172 # np.unique also sorts
173 times = np.unique(np.hstack(times))
175 duration = imod.util.time.timestep_duration(times, self.use_cftime) # type: ignore
176 # Generate time discretization, just rely on default arguments
177 # Probably won't be used that much anyway?
178 timestep_duration = xr.DataArray(
179 duration, coords={"time": np.array(times)[:-1]}, dims=("time",)
180 )
181 self["time_discretization"] = imod.mf6.TimeDiscretization(
182 timestep_duration=timestep_duration, validate=validate
183 )
185 def render(self, write_context: WriteContext):
186 """Renders simulation namefile"""
187 d: dict[str, Any] = {}
188 models = []
189 solutiongroups = []
190 for key, value in self.items():
191 if isinstance(value, Modflow6Model):
192 model_name_file = pathlib.Path(
193 write_context.root_directory / pathlib.Path(f"{key}", f"{key}.nam")
194 ).as_posix()
195 models.append((value.model_id(), model_name_file, key))
196 elif isinstance(value, Package):
197 if value._pkg_id == "tdis":
198 d["tdis6"] = f"{key}.tdis"
199 elif value._pkg_id == "ims":
200 slnnames = value["modelnames"].values
201 modeltypes = set()
202 for name in slnnames:
203 try:
204 modeltypes.add(type(self[name]))
205 except KeyError:
206 raise KeyError(f"model {name} of {key} not found")
208 if len(modeltypes) > 1:
209 raise ValueError(
210 "Only a single type of model allowed in a solution"
211 )
212 solutiongroups.append(("ims6", f"{key}.ims", slnnames))
214 d["models"] = models
215 if len(models) > 1:
216 d["exchanges"] = self.get_exchange_relationships()
218 d["solutiongroups"] = [solutiongroups]
219 return self._template.render(d)
221 @standard_log_decorator()
222 def write(
223 self,
224 directory=".",
225 binary=True,
226 validate: bool = True,
227 use_absolute_paths=False,
228 ):
229 """
230 Write Modflow6 simulation, including assigned groundwater flow and
231 transport models.
233 Parameters
234 ----------
235 directory: str, pathlib.Path
236 Directory to write Modflow 6 simulation to.
237 binary: ({True, False}, optional)
238 Whether to write time-dependent input for stress packages as binary
239 files, which are smaller in size, or more human-readable text files.
240 validate: ({True, False}, optional)
241 Whether to validate the Modflow6 simulation, including models, at
242 write. If True, erronous model input will throw a
243 ``ValidationError``.
244 absolute_paths: ({True, False}, optional)
245 True if all paths written to the mf6 inputfiles should be absolute.
246 """
247 # create write context
248 write_context = WriteContext(directory, binary, use_absolute_paths)
249 if self.is_split():
250 write_context.is_partitioned = True
252 # Check models for required content
253 for key, model in self.items():
254 # skip timedis, exchanges
255 if isinstance(model, Modflow6Model):
256 model._model_checks(key)
258 # Generate GWF-GWT exchanges
259 if gwfgwt_exchanges := self._generate_gwfgwt_exchanges():
260 self["gwtgwf_exchanges"] = gwfgwt_exchanges
262 directory = pathlib.Path(directory)
263 directory.mkdir(exist_ok=True, parents=True)
265 # Write simulation namefile
266 mfsim_content = self.render(write_context)
267 mfsim_path = directory / "mfsim.nam"
268 with open(mfsim_path, "w") as f:
269 f.write(mfsim_content)
271 # Write time discretization file
272 self["time_discretization"].write(directory, "time_discretization")
274 # Write individual models
275 status_info = NestedStatusInfo("Simulation validation status")
276 globaltimes = self["time_discretization"]["time"].values
277 for key, value in self.items():
278 model_write_context = write_context.copy_with_new_write_directory(
279 write_context.simulation_directory
280 )
281 # skip timedis, exchanges
282 if isinstance(value, Modflow6Model):
283 status_info.add(
284 value.write(
285 modelname=key,
286 globaltimes=globaltimes,
287 validate=validate,
288 write_context=model_write_context,
289 )
290 )
291 elif isinstance(value, Package):
292 if value._pkg_id == "ims":
293 ims_write_context = write_context.copy_with_new_write_directory(
294 write_context.simulation_directory
295 )
296 value.write(key, globaltimes, ims_write_context)
297 elif isinstance(value, list):
298 for exchange in value:
299 if isinstance(exchange, imod.mf6.exchangebase.ExchangeBase):
300 exchange.write(
301 exchange.package_name(), globaltimes, write_context
302 )
304 if status_info.has_errors():
305 raise ValidationError("\n" + status_info.to_string())
307 self.directory = directory
309 def run(self, mf6path: Union[str, Path] = "mf6") -> None:
310 """
311 Run Modflow 6 simulation. This method runs a subprocess calling
312 ``mf6path``. This argument is set to ``mf6``, which means the Modflow 6
313 executable is expected to be added to your PATH environment variable.
314 :doc:`See this writeup how to add Modflow 6 to your PATH on Windows </examples/mf6/index>`
316 Note that the ``write`` method needs to be called before this method is
317 called.
319 Parameters
320 ----------
321 mf6path: Union[str, Path]
322 Path to the Modflow 6 executable. Defaults to calling ``mf6``.
324 Examples
325 --------
326 Make sure you write your model first
328 >>> simulation.write(path/to/model)
329 >>> simulation.run()
330 """
331 if self.directory is None:
332 raise RuntimeError(f"Simulation {self.name} has not been written yet.")
333 with imod.util.cd(self.directory):
334 result = subprocess.run(mf6path, capture_output=True)
335 if result.returncode != 0:
336 raise RuntimeError(
337 f"Simulation {self.name}: {mf6path} failed to run with returncode "
338 f"{result.returncode}, and error message:\n\n{result.stdout.decode()} "
339 )
341 def open_head(
342 self,
343 dry_nan: bool = False,
344 simulation_start_time: Optional[np.datetime64] = None,
345 time_unit: Optional[str] = "d",
346 ) -> GridDataArray:
347 """
348 Open heads of finished simulation, requires that the ``run`` method has
349 been called.
351 The data is lazily read per timestep and automatically converted into
352 (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV
353 respectively. The conversion is done via the information stored in the
354 Binary Grid file (GRB).
356 Parameters
357 ----------
358 dry_nan: bool, default value: False.
359 Whether to convert dry values to NaN.
360 simulation_start_time : Optional datetime
361 The time and date correpsonding to the beginning of the simulation.
362 Use this to convert the time coordinates of the output array to
363 calendar time/dates. time_unit must also be present if this argument is present.
364 time_unit: Optional str
365 The time unit MF6 is working in, in string representation.
366 Only used if simulation_start_time was provided.
367 Admissible values are:
368 ns -> nanosecond
369 ms -> microsecond
370 s -> second
371 m -> minute
372 h -> hour
373 d -> day
374 w -> week
375 Units "month" or "year" are not supported, as they do not represent unambiguous timedelta values durations.
377 Returns
378 -------
379 head: Union[xr.DataArray, xu.UgridDataArray]
381 Examples
382 --------
383 Make sure you write and run your model first
385 >>> simulation.write(path/to/model)
386 >>> simulation.run()
388 Then open heads:
390 >>> head = simulation.open_head()
391 """
392 return self._open_output(
393 "head",
394 dry_nan=dry_nan,
395 simulation_start_time=simulation_start_time,
396 time_unit=time_unit,
397 )
399 def open_transport_budget(
400 self,
401 species_ls: Optional[list[str]] = None,
402 simulation_start_time: Optional[np.datetime64] = None,
403 time_unit: Optional[str] = "d",
404 ) -> GridDataArray | GridDataset:
405 """
406 Open transport budgets of finished simulation, requires that the ``run``
407 method has been called.
409 The data is lazily read per timestep and automatically converted into
410 (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV
411 respectively. The conversion is done via the information stored in the
412 Binary Grid file (GRB).
414 Parameters
415 ----------
416 species_ls: list of strings, default value: None.
417 List of species names, which will be used to concatenate the
418 concentrations along the ``"species"`` dimension, in case the
419 simulation has multiple species and thus multiple transport models.
420 If None, transport model names will be used as species names.
422 Returns
423 -------
424 budget: Dict[str, xr.DataArray|xu.UgridDataArray]
425 DataArray contains float64 data of the budgets, with dimensions ("time",
426 "layer", "y", "x").
428 """
429 return self._open_output(
430 "budget-transport",
431 species_ls=species_ls,
432 simulation_start_time=simulation_start_time,
433 time_unit=time_unit,
434 merge_to_dataset=True,
435 flowja=False,
436 )
438 def open_flow_budget(
439 self,
440 flowja: bool = False,
441 simulation_start_time: Optional[np.datetime64] = None,
442 time_unit: Optional[str] = "d",
443 ) -> GridDataArray | GridDataset:
444 """
445 Open flow budgets of finished simulation, requires that the ``run``
446 method has been called.
448 The data is lazily read per timestep and automatically converted into
449 (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV
450 respectively. The conversion is done via the information stored in the
451 Binary Grid file (GRB).
453 The ``flowja`` argument controls whether the flow-ja-face array (if
454 present) is returned in grid form as "as is". By default
455 ``flowja=False`` and the array is returned in "grid form", meaning:
457 * DIS: in right, front, and lower face flow. All flows are placed in
458 the cell.
459 * DISV: in horizontal and lower face flow.the horizontal flows are
460 placed on the edges and the lower face flow is placed on the faces.
462 When ``flowja=True``, the flow-ja-face array is returned as it is found in
463 the CBC file, with a flow for every cell to cell connection. Additionally,
464 a ``connectivity`` DataArray is returned describing for every cell (n) its
465 connected cells (m).
467 Parameters
468 ----------
469 flowja: bool, default value: False
470 Whether to return the flow-ja-face values "as is" (``True``) or in a
471 grid form (``False``).
473 Returns
474 -------
475 budget: Dict[str, xr.DataArray|xu.UgridDataArray]
476 DataArray contains float64 data of the budgets, with dimensions ("time",
477 "layer", "y", "x").
479 Examples
480 --------
481 Make sure you write and run your model first
483 >>> simulation.write(path/to/model)
484 >>> simulation.run()
486 Then open budgets:
488 >>> budget = simulation.open_flow_budget()
490 Check the contents:
492 >>> print(budget.keys())
494 Get the drainage budget, compute a time mean for the first layer:
496 >>> drn_budget = budget["drn]
497 >>> mean = drn_budget.sel(layer=1).mean("time")
498 """
499 return self._open_output(
500 "budget-flow",
501 flowja=flowja,
502 simulation_start_time=simulation_start_time,
503 time_unit=time_unit,
504 merge_to_dataset=True,
505 )
507 def open_concentration(
508 self,
509 species_ls: Optional[list[str]] = None,
510 dry_nan: bool = False,
511 simulation_start_time: Optional[np.datetime64] = None,
512 time_unit: Optional[str] = "d",
513 ) -> GridDataArray:
514 """
515 Open concentration of finished simulation, requires that the ``run``
516 method has been called.
518 The data is lazily read per timestep and automatically converted into
519 (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV
520 respectively. The conversion is done via the information stored in the
521 Binary Grid file (GRB).
523 Parameters
524 ----------
525 species_ls: list of strings, default value: None.
526 List of species names, which will be used to concatenate the
527 concentrations along the ``"species"`` dimension, in case the
528 simulation has multiple species and thus multiple transport models.
529 If None, transport model names will be used as species names.
530 dry_nan: bool, default value: False.
531 Whether to convert dry values to NaN.
533 Returns
534 -------
535 concentration: Union[xr.DataArray, xu.UgridDataArray]
537 Examples
538 --------
539 Make sure you write and run your model first
541 >>> simulation.write(path/to/model)
542 >>> simulation.run()
544 Then open concentrations:
546 >>> concentration = simulation.open_concentration()
547 """
548 return self._open_output(
549 "concentration",
550 species_ls=species_ls,
551 dry_nan=dry_nan,
552 simulation_start_time=simulation_start_time,
553 time_unit=time_unit,
554 )
556 def _open_output(self, output: str, **settings) -> GridDataArray | GridDataset:
557 """
558 Opens output of one or multiple models.
560 Parameters
561 ----------
562 output: str
563 Output variable name to open
564 **settings:
565 Extra settings that need to be passed through to the respective
566 output function.
567 """
568 modeltype = OUTPUT_MODEL_MAPPING[output]
569 modelnames = self.get_models_of_type(modeltype._model_id).keys()
570 if len(modelnames) == 0:
571 modeltype = OUTPUT_MODEL_MAPPING[output]
572 raise ValueError(
573 f"Could not find any models of appropriate type for {output}, "
574 f"make sure a model of type {modeltype} is assigned to simulation."
575 )
577 if output in ["head", "budget-flow"]:
578 return self._open_single_output(modelnames, output, **settings)
579 elif output in ["concentration", "budget-transport"]:
580 return self._concat_species(output, **settings)
581 else:
582 raise RuntimeError(
583 f"Unexpected error when opening {output} for {modelnames}"
584 )
585 return
587 def _open_single_output(
588 self, modelnames: list[str], output: str, **settings
589 ) -> GridDataArray | GridDataset:
590 """
591 Open single output, e.g. concentration of single species, or heads. This
592 can be output of partitioned models that need to be merged.
593 """
594 if len(modelnames) == 0:
595 modeltype = OUTPUT_MODEL_MAPPING[output]
596 raise ValueError(
597 f"Could not find any models of appropriate type for {output}, "
598 f"make sure a model of type {modeltype} is assigned to simulation."
599 )
600 elif len(modelnames) == 1:
601 modelname = next(iter(modelnames))
602 return self._open_single_output_single_model(modelname, output, **settings)
603 elif self.is_split():
604 if "budget" in output:
605 return self._merge_budgets(modelnames, output, **settings)
606 else:
607 return self._merge_states(modelnames, output, **settings)
608 raise ValueError("error in _open_single_output")
610 def _merge_states(
611 self, modelnames: list[str], output: str, **settings
612 ) -> GridDataArray:
613 state_partitions = []
614 for modelname in modelnames:
615 state_partitions.append(
616 self._open_single_output_single_model(modelname, output, **settings)
617 )
618 return merge_partitions(state_partitions)
620 def _merge_and_assign_exchange_budgets(self, cbc: GridDataset) -> GridDataset:
621 """
622 Merge and assign exchange budgets to cell by cell budgets:
623 cbc[[gwf-gwf_1, gwf-gwf_3]] to cbc[gwf-gwf]
624 """
625 exchange_names = [
626 key
627 for key in cast(Iterable[str], cbc.keys())
628 if (("gwf-gwf" in key) or ("gwt-gwt" in key))
629 ]
630 exchange_budgets = cbc[exchange_names].to_array().sum(dim="variable")
631 cbc = cbc.drop_vars(exchange_names)
632 # "gwf-gwf" or "gwt-gwt"
633 exchange_key = exchange_names[0].split("_")[0]
634 cbc[exchange_key] = exchange_budgets
635 return cbc
637 def _pad_missing_variables(self, cbc_per_partition: list[GridDataset]) -> None:
638 """
639 Boundary conditions can be missing in certain partitions, as do their
640 budgets, in which case we manually assign an empty grid of nans.
641 """
642 dims_per_unique_key = {
643 key: cbc[key].dims for cbc in cbc_per_partition for key in cbc.keys()
644 }
645 for cbc in cbc_per_partition:
646 missing_keys = set(dims_per_unique_key.keys()) - set(cbc.keys())
648 for missing in missing_keys:
649 missing_dims = dims_per_unique_key[missing]
650 missing_coords = {dim: cbc.coords[dim] for dim in missing_dims}
652 shape = tuple([len(missing_coords[dim]) for dim in missing_dims])
653 chunks = (1,) + shape[1:]
654 missing_data = dask.array.full(shape, np.nan, chunks=chunks)
656 missing_grid = xr.DataArray(
657 missing_data, dims=missing_dims, coords=missing_coords
658 )
659 if isinstance(cbc, xu.UgridDataset):
660 missing_grid = xu.UgridDataArray(
661 missing_grid,
662 grid=cbc.ugrid.grid,
663 )
664 cbc[missing] = missing_grid
666 def _merge_budgets(
667 self, modelnames: list[str], output: str, **settings
668 ) -> GridDataset:
669 if settings["flowja"] is True:
670 raise ValueError("``flowja`` cannot be set to True when merging budgets.")
672 cbc_per_partition = []
673 for modelname in modelnames:
674 cbc = self._open_single_output_single_model(modelname, output, **settings)
675 # Merge and assign exchange budgets to dataset
676 # FUTURE: Refactor to insert these exchange budgets in horizontal
677 # flows.
678 cbc = self._merge_and_assign_exchange_budgets(cbc)
679 if not is_unstructured(cbc):
680 cbc = cbc.where(self[modelname].domain, other=np.nan)
681 cbc_per_partition.append(cbc)
683 self._pad_missing_variables(cbc_per_partition)
685 return merge_partitions(cbc_per_partition)
687 def _concat_species(
688 self, output: str, species_ls: Optional[list[str]] = None, **settings
689 ) -> GridDataArray | GridDataset:
690 # groupby flow model, to somewhat enforce consistent transport model
691 # ordening. Say:
692 # F = Flow model, T = Transport model
693 # a = species "a", b = species "b"
694 # 1 = partition 1, 2 = partition 2
695 # then this:
696 # F1Ta1 F1Tb1 F2Ta2 F2Tb2 -> F1: [Ta1, Tb1], F2: [Ta2, Tb2]
697 # F1Ta1 F2Tb1 F1Ta1 F2Tb2 -> F1: [Ta1, Tb1], F2: [Ta2, Tb2]
698 tpt_models_per_flow_model = self._get_transport_models_per_flow_model()
699 all_tpt_names = list(tpt_models_per_flow_model.values())
701 # [[Ta_1, Tb_1], [Ta_2, Tb_2]] -> [[Ta_1, Ta_2], [Tb_1, Tb_2]]
702 # [[Ta, Tb]] -> [[Ta], [Tb]]
703 tpt_names_per_species = list(zip(*all_tpt_names))
705 if self.is_split():
706 # [[Ta_1, Tb_1], [Ta_2, Tb_2]] -> [Ta, Tb]
707 unpartitioned_modelnames = [
708 tpt_name.rpartition("_")[0] for tpt_name in all_tpt_names[0]
709 ]
710 else:
711 # [[Ta, Tb]] -> [Ta, Tb]
712 unpartitioned_modelnames = all_tpt_names[0]
714 if not species_ls:
715 species_ls = unpartitioned_modelnames
717 if len(species_ls) != len(tpt_names_per_species):
718 raise ValueError(
719 "species_ls does not equal the number of transport models, "
720 f"expected length {len(tpt_names_per_species)}, received {species_ls}"
721 )
723 if len(species_ls) == 1:
724 return self._open_single_output(
725 list(tpt_names_per_species[0]), output, **settings
726 )
728 # Concatenate species
729 outputs = []
730 for species, tpt_names in zip(species_ls, tpt_names_per_species):
731 output_data = self._open_single_output(list(tpt_names), output, **settings)
732 output_data = output_data.assign_coords(species=species)
733 outputs.append(output_data)
734 return concat(outputs, dim="species")
736 def _open_single_output_single_model(
737 self, modelname: str, output: str, **settings
738 ) -> GridDataArray | GridDataset:
739 """
740 Opens single output of single model
742 Parameters
743 ----------
744 modelname: str
745 Name of groundwater model from which output should be read.
746 output: str
747 Output variable name to open.
748 **settings:
749 Extra settings that need to be passed through to the respective
750 output function.
751 """
752 open_func = OUTPUT_FUNC_MAPPING[output]
753 expected_modeltype = OUTPUT_MODEL_MAPPING[output]
755 if self.directory is None:
756 raise RuntimeError(f"Simulation {self.name} has not been written yet.")
757 model_path = self.directory / modelname
759 # Get model
760 model = self[modelname]
761 if not isinstance(model, expected_modeltype):
762 raise TypeError(
763 f"{modelname} not a {expected_modeltype}, instead got {type(model)}"
764 )
765 # Get output file path
766 oc_key = model._get_pkgkey("oc")
767 oc_pkg = model[oc_key]
768 # Ensure "-transport" and "-flow" are stripped from "budget"
769 oc_output = output.split("-")[0]
770 output_path = oc_pkg._get_output_filepath(model_path, oc_output)
771 # Force path to always include simulation directory.
772 output_path = self.directory / output_path
774 grb_path = self._get_grb_path(modelname)
776 if not output_path.exists():
777 raise RuntimeError(
778 f"Could not find output in {output_path}, check if you already ran simulation {self.name}"
779 )
781 return open_func(output_path, grb_path, **settings)
783 def _get_flow_modelname_coupled_to_transport_model(
784 self, transport_modelname: str
785 ) -> str:
786 """
787 Get name of flow model coupled to transport model, throws error if
788 multiple flow models are couple to 1 transport model.
789 """
790 exchanges = self.get_exchange_relationships()
791 coupled_flow_models = [
792 i[2]
793 for i in exchanges
794 if (i[3] == transport_modelname) & (i[0] == "GWF6-GWT6")
795 ]
796 if len(coupled_flow_models) != 1:
797 raise ValueError(
798 f"Exactly one flow model must be coupled to transport model {transport_modelname}, got: {coupled_flow_models}"
799 )
800 return coupled_flow_models[0]
802 def _get_grb_path(self, modelname: str) -> Path:
803 """
804 Finds appropriate grb path belonging to modelname. Grb files are not
805 written for transport models, so this method always returns a path to a
806 flowmodel. In case of a transport model, it returns the path to the grb
807 file its coupled flow model.
808 """
809 model = self[modelname]
810 # Get grb path
811 if isinstance(model, GroundwaterTransportModel):
812 flow_model_name = self._get_flow_modelname_coupled_to_transport_model(
813 modelname
814 )
815 flow_model_path = self.directory / flow_model_name
816 else:
817 flow_model_path = self.directory / modelname
819 diskey = model._get_diskey()
820 dis_id = model[diskey]._pkg_id
821 return flow_model_path / f"{diskey}.{dis_id}.grb"
823 @standard_log_decorator()
824 def dump(
825 self, directory=".", validate: bool = True, mdal_compliant: bool = False
826 ) -> None:
827 directory = pathlib.Path(directory)
828 directory.mkdir(parents=True, exist_ok=True)
830 toml_content: DefaultDict[str, dict] = collections.defaultdict(dict)
831 for key, value in self.items():
832 cls_name = type(value).__name__
833 if isinstance(value, Modflow6Model):
834 model_toml_path = value.dump(directory, key, validate, mdal_compliant)
835 toml_content[cls_name][key] = model_toml_path.relative_to(
836 directory
837 ).as_posix()
838 elif key in ["gwtgwf_exchanges", "split_exchanges"]:
839 toml_content[key] = collections.defaultdict(list)
840 for exchange_package in self[key]:
841 exchange_type, filename, _, _ = exchange_package.get_specification()
842 exchange_class_short = type(exchange_package).__name__
843 path = f"{filename}.nc"
844 exchange_package.dataset.to_netcdf(directory / path)
845 toml_content[key][exchange_class_short].append(path)
847 else:
848 path = f"{key}.nc"
849 value.dataset.to_netcdf(directory / path)
850 toml_content[cls_name][key] = path
852 with open(directory / f"{self.name}.toml", "wb") as f:
853 tomli_w.dump(toml_content, f)
855 return
857 @staticmethod
858 def from_file(toml_path):
859 classes = {
860 item_cls.__name__: item_cls
861 for item_cls in (
862 GroundwaterFlowModel,
863 GroundwaterTransportModel,
864 imod.mf6.TimeDiscretization,
865 imod.mf6.Solution,
866 imod.mf6.GWFGWF,
867 imod.mf6.GWFGWT,
868 imod.mf6.GWTGWT,
869 )
870 }
872 toml_path = pathlib.Path(toml_path)
873 with open(toml_path, "rb") as f:
874 toml_content = tomli.load(f)
876 simulation = Modflow6Simulation(name=toml_path.stem)
877 for key, entry in toml_content.items():
878 if key not in ["gwtgwf_exchanges", "split_exchanges"]:
879 item_cls = classes[key]
880 for name, filename in entry.items():
881 path = toml_path.parent / filename
882 simulation[name] = item_cls.from_file(path)
883 else:
884 simulation[key] = []
885 for exchange_class, exchange_list in entry.items():
886 item_cls = classes[exchange_class]
887 for filename in exchange_list:
888 path = toml_path.parent / filename
889 simulation[key].append(item_cls.from_file(path))
891 return simulation
893 def get_exchange_relationships(self):
894 result = []
896 if "gwtgwf_exchanges" in self:
897 for exchange in self["gwtgwf_exchanges"]:
898 result.append(exchange.get_specification())
900 # exchange for splitting models
901 if self.is_split():
902 for exchange in self["split_exchanges"]:
903 result.append(exchange.get_specification())
904 return result
906 def get_models_of_type(self, modeltype):
907 return {
908 k: v
909 for k, v in self.items()
910 if isinstance(v, Modflow6Model) and (v.model_id() == modeltype)
911 }
913 def get_models(self):
914 return {k: v for k, v in self.items() if isinstance(v, Modflow6Model)}
916 def clip_box(
917 self,
918 time_min: Optional[cftime.datetime | np.datetime64 | str] = None,
919 time_max: Optional[cftime.datetime | np.datetime64 | str] = None,
920 layer_min: Optional[int] = None,
921 layer_max: Optional[int] = None,
922 x_min: Optional[float] = None,
923 x_max: Optional[float] = None,
924 y_min: Optional[float] = None,
925 y_max: Optional[float] = None,
926 states_for_boundary: Optional[dict[str, GridDataArray]] = None,
927 ) -> Modflow6Simulation:
928 """
929 Clip a simulation by a bounding box (time, layer, y, x).
931 Slicing intervals may be half-bounded, by providing None:
933 * To select 500.0 <= x <= 1000.0:
934 ``clip_box(x_min=500.0, x_max=1000.0)``.
935 * To select x <= 1000.0: ``clip_box(x_min=None, x_max=1000.0)``
936 or ``clip_box(x_max=1000.0)``.
937 * To select x >= 500.0: ``clip_box(x_min = 500.0, x_max=None.0)``
938 or ``clip_box(x_min=1000.0)``.
940 Parameters
941 ----------
942 time_min: optional
943 time_max: optional
944 layer_min: optional, int
945 layer_max: optional, int
946 x_min: optional, float
947 x_max: optional, float
948 y_min: optional, float
949 y_max: optional, float
950 states_for_boundary : optional, Dict[pkg_name:str, boundary_values:Union[xr.DataArray, xu.UgridDataArray]]
952 Returns
953 -------
954 clipped : Simulation
955 """
957 if self.is_split():
958 raise RuntimeError(
959 "Unable to clip simulation. Clipping can only be done on simulations that haven't been split."
960 + "Therefore clipping should be done before splitting the simulation."
961 )
962 if not self.has_one_flow_model():
963 raise ValueError(
964 "Unable to clip simulation. Clipping can only be done on simulations that have a single flow model ."
965 )
966 for model_name, model in self.get_models().items():
967 supported, error_with_object = model.is_clipping_supported()
968 if not supported:
969 raise ValueError(
970 f"simulation cannot be clipped due to presence of package '{error_with_object}' in model '{model_name}'"
971 )
973 clipped = type(self)(name=self.name)
974 for key, value in self.items():
975 state_for_boundary = (
976 None if states_for_boundary is None else states_for_boundary.get(key)
977 )
978 if isinstance(value, Modflow6Model):
979 clipped[key] = value.clip_box(
980 time_min=time_min,
981 time_max=time_max,
982 layer_min=layer_min,
983 layer_max=layer_max,
984 x_min=x_min,
985 x_max=x_max,
986 y_min=y_min,
987 y_max=y_max,
988 state_for_boundary=state_for_boundary,
989 )
990 elif isinstance(value, Package):
991 clipped[key] = value.clip_box(
992 time_min=time_min,
993 time_max=time_max,
994 layer_min=layer_min,
995 layer_max=layer_max,
996 x_min=x_min,
997 x_max=x_max,
998 y_min=y_min,
999 y_max=y_max,
1000 )
1001 else:
1002 raise ValueError(f"object of type {type(value)} cannot be clipped.")
1003 return clipped
1005 def split(self, submodel_labels: GridDataArray) -> Modflow6Simulation:
1006 """
1007 Split a simulation in different partitions using a submodel_labels array.
1009 The submodel_labels array defines how a simulation will be split. The array should have the same topology as
1010 the domain being split i.e. similar shape as a layer in the domain. The values in the array indicate to
1011 which partition a cell belongs. The values should be zero or greater.
1013 The method return a new simulation containing all the split models and packages
1014 """
1015 if self.is_split():
1016 raise RuntimeError(
1017 "Unable to split simulation. Splitting can only be done on simulations that haven't been split."
1018 )
1020 if not self.has_one_flow_model():
1021 raise ValueError(
1022 "splitting of simulations with more (or less) than 1 flow model currently not supported."
1023 )
1024 transport_models = self.get_models_of_type("gwt6")
1025 flow_models = self.get_models_of_type("gwf6")
1026 if not any(flow_models) and not any(transport_models):
1027 raise ValueError("a simulation without any models cannot be split.")
1029 original_models = {**flow_models, **transport_models}
1030 for model_name, model in original_models.items():
1031 supported, error_with_object = model.is_splitting_supported()
1032 if not supported:
1033 raise ValueError(
1034 f"simulation cannot be split due to presence of package '{error_with_object}' in model '{model_name}'"
1035 )
1037 original_packages = get_packages(self)
1039 partition_info = create_partition_info(submodel_labels)
1041 exchange_creator: ExchangeCreator_Unstructured | ExchangeCreator_Structured
1042 if is_unstructured(submodel_labels):
1043 exchange_creator = ExchangeCreator_Unstructured(
1044 submodel_labels, partition_info
1045 )
1046 else:
1047 exchange_creator = ExchangeCreator_Structured(
1048 submodel_labels, partition_info
1049 )
1051 new_simulation = imod.mf6.Modflow6Simulation(f"{self.name}_partioned")
1052 for package_name, package in {**original_packages}.items():
1053 new_simulation[package_name] = deepcopy(package)
1055 for model_name, model in original_models.items():
1056 solution_name = self.get_solution_name(model_name)
1057 new_simulation[solution_name].remove_model_from_solution(model_name)
1058 for submodel_partition_info in partition_info:
1059 new_model_name = f"{model_name}_{submodel_partition_info.id}"
1060 new_simulation[new_model_name] = slice_model(
1061 submodel_partition_info, model
1062 )
1063 new_simulation[solution_name].add_model_to_solution(new_model_name)
1065 exchanges: list[Any] = []
1067 for flow_model_name, flow_model in flow_models.items():
1068 exchanges += exchange_creator.create_gwfgwf_exchanges(
1069 flow_model_name, flow_model.domain.layer
1070 )
1072 if any(transport_models):
1073 for tpt_model_name in transport_models:
1074 exchanges += exchange_creator.create_gwtgwt_exchanges(
1075 tpt_model_name, flow_model_name, model.domain.layer
1076 )
1077 new_simulation._add_modelsplit_exchanges(exchanges)
1078 new_simulation._update_buoyancy_packages()
1079 new_simulation._set_flow_exchange_options()
1080 new_simulation._set_transport_exchange_options()
1081 new_simulation._update_ssm_packages()
1083 new_simulation._filter_inactive_cells_from_exchanges()
1084 return new_simulation
1086 def regrid_like(
1087 self,
1088 regridded_simulation_name: str,
1089 target_grid: GridDataArray,
1090 validate: bool = True,
1091 ) -> "Modflow6Simulation":
1092 """
1093 This method creates a new simulation object. The models contained in the new simulation are regridded versions
1094 of the models in the input object (this).
1095 Time discretization and solver settings are copied.
1097 Parameters
1098 ----------
1099 regridded_simulation_name: str
1100 name given to the output simulation
1101 target_grid: xr.DataArray or xu.UgridDataArray
1102 discretization onto which the models in this simulation will be regridded
1103 validate: bool
1104 set to true to validate the regridded packages
1106 Returns
1107 -------
1108 a new simulation object with regridded models
1109 """
1111 return _regrid_like(self, regridded_simulation_name, target_grid, validate)
1113 def _add_modelsplit_exchanges(self, exchanges_list: list[GWFGWF]) -> None:
1114 if not self.is_split():
1115 self["split_exchanges"] = []
1116 self["split_exchanges"].extend(exchanges_list)
1118 def _set_flow_exchange_options(self) -> None:
1119 # collect some options that we will auto-set
1120 for exchange in self["split_exchanges"]:
1121 if isinstance(exchange, GWFGWF):
1122 model_name_1 = exchange.dataset["model_name_1"].values[()]
1123 model_1 = self[model_name_1]
1124 exchange.set_options(
1125 save_flows=model_1["oc"].is_budget_output,
1126 dewatered=model_1["npf"].is_dewatered,
1127 variablecv=model_1["npf"].is_variable_vertical_conductance,
1128 xt3d=model_1["npf"].get_xt3d_option(),
1129 newton=model_1.is_use_newton(),
1130 )
1132 def _set_transport_exchange_options(self) -> None:
1133 for exchange in self["split_exchanges"]:
1134 if isinstance(exchange, GWTGWT):
1135 model_name_1 = exchange.dataset["model_name_1"].values[()]
1136 model_1 = self[model_name_1]
1137 advection_key = model_1._get_pkgkey("adv")
1138 dispersion_key = model_1._get_pkgkey("dsp")
1140 scheme = None
1141 xt3d_off = None
1142 xt3d_rhs = None
1143 if advection_key is not None:
1144 scheme = model_1[advection_key].dataset["scheme"].values[()]
1145 if dispersion_key is not None:
1146 xt3d_off = model_1[dispersion_key].dataset["xt3d_off"].values[()]
1147 xt3d_rhs = model_1[dispersion_key].dataset["xt3d_rhs"].values[()]
1148 exchange.set_options(
1149 save_flows=model_1["oc"].is_budget_output,
1150 adv_scheme=scheme,
1151 dsp_xt3d_off=xt3d_off,
1152 dsp_xt3d_rhs=xt3d_rhs,
1153 )
1155 def _filter_inactive_cells_from_exchanges(self) -> None:
1156 for ex in self["split_exchanges"]:
1157 for i in [1, 2]:
1158 self._filter_inactive_cells_exchange_domain(ex, i)
1160 def _filter_inactive_cells_exchange_domain(self, ex: GWFGWF, i: int) -> None:
1161 """Filters inactive cells from one exchange domain inplace"""
1162 modelname = ex[f"model_name_{i}"].values[()]
1163 domain = self[modelname].domain
1165 layer = ex.dataset["layer"] - 1
1166 id = ex.dataset[f"cell_id{i}"] - 1
1167 if is_unstructured(domain):
1168 exchange_cells = {
1169 "layer": layer,
1170 "mesh2d_nFaces": id,
1171 }
1172 else:
1173 exchange_cells = {
1174 "layer": layer,
1175 "y": id.sel({f"cell_dims{i}": f"row_{i}"}),
1176 "x": id.sel({f"cell_dims{i}": f"column_{i}"}),
1177 }
1178 exchange_domain = domain.isel(exchange_cells)
1179 active_exchange_domain = exchange_domain.where(exchange_domain.values > 0)
1180 active_exchange_domain = active_exchange_domain.dropna("index")
1181 ex.dataset = ex.dataset.sel(index=active_exchange_domain["index"])
1183 def get_solution_name(self, model_name: str) -> Optional[str]:
1184 for k, v in self.items():
1185 if isinstance(v, Solution):
1186 if model_name in v.dataset["modelnames"]:
1187 return k
1188 return None
1190 def __repr__(self) -> str:
1191 typename = type(self).__name__
1192 INDENT = " "
1193 attrs = [
1194 f"{typename}(",
1195 f"{INDENT}name={repr(self.name)},",
1196 f"{INDENT}directory={repr(self.directory)}",
1197 ]
1198 items = [
1199 f"{INDENT}{repr(key)}: {type(value).__name__},"
1200 for key, value in self.items()
1201 ]
1202 # Place the emtpy dict on the same line. Looks silly otherwise.
1203 if items:
1204 content = attrs + ["){"] + items + ["}"]
1205 else:
1206 content = attrs + ["){}"]
1207 return "\n".join(content)
1209 def _get_transport_models_per_flow_model(self) -> dict[str, list[str]]:
1210 flow_models = self.get_models_of_type("gwf6")
1211 transport_models = self.get_models_of_type("gwt6")
1212 # exchange for flow and transport
1213 result = collections.defaultdict(list)
1215 for flow_model_name in flow_models:
1216 flow_model = self[flow_model_name]
1217 for tpt_model_name in transport_models:
1218 tpt_model = self[tpt_model_name]
1219 if is_equal(tpt_model.domain, flow_model.domain):
1220 result[flow_model_name].append(tpt_model_name)
1221 return result
1223 def _generate_gwfgwt_exchanges(self) -> list[GWFGWT]:
1224 exchanges = []
1225 flow_transport_mapping = self._get_transport_models_per_flow_model()
1226 for flow_name, tpt_models_of_flow_model in flow_transport_mapping.items():
1227 if len(tpt_models_of_flow_model) > 0:
1228 for transport_model_name in tpt_models_of_flow_model:
1229 exchanges.append(GWFGWT(flow_name, transport_model_name))
1231 return exchanges
1233 def _update_ssm_packages(self) -> None:
1234 flow_transport_mapping = self._get_transport_models_per_flow_model()
1235 for flow_name, tpt_models_of_flow_model in flow_transport_mapping.items():
1236 flow_model = self[flow_name]
1237 for tpt_model_name in tpt_models_of_flow_model:
1238 tpt_model = self[tpt_model_name]
1239 ssm_key = tpt_model._get_pkgkey("ssm")
1240 if ssm_key is not None:
1241 old_ssm_package = tpt_model.pop(ssm_key)
1242 state_variable_name = old_ssm_package.dataset[
1243 "auxiliary_variable_name"
1244 ].values[0]
1245 ssm_package = SourceSinkMixing.from_flow_model(
1246 flow_model, state_variable_name, is_split=self.is_split()
1247 )
1248 if ssm_package is not None:
1249 tpt_model[ssm_key] = ssm_package
1251 def _update_buoyancy_packages(self) -> None:
1252 flow_transport_mapping = self._get_transport_models_per_flow_model()
1253 for flow_name, tpt_models_of_flow_model in flow_transport_mapping.items():
1254 flow_model = self[flow_name]
1255 flow_model.update_buoyancy_package(tpt_models_of_flow_model)
1257 def is_split(self) -> bool:
1258 return "split_exchanges" in self.keys()
1260 def has_one_flow_model(self) -> bool:
1261 flow_models = self.get_models_of_type("gwf6")
1262 return len(flow_models) == 1
1264 def mask_all_models(
1265 self,
1266 mask: GridDataArray,
1267 ):
1268 """
1269 This function applies a mask to all models in a simulation, provided they use
1270 the same discretization. The method parameter "mask" is an idomain-like array.
1271 Masking will overwrite idomain with the mask where the mask is 0 or -1.
1272 Where the mask is 1, the original value of idomain will be kept.
1273 Masking will update the packages accordingly, blanking their input where needed,
1274 and is therefore not a reversible operation.
1276 Parameters
1277 ----------
1278 mask: xr.DataArray, xu.UgridDataArray of ints
1279 idomain-like integer array. 1 sets cells to active, 0 sets cells to inactive,
1280 -1 sets cells to vertical passthrough
1281 """
1282 _mask_all_models(self, mask)