# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2019, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
# Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
# license information, respectively. Both files are also available online
# at the URL "https://github.com/IDAES/idaes-pse".

Supercritical CO2 Property Surrogate with OMLT Surrogate Object - Embedding Surrogate (Part 2)#

Maintainer: Javal Vyas

Author: Javal Vyas

Updated: 2024-01-24

1. Integration of Surrogate into Custom Property Package#

Here we shall see how to integrate the trained surrogate in the custom property package. One can read more about making a properties package from read the docs. To integrate the surrogate we first define the physical paramter block which will return the properties based on the state variables. State variables would be called from the State Block as Pyomo variables. We will define the surrogate input and output as pyomo variables as well. Once we have defined the variables in the state block then we define our surrogate block.

NOTE: For ease of explaination the property package is written in “.ipynb” format, ideally it should be in a python script. Each class of this package is separated in different cell for the same reason, in practive all the classes in this notebook should be part of the same python script. This folder includes “properties.py” file which is how embedding file should look like.

1.1 Steps in Creating a Property Package#

Creating a new property package can be broken down into the following steps, which will be demonstrated in the next part of this tutorial.

  1. Defining the units of measurement for the property package.

  2. Defining the properties supported by the property package and the associated metadata.

  3. Defining the phases and components of interest.

  4. Defining the necessary parameters required to calculate the properties of interest.

  5. Declaring the state variables to be used for the property package.

  6. Creating variables and constraints to describe the properties of interest.

  7. Creating an initialization routine for the property package.

  8. Defining interface methods used to couple the property package with unit models.

2. Importing libraries for making Property Package#

To begin with, we are going to need a number of components from the Pyomo modeling environment to construct the variables, constraints and parameters that will make up the property package, and we will also make use of the Pyomo units of measurement tools to define the units of our properties. We will also make use of a number of components and supporting methods from the IDAES modeling framework and libraries. We shall also use the Surrogate API in the IDAES framework to embed the trained surrogate in the property package.

# Import Python libraries
import logging

# Import Pyomo libraries
from pyomo.environ import Constraint, Param, \
      Reals, Set, value, Var, NonNegativeReals, units
from pyomo.opt import SolverFactory, TerminationCondition

# Import IDAES cores
from idaes.core import (declare_process_block_class,
from idaes.core.util.initialization import solve_indexed_blocks
from idaes.core.util.model_statistics import degrees_of_freedom
from idaes.core.util.misc import extract_data
from idaes.core.solvers import get_solver
from pyomo.util.check_units import assert_units_consistent
from idaes.core.surrogate.surrogate_block import SurrogateBlock
from idaes.core.surrogate.keras_surrogate import KerasSurrogate

from pyomo.util.model_size import build_model_size_report

# Some more information about this module
__author__ = "Javal Vyas"

# Set up logger
_log = logging.getLogger(__name__)
2024-03-18 23:17:33.337149: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.
2024-03-18 23:17:33.386565: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.
2024-03-18 23:17:33.387511: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-03-18 23:17:34.351388: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

3 Defining Classes#

We shall be going through each class of the property package in detail. Since there are not reactions occuring in the flowsheet we shall only write the Physical Parameter Block.

3.1 Physical Parameter Block#

The Physical Parameter Block serves as the central point of reference for all aspects of the property package, and needs to define a number of things about the package. These are summarized below:

  • Units of measurement

  • What properties are supported and how they are implemented

  • What components and phases are included in the packages

  • All the global parameters necessary for calculating properties

  • A reference to the associated State Block class, so that construction of the State Block components can be automated from the Physical Parameter Block

To assemble the above mentioned things in a class we need to follow the following steps:

  • Declaring the new class and inheriting from the PhysicalParameterBlock base class

  • Declaring any necessary configuration arguments

  • Writing the build method for our class

  • Creating a define_metadata method for the class.

The code below follows the above mentioned steps.

NOTE: The SCO2StateBlock will be discussed in the next section.

class PhysicalParameterData(PhysicalParameterBlock):
    Property Parameter Block Class

    Contains parameters and indexing sets associated with properties for
    supercritical CO2.

    def build(self):
        Callable method for Block construction.
        super(PhysicalParameterData, self).build()

        self._state_block_class  = SCO2StateBlock

        # List of valid phases in property package
        self.Liq = LiquidPhase()

        # Component list - a list of component identifiers
        self.CO2 = Component()

    def define_metadata(cls, obj):
                'flow_mol': {'method': None, 'units': 'kmol/s'},
                'pressure': {'method': None, 'units': 'MPa'},
                'temperature': {'method': None, 'units': 'K'},
                'enth_mol': {'method': None, 'units': 'kJ/kmol'},
                'entr_mol': {'method': None, 'units': 'kJ/kmol/K'}})

        obj.add_default_units({'time': units.s,
                               'length': units.m,
                               'mass': units.kg,
                               'amount': units.mol,
                               'temperature': units.K})

