Diafiltration Flowsheet Optimization Tutorial#
This tutorial demonstrates how to set up, initialize, and optimize a process flowsheet with costing from the PrOMMiS QGESS framework. We build on the example optimization script in a three-stage diafiltration process diafiltration_flowsheet_optimization_example.py.
Objectives:
Build and initialize the diafiltration flowsheet
Apply costing and constraints
Run baseline simulation (before optimization)
Perform optimization
Compare results before and after optimization
Step 1: Import and Setup#
Import the core packages from Pyomo, IDAES, and PrOMMiS that will be needed for the optimization. These modules give access to the flowsheet configuration, costing models, and optimization tools.
The script diafiltration.py defines the diafiltration flowsheet: it builds the model, initializes it, and solves for membrane lengths that satisfy product specifications. In this base flowsheet, the active constraints include product recovery rate and purity requirement. A detailed walk-through of the flowsheet configuration is provided in the tutorial on lithium and cobalt separation via diafiltration.
The script diafiltration_flowsheet_optimization_example.py extends this configuration by integrating the costing framework from PrOMMiS (using QGESSCosting in ree_plant_capcost.py. In later steps, this model will be extended with costing and optimization to determine stage lengths and operating parameters that minimize recovery cost while satisfying process constraints.
from pyomo.environ import value, units as pyunits
from idaes.core.util.model_diagnostics import DiagnosticsToolbox
from idaes.core.util.model_statistics import degrees_of_freedom
from prommis.nanofiltration.diafiltration import (
build_model,
initialize_model,
solve_model,
)
from prommis.uky.costing.diafiltration_flowsheet_optimization_example import (
build_costing,
build_optimization,
scale_and_solve_model,
apply_baseline_lengths,
print_stage_cuts,
print_io_snap,
)
Step 2: Build and Initialize Model#
In this step we construct the diafiltration flowsheet and prepare it for simulation. Building the model creates all unit operations (stages, mixers, product outlets, etc.) and links them together through material balances. Initialization then provides reasonable starting values for the solver, ensuring that nonlinear equations converge to a feasible solution.
By following this workflow—build → initialize → solve—we ensure that the model is robust and ready for the costing and optimization steps that follow.
m = build_model()
dt = DiagnosticsToolbox(m)
print("Degrees of freedom:", degrees_of_freedom(m))
initialize_model(m)
solve_model(m, tee=False)
dt.assert_no_numerical_warnings()
Degrees of freedom: 0
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,1].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,2].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,3].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,4].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,5].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,6].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,7].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,8].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,9].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage1.permeate[0.0,10].flow_vol' to a
numeric value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
2025-11-12 14:40:04 [INFO] idaes.init.fs.stage1: Stream Initialization Completed.
2025-11-12 14:40:04 [INFO] idaes.init.fs.stage1: Initialization Completed, optimal - <undefined>
2025-11-12 14:40:04 [INFO] idaes.init.fs.mix1: Initialization Complete: optimal - <undefined>
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,1].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,2].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,3].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,4].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,5].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,6].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,7].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,8].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,9].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage2.permeate[0.0,10].flow_vol' to a
numeric value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
2025-11-12 14:40:04 [INFO] idaes.init.fs.stage2: Stream Initialization Completed.
2025-11-12 14:40:04 [INFO] idaes.init.fs.stage2: Initialization Completed, optimal - <undefined>
2025-11-12 14:40:04 [INFO] idaes.init.fs.mix2: Initialization Complete: optimal - <undefined>
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,1].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,2].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,3].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,4].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,5].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,6].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,7].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,8].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,9].flow_vol' to a numeric
value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
WARNING (W1002): Setting Var 'fs.stage3.permeate[0.0,10].flow_vol' to a
numeric value `0.0` outside the bounds (1e-08, None).
See also https://pyomo.readthedocs.io/en/stable/errors.html#w1002
2025-11-12 14:40:05 [INFO] idaes.init.fs.stage3: Stream Initialization Completed.
2025-11-12 14:40:05 [INFO] idaes.init.fs.stage3: Initialization Completed, optimal - <undefined>
2025-11-12 14:40:05 [INFO] idaes.init.fs.stage3: Stream Initialization Completed.
2025-11-12 14:40:05 [INFO] idaes.init.fs.stage3: Initialization Completed, optimal - <undefined>
Step 3: Apply Costing#
In this step we apply the QGESS-based costing framework from PrOMMiS. In the build_costing(m) function, the costing model estimates capital and operating expenses for the diafiltration process, focusing on membrane modules, pumps, and associated installation and maintenance.
At this stage, only equipment-related costs (capital, installation labor, and maintenance) are included. No explicit operational labor costs, resource consumption costs, or product sales revenues for Li or Co are considered in the baseline. As a result, the costing report shows:
Zero annual technical, administrative and support labor cost
Zero variable chemicals cost
Zero sales revenue
The costing report therefore primarily reflects the capital expenditures (CapEx) for membranes and pumps and the fixed and variable O&M linked to equipment performance (e.g., pumping power, pressure-drop energy). Additional elements, such as product pricing or detailed resource usage, can be layered on in later analyses if required.
Tip: Run the cell and review the baseline costing report to understand what costs are included at this stage.
build_costing(m)
solve_model(m, tee=False)
dt.assert_no_numerical_warnings()
Step 4: Baseline Report (before Optimization)#
With the costing framework applied, it allows us to generate a baseline report before running any optimization. This report reflects the model’s performance and cost based on the initial guesses for decision variables (such as stage lengths and operating conditions).
Because the solver starts from these initial values, the baseline results can vary if users provide different initial guesses. The baseline is therefore not an optimized solution, but rather a reference point that illustrates how the system performs under the default initialization. Subsequent optimization will adjust these variables to minimize the recovery cost while satisfying process constraints.
Tip: Run the cell and compare the reported flows, purities, and costs. These results will serve as your reference point when evaluating optimization improvements.
apply_baseline_lengths(m, L1=754, L2=758, L3=756)
solve_model(m, tee=False) # lock in the baseline
print_io_snap(m.fs, tag="BEFORE OPTIMIZATION")
print_stage_cuts(m, label="STAGE CUTS — BEFORE OPTIMIZATION")
print(
"Stage lengths (before):",
[value(m.fs.stage1.length), value(m.fs.stage2.length), value(m.fs.stage3.length)],
pyunits.get_units(m.fs.stage1.length),
)
m.fs.costing.report()
========================================================================
I/O SNAPSHOT: BEFORE OPTIMIZATION
========================================================================
[FEED (initial; stage3.retentate_side_stream_state[0,10])]
flow_vol (m³/hr): 100.0
conc_Li (kg/m³): 1.7
conc_Co (kg/m³): 17.0
[DIAFILTRATE (initial; mix2.inlet_1)]
flow_vol (m³/hr): 30.0
conc_Li (kg/m³): 0.1
conc_Co (kg/m³): 0.2
[PRODUCT PERMEATE (stage3.permeate_outlet)]
flow_vol (m³/hr): 113.40000000000002
Li concentration (kg/m³): 1.441478886659757
Co concentration (kg/m³): 5.440970271237354
Li recovery fraction: 0.9615512102777439
purity_Li: 0.20944272214575888
purity_Co: 0.7905572778542411
[PRODUCT RETENTATE (stage1.retentate_outlet)]
flow_vol (m³/hr): 16.60000000000002
Li concentration (kg/m³): 0.5744755573954007
Co concentration (kg/m³): 65.60204646033733
Co recovery fraction: 0.6405846889656477
purity_Li: 0.00868095723195401
purity_Co: 0.9913190427680459
[PARAMETERS]
[PARAMETERS]
sieving_coefficient_Li : 1.3
sieving_coefficient_Co : 0.5
membrane_width (m.w): 1.5
operating_pressure (psi): 145.0
========================================================================
============================================================
STAGE CUTS — BEFORE OPTIMIZATION
============================================================
Stage 1 cut (Q_perm/Q_feed): 0.872012
Stage 2 cut (Q_perm/Q_feed): 0.467132
Stage 3 cut (Q_perm/Q_feed): 0.465326
============================================================
Stage lengths (before): [754, 758, 756] m
====================================================================================
costing
------------------------------------------------------------------------------------
Value
Plant Cost Units MUSD_2021
Total Plant Cost 0.52158
Total Bare Erected Cost 0.26079
Total Installation Cost 0.26079
Total Other Plant Costs 1.0000e-12
Total Fixed Operating & Maintenance Cost 0.049667
Total Annual Operating Labor Cost 1.1142e-12
Total Annual Technical Labor Cost 1.0932e-12
Summation of Annual Labor Costs 2.2604e-12
Total Maintenance and Material Cost 0.010432
Total Quality Assurance and Control Cost 1.5398e-13
Total Sales, Patenting and Research Cost 8.0682e-11
Summation of Sales, Admin and Insurance Cost 0.0052158
Total Admin Support and Labor Cost 2.6539e-13
Total Property Taxes and Insurance Cost 0.0052158
Total Other Fixed Costs 1.0000e-12
Total Variable Waste Cost 0.0000
Total Variable Chemicals Cost 0.0000
General Plant Overhead Cost 0.046936
Total Plant Overhead Cost, Including Maintenance & Quality Assurance 0.057367
Total Variable Operating & Maintenance Cost 0.23195
Total Land Cost 0.0000
Total Sales Revenue Cost 1.6128e-08
Cost of Recovery (USD/kg REE) 0.033803
====================================================================================
Notice that in the costing report, several categories are shown as zero. This is expected because:
Labor costs (technical or administrative) are not included at this stage.
Resource costs (e.g., utilities beyond pumping power) are not explicitly modeled.
Product revenues from Li and Co sales are also excluded.
As a result, the annual labor cost, administrative/support cost, variable chemical cost, and sales revenue all appear as zero in the baseline output.
At the same time, observe that the total variable operating and maintenance (O&M) cost is not zero. This comes primarily from two sources:
1. Membrane pressure drop, which contributes to energy consumption.
2. Pumping costs, driven by the operating pressure parameter defined in diafiltration.py.
Both of these are considered operating costs, and their dependence on the specified operating pressure explains why the total variable O&M cost remains non-zero even in the baseline case.
Step 5: Optimization#
The goal of the optimization is to minimize the cost of recovery for lithium and cobalt while still meeting the technical requirements for lithium and cobalt separation. Specifically, the optimization problem is subject to the following constraints:
Li recovery ≥ 0.945
Co recovery ≥ 0.635
Li concentration ≤ 20 kg/m³
Co concentration ≤ 200 kg/m³
Stage membrane lengths: [0.1, 10,000] m
Stage cuts: [0.01, 0.99]
Through this formulation, the optimization reallocates stage lengths and adjusts separation parameters to satisfy product recovery targets while reducing overall costs.
Tip: Run the solver and confirm that optimization completes successfully. You should see differences compared to the baseline results.
build_optimization(m)
scale_and_solve_model(m)
dt.assert_no_numerical_warnings()
WARNING: model contains export suffix 'scaling_factor' that contains 1 keys
that are not Var, Constraint, Objective, or the model. Skipping.
Step 6: Final Report (After Optimization)#
After the optimization completes, the model generates a final report summarizing the optimized solution. This includes updated values of the key decision variables (e.g., stage lengths, and stage cuts), important process variables (such as recovery purities of the two product streams), and the detailed costing breakdown.
Compare these outputs to the baseline report from Step 4 to see how the optimization shifted resources and reduced the overall cost of recovery.
print_io_snap(m.fs, tag="AFTER OPTIMIZATION")
print_stage_cuts(m, label="STAGE CUTS — AFTER OPTIMIZATION")
print(
"Stage lengths (after):",
[value(m.fs.stage1.length), value(m.fs.stage2.length), value(m.fs.stage3.length)],
pyunits.get_units(m.fs.stage1.length),
)
m.fs.costing.report()
========================================================================
I/O SNAPSHOT: AFTER OPTIMIZATION
========================================================================
[FEED (initial; stage3.retentate_side_stream_state[0,10])]
flow_vol (m³/hr): 100.0
conc_Li (kg/m³): 1.7
conc_Co (kg/m³): 17.0
[DIAFILTRATE (initial; mix2.inlet_1)]
flow_vol (m³/hr): 30.0
conc_Li (kg/m³): 0.1
conc_Co (kg/m³): 0.2
[PRODUCT PERMEATE (stage3.permeate_outlet)]
flow_vol (m³/hr): 113.66806527086233
Li concentration (kg/m³): 1.4382702904865812
Co concentration (kg/m³): 5.285890171746729
Li recovery fraction: 0.9616788309186519
purity_Li: 0.21389589058213598
purity_Co: 0.7861041094178639
[PRODUCT RETENTATE (stage1.retentate_outlet)]
flow_vol (m³/hr): 16.33193472913766
Li concentration (kg/m³): 0.5825763390331373
Co concentration (kg/m³): 67.66884078783221
Co recovery fraction: 0.6500959358489906
purity_Li: 0.008535739821346826
purity_Co: 0.9914642601786532
[PARAMETERS]
[PARAMETERS]
sieving_coefficient_Li : 1.3
sieving_coefficient_Co : 0.5
membrane_width (m.w): 1.5
operating_pressure (psi): 145.0
========================================================================
============================================================
STAGE CUTS — AFTER OPTIMIZATION
============================================================
Stage 1 cut (Q_perm/Q_feed): 0.916532
Stage 2 cut (Q_perm/Q_feed): 0.309411
Stage 3 cut (Q_perm/Q_feed): 0.522212
============================================================
Stage lengths (after): [1195.5698835330784, 584.4438796124447, 757.7871018057486] m
====================================================================================
costing
------------------------------------------------------------------------------------
Value
Plant Cost Units MUSD_2021
Total Plant Cost 0.55477
Total Bare Erected Cost 0.27739
Total Installation Cost 0.27739
Total Other Plant Costs 1.0000e-12
Total Fixed Operating & Maintenance Cost 0.054710
Total Annual Operating Labor Cost 1.0000e-12
Total Annual Technical Labor Cost 1.0000e-12
Summation of Annual Labor Costs 2.0000e-12
Total Maintenance and Material Cost 0.011095
Total Quality Assurance and Control Cost 1.0000e-13
Total Sales, Patenting and Research Cost 8.0640e-11
Summation of Sales, Admin and Insurance Cost 0.0055477
Total Admin Support and Labor Cost 2.0000e-13
Total Property Taxes and Insurance Cost 0.0055477
Total Other Fixed Costs 1.0000e-12
Total Variable Waste Cost 0.0000
Total Variable Chemicals Cost 0.0000
General Plant Overhead Cost 0.045387
Total Plant Overhead Cost, Including Maintenance & Quality Assurance 0.056483
Total Variable Operating & Maintenance Cost 0.21761
Total Land Cost 0.0000
Total Sales Revenue Cost 1.6128e-08
Cost of Recovery (USD/kg REE) 0.032835
====================================================================================
Key Observations#
By comparing the baseline and optimized results, you can observe that:
The optimization successfully identified a feasible optimal solution within the specified constraints.
Stage lengths were reallocated across the three stages, improving the efficiency of lithium and cobalt separation.
The cost of recovery decreased from 0.033803 to 0.032835 USD/kg REE, demonstrating clear economic improvement relative to the baseline.
Significant adjustments occurred in membrane length, stage cuts, reflecting how the model reallocates resources to meet recovery targets.
The optimized solution achieved slightly higher Li and Co purity in the permeate stream and Li, Co recovery fraction, while also reducing overall costs.
Takeaway: Optimization helps balance cost and purity by reallocating design parameters within realistic bounds.
Summary#
This case study demonstrates how the QGESS costing framework can be applied to optimize a diafiltration flowsheet. The optimization reduces the cost of recovery while meeting required recovery and purity constraints, and it highlights how stage lengths and operating parameters can be reallocated to improve overall process performance.
Next Steps: Future analyses could expand the costing framework by incorporating:
Product revenues from Li and Co sales,
Labor costs (technical and administrative), and
Resource costs such as utilities and consumables.
These additions would provide a more comprehensive economic assessment and enable direct comparisons between process alternatives under realistic market conditions.