Add Support for Pyomo/OMLT Models in Inverse (nlp_based) and Forward (AIS2AOS) Mapping

[hfsf]

This PR introduces native support for Pyomo and OMLT (Optimization and Machine Learning Toolkit) models within the opyrability framework.

Previously, the library treated process models strictly as "black-box" callables, relying on finite differences or external JAX wrappers for derivatives. This update enables an equation-oriented (glass-box) approach, allowing the ipopt solver to access exact Hessians and Jacobians directly from the algebraic structure.

This is particularly useful for integrating Neural Networks (via OMLT) or rigorous thermodynamic models (via IDAES) into the operability analysis workflow.

Key Changes

1. nlp_based_approach (Inverse Mapping)

• Added a logic branch to detect if the input model is a Pyomo constructor (e.g., has build_pyomo_constraints attribute) or an OMLT object.
• If detected, the function constructs a pyomo.ConcreteModel and solves it directly using the IPOPT interface, bypassing scipy.optimize.
• Benefit: Significantly improved convergence stability and speed for complex inverse problems compared to the finite-difference approach.

2. AIS2AOS_map (Forward Mapping)

• Implemented a Simulation Proxy to handle Pyomo models in the forward mapping loop.
• Added a new optional argument output_dim (required for Pyomo models) to correctly initialize result arrays.
• Performance Optimization:
  • Solver Persistence: The implementation attempts to use appsi_ipopt (Auto-Persistent Solver Interface) if available. This keeps the solver loaded in memory, reducing I/O overhead by orders of magnitude.
  • Fallback Mechanism: If appsi_ipopt is not installed/configured, it gracefully falls back to the standard ipopt executable.
  • Warm-Starting: Implemented a smart warm-start strategy where the solution of the previous grid point is used as the initialization for the current point, accelerating the simulation loop.

3. implicit_map

• Added an explicit NotImplementedError guard. If a user attempts to pass a Pyomo model to implicit_map, the function returns a clear error message explaining that JAX/JIT dependencies make it currently incompatible.

4. Examples (shower.py)

• Updated the shower.py example file to include build_shower* functions. These serve as clean examples of how users should define algebraic models for the new interface, preserving the original NumPy functions for backward compatibility.

Performance Note on AIS2AOS_map

When using Pyomo models for forward mapping, there is an inherent overhead compared to pure NumPy callables due to the Python-to-Solver I/O.

• With appsi_ipopt: Performance is comparable to fast simulation.
• With standard ipopt: There is a known I/O overhead per point (~0.05s). The implemented warm-start strategy mitigates the mathematical solution time, making it viable for standard grid resolutions.

How to Test

You can test the new functionality using the updated shower.py:

from opyrability import nlp_based_approach, AIS2AOS_map
from tests.shower pyomo_omlt import build_shower2x2
import numpy as np

# 1. Test Inverse Mapping (Optimization)
# The function automatically detects the Pyomo builder
fDIS, fDOS, msg = nlp_based_approach(
    model=build_shower2x2,
    DOS_bounds=[[180, 200], [90, 100]],
    DOS_resolution=[5, 5],
    u0=np.array([1.0, 1.0]),
    lb=np.array([0, 0]),
    ub=np.array([10, 10])
)

# 2. Test Forward Mapping (Simulation)
# Note the required 'output_dim' argument
AIS_bounds = np.array([[0, 10], [0, 10]])
AIS, AOS = AIS2AOS_map(
    model=build_shower2x2,
    AIS_bound=AIS_bounds,
    AIS_resolution=[10, 10],
    output_dim=2 
)

[victoraalves]

Hi @hfsf!

Thank you for this great addition! I've the plans to add this myself (but couldn't find the time to), so I appreciate your efforts in making it happen.

I'm currently reviewing your PR and should be able to incorporate it to the tool, give you the proper and deserved credit/acknowledgement in the next release, and plan to have an example in the Examples gallery.

Speaking of examples, could you prepare a more comprehensive OMLT/IDAES example, inspired by the "look-and-feel" of the ones present in the Examples Gallery, so we could have an end-to-end case-study? This can be done in a second PR later in the future, after I incorporate this one, or you could modify this PR. Your call.

Many thanks again for contributing!

[victoraalves]

Hey @hfsf, thanks for this contribution! The Pyomo/OMLT integration is a great addition. I tested both AIS2AOS_map and nlp_based_approach with a Pyomo shower model and the results match the NumPy-based model to machine precision. The warmstart strategy and APPSI detection work well too.

That said, there are a few things that need to be addressed before we can merge:

Critical

1. Pyomo missing from dependencies: import pyomo.environ as pyo is added at line 31, but Pyomo is not listed in pyproject.toml or setup.py. Please add it as a dependency in both.

2. Missing example file: The PR description references from tests.shower_pyomo_omlt import build_shower2x2, but that file isn't included in the PR. Could you add it?

Medium

3. Duplicate imports (lines 5-8): from itertools import permutations as perms and from typing import Callable, Union each appear twice.

4. Missing space + typo in error message (line 1459-1460): Currently produces "For Pyomo/OML models in forward mapping, you SHOULD providethe 'output_dim' argument". Missing a space between "provide" and "the", and "OMLT" is misspelled as "OML".

5. Duplicate docstring line (lines 1448-1449): "Discretized Available Output Set (AOS)." appears twice in the AIS2AOS_map returns section.

Minor

6. Bare except: (line 1495): Should be except Exception: at minimum.

7. Solver verbosity: For nlp_based_approach, the IPOPT print_level = 0 is commented out (line 948), so users will get verbose solver output on every iteration. Consider adding a verbose parameter (or similar) so users can control this.

Test results

Existing tests pass (3/4; test_dma_mr_design has a pre-existing tolerance drift, unrelated to this PR). The new Pyomo paths produce correct results:

• Forward mapping: Pyomo vs NumPy max difference ~1e-14
• Inverse mapping: Pyomo norms match existing cyipopt test values to ~1e-7

Overall the implementation logic is solid. Just needs the dependency fix and the cleanup items above. Let me know if you have any questions!

[hfsf]

Hi @victoraalves,

Sorry for the delay in responding and thank you for the detailed fix plan. I will implement the suggested fix and update this PR soon.