feat(advanced package support): refactoring for advanced package support

Author: jlarsen-usgs

modflowapi/extensions/advpaks.py

@@ -0,0 +1,107 @@
+import numpy as np
+
+from .data import ListInput
+from .pakbase import AdvancedPackage
+
+
+class SfrPakage(AdvancedPackage):
+    """
+    Container for SFR and SFR like packages
+
+    Parameters
+    ----------
+    model : ApiModel
+        modflowapi model object
+    pkg_type : str
+        package type. Ex. "SFR"
+    pkg_name : str
+        package name (in the mf6 variables)
+    sim_package : bool
+        boolean flag for simulation level packages. Ex. TDIS, IMS
+    """
+
+    def __init__(self, model, pkg_type, pkg_name, sim_package=False):
+        super().__init__(model, pkg_type, pkg_name, sim_package)
+
+        self._diversion_var_arrs = []
+        self._set_advanced_variable_addrs("diversions", "_diversion_var_addrs")
+        self._diversion_vars = ListInput(self, self._diversion_var_arrs, spd=False)
+
+    @property
+    def diversions(self):
+        return self._diversion_vars
+
+    @diversions.setter
+    def diversions(self, recarray):
+        """
+        Setter object to update the diversions data
+
+        """
+        if isinstance(recarray, np.recarray):
+            self._diversion_vars.values = recarray
+        elif isinstance(recarray, ListInput):
+            self._diversion_vars.values = recarray.values
+        elif recarray is None:
+            self._diversion_vars.values = recarray
+        else:
+            raise TypeError(f"{type(recarray)} is not a supported diversions type")
+
+
+class LakPackage(AdvancedPackage):
+    """
+    Container for LAK and LAK like packages
+
+    Parameters
+    ----------
+    model : ApiModel
+        modflowapi model object
+    pkg_type : str
+        package type. Ex. "LAK"
+    pkg_name : str
+        package name (in the mf6 variables)
+    sim_package : bool
+        boolean flag for simulation level packages. Ex. TDIS, IMS
+    """
+
+    def __init__(self, model, pkg_type, pkg_name, sim_package=False):
+        super().__init__(model, pkg_type, pkg_name, sim_package)
+
+
+class MawPackage(AdvancedPackage):
+    """
+    Container for MAW and MAW like packages
+
+    Parameters
+    ----------
+    model : ApiModel
+        modflowapi model object
+    pkg_type : str
+        package type. Ex. "MAW"
+    pkg_name : str
+        package name (in the mf6 variables)
+    sim_package : bool
+        boolean flag for simulation level packages. Ex. TDIS, IMS
+    """
+
+    def __init__(self, model, pkg_type, pkg_name, sim_package=False):
+        super().__init__(model, pkg_type, pkg_name, sim_package)
+
+
+class UzfPackage(AdvancedPackage):
+    """
+    Container for UZF and UZF like packages
+
+    Parameters
+    ----------
+    model : ApiModel
+        modflowapi model object
+    pkg_type : str
+        package type. Ex. "UZF"
+    pkg_name : str
+        package name (in the mf6 variables)
+    sim_package : bool
+        boolean flag for simulation level packages. Ex. TDIS, IMS
+    """
+
+    def __init__(self, model, pkg_type, pkg_name, sim_package=False):
+        super().__init__(model, pkg_type, pkg_name, sim_package)

modflowapi/extensions/apimodel.py

@@ -1,9 +1,8 @@
 import numpy as np
 
+from .datamodel import get_package_type, gridshape
 from .pakbase import AdvancedPackage, ArrayPackage, ListPackage, package_factory
 
-gridshape = {"dis": ["nlay", "nrow", "ncol"], "disu": ["nlay", "ncpl"]}
-
 
 class ApiMbase:
     """
@@ -15,16 +14,16 @@ class ApiMbase:
         initialized ModflowApi object
     name : str
         modflow model name. ex. "GWF_1", "GWF-GWF_1"
-    pkg_types : dict
-        dictionary of package types and ApiPackage class types
+    pkg_types : None, dict
+        optional dictionary of package types and ApiPackage class types
     """
 
-    def __init__(self, mf6, name, pkg_types):
+    def __init__(self, mf6, name, pkg_types=None):
         self.mf6 = mf6
         self.name = name
         self._pkg_names = None
         self._pak_type = None
-        self.pkg_types = pkg_types
+        self._pkg_types = pkg_types
         self.package_dict = {}
         self._set_package_names()
         self._create_package_list()
@@ -73,10 +72,13 @@ def _create_package_list(self):
         """
         for ix, pkg_name in enumerate(self._pkg_names):
             pkg_type = self._pak_type[ix].lower()
-            if pkg_type in self.pkg_types:
-                basepackage = self.pkg_types[pkg_type]
+            if self._pkg_types is None:
+                basepackage = get_package_type(pkg_type)
             else:
-                basepackage = AdvancedPackage
+                if pkg_type in self._pkg_types:
+                    basepackage = self._pkg_types[pkg_type]
+                else:
+                    basepackage = AdvancedPackage
 
             package = package_factory(pkg_type, basepackage)
             adj_pkg_name = "".join(pkg_type.split("-"))
@@ -135,41 +137,14 @@ def __init__(self, mf6, name):
         else:
             raise AssertionError(f"Unrecognized discretization type {grid_type}")
 
-        pkg_types = {
-            "dis": ArrayPackage,
-            "chd": ListPackage,
-            "drn": ListPackage,
-            "evt": ListPackage,
-            "ghb": ListPackage,
-            "ic": ArrayPackage,
-            "npf": ArrayPackage,
-            "rch": ListPackage,
-            "riv": ListPackage,
-            "sto": ArrayPackage,
-            "wel": ListPackage,
-            # gwt
-            "dsp": ArrayPackage,
-            "cnc": ListPackage,
-            "ist": ArrayPackage,
-            "mst": ArrayPackage,
-            "src": ListPackage,
-            # gwe
-            "cnd": ArrayPackage,
-            "est": ArrayPackage,
-            "cpt": ListPackage,
-            "esl": ListPackage,
-            # prt
-            "mip": ArrayPackage,
-        }
-
         self.allow_convergence = True
         self._shape = None
         self._size = None
         self._nodetouser = None
         self._usertonode = None
         self._iteration = 0
 
-        super().__init__(mf6, name, pkg_types)
+        super().__init__(mf6, name)
 
     def __repr__(self):
         s = f"{self.name}, "

modflowapi/extensions/apisimulation.py

@@ -311,7 +311,7 @@ def load(mf6):
             i[:-1].lower() for ix, i in enumerate(mf6.get_value("__INPUT__/SIM/NAM/SLNTYPE")) if idp_names[ix]
         ]
 
-        tmpmdl = ApiMbase(mf6, "", {})
+        tmpmdl = ApiMbase(mf6, "")
         solution_names = list(set(solution_names))
         solution_dict = {}
         for name in solution_names: