Source code for watertap.ui.fsapi

"""
Simple flowsheet interface API
"""

__author__ = "Dan Gunter"

# stdlib
import logging
from collections import namedtuple
from enum import Enum
from typing import Any, Callable, Optional, Dict, Union, TypeVar
from types import ModuleType
from uuid import uuid4

try:
    from importlib import metadata
except ImportError:
    import importlib_metadata as metadata

# third-party
import idaes.logger as idaeslog
from pydantic import BaseModel, validator, Field
import pyomo.environ as pyo

#: Forward-reference to a FlowsheetInterface type, used in
#: :meth:`FlowsheetInterface.find`
FSI = TypeVar("FSI", bound="FlowsheetInterface")


_log = idaeslog.getLogger(__name__)


[docs]class UnsupportedObjType(TypeError):
[docs] def __init__( self, obj: Any, supported: Optional = None, ): msg = f"Object '{obj}' of type '{type(obj)}' is not supported." if supported is not None: msg += f"\nSupported: {supported}" super().__init__(msg) self.obj = obj self.supported = supported
[docs]class ModelExport(BaseModel): """A variable, expression, or parameter.""" _SupportedObjType = Union[ pyo.Var, pyo.Expression, pyo.Param, ] "Used for type hints and as a shorthand in error messages (i.e. not for runtime checks)" # TODO: if Optional[_SupportedObjType] is used for the `obj` type hint, # pydantic will run the runtime instance check which is not what we want # (as we want/need to use the pyomo is_xxx_type() methods instead) # so we're using Optional[object] unless we find a way to tell pydantic to skip this check obj: Optional[object] = Field(default=None, exclude=True) name: str = "" value: float = 0.0 ui_units: object = Field(default=None, exclude=True) display_units: str = "" rounding: float = 0 description: str = "" is_input: bool = True is_output: bool = True is_readonly: bool = None input_category: Optional[str] output_category: Optional[str] obj_key: str = None class Config: arbitrary_types_allowed = True @validator("obj", always=True, pre=True) def ensure_obj_is_supported(cls, v): if v is not None: cls._ensure_supported_type(v) return v @classmethod def _ensure_supported_type(cls, obj: object): is_valid = ( obj.is_variable_type() or obj.is_expression_type() or obj.is_parameter_type() # TODO: add support for numbers with pyo.numvalue.is_numeric_data() ) if is_valid: return True raise UnsupportedObjType(obj, supported=cls._SupportedObjType) @classmethod def _get_supported_obj( cls, values: dict, field_name: str = "obj", allow_none: bool = False ): obj = values.get(field_name, None) if not allow_none and obj is None: raise TypeError(f"'{field_name}' is None but allow_none is False") cls._ensure_supported_type(obj) return obj # NOTE: IMPORTANT: all validators used to set a dynamic default value # should have the `always=True` option, or the validator won't be called # when the value for that field is not passed # (which is precisely when we need the default value) # additionally, `pre=True` should be given if the field can at any point # have a value that doesn't match its type annotation # (e.g. `None` for a strict (non-`Optional` `bool` field) # Get value from object @validator("value", always=True) def validate_value(cls, v, values): if values.get("obj", None) is None: return v obj = cls._get_supported_obj(values, allow_none=False) return pyo.value(obj) # Derive display_units from ui_units @validator("display_units", always=True) def validate_units(cls, v, values): if not v: u = values.get("ui_units", pyo.units.dimensionless) v = str(pyo.units.get_units(u)) return v # set name dynamically from object @validator("name", always=True) def validate_name(cls, v, values): if not v: obj = cls._get_supported_obj(values, allow_none=False) try: v = obj.name except AttributeError: pass return v @validator("is_readonly", always=True, pre=True) def set_readonly_default(cls, v, values): if v is None: v = True obj = cls._get_supported_obj(values, allow_none=False) if obj.is_variable_type() or (obj.is_parameter_type() and obj.mutable): v = False return v @validator("obj_key", always=True, pre=True) def set_obj_key_default(cls, v, values): if v is None: obj = cls._get_supported_obj(values, allow_none=False) v = str(obj) return v
[docs]class FlowsheetExport(BaseModel): """A flowsheet and its contained exported model objects.""" obj: object = Field(default=None, exclude=True) name: str = "" description: str = "" model_objects: Dict[str, ModelExport] = {} version: int = 2 requires_idaes_solver: bool = False # set name dynamically from object @validator("name", always=True) def validate_name(cls, v, values): if not v: try: v = values["obj"].name except (KeyError, AttributeError): pass if not v: v = "default" return v @validator("description", always=True) def validate_description(cls, v, values): if not v: try: v = values["obj"].doc except (KeyError, AttributeError): v = f"{values['name']} flowsheet" return v
[docs] def add(self, *args, data: Union[dict, ModelExport] = None, **kwargs) -> object: """Add a new variable (or other model object). There are a few different ways of invoking this function. Users will typically use this form:: add(obj=<pyomo object>, name="My value name", ..etc..) If these same name/value pairs are already in a dictionary, this form is more convenient:: add(data=my_dict_of_name_value_pairs) If you have an existing ModelExport object, you can add it more directly with:: add(my_object) # -- OR -- add(data=my_object) Args: *args: If present, should be a single non-named argument, which is a ModelExport object. Create by adding it. data: If present, create from this argument. If it's a dict, create from its values just as from the kwargs. Otherwise it should be a ModelExport object, and create by adding it. kwargs: Name/value pairs to create a ModelExport object. Raises: KeyError: If the name of the Pyomo object is the same as an existing one, i.e. refuse to overwrite. """ if len(args) > 1: raise ValueError(f"At most one non-keyword arg allowed. Got: {args}") if len(args) == 1: model_export = args[0] elif data is None: _log.debug(f"Create ModelExport from args: {kwargs}") model_export = ModelExport.parse_obj(kwargs) else: if isinstance(data, dict): model_export = ModelExport.parse_obj(data) else: model_export = data key = model_export.obj_key if key in self.model_objects: raise KeyError( f"Adding ModelExport object failed: duplicate key '{key}' (model_export={model_export})" ) if _log.isEnabledFor(logging.DEBUG): # skip except in debug mode _log.debug( f"Adding ModelExport object with key={key}: {model_export.dict()}" ) self.model_objects[key] = model_export return model_export
[docs]class Actions(str, Enum): """Known actions that can be run. Actions that users should not run directly (unless they know what they are doing) are prefixed with an underscore. """ build = "build" solve = "solve" export = "_export"
[docs]class FlowsheetInterface: """Interface between users, UI developers, and flowsheet models.""" #: Function to look for in modules. See :meth:`find`. UI_HOOK = "export_to_ui" #: Type of item in list ``MissingObjectError.missing``. #: ``key`` is the unique key assigned to the variable, #: ``name`` is the variable name in the flowsheet MissingObject = namedtuple("MissingObject", "key name")
[docs] class MissingObjectError(Exception): """Error returned if data in `load` refers to a variable not found in the target object. Use the `.missing` attribute of the error object to get the list of MissingObjects. """
[docs] def __init__(self, missing): num = len(missing) plural = "" if num == 1 else "s" things = [f"{m[1]}" for m in missing] super().__init__( f"{num} object{plural} not found in the model: {', '.join(things)}" ) self.missing = [ FlowsheetInterface.MissingObject(key=m[0], name=m[1]) for m in missing ]
[docs] def __init__( self, fs: FlowsheetExport = None, do_build: Callable = None, do_export: Callable = None, do_solve: Callable = None, **kwargs, ): """Constructor. Args: fs: An existing wrapper to a flowsheet object. If this is not provided, then one will be constructed by passing the keyword arguments to the built-in pydantic ``parse_obj()`` method of :class:`FlowsheetExport`. do_build: Function to call to build the flowsheet. It should build the flowsheet model and return the `FlowsheetBlock`, which is typically the `fs` attribute of the model object. **Required** do_export: Function to call to export variables after the model is built. This will be called automatically by :meth:`build()`. **Required** do_solve: Function to solve the model. It should return the result that the solver itself returns. **Required** **kwargs: See `fs` arg. If the `fs` arg *is* provided, these are ignored. """ if fs is None: self.fs_exp = FlowsheetExport.parse_obj(kwargs) else: self.fs_exp = fs self._actions = {} for arg, name in ( (do_export, "export"), (do_build, "build"), (do_solve, "solve"), ): if arg: if not callable(arg): raise TypeError(f"'do_{name}' argument must be callable") self.add_action(getattr(Actions, name), arg) else: raise ValueError(f"'do_{name}' argument is required")
[docs] def build(self, **kwargs): """Build flowsheet Args: **kwargs: User-defined values Returns: None Raises: RuntimeError: If the build fails """ try: self.run_action(Actions.build, **kwargs) except Exception as err: raise RuntimeError(f"Building flowsheet: {err}") from err return
[docs] def solve(self, **kwargs): """Solve flowsheet. Args: **kwargs: User-defined values Returns: Return value of the underlying solve function Raises: RuntimeError: if the solver did not terminate in an optimal solution """ try: result = self.run_action(Actions.solve, **kwargs) except Exception as err: raise RuntimeError(f"Solving flowsheet: {err}") from err return result
[docs] def dict(self) -> Dict: """Serialize. Returns: Serialized contained FlowsheetExport object """ return self.fs_exp.dict(exclude={"obj"})
[docs] def load(self, data: Dict): """Load values from the data into corresponding variables in this instance's FlowsheetObject. Args: data: The input flowsheet (probably deserialized from JSON) """ u = pyo.units fs = FlowsheetExport.parse_obj(data) # new instance from data # Set the value for each input variable missing = [] # 'src' is the data source and 'dst' is this flowsheet (destination) for key, src in fs.model_objects.items(): # get corresponding exported variable try: dst = self.fs_exp.model_objects[key] except KeyError: missing.append((key, src.name)) continue # set value in this flowsheet ui_units = dst.ui_units if dst.is_input and not dst.is_readonly: # create a Var so Pyomo can do the unit conversion for us tmp = pyo.Var(initialize=src.value, units=ui_units) tmp.construct() # Convert units when setting value in the model dst.obj.value = u.convert(tmp, to_units=u.get_units(dst.obj)) # Don't convert units when setting the exported value dst.value = src.value if missing: raise self.MissingObjectError(missing)
[docs] def add_action(self, action_name: str, action_func: Callable): """Add an action for the flowsheet. Args: action_name: Name of the action to take (see :class:`Actions`) action_func: Function to call for the action Returns: None """ def action_wrapper(**kwargs): if action_name == Actions.build: # set new model object from return value of build action action_result = action_func(**kwargs) if action_result is None: raise RuntimeError( f"Flowsheet `{Actions.build}` action failed. " f"See logs for details." ) self.fs_exp.obj = action_result # [re-]create exports (new model object) if Actions.export not in self._actions: raise KeyError( "Error in 'build' action: no export action defined. " "Add `do_export=<function>` to FlowsheetInterface " "constructor or call `add_action(Actions.export, <function>)` " "on FlowsheetInterface instance." ) # clear model_objects dict, since duplicates not allowed self.fs_exp.model_objects.clear() # use get_action() since run_action() will refuse to call it directly self.get_action(Actions.export)(exports=self.fs_exp) result = None elif self.fs_exp.obj is None: raise RuntimeError( f"Cannot run any flowsheet action (except " f"'{Actions.build}') before flowsheet is built" ) else: result = action_func(flowsheet=self.fs_exp.obj, **kwargs) # Issue 755: Report optimization errors if action_name == Actions.solve: _log.debug(f"Solve result: {result}") if result is None: raise RuntimeError("Solver did not return a result") if not pyo.check_optimal_termination(result): raise RuntimeError(f"Solve failed: {result}") # Sync model with exported values if action_name in (Actions.build, Actions.solve): self.export_values() return result self._actions[action_name] = action_wrapper
[docs] def get_action(self, name: str) -> Union[Callable, None]: """Get the function for an ``add()``-ed action. Args: name: Name of the action (see :class:`Actions`) Returns: Function for this action Raises: KeyError, if no such action is defined """ return self._actions[name]
[docs] def run_action(self, name, **kwargs): """Run the named action.""" func = self.get_action(name) if name.startswith("_"): raise ValueError( f"Refusing to call '{name}' action directly since its " f"name begins with an underscore" ) return func(**kwargs)
[docs] def export_values(self): """Copy current values in underlying Pyomo model into exported model. Side-effects: Attribute ``fs_exp`` is modified. """ _log.info("Exporting values from flowsheet model to UI") u = pyo.units for key, mo in self.fs_exp.model_objects.items(): mo.value = pyo.value(u.convert(mo.obj, to_units=mo.ui_units))
[docs] @classmethod def from_installed_packages( cls, group_name: str = "watertap.flowsheets" ) -> Dict[str, "FlowsheetInterface"]: """Get all flowsheet interfaces defined as entry points within the Python packages installed in the environment. This uses the :func:`importlib.metadata.entry_points` function to fetch the list of flowsheets declared as part of a Python package distribution's `entry points <https://docs.python.org/3/library/importlib.metadata.html#entry-points>`_ under the group ``group_name``. To set up a flowsheet interface for discovery, locate your Python package distribution's file (normally :file:`setup.py`, :file:`pyproject.toml`, or equivalent) and add an entry in the ``entry_points`` section. For example, to add a flowsheet defined in :file:`watertap/examples/flowsheets/my_flowsheet.py` so that it can be discovered with the name ``my_flowsheet`` wherever the ``watertap`` package is installed, the following should be added to WaterTAP's :file:`setup.py`:: setup( name="watertap", # other setup() sections entry_points={ "watertap.flowsheets": [ # other flowsheet entry points "my_flowsheet = watertap.examples.flowsheets.my_flowsheet", ] } ) Args: group_name: The entry_points group from which the flowsheet interface modules will be populated. Returns: Mapping with keys the module names and values FlowsheetInterface objects """ eps = metadata.entry_points() try: # this happens for Python 3.7 (via importlib_metadata) and Python 3.10+ entry_points = list(eps.select(group=group_name)) except AttributeError: # this will happen on Python 3.8 and 3.9, where entry_points() has dict-like group selection entry_points = list(eps[group_name]) if not entry_points: _log.error(f"No interfaces found for entry points group: {group_name}") return {} interfaces = {} _log.debug(f"Loading {len(entry_points)} entry points") for ep in entry_points: _log.debug(f"ep = {ep}") module_name = ep.value try: module = ep.load() except ImportError as err: _log.error(f"Cannot import module '{module_name}': {err}") continue interface = cls.from_module(module) if interface: interfaces[module_name] = interface return interfaces
[docs] @classmethod def from_module( cls, module: Union[str, ModuleType] ) -> Optional["FlowsheetInterface"]: """Get a a flowsheet interface for module. Args: module: The module Returns: A flowsheet interface or None if it failed """ if not isinstance(module, ModuleType): module = importlib.import_module(module) # Get function that creates the FlowsheetInterface func = getattr(module, cls.UI_HOOK, None) if func is None: _log.warning( f"Interface for module '{module}' is missing UI hook function: " f"{cls.UI_HOOK}()" ) return None # Call the function that creates the FlowsheetInterface try: interface = func() except Exception as err: _log.error( f"Cannot get FlowsheetInterface object for module '{module}': {err}" ) return None # Return created FlowsheetInterface return interface