Simplified API and AI skill for structured flowsheets

idaes/core/util/structfs/__init__.py

@@ -21,7 +21,7 @@
 ## Overview
 
 The core idea of the
-{py:class}`FlowsheetRunner <structfs.fsrunner.FlowsheetRunner>` class is
+{py:class}`StructuredFlowsheet <structfs.fsrunner.StructuredFlowsheet>` class is
 that flowsheets should follow a standard set of "steps". By standardizing the
 naming and ordering of these steps, it becomes easier to build tools that run
 and inspect flowsheets. The Python mechanics of this are to put each step in a
@@ -43,7 +43,7 @@
 It is assumed here that you have Python code to build, configure, and run an
 IDAES flowsheet. You will first arrange this code to follow the standard "steps"
 of a flowsheet workflow, which are listed in the
-{py:class}`BaseFlowsheetRunner <structfs.fsrunner.BaseFlowsheetRunner>`
+{py:class}`BaseStructuredFlowsheet <structfs.fsrunner.BaseStructuredFlowsheet>`
 class' `STEPS` attribute. Not all the steps need to be defined: the API will
 skip over steps with no definition when executing a range of steps. To make the
 code more structured you can also define internal sub-steps, as described later.
@@ -118,7 +118,7 @@ def solve(m):
 #### After
 
 In order to make this into a
-{py:class}`FlowsheetRunner <structfs.fsrunner.FlowsheetRunner>`-wrapped
+{py:class}`StructuredFlowsheet <structfs.fsrunner.StructuredFlowsheet>`-wrapped
 flowsheet, we need to do make a few changes. The modified file is shown below,
 with changed lines highlighted and descriptions below.
 
@@ -134,11 +134,11 @@ def solve(m):
 import ( BTXParameterBlock, )
 from idaes.models.unit_models import Flash
 
-from idaes.core.util.structfs.fsrunner import FlowsheetRunner
+from idaes.core.util.structfs import StructuredFlowsheet
 
-FS = FlowsheetRunner()
+flowsheet = StructuredFlowsheet()
 
-@FS.step("build") 
+@flowsheet.step.build
 def build_model(ctx):
     """Build the model."""
     m = ConcreteModel()
@@ -152,7 +152,7 @@ def build_model(ctx):
     # assert degrees_of_freedom(m) == 7
     ctx.model = m
 
-@FS.step("set_operating_conditions")
+@flowsheet.step.set_operating_conditions
 def set_operating_conditions(ctx):
     """Set operating conditions."""
     m = ctx.model
@@ -164,27 +164,27 @@ def set_operating_conditions(ctx):
     m.fs.flash.heat_duty.fix(0)
     m.fs.flash.deltaP.fix(0)
 
-@FS.step("initialize") 
+@flowsheet.step.initialize 
 def init_model(ctx): 
     """ "Initialize the model."""
     m = ctx.model
     m.fs.flash.initialize()
 
-@FS.step("set_solver") 
+@flowsheet.step.set_solver 
 def set_solver(ctx):
     """Set the solver."""
     ctx.solver = SolverFactory("ipopt")
 
-@FS.step("solve_optimization")
+@flowsheet.step.solve_optimization
 def solve_opt(ctx):
     ctx["results"] = ctx.solver.solve(ctx.model, tee=ctx["tee"])
 ```
 
 Details on the changes:
 
-* **7**: Import the FlowsheetRunner class.
-* **9**: Create a global {py:class}`FlowsheetRunner <structfs.fsrunner.FlowsheetRunner>` object, here called `FS`.
-* **11, 25, 37, 43, 48**: Add a `@FS.step()` decorator in front of each function
+* **7**: Import the StructuredFlowsheet class.
+* **9**: Create a global {py:class}`StructuredFlowsheet <structfs.fsrunner.StructuredFlowsheet>` object, here called `flowsheet`.
+* **11, 25, 37, 43, 48**: Add a `@flowsheet.step.step-name` decorator in front of each function
   with the name of the associated step.
 * **12, 26, 38, 44, 49**: Make each function take a single argument which is a {py:class}`fsrunner.Context <structfs.fsrunner.Context>` instance used to
   pass state information between functions (here, that argument is named `ctx`).
@@ -208,11 +208,11 @@ def solve_opt(ctx):
 could do this:
 
 ```{code}
-FS.run_steps()
-assert FS.results.solver.status == SolverStatus.ok
+flowsheet.run()
+assert flowsheet.results.solver.status == SolverStatus.ok
 ```
 
-Some more examples of using the FlowsheetRunner are shown in the
+Some more examples of using the StructuredFlowsheet are shown in the
 example notebooks found under the `docs/examples/structfs` directory
 ([docs link](/examples/structfs/index)).
 
@@ -226,13 +226,14 @@ def solve_opt(ctx):
 You can also 'annotate' variables for special 
 treatment in display, etc. with the
 `annotate_var` function in the 
-{py:class}`FlowsheetRunner <structfs.fsrunner.FlowsheetRunner>` class.
+{py:class}`StructuredFlowsheet <structfs.fsrunner.StructuredFlowsheet>` class.
 
-```{autodoc2-object} structfs.fsrunner.FlowsheetRunner.annotate_var
+```{autodoc2-object} structfs.fsrunner.StructuredFlowsheet.annotate_var
 ```
 
