Create MetaModNewton

imod_coupler/config.py

@@ -24,3 +24,4 @@ class BaseConfig(BaseModel):
     timing: bool = False
     driver_type: DriverType
     driver: BaseModel
+    modflow_newton_formulation: bool = False

imod_coupler/drivers/driver.py

@@ -64,15 +64,18 @@ def get_driver(
     config_dict: dict[str, Any], config_dir: Path, base_config: BaseConfig
 ) -> Driver:
     from imod_coupler.drivers.metamod.config import MetaModConfig
-    from imod_coupler.drivers.metamod.metamod import MetaMod
+    from imod_coupler.drivers.metamod.metamod import MetaMod, MetaModNewton
     from imod_coupler.drivers.ribametamod.config import RibaMetaModConfig
     from imod_coupler.drivers.ribametamod.ribametamod import RibaMetaMod
     from imod_coupler.drivers.ribamod.config import RibaModConfig
     from imod_coupler.drivers.ribamod.ribamod import RibaMod
 
     if base_config.driver_type == "metamod":
         metamod_config = MetaModConfig(config_dir=config_dir, **config_dict["driver"])
-        return MetaMod(base_config, metamod_config)
+        if base_config.modflow_newton_formulation:
+            return MetaModNewton(base_config, metamod_config)
+        else:
+            return MetaMod(base_config, metamod_config)
     elif base_config.driver_type == "ribamod":
         ribamod_config = RibaModConfig(config_dir=config_dir, **config_dict["driver"])
         return RibaMod(base_config, ribamod_config)

imod_coupler/drivers/metamod/metamod.py

@@ -17,6 +17,11 @@
 from imod_coupler.config import BaseConfig
 from imod_coupler.drivers.driver import Driver
 from imod_coupler.drivers.metamod.config import Coupling, MetaModConfig
+from imod_coupler.kernelwrappers.mf6_newton_wrapper import (
+    PhreaticHeads,
+    PhreaticRecharge,
+    PhreaticStorage,
+)
 from imod_coupler.kernelwrappers.mf6_wrapper import Mf6Wrapper
 from imod_coupler.kernelwrappers.msw_wrapper import MswWrapper
 from imod_coupler.logging.exchange_collector import ExchangeCollector
@@ -39,7 +44,7 @@ class MetaMod(Driver):
 
     mf6_head: NDArray[Any]  # the hydraulic head array in the coupled model
     mf6_recharge: NDArray[np.float64]  # the coupled recharge array from the RCH package
-    mf6_storage: NDArray[Any]  # the specific storage array (ss)
+    mf6_ss: NDArray[Any]  # the specific storage array (ss)
     mf6_has_sc1: bool  # when true, specific storage in mf6 is given as a storage coefficient (sc1)
     mf6_area: NDArray[Any]  # cell area (size:nodes)
     mf6_top: NDArray[Any]  # top of cell (size:nodes)
@@ -60,6 +65,9 @@ class MetaMod(Driver):
     # dict. with mask arrays for msw=>mod coupling
     mask_msw2mod: dict[str, NDArray[Any]] = {}
 
+    node_idx: NDArray[Any]  # mf6 nodes coupled to msw
+    msw_idx: NDArray[Any]  # msw nodes coupled to MF6
+
     def __init__(self, base_config: BaseConfig, metamod_config: MetaModConfig):
         """Constructs the `MetaMod` object"""
         self.base_config = base_config
@@ -106,7 +114,7 @@ def couple(self) -> None:
         self.mf6_recharge = self.mf6.get_recharge(
             self.coupling.mf6_model, self.coupling.mf6_msw_recharge_pkg
         )
-        self.mf6_storage = self.mf6.get_storage(self.coupling.mf6_model)
+        self.mf6_ss = self.mf6.get_ss(self.coupling.mf6_model)
         self.mf6_has_sc1 = self.mf6.has_sc1(self.coupling.mf6_model)
         self.mf6_area = self.mf6.get_area(self.coupling.mf6_model)
         self.mf6_top = self.mf6.get_top(self.coupling.mf6_model)
@@ -136,17 +144,19 @@ def couple(self) -> None:
         table_node2svat: NDArray[np.int32] = np.loadtxt(
             self.coupling.mf6_msw_node_map, dtype=np.int32, ndmin=2
         )
