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

1 import abc

2 import pathlib

3 import warnings

4 from copy import copy, deepcopy

5 from typing import Mapping, Optional, Union

6

7 import numpy as np

8 import xarray as xr

9 import xugrid as xu

10

11 from imod.mf6.auxiliary_variables import (

12 expand_transient_auxiliary_variables,

13 get_variable_names,

14)

15 from imod.mf6.package import Package

16 from imod.mf6.utilities.package import get_repeat_stress

17 from imod.mf6.write_context import WriteContext

18 from imod.typing.grid import GridDataArray

19

20

21 def _dis_recarr(arrdict, layer, notnull):

22 # Define the numpy structured array dtype

23 index_spec = [("layer", np.int32), ("row", np.int32), ("column", np.int32)]

24 field_spec = [(key, np.float64) for key in arrdict]

25 sparse_dtype = np.dtype(index_spec + field_spec)

26 # Initialize the structured array

27 nrow = notnull.sum()

28 recarr = np.empty(nrow, dtype=sparse_dtype)

29 # Fill in the indices

30 if notnull.ndim == 2:

31 recarr["row"], recarr["column"] = (np.argwhere(notnull) + 1).transpose()

32 recarr["layer"] = layer

33 else:

34 ilayer, irow, icolumn = np.argwhere(notnull).transpose()

35 recarr["row"] = irow + 1

36 recarr["column"] = icolumn + 1

37 recarr["layer"] = layer[ilayer]

38 return recarr

39

40

41 def _disv_recarr(arrdict, layer, notnull):

42 # Define the numpy structured array dtype

43 index_spec = [("layer", np.int32), ("cell2d", np.int32)]

44 field_spec = [(key, np.float64) for key in arrdict]

45 sparse_dtype = np.dtype(index_spec + field_spec)

46 # Initialize the structured array

47 nrow = notnull.sum()

48 recarr = np.empty(nrow, dtype=sparse_dtype)

49 # Fill in the indices

50 if notnull.ndim == 1 and layer.size == 1:

51 recarr["cell2d"] = (np.argwhere(notnull) + 1).transpose()

52 recarr["layer"] = layer

53 else:

54 ilayer, icell2d = np.argwhere(notnull).transpose()

55 recarr["cell2d"] = icell2d + 1

56 recarr["layer"] = layer[ilayer]

57 return recarr

58

59

60 class BoundaryCondition(Package, abc.ABC):

61 """

62 BoundaryCondition is used to share methods for specific stress packages

63 with a time component.

64

65 It is not meant to be used directly, only to inherit from, to implement new

66 packages.

67

68 This class only supports `list input

69 <https://water.usgs.gov/water-resources/software/MODFLOW-6/mf6io_6.0.4.pdf#page=19>`_,

70 not the array input which is used in :class:`Package`.

71 """

72

73 def __init__(self, allargs: Mapping[str, GridDataArray | float | int | bool | str]):

74 super().__init__(allargs)

75 if "concentration" in allargs.keys() and allargs["concentration"] is None:

76 # Remove vars inplace

77 del self.dataset["concentration"]

78 del self.dataset["concentration_boundary_type"]

79 else:

80 expand_transient_auxiliary_variables(self)

81

82 def set_repeat_stress(self, times: dict[np.datetime64, np.datetime64]) -> None:

83 """

84 Set repeat stresses: re-use data of earlier periods.

85

86 Parameters

87 ----------

88 times: Dict of datetime-like to datetime-like.

89 The data of the value datetime is used for the key datetime.

90 """

91 warnings.warn(

92 f"""{self.__class__.__name__}.set_repeat_stress(...) is deprecated.

93 In the future, add repeat stresses as constructor parameters. An

94 object containing them can be created using 'get_repeat_stress', as

95 follows:

96

97 from imod.mf6.utilities.package_utils import get_repeat_stress

98

99 repeat_stress = get_repeat_stress(repeat_periods) # args before provided to River.set_repeat_stress

100 riv = imod.mf6.River(..., repeat_stress=repeat_stress)

101

102 Note that the location of get_repeat_stress (imod.mf6.utilities.package_utils)

103 may change in the future

104 """,

105 DeprecationWarning,

106 )

107

108 self.dataset["repeat_stress"] = get_repeat_stress(times)

109

110 def _max_active_n(self):

111 """

112 Determine the maximum active number of cells that are active

113 during a stress period.

114 """

115 da = self.dataset[self.get_period_varnames()[0]]

116 if "time" in da.coords:

117 nmax = int(da.groupby("time").count(xr.ALL_DIMS).max())

118 else:

119 nmax = int(da.count())

120 return nmax

121

122 def _write_binaryfile(self, outpath, struct_array):

123 with open(outpath, "w") as f:

124 struct_array.tofile(f)

125

