Coverage for C:\src\imod-python\imod\mf6\lak.py: 98%

365 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1""" 

2This source file contains the Lake Package and interface objects to the lake 

3package. Usage: create instances of the LakeData class, and optionally 

4instances of the Outlets class, and use the method "from_lakes_and_outlets" to 

5create a lake package. 

6""" 

7 

8import pathlib 

9import textwrap 

10from collections import defaultdict 

11 

12import jinja2 

13import numpy as np 

14import pandas as pd 

15import xarray as xr 

16 

17from imod import mf6 

18from imod.logging import init_log_decorator 

19from imod.mf6.boundary_condition import BoundaryCondition 

20from imod.mf6.package import Package 

21from imod.mf6.pkgbase import PackageBase 

22from imod.mf6.write_context import WriteContext 

23from imod.schemata import AllValueSchema, DimsSchema, DTypeSchema 

24 

25CONNECTION_DIM = "connection_dim" 

26OUTLET_DIM = "outlet_dim" 

27LAKE_DIM = "lake_dim" 

28 

29 

30class LakeApi_Base(PackageBase): 

31 """ 

32 Base class for lake and outlet object. 

33 """ 

34 

35 def __init__(self, allargs): 

36 super().__init__(allargs) 

37 

38 

39class LakeData(LakeApi_Base): 

40 """ 

41 This class is used to initialize the lake package. It contains data 

42 relevant to one lake. The input needed to create an instance of this 

43 consists of a few scalars (name, starting stage), some xarray data-arrays, 

44 and time series. The xarray data-arrays should be of the same size as the 

45 grid, and contain missing values in all locations where the lake does not 

46 exist (similar to the input of the other grid based packages, such a 

47 ConstantHead, River, GeneralHeadBounandary, etc.). 

48 

49 Parameters 

50 ---------- 

51 

52 starting_stage: float, 

53 boundname: str, 

54 connection_type: xr.DataArray of integers. 

55 bed_leak: xr.DataArray of floats. 

56 top_elevation: xr.DataArray of floats. 

57 bot_elevation: xr.DataArray of floats. 

58 connection_length: xr.DataArray of floats. 

59 connection_width: xr.DataArray of floats. 

60 status, 

61 stage: timeseries of float numbers 

62 rainfall: timeseries of float numbers 

63 evaporation: timeseries of float numbers 

64 runoff: timeseries of float numbers 

65 inflow: timeseries of float numbers 

66 withdrawal: timeseries of float numbers 

67 auxiliary: timeseries of float numbers 

68 

69 """ 

70 

71 timeseries_names = [ 

72 "status", 

73 "stage", 

74 "rainfall", 

75 "evaporation", 

76 "runoff", 

77 "inflow", 

78 "withdrawal", 

79 "auxiliary", 

80 ] 

81 

82 def __init__( 

83 self, 

84 starting_stage: float, 

85 boundname: str, 

86 connection_type, 

87 bed_leak=None, 

88 top_elevation=None, 

89 bot_elevation=None, 

90 connection_length=None, 

91 connection_width=None, 

92 status=None, 

93 stage=None, 

94 rainfall=None, 

95 evaporation=None, 

96 runoff=None, 

97 inflow=None, 

98 withdrawal=None, 

99 auxiliary=None, 

100 lake_table=None, 

101 ): 

102 dict_dataset = { 

103 "starting_stage": starting_stage, 

104 "boundname": boundname, 

105 "connection_type": (connection_type.dims, connection_type.values), 

106 "bed_leak": (bed_leak.dims, bed_leak.values), 

107 "top_elevation": (top_elevation.dims, top_elevation.values), 

108 "bottom_elevation": (bot_elevation.dims, bot_elevation.values), 

109 "connection_length": (connection_length.dims, connection_length.values), 

110 "connection_width": (connection_width.dims, connection_width.values), 

111 "lake_table": lake_table, 

112 } 

113 super().__init__(dict_dataset) 

114 

115 # timeseries data 

116 times = [] 

117 timeseries_dict = { 

118 "status": status, 

119 "stage": stage, 

120 "rainfall": rainfall, 

121 "evaporation": evaporation, 

122 "runoff": runoff, 

123 "inflow": inflow, 

124 "withdrawal": withdrawal, 

125 "auxiliary": auxiliary, 

126 } 

127 for _, timeseries in timeseries_dict.items(): 

128 if timeseries is not None: 

129 if "time" in timeseries.coords: 

130 times.extend(list(timeseries.coords["time"].values)) 

131 times = sorted(set(times)) 

132 self.dataset.assign_coords({"time": times}) 

133 for ts_name, ts_object in timeseries_dict.items(): 

134 if ts_object is not None: 

135 fillvalue: float | str = np.nan 

136 if not pd.api.types.is_numeric_dtype(ts_object.dtype): 

137 fillvalue = "" 

138 self.dataset[ts_name] = ts_object.reindex( 

139 {"time": times}, fill_value=fillvalue 

140 ) 

141 else: 

142 self.dataset[ts_name] = None 

143 

144 

145class OutletBase(LakeApi_Base): 

146 """ 

147 Base class for the different kinds of outlets 

148 """ 

149 

150 timeseries_names = ["rate", "invert", "rough", "width", "slope"] 

151 

152 def __init__(self, allargs): 

153 super().__init__(allargs) 

154 

155 

156class OutletManning(OutletBase): 

157 """ 

158 Lake outlet which discharges via a rectangular outlet that uses Manning's 

159 equation. 

160 """ 

161 

162 _couttype = "manning" 

163 

164 def __init__( 

165 self, 

166 lakein: str, 

167 lakeout: str, 

168 invert, 

169 width, 

170 roughness, 

171 slope, 

172 ): 

173 dict_dataset = { 

174 "lakein": lakein, 

175 "lakeout": lakeout, 

176 "invert": invert, 

177 "width": width, 

178 "roughness": roughness, 

179 "slope": slope, 

180 } 

181 super().__init__(dict_dataset) 

182 

183 

184class OutletWeir(OutletBase): 

185 """ 

186 Lake outlet which discharges via a sharp-crested weir. 

187 """ 

188 

189 _couttype = "weir" 

190 

