Coverage for C:\src\imod-python\imod\mf6\gwfgwf.py: 100%

1 from pathlib import Path

2 from typing import Optional, Union

3

4 import cftime

5 import numpy as np

6 import xarray as xr

7

8 from imod.logging import init_log_decorator

9 from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables

10 from imod.mf6.exchangebase import ExchangeBase

11 from imod.mf6.package import Package

12 from imod.typing import GridDataArray

13 from imod.typing.grid import ones_like

14

15

16 class GWFGWF(ExchangeBase):

17 """

18 This package is for writing an exchange file, used for splitting up a model

19 into different submodels (that can be solved in parallel). It (usually)

20 is not instantiated by users, but created by the "split" method of the

21 simulation class."""

22

23 _auxiliary_data = {"auxiliary_data": "variable"}

24 _pkg_id = "gwfgwf"

25 _template = Package._initialize_template(_pkg_id)

26

27 @init_log_decorator()

28 def __init__(

29 self,

30 model_id1: str,

31 model_id2: str,

32 cell_id1: xr.DataArray,

33 cell_id2: xr.DataArray,

34 layer: xr.DataArray,

35 cl1: xr.DataArray,

36 cl2: xr.DataArray,

37 hwva: xr.DataArray,

38 angldegx: Optional[xr.DataArray] = None,

39 cdist: Optional[xr.DataArray] = None,

40 ):

41 dict_dataset = {

42 "cell_id1": cell_id1,

43 "cell_id2": cell_id2,

44 "layer": layer,

45 "model_name_1": model_id1,

46 "model_name_2": model_id2,

47 "ihc": ones_like(cl1, dtype=int),

48 "cl1": cl1,

49 "cl2": cl2,

50 "hwva": hwva,

51 }

52 super().__init__(dict_dataset)

53

54 auxiliary_variables = [var for var in [angldegx, cdist] if var is not None]

55 if auxiliary_variables:

56 self.dataset["auxiliary_data"] = xr.merge(auxiliary_variables).to_array(

57 name="auxiliary_data"

58 )

59 expand_transient_auxiliary_variables(self)

60

61 def set_options(

62 self,

63 print_input: Optional[bool] = None,

64 print_flows: Optional[bool] = None,

65 save_flows: Optional[bool] = None,

66 cell_averaging: Optional[str] = None,

67 dewatered: Optional[bool] = None,

68 variablecv: Optional[bool] = None,

69 xt3d: Optional[bool] = None,

70 newton: Optional[bool] = None,

71 ):

72 self.dataset["print_input"] = print_input

73 self.dataset["print_flows"] = print_flows

74 self.dataset["save_flows"] = save_flows

75 self.dataset["cell_averaging"] = cell_averaging

76 self.dataset["dewatered"] = dewatered

77 self.dataset["variablecv"] = variablecv

78 self.dataset["xt3d"] = xt3d

79 self.dataset["newton"] = newton

80

81 def clip_box(

82 self,

83 time_min: Optional[cftime.datetime | np.datetime64 | str] = None,

84 time_max: Optional[cftime.datetime | np.datetime64 | str] = None,

85 layer_min: Optional[int] = None,

86 layer_max: Optional[int] = None,

87 x_min: Optional[float] = None,

88 x_max: Optional[float] = None,

89 y_min: Optional[float] = None,

90 y_max: Optional[float] = None,

91 top: Optional[GridDataArray] = None,

92 bottom: Optional[GridDataArray] = None,

93 state_for_boundary: Optional[GridDataArray] = None,

94 ) -> Package:

95 raise NotImplementedError("this package cannot be clipped")

96

97 def render(

98 self,

99 directory: Path,

100 pkgname: str,

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

102 binary: bool,

103 ) -> str:

104 return self.render_with_geometric_constants(

105 directory, pkgname, globaltimes, binary

106 )