Coverage for C:\src\imod-python\imod\flow\model.py: 83%

259 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1import abc 

2import collections 

3import os 

4import pathlib 

5import warnings 

6 

7import cftime 

8import jinja2 

9import numpy as np 

10import pandas as pd 

11import xarray as xr 

12 

13import imod 

14from imod.flow.pkgbase import BoundaryCondition 

15from imod.flow.pkggroup import PackageGroups 

16from imod.flow.timeutil import insert_unique_package_times 

17from imod.util.nested_dict import append_nested_dict, initialize_nested_dict 

18from imod.util.time import _compose_timestring, timestep_duration, to_datetime_internal 

19 

20 

21class IniFile(collections.UserDict, abc.ABC): 

22 """ 

23 Some basic support for iMOD ini files here 

24 

25 These files contain the settings that iMOD uses to run its batch 

26 functions. For example to convert its model description -- a projectfile 

27 containing paths to respective .IDFs for each package -- to a Modflow6 

28 model. 

29 """ 

30 

31 # TODO: Create own key mapping to avoid keys like "edate"? 

32 _template = jinja2.Template( 

33 "{%- for key, value in settings %}\n" "{{key}}={{value}}\n" "{%- endfor %}\n" 

34 ) 

35 

36 def _format_datetimes(self): 

37 for timekey in ["sdate", "edate"]: 

38 if timekey in self.keys(): 

39 # If not string assume it is in some kind of datetime format 

40 if type(self[timekey]) is not str: 

41 self[timekey] = _compose_timestring(self[timekey]) 

42 

43 def render(self): 

44 self._format_datetimes() 

45 return self._template.render(settings=self.items()) 

46 

47 

48def _relpath(path, to): 

49 # Wraps os.path.relpath 

50 try: 

51 return pathlib.Path(os.path.relpath(path, to)) 

52 except ValueError: 

53 # Fails to switch between drives e.g. 

54 return pathlib.Path(os.path.abspath(path)) 

55 

56 

57# This class allows only imod packages as values 

58class Model(collections.UserDict): 

59 def __setitem__(self, key, value): 

60 # TODO: raise ValueError on setting certain duplicates 

61 # e.g. two solvers 

62 if self.check == "eager": 

63 value._pkgcheck() 

64 super().__setitem__(key, value) 

65 

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

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

68 self[k] = v 

69 

70 def _delete_empty_packages(self, verbose=False): 

71 to_del = [] 

72 for pkg in self.keys(): 

73 dv = list(self[pkg].dataset.data_vars)[0] 

74 if not self[pkg][dv].notnull().any().compute(): 

75 if verbose: 

76 warnings.warn( 

77 f"Deleting package {pkg}, found no data in parameter {dv}" 

78 ) 

79 to_del.append(pkg) 

80 for pkg in to_del: 

81 del self[pkg] 

82 

83 

84class ImodflowModel(Model): 

85 """ 

86 Class representing iMODFLOW model input. Running it requires iMOD5. 

87 

88 `Download iMOD5 here <https://oss.deltares.nl/web/imod/download-imod5>`_ 

89 

90 Attributes 

91 ---------- 

92 modelname : str check : str, optional 

93 When to perform model checks {None, "defer", "eager"}. Defaults to 

94 "defer". 

95 

96 Examples 

97 -------- 

98 

99 >>> m = Imodflow("example") 

100 >>> m["riv"] = River(...) 

101 >>> # ...etc. 

102 >>> m.create_time_discretization(endtime) 

103 >>> m.write() 

104 """ 

105 

106 # These templates end up here since they require global information 

107 # from more than one package 

108 _PACKAGE_GROUPS = PackageGroups 

109 

110 def __init__(self, modelname, check="defer"): 

111 super().__init__() 

112 self.modelname = modelname 

113 self.check = check 

114 

115 def _get_pkgkey(self, pkg_id): 

116 """ 

117 Get package key that belongs to a certain pkg_id, since the keys are 

118 user specified. 

119 """ 

120 key = [pkgname for pkgname, pkg in self.items() if pkg._pkg_id == pkg_id] 