126 def _write_textfile(self, outpath, struct_array):

127 fields = struct_array.dtype.fields

128 fmt = [self._number_format(field[0]) for field in fields.values()]

129 header = " ".join(list(fields.keys()))

130 with open(outpath, "w") as f:

131 np.savetxt(fname=f, X=struct_array, fmt=fmt, header=header)

132

133 def _write_datafile(self, outpath, ds, binary):

134 """

135 Writes a modflow6 binary data file

136 """

137 layer = ds["layer"].values if "layer" in ds.coords else None

138 arrdict = self._ds_to_arrdict(ds)

139 struct_array = self._to_struct_array(arrdict, layer)

140 outpath.parent.mkdir(exist_ok=True, parents=True)

141 if binary:

142 self._write_binaryfile(outpath, struct_array)

143 else:

144 self._write_textfile(outpath, struct_array)

145

146 def _ds_to_arrdict(self, ds):

147 for datavar in ds.data_vars:

148 if ds[datavar].shape == ():

149 raise ValueError(

150 f"{datavar} in {self._pkg_id} package cannot be a scalar"

151 )

152

153 arrdict = {}

154 for datavar in ds.data_vars:

155 arrdict[datavar] = ds[datavar].values

156

157 return arrdict

158

159 def _to_struct_array(self, arrdict, layer):

160 """Convert from dense arrays to list based input"""

161 # TODO stream the data per stress period

162 # TODO add pkgcheck that period table aligns

163 # Get the number of valid values

164 if layer is None:

165 raise ValueError("Layer should be provided")

166

167 data = next(iter(arrdict.values()))

168 notnull = ~np.isnan(data)

169

170 if isinstance(self.dataset, xr.Dataset):

171 recarr = _dis_recarr(arrdict, layer, notnull)

172 elif isinstance(self.dataset, xu.UgridDataset):

173 recarr = _disv_recarr(arrdict, layer, notnull)

174 else:

175 raise TypeError(

176 "self.dataset should be xarray.Dataset or xugrid.UgridDataset,"

177 f" is {type(self.dataset)} instead"

178 )

179 # Fill in the data

180 for key, arr in arrdict.items():

181 values = arr[notnull].astype(np.float64)

182 recarr[key] = values

183

184 return recarr

185

186 def _period_paths(self, directory, pkgname, globaltimes, bin_ds, binary):

187 directory = pathlib.Path(directory) / pkgname

188

189 if binary:

190 ext = "bin"

191 else:

192 ext = "dat"

193

194 periods = {}

195 if "time" in bin_ds: # one of bin_ds has time

196 package_times = bin_ds.coords["time"].values

197 starts = np.searchsorted(globaltimes, package_times) + 1

198 for i, start in enumerate(starts):

199 path = directory / f"{self._pkg_id}-{i}.{ext}"

200 periods[start] = path.as_posix()

201

202 repeat_stress = self.dataset.get("repeat_stress")

203 if repeat_stress is not None and repeat_stress.values[()] is not None:

204 keys = repeat_stress.isel(repeat_items=0).values

205 values = repeat_stress.isel(repeat_items=1).values

206 repeat_starts = np.searchsorted(globaltimes, keys) + 1

207 values_index = np.searchsorted(globaltimes, values) + 1

208 for i, start in zip(values_index, repeat_starts):

209 periods[start] = periods[i]

210 # Now make sure the periods are sorted by key.

211 periods = dict(sorted(periods.items()))

212 else:

213 path = directory / f"{self._pkg_id}.{ext}"

214 periods[1] = path.as_posix()

215

216 return periods

217

218 def _get_options(

219 self, predefined_options: dict, not_options: Optional[list] = None

220 ):

221 options = copy(predefined_options)

222

223 if not_options is None:

224 not_options = self.get_period_varnames()

225

226 for varname in self.dataset.data_vars.keys(): # pylint:disable=no-member

227 if varname in not_options:

228 continue

229 v = self.dataset[varname].values[()]

230 if self._valid(v): # skip None and False

231 options[varname] = v

232 return options

233

234 def _get_bin_ds(self):

235 """

236 Get binary dataset data for stress periods, this data will be written to

237 datafiles. This method can be overriden to do some extra operations on

238 this dataset before writing.

239 """

240 return self[self.get_period_varnames()]

241

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

243 """Render fills in the template only, doesn't write binary data"""

244 d = {"binary": binary}

245 bin_ds = self._get_bin_ds()

246 d["periods"] = self._period_paths(

247 directory, pkgname, globaltimes, bin_ds, binary

248 )

249 # construct the rest (dict for render)

250 d = self._get_options(d)

251 d["maxbound"] = self._max_active_n()

252

253 if (hasattr(self, "_auxiliary_data")) and (names := get_variable_names(self)):

254 d["auxiliary"] = names

255

