An agent is only as good as its environment

The simplest way to describe our approach is that we spend our effort on the environments that we work in and that the agent works in, rather than on the agent itself. Three ideas run through everything below.

The first is that plain text is ready for agents. Our notes, code, data and manuscripts live in org-mode and Markdown files that an agent can read, search and edit with the same tools it uses on source code. Nothing has to be exported or scraped first.

The second is that abilities should be installable. When an agent needs to know how to set up a VASP calculation or diagnose a stalled interior point solve, that knowledge should not live in a prompt someone pasted once. It should ship as a skill or an MCP server, versioned alongside the code it describes, and be installable with one command.

The third, and the most important, is that software should report its work. Agents are fluent, but fluency is not correctness. The tools we build for agents return structured results, validate their input, raise typed errors instead of returning a plausible wrong answer, and are tested against references we trust. That gives us something to check the agent’s claims against.

The diagram shows how the pieces fit together. You set the goals and check the results. The agent does the work through abilities that knowhow installs, using four groups of tools. Underneath is plain text, and around all of it are checks the agent cannot talk its way past. The rest of the post follows the diagram from top to bottom.

The workbench: scimax for VS Code

scimax for VS Code is where the work happens. It brings the scimax environment I built in Emacs to VS Code, and it is one plain-text environment for the whole research cycle. I think in it, with a journal, notes, an agenda, and a searchable database of every org file I have. I compute in it, with runnable Python and Jupyter code blocks inside the document. I write in it, with org-ref citations, BibTeX tools and manuscript publishing for journal submission. I publish from it, and I present from it, using Marp decks with a pen, a laser pointer, a timer and live Python.

Because all of this is plain text, an agent can work with it as easily as I can. A TODO is just a headline in an org file, and scimax gathers those headlines into an agenda. Claude can read that agenda and update it too.

The agent-facing part of scimax is a command-line tool and four skills. The scimax CLI covers the agenda, search, the org database, the journal, tasks and projects, citation checking, export, tangling, publishing and slide builds. scimax skill install adds skills that teach an agent how to use it, so it can act as an executive assistant over my notes: a morning briefing, triaging tasks, finding what is blocked, or working out the critical path through a project. Sites built with scimax publish can include an “Ask” chat that finds the relevant sections and writes an answer with citations, using a language model that runs in the reader’s browser, so there is no server to maintain. This site has one. scimax for VS Code has had 838 commits since January 2026, about 72% of them co-authored with AI.

Giving agents abilities: knowhow

The skills that come with scimax are one example of a broader pattern, and knowhow is how we manage it. knowhow is a package manager for agent capabilities. uvx knowhow add owner/repo --skill name installs skills, slash commands, subagents, hooks and MCP servers from any git repository. It works with Claude Code, Codex, OpenCode, Gemini, Cursor and Copilot, and it records what it installed in a lock file so a project’s agent setup can be reproduced. knowhow replaces an earlier tool of ours called skillz.

knowhow also carries a library of our own: 44 skills for scientific, laboratory and programming work (ASE, VASP, pymatgen, Opentrons robots, design of experiments, electronic lab notebooks and scientific review, among others), along with subagents, commands and hooks. Some of the hooks are guardrails rather than abilities. One keeps an automatic lab-notebook record of what the agent did, and another keeps it from writing secrets into files.

The most useful thing about knowhow may be the convention it set. Every package in the rest of this post ships its own skill in skills/<name>/SKILL.md, next to the code it describes. The instructions change when the code changes, and knowhow add can install any of them.

Knowing what we know: litdb and crucible

Research starts with the literature, and agents are notorious for inventing it. Our answer is to give the agent a library that we curate, so it searches real papers instead of what it half-remembers.

litdb is a local database of the papers I care about, with natural-language search. Its MCP server is read-only by design. An agent can run semantic and full-text searches, find papers similar to one it already has, pull new work from OpenAlex, and return citations and BibTeX entries that point to real records. It can even generate the collaborator list an NSF proposal requires, which is the kind of tedious, checkable task that suits an agent well.

crucible works one level up. An LLM distills papers, notebooks and data into an interlinked org-mode wiki, one article per concept, with an SQLite graph of the links behind it. Ingesting happens in two steps: the agent first proposes which articles to write or update, then writes them. Once the wiki exists, an agent can use the crucible CLI to search it, follow backlinks, find related articles, and look for gaps such as orphan articles and sources that have not been distilled yet. crucible publish turns the wiki into a website using scimax.

These tools are not wired together in code, but they make one workflow. litdb finds the papers, crucible distills them into concepts, and scimax is where I write with them and publish the result. org-ref citation keys run through all three, so a claim in a draft can be traced back to a paper in the database.

Computing with agents: fairchem-mcp and discopt-calphad

Most atomistic simulation workflows were designed for people. You write a script, submit it, wait, and read the output. When an agent drives that workflow, it does the same thing in a loop: write a script, run it, wait, read a log file, and guess what went wrong. Every step reloads the model, and the agent finds out that a calculation went badly only after it is over. We wanted computation that an agent could work with interactively, and results that carry their own evidence of correctness.

uvx knowhow add jkitchin/fairchem-mcp --skill fairchem-mcp --mcp fairchem

Solvers that explain themselves

Since December 2025 we have started 15 public repositories in differentiable programming and optimization. Together they hold about 1.8 million lines of source and tests, roughly 1.1 million in Python and 0.66 million in Rust, and they have had about 63,000 downloads from PyPI in the last 180 days. Projects like these usually take a team years. The map above, from an earlier post, shows how they relate.

Optimization turns out to be a good domain for AI-assisted development because the answers can be checked. A reported optimum either satisfies the optimality conditions and matches the reference solution, or it does not, and there are large public benchmark sets that say what the right answer is. The same property makes these solvers good tools for agents to use, as long as the solvers explain what they did. That was a design goal from the start.

difflow is a JAX framework for differentiable chemical process flowsheets, for simulation, optimization and technoeconomic analysis. Its MCP server lets an agent build a flowsheet one step at a time: choose the species, add units, connect streams and set feeds, with undo and redo when it makes a mistake. Then it can solve the flowsheet, run sensitivity analyses and sweeps, optimize it, propagate uncertainty, and estimate costs. When a flowsheet does not converge, which is the normal state of a new flowsheet, the diagnosis tools return structured findings and suggested remedies rather than a stack trace, and converge tries the remedies in turn. Each tool is labeled as reading, editing or executing code, and one server configuration removes arbitrary Python execution entirely, for when you want the agent to stay inside the modeling tools. difflow can use pounce and discopt as its solvers.

pounce is a family of interior point solvers written in Rust. It began as a port of Ipopt and now covers nonlinear, conic and quadratic programs. Every solve writes a JSON solve report, so an agent reads structured data instead of parsing an iteration log. The pounce-studio MCP server has about 30 tools built on those reports. Before a solve, an agent can analyze the problem and check the starting point. After a solve, it can diagnose the result, find where progress stalled, look at the convergence trace and the restoration phases, compare two runs, and verify the solution independently. If that is not enough, it can step through a solve in a live debugger. pounce also includes commands for adversarial testing and for reviewing sensitivity results.

discopt is a modeling language and solver for mixed-integer nonlinear programs, and it has more agent-facing features than anything else we have built. It ships 17 expert subagents (for MINLP solving, convexity detection, Ipopt, benchmarking and more) and commands for each stage of modeling: formulating, diagnosing, debugging, reformulating, explaining a model, converting it between formats, and reporting on benchmarks. It also has a tool-call interface for building models from a natural-language description. The agent creates variables and adds constraints and logical conditions through tool calls, and the expressions it writes go through a safe parser instead of being executed as code. When a model is infeasible or a solve goes wrong, discopt produces reports on infeasibility, convexity and provenance that an agent can act on.

Around discopt is a growing set of plugins that bring the same approach to particular domains. Each one checks its input against a schema, solves the problem, and reports the result, through a skill or an MCP server. discopt-mkm does microkinetic modeling for catalysis, including degree of rate control and limiting potentials. discopt-grid does power systems optimization, and it was built largely by an autonomous agent loop working from a written specification and implementation plan. discopt-doe adds design of experiments, with agents for estimability and identifiability analysis. discopt-flowsheet raises typed errors instead of returning silently wrong answers. discopt-course is an optimization course in which Claude Code is the teaching assistant, grading against rubrics. discopt-calphad, covered in the computing section, is one more plugin.

Where this is going

The direction is agents that run the whole loop, from literature to computation to experiment to paper, with people setting the goals and checking the results. The recent roadmap for AI in heterogeneous catalysis describes that future for one field (Xin et al., 2026). The software here is our attempt to build the environment it needs.

If I had to compress the approach into advice for students, it would be four things. Own the verification: you are responsible for every number, line of code and citation, so learn enough of the fundamentals to know when an answer is wrong. Write things down in plain text, where an agent can read them. Specify, then delegate: the skill that matters now is stating a problem precisely, with tests or checks that define when it is done. And be ambitious, and disclose. Attempt projects that were out of reach, and take responsibility for what you generate.

AI does not reduce the need for expertise. It makes judgment more valuable.

All of the software is on our software page and at github.com/jkitchin.

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

How to choose

There are many things you should consider in choosing a surrogate:

1. Accuracy: can it fit the data well enough for your task?
2. Mathematical / optimization properties: is it convex, monotonic, periodic, piecewise-linear, etc.?
3. Interpretability: can you read the model, check it against physics, and publish it?
4. Uncertainty: does it give honest intervals, and do they cover the truth?

A lot depends on what you do with the model after you make it. Here are some guidelines for picking a surrogate, and the packages that implement them.

A few rules of thumb sit behind the table.

• Structure you can guarantee beats structure you hope the data teach. If you know a response is monotonic or convex, use a model that cannot be otherwise. Constraints enforced only at training points can still be violated between them.
• If the surrogate goes into an optimizer, pick it for the optimizer. For a MILP use piecewise-linear trees. For a convex problem use an ICNN. For a nonlinear program, any smooth model with good derivatives will do, and a short symbolic expression is easiest on the solver.
• Interpretability and flexibility trade off. Symbolic models and linear trees can be read. Networks usually fit more but explain less. KANs sit in between, because you can plot each learned edge function.
• Check the uncertainty, not just the fit. Each model gets its σ a different way, and they fail differently. The UQ comparison is a good place to start.

The symbolic models

The trees

The neural networks

The networks fall into two groups. Some build a property into the architecture so it holds everywhere. Others add a cheap uncertainty estimate to an otherwise ordinary network.

Distributions

Across the zoo

Finishing thoughts

Hopefully this gives you some insights into the breadth of differentiable surrogates that know their uncertainty, and where to find an example of each. The zoo is still growing, and we welcome feedback. If you find bugs, please report them!

Why differentiable models?

If there is one nearly unifying principle in optimization, it is that derivatives play a central role in much of it (I know not all of it; integers are not differentiable, and you often have to estimate derivatives, etc.). Optimization means many things as illustrated above. We can use it to fit parameters, simulate, reconcile data, plan, control, design experiments and more! These all rely on the central model, and different combinations of what know, what you want, and how you get it. What we aim for in our ecosystem shown above is for a single model to be reusable across many tasks.

The graphic above is interactive, and you can see how the task is implemented in many of the tasks through notebook examples. We have not fully integrated everything seamlessly yet, but we are working on it and making good progress. If you find bugs or have suggestions, please let us know by opening an issue on GitHub!

The big pivot into optimization

