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

1 from __future__ import annotations

2

3 from typing import Optional

4

5 import cftime

6 import numpy as np

7

8 from imod.logging import init_log_decorator

9 from imod.mf6 import ConstantHead

10 from imod.mf6.clipped_boundary_condition_creator import create_clipped_boundary

11 from imod.mf6.model import Modflow6Model

12 from imod.typing import GridDataArray

13

14

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