191 def __init__(self, lakein: str, lakeout: str, invert, width): 

192 dict_dataset = { 

193 "lakein": lakein, 

194 "lakeout": lakeout, 

195 "invert": invert, 

196 "width": width, 

197 "roughness": None, 

198 "slope": None, 

199 } 

200 super().__init__(dict_dataset) 

201 

202 

203class OutletSpecified(OutletBase): 

204 """ 

205 Lake outlet which discharges a specified outflow. 

206 """ 

207 

208 _couttype = "specified" 

209 

210 def __init__(self, lakein: str, lakeout: str, rate): 

211 dict_dataset = { 

212 "lakein": lakein, 

213 "lakeout": lakeout, 

214 "invert": None, 

215 "width": None, 

216 "roughness": None, 

217 "slope": None, 

218 "rate": rate, 

219 } 

220 super().__init__(dict_dataset) 

221 

222 

223def create_connection_data(lakes): 

224 connection_vars = [ 

225 "bed_leak", 

226 "top_elevation", 

227 "bottom_elevation", 

228 "connection_length", 

229 "connection_width", 

230 ] 

231 

232 # This will create a statically sized array of strings (dtype="<U10") 

233 claktype_string = np.array( 

234 [ 

235 "horizontal", # 0 

236 "vertical", # 1 

237 "embeddedh", # 2 

238 "embeddedv", # 3 

239 ] 

240 ) 

241 

242 connection_data = defaultdict(list) 

243 cell_ids = [] 

244 for lake in lakes: 

245 notnull = lake["connection_type"].notnull() 

246 indices = np.argwhere( 

247 notnull.values 

248 ) # -1 to convert mf6's 1 based indexing to np's 0-based indexing 

249 xr_indices = { 

250 dim: xr.DataArray(index, dims=("cell_id",)) 

251 for dim, index in zip(notnull.dims, indices.T) 

252 } 

253 

254 # There should be no nodata values in connection_type, so we can use it to index. 

255 type_numeric = ( 

256 lake.dataset["connection_type"].isel(**xr_indices).astype(int).values 

257 ) 

258 type_string = claktype_string[type_numeric] 

259 connection_data["connection_type"].append(type_string) 

260 

261 selection = lake.dataset[connection_vars].isel(**xr_indices) 

262 for var, da in selection.items(): 

263 if not var.startswith("connection_"): 

264 var = f"connection_{var}" 

265 connection_data[var].append(da.values) 

266 

267 # Offset by one since MODFLOW is 1-based! 

268 cell_id = xr.DataArray( 

269 data=indices + 1, 

270 coords={"celldim": list(xr_indices.keys())}, 

271 dims=(CONNECTION_DIM, "celldim"), 

272 ) 

273 cell_ids.append(cell_id) 

274 

275 connection_data = { 

276 k: xr.DataArray(data=np.concatenate(v), dims=[CONNECTION_DIM]) 

277 for k, v in connection_data.items() 

278 } 

279 

280 connection_data["connection_cell_id"] = xr.concat(cell_ids, dim=CONNECTION_DIM) 

281 return connection_data 

282 

283 

284def create_outlet_data(outlets, name_to_number): 

285 outlet_vars = [ 

286 "invert", 

287 "roughness", 

288 "width", 

289 "slope", 

290 ] 

291 outlet_data = defaultdict(list) 

292 for outlet in outlets: 

293 outlet_data["outlet_couttype"].append(outlet._couttype) 

294 

295 # Convert names to numbers 

296 for var in ("lakein", "lakeout"): 

297 name = outlet.dataset[var].item() 

298 if ( 

299 var == "lakeout" and name == "" 

300 ): # the outlet lakeout can be outside of the model 

301 lake_number = 0 

302 else: 

303 try: 

304 lake_number = name_to_number[name] 

305 except KeyError: 

306 names = ", ".join(name_to_number.keys()) 

307 raise KeyError( 

308 f"Outlet lake name {name} not found among lake names: {names}" 

309 ) 

310 outlet_data[f"outlet_{var}"].append(lake_number) 

311 

312 # For other values: fill in NaN if not applicable. 

313 for var in outlet_vars: 

314 if var in outlet.dataset: 

315 if "time" in outlet.dataset[var].dims: 

316 value = 0.0 

317 else: 

318 value = outlet.dataset[var].item() 

319 if value is None: 

320 value = np.nan 

321 else: 

322 value = np.nan 

323 outlet_data[f"outlet_{var}"].append(value) 

324 

325 outlet_data = { 

326 k: xr.DataArray(data=v, dims=[OUTLET_DIM]) for k, v in outlet_data.items() 

327 } 

328 return outlet_data 

329 

330 

331def concatenate_timeseries(list_of_lakes_or_outlets, timeseries_name): 

332 """ 

333 In this function we create a dataarray with a given time coordinate axis. We 

334 add all the timeseries of lakes or outlets with the given name. We also 

335 create a dimension to specify the lake or outlet number. 

336 """ 

337 if list_of_lakes_or_outlets is None: 

338 return None 

339 

340 list_of_dataarrays = [] 

341 list_of_indices = [] 

342 for index, lake_or_outlet in enumerate(list_of_lakes_or_outlets): 

343 if timeseries_name in lake_or_outlet.dataset: 

344 da = lake_or_outlet.dataset[timeseries_name] 

345 if "time" in da.coords: 

346 list_of_dataarrays.append(da) 

347 list_of_indices.append(index + 1) 

348 

349 index = index + 1 

350 if len(list_of_dataarrays) == 0: 

351 return None 

352 fill_value = np.nan 

353 if not pd.api.types.is_numeric_dtype(list_of_dataarrays[0].dtype): 

354 fill_value = "" 

355 out = xr.concat( 

356 list_of_dataarrays, join="outer", dim="index", fill_value=fill_value 

357 ) 

358 out = out.assign_coords(index=list_of_indices) 

359 return out 

360 

361 

362def join_lake_tables(lake_numbers, lakes): 

363 """ 

364 merges all the lake tables into a single dataarray. The lake number corresponding to each table 

365 is added in a new dimension 

366 """ 

367 nr_lakes = len(lakes) 

368 

369 any_lake_table = any(not da_is_none(lake["lake_table"]) for lake in lakes) 

