Deriving envelopes & relaxations with SymPy

Deriving envelopes & relaxations with SymPy#

discopt can derive tight convex/concave envelopes for nonlinear atoms symbolically, verify their soundness, and code-generate pure-JAX relaxation closures. It is a design tool for studying an atom’s envelope before it is implemented in the solver: it extends the McCormick construction [McCormick, 1976, Tsoukalas and Mitsos, 2014] to terms arising in gas-network and chemical-engineering optimization.

The toolkit lives in discopt._relax.symbolic and is a design-time dependency (pip install 'discopt[sympy,jax]'): SymPy derives the formulas, and the generated closures run on jax.numpy.

Note

No solver path runs these closures. The default relaxation engine is numpy and relaxes each atom with its own envelopes (inventoried in docs/design/relaxation-catalog.md); the McCormick-NLP route that generated closures and the hand-written domain packs used to plug into was retired in #1873, along with those packs, their registry and the learned-envelope prototype. What remains is the derive -> code-generate -> verify loop shown here. A default solve imports no jax module; this notebook does, because the generated closures are JAX.

Soundness — cv(x) <= f(x) <= cc(x) with cv convex and cc concave — is the non-negotiable invariant, checked for every atom.

1. Deriving a univariate envelope#

For a function f(x) on a box [a, b], the engine classifies the curvature and derives the envelope. For an odd cubic x^3 (concave then convex through an inflection at 0) it finds the convex-underestimator tangent point in closed form.

import jax
jax.config.update('jax_enable_x64', True)
import sympy as sp
from discopt._relax.symbolic import derive_envelope, lambdify_envelope, verify_envelope

x = sp.Symbol('x', real=True)
res = derive_envelope(x**3, x)
print('curvature        :', res.curvature.value)
print('inflection point :', res.inflection)
print('cv tangent point :', res.cv_tangent.point)   # closed form in the lower bound a

curvature        : concavo_convex
inflection point : 0
cv tangent point : -a/2

lambdify_envelope turns the result into a JAX closure (x, lb, ub) -> (cv, cc), and verify_envelope certifies soundness over many randomized boxes.

relax = lambdify_envelope(res)
report = verify_envelope(relax, lambda t: t**3, domain=(-2.0, 2.0), n_boxes=400)
print('sound           :', report.sound)
print('max lower viol  :', f'{report.max_lower_violation:.2e}')
print('max upper viol  :', f'{report.max_upper_violation:.2e}')
print('mean gap (tight) :', f'{report.mean_gap:.3f}')

sound           : True
max lower viol  : 0.00e+00
max upper viol  : 0.00e+00
mean gap (tight) : 0.687

2. Transcendental atoms: numeric tangent solving#

When the tangent equation has no closed form (e.g. a sigmoid), code-gen solves the tangent point per box with a JAX bisection baked into the closure — still jit/vmap-compatible, still sound.

sig = 1 / (1 + sp.exp(-x))
res_sig = derive_envelope(sig, x)
print('curvature   :', res_sig.curvature.value)
print('closed form :', res_sig.cv_tangent.point)   # None -> numeric bisection
relax_sig = lambdify_envelope(res_sig)
rep = verify_envelope(relax_sig, lambda t: 1/(1+jax.numpy.exp(-t)), domain=(-5.0, 5.0))
print('sound       :', rep.sound)

curvature   : convexo_concave
closed form : None

sound       : True

3. Domain atoms, derived from their formulas#

The same three calls handle atoms from application domains. Steady-state gas flow couples squared nodal pressures to flow through the Weymouth term \(f\,|f| = \operatorname{sign}(f) f^2\), concave for \(f<0\) and convex for \(f>0\) — a single-inflection atom. The Arrhenius rate \(\exp(-E/RT)\) and Langmuir/Monod saturation \(x/(K+x)\) come from reaction engineering.

import jax.numpy as jnp

atoms = [
    ("Weymouth f|f|", x * sp.Abs(x), lambda t: t * jnp.abs(t), (-10.0, 10.0)),
    ("Arrhenius, E/R=6000", sp.exp(-6000 / x), lambda t: jnp.exp(-6000.0 / t), (250.0, 900.0)),
    ("saturation, K=1", x / (1 + x), lambda t: t / (1.0 + t), (0.0, 20.0)),
]
for name, expr, f, dom in atoms:
    r = derive_envelope(expr, x)
    rep = verify_envelope(lambdify_envelope(r), f, domain=dom, n_boxes=200)
    print(f"{name:22s} curvature={r.curvature.value:15s} sound={rep.sound}  mean_gap={rep.mean_gap:.3f}")

Weymouth f|f|          curvature=concavo_convex  sound=True  mean_gap=7.749
Arrhenius, E/R=6000    curvature=convex          sound=True  mean_gap=0.000

saturation, K=1        curvature=concave         sound=True  mean_gap=0.032

The derivation is honest about its limits: an atom with more than one inflection on the domain (for example \(\sin\) or \(\cos\) over a wide angle box, the trigonometric terms of AC power flow [Coffrin et al., 2016, Kröger et al., 2018]) is refused rather than given an envelope it cannot certify.

from discopt._relax.symbolic.envelope_deriver import EnvelopeDerivationError

try:
    derive_envelope(sp.cos(x), x)
except EnvelopeDerivationError as exc:
    print("refused:", exc)

refused: cos(x) has 2 inflection points; multi-inflection envelopes are handled by a later phase

Summary#

  • The SymPy engine derives convex/concave envelopes, code-generates them as JAX closures (codegen.py over the sympy-free runtime.py), and verify_envelope certifies soundness over randomized boxes.

  • Single-inflection atoms use a tangent/secant construction with a closed-form or numeric tangent point and an automatic secant fallback on asymmetric boxes; multi-inflection atoms are refused.

  • It is a design tool: no solver path runs the generated closures. The solver’s own envelopes are numpy, listed in docs/design/relaxation-catalog.md.

See design/symbolic-envelopes-plan.md for the original roadmap.