Why global optimization of the surrogate matters

The usual way to optimize a neural network or a smooth tree is gradient-based: start somewhere, follow the gradient, restart a few times. That works when the surrogate is nearly convex. When the surrogate has several basins, which is common for exactly the response surfaces we care about (selectivity versus temperature and pressure, cost versus a handful of design variables) you have little control over which one you find. The jax-difftree docs say it plainly: inverse design over a soft tree is local, and multistart is the practical fallback.

A MIP solver works differently. It bounds the objective over regions of the input space and discards regions that cannot contain the optimum. When it finishes, the optimality gap is zero (or a tolerance you chose), and that is a statement about the whole box, not about the neighborhood of a starting point. For smooth networks you can get the same guarantee from deterministic global MINLP solvers (Schweidtmann and Mitsos 2019), but for piecewise-linear models the problem is a MILP, and MILP solvers are fast, mature and free (HiGHS (Huangfu and Hall 2018) solves every MILP in this post).

There is a second reason that matters just as much in practice. A surrogate embedded as constraints composes with everything else. You can put two unit-operation surrogates in one flowsheet, add a budget constraint, choose between discrete equipment options, and still get a global optimum of the combined problem. A predict function cannot do that.

What “MIP-representable” means

The building block is the ReLU, . If you know bounds with , one binary variable encodes it exactly:

When the constraints force , and when they force . This is the big-M formulation, and it has been used for trained ReLU networks for some time (Fischetti and Jo 2018; Tjeng et al. 2019; Grimstad and Andersson 2019). Stronger formulations exist (Anderson et al. 2020), and Vielma (2015) is a good general reference. The same idea covers hinge terms , absolute values, linear splines (with SOS2 or segment binaries), and tree routing (one binary per leaf).

Two consequences follow from that formulation and show up in every section below. First, the bounds and come from bounds on the inputs, propagated through the model. Every input must have finite bounds, and tight bounds give small big-M constants, which gives fast solves. Loose bounds are the most common cause of a slow MILP. Second, the formulation is exact only for piecewise-linear pieces. A SiLU, tanh or softplus activation, a cubic spline, or a polynomial term is not piecewise linear, and you either retrain with a piecewise-linear choice or move to a nonlinear solver.

It helps to name the levels of representability, since not every model lands in the same place.

Uncertainty adds one more ingredient. For several pycse models the predicted standard deviation is a second-order cone in the network’s internal variables, so robust objectives like become mixed-integer second-order cone programs. We come back to this in the pycse section.

Getting a surrogate into a solver

There are three routes from a trained model to an optimization problem, and our packages use all of them.

Direct Pyomo embedding. The model writes its own variables and constraints into a Pyomo Block (Bynum et al. 2021). You pass in your own Pyomo variables as inputs, and the block links them to the surrogate. This is the most flexible route and the one every package here supports (to_pyomo in pycse, jaxsr and jax-difftree, embed_in_pyomo_model in jax-ldt). Any Pyomo MILP solver works: HiGHS, Gurobi, CPLEX, SCIP.

OMLT. OMLT (Ceccon et al. 2022) reads a neural network (usually from ONNX) and builds a Pyomo block with a choice of formulations, such as ReluBigMFormulation. OMLT also has tree formulations, including linear model decision trees (Ammari et al. 2023). In our packages, pycse’s LLPRRegressor exports to ONNX and directly to an OMLT NetworkDefinition (to_onnx(), to_omlt()). One caveat: jax-ldt also writes ONNX, but OMLT cannot read it. OMLT’s ONNX reader understands neural-network layers (Gemm, MatMul, ReLU), and the tree-routing operations (Gather, LessOrEqual) are rejected. That is why the tree packages ship their own adapters.

discopt. discopt is our hybrid MINLP solver with a JAX backend. jax-ldt (embed_in_discopt_model) and jaxsr (to_discopt) write directly into a discopt.Model, and jax-difftree reaches discopt through the jax-ldt spec (more on that below). The advantage of discopt is that the same model can hold MILP surrogate blocks and nonlinear constraints, which is where jaxsr’s allow_nonlinear=True and pycse’s SISSO export are headed. pycse itself writes Pyomo, but the ONNX file from LLPRRegressor.to_onnx() can also be read by discopt.ml.load_onnx.

Whatever the route, the pattern is the same: declare bounded input variables, embed the surrogate, add your constraints, set an objective on the surrogate output, solve, and then check the answer against the model’s own predict.

pycse.sklearn models

pycse has the broadest coverage. A new module in pycse 2.13.0, pycse.sklearn.mip, holds the formulation machinery once, and each estimator only describes its trained weights. A model implements mip_network(), which returns its weights as plain numpy layer specs (Affine, ReLU, PiecewiseLinear), a TreeModel of leaf boxes, or a symbolic expression. A single function, add_to_pyomo, turns any of these into constraints, so the bound propagation, the big-M formulation and the tests are shared across models. For the user it is one call:

import pyomo.environ as pyo
from pycse.sklearn import DPOSE

model = DPOSE(layers=(2, 32, 16)).fit(X, y)       # ReLU by default

