Saving, Loading and Exporting Flowsheets

Saving, Loading and Exporting Flowsheets#

A difflow flowsheet is built in Python, and until it is written down the code that builds it is the model. difflow.serialize gives it a file format, so it can be saved, diffed under version control, sent to a service, or opened by something that never imported the module that built it. difflow.codegen gives it a source form instead: the same model as an editable Python script.

This notebook takes a reaction–separation recycle through both:

  1. build and solve it;

  2. save it as JSON and look at what is in the file;

  3. reload it and check the answer and the gradient survive exactly;

  4. see what cannot be written (a lambda, a function defined in the notebook), and the ways out: refs=, extras=, and your own module with trust=True;

  5. keep the solved results with the model, and warm-start from them;

  6. save other models: a single object, a dynamic flowsheet, a power network;

  7. export it as a Python script.

The reference is Saving and Loading Flowsheets.

import json
import tempfile
from pathlib import Path

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

from difflow import codegen, serialize
from difflow.database import get_species_data
from difflow.flowsheet import Flowsheet, Unit
from difflow.kinetics import mass_action_kinetics
from difflow.streams import make_stream
from difflow.thermo import IdealThermo
from difflow.units.cstr import CSTR, CSTRParams
from difflow.units.flash import Flash, FlashParams, Mixer, Splitter

workdir = Path(tempfile.mkdtemp())

1. A flowsheet to save#

n-hexane isomerises to 2-methylpentane in a CSTR; a flash and a splitter send most of the liquid back. The rate law is built with mass_action_kinetics rather than written as a Python function. That matters for saving: the factory records the reaction data it was built from, and that data can go in a file, where a closure cannot.

species = ["n_hexane", "2_methylpentane"]
thermo = IdealThermo({s: get_species_data(s) for s in species})

kin = mass_action_kinetics(
    [{"reactants": {species[0]: 1.0}, "products": {species[1]: 1.0},
      "rate_params": {"A": 0.05}}],                    # k = 0.05 1/s
    species_order=species,
)


def build():
    fs = Flowsheet(species_order=species)
    fs.add_feed("fresh_feed", make_stream({species[0]: 1.0, species[1]: 0.0},
                                          T=340.0, P=101325.0))
    fs.add_unit(Unit("mixer", Mixer(species, thermo),
                     ["fresh_feed", "recycle"], ["reactor_feed"]))
    fs.add_unit(Unit("reactor",
                     CSTR(CSTRParams(V=1.0, molar_density=10.0, **kin.params_kwargs()), thermo),
                     ["reactor_feed"], ["reactor_effluent"], params={"T_spec": 340.0}))
    fs.add_unit(Unit("flash", Flash(FlashParams(species_order=species), thermo),
                     ["reactor_effluent"], ["liquid", "vapor_product"], params={"T": 337.0}))
    fs.add_unit(Unit("splitter", Splitter(species),
                     ["liquid"], ["recycle", "liquid_product"], params={"split_frac": 0.8}))
    fs.add_recycle("recycle", "recycle")
    return fs


fs = build()
solved = fs.solve(tol=1e-10, max_iter=200)
for name in ("reactor_feed", "liquid_product"):
    s = solved[name]
    print(f"{name:15s} T = {float(s['T']):.4f} K   "
          + "  ".join(f"{k} = {float(v):.4f}" for k, v in s.items() if k.startswith("F_")))

reactor_feed    T = 337.6018 K   F_n_hexane = 3.6667  F_2_methylpentane = 1.3333
liquid_product  T = 337.0000 K   F_n_hexane = 0.6667  F_2_methylpentane = 0.3333

2. Save it, and look inside#

serialize.save writes JSON. The top level holds the format version, the species, the feeds, the units and the recycles.

path = serialize.save(fs, workdir / "isomerisation.json")
data = json.loads(path.read_text())
print(f"{path.name}: {path.stat().st_size} bytes")
print("top-level keys:", list(data))
print("format_version:", data["format_version"], "  recycles:", data["recycles"])

