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

1"""

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

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

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

5 create a lake package.

6"""

7

8 import pathlib

9 import textwrap

10 from collections import defaultdict

11

12 import jinja2

13 import numpy as np

14 import pandas as pd

15 import xarray as xr

16

17 from imod import mf6

18 from imod.logging import init_log_decorator

19 from imod.mf6.boundary_condition import BoundaryCondition

20 from imod.mf6.package import Package

21 from imod.mf6.pkgbase import PackageBase

22 from imod.mf6.write_context import WriteContext

23 from imod.schemata import AllValueSchema, DimsSchema, DTypeSchema

24

25 CONNECTION_DIM = "connection_dim"

26 OUTLET_DIM = "outlet_dim"

27 LAKE_DIM = "lake_dim"

28

29

30 class 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

39 class 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

145 class 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

156 class 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

184 class 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

203 class 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

223 def 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

284 def 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

331 def 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

362 def 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

386 def 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

393 class 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