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

36 statements  

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

1from pathlib import Path 

2from typing import Optional, Union 

3 

4import cftime 

5import numpy as np 

6import xarray as xr 

7 

8from imod.logging import init_log_decorator 

9from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables 

10from imod.mf6.exchangebase import ExchangeBase 

11from imod.mf6.package import Package 

12from imod.typing import GridDataArray 

13from imod.typing.grid import ones_like 

14 

15 

16class 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 )