-        node_idx = table_node2svat[:, 0] - 1
-        msw_idx = [
-            svat_lookup[table_node2svat[ii, 1], table_node2svat[ii, 2]]
-            for ii in range(len(table_node2svat))
-        ]
+        self.node_idx = table_node2svat[:, 0] - 1
+        self.msw_idx = np.array(
+            [
+                svat_lookup[table_node2svat[ii, 1], table_node2svat[ii, 2]]
+                for ii in range(len(table_node2svat))
+            ]
+        )
 
         self.map_msw2mod["storage"], self.mask_msw2mod["storage"] = create_mapping(
-            msw_idx,
-            node_idx,
+            self.msw_idx,
+            self.node_idx,
             self.msw_storage.size,
-            self.mf6_storage.size,
+            self.mf6_ss.size,
             "sum",
         )
 
@@ -167,8 +177,8 @@ def couple(self) -> None:
         self.map_msw2mod["storage"] = conversion_matrix * self.map_msw2mod["storage"]
 
         self.map_mod2msw["head"], self.mask_mod2msw["head"] = create_mapping(
-            node_idx,
-            msw_idx,
+            self.node_idx,
+            self.msw_idx,
             self.mf6_head.size,
             self.msw_head.size,
             "avg",
@@ -178,13 +188,15 @@ def couple(self) -> None:
             self.coupling.mf6_msw_recharge_map, dtype=np.int32, ndmin=2
         )
         rch_idx = table_rch2svat[:, 0] - 1
-        msw_idx = [
-            svat_lookup[table_rch2svat[ii, 1], table_rch2svat[ii, 2]]
-            for ii in range(len(table_rch2svat))
-        ]
+        self.msw_idx = np.array(
+            [
+                svat_lookup[table_rch2svat[ii, 1], table_rch2svat[ii, 2]]
+                for ii in range(len(table_rch2svat))
+            ]
+        )
 
         self.map_msw2mod["recharge"], self.mask_msw2mod["recharge"] = create_mapping(
-            msw_idx,
+            self.msw_idx,
             rch_idx,
             self.msw_volume.size,
             self.mf6_recharge.size,
@@ -205,16 +217,18 @@ def couple(self) -> None:
                 ndmin=2,
             )
             well_idx = table_well2svat[:, 0] - 1