-```{autodoc2-docstring} structfs.fsrunner.FlowsheetRunner.annotate_var
+```{autodoc2-docstring} structfs.fsrunner.StructuredFlowsheet.annotate_var
 :parser: myst
 ```
 
 '''
+from .fsrunner import StructuredFlowsheet

idaes/core/util/structfs/fsrunner.py

@@ -12,7 +12,7 @@
 #################################################################################
 """
 Specialize the generic `Runner` class to running a flowsheet,
-in `FlowsheetRunner`.
+in `StructuredFlowsheet`.
 """
 
 # stdlib
@@ -34,7 +34,7 @@
 
 class Context(dict):
     """Syntactic sugar for the dictionary for the 'context' passed into each
-    step of the `FlowsheetRunner` class.
+    step of the `StructuredFlowsheet` class.
     """
 
     @property
@@ -58,7 +58,7 @@ def solver(self, value):
         self["solver"] = value
 
 
-class BaseFlowsheetRunner(Runner):
+class BaseStructuredFlowsheet(Runner):
     """Specialize the base `Runner` to handle IDAES flowsheets.
 
     This class pre-determine the name and order of steps to run
@@ -182,10 +182,10 @@ def annotate_var(
 
         ```{code}
         :name: annotate_vars
-        from idaes.core.util.structfs.fsrunner import FlowsheetRunner
+        from idaes.core.util.structfs.fsrunner import StructuredFlowsheet
         from pyomo.environ import *
 
-        def example(f: FlowsheetRunner):
+        def example(f: StructuredFlowsheet):
             v = Var()
             v.construct()
             f.annotate_var(v, key="example", title="Example variable").fix(1)
@@ -195,7 +195,7 @@ def example(f: FlowsheetRunner):
         property:
 
         ```{code}
-        example(fr := FlowsheetRunner())
+        example(fr := StructuredFlowsheet())
         print(fr.annotated_vars)
         # prints something like this:
         # {'example': {'var': <pyomo.core.base.var.ScalarVar object at 0x762ffb124b40>,
@@ -233,7 +233,7 @@ def annotated_vars(self) -> dict[str,]:
         return self._ann.copy()
 
 
-class FlowsheetRunner(BaseFlowsheetRunner):
+class StructuredFlowsheet(BaseStructuredFlowsheet):
     """Interface for running and inspecting IDAES flowsheets."""
 
     class DegreesOfFreedom:

idaes/core/util/structfs/runner.py

@@ -17,6 +17,7 @@
 # stdlib
 from abc import ABC, abstractmethod
 import logging
+from types import FunctionType
 from typing import Callable, Optional, Tuple, Sequence, TypeVar
 
 # third party
@@ -63,6 +64,32 @@ class Runner:
 
     STEP_ANY = "-"
 
+    class StepHelper:
+        """Instances of this class will return the decorator function when an
+        attribute is invoked, allowing @decorator syntax like `@flowsheet.step.build`.
+        Also allows `@flowsheet.step("build")`.
+        """
+
+        def __init__(self, method):
+            """Constructor with `method` to invoke with given name."""
+            self._method = method
+
+        def __getattr__(self, name: str):
+            self._check_type(name)
+            return self._method(name)
+
+        def __call__(self, name: str):
+            self._check_type(name, call=True)
+            return self._method(name)
+
+        def _check_type(self, name, call=False):
+            if isinstance(name, str):
+                return
+            message = f"Step name must be a string, got '{type(name)}'"
+            if call and isinstance(name, FunctionType):
+                message += ". You may have '.step' instead of '.step.name'"
+            raise TypeError(message)
+
     def __init__(self, steps: Sequence[str]):
         """Constructor.
 
@@ -73,6 +100,7 @@ def __init__(self, steps: Sequence[str]):
         self._actions: dict[str, ActionType] = {}
         self._step_names = list(steps)
         self._steps: dict[str, Step] = {}
+        self.step = self.StepHelper(self._step)
         self.reset()
 
     def __getitem__(self, key):
@@ -132,6 +160,10 @@ def run_step(self, name):
         """Syntactic sugar for calling `run_steps` for a single step."""
         self.run_steps(first=name, last=name)
 
+    def run(self):
+        """Syntactic sugar for run_steps() with no args."""
+        return self.run_steps()
+
     def run_steps(
         self, first: str = "", last: str = "", after: str = "", before: str = ""
     ):
@@ -306,7 +338,7 @@ def _substep_end(self, base: str, name: str):
         for action in self._actions.values():
             action.after_substep(base, name)
 
-    def step(self, name: str):
+    def _step(self, name: str):
         """Decorator function for creating a new step.
 
         Args:

idaes/core/util/structfs/runner_actions.py

@@ -38,7 +38,7 @@
 from idaes.core.util.model_statistics import degrees_of_freedom
 from idaes.core.base.unit_model import ProcessBlockData
 from .runner import Action
-from .fsrunner import FlowsheetRunner
+from .fsrunner import StructuredFlowsheet
 
 
 class Timer(Action):
@@ -198,7 +198,7 @@ class Report(BaseModel):
 
     def __init__(
         self,
-        runner: FlowsheetRunner,
+        runner: StructuredFlowsheet,
         flowsheet: str,
         steps: Union[str, list[str]],
         step_func: Optional[Callable[[str, UnitDofType], None]] = None,
@@ -415,7 +415,7 @@ class Report(BaseModel):
         variables: dict = Field(default={})
 
     def __init__(self, runner, **kwargs):
-        assert isinstance(runner, FlowsheetRunner)  # makes no sense otherwise
+        assert isinstance(runner, StructuredFlowsheet)  # makes no sense otherwise
         super().__init__(runner, **kwargs)
 
     def after_run(self):

idaes/core/util/structfs/tests/flash_flowsheet.py

@@ -23,9 +23,9 @@
     BTXParameterBlock,
 )
 from idaes.models.unit_models import Flash
-from idaes.core.util.structfs.fsrunner import FlowsheetRunner
+from idaes.core.util.structfs.fsrunner import StructuredFlowsheet
 
-FS = FlowsheetRunner()
+FS = StructuredFlowsheet()
 
 # # Flash Unit Model
 #