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

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 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 ) -> Package:

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

92

93 def render(

94 self,

95 directory: Path,

96 pkgname: str,

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

98 binary: bool,

99 ) -> str:

100 return self.render_with_geometric_constants(

101 directory, pkgname, globaltimes, binary

102 )