Optimizing over a fitted model¶
An iBART model is a short formula: an intercept plus a few closed-form descriptors, each with a coefficient. That makes it more than a predictor. You can write it as exact algebra in an optimization model and ask a global solver questions like:
- which inputs give the largest (or smallest) prediction?
- what is the cheapest design that the model says reaches a target?
IBart.to_discopt does the translation for
discopt, a global MINLP solver in JAX. It
walks each descriptor's sympy tree and emits discopt's own operators (no
black-box callbacks), so the solver can build relaxations and certify that
its answer is the global optimum, not just a local one.
discopt needs Python 3.12 or newer: pip install "jax-ibart[discopt]".
A model to optimize¶
We use a problem where we know the answer: y = 3ab + 2c/d plus noise, with a
distractor feature e. One binary round finds both descriptors.
import warnings
import numpy as np
from jax_ibart import IBart
rng = np.random.default_rng(0)
names = ["a", "b", "c", "d", "e"]
lo = np.array([0.5, 0.5, 0.0, 0.5, 0.0])
hi = np.array([3.0, 1.5, 1.0, 2.0, 1.0])
X = rng.uniform(lo, hi, size=(80, 5))
a, b, c, d = X[:, 0], X[:, 1], X[:, 2], X[:, 3]
y = 3 * a * b + 2 * c / d + rng.normal(0, 0.05, 80)
model = IBart(backend="diff", schedule=("b",), max_dim=2, seed=0, steps=1200)
model.fit(X, y, feature_names=names)
model
Maximize the prediction¶
to_discopt adds one continuous variable per input the model actually uses
(e is not used, so it does not appear), bounded to the range of the training
data. It returns those inputs and a variable y that is constrained to equal the
prediction. Then we maximize y.
import discopt.modeling as dm
m = dm.Model("max_y")
x, yv = model.to_discopt(m)
m.maximize(yv)
with warnings.catch_warnings():
warnings.simplefilter("ignore") # discopt's solver diagnostics
res = m.solve(time_limit=60)
best = {k: round(float(res.value(v)), 4) for k, v in x.items()}
print(res.status, "certified:", res.gap_certified)
print("best inputs:", best)
print("max prediction:", round(res.objective, 4))
optimal certified: True
best inputs: {'a': 2.9546, 'b': 1.4972, 'c': 0.9951, 'd': 0.524}
max prediction: 17.0415
The maximum is at a corner of the box: a and b as large as the data
allow, c as large as possible, d as small as possible. That is what the
formula says, and the solver proved there is nothing better in the box.
A useful check is that the solver's answer is a genuine prediction of the model:
plug the inputs back into predict. It should also beat every training point.
row = np.array([[best.get(n, 0.0) for n in names]])
print("predict at the optimum:", round(float(model.predict(row)[0]), 4))
print("best training prediction:", round(float(model.predict(X).max()), 4))
predict at the optimum: 17.0415 best training prediction: 12.4283
Use the model as a constraint¶
Optimizing the prediction alone usually just drives the inputs to their bounds.
A more realistic question puts the model inside a larger design problem. Say each
input costs something to raise, and we want the cheapest design that the model
predicts reaches y >= 6:
$$\min_{a,b,c,d}\; a + 4b + 3c - d \quad\text{s.t.}\quad \hat y(a,b,c,d) \ge 6$$
(-d: raising d saves money, but it also lowers y through c/d.) The
answer is a trade-off, not a corner.
m = dm.Model("cheapest")
x, yv = model.to_discopt(m)
m.subject_to(yv >= 6.0, name="target")
m.minimize(x["a"] + 4 * x["b"] + 3 * x["c"] - x["d"])
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res = m.solve(time_limit=60)
design = {k: round(float(res.value(v)), 4) for k, v in x.items()}
row = np.array([[design.get(n, 0.0) for n in names]])
print(res.status, "certified:", res.gap_certified)
print("design:", design, " cost:", round(res.objective, 4))
print("predicted y at the design:", round(float(model.predict(row)[0]), 4))
optimal certified: True
design: {'a': 2.8227, 'b': 0.7057, 'c': 0.0147, 'd': 1.9955} cost: 3.6939
predicted y at the design: 6.0004
The target constraint is active (predicted y is 6), and the solver
certified the global optimum of a nonconvex problem: ab is bilinear and c/d
is a ratio, so a local solver could stop at a worse design.
Guarding the model's domain¶
A descriptor like c/d has a pole at d = 0, and log or sqrt descriptors
are undefined for some inputs. If you widen the bounds beyond the data, an
optimizer will happily walk to a pole and report an enormous "prediction".
So by default to_discopt also adds domain constraints taken from the
training data: a denominator keeps the sign it had on the data and stays at
least as far from zero as it ever was there; a log argument stays at least as
large as its smallest training value; a sqrt argument stays non-negative.
Here we let d go all the way down to 0. With the guard, the optimizer stops
at the smallest d seen in training. With domain=False the box would contain
the pole at d = 0.
wide = (np.zeros(5), hi)
m = dm.Model("wide")
x, yv = model.to_discopt(m, bounds=wide)
m.maximize(yv)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res = m.solve(time_limit=60)
print(res.status, "d at optimum:", round(float(res.value(x["d"])), 4),
" smallest training d:", round(X[:, 3].min(), 4))
optimal d at optimum: 0.524 smallest training d: 0.524
What the box does and does not protect against¶
By default the inputs are bounded to the box spanned by the training data (each feature's min and max). That is a coarse notion of "where the model is valid". If your features are correlated, the box contains corners far from any training point, and the optimum often lands in one of them. Here the features were sampled independently, so the box is filled; with real data, check how far the optimum is from the nearest training point before trusting it, or add your own constraints that keep the design near the data.
Other options¶
x=...uses your own discopt variables or expressions (a dict by feature name, or a sequence in feature order), so the model can share inputs with the rest of a larger model. The box is then yours to set, and you can passy_bounds=(lb, ub)to bound the prediction.y_bounds="auto"(the default whento_discoptcreates the inputs) boundsyby interval arithmetic over the box. These bounds are conservative, so they never cut off a reachable prediction, and they make the solve faster.X_ref=...takes the bounds and domain guards from other reference data.name=...names the prediction variable, so several models can live in one optimization problem (for example, two properties of one material).- Supported functions:
+ - * /, integer and fractional powers,exp,log,sqrt,abs,sin,cos,tan,tanh. Anything else raisesNotImplementedErrornaming the function.