Interactive Flowsheet Visualization

Interactive Flowsheet Visualization#

This notebook demonstrates the visualization capabilities of difflow for creating interactive process flow diagrams (PFDs) that look like chemical engineering flowsheets.

Features#

  • Custom P&ID-style icons for each unit operation type

  • Left-to-right hierarchical layout (typical for process flowsheets)

  • Interactive drag-and-drop to rearrange nodes

  • Tooltips showing stream data (composition, T, P) and unit parameters

  • Export to HTML for publications and sharing

Requirements#

pip install difflow[visualization]
# or
pip install ipycytoscape networkx

Setup and Imports#

import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

from difflow import (
    Flowsheet, Unit, make_stream,
    CSTR, Flash, Mixer, Heater, Cooler,
    IdealThermo, SpeciesData,
)
from difflow.visualization import (
    FlowsheetVisualizer,
    visualize_flowsheet,
    get_icon_svg,
    list_available_icons,
)

print("Imports successful!")

WARNING:2026-02-28 15:45:33,104:jax._src.xla_bridge:905: Platform 'mps' is experimental and not all JAX functionality may be correctly supported!

Imports successful!

Preview Available Icons#

The visualization module includes P&ID-style icons for all common unit operations. Let’s see what icons are available:

# List all available icon types
icons = list_available_icons()
print("Available icon types:")
for i, icon in enumerate(icons):
    print(f"  {icon}", end="\n" if (i+1) % 5 == 0 else ", ")

Available icon types:
  CSTR,   Centrifuge,   CoCurrentHX,   ContinuousBioreactor,   Cooler
  CounterCurrentHX,   Diafiltration,   DifferentialContactor,   DiscStackCentrifuge,   DistillationColumn
  FedBatchBioreactor,   FedBatchReactor,   Feed,   Flash,   GasPFR
  Heater,   IonExchangeChromatography,   LLEEquilibrium,   Mixer,   MultistageCascade
  PFR,   Product,   ProteinAChromatography,   SemiBatchReactor,   ShortcutColumn
  SizeExclusionChromatography,   Splitter,   TFF,   Ultrafiltration,   bioreactor
  cascade,   centrifuge,   chromatography,   cooler,   cstr
  distillation,   fed_batch,   feed,   filtration,   flash
  generic,   heat_exchanger,   heater,   mixer,   pfr
  product,   splitter, 

# Display some icons (as HTML in Jupyter)
from IPython.display import HTML, display

icon_types = ["CSTR", "PFR", "Flash", "Mixer", "Splitter", "Heater", "Cooler", 
              "distillation", "heat_exchanger", "bioreactor", "feed", "product"]

html = '<div style="display: flex; flex-wrap: wrap; gap: 20px;">'
for icon_type in icon_types:
    svg = get_icon_svg(icon_type)
    html += f'<div style="text-align: center;"><div style="width: 60px; height: 60px;">{svg}</div><small>{icon_type}</small></div>'
html += '</div>'

display(HTML(html))
CSTR
PFR
V L
Flash
Mixer
Splitter
Heater
Cooler
distillation
heat_exchanger
bioreactor
feed
product

Example 1: Simple CSTR + Flash Flowsheet#

Let’s visualize the classic reaction-separation flowsheet:

Feed → CSTR → Flash → Product
         ↑      ↓
         └──────┘ (recycle)

# Define thermodynamic properties
species_data = {
    "A": SpeciesData(
        name="A", MW=92.0,
        Cp_coeffs=(75.0, 0.0, 0.0, 0.0),
        Hvap_coeffs=(35000.0, 0.38, 590.0),
        antoine_coeffs=(10.0, 2000.0, -40.0),
        Hf=0.0,
    ),
    "B": SpeciesData(
        name="B", MW=78.0,
        Cp_coeffs=(50.0, 0.0, 0.0, 0.0),
        Hvap_coeffs=(30000.0, 0.38, 560.0),
        antoine_coeffs=(10.0, 1500.0, -40.0),
        Hf=-50000.0,
    ),
}
thermo = IdealThermo(species_data)
species_order = ["A", "B"]

# Create reaction kinetics
def rate_fn(C, T, params):
    k = params["A"] * jnp.exp(-params["Ea"] / (8.314 * T))
    return jnp.array([k * C["A"]])

stoich = jnp.array([[-1.0], [1.0]])

print("Thermodynamics and kinetics defined.")