m = pyo.ConcreteModel()
m.T = pyo.Var(bounds=(300, 500))
m.P = pyo.Var(bounds=(1, 10))
m.budget = pyo.Constraint(expr=m.T / 100 + m.P <= 12)   # your own constraints
m.nn = pyo.Block()
model.to_pyomo(block=m.nn, inputs=[m.T, m.P])        # input bounds come from your vars
m.obj = pyo.Objective(expr=m.nn.y[0], sense=pyo.maximize)
pyo.SolverFactory("highs").solve(m)

Internal input and output scaling is folded into the exported weights, so m.T, m.P and m.nn.y are in your units.

jax-ldt: linear and hyperplane decision trees

jax-ldt fits linear-model decision trees, where every leaf holds a linear regression instead of a constant, and hyperplane trees, where the splits are oblique hyperplanes chosen from sparse combinations of the inputs. Both are piecewise linear by construction, so they were designed from the start to be embedded. The package was motivated by our work with the Laird group on hyperplane trees as surrogates for chemical process design (Sunshine et al. 2025), which builds on linear model decision trees in OMLT (Ammari et al. 2023). Optimization over tree models more generally goes back to Mišić (2020) and ENTMOOT (Thebelt et al. 2021).

The MILP formulation lifts the features into the transformed (hyperplane) space, uses one binary per leaf, writes the routing as big-M constraints, and linearizes the leaf outputs with big-M as well. The result is an exact MILP. The same formulation is available in Pyomo and in discopt:

import discopt
from jax_ldt import HyperplaneTreeRegressor
from jax_ldt.export import embed_in_discopt_model

model = HyperplaneTreeRegressor(max_depth=4, max_bins=6, min_samples_leaf=15)
model.fit(X, y)

m = discopt.Model("min-y")
x = m.continuous("x", shape=(2,), lb=-2, ub=2)
y_expr = embed_in_discopt_model(model, m, x)
m.minimize(y_expr[0])
result = m.solve()

For Pyomo, embed_in_pyomo_model(model, m, m.x) builds the same formulation on a sub-block, so several trees can share one model. For axis-aligned trees with constant leaves there is also a lighter adapter, to_discopt_decision_tree, which uses discopt’s built-in tree formulation.

The canonical export is a JSON spec (to_json / from_json), and every other format is regenerated from it. That turns out to be useful beyond jax-ldt, because jax-difftree writes the same spec.

Caveats. Big-M constants are derived per node and per leaf from the input bounds, so bounds are required. You can override them with big_m=, and they are validated when the inputs are bounded. The ONNX export is for inference with onnxruntime, not for OMLT, as noted above. The tree’s uncertainty estimates are stored in the spec but are not part of the MILP.

Docs: MIP export, discopt integration. Examples: 04_mip_optimization.ipynb, 07_pyomo_embedding.ipynb.

jax-difftree: differentiable soft trees

jax-difftree takes a different route to a tree. Instead of growing it greedily, it fits a fixed-depth soft tree end to end with gradients (Hehn et al. 2020), with oblique splits and sigmoid routing. Soft trees are smooth and train well, and they can match derivative data (Jacobians from a simulator) as well as values. They are not piecewise linear, though, so getting a MILP takes one extra step.

The workflow is to train with affine leaves, harden the tree so that each point routes to exactly one leaf, and then export the crisp piecewise-affine model:

from difftree import DifferentiableTreeRegressor, export
import pyomo.environ as pyo

model = DifferentiableTreeRegressor(
    depth=3, leaf_type="affine", jacobian_weight=1.0, hardening=True
)
model.fit(X_tr, Y_tr, dydx=J_tr)
print(model.hardening_report_)            # soft vs hard error, before and after

m = pyo.ConcreteModel()
m.x = pyo.Var(range(2))
m.y = pyo.Var(range(2))
tree = export.to_pyomo(model, m, x=m.x, y=m.y, bounds=(lo, hi), method="bigm")
m.obj = pyo.Objective(expr=m.y[0])
pyo.SolverFactory("appsi_highs").solve(m)

The Pyomo export is written as a generalized disjunctive program: each leaf is a Disjunct holding its polytope and its affine model, and a Disjunction makes exactly one leaf active, which Pyomo then transforms (here with big-M). For discopt, export.to_json writes the jax-ldt v1 spec, so jax-ldt’s from_json and embed_in_discopt_model consume a difftree model unchanged. All outputs share one partition, which means a multi-output tree for a unit operation needs only one set of binaries.

Caveats. The soft/hard gap is the main one. The MILP optimizes the hardened tree, not the soft one you trained, so check soft_hard_gap (and train with hardening=True) before trusting the optimum. The leaf polytopes are closed, as usual in big-M formulations, so an optimum can sit on a split boundary using either neighbor’s model; export.repair_solution moves such a point into the chosen leaf. Continuity across oblique splits needs tied_splits=True, which gives up flexibility (a depth- tree then has hyperplanes instead of ). The fixed full-binary structure caps practical depth at around 10 to 12. A pruned export matches predict(hard=True) only inside the export box.

