jaxsr.mip

jaxsr.mip#

Embed fitted JAXSR models in mixed-integer optimization models.

A JAXSR model is a weighted sum of basis terms, so it can be written into an optimization model term by term. When every selected term is constant, linear, a hinge max(0, x - k), abs(x), a categorical indicator, or an indicator times a continuous feature, the model is piecewise linear and the embedding is an exact mixed-integer linear program (MILP):

Term

Formulation

1, x

written directly

max(0, x - k)

aux h >= 0, h >= x - k, plus a binary and two big-M rows forcing equality

abs(x)

aux t >= x, t >= -x, plus a binary and two big-M rows forcing equality

I(c=v)

one-hot binary for category v of c

I(c=v)*x

exact McCormick product of binary and x

The big-M constants come from the input bounds, which are therefore required and are enforced by the embedding (the surrogate is only valid inside them). Polynomial, interaction, transcendental, power-law and ratio terms are not MILP-representable; pass allow_nonlinear=True to write them as nonlinear expressions (giving an MIQCP or MINLP).

Two backends are supported: Pyomo (to_pyomo()) and discopt (to_discopt()). Both share one formulation, so they produce the same constraints.

jaxsr.mip.MILP_TERM_TYPES = frozenset({'categorical_interaction', 'constant', 'hinge', 'indicator', 'linear'})#

Term types that have an exact MILP formulation.

jaxsr.mip.NONLINEAR_TERM_TYPES = frozenset({'interaction', 'polynomial', 'polynomial_interaction', 'power', 'ratio', 'transcendental'})#

Term types that can be written as nonlinear expressions (allow_nonlinear=True).

jaxsr.mip.milp_incompatible_terms(model: SymbolicRegressor) → list[str]#

List the selected terms of a fitted model that are not MILP-representable.

Parameters:

model (SymbolicRegressor) – A fitted model.

Returns:

names – Names of selected terms with no exact MILP formulation. An empty list means to_pyomo() and to_discopt() produce a pure MILP.

Return type:

list of str

Raises:

  • TypeError – If model is not a SymbolicRegressor.

  • ValueError – If the model is not fitted or uses parametric basis functions.

jaxsr.mip.to_pyomo(model: SymbolicRegressor, block: Any = None, inputs: dict[str, Any] | None = None, output: Any = None, bounds: dict[str, tuple[float, float]] | list[tuple[float, float]] | None = None, output_sense: str | None = None, allow_nonlinear: bool = False) → Any#

Write a fitted model into a Pyomo block.

Parameters:

  • model (SymbolicRegressor) – A fitted model.

  • block (pyomo Block, optional) – Block to populate, e.g. m.surrogate = pyo.Block(). If None, a new ConcreteModel is created and returned.

  • inputs (dict, optional) – Maps feature names to your own Pyomo variables or expressions; each is linked to the block’s input with an equality constraint. Categorical inputs are integer-coded 0..K-1.

  • output (pyomo Var or expression, optional) – Linked to the block’s output with an equality constraint.

  • bounds (dict or list, optional) – Input bounds, as {name: (lo, hi)} or a list in feature order. Defaults to the library’s feature_bounds, then the training data range. The block enforces them.

  • output_sense ({None, "min", "max"}) – If the output is only ever minimized (or only maximized) in your problem, hinge and abs terms whose coefficient makes them convex in that direction are written without binaries. Leave None (the default) for an exact formulation that is valid in any problem.

  • allow_nonlinear (bool) – If True, write non-MILP terms as nonlinear Pyomo expressions.

Returns:

block – The populated block, with components aux (inputs aux["in_<name>"], output aux["out"] and auxiliary variables), binaries, integers and cons. block.jaxsr_inputs and block.jaxsr_output point to the input and output variables.

Return type:

pyomo Block

Raises:

  • ImportError – If Pyomo is not installed.

  • ValueError – If the model has terms that cannot be embedded, has non-MILP terms and allow_nonlinear is False, or bounds are missing or invalid.

Examples

>>> import pyomo.environ as pyo
>>> m = pyo.ConcreteModel()
>>> m.x = pyo.Var(bounds=(0, 1))
>>> m.sr = pyo.Block()
>>> to_pyomo(model, m.sr, inputs={"x": m.x})
>>> m.obj = pyo.Objective(expr=m.sr.jaxsr_output)
class jaxsr.mip.DiscoptEmbedding(inputs: dict[str, Any], output: Any, onehot: dict[tuple[str, Any], Any], binaries: list[Any])#

Bases: object

Variables created by to_discopt().

inputs#

Feature name to input variable (integer-coded for categorical features).

Type:

dict

output#

Variable equal to the model prediction.

Type:

discopt Variable

onehot#

(feature name, category) to the one-hot binary for that category.

Type:

dict

binaries#

All binary variables the embedding created.

Type:

list

inputs: dict[str, Any]#
output: Any#
onehot: dict[tuple[str, Any], Any]#
binaries: list[Any]#
jaxsr.mip.to_discopt(model: SymbolicRegressor, dm_model: Any, inputs: dict[str, Any] | None = None, output: Any = None, bounds: dict[str, tuple[float, float]] | list[tuple[float, float]] | None = None, output_sense: str | None = None, allow_nonlinear: bool = False, prefix: str = 'jaxsr') → DiscoptEmbedding#

Write a fitted model into a discopt model.

Parameters:

  • model (SymbolicRegressor) – A fitted model.

  • dm_model (discopt.Model) – Model to add variables and constraints to.

  • inputs (dict, optional) – Maps feature names to your own discopt variables or expressions; each is linked to the embedding’s input with an equality constraint. Categorical inputs are integer-coded 0..K-1.

  • output (discopt Variable or expression, optional) – Linked to the embedding’s output with an equality constraint.

  • bounds (dict or list, optional) – Input bounds, as {name: (lo, hi)} or a list in feature order. Defaults to the library’s feature_bounds, then the training data range. The embedding enforces them.

  • output_sense ({None, "min", "max"}) – See to_pyomo().

  • allow_nonlinear (bool) – If True, write non-MILP terms as nonlinear discopt expressions.

  • prefix (str) – Prefix for every variable name, so several models can be embedded in one discopt model (default "jaxsr").

Returns:

embedding – The input, output and binary variables created.

Return type:

DiscoptEmbedding

Raises:

  • ImportError – If discopt is not installed.

  • ValueError – If the model has terms that cannot be embedded, has non-MILP terms and allow_nonlinear is False, or bounds are missing or invalid.

Examples

>>> import discopt
>>> m = discopt.Model("design")
>>> emb = to_discopt(model, m)
>>> m.minimize(emb.output)
>>> result = m.solve()