121 nkey = len(key) 

122 if nkey > 1: 

123 raise ValueError(f"Multiple instances of {key} detected") 

124 elif nkey == 1: 

125 return key[0] 

126 else: 

127 return None 

128 

129 def _group(self): 

130 """ 

131 Group multiple systems of a single package E.g. all river or drainage 

132 sub-systems 

133 """ 

134 groups = collections.defaultdict(dict) 

135 groupable = set(self._PACKAGE_GROUPS.__members__.keys()) 

136 for key, package in self.items(): 

137 pkg_id = package._pkg_id 

138 if pkg_id in groupable: 

139 groups[pkg_id][key] = package 

140 

141 package_groups = [] 

142 for pkg_id, group in groups.items(): 

143 # Create PackageGroup for every package 

144 # RiverGroup for rivers, DrainageGroup for drainage, etc. 

145 package_groups.append(self._PACKAGE_GROUPS[pkg_id].value(**group)) 

146 

147 return package_groups 

148 

149 def _use_cftime(self): 

150 """ 

151 Also checks if datetime types are homogeneous across packages. 

152 """ 

153 types = [] 

154 for pkg in self.values(): 

155 if pkg._hastime(): 

156 types.append(type(np.atleast_1d(pkg["time"].values)[0])) 

157 

158 # Types will be empty if there's no time dependent input 

159 set_of_types = set(types) 

160 if len(set_of_types) == 0: 

161 return None 

162 else: # there is time dependent input 

163 if not len(set_of_types) == 1: 

164 raise ValueError( 

165 f"Multiple datetime types detected: {set_of_types}. " 

166 "Use either cftime or numpy.datetime64[ns]." 

167 ) 

168 # Since we compare types and not instances, we use issubclass 

169 if issubclass(types[0], cftime.datetime): 

170 return True 

171 elif issubclass(types[0], np.datetime64): 

172 return False 

173 else: 

174 raise ValueError("Use either cftime or numpy.datetime64[ns].") 

175 

176 def time_discretization(self, times): 

177 warnings.warn( 

178 f"{self.__class__.__name__}.time_discretization() is deprecated. " 

179 f"In the future call {self.__class__.__name__}.create_time_discretization().", 

180 DeprecationWarning, 

181 ) 

182 self.create_time_discretization(additional_times=times) 

183 

184 def create_time_discretization(self, additional_times): 

185 """ 

186 Collect all unique times from model packages and additional given `times`. These 

187 unique times are used as stress periods in the model. All stress packages must 

188 have the same starting time. 

189 

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

191 

192 - The datetimes of all packages you send in are always respected 

193 - Subsequently, the input data you use is always included fully as well 

194 - All times are treated as starting times for the stress: a stress is 

195 always applied until the next specified date 

196 - For this reason, a final time is required to determine the length of 

197 the last stress period 

198 - Additional times can be provided to force shorter stress periods & 

199 more detailed output 

200 - Every stress has to be defined on the first stress period (this is a 

201 modflow requirement) 

202 

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

204 

205 >>> recharge a - b - c - d - e - f 

206 >>> river g - - - - h - - - - j 

207 >>> times - - - - - - - - - - - i 

208 >>> model a - b - c h d - e - f i 

209 

210 

211 with the stress periods defined between these dates. I.e. the model times are the set of all times you include in the model. 

212 

213 Parameters 

214 ---------- 

215 times : str, datetime; or iterable of str, datetimes. 

216 Times to add to the time discretization. At least one single time 

217 should be given, which will be used as the ending time of the 

218 simulation. 

219 

220 Examples 

221 -------- 

222 Add a single time: 

223 

224 >>> m.create_time_discretization("2001-01-01") 

225 

226 Add a daterange: 

227 

228 >>> m.create_time_discretization(pd.daterange("2000-01-01", "2001-01-01")) 

229 

230 Add a list of times: 

231 

232 >>> m.create_time_discretization(["2000-01-01", "2001-01-01"]) 

233 

234 """ 

235 