370 if not any_lake_table: 

371 return None 

372 

373 lake_tables = [] 

374 for i in range(nr_lakes): 

375 if not da_is_none(lakes[i]["lake_table"]): 

376 lake_number = lake_numbers[i] 

377 lakes[i]["lake_table"] = lakes[i]["lake_table"].expand_dims( 

378 dim={"laketable_lake_nr": [lake_number]} 

379 ) 

380 lake_tables.append(lakes[i]["lake_table"]) 

381 

382 result = xr.merge(lake_tables, compat="no_conflicts") 

383 return result["lake_table"] 

384 

385 

386def da_is_none(array): 

387 """ 

388 when the None value is appended as a dataarray it gets converted to an array with zero values. 

389 """ 

390 return array is None or array.values[()] is None 

391 

392 

393class Lake(BoundaryCondition): 

394 """ 

395 Lake (LAK) Package. 

396 

397 Parameters 

398 ---------- 

399 lake_number: array of integers (xr.DataArray)- dimension number of lakes: 

400 integer used as identifier for the lake. 

401 lake_starting_stage: array of floats (xr.DataArray)- dimension number of lakes: 

402 starting lake stage. 

403 lake_boundname: array of strings (xr.DataArray)- dimension number of lakes: 

404 name of the lake 

405 

406 connection_lake_number: array of floats (xr.DataArray)- dimension number of connections 

407 lake number for the current lake-to-aquifer connection. 

408 connection_cell_id: array of integers (xr.DataArray)- dimension number of connections 

409 

410 connection_type: array of strings (xr.DataArray)- dimension number of connections 

411 indicates if connection is horizontal, vertical, embeddedH or embeddedV 

412 connection_bed_leak: array of floats (xr.DataArray)- dimension number of connections 

413 defines the bed leakance for the lake-GWF connection. 

414 BEDLEAK must be greater than or equal to zero or specified to be np.nan. If BEDLEAK is specified to 

415 be np.nan, the lake-GWF connection conductance is solely a function of aquifer properties in the 

416 connected GWF cell and lakebed sediments are assumed to be absent. 

417 connection_bottom_elevation: array of floats (xr.DataArray, optional)- dimension number of connections 

418 defines the bottom elevation for a horizontal lake-GWF connection. 

419 If not provided, will be set to the bottom elevation of the cell it is connected to. 

420 connection_top_elevation:array of floats (xr.DataArray, optional)- dimension number of connections 

421 defines the top elevation for a horizontal lake-GWF connection. 

422 If not provided, will be set to the top elevation of the cell it is connected to. 

423 connection_width: array of floats (xr.DataArray, optional) 

424 defines the connection face width for a horizontal lake-GWF connection. 

425 connwidth must be greater than zero for a horizontal lake-GWF connection. Any value can be 

426 specified if claktype is vertical, embeddedh, or embeddedv. If not set, will be set to dx or dy. 

427 connection_length: array of floats (xr.DataArray, optional) 

428 defines the distance between the connected GWF cellid node and the lake 

429 for a horizontal, embeddedh, or embeddedv lake-GWF connection. connlen must be greater than 

430 zero for a horizontal, embeddedh, or embeddedv lake-GWF connection. Any value can be specified 

431 if claktype is vertical. If not set, will be set to dx or dy. 

432 

433 

434 outlet_lakein: array of integers (xr.DataArray, optional) 

435 integer defining the lake number that outlet is connected to. Must be 

436 greater than zero. 

437 outlet_lakeout: array of integers (xr.DataArray, optional) 

438 integer value that defines the lake number that outlet discharge from lake outlet OUTLETNO 

439 is routed to. Must be greater than or equal to zero. 

440 If zero, outlet discharge from lake outlet OUTLETNO is discharged to an external 

441 boundary. 

442 outlet_couttype: array of string (xr.DataArray, optional) 

443 character string that defines the outlet type for the outlet OUTLETNO. Possible 

444 strings include: SPECIFIED–character keyword to indicate the outlet is defined as a specified 

445 flow. MANNING–character keyword to indicate the outlet is defined using Manning’s equation. 

446 WEIR–character keyword to indicate the outlet is defined using a sharp weir equation. 

447 outlet_invert: array of floats (xr.DataArray, optional): 

448 float or character value that defines the invert elevation for the lake outlet. A specified 

449 INVERT value is only used for active lakes if outlet_type for lake outlet OUTLETNO is not 

450 SPECIFIED. 

451 outlet_roughness: array of floats (xr.DataArray, optional) 

452 defines the roughness coefficient for the lake outlet. Any value can be specified 

453 if outlet_type is not MANNING. 

454 outlet_width: array of floats (xr.DataArray, optional) 

455 float or character value that defines the width of the lake outlet. A specified WIDTH value is 

456 only used for active lakes if outlet_type for lake outlet OUTLETNO is not SPECIFIED. 

457 outlet_slope: array of floats (xr.DataArray, optional) 

458 float or character value that defines the bed slope for the lake outlet. A specified SLOPE value is 

459 only used for active lakes if outlet_type for lake outlet OUTLETNO is MANNING. 

460 

461 

462 #time series (lake) 

463 ts_status: array of strings (xr.DataArray, optional) 

464 timeserie used to indicate lake status. Can be ACTIVE, INACTIVE, or CONSTANT. 

465 By default, STATUS is ACTIVE. 

466 ts_stage: array of floats (xr.DataArray, optional) 

467 timeserie used to specify the stage of the lake. The specified STAGE is only applied if 

468 the lake is a constant stage lake 

469 ts_rainfall: array of floats (xr.DataArray, optional) 

470 timeserie used to specify the rainfall rate (LT-1) for the lake. Value must be greater than or equal to zero. 

471 ts_evaporation: array of floats (xr.DataArray, optional) 

472 timeserie used to specify the the maximum evaporation rate (LT-1) for the lake. Value must be greater than or equal to zero. I 

473 ts_runoff: array of floats (xr.DataArray, optional) 

474 timeserie used to specify the the runoff rate (L3 T-1) for the lake. Value must be greater than or equal to zero. 

475 ts_inflow: array of floats (xr.DataArray, optional) 

476 timeserie used to specify the volumetric inflow rate (L3 T-1) for the lake. Value must be greater than or equal to zero. 

477 ts_withdrawal: array of floats (xr.DataArray, optional) 

478 timeserie used to specify the maximum withdrawal rate (L3 T-1) for the lake. Value must be greater than or equal to zero. 

479 ts_auxiliary: array of floats (xr.DataArray, optional) 

480 timeserie used to specify value of auxiliary variables of the lake 

481 

482 #time series (outlet) 

483 ts_rate: array of floats (xr.DataArray, optional) 

484 timeserie used to specify the extraction rate for the lake outflow. A positive value indicates 

485 inflow and a negative value indicates outflow from the lake. RATE only applies to active 

486 (IBOUND > 0) lakes. A specified RATE is only applied if COUTTYPE for the OUTLETNO is SPECIFIED 

487 ts_invert: array of floats (xr.DataArray, optional) 

488 timeserie used to specify the invert elevation for the lake outlet. A specified INVERT value is only used for 

489 active lakes if COUTTYPE for lake outlet OUTLETNO is not SPECIFIED. 

490 ts_rough: array of floats (xr.DataArray, optional) 

491 timeserie used to specify defines the roughness coefficient for the lake outlet. Any value can be 

492 specified if COUTTYPE is not MANNING. 

493 ts_width: array of floats (xr.DataArray, optional) 

494 timeserie used to specify the width of the lake outlet. A specified WIDTH value is only used for active lakes if 

495 COUTTYPE for lake outlet OUTLETNO is not SPECIFIED. 

496 ts_slope: array of floats (xr.DataArray, optional) 

497 timeserie used to specify the bed slope for the lake outlet. A specified SLOPE value is only used for active lakes 

498 if COUTTYPE for lake outlet OUTLETNO is MANNING. 

499 

500 lake_tables: array of floats (xr.DataArray, optional) 

501 the array contains the lake tables- for those lakes that have them 

502 

503 

504 print_input: ({True, False}, optional) 

505 keyword to indicate that the list of constant head information will 

506 be written to the listing file immediately after it is read. Default is 

507 False. 

508 print_stage: ({True, False}, optional) 

509 Keyword to indicate that the list of lake stages will be printed to the listing file for every 

510 stress period in which "HEAD PRINT" is specified in Output Control. If there is no Output Control 

511 option and PRINT_STAGE is specified, then stages are printed for the last time step of each 

512 stress period. 

513 print_flows: ({True, False}, optional) 

514 Indicates that the list of constant head flow rates will be printed to 

515 the listing file for every stress period time step in which "BUDGET 

516 PRINT" is specified in Output Control. If there is no Output Control 

517 option and PRINT FLOWS is specified, then flow rates are printed for the 

518 last time step of each stress period. 

519 Default is False. 

520 save_flows: ({True, False}, optional) 

521 Indicates that constant head flow terms will be written to the file 

522 specified with "BUDGET FILEOUT" in Output Control. Default is False. 

523 

524 

525 stagefile: (String, optional) 

526 name of the binary output file to write stage information. 

527 budgetfile: (String, optional) 

528 name of the binary output file to write budget information. 

529 budgetcsvfile: (String, optional) 

530 name of the comma-separated value (CSV) output file to write budget summary information. 

531 A budget summary record will be written to this file for each time step of the simulation. 

532 package_convergence_filename: (String, optional) 

533 name of the comma spaced values output file to write package convergence information. 

534 ts6_filename: String, optional 

535 defines a time-series file defining time series that can be used to assign time-varying values. 

536 See the "Time-Variable Input" section for instructions on using the time-series capability. 

537 time_conversion: float 

538 value that is used in converting outlet flow terms that use Manning’s equation or gravitational 

539 acceleration to consistent time units. TIME_CONVERSION should be set to 1.0, 60.0, 3,600.0, 

540 86,400.0, and 31,557,600.0 when using time units (TIME_UNITS) of seconds, minutes, hours, days, 

541 or years in the simulation, respectively. CONVTIME does not need to be specified if no lake 

542 outlets are specified or TIME_UNITS are seconds., 

543 length_conversion: float 

544 float value that is used in converting outlet flow terms that use Manning’s equation or gravitational 

545 acceleration to consistent length units. LENGTH_CONVERSION should be set to 3.28081, 1.0, and 100.0 

546 when using length units (LENGTH_UNITS) of feet, meters, or centimeters in the simulation, 

547 respectively. LENGTH_CONVERSION does not need to be specified if no lake outlets are specified or 

548 LENGTH_UNITS are meters. 

549 validate: {True, False} 

550 Flag to indicate whether the package should be validated upon 

551 initialization. This raises a ValidationError if package input is 

552 provided in the wrong manner. Defaults to True. 

553 """ 

