Create RibaMetaMod as a copy of RibaMod (#202)

* add RibaMetaMod to premod

* Initial copy of the RibaMod-driver

* add copy of ribamod tests

* linter

* Ruff

* black

* Fix type-hints

* Remove unused import

---------

Co-authored-by: Hofer-Julian <[email protected]>

imod_coupler/drivers/driver.py

@@ -65,6 +65,8 @@ def get_driver(
 ) -> Driver:
     from imod_coupler.drivers.metamod.config import MetaModConfig
     from imod_coupler.drivers.metamod.metamod import MetaMod
+    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
 
@@ -74,5 +76,10 @@ def get_driver(
     elif base_config.driver_type == "ribamod":
         ribamod_config = RibaModConfig(config_dir=config_dir, **config_dict["driver"])
         return RibaMod(base_config, ribamod_config)
+    elif base_config.driver_type == "ribametamod":
+        ribametamod_config = RibaMetaModConfig(
+            config_dir=config_dir, **config_dict["driver"]
+        )
+        return RibaMetaMod(base_config, ribametamod_config)
     else:
         raise ValueError(f"Driver type {base_config.driver_type} is not supported.")

imod_coupler/drivers/metamod/metamod.py

@@ -12,7 +12,6 @@
 from loguru import logger
 from numpy.typing import NDArray
 from scipy.sparse import csr_matrix, dia_matrix
-from xmipy import XmiWrapper
 
 from imod_coupler.config import BaseConfig
 from imod_coupler.drivers.driver import Driver
@@ -27,20 +26,24 @@ class MetaMod(Driver):
     """The driver coupling MetaSWAP and MODFLOW 6"""
 
     base_config: BaseConfig  # the parsed information from the configuration file
-    metamod_config: MetaModConfig  # the parsed information from the configuration file specific to MetaMod
+    metamod_config: (
+        MetaModConfig
+    )  # the parsed information from the configuration file specific to MetaMod
     coupling: Coupling  # the coupling information
 
     timing: bool  # true, when timing is enabled
-    mf6: XmiWrapper  # the MODFLOW 6 XMI kernel
-    msw: XmiWrapper  # the MetaSWAP XMI kernel
+    mf6: Mf6Wrapper  # the MODFLOW 6 XMI kernel
+    msw: MswWrapper  # the MetaSWAP XMI kernel
 
     max_iter: NDArray[Any]  # max. nr outer iterations in MODFLOW kernel
     delt: float  # time step from MODFLOW 6 (leading)
 
     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_storage: 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_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)
     mf6_bot: NDArray[Any]  # bottom of cell (size:nodes)

imod_coupler/drivers/ribametamod/config.py

@@ -0,0 +1,144 @@
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, DirectoryPath, FilePath, validator
+
+
+class Modflow6(BaseModel):
+    dll: FilePath
+    dll_dep_dir: Optional[DirectoryPath]
+    work_dir: DirectoryPath
+
+    @validator("dll")
+    def resolve_dll(cls, dll: FilePath) -> FilePath:
+        return dll.resolve()
+
+    @validator("dll_dep_dir")
+    def resolve_dll_dep_dir(
+        cls, dll_dep_dir: Optional[DirectoryPath]
+    ) -> Optional[DirectoryPath]:
+        if dll_dep_dir is not None:
+            dll_dep_dir = dll_dep_dir.resolve()
+        return dll_dep_dir
+
+
+class Ribasim(BaseModel):
+    dll: FilePath
+    dll_dep_dir: DirectoryPath
+    config_file: FilePath
+
+    @validator("dll")
+    def resolve_dll(cls, dll: FilePath) -> FilePath:
+        return dll.resolve()
+
+    @validator("dll_dep_dir")
+    def resolve_dll_dep_dir(
+        cls, dll_dep_dir: Optional[DirectoryPath]
+    ) -> Optional[DirectoryPath]:
+        if dll_dep_dir is not None:
+            dll_dep_dir = dll_dep_dir.resolve()
+        return dll_dep_dir
+
+
+class Metaswap(BaseModel):
+    dll: FilePath
+    dll_dep_dir: DirectoryPath
+    config_file: FilePath
+
+    @validator("dll")
+    def resolve_dll(cls, dll: FilePath) -> FilePath:
+        return dll.resolve()
+
+    @validator("dll_dep_dir")
+    def resolve_dll_dep_dir(
+        cls, dll_dep_dir: Optional[DirectoryPath]
+    ) -> Optional[DirectoryPath]:
+        if dll_dep_dir is not None:
+            dll_dep_dir = dll_dep_dir.resolve()
+        return dll_dep_dir
+
+
+class Kernels(BaseModel):
+    modflow6: Modflow6
+    ribasim: Ribasim
+    metaswap: Metaswap
+
+
+class Coupling(BaseModel):
+    mf6_model: str  # the MODFLOW 6 model that will be coupled
+    mf6_active_river_packages: Dict[str, str]
+    mf6_active_drainage_packages: Dict[str, str]
+    mf6_passive_river_packages: Dict[str, str]
+    mf6_passive_drainage_packages: Dict[str, str]
+
+    enable_sprinkling: bool  # true whemn sprinkling is active
+    mf6_msw_recharge_pkg: str  # the recharge package that will be used for coupling
+    mf6_msw_well_pkg: Optional[
+        str
+    ] = None  # the well package that will be used for coupling when sprinkling is active
+    mf6_msw_node_map: FilePath  # the path to the node map file
+    mf6_msw_recharge_map: FilePath  # the pach to the recharge map file
+    mf6_msw_sprinkling_map: Optional[
+        FilePath
+    ] = None  # the path to the sprinkling map file
+    output_config_file: Optional[FilePath] = None
+
+    class Config:
+        arbitrary_types_allowed = True  # Needed for `mf6_msw_sprinkling_map`
+
+    @validator("mf6_msw_well_pkg")
+    def validate_mf6_msw_well_pkg(
+        cls, mf6_msw_well_pkg: Optional[str], values: Any
+    ) -> Optional[str]:
+        if values.get("enable_sprinkling") and mf6_msw_well_pkg is None:
+            raise ValueError(
+                "If `enable_sprinkling` is True, then `mf6_msw_well_pkg` needs to be set."
+            )
+        return mf6_msw_well_pkg
+
+    @validator(
+        "output_config_file",
+        "mf6_msw_node_map",
+        "mf6_msw_recharge_map",
+        "output_config_file",
+    )
+    def resolve_file_path(cls, file_path: FilePath) -> FilePath:
+        return file_path.resolve()
+
+    @validator("mf6_msw_sprinkling_map")
+    def validate_mf6_msw_sprinkling_map(
+        cls, mf6_msw_sprinkling_map: Optional[FilePath], values: Any
+    ) -> Optional[FilePath]:
+        if mf6_msw_sprinkling_map is not None:
+            return mf6_msw_sprinkling_map.resolve()
+        elif values.get("enable_sprinkling"):
+            raise ValueError(
+                "If `enable_sprinkling` is True, then `mf6_msw_sprinkling_map` needs to be set."
+            )
+        return mf6_msw_sprinkling_map
+
+
+class RibaMetaModConfig(BaseModel):
+    kernels: Kernels
+    coupling: List[Coupling]
+
+    def __init__(self, config_dir: Path, **data: Any) -> None:
+        """Model for the Ribamod config validated by pydantic
+
+        The validation expects current working directory at config file level
+        so it is changed during initialization
+
+        Args:
+            config_dir (Path): Directory where the config file resides
+        """
+        os.chdir(config_dir)
+        super().__init__(**data)
+
+    @validator("coupling")
+    def restrict_coupling_count(cls, coupling: List[Coupling]) -> List[Coupling]:
+        if len(coupling) == 0:
+            raise ValueError("At least one coupling has to be defined.")
+        if len(coupling) > 1:
+            raise ValueError("Multi-model coupling is not yet supported.")
+        return coupling