236 # Make sure it's an iterable 

237 if not isinstance( 

238 additional_times, (np.ndarray, list, tuple, pd.DatetimeIndex) 

239 ): 

240 additional_times = [additional_times] 

241 

242 # Loop through all packages, check if cftime is required. 

243 self.use_cftime = self._use_cftime() 

244 # use_cftime is None if you no datetimes are present in packages 

245 # use_cftime is False if np.datetimes present in packages 

246 # use_cftime is True if cftime.datetime present in packages 

247 for time in additional_times: 

248 if issubclass(type(time), cftime.datetime): 

249 if self.use_cftime is None: 

250 self.use_cftime = True 

251 if self.use_cftime is False: 

252 raise ValueError( 

253 "Use either cftime or numpy.datetime64[ns]. " 

254 f"Received: {type(time)}." 

255 ) 

256 if self.use_cftime is None: 

257 self.use_cftime = False 

258 

259 times = [ 

260 to_datetime_internal(time, self.use_cftime) for time in additional_times 

261 ] 

262 times, first_times = insert_unique_package_times(self.items(), times) 

263 

264 # Check if every transient package commences at the same time. 

265 for key, first_time in first_times.items(): 

266 time0 = times[0] 

267 if (first_time != time0) and not self[key]._is_periodic(): 

268 raise ValueError( 

269 f"Package {key} does not have a value specified for the " 

270 f"first time: {time0}. Every input must be present in the " 

271 "first stress period. Values are only filled forward in " 

272 "time." 

273 ) 

274 

275 duration = timestep_duration(times, self.use_cftime) 

276 # Generate time discretization, just rely on default arguments 

277 # Probably won't be used that much anyway? 

278 times = np.array(times) 

279 timestep_duration_da = xr.DataArray( 

280 duration, coords={"time": times[:-1]}, dims=("time",) 

281 ) 

282 self["time_discretization"] = imod.flow.TimeDiscretization( 

283 timestep_duration=timestep_duration_da, endtime=times[-1] 

284 ) 

285 

286 def _calc_n_entry(self, composed_package, is_boundary_condition): 

287 """Calculate amount of entries for each timestep and variable.""" 

288 

289 def first(d): 

290 """Get first value of dictionary values""" 

291 return next(iter(d.values())) 

292 

293 if is_boundary_condition: 

294 first_variable = first(first(composed_package)) 

295 n_entry = 0 

296 for sys in first_variable.values(): 

297 n_entry += len(sys) 

298 

299 return n_entry 

300 

301 else: # No time and no systems in regular packages 

302 first_variable = first(composed_package) 

303 return len(first_variable) 

304 

305 def _compose_timestrings(self, globaltimes): 

306 time_format = "%Y-%m-%d %H:%M:%S" 

307 time_composed = self["time_discretization"]._compose_values_time( 

308 "time", globaltimes 

309 ) 

310 time_composed = dict( 

311 [ 

312 (timestep_nr, _compose_timestring(time, time_format=time_format)) 

313 for timestep_nr, time in time_composed.items() 

314 ] 

315 ) 

316 return time_composed 

317 

318 def _compose_periods(self): 

319 periods = {} 

320 

321 for key, package in self.items(): 

322 if package._is_periodic(): 

323 # Periodic stresses are defined for all variables 

324 first_var = list(package.dataset.data_vars)[0] 

325 periods.update(package.dataset[first_var].attrs["stress_periodic"]) 

326 

327 # Create timestrings for "Periods" section in projectfile 

328 # Basically swap around period attributes and compose timestring 

329 # Note that the timeformat for periods in the Projectfile is different 

330 # from that for stress periods 

331 time_format = "%d-%m-%Y %H:%M:%S" 

332 periods_composed = dict( 

333 [ 

334 (value, _compose_timestring(time, time_format=time_format)) 

335 for time, value in periods.items() 

336 ] 

337 ) 

338 return periods_composed 

339 

340 def _compose_all_packages(self, directory, globaltimes): 