-            msw_idx = [
-                svat_lookup[table_well2svat[ii, 1], table_well2svat[ii, 2]]
-                for ii in range(len(table_well2svat))
-            ]
+            self.msw_idx = np.array(
+                [
+                    svat_lookup[table_well2svat[ii, 1], table_well2svat[ii, 2]]
+                    for ii in range(len(table_well2svat))
+                ]
+            )
 
             (
                 self.map_msw2mod["sprinkling"],
                 self.mask_msw2mod["sprinkling"],
             ) = create_mapping(
-                msw_idx,
+                self.msw_idx,
                 well_idx,
                 self.msw_volume.size,
                 self.mf6_sprinkling_wells.size,
@@ -244,6 +258,10 @@ def update(self) -> None:
         self.mf6.finalize_time_step()
         self.msw.finalize_time_step()
 
+        if not has_converged:
+            logger.debug("MF6-MSW did not converge")
+            # raise ValueError("MF6-MSW did not converge")
+
     def finalize(self) -> None:
         self.mf6.finalize()
         self.msw.finalize()
@@ -257,12 +275,12 @@ def get_end_time(self) -> float:
 
     def exchange_msw2mod(self) -> None:
         """Exchange Metaswap to Modflow"""
-        self.mf6_storage[:] = (
-            self.mask_msw2mod["storage"][:] * self.mf6_storage[:]
+        self.mf6_ss[:] = (
+            self.mask_msw2mod["storage"][:] * self.mf6_ss[:]
             + self.map_msw2mod["storage"].dot(self.msw_storage)[:]
         )
         self.exchange_logger.log_exchange(
-            "mf6_storage", self.mf6_storage, self.get_current_time()
+            "mf6_storage", self.mf6_ss, self.get_current_time()
         )
         self.exchange_logger.log_exchange(
             "msw_storage", self.msw_storage, self.get_current_time()
@@ -306,3 +324,169 @@ def report_timing_totals(self) -> None:
         total_msw = self.msw.report_timing_totals()
         total = total_mf6 + total_msw
         logger.info(f"Total elapsed time in numerical kernels: {total:0.4f} seconds")
+
+
+class MetaModNewton(MetaMod):
+    """
+    MetaModNewton: the coupling between MetaSWAP and MODFLOW 6, for the Newton formulation of MODFLOW 6
+    """
+
+    def __init__(self, base_config: BaseConfig, metamod_config: MetaModConfig):
+        super().__init__(base_config, metamod_config)
+
+    def couple(self) -> None:
+        super().couple()
+        # Mapping for heads, SY and SS based on the top layer subset of MODFLOW 6 model arrays.
+        # The PhreaticHeads + PhreaticStorage methods get and set values to the corresponding
+        # phreatic nodes of potential deeper layers.
+        first_layer_nodes = self.get_first_layer_nodes()
+        self.map_mod2msw["head"], self.mask_mod2msw["head"] = create_mapping(
+            self.node_idx,
+            self.msw_idx,
+            first_layer_nodes.size,
+            self.msw_head.size,
+            "avg",
+        )
+
+        # mapping for SS now to top layer subset
+        self.map_msw2mod["storage"], self.mask_msw2mod["storage"] = create_mapping(
+            self.msw_idx,
+            self.node_idx,
+            self.msw_storage.size,
+            first_layer_nodes.size,
+            "sum",
+        )
+        conversion_matrix = self.sto_conversion_terms(
+            self.mf6_has_sc1, first_layer_nodes
+        )
+        self.map_msw2mod["storage"] = conversion_matrix * self.map_msw2mod["storage"]
+        # Create extra mapping for SY, since the SS mapping could contains a different conversion term
+        (
+            self.map_msw2mod["storage_sy"],
+            self.mask_msw2mod["storage_sy"],
+        ) = create_mapping(
+            self.msw_idx,
+            self.node_idx,
+            self.msw_storage.size,
+            first_layer_nodes.size,
+            "sum",
+        )
+        # For exchange to SY, act as is sc1
+        conversion_matrix = self.sto_conversion_terms(True, first_layer_nodes)
+        self.map_msw2mod["storage_sy"] = (
+            conversion_matrix * self.map_msw2mod["storage_sy"]
+        )
+        # Add Newton related phreatic exchange classes
+        self.couple_phreatic(first_layer_nodes)
+
+    def couple_phreatic(self, first_layer_nodes: NDArray[Any]) -> None:
+        userid = self.mf6_get_userid()
+        saturation = self.mf6.get_saturation(self.coupling.mf6_model)
+        self.sy = self.mf6.get_sy(self.coupling.mf6_model)
+        self.sy_top = self.sy[first_layer_nodes]
+        self.recharge_nodelist = self.mf6.get_recharge_nodes(
+            self.coupling.mf6_model, self.coupling.mf6_msw_recharge_pkg
+        )
+        self.hds = PhreaticHeads(
+            self.mf6.get_dis_shape(self.coupling.mf6_model),
+            userid,
+            saturation,
+            self.mf6_head,
+            first_layer_nodes,
+        )
+        self.sto = PhreaticStorage(
+            self.mf6.get_dis_shape(self.coupling.mf6_model),
+            userid,
+            saturation,
+            self.sy,
+            self.mf6_ss,
+            first_layer_nodes,
+        )
+        self.rch = PhreaticRecharge(
+            self.mf6.get_dis_shape(self.coupling.mf6_model),
+            userid,
+            saturation,
+            self.mf6_recharge,
+            self.recharge_nodelist,
+        )
+
+    def exchange_msw2mod(self) -> None:
+        """Exchange Metaswap to Modflow"""
+        new_sy = (
+            self.mask_msw2mod["storage_sy"][:] * self.sy_top[:]
+            + self.map_msw2mod["storage_sy"].dot(self.msw_storage)[:]
+        )
+        self.sto.reset()
+        self.sto.set(new_sy)
+        self.exchange_logger.log_exchange("mf6_sy", new_sy, self.get_current_time())
+        self.exchange_logger.log_exchange(
+            "msw_storage", self.msw_storage, self.get_current_time()
+        )
+
+        # Set recharge values
+        new_recharge = (
+            self.mask_msw2mod["recharge"][:] * self.mf6_recharge[:]
+            + self.map_msw2mod["recharge"].dot(self.msw_volume)[:] / self.delt
+        ) / self.mf6_area[self.recharge_nodelist - 1]
+        self.rch.set(new_recharge)
+
+        if self.enable_sprinkling_groundwater:
+            self.mf6_sprinkling_wells[:] = (
+                self.mask_msw2mod["sprinkling"][:] * self.mf6_sprinkling_wells[:]
+                + self.map_msw2mod["sprinkling"].dot(self.msw_volume)[:] / self.delt
+            )
+
+    def exchange_mod2msw(self) -> None:
+        """Exchange Modflow to Metaswap"""
+        mf6_head = self.hds.get()
+        self.msw_head[:] = (
+            self.mask_mod2msw["head"][:] * self.msw_head[:]
+            + self.map_mod2msw["head"].dot(mf6_head)[:]
+        )
+
+    def do_iter(self, sol_id: int) -> bool:
+        """Execute a single iteration"""
+        self.msw.prepare_solve(0)
+        self.msw.solve(0)
+        self.exchange_msw2mod()
+        has_converged = self.mf6.solve(sol_id)
+        self.exchange_mod2msw()
+        self.msw.finalize_solve(0)
+        # update nodelist rch-package
+        self.rch.set_nodes()
+        return has_converged
+
+    def get_first_layer_nodes(self) -> NDArray[Any]:
+        _, nrow, ncol = self.mf6.get_dis_shape(self.coupling.mf6_model)
+        userid = self.mf6_get_userid()
+        first_layer_ids = userid[userid <= (nrow * ncol)]
+        if self.node_idx.max() > first_layer_ids.max():
+            raise ValueError(
+                "MetaSWAP could only be coupled to the first model layer of MODFLOW 6"
+            )
+        return first_layer_ids
+
+    def mf6_get_userid(self) -> NDArray[Any]:
+        nlay, nrow, ncol = self.mf6.get_dis_shape(self.coupling.mf6_model)
+        userid = self.mf6.get_nodeuser(self.coupling.mf6_model)
+        if userid.size == 1:
+            # no reduced domain, set userid to modelid
+            # TODO: find out if there is a flag that indicates that usernodes == modelnodes
+            userid = np.arange(nlay * nrow * ncol) + 1
+        return userid
+
+    def sto_conversion_terms(
+        self, mf6_has_sc1: bool, indices: NDArray[Any]
+    ) -> NDArray[Any]:
+        if mf6_has_sc1:
+            conversion_terms = 1.0 / self.mf6_area[indices]
+        else:
+            conversion_terms = 1.0 / (
+                self.mf6_area[indices] * (self.mf6_top[indices] - self.mf6_bot[indices])
+            )
+        conversion_matrix = dia_matrix(
+            (conversion_terms, [0]),
+            shape=(self.mf6_area[indices].size, self.mf6_area[indices].size),
+            dtype=self.mf6_area.dtype,
+        )
+        return conversion_matrix

imod_coupler/drivers/ribametamod/ribametamod.py

@@ -51,7 +51,7 @@ class RibaMetaMod(Driver):
     mf6_head: NDArray[Any]  # the hydraulic head array in the coupled model
     mf6_recharge: NDArray[Any]  # the coupled recharge array from the RCH package
     mf6_recharge_nodes: NDArray[Any]  # node selection of rch nodes
-    mf6_storage: NDArray[Any]  # the specific storage array (ss)
+    mf6_ss: NDArray[Any]  # the specific storage array (ss)
     mf6_has_sc1: bool  # when true, specific storage in mf6 is given as a storage coefficient (sc1)
     mf6_area: NDArray[Any]  # cell area (size:nodes)
     mf6_top: NDArray[Any]  # top of cell (size:nodes)
@@ -220,7 +220,7 @@ def couple_metaswap(self) -> dict[str, Any]:
             self.mf6_recharge_nodes = self.mf6.get_recharge_nodes(
                 self.coupling.mf6_model, self.coupling.mf6_msw_recharge_pkg
             )
-            self.mf6_storage = self.mf6.get_storage(self.coupling.mf6_model)
+            self.mf6_ss = self.mf6.get_ss(self.coupling.mf6_model)
             self.mf6_has_sc1 = self.mf6.has_sc1(self.coupling.mf6_model)
             self.mf6_area = self.mf6.get_area(self.coupling.mf6_model)
             self.mf6_top = self.mf6.get_top(self.coupling.mf6_model)
@@ -235,7 +235,7 @@ def couple_metaswap(self) -> dict[str, Any]:
             arrays["msw_storage"] = self.msw_storage
             arrays["mf6_recharge"] = self.mf6_recharge
             arrays["mf6_head"] = self.mf6_head
-            arrays["mf6_storage"] = self.mf6_storage
+            arrays["mf6_storage"] = self.mf6_ss
             arrays["mf6_has_sc1"] = self.mf6_has_sc1
             arrays["mf6_area"] = self.mf6_area
             arrays["mf6_top"] = self.mf6_top
@@ -444,12 +444,12 @@ def exchange_stage_rib2mod(self) -> None:
 
     def exchange_msw2mod(self) -> None:
         """Exchange Metaswap to Modflow"""
-        self.mf6_storage[:] = (
-            self.mapping.msw2mod["storage_mask"][:] * self.mf6_storage[:]
+        self.mf6_ss[:] = (
+            self.mapping.msw2mod["storage_mask"][:] * self.mf6_ss[:]
             + self.mapping.msw2mod["storage"].dot(self.msw_storage)[:]
         )
         self.exchange_logger.log_exchange(
-            "mf6_storage", self.mf6_storage, self.get_current_time()
+            "mf6_storage", self.mf6_ss, self.get_current_time()
         )
         self.exchange_logger.log_exchange(
             "msw_storage", self.msw_storage, self.get_current_time()