256 return self._template.render(d)

257

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

259 if len(self.get_period_varnames()) == 0:

260 return

261 bin_ds = self._get_bin_ds()

262

263 if binary:

264 ext = "bin"

265 else:

266 ext = "dat"

267

268 if "time" in bin_ds: # one of bin_ds has time

269 for i in range(len(self.dataset.time)):

270 path = directory / pkgname / f"{self._pkg_id}-{i}.{ext}"

271 self._write_datafile(

272 path, bin_ds.isel(time=i), binary=binary

273 ) # one timestep

274 else:

275 path = directory / pkgname / f"{self._pkg_id}.{ext}"

276 self._write_datafile(path, bin_ds, binary=binary)

277

278 def write(

279 self,

280 pkgname: str,

281 globaltimes: Union[list[np.datetime64], np.ndarray],

282 write_context: WriteContext,

283 ):

284 """

285 writes the blockfile and binary data

286

287 directory is modelname

288 """

289

290 super().write(pkgname, globaltimes, write_context)

291 directory = write_context.write_directory

292

293 self._write_perioddata(

294 directory=directory,

295 pkgname=pkgname,

296 binary=write_context.use_binary,

297 )

298

299 def get_period_varnames(self):

300 result = []

301 if hasattr(self, "_period_data"):

302 result.extend(self._period_data)

303 if hasattr(self, "_auxiliary_data"):

304 result.extend(get_variable_names(self))

305

306 return result

307

308

309 class AdvancedBoundaryCondition(BoundaryCondition, abc.ABC):

310 """

311 Class dedicated to advanced boundary conditions, since MF6 does not support

312 binary files for Advanced Boundary conditions.

313

314 The advanced boundary condition packages are: "uzf", "lak", "maw", "sfr".

315

316 """

317

318 def _get_field_spec_from_dtype(self, recarr):

319 """

320 From https://stackoverflow.com/questions/21777125/how-to-output-dtype-to-a-list-or-dict

321 """

322 return [

323 (x, y[0])

324 for x, y in sorted(recarr.dtype.fields.items(), key=lambda k: k[1])

325 ]

326

327 def _write_file(self, outpath, sparse_data):

328 """

329 Write to textfile, which is necessary for Advanced Stress Packages

330 """

331 fields = sparse_data.dtype.fields

332 fmt = [self._number_format(field[0]) for field in fields.values()]

333 header = " ".join(list(fields.keys()))

334 np.savetxt(fname=outpath, X=sparse_data, fmt=fmt, header=header)

335

336 @abc.abstractmethod

337 def _package_data_to_sparse(self):

338 """

339 Get packagedata, override with function for the advanced boundary

340 condition in particular

341 """

342 return

343

344 def write_packagedata(self, directory, pkgname, binary):

345 outpath = directory / pkgname / f"{self._pkg_id}-pkgdata.dat"

346 outpath.parent.mkdir(exist_ok=True, parents=True)

347 package_data = self._package_data_to_sparse()

348 self._write_file(outpath, package_data)

349

350 def write(

351 self,

352 pkgname: str,

353 globaltimes: Union[list[np.datetime64], np.ndarray],

354 write_context: WriteContext,

355 ):

356 boundary_condition_write_context = deepcopy(write_context)

357 boundary_condition_write_context.use_binary = False

358

359 self.fill_stress_perioddata()

360 super().write(pkgname, globaltimes, boundary_condition_write_context)

361

362 directory = boundary_condition_write_context.write_directory

363 self.write_packagedata(directory, pkgname, binary=False)

364

365 @abc.abstractmethod

366 def fill_stress_perioddata(self):

367 raise NotImplementedError

368

369

370 class DisStructuredBoundaryCondition(BoundaryCondition):

371 def _to_struct_array(self, arrdict, layer):

372 spec = []

373 for key in arrdict:

374 if key in ["layer", "row", "column"]:

375 spec.append((key, np.int32))

376 else:

377 spec.append((key, np.float64))

378

379 sparse_dtype = np.dtype(spec)

380 nrow = next(iter(arrdict.values())).size

381 recarr = np.empty(nrow, dtype=sparse_dtype)

382 for key, arr in arrdict.items():

383 recarr[key] = arr

384 return recarr

385

386

387 class DisVerticesBoundaryCondition(BoundaryCondition):

388 def _to_struct_array(self, arrdict, layer):

389 spec = []

390 for key in arrdict:

391 if key in ["layer", "cell2d"]:

392 spec.append((key, np.int32))

393 else:

394 spec.append((key, np.float64))

395

396 sparse_dtype = np.dtype(spec)

397 nrow = next(iter(arrdict.values())).size

398 recarr = np.empty(nrow, dtype=sparse_dtype)

399 for key, arr in arrdict.items():

400 recarr[key] = arr

401 return recarr