554 

555 _pkg_id = "lak" 

556 _template = Package._initialize_template(_pkg_id) 

557 

558 _period_data_lakes = ( 

559 "ts_status", 

560 "ts_stage", 

561 "ts_rainfall", 

562 "ts_evaporation", 

563 "ts_runoff", 

564 "ts_inflow", 

565 "ts_withdrawal", 

566 "ts_auxiliary", 

567 ) 

568 _period_data_outlets = ( 

569 "ts_rate", 

570 "ts_invert", 

571 "ts_rough", 

572 "ts_width", 

573 "ts_slope", 

574 ) 

575 

576 _period_data = _period_data_lakes + _period_data_outlets 

577 

578 _init_schemata = { 

579 "lake_number": [DTypeSchema(np.integer), DimsSchema(LAKE_DIM)], 

580 "lake_starting_stage": [DTypeSchema(np.floating), DimsSchema(LAKE_DIM)], 

581 "lake_boundname": [DTypeSchema(str), DimsSchema(LAKE_DIM)], 

582 "connection_lake_number": [DTypeSchema(np.integer), DimsSchema(CONNECTION_DIM)], 

583 "connection_type": [DTypeSchema(str), DimsSchema(CONNECTION_DIM)], 

584 "connection_bed_leak": [DTypeSchema(np.floating), DimsSchema(CONNECTION_DIM)], 

585 "connection_bottom_elevation": [ 

586 DTypeSchema(np.floating), 

587 DimsSchema(CONNECTION_DIM), 

588 ], 

589 "connection_top_elevation": [ 

590 DTypeSchema(np.floating), 

591 DimsSchema(CONNECTION_DIM), 

592 ], 

593 "connection_width": [DTypeSchema(np.floating), DimsSchema(CONNECTION_DIM)], 

594 "connection_length": [DTypeSchema(np.floating), DimsSchema(CONNECTION_DIM)], 

595 "outlet_lakein": [ 

596 DTypeSchema(np.integer), 

597 DimsSchema(OUTLET_DIM) | DimsSchema(), 

598 ], 

599 "outlet_lakeout": [ 

600 DTypeSchema(np.integer), 

601 DimsSchema(OUTLET_DIM) | DimsSchema(), 

602 ], 

603 "outlet_couttype": [DTypeSchema(str), DimsSchema(OUTLET_DIM) | DimsSchema()], 

604 "outlet_invert": [ 

605 DTypeSchema(np.floating), 

606 DimsSchema(OUTLET_DIM) | DimsSchema(), 

607 ], 

608 "outlet_roughness": [ 

609 DTypeSchema(np.floating), 

610 DimsSchema(OUTLET_DIM) | DimsSchema(), 

611 ], 

612 "outlet_width": [ 

613 DTypeSchema(np.floating), 

614 DimsSchema(OUTLET_DIM) | DimsSchema(), 

615 ], 

616 "outlet_slope": [ 

617 DTypeSchema(np.floating), 

618 DimsSchema(OUTLET_DIM) | DimsSchema(), 

619 ], 

620 "ts_status": [DTypeSchema(str), DimsSchema("index", "time") | DimsSchema()], 

621 "ts_stage": [ 

622 DTypeSchema(np.floating), 

623 DimsSchema("index", "time") | DimsSchema(), 

624 ], 

625 "ts_rainfall": [ 

626 DTypeSchema(np.floating), 

627 DimsSchema("index", "time") | DimsSchema(), 

628 ], 

629 "ts_evaporation": [ 

630 DTypeSchema(np.floating), 

631 DimsSchema("index", "time") | DimsSchema(), 

632 ], 

633 "ts_runoff": [ 

634 DTypeSchema(np.floating), 

635 DimsSchema("index", "time") | DimsSchema(), 

636 ], 

637 "ts_inflow": [ 

638 DTypeSchema(np.floating), 

639 DimsSchema("index", "time") | DimsSchema(), 

640 ], 

641 "ts_withdrawal": [ 

642 DTypeSchema(np.floating), 

643 DimsSchema("index", "time") | DimsSchema(), 

644 ], 

645 "ts_auxiliary": [ 

646 DTypeSchema(np.floating), 

647 DimsSchema("index", "time") | DimsSchema(), 

648 ], 

649 "ts_rate": [ 

650 DTypeSchema(np.floating), 

651 DimsSchema("index", "time") | DimsSchema(), 

652 ], 

653 "ts_invert": [ 

654 DTypeSchema(np.floating), 

655 DimsSchema("index", "time") | DimsSchema(), 

656 ], 

657 "ts_rough": [ 

658 DTypeSchema(np.floating), 

659 DimsSchema("index", "time") | DimsSchema(), 

660 ], 

661 "ts_width": [ 

662 DTypeSchema(np.floating), 

663 DimsSchema("index", "time") | DimsSchema(), 

664 ], 

665 "ts_slope": [ 

666 DTypeSchema(np.floating), 

667 DimsSchema("index", "time") | DimsSchema(), 

668 ], 

669 } 

