pyepo.model.mpax.mpaxmodel

Abstract optimization model based on MPAX

Classes

optMpaxModel

Abstract base class for MPAX-backed (JAX) linear / quadratic program models.

Module Contents

class pyepo.model.mpax.mpaxmodel.optMpaxModel

Bases: pyepo.model.opt.optModel

Abstract base class for MPAX-backed (JAX) linear / quadratic program models.

MPAX is a JAX implementation of the PDHG (Primal-Dual Hybrid Gradient) first-order solver, designed for large-scale continuous programs that benefit from GPU acceleration and vmap-batched solving. Unlike the Gurobi / COPT / Pyomo / OR-Tools backends, an MPAX model has no explicit solver model object – the constraint matrices and bounds are the model. Subclasses populate them inside _getModel and return (None, []):

def _getModel(self):
    self.A = jnp.array(...)   # equality A x = b
    self.b = jnp.array(...)
    self.G = jnp.array(...)   # inequality G x >= h
    self.h = jnp.array(...)
    self.l = jnp.array(...)   # variable lower bound
    self.u = jnp.array(...)   # variable upper bound
    # optional: leave None for LP, set for convex QP
    self.Q = jnp.array(...)   # PSD; objective is 0.5 xᵀQx + cᵀx
    return None, []

LP vs QP is selected automatically from self.Q: None (default) keeps the LP code path via create_lp, any other value routes through create_qp. Q must be PSD; MPAX supports quadratic objective only – constraints stay linear (this is a hard MPAX limit; e.g. a quadratic risk-budget constraint cannot be expressed).

Objective sense follows self.modelSense (set by a problem-level base such as knapsackBase or directly in _getModel; defaults to minimization). Dense vs sparse matrices can be toggled by overriding the class attribute use_sparse_matrix (default True).

A jitted single-instance solver and a vmap-batched solver (batch_optimize) are pre-compiled on construction, so optDataset can solve every training instance in a single dispatch.

Variables:

  • A (jnp.ndarray) – equality-constraint matrix (Ax = b)

  • b (jnp.ndarray) – equality-constraint right-hand side

  • G (jnp.ndarray) – inequality-constraint matrix (Gx >= h)

  • h (jnp.ndarray) – inequality-constraint right-hand side

  • l (jnp.ndarray) – variable lower bounds

  • u (jnp.ndarray) – variable upper bounds

  • Q (jnp.ndarray | None) – PSD quadratic-objective matrix; None ⇒ LP

  • use_sparse_matrix (bool) – whether to use sparse matrices

use_sparse_matrix: bool = True
Q = None
device = None
property num_cost: int

number of costs to be predicted

setObj(c: numpy.ndarray | torch.Tensor | list) → None

A method to set the objective function

Parameters:

c – cost of objective function

solve() → tuple[torch.Tensor, float]

A method to solve the model

Returns:

optimal solution (torch.Tensor) and objective value (float)

Return type:

tuple

copy() → Self

A method to copy the model

Returns:

new copied model

Return type:

optModel

addConstr(coefs: numpy.ndarray | torch.Tensor | list, rhs: float) → Self

A method to add a new constraint

Parameters:

  • coefs – coefficients of new constraint

  • rhs – right-hand side of new constraint

Returns:

new model with the added constraint

Return type:

optModel