Docs: surrogates for MILP, limitations. Notebooks: piecewise_linear_milp, pyomo_embedding, multivariate_inverse_design.

jaxsr: symbolic regression with hinge functions

jaxsr fits sparse symbolic models, a weighted sum of terms chosen from a basis library by an information criterion, in the spirit of ALAMO (Cozad et al. 2014). The new capability is a family of piecewise-linear basis terms, most usefully hinges and , which give MARS-style fits (Friedman 1991; Breiman 1993). If every selected term is piecewise linear, the fitted model is an exact MILP, and unlike the other models in this post it is also a short, readable equation.

from jaxsr import BasisLibrary, SymbolicRegressor, milp_incompatible_terms

library = (
    BasisLibrary(n_features=2, feature_names=["T", "P"], feature_bounds=[(0, 4), (0, 4)])
    .add_constant()
    .add_linear()
    .add_hinges(X=X, n_knots=7)
    .add_transcendental(["abs"])
)
model = SymbolicRegressor(basis_library=library, max_terms=6, information_criterion="bic")
model.fit(X, cost)
print(model.expression_)
print(milp_incompatible_terms(model))   # [] means a pure MILP

The MILP-representable terms are constants and linear terms (written directly), hinges and abs (one binary each, big-M), categorical indicators (one-hot binaries), and categorical-continuous interactions (an exact McCormick product). The model then goes to either backend with one call, linked to your own variables:

model.to_pyomo(m.sr, inputs={"T": m.T, "P": m.P}, output=m.cost)   # Pyomo block

emb = model.to_discopt(dm_model)                                    # discopt
dm_model.subject_to(emb.inputs["T"] + emb.inputs["P"] >= 5)
dm_model.minimize(emb.output)

Polynomials, interactions, log, exp and the rest are not piecewise linear. You can still embed them with allow_nonlinear=True, which gives an MIQCP or MINLP for discopt, BARON, SCIP or Gurobi.

Caveats. Bounds are required and are enforced on the embedded inputs, on purpose: outside the training range the model is extrapolating, and the right fix is more data, not wider bounds. Knot placement matters more than knot count, since the fit can only bend where a knot exists, so pass known breakpoints (a phase boundary, a saturation point) explicitly with knots=. Parametric and custom basis functions, rational forms and SISSO-style features cannot be embedded. MultiOutputSymbolicRegressor is not supported directly, but each estimators_[i] can be embedded separately. The embedding is built from the fitted coefficients, so rebuild it after a refit.

Guide: Embedding models in MILPs. Example: mip_surrogate_optimization.ipynb, which includes a refine-and-resolve loop and a design problem solved in both backends.

Choosing a model

The packages overlap, but they are good at different things. If you want an interpretable equation, or you know where the response bends, jaxsr’s hinge models are the natural choice, and they stay small. If you want a flexible fit with uncertainty and the option of a robust or exploratory objective, the pycse ReLU networks are the most complete, and the input-convex networks are the cheapest to optimize when convexity is physically justified. Trees are attractive when the response has genuinely different regimes: jax-ldt grows them greedily and exports directly, while jax-difftree trains them end to end, can use derivative data, and needs the hardening step before export.

	pycse ReLU nets	pycse KAN	jax-ldt	jax-difftree	jaxsr hinges
Binaries scale with	hidden units	grid × inputs	leaves	leaves ()	hinge/abs terms
Interpretable	no	partly	partly	partly	yes
Std in the MIP	yes (SOCP)	with LLPR/ensemble	no	no	no
Pyomo	yes	yes	yes	yes (GDP)	yes
discopt	LLPR via ONNX	no	yes	via jax-ldt spec	yes
OMLT	LLPR via ONNX	no	no	no	no

The caveat that applies to all of them

A global optimizer finds your model’s mistakes. Our first run of the pycse benchmark used 150 samples and shorter training. Two of the models came out poorly fitted, and their global optima sat on the edge of the box, where there was little data and the models were extrapolating: one predicted at a point where the true is . A local optimizer started in the middle might never go there, and the MIP always will. The global optimum of the surrogate is only as good as the surrogate is in that region. So validate the optimum against the true model, and then either add data there and refit (jaxsr’s notebook has a refine-and-resolve loop, and our earlier active-learning post covers the data side), penalize the std, or tighten the bounds to the region you trust.

Installing

The MIP export is in pycse 2.13.0 and later. Pyomo and HiGHS come with the mip extra; the ONNX and OMLT exports of LLPRRegressor also need onnx and omlt.

pip install "pycse[mip]"

jaxsr 0.3.0 is on PyPI, but the hinge terms and MILP embedding are newer than that release, so install from GitHub for now. jax-ldt and jax-difftree are not on PyPI. (The difftree package on PyPI is an unrelated project, so install jax-difftree from GitHub as shown.)

pip install "jaxsr[mip] @ git+https://github.com/jkitchin/jaxsr.git"
pip install "jax-ldt[pyomo] @ git+https://github.com/jkitchin/jax-ldt.git"
pip install "difftree[pyomo] @ git+https://github.com/jkitchin/jax-difftree.git"
pip install discopt omlt onnx

References

One layer: combine, then activate