670 

671 _write_schemata = { 

672 "lake_number": [AllValueSchema(">", 0)], 

673 "connection_lake_number": [AllValueSchema(">", 0)], 

674 "connection_cell_id": [AllValueSchema(">", 0)], 

675 "connection_width": [AllValueSchema(">", 0)], 

676 "connection_length": [AllValueSchema(">", 0)], 

677 "outlet_lakein": [AllValueSchema(">", 0)], 

678 "outlet_lakeout": [AllValueSchema(">", 0)], 

679 "outlet_width": [AllValueSchema(">", 0)], 

680 } 

681 

682 @init_log_decorator() 

683 def __init__( 

684 # lake 

685 self, 

686 lake_number, 

687 lake_starting_stage, 

688 lake_boundname, 

689 # connection 

690 connection_lake_number, 

691 connection_cell_id, 

692 connection_type, 

693 connection_bed_leak, 

694 connection_bottom_elevation, 

695 connection_top_elevation, 

696 connection_width, 

697 connection_length, 

698 # outlet 

699 outlet_lakein=None, 

700 outlet_lakeout=None, 

701 outlet_couttype=None, 

702 outlet_invert=None, 

703 outlet_roughness=None, 

704 outlet_width=None, 

705 outlet_slope=None, 

706 # time series (lake) 

707 ts_status=None, 

708 ts_stage=None, 

709 ts_rainfall=None, 

710 ts_evaporation=None, 

711 ts_runoff=None, 

712 ts_inflow=None, 

713 ts_withdrawal=None, 

714 ts_auxiliary=None, 

715 # time series (outlet) 

716 ts_rate=None, 

717 ts_invert=None, 

718 ts_rough=None, 

719 ts_width=None, 

720 ts_slope=None, 

721 # lake tables 

722 lake_tables=None, 

723 # options 

724 print_input=False, 

725 print_stage=False, 

726 print_flows=False, 

727 save_flows=False, 

728 stagefile=None, 

729 budgetfile=None, 

730 budgetcsvfile=None, 

731 package_convergence_filename=None, 

732 ts6_filename=None, 

733 time_conversion=None, 

734 length_conversion=None, 

735 validate=True, 

736 ): 

737 if ts_status is not None: 

738 ts_status = self._convert_to_string_dataarray(xr.DataArray(ts_status)) 

739 

