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

1from __future__ import annotations 

2 

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 

10 

11import cftime 

12import dask 

13import jinja2 

14import numpy as np 

15import tomli 

16import tomli_w 

17import xarray as xr 

18import xugrid as xu 

19 

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) 

52 

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} 

59 

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} 

68 

69 

70def get_models(simulation: Modflow6Simulation) -> dict[str, Modflow6Model]: 

71 return {k: v for k, v in simulation.items() if isinstance(v, Modflow6Model)} 

72 

73 

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 } 

80 

81 

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") 

87 

88 def __init__(self, name): 

89 super().__init__() 

90 self.name = name 

91 self.directory = None 

92 self._initialize_template() 

93 

94 def __setitem__(self, key, value): 

95 super().__setitem__(key, value) 

96 

97 def update(self, *args, **kwargs): 

98 for k, v in dict(*args, **kwargs).items(): 

99 self[k] = v 

100 

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) 

108 

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"] 

115 

116 The time discretization in imod-python works as follows: 

117 

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) 

128 

129 Or visually (every letter a date in the time axes): 

130 

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 

135 

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. 

138 

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. 

145 

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. 

150 

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 ) 

163 

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()) 

171 

172 # np.unique also sorts 

173 times = np.unique(np.hstack(times)) 

174 

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 ) 

184 

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") 

207 

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)) 

213 

214 d["models"] = models 

215 if len(models) > 1: 

216 d["exchanges"] = self.get_exchange_relationships() 

217 

218 d["solutiongroups"] = [solutiongroups] 

219 return self._template.render(d) 

220 

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. 

232 

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 

251 

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) 

257 

258 # Generate GWF-GWT exchanges 

259 if gwfgwt_exchanges := self._generate_gwfgwt_exchanges(): 

260 self["gwtgwf_exchanges"] = gwfgwt_exchanges 

261 

262 directory = pathlib.Path(directory) 

263 directory.mkdir(exist_ok=True, parents=True) 

264 

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) 

270 

271 # Write time discretization file 

272 self["time_discretization"].write(directory, "time_discretization") 

273 

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 ) 

303 

304 if status_info.has_errors(): 

305 raise ValidationError("\n" + status_info.to_string()) 

306 

307 self.directory = directory 