A single neuron in the first layer does two things. It forms a weighted sum of the inputs plus a bias, and then it applies an activation function :

Let’s draw that. In the figures below the colored squares are input features and the grey squares are weights. A grey square next to a colored square means “multiply,” and the plus signs mean “add.”

One neuron. Every colored square is multiplied by a grey square, the products are added to a bias, and the sum is bent by σ.

A layer is many neurons that look at the same inputs with different weights. In matrix form , and is applied element by element. Each is a mixture of the inputs, and we draw it as a square striped with the colors of the inputs in proportion to how strongly it weights each one.

A layer with three inputs and four neurons. Hover over (or tap) a neuron to see the row of weights it uses. Each pre-activation zⱼ is a different blend of the same three inputs.

That is everything a layer can do. Notice what is not in the picture. There is no square where meets , so nothing like , or is ever formed by the linear part. The pre-activations are exactly linear in the inputs, and so is any linear combination of them. We can check that with a random layer.

import numpy as np

rng = np.random.default_rng(0)
X = rng.uniform(-1, 1, size=(500, 3))  # 500 samples, 3 inputs
W = rng.normal(size=(4, 3))            # 4 neurons
b = rng.normal(size=4)

Z = X @ W.T + b                        # pre-activations
H = np.tanh(Z)                         # features

def r2(A, y):
    "R² of a least squares fit of y on A plus a constant."
    A = np.column_stack([A, np.ones(len(y))])
    c = np.linalg.lstsq(A, y, rcond=None)[0]
    return 1 - (y - A @ c).var() / y.var()

x1, x2 = X[:, 0], X[:, 1]
tests = [("each z_j from x1, x2, x3", min(r2(X, z) for z in Z.T)),
         ("x1*x2 from the z's", r2(Z, x1 * x2)),
         ("x1/(x2+2) from the z's", r2(Z, x1 / (x2 + 2))),
         ("x1*x2 from the h's", r2(H, x1 * x2))]
for name, val in tests:
    print(f"{name:24s} R² = {val:.6f}")

each z_j from x1, x2, x3 R² = 1.000000
x1*x2 from the z's       R² = 0.004774
x1/(x2+2) from the z's   R² = 0.912157
x1*x2 from the h's       R² = 0.142195

The pre-activations are perfectly linear in . No combination of them explains any of the variance of . (The ratio scores well only because is nearly linear on this range.) After the activation, the features explain a little of the product, and that little bit is the subject of the next section.

There is a geometric way to say the same thing. A neuron’s output depends on only through the single number . It is constant along every line (in 3D, every plane) perpendicular to . Functions like that are called ridge functions: they look like a wave, a step or a ramp that is perfectly straight in one direction. A product and a ratio are not ridge functions. Their contours are hyperbolas and rays.

Left: one neuron σ(w₁x₁ + w₂x₂ + b). Its contours are always parallel straight lines, whatever the weights, bias and activation. Middle and right: the contours of x₁x₂ are hyperbolas and the contours of x₁/x₂ are rays from the origin. No single neuron can look like either one.

Where the nonlinearity comes from

If the linear part never multiplies inputs, a network must get products from the activation function. We can make that precise. For a network with one hidden layer,

take two derivatives. The chain rule gives

The mixed partial derivative is a good measure of how much “product” there is in a function. It is exactly 1 everywhere for , and exactly 0 for anything of the form , and in particular for anything linear. The formula says that every bit of it in a neural network comes from , the curvature of the activation function. The weights only decide where that curvature is applied and how it is shared between the inputs.

There are two common ways for to be nonzero, and they lead to quite different networks.

def g(s, eps, b=0.66):
    return np.tanh(b + eps * s)

def tanh2(b):
    "Second derivative of tanh."
    t = np.tanh(b)
    return -2 * t * (1 - t**2)

grid = np.linspace(-1, 1, 101)
x1, x2 = np.meshgrid(grid, grid)
u, v = x1 + x2, x1 - x2
for eps in [1.0, 0.3, 0.1, 0.03]:
    c = 1 / (4 * eps**2 * tanh2(0.66))
    approx = c * (g(u, eps) + g(-u, eps)
                  - g(v, eps) - g(-v, eps))
    err = np.abs(approx - x1 * x2).max()
    print(f"eps = {eps:4.2f}  |output weight| = {abs(c):6.1f}"
          f"  max error = {err:.1e}")

eps = 1.00  |output weight| =    0.3  max error = 6.6e-01
eps = 0.30  |output weight| =    3.6  max error = 1.1e-01
eps = 0.10  |output weight| =   32.5  max error = 1.3e-02
eps = 0.03  |output weight| =  360.8  max error = 1.2e-03

The sin of adding kelvins to pascals

Let’s give the inputs physical meaning. Suppose a neuron sees a temperature in kelvin, a pressure in pascal and a concentration in mol/L:

In physics you cannot add quantities with different units. “300 K + 101325 Pa” is not a number; it is an error. A neural network adds them anyway, because all it sees are the numbers 300 and 101325. The only way the sum makes sense is if each weight carries units that convert its term into a common unit:

What is ? It is the argument of . A transcendental function like needs a dimensionless argument, which you can see from its series : you cannot add and unless is a pure number. So is dimensionless, the bias is dimensionless, and every weight has the inverse units of its input:

Each weight is a reciprocal characteristic scale. A neuron is implicitly asking “how many is this temperature, how many is this pressure,” and adding the answers. Learning the weights means learning those scales, one set per neuron.

The neuron from before with units. Each grey square must carry the inverse units of the colored square it multiplies, so that every product is dimensionless and the sum is legal.

The same bookkeeping goes through the whole network. The features are dimensionless, so the output weights carry the units of . In the mixed partial derivative , the units are , exactly the units of . The network is dimensionally consistent, but only because the weights silently carry units. ReLU is slightly more forgiving because for , so its argument may have units, but the terms of the sum still have to agree, so the weights still convert between input units.

The software does not know any of this, and three things go wrong when we forget it.

Initialization. Weights are usually initialized as random numbers of order one. With raw inputs, that means K and Pa. A pressure of a few bar is then hundreds of thousands of “characteristic pressures,” and every neuron is saturated before training starts.

n = 500
T = rng.uniform(300, 500, n)    # K
P = rng.uniform(1e5, 1e6, n)    # Pa
X = np.column_stack([T, P])

W = rng.normal(size=(2, 100))   # 100 neurons, weights ~ 1
b = rng.normal(size=100)

Xs = (X - X.mean(axis=0)) / X.std(axis=0)
for name, A in [("raw (K, Pa)", X), ("standardized", Xs)]:
    sat = np.mean(np.abs(A @ W + b) > 3)
    print(f"{name:13s} saturated (|z| > 3): {sat:.2f}")

raw (K, Pa)   saturated (|z| > 3): 1.00
standardized  saturated (|z| > 3): 0.09

A saturated tanh unit has and : it contributes no gradient and none of the curvature we needed for products.

Regularization. Weight decay adds to the loss. With mixed inputs that is a sum of terms with units of and , which is another sin. In practice it means the answer depends on your choice of units. Here is ridge regression (the linear model is enough to show it) fitted to the same data with pressure in Pa, bar and MPa, and with temperature in K and °C.

y = 2 + 0.01 * T + 3e-6 * P + rng.normal(0, 0.1, n)

def ridge(X, y, lam=100.0):
    A = np.column_stack([X, np.ones(len(y))])
    I = np.eye(A.shape[1])
    I[-1, -1] = 0                  # do not penalize the bias
    theta = np.linalg.solve(A.T @ A + lam * I, A.T @ y)
    return A @ theta

for name, Tu, Pu in [("K, Pa", T, P),
                     ("K, bar", T, P / 1e5),
                     ("°C, bar", T - 273.15, P / 1e5),
                     ("K, MPa", T, P / 1e6)]:
    yp = ridge(np.column_stack([Tu, Pu]), y)
    print(f"{name:8s} first predictions:", np.round(yp[:3], 3))

K, Pa    first predictions: [9.133 8.565 6.603]
K, bar   first predictions: [9.102 8.54  6.63 ]
°C, bar  first predictions: [9.102 8.54  6.63 ]
K, MPa   first predictions: [8.301 7.925 7.324]

The same data and the same model give different predictions when we change the pressure units. Changing kelvin to Celsius does nothing, because a shift of origin is absorbed by the bias, which is not penalized. A change of scale is not absorbed by anything.

Standardization is a units fix. Replacing each input by divides it by a scale with the same units, so the network sees dimensionless numbers and the weights become dimensionless too. That is why standardizing inputs helps so much. It is worth being clear about what it does, though: it chooses the characteristic scales from the spread of the training data, not from the physics. A dimensionless group chosen from the physics, such as , or a Reynolds number, does the same job with scales that mean something and that transfer to new conditions.

There is one more consequence. Physical laws are full of products and ratios (, , a rate ), and products and ratios are dimensionally fine: K·Pa is a perfectly good unit. Sums are what require matching units. So the network’s single native operation, the sum, is the one operation physics constrains most, and the operations physics uses most freely are the ones the network has to approximate.

Composition: features of features

A deep network stacks layers. The second layer does exactly what the first one did, with the first layer’s features as its inputs:

In the square picture, each colored input square of the second layer is a whole first-layer neuron. The structure repeats: grey times feature, add, bend.

A second-layer neuron (top). Its colored squares are first-layer features, and each of those is itself a bent sum of grey-times-color terms (below). However deep the nesting goes, every colored square sits next to a grey square and never next to another colored square. The product and ratio at the bottom are the patterns that never appear.

Written out, a second-layer pre-activation is

There is still no anywhere in it. Depth does not add a new operation; it applies the same two operations again. What depth changes is how the bends combine. Take two derivatives of a two-layer network:

Every term still contains a , either from the second layer acting on its (now curved) inputs, or passed through from the first layer. Depth gives the network more ways to place curvature, but curvature is still the only source of interaction between inputs. Let’s check the formula against finite differences for a random network, and also try the same network with ReLU.