341 """ 

342 Compose all transient packages before rendering. 

343 

344 Required because of outer timeloop 

345 

346 Returns 

347 ------- 

348 A tuple with lists of respectively the composed packages and boundary conditions 

349 """ 

350 bndkey = self._get_pkgkey("bnd") 

351 nlayer = self[bndkey]["layer"].size 

352 

353 composition = initialize_nested_dict(5) 

354 

355 group_packages = self._group() 

356 

357 # Get get pkg_id from first value in dictionary in group list 

358 group_pkg_ids = [next(iter(group.values()))._pkg_id for group in group_packages] 

359 

360 for group in group_packages: 

361 group_composition = group.compose( 

362 directory, 

363 globaltimes, 

364 nlayer, 

365 ) 

366 append_nested_dict(composition, group_composition) 

367 

368 for key, package in self.items(): 

369 if package._pkg_id not in group_pkg_ids: 

370 package_composition = package.compose( 

371 directory.joinpath(key), 

372 globaltimes, 

373 nlayer, 

374 ) 

375 append_nested_dict(composition, package_composition) 

376 

377 return composition 

378 

379 def _render_periods(self, periods_composed): 

380 _template_periods = jinja2.Template( 

381 "Periods\n" 

382 "{%- for key, timestamp in periods.items() %}\n" 

383 "{{key}}\n{{timestamp}}\n" 

384 "{%- endfor %}\n" 

385 ) 

386 

387 return _template_periods.render(periods=periods_composed) 

388 

389 def _render_projectfile(self, directory): 

390 """ 

391 Render projectfile. The projectfile has the hierarchy: 

392 package - time - system - layer 

393 """ 

394 diskey = self._get_pkgkey("dis") 

395 globaltimes = self[diskey]["time"].values 

396 

397 content = [] 

398 

399 composition = self._compose_all_packages(directory, globaltimes) 

400 

401 times_composed = self._compose_timestrings(globaltimes) 

402 

403 periods_composed = self._compose_periods() 

404 

405 # Add period strings to times_composed 

406 # These are the strings atop each stress period in the projectfile 

407 times_composed.update({key: key for key in periods_composed.keys()}) 

408 

409 # Add steady-state for packages without time specified 

410 times_composed["steady-state"] = "steady-state" 

411 

412 rendered = [] 

413 ignored = ["dis", "oc"] 

414 

415 for key, package in self.items(): 

416 pkg_id = package._pkg_id 

417 

418 if (pkg_id in rendered) or (pkg_id in ignored): 

419 continue # Skip if already rendered (for groups) or not necessary to render 

420 

421 kwargs = dict( 

422 pkg_id=pkg_id, 

423 name=package.__class__.__name__, 

424 variable_order=package._variable_order, 

425 package_data=composition[pkg_id], 

426 ) 

427 

428 if isinstance(package, BoundaryCondition): 

429 kwargs["n_entry"] = self._calc_n_entry(composition[pkg_id], True) 

430 kwargs["times"] = times_composed 

431 else: 

432 kwargs["n_entry"] = self._calc_n_entry(composition[pkg_id], False) 

433 

434 content.append(package._render_projectfile(**kwargs)) 

435 rendered.append(pkg_id) 

436 

437 # Add periods definition 

438 content.append(self._render_periods(periods_composed)) 

439 

440 return "\n\n".join(content) 

441 

442 def _render_runfile(self, directory): 

443 """ 

444 Render runfile. The runfile has the hierarchy: 

445 time - package - system - layer 

446 """ 

447 raise NotImplementedError("Currently only projectfiles can be rendered.") 

448 

449 def render(self, directory, render_projectfile=True): 

450 """ 

451 Render the runfile as a string, package by package. 

452 """ 

453 if render_projectfile: 

454 return self._render_projectfile(directory) 

455 else: 

456 return self._render_runfile(directory) 

457 

458 def _model_path_management( 

459 self, directory, result_dir, resultdir_is_workdir, render_projectfile 

460 ): 

461 # Coerce to pathlib.Path 

462 directory = pathlib.Path(directory) 

463 if result_dir is None: 

464 result_dir = pathlib.Path("results") 