740 dict_dataset = { 

741 "lake_boundname": lake_boundname, 

742 "lake_number": lake_number, 

743 "lake_starting_stage": lake_starting_stage, 

744 "connection_lake_number": connection_lake_number, 

745 "connection_cell_id": connection_cell_id, 

746 "connection_type": connection_type, 

747 "connection_bed_leak": connection_bed_leak, 

748 "connection_bottom_elevation": connection_bottom_elevation, 

749 "connection_top_elevation": connection_top_elevation, 

750 "connection_width": connection_width, 

751 "connection_length": connection_length, 

752 "outlet_lakein": outlet_lakein, 

753 "outlet_lakeout": outlet_lakeout, 

754 "outlet_couttype": outlet_couttype, 

755 "outlet_invert": outlet_invert, 

756 "outlet_roughness": outlet_roughness, 

757 "outlet_width": outlet_width, 

758 "outlet_slope": outlet_slope, 

759 "print_input": print_input, 

760 "print_stage": print_stage, 

761 "print_flows": print_flows, 

762 "save_flows": save_flows, 

763 "stagefile": stagefile, 

764 "budgetfile": budgetfile, 

765 "budgetcsvfile": budgetcsvfile, 

766 "package_convergence_filename": package_convergence_filename, 

767 "ts6_filename": ts6_filename, 

768 "time_conversion": time_conversion, 

769 "length_conversion": length_conversion, 

770 "ts_status": ts_status, 

771 "ts_stage": ts_stage, 

772 "ts_rainfall": ts_rainfall, 

773 "ts_evaporation": ts_evaporation, 

774 "ts_runoff": ts_runoff, 

775 "ts_inflow": ts_inflow, 

776 "ts_withdrawal": ts_withdrawal, 

777 "ts_auxiliary": ts_auxiliary, 

778 "ts_rate": ts_rate, 

779 "ts_invert": ts_invert, 

780 "ts_rough": ts_rough, 

781 "ts_width": ts_width, 

782 "ts_slope": ts_slope, 

783 "lake_tables": lake_tables, 

784 } 

785 

786 super().__init__(dict_dataset) 

787 

788 self._validate_init_schemata(validate) 

789 

790 @staticmethod 

791 def from_lakes_and_outlets( 

792 lakes, 

793 outlets=None, 

794 print_input=False, 

795 print_stage=False, 

796 print_flows=False, 

797 save_flows=False, 

798 stagefile=None, 

799 budgetfile=None, 

800 budgetcsvfile=None, 

801 package_convergence_filename=None, 

802 time_conversion=None, 

803 length_conversion=None, 

804 ): 

805 package_content = {} 

806 name_to_number = { 

807 lake["boundname"].item(): i + 1 for i, lake in enumerate(lakes) 

808 } 

809 

810 # Package data 

811 lake_numbers = list(name_to_number.values()) 

812 n_connection = [lake["connection_type"].count().values[()] for lake in lakes] 

813 package_content["lake_starting_stage"] = xr.DataArray( 

814 data=[lake["starting_stage"].item() for lake in lakes], 

815 dims=[LAKE_DIM], 

816 ) 

817 package_content["lake_number"] = xr.DataArray( 

818 data=lake_numbers, dims=[LAKE_DIM] 

819 ) 

820 package_content["lake_boundname"] = xr.DataArray( 

821 list(name_to_number.keys()), dims=[LAKE_DIM] 

822 ) 

823 

824 # Connection data 

825 package_content["connection_lake_number"] = xr.DataArray( 

826 data=np.repeat(lake_numbers, n_connection), 

827 dims=[CONNECTION_DIM], 

828 ) 

829 connection_data = create_connection_data(lakes) 

830 package_content.update(connection_data) 

831 for ts_name in Lake._period_data_lakes: 

832 shortname = ts_name[3:] 

833 package_content[ts_name] = concatenate_timeseries(lakes, shortname) 

834 

835 for ts_name in Lake._period_data_outlets: 

836 shortname = ts_name[3:] 

837 package_content[ts_name] = concatenate_timeseries(outlets, shortname) 

838 

839 # Align timeseries variables 

840 ts_vars = Lake._period_data_lakes + Lake._period_data_outlets 

841 assigned_ts_vars = [ 

842 ts_var for ts_var in ts_vars if package_content[ts_var] is not None 

843 ] 

844 aligned_timeseries = xr.align( 

845 *(package_content[ts_var] for ts_var in assigned_ts_vars), join="outer" 

846 ) 

847 package_content.update(dict(zip(assigned_ts_vars, aligned_timeseries))) 

848 

849 if outlets is not None: 

850 outlet_data = create_outlet_data(outlets, name_to_number) 

851 package_content.update(outlet_data) 

852 package_content["print_input"] = print_input 

853 package_content["print_stage"] = print_stage 

854 package_content["print_flows"] = print_flows 

855 package_content["save_flows"] = save_flows 

856 package_content["stagefile"] = stagefile 

857 package_content["budgetfile"] = budgetfile 

858 package_content["budgetcsvfile"] = budgetcsvfile 

859 package_content["package_convergence_filename"] = package_convergence_filename 

860 package_content["time_conversion"] = time_conversion 

861 package_content["length_conversion"] = length_conversion 

862 package_content["lake_tables"] = join_lake_tables(lake_numbers, lakes) 

863 return mf6.Lake(**package_content) 

864 

865 def _has_outlets(self): 

866 # item() will not work here if the object is an array. 

867 # .values[()] will simply return the full numpy array. 

868 outlet_lakein = self.dataset["outlet_lakein"].values[()] 

869 if outlet_lakein is None or ( 

870 np.isscalar(outlet_lakein) and np.isnan(outlet_lakein) 

871 ): 

872 return False 

873 return True 

874 

875 def _has_laketables(self): 

876 # item() will not work here if the object is an array. 

877 # .values[()] will simply return the full numpy array. 

878 tables = self.dataset["lake_tables"].values[()] 

879 if tables is None: 

880 return False 

881 if any(pd.api.types.is_numeric_dtype(t) for t in tables): 

882 return True 

883 return False 

884 

885 def _has_timeseries(self): 

886 for name in self._period_data: 

887 if "time" in self.dataset[name].coords: 

888 if len(self.dataset[name].coords["time"]) > 0: 

889 return True 

890 return False 

891 

892 @classmethod 

893 def _convert_to_string_dataarray(cls, x: xr.DataArray) -> xr.DataArray: 

894 # when adding a string dataarray to a dataset with more coordinates, the 