W1, b1 = rng.normal(size=(5, 2)), rng.normal(size=5)  # 2 -> 5
W2, b2 = rng.normal(size=(4, 5)), rng.normal(size=4)  # 5 -> 4
c = rng.normal(size=4)

relu = lambda z: np.maximum(0, z)

def f(x, act=np.tanh):
    return c @ act(W2 @ act(W1 @ x + b1) + b2)

def mixed_fd(x, act=np.tanh, h=1e-4):
    "∂²f/∂x1∂x2 by central differences."
    e1, e2 = np.array([h, 0]), np.array([0, h])
    return (f(x + e1 + e2, act) - f(x + e1 - e2, act)
            - f(x - e1 + e2, act) + f(x - e1 - e2, act)
            ) / (4 * h**2)

def mixed_formula(x):
    "∂²f/∂x1∂x2 for tanh from the chain rule above."
    a1 = np.tanh(W1 @ x + b1)
    s1, s1pp = 1 - a1**2, -2 * a1 * (1 - a1**2)   # σ', σ''
    a2 = np.tanh(W2 @ a1 + b2)
    s2, s2pp = 1 - a2**2, -2 * a2 * (1 - a2**2)
    dz2 = W2 @ (s1[:, None] * W1)              # ∂z2/∂x
    d2z2 = W2 @ (s1pp * W1[:, 0] * W1[:, 1])   # ∂²z2/∂x1∂x2
    return c @ (s2pp * dz2[:, 0] * dz2[:, 1] + s2 * d2z2)

print("x               tanh formula  tanh FD    relu FD")
for x in rng.uniform(-1, 1, size=(3, 2)):
    print(f"{np.round(x, 2)!s:14s}  {mixed_formula(x):+.6f}  "
          f"{mixed_fd(x):+.6f}  {mixed_fd(x, relu):+.6f}")

x               tanh formula  tanh FD    relu FD
[-0.51  0.89]   -0.072855  -0.072855  +0.000000
[ 0.82 -0.74]   +0.057736  +0.057736  +0.000000
[-0.9   0.42]   +0.276423  +0.276423  +0.000000

For tanh the formula and the finite differences agree. For ReLU the mixed partial is zero, at these points and almost everywhere else, because between the kinks. A ReLU network of any depth is still piecewise linear. What depth buys it is the number of pieces: each layer can fold the regions created by the layer before, so the number of linear regions can grow much faster with depth than with width. A ReLU network that approximates does so by tiling the plane with many small flat facets that bend along creases.

Feature engineering, revisited

Put the pieces together. A neuron can only add scaled inputs and bend the sum. Products and ratios have to be approximated with kinks or curvature, using many neurons, large canceling weights, or both, and only over the range where the network was trained. Units are carried silently by the weights, and the software does not check them.

So what does “the network learns the features” mean? It learns features in the vocabulary of sums and bends. If the relationship you are modeling is naturally written in that vocabulary, the network will find it efficiently. If it is naturally a product of powers, which describes a lot of physical science, you can change the vocabulary for it. Take logarithms of positive inputs (after dividing by reference values, so the logarithms are dimensionless):

A product of powers is a linear combination of logarithms, which is the one thing the network does exactly. Here is a small test: learn from 300 samples with , once with a tanh network on the (standardized) raw inputs and once with plain linear regression on . Then test both outside the training range, on .

from scipy.optimize import minimize

def sample(lo, hi, n):
    X = rng.uniform(lo, hi, size=(n, 3))
    return X, X[:, 0] * X[:, 1] / X[:, 2]

Xtr, ytr = sample(1, 3, 300)
Xte, yte = sample(3, 5, 200)
mu, sd = Xtr.mean(0), Xtr.std(0)
ym, ys = ytr.mean(), ytr.std()

m = 10                                  # hidden neurons
def net(p, X):
    W, b = p[:3 * m].reshape(m, 3), p[3 * m:4 * m]
    c, c0 = p[4 * m:5 * m], p[5 * m]
    H = np.tanh((X - mu) / sd @ W.T + b)
    return (H @ c + c0) * ys + ym

loss = lambda p: np.mean(((net(p, Xtr) - ytr) / ys)**2)
p0 = 0.5 * rng.normal(size=5 * m + 1)
p = minimize(loss, p0, method="BFGS",
             options={"maxiter": 3000}).x

A = lambda X: np.column_stack([np.log(X), np.ones(len(X))])
theta = np.linalg.lstsq(A(Xtr), np.log(ytr), rcond=None)[0]
loglin = lambda X: np.exp(A(X) @ theta)

rmse = lambda a, b: np.sqrt(np.mean((a - b)**2))
print("model           params  train RMSE  test RMSE")
models = [("tanh network", 5 * m + 1, lambda X: net(p, X)),
          ("log-linear", 4, loglin)]
for name, k, model in models:
    print(f"{name:14s}  {k:6d}  {rmse(model(Xtr), ytr):10.1e}"
          f"  {rmse(model(Xte), yte):9.1e}")
print("exponents:", np.round(theta[:3], 6))

model           params  train RMSE  test RMSE
tanh network        51     3.9e-03    1.6e-01
log-linear           4     5.7e-16    2.6e-15
exponents: [ 1.  1. -1.]