Over the past year I have been making a big pivot in research into optimization. I have always learned by doing, and this post documents some of that journey. This transition has been happening for a while, really starting about a decade ago when I started using machine learning to build surrogate models for DFT. That work led to this perspective on differentiable programming in scientific programming. That has led to this year’s work in differentiable optimization. My goal is to build an end-to-end differentiable optimization stack that integrates differentiable programming, differentiable solvers, and differentiable surrogates.

Why do we need new solvers at all?

There are a plethora of options, very good commercial ones, and open-source ones. I work in higher education, and I think it is important that free, open-source solvers exist for students and researchers. I got into computational catalysis as a student because there was a free code I could use, read, modify and learn from. It is often not possible to do this with closed codes. Even if my codes aren’t as good as other codes, if they open the door for others to learn then a part of the mission is accomplished.

New starts mean new things are possible. pounce is more comprehensive than Ipopt in interior point methods, with integrated sensitivity and other great features that are 1) pip-installable, 2) capable of running in your browser, and with interfaces to use in other interfaces. discopt led to new integrations of ideas, and eventually to the development of jax-go which bridges the gap between program and equation-based optimization. difflow is a new way to think about optimization, and the differentiable surrogates are a new way to think about optimization with uncertainty quantification.

Agentic optimization is a new thing with these solvers. The tight integration of agentic coding with differentiable programming and differentiable solvers is a new paradigm that is more possible with these new solvers because they were designed with this in mind. We have integrated agentic debugging, and agentic optimization.

What is next?

I have laid out an ambitious vision for the end-to-end differentiable optimization stack, and I think we are making good progress on it. It is a lot, and we don’t use all of them every day. Some of these tools are stabilizing. Feral is pretty stable now. I probably work on pounce once a week or so. discopt and difflow are still changing fast. It would probably be a lie to say there won’t be new packages coming. I still have many private projects that haven’t seen the light of day yet, but will eventually.

We (the Kitchin group) are beginning to enter a new era of applications. Inspired by the algebraic modeling language we made discopt-mkm for microkinetic modeling, and discopt-calphad for alloy phase behavior. discopt-doe is a new direction in design of experiments. We have projects in process modeling, alloy design, rare earth separation, and more.

It is always dangerous to predict the future, but I think 5 years from now we will have a much more integrated landscape of differentiable optimization tools. Right now programmatic automatic differentiation is mostly only available as addons like jax, pytorch in Python, and something like that in Julia. AD is coming to Rust as a native feature in a few years, and predict this will be a game changer for optimization. The pace of development I have demonstrated this year means it will probably only take a year or so to transition from a Python or Julia focused solution space to Rust if we decide to. These changes come with 20/20 hindsight and lessons learned from everything we know now, so they happen very quickly. I think the next 5 years will be very exciting for differentiable optimization.

I still think we are not thinking big enough as a community. I am grateful to a few people who made substantial contributions to these codes: Devin Griffith for major contributions to sensitivity in pounce, Germán Heim for integrating pounce into his modeling language oximo. Milan Rother forked feral into RSLAB and was the inspiration for pounce in the browser. Victor Alves for using pounce before anyone else. David Bernal for major contributions in global optimization in discopt, and ideas linking it to pounce. All this work was done by a very small number of people, and I think we can do better as a community.

It is an open challenge to figure out how to sustain these efforts. Pyomo and IDAES were built over two decades, with generous federal funding. This is a challenge today as that funding becomes more scarce. We need new models to support this kind of work that balances keeping open-access to the code, and generating revenue to support the people who build it.

Along this journey I have come to appreciate that 1) I still have a lot to learn in this space, and 2) there is not one tool to rule in optimization; there are many ways to solve problems, and some ways are better than others for certain problems. That means you have to have a landscape of pathways to choose, and historically I have chosen the roads less traveled. I have learned a lot along the way, and I hope this post helps others to learn from my journey.

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.