sketching

Author: wpbonelli

docs/examples/quickstart.py

@@ -1,9 +1,9 @@
-from flopy4.mf6 import Simulation, Tdis
+from flopy4.mf6 import Sim, Tdis
 from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
 
 ws = "./mymodel"
 name = "mymodel"
-sim = Simulation(name=name, path=ws, exe="mf6")
+sim = Sim(name=name, path=ws, exe="mf6")
 tdis = Tdis(sim)
 gwf = Gwf(sim, name=name, save_flows=True)
 dis = Dis(gwf, nrow=10, ncol=10)

flopy4/__init__.py

@@ -95,7 +95,7 @@ def _bind_tree(self, parent):
         _bind_tree(parent, grandparent)
 
 
-def init_tree(self, parent=None, **kwargs):
+def init_tree(self, parent=None, children=None, **kwargs):
     """
     Initialize a data tree for a component instance.
     """
@@ -144,7 +144,12 @@ def init_tree(self, parent=None, **kwargs):
         value = kwargs.get(name, attr.default)
         data[name] = value
 
-    self.data = DataTree(data, name=cls_name)
+    children = children or {}
+    self.data = DataTree(
+        data,
+        name=cls_name,
+        children={n: c for n, c in children.items() if c is not None},
+    )
     if parent is not None:
         self.parent = parent
         _bind_tree(self, parent)

flopy4/mf6/__init__.py

@@ -8,7 +8,16 @@
 
 from flopy4 import component, init_tree, setattribute
 
-__all__ = ["Component", "Package", "Model", "Sim", "Simulation"]
+__all__ = [
+    "Component",
+    "Package",
+    "Model",
+    "Simulation",
+    "Sim",
+    "COMPONENTS",
+]
+
+COMPONENTS = {}  # component registry
 
 
 class Component(ABC):
@@ -19,6 +28,10 @@ def __init__(self, name=None, path=None):
         self.name = name
         self.path = path
 
+    @classmethod
+    def __attrs_init_subclass__(cls):
+        COMPONENTS[cls.__name__.lower()] = cls
+
 
 class Package(Component):
     def __init__(self, name=None, path=None):
@@ -30,14 +43,6 @@ def __init__(self, name=None, path=None):
         super().__init__(name, path)
 
 
-class Sim(Component):
-    exe: Path
-
-    def __init__(self, name=None, path=None, exe=None):
-        super().__init__(name, path)
-        self.exe = exe or which("mf6")
-
-
 @component
 @define(init=False, slots=False, on_setattr=setattribute)
 class Tdis(Package):
@@ -80,9 +85,66 @@ def __init__(
         )
 
 
-@component
-@define(init=False, slots=False, on_setattr=setattribute)
-class Simulation(Sim):
+class Solution(Package):
+    def __init__(self, name=None, path=None):
+        super().__init__(name, path)
+
+
+@define(init=False, slots=False)
+class Exchange(Package):
+    exgtype: type = field()
+    exgfile: Path = field()
+    exgmnamea: Optional[str] = field(default=None)
+    exgmnameb: Optional[str] = field(default=None)
+
+    def __init__(
+        self,
+        name=None,
+        path=None,
+        mnamea=None,
+        mnameb=None,
+    ):
+        super().__init__(name, path)
+        self.exgtype = type(self)
+        self.exgfile = path
+        self.exgmnamea = mnamea
+        self.exgmnameb = mnameb
+
+
+class Simulation(Component):
+    exe: Path
+
     def __init__(self, name=None, path=None, exe=None):
+        super().__init__(name, path)
+        self.exe = exe or which("mf6")
+
+
+@component
+@define(init=False, slots=False)
+class Sim(Simulation):
+    # tdis: Tdis = field(metadata={"block": "timing"})
+    # models: dict[str, Model] = field(metadata={"block": "models"})
+    # exchanges: dict[str, Exchange] = field(metadata={"block": "exchanges"})
+    # solutions: dict[str, Solution] = field(metadata={"block": "solutions"})
+
+    def __init__(
+        self,
+        name=None,
+        path=None,
+        exe=None,
+        tdis=None,
+        models=None,
+        exchanges=None,
+        solutions=None,
+    ):
         super().__init__(name, path, exe)
-        init_tree(self)
+        # TODO pull init_tree(self) call into @component definition
+        init_tree(
+            self,
+            # children={
+            #     "tdis": tdis,
+            #     "models": models,
+            #     "exchanges": exchanges,
+            #     "solutions": solutions,
+            # },
+        )

test/test_component.py

@@ -1,23 +1,29 @@
 import numpy as np
 from xarray import DataTree
 
-from flopy4.mf6 import Simulation, Tdis
+from flopy4.mf6 import COMPONENTS, Sim, Tdis
 from flopy4.mf6.gwf import Dis, Gwf, Ic, Npf, Oc
 
 
 def test_components():
-    # Create a simulation.
-    sim = Simulation()
+    assert COMPONENTS["sim"] is Sim
+    assert COMPONENTS["tdis"] is Tdis
+    assert COMPONENTS["gwf"] is Gwf
+    assert COMPONENTS["npf"] is Npf
+
+
+def test_sim():
+    sim = Sim()
     tdis = Tdis(sim=sim, nper=1, perioddata=[Tdis.PeriodData()])
     gwf = Gwf(sim=sim)
     dis = Dis(model=gwf)
     ic = Ic(model=gwf, strt=1.0)
     oc = Oc(model=gwf, perioddata=[Oc.Steps()])
     npf = Npf(model=gwf, icelltype=0, k=1.0)
 
-    # View the data tree.
-    # sim.data
     assert isinstance(sim.data, DataTree)
+    # sim.data  # view the tree
+
     assert "tdis" in sim.data.children
     assert "gwf" in sim.data.children
     assert "dis" in sim.data.children["gwf"].children
@@ -28,3 +34,7 @@ def test_components():
     assert np.array_equal(
         sim.data.children["gwf"].children["npf"].k, np.ones((4))
     )
+    assert sim.data.children["gwf"] is gwf.data
+
+    # why fails?
+    # assert gwf.parent.data.children["gwf"].children["npf"] is npf.data