Both models fit the training data well. The network, with 51 parameters, built an approximation to the product out of curvature, and it degrades as soon as it leaves the region where that approximation was built. The log-linear model has four parameters, recovers the exponents exactly, and extrapolates perfectly, because in log space the true function is a linear combination of the features. Of course this example was chosen to make the point, and a real problem will not be an exact power law. The lesson is not “always use logarithms.” It is that feature engineering and neural networks are not alternatives. The network builds features from sums and bends, and a good input representation (log transforms for multiplicative effects, dimensionless groups for units, physically meaningful scales) puts the structure you already know into the one operation the network does exactly. Then the network can spend its neurons and its curvature on the parts you do not know.

Part 1: a basis set expansion

Let’s start with something familiar: the Fourier series. On the interval we can write a function as

This is a basis set expansion: is approximated by a weighted sum of fixed functions . The only unknowns are the coefficients and .

import numpy as np
from scipy.integrate import quad

def f(x):
    return np.exp(-2 * (x - 1)**2) + 0.6 * np.exp(-4 * (x + 1.5)**2)

def a(n):
    return quad(lambda x: f(x) * np.cos(n * x), -np.pi, np.pi)[0] / np.pi

def b(n):
    return quad(lambda x: f(x) * np.sin(n * x), -np.pi, np.pi)[0] / np.pi

def fourier(x, N):
    y = a(0) / 2 * np.ones_like(x)
    for n in range(1, N + 1):
        y += a(n) * np.cos(n * x) + b(n) * np.sin(n * x)
    return y

xg = np.linspace(-np.pi, np.pi, 2001)[:-1]   # one period, evenly spaced
yg = f(xg)

def fourier_basis(x, N):
    cols = [np.ones_like(x)]
    for n in range(1, N + 1):
        cols += [np.cos(n * x), np.sin(n * x)]
    return np.column_stack(cols)

def poly_basis(x, N):
    return np.column_stack([x**k for k in range(N + 1)])

for name, basis, Ns in [("Fourier", fourier_basis, (3, 6)),
                        ("polynomial", poly_basis, (3, 6))]:
    print(name)
    for N in Ns:
        c = np.linalg.lstsq(basis(xg, N), yg, rcond=None)[0]
        print(f"  N = {N}: first four coefficients", np.round(c[:4], 4))

Fourier
  N = 3: first four coefficients [ 0.2841  0.2015  0.1376 -0.2312]
  N = 6: first four coefficients [ 0.2841  0.2015  0.1376 -0.2312]
polynomial
  N = 3: first four coefficients [ 0.4682  0.1336 -0.056  -0.0188]
  N = 6: first four coefficients [ 0.2641  0.3645  0.3232 -0.1281]

Part 2: a feature generator for linear regression

Here is another way to read the same picture. A neural network is a nonlinear transformation of the input into a new, usually higher dimensional, vector space, followed by linear regression in that space. The hidden layers generate features. The last layer is plain old linear regression on those features.

We have actually already done both of these in Part 1. Let’s look at the Fourier network again, one step at a time, for a single input .

Step 1: expand the dimension. The first layer maps the scalar to a vector with one element per hidden unit. For that is a vector in :

This step is linear. It does not add any information; it just makes copies of that are stretched and shifted by different amounts. (The constant unit is .)

Step 2: transform nonlinearly. Apply the activation function to each element:

Now the seven elements are genuinely different functions of . These are the features.

Step 3: linear regression. The output is linear in the features,

If we have data points, we stack the feature vectors as rows of a matrix (), and finding is the linear least squares problem , with solution . That is the problem np.linalg.lstsq solved with fourier_basis in Part 1. It uses an SVD rather than forming that inverse, but the answer is the same. In the Fourier case happens to be diagonal, which is why the integrals were enough.

Here are the three steps written as a network layer, with matrices instead of loops.

w = np.array([0, 1, 1, 2, 2, 3, 3])
beta = np.array([0, 0, -np.pi / 2, 0, -np.pi / 2, 0, -np.pi / 2])

Z = xg[:, None] * w + beta    # step 1: (n, 1) -> (n, 7)
H = np.cos(Z)                 # step 2: nonlinear activation
c = np.linalg.lstsq(H, yg, rcond=None)[0]   # step 3: linear regression

print(H.shape)
print(np.round(c, 4))

(2000, 7)
[ 0.2841  0.2015  0.1376 -0.2312  0.2014 -0.1485  0.1126]

These are the same coefficients as before.

from scipy.optimize import minimize

sizes = [1, 6, 6]   # input, hidden layer 1, hidden layer 2
shapes = []
for m, k in zip(sizes[:-1], sizes[1:]):
    shapes += [(m, k), (k,)]      # W, b for each hidden layer
shapes += [(sizes[-1],), ()]      # c, c0 for the output layer

def unpack(p):
    out, i = [], 0
    for s in shapes:
        k = int(np.prod(s))
        out.append(p[i:i + k].reshape(s))
        i += k
    return out