308 

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>` 

315 

316 Note that the ``write`` method needs to be called before this method is 

317 called. 

318 

319 Parameters 

320 ---------- 

321 mf6path: Union[str, Path] 

322 Path to the Modflow 6 executable. Defaults to calling ``mf6``. 

323 

324 Examples 

325 -------- 

326 Make sure you write your model first 

327 

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 ) 

340 

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. 

350 

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). 

355 

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. 

376 

377 Returns 

378 ------- 

379 head: Union[xr.DataArray, xu.UgridDataArray] 

380 

381 Examples 

382 -------- 

383 Make sure you write and run your model first 

384 

385 >>> simulation.write(path/to/model) 

386 >>> simulation.run() 

387 

388 Then open heads: 

389 

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 ) 

398 

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. 

408 

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). 

413 

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. 

421 

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"). 

427 

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 ) 

437 

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. 

447 

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). 

452 

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: 

456 

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. 

461 

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). 

466 

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``). 

472 

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"). 

478 

479 Examples 

480 -------- 

481 Make sure you write and run your model first 

482 

483 >>> simulation.write(path/to/model) 

484 >>> simulation.run() 

485 

486 Then open budgets: 

487 

488 >>> budget = simulation.open_flow_budget() 

489 

490 Check the contents: 

491 

492 >>> print(budget.keys()) 

493 

494 Get the drainage budget, compute a time mean for the first layer: 

495 

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 ) 

506 

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. 

517 

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). 

522 

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. 

532 

533 Returns 

534 ------- 

535 concentration: Union[xr.DataArray, xu.UgridDataArray] 

536 

537 Examples 

538 -------- 

539 Make sure you write and run your model first 

540 

541 >>> simulation.write(path/to/model) 

542 >>> simulation.run() 

543 

544 Then open concentrations: 

545 

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 ) 

555 

556 def _open_output(self, output: str, **settings) -> GridDataArray | GridDataset: 

557 """ 

558 Opens output of one or multiple models. 

559 

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 ) 

576 

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 

586 

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") 

609 

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) 

619 

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 

636 

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()) 

647 

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} 

651 

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) 

655 

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 

665 

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.") 

671 

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) 

682 

683 self._pad_missing_variables(cbc_per_partition) 

684 

685 return merge_partitions(cbc_per_partition) 

686 

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()) 

700 

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)) 

704 

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] 

713 

714 if not species_ls: 

715 species_ls = unpartitioned_modelnames 

716 

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 ) 

722 

723 if len(species_ls) == 1: 

724 return self._open_single_output( 

725 list(tpt_names_per_species[0]), output, **settings 

726 ) 

727 

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") 

735 

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 

741 

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] 

754 

755 if self.directory is None: 

756 raise RuntimeError(f"Simulation {self.name} has not been written yet.") 

757 model_path = self.directory / modelname 

758 

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 

773 

774 grb_path = self._get_grb_path(modelname) 

775 

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 ) 

780 

781 return open_func(output_path, grb_path, **settings) 

782 

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] 

801 

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 

818 

819 diskey = model._get_diskey() 

820 dis_id = model[diskey]._pkg_id 

821 return flow_model_path / f"{diskey}.{dis_id}.grb" 

822 

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) 

829 

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) 

846 

847 else: 

848 path = f"{key}.nc" 

849 value.dataset.to_netcdf(directory / path) 

850 toml_content[cls_name][key] = path 

851 

852 with open(directory / f"{self.name}.toml", "wb") as f: 

853 tomli_w.dump(toml_content, f) 

854 

855 return 

856 

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 } 

871 

872 toml_path = pathlib.Path(toml_path) 

873 with open(toml_path, "rb") as f: 

874 toml_content = tomli.load(f) 

875 

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)) 

890 

891 return simulation 

892 

893 def get_exchange_relationships(self): 

894 result = [] 

895 

896 if "gwtgwf_exchanges" in self: 

897 for exchange in self["gwtgwf_exchanges"]: 

898 result.append(exchange.get_specification()) 

899 

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 

905 

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 } 

912 

913 def get_models(self): 

914 return {k: v for k, v in self.items() if isinstance(v, Modflow6Model)} 

915 

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). 

930 

931 Slicing intervals may be half-bounded, by providing None: 

932 

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)``. 

939 

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]] 

951 

952 Returns 

953 ------- 

954 clipped : Simulation 

955 """ 

956 

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 ) 

972 

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 

1004 

1005 def split(self, submodel_labels: GridDataArray) -> Modflow6Simulation: 

1006 """ 

1007 Split a simulation in different partitions using a submodel_labels array. 

1008 

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. 

1012 

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 ) 

1019 

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.") 

1028 

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 ) 

1036 

1037 original_packages = get_packages(self) 

1038 

1039 partition_info = create_partition_info(submodel_labels) 

1040 

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 ) 

1050 

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) 

1054 

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) 

1064 

1065 exchanges: list[Any] = [] 

1066 

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 ) 

1071 

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() 

1082 

1083 new_simulation._filter_inactive_cells_from_exchanges() 

1084 return new_simulation 

1085 

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. 

1096 

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 

1105 

1106 Returns 

1107 ------- 

1108 a new simulation object with regridded models 

1109 """ 

1110 

1111 return _regrid_like(self, regridded_simulation_name, target_grid, validate) 

1112 

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) 

1117 

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 ) 

1131 

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") 

1139 

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 ) 

1154 

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) 

1159 

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 

1164 

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"]) 

1182 

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 

1189 

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) 

1208 

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) 

1214 

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 

1222 

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)) 

1230 

1231 return exchanges 

1232 

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 

1250 

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) 

1256 

1257 def is_split(self) -> bool: 

1258 return "split_exchanges" in self.keys() 

1259 

1260 def has_one_flow_model(self) -> bool: 

1261 flow_models = self.get_models_of_type("gwf6") 

1262 return len(flow_models) == 1 

1263 

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. 

1275 

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)