895 # values for coordinates in the dataset that are not present in the dataarray 

896 # are set to NaN, and the dataarray type changes to obj (because it now has both strings and NaNs) 

897 # This function can be used to convert such a dataarray back to string type. to detect nan's we cannot use np.isnan because that works only on numeric types. 

898 # instead we use the property that any equality check is false for nan's 

899 

900 idx = np.where(x != x) 

901 x[idx[0][:]] = "" 

902 return x.astype(str) 

903 

904 def render(self, directory, pkgname, globaltimes, binary): 

905 d = {} 

906 for var in ( 

907 "print_input", 

908 "print_stage", 

909 "print_flows", 

910 "save_flows", 

911 "stagefile", 

912 "budgetfile", 

913 "budgetcsvfile", 

914 "package_convergence_filename", 

915 "ts6_filename", 

916 "time_conversion", 

917 "length_conversion", 

918 ): 

919 value = self[var].item() 

920 if self._valid(value): 

921 d[var] = value 

922 

923 d["nlakes"] = len(self.dataset["lake_number"]) 

924 d["noutlets"] = 0 

925 if self._has_outlets(): 

926 d["noutlets"] = len(self.dataset["outlet_lakein"]) 

927 

928 d["ntables"] = 0 

929 if self._has_laketables(): 

930 d["ntables"] = len(self.dataset["lake_tables"].coords["laketable_lake_nr"]) 

931 

932 packagedata = [] 

933 for name, number, stage in zip( 

934 self.dataset["lake_boundname"], 

935 self.dataset["lake_number"], 

936 self.dataset["lake_starting_stage"], 

937 ): 

938 nconn = (self.dataset["connection_lake_number"] == number).sum() 

939 row = tuple(a.item() for a in (number, stage, nconn, name)) 

940 packagedata.append(row) 

941 d["packagedata"] = packagedata 

942 

943 return self._template.render(d) 

944 

945 def _get_iconn(self, lake_numbers_per_connection): 

946 iconn = np.full_like(lake_numbers_per_connection, dtype=np.int64, fill_value=0) 

947 maxlake = lake_numbers_per_connection.max() 

948 connections_per_lake = np.zeros(maxlake + 1) 

949 for i in range(np.size(lake_numbers_per_connection)): 

950 lakeno = lake_numbers_per_connection[i] 

951 connections_per_lake[lakeno] += 1 

952 iconn[i] = connections_per_lake[lakeno] 

953 return iconn 

954 

955 def _connection_dataframe(self) -> pd.DataFrame: 

956 connection_vars = [ 

957 "connection_lake_number", 

958 "connection_type", 

959 "connection_bed_leak", 

960 "connection_bottom_elevation", 

961 "connection_top_elevation", 

962 "connection_width", 

963 "connection_length", 

964 ] 

965 data_df = self.dataset[connection_vars].to_dataframe() 

966 lake_numbers_per_connection = self.dataset["connection_lake_number"].values 

967 

968 data_df["iconn"] = self._get_iconn(lake_numbers_per_connection) 

969 cell_id_df = self.dataset["connection_cell_id"].to_pandas() 

970 order = ( 

971 ["connection_lake_number", "iconn"] 

972 + list(cell_id_df.columns) 

973 + connection_vars[1:] 

974 ) 

975 return pd.concat([data_df, cell_id_df], axis=1)[order] 

976 

977 def _outlet_dataframe(self) -> pd.DataFrame: 

978 outlet_vars = [ 

979 "outlet_lakein", 

980 "outlet_lakeout", 

981 "outlet_couttype", 

982 "outlet_invert", 

983 "outlet_roughness", 

984 "outlet_width", 

985 "outlet_slope", 

986 ] 

987 df = self.dataset[outlet_vars].to_dataframe() 

988 return df 

989 

990 def write_blockfile(self, pkgname, globaltimes, write_context: WriteContext): 

991 renderdir = pathlib.Path(write_context.write_directory.stem) 

992 content = self.render( 

993 directory=renderdir, 

994 pkgname=pkgname, 

995 globaltimes=globaltimes, 

996 binary=write_context.use_binary, 

997 ) 

998 filename = write_context.write_directory / f"{pkgname}.{self._pkg_id}" 

999 with open(filename, "w") as f: 

1000 f.write(content) 

1001 f.write("\n") 

1002 

1003 self._write_table_section( 

1004 f, 

1005 self._connection_dataframe(), 

1006 "connectiondata", 

1007 ) 

1008 if self._has_laketables(): 

1009 lake_number_to_filename = self._write_laketable_filelist_section(f) 

1010 self._write_laketable_files( 

1011 write_context.write_directory, lake_number_to_filename 

1012 ) 

1013 

1014 if self._has_outlets(): 

1015 f.write("\n") 

1016 self._write_table_section(f, self._outlet_dataframe(), "outlets") 

1017 

1018 if self._has_timeseries(): 

1019 self._write_period_section(f, globaltimes) 

1020 

1021 return 

1022 

1023 def _write_period_section(self, f, globaltimes): 

1024 class Period_internal: 

1025 """ 

1026 The Period_internal class is used for rendering the lake package in jinja. 

1027 There is no point in instantiating this class as a user. 

1028 """ 

1029 

1030 def __init__(self, period_number): 

1031 self.period_number = period_number 

1032 self.nr_values = 0 

1033 self.lake_or_outlet_number = [] 

1034 self.series_name = [] 

1035 self.value = [] 

1036 

1037 period_data_list = [] 

1038 

1039 period_data_name_list = list(self._period_data) 

1040 timeseries_dataset = self.dataset[period_data_name_list] 

1041 timeseries_times = self.dataset.coords["time"] 

1042 iperiods = np.searchsorted(globaltimes, timeseries_times) + 1 

1043 for iperiod, (_, period_data) in zip( 

1044 iperiods, timeseries_dataset.groupby("time") 

1045 ): 

1046 period_data_list.append(Period_internal(iperiod)) 

1047 for tssname in self._period_data: 

1048 if len(period_data[tssname].dims) > 0: 

1049 for index in period_data.coords["index"].values: 

1050 value = period_data[tssname].sel(index=index).values[()] 

1051 isvalid = False 