465 else: 

466 result_dir = pathlib.Path(result_dir) 

467 

468 # Create directories if necessary 

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

470 result_dir.mkdir(exist_ok=True, parents=True) 

471 

472 if render_projectfile: 

473 ext = ".prj" 

474 else: 

475 ext = ".run" 

476 

477 runfilepath = directory / f"{self.modelname}{ext}" 

478 results_runfilepath = result_dir / f"{self.modelname}{ext}" 

479 

480 # Where will the model run? 

481 # Default is inputdir, next to runfile: 

482 # in that case, resultdir is relative to inputdir 

483 # If resultdir_is_workdir, inputdir is relative to resultdir 

484 # render_dir is the inputdir that is printed in the runfile. 

485 # result_dir is the resultdir that is printed in the runfile. 

486 # caching_reldir is from where to check for files. This location 

487 # is the same as the eventual model working dir. 

488 if resultdir_is_workdir: 

489 caching_reldir = result_dir 

490 if not directory.is_absolute(): 

491 render_dir = _relpath(directory, result_dir) 

492 else: 

493 render_dir = directory 

494 result_dir = pathlib.Path(".") 

495 else: 

496 caching_reldir = directory 

497 render_dir = pathlib.Path(".") 

498 if not result_dir.is_absolute(): 

499 result_dir = _relpath(result_dir, directory) 

500 

501 return result_dir, render_dir, runfilepath, results_runfilepath, caching_reldir 

502 

503 def write( 

504 self, 

505 directory=pathlib.Path("."), 

506 result_dir=None, 

507 resultdir_is_workdir=False, 

508 convert_to="mf2005_namfile", 

509 ): 

510 """ 

511 Writes model input files. 

512 

513 Parameters 

514 ---------- 

515 directory : str, pathlib.Path 

516 Directory into which the model input will be written. The model 

517 input will be written into a directory called modelname. 

518 result_dir : str, pathlib.Path 

519 Path to directory in which output will be written when running the 

520 model. Is written as the value of the ``result_dir`` key in the 

521 runfile. See the examples. 

522 resultdir_is_workdir: boolean, optional 

523 Wether the set all input paths in the runfile relative to the output 

524 directory. Because iMOD-wq generates a number of files in its 

525 working directory, it may be advantageous to set the working 

526 directory to a different path than the runfile location. 

527 convert_to: str 

528 The type of object to convert the projectfile to in the 

529 configuration ini file. Should be one of ``["mf2005_namfile", 

530 "mf6_namfile", "runfile"]``. 

531 

532 Returns 

533 ------- 

534 None 

535 

536 Examples 

537 -------- 

538 Say we wish to write the model input to a file called input, and we 

539 desire that when running the model, the results end up in a directory 

540 called output. We may run: 

541 

542 >>> model.write(directory="input", result_dir="output") 

543 

544 And in the ``config_run.ini``, a value of ``../../output`` will be 

545 written for ``result_dir``. This ``config_run.ini`` has to be called 

546 with iMOD 5 to convert the model projectfile to a Modflow 2005 namfile. 

547 To specify a conversion to a runfile, run: 

548 

549 >>> model.write(directory="input", convert_to="runfile") 

550 

551 You can then run the following command to convert the projectfile to a runfile: 

552 

553 >>> path/to/iMOD5.exe ./input/config_run.ini 

554 

555 `Download iMOD5 here <https://oss.deltares.nl/web/imod/download-imod5>`_ 

556 

557 """ 

558 directory = pathlib.Path(directory) 

559 

560 allowed_conversion_settings = ["mf2005_namfile", "mf6_namfile", "runfile"] 

561 if convert_to not in allowed_conversion_settings: 

562 raise ValueError( 

563 f"Got convert_setting: '{convert_to}', should be one of: {allowed_conversion_settings}" 

564 ) 

565 

566 # Currently only supported, no runfile can be directly written by iMOD Python 

567 # TODO: Add runfile support 

568 render_projectfile = True 

569 

570 # TODO: Find a cleaner way to pack and unpack these paths 

