Coverage for C:\src\imod-python\imod\mf6\model_gwf.py: 91%

47 statements  

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

1from __future__ import annotations 

2 

3from typing import Optional 

4 

5import cftime 

6import numpy as np 

7 

8from imod.logging import init_log_decorator 

9from imod.mf6 import ConstantHead 

10from imod.mf6.clipped_boundary_condition_creator import create_clipped_boundary 

11from imod.mf6.model import Modflow6Model 

12from imod.typing import GridDataArray 

13 

14 

15class GroundwaterFlowModel(Modflow6Model): 

16 _mandatory_packages = ("npf", "ic", "oc", "sto") 

17 _model_id = "gwf6" 

18 _template = Modflow6Model._initialize_template("gwf-nam.j2") 

19 

20 @init_log_decorator() 

21 def __init__( 

22 self, 

23 listing_file: Optional[str] = None, 

24 print_input: bool = False, 

25 print_flows: bool = False, 

26 save_flows: bool = False, 

27 newton: bool = False, 

28 under_relaxation: bool = False, 

29 ): 

30 super().__init__() 

31 self._options = { 

32 "listing_file": listing_file, 

33 "print_input": print_input, 

34 "print_flows": print_flows, 

35 "save_flows": save_flows, 

36 "newton": newton, 

37 "under_relaxation": under_relaxation, 

38 } 

39 

40 def clip_box( 

41 self, 

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

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

44 layer_min: Optional[int] = None, 

45 layer_max: Optional[int] = None, 

46 x_min: Optional[float] = None, 

47 x_max: Optional[float] = None, 

48 y_min: Optional[float] = None, 

49 y_max: Optional[float] = None, 

50 state_for_boundary: Optional[GridDataArray] = None, 

51 ): 

52 clipped = super().clip_box( 

53 time_min, time_max, layer_min, layer_max, x_min, x_max, y_min, y_max 

54 ) 

55 

56 clipped_boundary_condition = self.__create_boundary_condition_clipped_boundary( 

57 self, clipped, state_for_boundary 

58 ) 

59 if clipped_boundary_condition is not None: 

60 clipped["chd_clipped"] = clipped_boundary_condition 

61 

62 clipped.purge_empty_packages() 

63 

64 return clipped 

65 

66 def __create_boundary_condition_clipped_boundary( 

67 self, 

68 original_model: Modflow6Model, 

69 clipped_model: Modflow6Model, 

70 state_for_boundary: Optional[GridDataArray], 

71 ): 

72 unassigned_boundary_original_domain = ( 

73 self.__create_boundary_condition_for_unassigned_boundary( 

74 original_model, state_for_boundary 

75 ) 

76 ) 

77 

78 return self.__create_boundary_condition_for_unassigned_boundary( 

79 clipped_model, state_for_boundary, [unassigned_boundary_original_domain] 

80 ) 

81 

82 @staticmethod 

83 def __create_boundary_condition_for_unassigned_boundary( 

84 model: Modflow6Model, 

85 state_for_boundary: Optional[GridDataArray], 

86 additional_boundaries: Optional[list[ConstantHead]] = None, 

87 ): 

88 if state_for_boundary is None: 

89 return None 

90 

91 constant_head_packages = [ 

92 pkg for name, pkg in model.items() if isinstance(pkg, ConstantHead) 

93 ] 

94 

95 additional_boundaries = [ 

96 item for item in additional_boundaries or [] if item is not None 

97 ] 

98 

99 constant_head_packages.extend(additional_boundaries) 

100 

101 return create_clipped_boundary( 

102 model.domain, state_for_boundary, constant_head_packages 

103 ) 

104 

105 def is_use_newton(self): 

106 return self._options["newton"] 

107 

108 def set_newton(self, is_newton: bool) -> None: 

109 self._options["newton"] = is_newton 

110 

111 def update_buoyancy_package(self, transport_models_per_flow_model) -> None: 

112 """ 

113 If the simulation is partitioned, then the buoyancy package, if present, 

114 must be updated for the renamed transport models. 

115 """ 

116 buoyancy_key = self._get_pkgkey("buy") 

117 if buoyancy_key is None: 

118 return 

119 buoyancy_package = self[buoyancy_key] 

120 transport_models_old = buoyancy_package.get_transport_model_names() 

121 if len(transport_models_old) == len(transport_models_per_flow_model): 

122 buoyancy_package.update_transport_models(transport_models_per_flow_model)