1052 if isinstance(value, str): 

1053 isvalid = value != "" 

1054 else: 

1055 isvalid = ~np.isnan(value) 

1056 

1057 if isvalid: 

1058 period_data_list[-1].nr_values += 1 

1059 period_data_list[-1].lake_or_outlet_number.append(index) 

1060 period_data_list[-1].series_name.append(tssname[3:]) 

1061 period_data_list[-1].value.append(value) 

1062 

1063 _template = jinja2.Template( 

1064 textwrap.dedent( 

1065 """ 

1066 {% if nperiod > 0 -%} 

1067 {% for iperiod in range(0, nperiod) %} 

1068 {% if periods[iperiod].nr_values > 0 -%} 

1069 begin period {{periods[iperiod].period_number}}{% for ivalue in range(0, periods[iperiod].nr_values) %} 

1070 {{periods[iperiod].lake_or_outlet_number[ivalue]}} {{periods[iperiod].series_name[ivalue]}} {{periods[iperiod].value[ivalue]}}{% endfor %} 

1071 end period 

1072 {% endif -%} 

1073 {% endfor -%} 

1074 {% endif -%}""" 

1075 ) 

1076 ) 

1077 

1078 d = {} 

1079 d["nperiod"] = len(period_data_list) 

1080 d["periods"] = period_data_list 

1081 

1082 period_block = _template.render(d) 

1083 f.write(period_block) 

1084 

1085 def _package_data_to_sparse(self): 

1086 return 

1087 

1088 def fill_stress_perioddata(self): 

1089 # this function is called from packagebase and should do nothing in this context 

1090 return 

1091 

1092 def _write_perioddata(self, directory, pkgname, binary): 

1093 # this function is called from packagebase and should do nothing in this context 

1094 return 

1095 

1096 def _write_laketable_filelist_section( 

1097 self, 

1098 f, 

1099 ): 

1100 """ 

1101 Writes a section of the lake package input file which lists the referenced laketable files. 

1102 Returns a dictionary with as key the lake number and as value the laketable filename- for those 

1103 lakes that have a laketable. 

1104 """ 

1105 template = jinja2.Template( 

1106 textwrap.dedent( 

1107 """\n\ 

1108 begin tables{% for lakenumber, lakefile in file_dict.items() %} 

1109 {{lakenumber}} TAB6 FILEIN {{lakefile}}{% endfor %} 

1110 end tables\n 

1111 """ 

1112 ) 

1113 ) 

1114 

1115 lake_number_to_lake_table_filename = {} 

1116 d = {} 

1117 

1118 for name, number in zip( 

1119 self.dataset["lake_boundname"], 

1120 self.dataset["lake_number"], 

1121 ): 

1122 lake_number = number.values[()] 

1123 lake_name = name.values[()] 

1124 

1125 if ( 

1126 lake_number 

1127 in self.dataset["lake_tables"].coords["laketable_lake_nr"].values 

1128 ): 

1129 table_file = lake_name + ".ltbl" 

1130 lake_number_to_lake_table_filename[lake_number] = table_file 

1131 

1132 d["file_dict"] = lake_number_to_lake_table_filename 

1133 tables_block = template.render(d) 

1134 f.write(tables_block) 

1135 

1136 return lake_number_to_lake_table_filename 

1137 

1138 def _write_laketable_files(self, directory, lake_number_to_filename): 

1139 """ 

1140 writes a laketable file, containing a table which specifies the relation between stage, 

1141 volume and area for one single lake. 

1142 """ 

1143 

1144 template = Package._initialize_template("laketable") 

1145 

1146 for num, file in lake_number_to_filename.items(): 

1147 d = {} 

1148 table = self.dataset["lake_tables"].sel( 

1149 { 

1150 "laketable_lake_nr": num, 

1151 } 

1152 ) 

1153 

1154 # count number of rows 

1155 stage_col = table.sel({"column": "stage"}) 

1156 d["nrow"] = ( 

1157 stage_col.where(pd.api.types.is_numeric_dtype).count().values[()] 

1158 ) 

1159 

1160 # check if the barea column is present for this table (and not filled with nan's) 

1161 has_barea_column = "barea" in table.coords["column"] 

1162 if has_barea_column: 

1163 barea_column = table.sel({"column": "barea"}) 

1164 has_barea_column = ( 

1165 barea_column.where(pd.api.types.is_numeric_dtype).count().values[()] 

1166 > 0 

1167 ) 

1168 

1169 columns = ["stage", "sarea", "volume"] 

1170 if has_barea_column: 

1171 columns.append("barea") 

1172 d["ncol"] = len(columns) 

1173 table_dataframe = pd.DataFrame(table.sel({"column": columns}).transpose()) 

1174 

1175 string_table = table_dataframe.iloc[ 

1176 range(d["nrow"]), range(d["ncol"]) 

1177 ].to_csv( 

1178 header=False, 

1179 index=False, 

1180 sep=" ", 

1181 lineterminator="\n", 

1182 ) 

1183 

1184 d["table"] = string_table 

1185 # write lake table to file 

1186 fullpath_laketable = directory / file 

1187 laketable_file = template.render(d) 

1188 with open(fullpath_laketable, "w") as f: 

1189 f.write(laketable_file) 

1190 

1191 def _write_table_section( 

1192 self, f, dataframe: pd.DataFrame, title: str, index: bool = False 

1193 ) -> None: 

1194 """ 

1195 writes a dataframe to file. Used for the connection data and for the outlet data. 

1196 """ 

1197 f.write(f"begin {title}\n") 

1198 block = dataframe.to_csv( 

1199 index=index, 

1200 header=False, 

1201 sep=" ", 

1202 lineterminator="\n", 

1203 ) 

1204 trimmedlines = [line.strip() for line in block.splitlines()] 

1205 trimmedblock = "\n".join(map(str, trimmedlines)) + "\n" 

1206 f.write(trimmedblock) 

1207 f.write(f"end {title}\n") 

1208 return 

1209 

1210 def is_splitting_supported(self) -> bool: 

1211 return False 

1212 

1213 def is_regridding_supported(self) -> bool: 

1214 return False 

1215 

1216 def is_clipping_supported(self) -> bool: 

1217 return False