Thermodynamics and kinetics defined.

from difflow.units.cstr import CSTRParams
from difflow.units.flash import FlashParams

# Create the flowsheet
fs = Flowsheet(species_order)

# Add feed
feed = make_stream({"A": 10.0, "B": 0.0}, T=300.0, P=101325.0)
fs.add_feed("feed", feed)

# Create unit operations
cstr_params = CSTRParams(
    V=1.0,
    rate_fn=rate_fn,
    stoich=stoich,
    rate_params={"A": 1e8, "Ea": 50000.0},
    species_order=species_order,
)
cstr = CSTR(cstr_params, thermo=thermo, mode="isothermal")

flash_params = FlashParams(species_order=species_order)
flash = Flash(flash_params, thermo=thermo)

mixer = Mixer(species_order, thermo=thermo)

# Add units to flowsheet
fs.add_unit(Unit(
    name="Mixer",
    operation=mixer,
    inlet_names=["feed", "recycle"],
    outlet_names=["mixer_out"],
))

fs.add_unit(Unit(
    name="Reactor",
    operation=lambda s, **kw: cstr(s, T_spec=400.0),
    inlet_names=["mixer_out"],
    outlet_names=["reactor_out"],
    params={"V": 1.0, "T": 400.0},
))

fs.add_unit(Unit(
    name="Flash",
    operation=lambda s, **kw: flash(s, T=350.0, P=101325.0),
    inlet_names=["reactor_out"],
    outlet_names=["liquid", "vapor"],
    params={"T": 350.0, "P": 101325.0},
))

# Add recycle
fs.add_recycle("liquid", "recycle")

print(f"Flowsheet created with {len(fs.units)} units")
print(f"Feeds: {list(fs.feeds.keys())}")
print(f"Recycles: {fs.recycles}")

Flowsheet created with 3 units
Feeds: ['feed']
Recycles: {'liquid': 'recycle'}

# Solve the flowsheet
recycle_init = {
    "recycle": make_stream({"A": 1.0, "B": 0.1}, T=350.0, P=101325.0)
}

print("Solving flowsheet...")
streams = fs.solve(tear_initial=recycle_init)
print("Solved!")

# Show stream summary
for name, stream in streams.items():
    total = float(stream["F_A"]) + float(stream["F_B"])
    print(f"  {name}: F={total:.2f} mol/s, T={float(stream['T']):.0f} K")

Solving flowsheet...

Solved!
  feed: F=10.00 mol/s, T=300 K
  recycle: F=0.00 mol/s, T=350 K
  mixer_out: F=10.00 mol/s, T=300 K
  reactor_out: F=10.00 mol/s, T=400 K
  liquid: F=0.00 mol/s, T=350 K
  vapor: F=10.00 mol/s, T=350 K

Visualize the Flowsheet#

Now let’s create an interactive visualization. You can:

  • Drag nodes to rearrange the layout

  • Hover over nodes to see unit parameters

  • Hover over edges to see stream compositions

  • Zoom and pan using mouse scroll and drag

# Quick visualization with one line
visualize_flowsheet(fs, streams=streams)

Advanced Visualization Options#

For more control, use the FlowsheetVisualizer class directly:

# Create visualizer with custom options
viz = FlowsheetVisualizer(
    fs,
    show_feeds=True,
    show_products=True,
    layout="dagre",  # Left-to-right hierarchical layout
    layout_options={
        "rankDir": "LR",  # Left to right
        "nodeSep": 100,   # Vertical spacing
        "rankSep": 150,   # Horizontal spacing
    },
)

# Show with custom size
widget = viz.show(streams=streams, height="400px")
widget

Export to HTML#

Export the visualization as a standalone HTML file for publications, reports, or sharing:

# Export to HTML
viz.export_html(
    "cstr_flash_flowsheet.html",
    title="CSTR + Flash Recycle Process",
)
print("Exported to: cstr_flash_flowsheet.html")
print("Open this file in a browser for full interactivity!")

Exported to: cstr_flash_flowsheet.html
Open this file in a browser for full interactivity!

Example 2: More Complex Flowsheet#

Let’s create a more complex flowsheet with multiple unit operations:

from difflow.units.heat_exchanger import Heater, HeaterParams, Cooler, CoolerParams

# Create a more complex flowsheet
fs2 = Flowsheet(species_order)

