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
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1from __future__ import annotations
3from typing import Optional
5import cftime
6import numpy as np
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
15class GroundwaterFlowModel(Modflow6Model):
16 _mandatory_packages = ("npf", "ic", "oc", "sto")
17 _model_id = "gwf6"
18 _template = Modflow6Model._initialize_template("gwf-nam.j2")
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 }
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 )
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
62 clipped.purge_empty_packages()
64 return clipped
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 )
78 return self.__create_boundary_condition_for_unassigned_boundary(
79 clipped_model, state_for_boundary, [unassigned_boundary_original_domain]
80 )
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
91 constant_head_packages = [
92 pkg for name, pkg in model.items() if isinstance(pkg, ConstantHead)
93 ]
95 additional_boundaries = [
96 item for item in additional_boundaries or [] if item is not None
97 ]
99 constant_head_packages.extend(additional_boundaries)
101 return create_clipped_boundary(
102 model.domain, state_for_boundary, constant_head_packages
103 )
105 def is_use_newton(self):
106 return self._options["newton"]
108 def set_newton(self, is_newton: bool) -> None:
109 self._options["newton"] = is_newton
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)