def features(p, X):
    "The hidden layers: H = tanh(... tanh(X W1 + b1) ... W + b)."
    *hidden, c, c0 = unpack(p)
    H = X
    for W, b in zip(hidden[::2], hidden[1::2]):
        H = np.tanh(H @ W + b)
    return H

def nn(p, X):
    *_, c, c0 = unpack(p)
    return features(p, X) @ c + c0

xt = np.linspace(-np.pi, np.pi, 200)
Xt, yt = xt[:, None], f(xt)

def loss(p):
    return np.mean((nn(p, Xt) - yt)**2)

p0 = np.random.default_rng(42).normal(size=sum(int(np.prod(s)) for s in shapes))
sol = minimize(loss, p0, method="BFGS", options={"maxiter": 5000})
print(f"{len(p0)} parameters, MSE = {sol.fun:.2e}")

61 parameters, MSE = 4.81e-06

H2 = features(sol.x, Xt)
A = np.column_stack([H2, np.ones(len(xt))])
c_lstsq = np.linalg.lstsq(A, yt, rcond=None)[0]

*_, c_nn, c0_nn = unpack(sol.x)
print("network:", np.round(np.append(c_nn, c0_nn), 3))
print("lstsq:  ", np.round(c_lstsq, 3))
print(f"cond(A) = {np.linalg.cond(A):.0f}")

network: [-1.044  0.962 -1.522  1.387 -0.859  0.774 -0.994]
lstsq:   [-1.044  0.962 -1.523  1.387 -0.859  0.776 -0.994]
cond(A) = 159

Changes that may affect your code

pycse now requires Python 3.12 or newer, and the JAX stack moves up to jax>=0.10, flax>=0.12.7 and optax>=0.2.5.

Prediction intervals from predict() and polyval() now have shape (n, 2) instead of (2, n), the same as nlpredict() has always returned. Column 0 is the lower bound and column 1 the upper bound. If you unpack the intervals, the fix is one transpose:

y, yint, se = predict(X, y, pars, XX)
lower, upper = yint.T    # was: lower, upper = yint

The optimizer="bfgs" and "lbfgs" options in the JAX models (DPOSE, KAN, KANLLPR, KfoldNN and SISSOEnsemble) now run a real L-BFGS with a line search from optax. Before, they ran Adam with a cosine schedule. Your fits will be different, and they should actually be converged now. If you need the old behavior, use optimizer="adam_cosine".

JAXMonotonicRegressor is now only monotone by default. It used to be convex as well, which meant concave or S-shaped monotone data could not be fitted. Pass convexity="convex" to get the old behavior, or convexity="concave" for the other case.

What is new

The largest addition is pycse.sklearn.zenn, an independent JAX/Flax implementation of Zentropy-Enhanced Neural Networks from Wang, Shang, Liu and Hao, “ZENN: A thermodynamics-inspired computational framework for heterogeneous data-driven modeling,” PNAS 123(1), e2511227122 (2026). It provides ZENNRegressor, ZENNRegressorNLL and ZENNClassifier. Like the other models in pycse.sklearn, they follow the scikit-learn estimator API.

LeafModelRegressor(refine_splits=True) chooses each split threshold for the leaf model instead of for constant leaves, so the breakpoints land in the right place for piecewise-linear data. The active learning tools picked up a penalty_scale option for penalized batches, and ActiveSurrogate.build() and LatinSquare take a random_state so results are reproducible.

The MCP server works with both the 1.x and 2.x versions of mcp, and the Claude Code skill installed by pycse skill install was rewritten. It now covers the regression functions, all of the sklearn models, design of experiments and active learning.

Fixes

pycse.__version__ finally reports the installed version. nlpredict() works for a single prediction point, NNGMM works on Python 3.12, and NNBR reports the correct standard deviation of the intercept. In active learning, penalized batches no longer contain duplicate points and ThompsonSampling makes a new draw on every call. The hashcache backends return NumPy arrays and tuples with their original types, and pycse.sklearn now imports its models lazily, so import pycse.sklearn no longer loads every backend.

The full list is in the changelog. As always, issues and pull requests are welcome on GitHub.

A small example

Here is the first-order reaction in a batch reactor, , solved numerically and compared with the analytical solution.

import numpy as np
from scipy.integrate import solve_ivp

k = 0.23   # 1/min
Ca0 = 2.0  # mol/L

sol = solve_ivp(lambda t, Ca: -k * Ca, (0, 20), [Ca0], dense_output=True)
t = np.linspace(0, 20, 200)
Ca = sol.sol(t)[0]

print(f"C_A(20 min) = {Ca[-1]:.4f} mol/L (exact {Ca0 * np.exp(-k * 20):.4f})")

C_A(20 min) = 0.0201 mol/L (exact 0.0201)

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(6, 3.5))
ax.plot(t, Ca, label="solve_ivp")
ax.plot(t[::15], Ca0 * np.exp(-k * t[::15]), "o", label="exact")
ax.set_xlabel("t (min)")
ax.set_ylabel("$C_A$ (mol/L)")
ax.legend()
plt.tight_layout()

Figure 1 shows the two agree. Equations, figures and cross-references all come from Quarto, and code blocks have a copy button.