571 ( 

572 result_dir, 

573 render_dir, 

574 runfilepath, 

575 results_runfilepath, 

576 caching_reldir, 

577 ) = self._model_path_management( 

578 directory, result_dir, resultdir_is_workdir, render_projectfile 

579 ) 

580 

581 directory = directory.resolve() # Force absolute paths 

582 

583 # TODO 

584 # Check if any caching packages are present, and set necessary states. 

585 # self._set_caching_packages(caching_reldir) 

586 

587 if self.check is not None: 

588 self.package_check() 

589 

590 # TODO Necessary? 

591 # Delete packages without data 

592 # self._delete_empty_packages(verbose=True) 

593 

594 runfile_content = self.render( 

595 directory=directory, render_projectfile=render_projectfile 

596 ) 

597 

598 # Start writing 

599 # Write the runfile 

600 with open(runfilepath, "w") as f: 

601 f.write(runfile_content) 

602 # Also write the runfile in the workdir 

603 if resultdir_is_workdir: 

604 with open(results_runfilepath, "w") as f: 

605 f.write(runfile_content) 

606 

607 # Write iMOD TIM file 

608 diskey = self._get_pkgkey("dis") 

609 time_path = directory / f"{diskey}.tim" 

610 self[diskey].save(time_path) 

611 

612 # Create and write INI file to configure conversion/simulation 

613 ockey = self._get_pkgkey("oc") 

614 bndkey = self._get_pkgkey("bnd") 

615 nlayer = self[bndkey]["layer"].size 

616 

617 if ockey is None: 

618 raise ValueError("No OutputControl was specified for the model") 

619 else: 

620 oc_configuration = self[ockey]._compose_oc_configuration(nlayer) 

621 

622 outfilepath = directory / runfilepath 

623 

624 RUNFILE_OPTIONS = { 

625 "mf2005_namfile": dict( 

626 sim_type=2, namfile_out=outfilepath.with_suffix(".nam") 

627 ), 

628 "runfile": dict(sim_type=1, runfile_out=outfilepath.with_suffix(".run")), 

629 "mf6_namfile": dict( 

630 sim_type=3, namfile_out=outfilepath.with_suffix(".nam") 

631 ), 

632 } 

633 

634 conversion_settings = RUNFILE_OPTIONS[convert_to] 

635 

636 config = IniFile( 

637 function="runfile", 

638 prjfile_in=directory / runfilepath.name, 

639 iss=1, 

640 timfname=directory / time_path.name, 

641 output_folder=result_dir, 

642 **conversion_settings, 

643 **oc_configuration, 

644 ) 

645 config_content = config.render() 

646 

647 with open(directory / "config_run.ini", "w") as f: 

648 f.write(config_content) 

649 

650 # Write all IDFs and IPFs 

651 for pkgname, pkg in self.items(): 

652 if ( 

653 "x" in pkg.dataset.coords and "y" in pkg.dataset.coords 

654 ) or pkg._pkg_id in ["wel", "hfb"]: 

655 try: 

656 pkg.save(directory=directory / pkgname) 

657 except Exception as error: 

658 raise type(error)( 

659 f"{error}/nAn error occured during saving of package: {pkgname}." 

660 ) 

661 

662 def _check_top_bottom(self): 

663 """Check whether bottom of a layer does not exceed a top somewhere.""" 

664 basic_ids = ["top", "bot"] 

665 

666 topkey, botkey = [self._get_pkgkey(pkg_id) for pkg_id in basic_ids] 

667 top, bot = [self[key] for key in (topkey, botkey)] 

668 

669 if (top["top"] < bot["bottom"]).any(): 

670 raise ValueError( 

671 f"top should be larger than bottom in {topkey} and {botkey}" 

672 ) 

673 

674 def package_check(self): 

675 bndkey = self._get_pkgkey("bnd") 

676 active_cells = self[bndkey]["ibound"] != 0 

677 

678 self._check_top_bottom() 

679 

680 for pkg in self.values(): 

681 pkg._pkgcheck(active_cells=active_cells)