3.2 State Block#

After the Physical Parameter Block class has been created, the next step is to write the code necessary to create the State Blocks that will be used through out the flowsheet.

For this example, we will begin by describing the content of the StateBlockData objects, as this is where we create the variables and constraints that describe how to calculate the thermophysical properties of the material.

We start by defining the 5 state variables: flow_mol, pressure, temperature, enth_mol and entr_mol as the Pyomo Var, each of this variable has a unit for unit consistency. This is done in _make_state_vars function. We get the enth_mol and entr_mol variables from trained surrogate which we define in this function as well. To get the output variables from the surrogate:

  1. Define the input and output variables to the trained surrogate

  2. Load the surrogate from the folder it is saved in, here it is saved in the folder called keras_surrogate (look at the keras_training_doc.md file) using the keras Surrogate API of IDAES package

  3. Define a SurrogateBlock and call the build_model method on the block with the input variables, output variables, model formulation and the loaded surrogate as the arguements.

  4. Define the constraints necessary for ensuring physical feasibility of the system like the mass balance and energy balance. Check for the state variables to be within the bounds.

class SCO2StateBlockData(StateBlockData):
    An example property package for ideal gas properties with Gibbs energy

    def build(self):
        Callable method for Block construction
        super(SCO2StateBlockData, self).build()

    def _make_state_vars(self):        
        self.flow_mol = Var(domain=NonNegativeReals,
                            doc='Total molar flowrate [kmol/s]')
        self.pressure = Var(domain=NonNegativeReals,
                            bounds=(7.38, 40),
                            doc='State pressure [MPa]')
        self.temperature = Var(domain=NonNegativeReals,
                               bounds=(304.2, 760+273.15),
                               doc='State temperature [K]')
        self.entr_mol = Var(domain=Reals,
                            doc='Entropy [kJ/kmol/K]')
        self.enth_mol = Var(domain=Reals,
                            doc='Enthalpy [kJ/ kmol]')
        self.keras_surrogate = KerasSurrogate.load_from_folder("sco2_keras_surr")
        self.surrogate_enth = SurrogateBlock()
    def get_material_flow_terms(self, p, j):
        return self.flow_mol

    def get_enthalpy_flow_terms(self, p):
        return self.flow_mol*self.enth_mol

    def default_material_balance_type(self):
        return MaterialBalanceType.componentTotal

    def default_energy_balance_type(self):
        return EnergyBalanceType.enthalpyTotal

    def define_state_vars(self):
        return {"flow_mol": self.flow_mol,
                "temperature": self.temperature,
                "pressure": self.pressure}

    def model_check(blk):
        Model checks for property block
        # Check temperature bounds
        if value(blk.temperature) < blk.temperature.lb:
            _log.error('{} Temperature set below lower bound.'
        if value(blk.temperature) > blk.temperature.ub:
            _log.error('{} Temperature set above upper bound.'

        # Check pressure bounds
        if value(blk.pressure) < blk.pressure.lb:
            _log.error('{} Pressure set below lower bound.'.format(blk.name))
        if value(blk.pressure) > blk.pressure.ub:
            _log.error('{} Pressure set above upper bound.'.format(blk.name))

3.3 Define Initialization Routine#

After defining the variables and constraints required to describe the properties of interest for S-CO2, we need to provide them with a good initial guess. It is often the case that the default values provided to the variables while creating the model are not likely the actual conditions the user would simulate. Given the highly non-linear nature of the physical property calculations, it is more often than not impossible to solve a State Block without providing a set of good initial values for all the variables we have declared.

Any initialization routine can be written by following a 3 step process:

  1. Fix the state of the model such that there are no degrees of freedom. For State Blocks, it should only be necessary to fix the state variables to a set of initial guesses provided by the user or unit model, as well as deactivating any constraints like the sum of mole fractions.

  2. Iteratively build up a solution for the full model. This often involves multiple steps and can involve deactivating constraints and fixing some variables to reduce complexity, as well as analytically calculating values for variables based on the known state (and any previously calculated variables). Solvers can be called as part of any step to efficiently initialize large numbers of variables simultaneously.

  3. Return the state of the model to where it originally started (with the exception of variable values). Any variable that was fixed or constraint that was deactivated during initialization should be unfixed or reactivated, so that the degrees of freedom are restored to what they were before the initialization began.

Thus, we start with fixing the state variables. Here since enth_mol and entr_mol are a function of pressure and temperature, we do not fix them as fixing pressure and temperature would interm fix them. So, we check if a state variable if fixed or not, if it is fixed then we do not change them, if they are not fixed then we check for an initial guess from the state_args, if we get a value then we fix the varible with state_args, else we fix it with the value provided by the user. This should bring the degrees of freedom to 0. Here since we do not have any variable/constrained that we have unfixed/deactivated we can skip step 2 and move to step 3. We unfix the variables that were fixed in step 1 using the release_state function.

class _StateBlock(StateBlock):
    This Class contains methods which should be applied to Property Blocks as a
    whole, rather than individual elements of indexed Property Blocks.
    def initialize(blk, state_args=None, hold_state=False, outlvl=1,
                   state_vars_fixed=False, solver='ipopt',
                   optarg={'tol': 1e-8}):

        Initialisation routine for property package.

        Keyword Arguments:
            flow_mol : value at which to initialize component flows
            pressure : value at which to initialize pressure (default=None)
            temperature : value at which to initialize temperature
            outlvl : sets output level of initialisation routine

                     * 0 = no output (default)
                     * 1 = return solver state for each step in routine
                     * 2 = include solver output infomation (tee=True)
            state_vars_fixed: Flag to denote if state vars have already been
                              - True - states have already been fixed by the
                                       control volume 1D. Control volume 0D
                                       does not fix the state vars, so will
                                       be False if this state block is used
                                       with 0D blocks.
                             - False - states have not been fixed. The state
                                       block will deal with fixing/unfixing.
            optarg : solver options dictionary object (default=None)
            solver : str indicating whcih solver to use during
                     initialization (default = 'ipopt')
            hold_state : flag indicating whether the initialization routine
                         should unfix any state variables fixed during
                         initialization (default=False).
                         - True - states varaibles are not unfixed, and
                                 a dict of returned containing flags for
                                 which states were fixed during
                        - False - state variables are unfixed after
                                 initialization by calling the
                                 relase_state method

            If hold_states is True, returns a dict containing flags for
            which states were fixed during initialization.
        if state_vars_fixed is False:
            # Fix state variables if not already fixed
            Fcflag = {}
            Pflag = {}
            Tflag = {}

            for k in blk.keys():
                if blk[k].flow_mol.fixed is True:
                    Fcflag[k] = True
                    Fcflag[k] = False
                    if state_args is None:

                if blk[k].pressure.fixed is True:
                    Pflag[k] = True
                    Pflag[k] = False
                    if state_args is None:

                if blk[k].temperature.fixed is True:
                    Tflag[k] = True
                    Tflag[k] = False
                    if state_args is None:

            # If input block, return flags, else release state
            flags = {"Fcflag": Fcflag, "Pflag": Pflag,
                     "Tflag": Tflag}

            # Check when the state vars are fixed already result in dof 0
            for k in blk.keys():
                if degrees_of_freedom(blk[k]) != 0:
                    raise Exception("State vars fixed but degrees of freedom "
                                    "for state block is not zero during "
        if state_vars_fixed is False:
            if hold_state is True:
                return flags

    def release_state(blk, flags, outlvl=0):
        Method to relase state variables fixed during initialisation.

        Keyword Arguments:
            flags : dict containing information of which state variables
                    were fixed during initialization, and should now be
                    unfixed. This dict is returned by initialize if
            outlvl : sets output level of of logging
        if flags is None:

        # Unfix state variables
        for k in blk.keys():
            if flags['Fcflag'][k] is False:
            if flags['Pflag'][k] is False:
            if flags['Tflag'][k] is False:

        if outlvl > 0:
            if outlvl > 0:
                _log.info('{} State Released.'.format(blk.name))

Now we have our property package ready for being used in the flowsheet for optimization. We shall see that in the next part of this tutorial, flowsheet_optimization. To learn in detail about making a custom property package, one should go through Property Package Example.