Coverage for C:\src\imod-python\imod\flow\oc.py: 97%

30 statements  

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

1import numpy as np 

2 

3from imod.flow.pkgbase import Package 

4 

5 

6class OutputControl(Package): 

7 """ 

8 The Output Control Option determines how output is printed to the listing 

9 file and/or written to a separate binary output file. 

10 

11 Parameters 

12 ---------- 

13 save_head : {xr.DataArray, integer} 

14 xr.DataArray of integers or integer indicating the layer at which heads should 

15 be saved. If set to -1, output is saved for all layers. If set to 0, 

16 nothing is saved. 

17 save_flux : {xr.DataArray, integer} 

18 xr.DataArray of integers or integer indicating the layer at which spatial fluxes 

19 should be saved. If set to -1, output is saved for all layers. If set to 

20 0, nothing is saved. 

21 save_ghb : {xr.DataArray, integer} 

22 xr.DataArray of integers or integer indicating the layer at which ghb budget 

23 should be saved. If set to -1, output is saved for all layers. If set to 

24 0, nothing is saved. 

25 save_drn : {xr.DataArray, integer} 

26 xr.DataArray of integers or integer indicating the layer at which drn budget 

27 should be saved. If set to -1, output is saved for all layers. If set to 

28 0, nothing is saved. 

29 save_wel : {xr.DataArray, integer} 

30 xr.DataArray of integers or integer indicating the layer at which wel budget 

31 should be saved. If set to -1, output is saved for all layers. If set to 

32 0, nothing is saved. 

33 save_riv : {xr.DataArray, integer} 

34 xr.DataArray of integers or integer indicating the layer at which riv budget 

35 should be saved. If set to -1, output is saved for all layers. If set to 

36 0, nothing is saved. 

37 save_rch : {xr.DataArray, integer} 

38 xr.DataArray of integers or integer indicating the layer at which rch budget 

39 should be saved. If set to -1, output is saved for all layers. If set to 

40 0, nothing is saved. 

41 save_evt : {xr.DataArray, integer} 

42 xr.DataArray of integers or integer indicating the layer at which evt budget 

43 should be saved. If set to -1, output is saved for all layers. If set to 

44 0, nothing is saved. 

45 save_hfb : {xr.DataArray, integer} 

46 xr.DataArray of integers or integer indicating the layer at which hfb budget 

47 should be saved. If set to -1, output is saved for all layers. If set to 

48 0, nothing is saved. 

49 

50 Examples 

51 -------- 

52 

53 Don't save heads: 

54 

55 >>> oc = imod.flow.OutputControl(save_head=0) 

56 

57 Save heads for all layers: 

58 

59 >>> oc = imod.flow.OutputControl(save_head=-1) 

60 

61 Save heads for specific layers, in this case layer 1 and 4: 

62 

63 >>> specific_layers = xr.DataArray([1,4], coords={"layer": [1,4], dims=("layer",)}) 

64 >>> oc = imod.flow.OutputControl(save_head=specific_layers) 

65 

66 

67 """ 

68 

69 _pkg_id = "oc" 

70 _variable_order = [ 

71 "save_head", 

72 "save_flux", 

73 "save_ghb", 

74 "save_drn", 

75 "save_wel", 

76 "save_riv", 

77 "save_rch", 

78 "save_evt", 

79 "save_hfb", 

80 ] 

81 

82 def __init__( 

83 self, 

84 save_head=0, 

85 save_flux=0, 

86 save_ghb=0, 

87 save_drn=0, 

88 save_wel=0, 

89 save_riv=0, 

90 save_rch=0, 

91 save_evt=0, 

92 save_hfb=0, 

93 ): 

94 super().__init__() 

95 self.dataset["save_head"] = save_head 

96 self.dataset["save_flux"] = save_flux 

97 self.dataset["save_ghb"] = save_ghb 

98 self.dataset["save_drn"] = save_drn 

99 self.dataset["save_wel"] = save_wel 

100 self.dataset["save_riv"] = save_riv 

101 self.dataset["save_rch"] = save_rch 

102 self.dataset["save_evt"] = save_evt 

103 self.dataset["save_hfb"] = save_hfb 

104 

105 def _compose_oc_configuration(self, nlayer): 

106 # Create mapping to convert e.g. 'save_ghb' to 'saveghb' 

107 conf_mapping = dict( 

108 [(var, var.replace("_", "")) for var in self._variable_order] 

109 ) 

110 conf_mapping["save_head"] = "saveshd" 

111 conf_mapping["save_flux"] = "saveflx" 

112 

113 pkg_composition = self.compose(None, None, nlayer)["oc"] 

114 

115 oc_composition = {} 

116 

117 for var in self._variable_order: 

118 values = np.array(list(pkg_composition[var].values())) 

119 # integer was provided 

120 if np.all(values == values[0]): 

121 oc_composition[conf_mapping[var]] = values[0] 

122 elif np.any(values == 0) or np.any(values == -1): 

123 raise ValueError( 

124 f"Received invalid layer numbers, layer array cannot contain 0 or -1, got {values} for {var}" 

125 ) 

126 else: 

127 oc_composition[conf_mapping[var]] = ",".join( 

128 str(value) for value in values 

129 ) 

130 

131 return oc_composition