isomerisation.json: 10323 bytes
top-level keys: ['format_version', 'difflow_version', 'species_order', 'defaults', 'feeds', 'units', 'recycles']
format_version: 2   recycles: {'recycle': 'recycle'}

Each unit is written as the name its class is registered under in the operation catalog, plus its parameters. The reactor shows the three things JSON has no native form for, each written as a tagged value:

  • $callable: the rate law, stored as the mass_action_kinetics call that rebuilds it;

  • $array: JAX arrays such as stoich;

  • $float: non-finite numbers. K_eq is inf for an irreversible reaction, and a bare Infinity is not valid JSON.

reactor = next(u for u in data["units"] if u["name"] == "reactor")
print("operation:", reactor["operation"])
print("rate_fn  :", json.dumps(reactor["params"]["rate_fn"])[:160], "...")
print("stoich   :", reactor["params"]["stoich"])
print("K_eq     :", reactor["params"]["rate_params"]["K_eq"])
print("constructor arguments:", list(reactor["constructor"]))

operation: CSTR
rate_fn  : {"$callable": {"factory": "mass_action_kinetics", "attr": "rate_fn", "kwargs": {"reactions": [{"reactants": {"n_hexane": 1.0}, "products": {"2_methylpentane": 1 ...
stoich   : {'$array': [[-1.0], [1.0]]}
K_eq     : {'$array': [{'$float': 'inf'}]}
constructor arguments: ['thermo']

Objects the constructor takes outside Params (here the reactor’s thermo) go under constructor. IdealThermo is written as its species data. An optional constructor argument is written only when it was set away from its default. The mixer’s thermo is one, and without it the mixer would mix temperatures without an enthalpy balance:

mixer = next(u for u in data["units"] if u["name"] == "mixer")
print("mixer constructor arguments:", list(mixer["constructor"]))

mixer constructor arguments: ['species_order', 'thermo']

3. Reload it: same answer, same gradient#

A reloaded flowsheet should solve to exactly the original numbers, with no tolerance. Saving it again should also give the same bytes.

fs2 = serialize.load(path)
reloaded = fs2.solve(tol=1e-10, max_iter=200)

worst = max(
    float(jnp.max(jnp.abs(solved[s][k] - reloaded[s][k])))
    for s in solved for k in solved[s] if not isinstance(solved[s][k], str)
)
print("largest difference over every stream variable:", worst)
print("re-saving gives the same bytes:", serialize.to_json(fs2) == serialize.to_json(fs))

largest difference over every stream variable: 0.0
re-saving gives the same bytes: True

The file holds data, not a frozen solution. The reloaded flowsheet is an ordinary difflow model, so it is still differentiable end to end. make_objective_fn gives the sensitivity of the 2-methylpentane production rate to reactor volume:

def product(streams):
    return streams["liquid_product"]["F_2_methylpentane"] + streams["vapor_product"]["F_2_methylpentane"]

for label, model in (("original", fs), ("reloaded", fs2)):
    obj = model.make_objective_fn(product, tol=1e-10, max_iter=200)
    print(f"{label:9s} d(product)/dV = {float(jax.grad(obj)({'reactor.V': 1.0})['reactor.V']):.12f} mol/s per m^3")

original  d(product)/dV = 0.222222222193 mol/s per m^3

reloaded  d(product)/dV = 0.222222222193 mol/s per m^3

4. What cannot be written, and the ways out#

A function with no import path is refused#

A rate law written here, in the notebook, has no data form and no import path that would find it again in another session (the notebook is __main__). It is refused rather than dropped, because a file that quietly lost a reactor’s kinetics would reload into a different model that still looked plausible.

def first_order(c, T, p):
    return jnp.array([p["k"] * c[species[0]]])

custom = build()
reactor_op = custom.units[1].operation
reactor_op.params = reactor_op.params.update(rate_fn=first_order, rate_params={"k": 0.05})

try:
    serialize.to_json(custom)
except serialize.SerializationError as exc:
    print("SerializationError:", exc)

SerializationError: unit 'reactor' field 'rate_fn' holds a callable ('first_order'), which cannot be written to a file: a lambda, a closure or a function defined in a script or notebook has no import path to be found again by. Define it at the top level of a module, build it from data (difflow.kinetics.mass_action_kinetics for rate laws), or pass it by name with refs={'name': fn} on both save and load.

refs=: write the object by name#

If the function lives in code that travels with the file, write a reference to it. Pass the same name→object mapping on save and on load. The file then holds {"$ref": "first_order"}, and loading puts back the very same object.

refs = {"first_order": first_order}
text = serialize.to_json(custom, refs=refs)
print([line.strip() for line in text.splitlines() if "$ref" in line])

try:
    serialize.from_json(text)                     # forgot the namespace
except serialize.SerializationError as exc:
    print("without refs:", str(exc)[:110], "...")

back = serialize.from_json(text, refs=refs)
print("same function object:", back.units[1].operation.params.rate_fn is first_order)

['"$ref": "first_order"']
without refs: this flowsheet refers to 'first_order', which is defined in its code context rather than in the file. Supply i ...
same function object: True

extras=: supply or override constructor objects on load#

IdealThermo, CubicThermo, ActivityThermo and the cubic equations of state PengRobinson and SRK are written as data. A unit that needs any other property package is supplied with one on load, keyed by unit name. The same argument overrides what the file holds, which is how you re-run a saved flowsheet against a different property package.

fs3 = serialize.load(path, extras={"flash": {"thermo": thermo}})
print("flash uses the object passed in:", fs3.units[2].operation.thermo is thermo)

flash uses the object passed in: True

Format versions#

Every file carries format_version. difflow reads every version it still supports, so older files keep opening, and writes the lowest version that can hold the file: this flowsheet uses nothing newer than version 2, so a release that reads only 2 still opens it. A file from a newer release is rejected with a clear error rather than misread:

print("this file:", data["format_version"], " supported:", serialize.SUPPORTED_VERSIONS,
      " newest written:", serialize.FORMAT_VERSION)
future = dict(data, format_version=99)
try:
    serialize.from_dict(future)
except serialize.SerializationError as exc:
    print("SerializationError:", exc)

this file: 2  supported: (1, 2, 3)  newest written: 3
SerializationError: format version 99 is not supported; this difflow reads 1, 2, 3 and writes up to 3.

Your own module: written by import path, loaded with trust=True#

A function or class defined at the top level of a module has an import path, and is written by it. Here is a small module of our own, with a rate law in it:

import importlib
import sys

(workdir / "mylab.py").write_text('''
import jax.numpy as jnp

def first_order(c, T, p):
    """Isomerisation, first order in n-hexane."""
    return jnp.array([p["k"] * c["n_hexane"]])
''')
sys.path.insert(0, str(workdir))
mylab = importlib.import_module("mylab")

mine = build()
reactor_op = mine.units[1].operation
reactor_op.params = reactor_op.params.update(rate_fn=mylab.first_order, rate_params={"k": 0.05})
text = serialize.to_json(mine)
print([line.strip() for line in text.splitlines() if "mylab" in line])

['"function": "mylab:first_order"']

Loading it means importing mylab, which runs that module’s code. A file must not be able to do that unasked, so it is refused until you say you trust it:

try:
    serialize.from_json(text)
except serialize.SerializationError as exc:
    print("SerializationError:", str(exc)[:160], "...")

trusted = serialize.from_json(text, trust=True)
print("same function:", trusted.units[1].operation.params.rate_fn is mylab.first_order)

SerializationError: the file names 'mylab:first_order', from the module 'mylab', which is not part of difflow. Loading it would import that module and run its code, so it is only d ...
same function: True

5. Keep the solved results, and warm-start from them#

save(..., results=streams) keeps a solve’s streams with the model. They come back as saved_results, with what the solve reported, and warm_start hands the saved tear streams to the next solve:

path = serialize.save(fs, workdir / "with_results.json", results=solved)
fs4 = serialize.load(path)
print("saved product:", float(fs4.saved_results["liquid_product"]["F_2_methylpentane"]))
print("saved solve:", fs4.saved_solve_info)

fs.solve(tol=1e-10, max_iter=200)
print("cold start iterations:", fs.last_solve_iterations)
fs4.solve(tol=1e-10, max_iter=200, tear_initial=serialize.warm_start(fs4))
print("warm start iterations:", fs4.last_solve_iterations)

saved product: 0.33333333333333354
saved solve: {'converged': True, 'iterations': 8, 'residual': 4.440892098500626e-16, 'method': 'anderson', 'tear_streams': ['recycle']}

cold start iterations: 8

warm start iterations: 0

6. Other models#

A single object#

save_object / load_object (and dumps_object / loads_object) write one object: a Params, a thermo, an equation of state, a ReactionSet. A reloaded object compares equal to the one written:

from difflow import CubicThermo, PengRobinson, get_critical_props

cubic = CubicThermo(thermo, PengRobinson({s: get_critical_props(s) for s in species}))
path = serialize.save_object(cubic, workdir / "cubic_thermo.json")
print("CubicThermo equal after reload:", serialize.load_object(path) == cubic)

kin_back = serialize.loads_object(serialize.dumps_object(kin))
print("ReactionSet equations:", kin_back.equations)

CubicThermo equal after reload: True
ReactionSet equations: ['r = k(T) \\, C_{n_hexane}']

A dynamic flowsheet#

A DynamicFlowsheet is written unit by unit, each as its class and the arguments it was built with. The rate law here comes from our module, so loading needs trust=True again:

from difflow.dynamic import DynamicCSTR, DynamicFlowsheet, DynamicTank

dyn = DynamicFlowsheet(species_order=species)
dyn.add_feed("feed", make_stream({species[0]: 1.0, species[1]: 0.0}, T=340.0, P=101325.0))
dyn.add_unit(DynamicCSTR(volume=1.0, rate_fn=mylab.first_order, stoich=kin.stoich,
                         species_order=species, rate_params={"k": 0.05}, name="reactor"),
             inlet_names=["feed"], outlet_names=["reactor_out"])
dyn.add_unit(DynamicTank(max_volume=10.0, species_order=species, name="storage"),
             inlet_names=["reactor_out"], outlet_names=["product"])

dyn2 = serialize.from_json(serialize.to_json(dyn), trust=True)
a = dyn.simulate(t_span=(0.0, 100.0), method="RK4", n_steps=100)
b = dyn2.simulate(t_span=(0.0, 100.0), method="RK4", n_steps=100)
print("reloaded trajectory identical:", bool(jnp.all(a.y_final == b.y_final)))

reloaded trajectory identical: True

A power network#

A PowerNetwork is data all the way down. Its buses keep their insertion order, which the power-flow numbering depends on:

from difflow_power import case9
from difflow_power.powerflow import solve_power_flow

net = case9()
net2 = serialize.load(serialize.save(net, workdir / "case9.json"))
print("network equal:", net2 == net, "| bus order kept:", list(net2.buses) == list(net.buses))
print("power flow identical:", bool(jnp.all(solve_power_flow(net).vm == solve_power_flow(net2).vm)))

network equal: True | bus order kept: True

power flow identical: True

7. Export as Python#

codegen.to_python writes the same model as the code you would have typed: imports, thermodynamics, kinetics, then the flowsheet. Use it to leave the file format, for example to add something the operation palette does not offer. It refuses exactly what serialize refuses.

source = codegen.to_python(fs)
print(source)

"""Flowsheet generated by difflow 0.3.0.

Edit freely --- this is ordinary difflow code.
"""

from difflow import (
    CSTR,
    CSTRParams,
    Flash,
    FlashParams,
    Flowsheet,
    IdealThermo,
    Mixer,
    Splitter,
    Unit,
    get_species_data,
    make_stream,
    mass_action_kinetics,
)

thermo_mixer = IdealThermo({s: get_species_data(s) for s in ['n_hexane', '2_methylpentane']})
kinetics_reactor = mass_action_kinetics(reactions=[{'reactants': {'n_hexane': 1.0}, 'products': {'2_methylpentane': 1.0}, 'rate_params': {'A': 0.05}}], species_order=['n_hexane', '2_methylpentane'], reverse='error', orders=None)
thermo_reactor = IdealThermo({s: get_species_data(s) for s in ['n_hexane', '2_methylpentane']})
thermo_flash = IdealThermo({s: get_species_data(s) for s in ['n_hexane', '2_methylpentane']})

fs = Flowsheet(species_order=['n_hexane', '2_methylpentane'], default_flow=0.01, default_T=None, default_P=None)

fs.add_feed('fresh_feed', make_stream({'n_hexane': 1.0, '2_methylpentane': 0.0}, T=340.0, P=101325.0))

fs.add_unit(Unit('mixer', Mixer(['n_hexane', '2_methylpentane'], thermo=thermo_mixer), ['fresh_feed', 'recycle'], ['reactor_feed']))
fs.add_unit(Unit('reactor', CSTR(CSTRParams(**kinetics_reactor.params_kwargs(), V=1.0, dH_rxn=None, T_damping=0.3, molar_density=10.0, eos=None, reaction_phase=None, H_mix_fn=None, K_eq_fn=None, outlet_volumetric_basis=False), thermo=thermo_reactor), ['reactor_feed'], ['reactor_effluent'], params={'T_spec': 340.0}))
fs.add_unit(Unit('flash', Flash(FlashParams(species_order=['n_hexane', '2_methylpentane']), thermo_flash), ['reactor_effluent'], ['liquid', 'vapor_product'], params={'T': 337.0}))
fs.add_unit(Unit('splitter', Splitter(['n_hexane', '2_methylpentane']), ['liquid'], ['recycle', 'liquid_product'], params={'split_frac': 0.8}))

fs.add_recycle('recycle', 'recycle')


if __name__ == "__main__":
    streams = fs.solve()
    for name, stream in streams.items():
        print(name, stream)

namespace = {}
exec(source, namespace)
from_script = namespace["fs"].solve(tol=1e-10, max_iter=200)
print("script rebuilds the same answer:",
      all(float(jnp.max(jnp.abs(solved[s][k] - from_script[s][k]))) == 0.0
          for s in solved for k in solved[s] if not isinstance(solved[s][k], str)))

script rebuilds the same answer: True

Summary#

Goal

Use

Save / load a model as data

serialize.save(model, path), serialize.load(path) (or to_json / from_json, to_dict / from_dict)

A rate law that can be saved

build it with mass_action_kinetics, or define it in a module

Your own functions and classes

define them in a module; load with trust=True

A lambda or notebook function

refs={"name": obj} on both save and load

Another property package on load

load(path, extras={"unit": {"thermo": obj}})

Keep results, warm-start

save(fs, path, results=streams); fs.solve(tear_initial=serialize.warm_start(fs))

One object

save_object / load_object

Dynamic flowsheets, power networks

serialize.save / serialize.load, as for a flowsheet

An editable script

codegen.to_python(fs) / codegen.save_script(fs, path)

A grid of solves with exact gradients

difflow.sweep (the book’s app turns one into a static interactive page)

The round trip keeps the answer and the gradient exactly, because the file stores the model’s inputs (parameters, kinetics data, topology) and difflow rebuilds the same JAX computation from them.