Abstract
Dynamic simulation of distillation columns and chemical reactors remains essential for plant design, controllability analysis, and economic optimization. High-purity separations of close-boiling mixtures present significant computational challenges due to nonlinear thermodynamic behavior and stiff differential-algebraic equation (DAE) systems. This work presents a unified Python/JAX framework integrating four computational modules: (1) Peng-Robinson thermodynamics with complex-step differentiation, (2) nonlinear solvers (Newton, Broyden, Newton-Krylov) with automatic Curtis-Reid scaling, (3) DAE solver with Radau IIA collocation and intelligent auto-selection, and (4) constrained optimization using the Augmented Lagrangian Method with JAX automatic differentiation. The framework leverages JAX's just-in-time compilation (JIT), vectorization (vmap), and automatic differentiation (AD) to achieve near-compiled-language performance. Validation includes: nonlinear solver benchmarks with Newton-Krylov achieving 100% success across seven problems (n=2 to 5000), Williams-Otto reactor optimization with 0.06% deviation from published literature and 1325× real-time speedup, and a 180-stage propylene-propane splitter with 18× real-time performance under three concurrent disturbances. The framework shows that is possible to have an open-source alternative suitable for real-time optimization, operator training, and advanced process control applications.