Coverage for C:\src\imod-python\imod\mf6\gwtgwt.py: 97%

32 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 GWTGWT(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 

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

25 _pkg_id = "gwtgwt" 

26 _template = Package._initialize_template(_pkg_id) 

27 

28 @init_log_decorator() 

29 def __init__( 

30 self, 

31 transport_model_id1: str, 

32 transport_model_id2: str, 

33 flow_model_id1: str, 

34 flow_model_id2: str, 

35 cell_id1: xr.DataArray, 

36 cell_id2: xr.DataArray, 

37 layer: xr.DataArray, 

38 cl1: xr.DataArray, 

39 cl2: xr.DataArray, 

40 hwva: xr.DataArray, 

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

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

43 ): 

44 dict_dataset = { 

45 "cell_id1": cell_id1, 

46 "cell_id2": cell_id2, 

47 "layer": layer, 

48 "model_name_1": transport_model_id1, 

49 "model_name_2": transport_model_id2, 

50 "flow_model_name_1": flow_model_id1, 

51 "flow_model_name_2": flow_model_id2, 

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

53 "cl1": cl1, 

54 "cl2": cl2, 

55 "hwva": hwva, 

56 } 

57 super().__init__(dict_dataset) 

58 

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

60 if auxiliary_variables: 

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

62 name="auxiliary_data" 

63 ) 

64 expand_transient_auxiliary_variables(self) 

65 

66 def set_options( 

67 self, 

68 save_flows: Optional[bool] = None, 

69 adv_scheme: Optional[str] = None, 

70 dsp_xt3d_off: Optional[bool] = None, 

71 dsp_xt3d_rhs: Optional[bool] = None, 

72 ): 

73 self.dataset["save_flows"] = save_flows 

74 self.dataset["adv_scheme"] = adv_scheme 

75 self.dataset["dsp_xt3d_off"] = dsp_xt3d_off 

76 self.dataset["dsp_xt3d_rhs"] = dsp_xt3d_rhs 

77 

78 def clip_box( 

79 self, 

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

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

82 layer_min: Optional[int] = None, 

83 layer_max: Optional[int] = None, 

84 x_min: Optional[float] = None, 

85 x_max: Optional[float] = None, 

86 y_min: Optional[float] = None, 

87 y_max: Optional[float] = None, 

88 top: Optional[GridDataArray] = None, 

89 bottom: Optional[GridDataArray] = None, 

90 state_for_boundary: Optional[GridDataArray] = None, 

91 ) -> Package: 

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

93 

94 def render( 

95 self, 

96 directory: Path, 

97 pkgname: str, 

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

99 binary: bool, 

100 ) -> str: 

101 return self.render_with_geometric_constants( 

102 directory, pkgname, globaltimes, binary 

103 )