# Add feed
fs2.add_feed("raw_feed", make_stream({"A": 10.0, "B": 0.5}, T=298.0, P=101325.0))

# Heater to preheat feed
heater_params = HeaterParams(Cp=75.0)  # Liquid Cp
heater = Heater(heater_params)

fs2.add_unit(Unit(
    name="Feed_Heater",
    operation=lambda s, **kw: heater(s, T_out=350.0),
    inlet_names=["raw_feed"],
    outlet_names=["heated_feed"],
    params={"T_out": 350.0},
))

# Mixer
fs2.add_unit(Unit(
    name="Feed_Mixer",
    operation=mixer,
    inlet_names=["heated_feed", "recycle"],
    outlet_names=["mixed_feed"],
))

# CSTR
fs2.add_unit(Unit(
    name="CSTR_1",
    operation=lambda s, **kw: cstr(s, T_spec=400.0),
    inlet_names=["mixed_feed"],
    outlet_names=["cstr_out"],
    params={"V": 1.0, "T": 400.0},
))

# Cooler before flash
cooler_params = CoolerParams(Cp=75.0)
cooler = Cooler(cooler_params)

fs2.add_unit(Unit(
    name="Pre_Flash_Cooler",
    operation=lambda s, **kw: cooler(s, T_out=350.0),
    inlet_names=["cstr_out"],
    outlet_names=["cooled_stream"],
    params={"T_out": 350.0},
))

# Flash separator
fs2.add_unit(Unit(
    name="Flash_Drum",
    operation=lambda s, **kw: flash(s, T=350.0, P=101325.0),
    inlet_names=["cooled_stream"],
    outlet_names=["bottoms", "overhead"],
    params={"T": 350.0, "P": 101325.0},
))

# Add recycle
fs2.add_recycle("bottoms", "recycle")

print(f"Complex flowsheet with {len(fs2.units)} units")

Complex flowsheet with 5 units

# Solve
recycle_init2 = {
    "recycle": make_stream({"A": 1.0, "B": 0.1}, T=350.0, P=101325.0)
}
streams2 = fs2.solve(tear_initial=recycle_init2)
print("Solved!")

# Visualize
visualize_flowsheet(fs2, streams=streams2, height="450px")

Solved!

Export Graph Data#

You can also export the graph data for use with other visualization tools:

viz2 = FlowsheetVisualizer(fs2)

# Get raw graph data
graph_data = viz2.get_graph_data()
print(f"Nodes: {len(graph_data['nodes'])}")
print(f"Edges: {len(graph_data['edges'])}")

print("\nNode IDs:")
for node in graph_data['nodes']:
    print(f"  - {node['data']['id']} ({node['data']['type']})")

Nodes: 7
Edges: 7

Node IDs:
  - feed_raw_feed (feed)
  - Feed_Heater (<lambda>)
  - Feed_Mixer (Mixer)
  - CSTR_1 (<lambda>)
  - Pre_Flash_Cooler (<lambda>)
  - Flash_Drum (<lambda>)
  - product_overhead (product)

# Convert to NetworkX for further analysis
try:
    import networkx as nx
    G = viz2.to_networkx()
    print(f"NetworkX DiGraph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")
    
    # Check for cycles (recycle loops)
    if nx.is_directed_acyclic_graph(G):
        print(f"\nTopological order: {list(nx.topological_sort(G))}")
    else:
        cycles = list(nx.simple_cycles(G))
        print(f"\nGraph contains {len(cycles)} cycle(s) (recycle loops)")
        for i, cycle in enumerate(cycles):
            print(f"  Cycle {i+1}: {' -> '.join(cycle)}")
except ImportError:
    print("NetworkX not installed. Run: pip install networkx")

NetworkX DiGraph: 7 nodes, 7 edges

Graph contains 1 cycle(s) (recycle loops)
  Cycle 1: CSTR_1 -> Pre_Flash_Cooler -> Flash_Drum -> Feed_Mixer

Summary#

The visualization module provides:

  1. P&ID-style icons for all common unit operations

  2. Interactive drag-and-drop layout adjustment

  3. Rich tooltips with stream compositions and unit parameters

  4. HTML export for publications and sharing

  5. NetworkX integration for graph analysis

Tips for Best Results#

  • Use the dagre layout with rankDir: "LR" for left-to-right flow

  • Adjust nodeSep and rankSep to control spacing

  • Export to HTML for full interactivity in publications

  • Pass solved streams to show() for detailed tooltips