nlpoptnet 0.1.0__py3-none-any.whl

jaxmodel/__init__.py

@@ -0,0 +1,54 @@
+"""Public exports for the symbolic JAX modeling layer used by NLPOptNet."""
+
+from .variables import VariableBuilder, VariableSpec
+from .parameters import ParameterSpec
+from .bounds import BoundSpec
+from .builder import NLPBuilder
+from .highlevel import HighLevelNLPBuilder
+from .objective import QuadraticObjective
+from .constraints import (
+    ConstraintEntry,
+    eq_scalar,
+    ineq_scalar,
+    eq_block,
+    ineq_block,
+    quadratic_eq_scalar,
+    quadratic_ineq_scalar,
+)
+from .blocks import (
+    affine_eq_from_parts,
+    affine_ineq_from_parts,
+    nonlinear_eq_from_parts,
+    nonlinear_ineq_from_parts,
+)
+from .approximations import (
+    LinearizationData,
+    QuadraticObjectiveData,
+    SQPSubproblemData,
+)
+from .model import JaxNLPModel
+
+__all__ = [
+    "VariableBuilder",
+    "VariableSpec",
+    "ParameterSpec",
+    "BoundSpec",
+    "NLPBuilder",
+    "HighLevelNLPBuilder",
+    "QuadraticObjective",
+    "ConstraintEntry",
+    "eq_scalar",
+    "ineq_scalar",
+    "eq_block",
+    "ineq_block",
+    "quadratic_eq_scalar",
+    "quadratic_ineq_scalar",
+    "affine_eq_from_parts",
+    "affine_ineq_from_parts",
+    "nonlinear_eq_from_parts",
+    "nonlinear_ineq_from_parts",
+    "LinearizationData",
+    "QuadraticObjectiveData",
+    "SQPSubproblemData",
+    "JaxNLPModel",
+]

jaxmodel/approximations.py

@@ -0,0 +1,195 @@
+"""Linear and quadratic approximations used to form SQP-style subproblems."""
+
+import jax.numpy as jnp
+from typing import NamedTuple, Optional
+
+
+class LinearizationData(NamedTuple):
+    """Affine approximation data for equality and inequality residuals."""
+
+    A: jnp.ndarray
+    b: jnp.ndarray
+    C: jnp.ndarray
+    d: jnp.ndarray
+    h_val: jnp.ndarray
+    g_val: jnp.ndarray
+
+
+class QuadraticObjectiveData(NamedTuple):
+    """Quadratic approximation data for a local objective model."""
+
+    grad_f: jnp.ndarray
+    Q: jnp.ndarray
+    Q_diag: jnp.ndarray
+    c: jnp.ndarray
+
+
+class SQPSubproblemData(NamedTuple):
+    """Structured data needed to solve a local SQP subproblem."""
+
+    objective: QuadraticObjectiveData
+    constraints: LinearizationData
+    l: Optional[jnp.ndarray]
+    u: Optional[jnp.ndarray]
+
+
+def linearize_constraints_data(eq_fun, ineq_fun, jac_eq, jac_ineq, params, y):
+    """Return matrix data for affine residual linearizations at ``y``."""
+    A = jac_eq(y, params)
+    h = eq_fun(params, y)
+    b = A @ y - h
+
+    C = jac_ineq(y, params)
+    g = ineq_fun(params, y)
+    d = C @ y - g
+
+    return LinearizationData(A=A, b=b, C=C, d=d, h_val=h, g_val=g)
+
+
+def linearize_constraints(eq_fun, ineq_fun, jac_eq, jac_ineq, params, y):
+    """Return affine constraint linearizations as a dictionary."""
+    data = linearize_constraints_data(
+        eq_fun=eq_fun,
+        ineq_fun=ineq_fun,
+        jac_eq=jac_eq,
+        jac_ineq=jac_ineq,
+        params=params,
+        y=y,
+    )
+    return dict(data._asdict())
+
+
+def quadraticize_objective_data(
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    """Return the quadratic approximation data for the objective at ``y``."""
+    grad = grad_fun(y, params)
+
+    if use_diagonal_hessian:
+        qdiag = diag_hess_fun(params, y)
+        if diag_floor is not None:
+            qdiag = jnp.maximum(qdiag, diag_floor)
+        Q = jnp.diag(rho * qdiag)
+    else:
+        Q = rho * hess_fun(y, params)
+
+    c = grad - Q @ y
+
+    return QuadraticObjectiveData(grad_f=grad, Q=Q, Q_diag=jnp.diag(Q), c=c)
+
+
+def quadraticize_objective(
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    """Return the quadratic approximation of the objective as a dictionary."""
+    data = quadraticize_objective_data(
+        grad_fun=grad_fun,
+        hess_fun=hess_fun,
+        diag_hess_fun=diag_hess_fun,
+        params=params,
+        y=y,
+        rho=rho,
+        use_diagonal_hessian=use_diagonal_hessian,
+        diag_floor=diag_floor,
+    )
+    return dict(data._asdict())
+
+
+def build_sqp_subproblem_data(
+    eq_fun,
+    ineq_fun,
+    jac_eq,
+    jac_ineq,
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    lower_fun,
+    upper_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    """Build a structured SQP subproblem approximation at ``y``."""
+    qdata = quadraticize_objective_data(
+        grad_fun=grad_fun,
+        hess_fun=hess_fun,
+        diag_hess_fun=diag_hess_fun,
+        params=params,
+        y=y,
+        rho=rho,
+        use_diagonal_hessian=use_diagonal_hessian,
+        diag_floor=diag_floor,
+    )
+
+    cdata = linearize_constraints_data(
+        eq_fun=eq_fun,
+        ineq_fun=ineq_fun,
+        jac_eq=jac_eq,
+        jac_ineq=jac_ineq,
+        params=params,
+        y=y,
+    )
+
+    return SQPSubproblemData(
+        objective=qdata,
+        constraints=cdata,
+        l=None if lower_fun is None else lower_fun(params),
+        u=None if upper_fun is None else upper_fun(params),
+    )
+
+
+def build_sqp_data(
+    eq_fun,
+    ineq_fun,
+    jac_eq,
+    jac_ineq,
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    lower_fun,
+    upper_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    """Build SQP approximation data as a flat dictionary."""
+    data = build_sqp_subproblem_data(
+        eq_fun=eq_fun,
+        ineq_fun=ineq_fun,
+        jac_eq=jac_eq,
+        jac_ineq=jac_ineq,
+        grad_fun=grad_fun,
+        hess_fun=hess_fun,
+        diag_hess_fun=diag_hess_fun,
+        lower_fun=lower_fun,
+        upper_fun=upper_fun,
+        params=params,
+        y=y,
+        rho=rho,
+        use_diagonal_hessian=use_diagonal_hessian,
+        diag_floor=diag_floor,
+    )
+    return {
+        **dict(data.objective._asdict()),
+        **dict(data.constraints._asdict()),
+        "l": data.l,
+        "u": data.u,
+    }

jaxmodel/autodiff.py

@@ -0,0 +1,30 @@
+"""Automatic-differentiation helpers for flat JAX objectives and constraints."""
+
+import jax
+import jax.numpy as jnp
+
+
+def build_grad_y_objective(obj_flat):
+    """Build the gradient of a flat objective with respect to variables."""
+    return jax.grad(lambda y, p: obj_flat(p, y), argnums=0)
+
+
+def build_hess_y_objective(obj_flat):
+    """Build the Hessian of a flat objective with respect to variables."""
+    return jax.hessian(lambda y, p: obj_flat(p, y), argnums=0)
+
+
+def build_diag_hess_y_objective(grad_fun):
+    """Build a diagonal-Hessian evaluator from a gradient function."""
+    def diag_hess(params, y):
+        n = y.shape[0]
+        basis = jnp.eye(n, dtype=y.dtype)
+        _, jvp_fun = jax.linearize(lambda yy: grad_fun(yy, params), y)
+        cols = jax.vmap(jvp_fun)(basis)
+        return jnp.diag(cols.T)
+    return diag_hess
+
+
+def build_jac_y_vector_fun(vec_fun):
+    """Build the Jacobian of a vector-valued residual with respect to variables."""
+    return jax.jacrev(lambda y, p: vec_fun(p, y), argnums=0)

jaxmodel/blocks.py

@@ -0,0 +1,171 @@
+"""Helpers for assembling structured affine and nonlinear residual blocks."""
+
+from typing import Optional, Sequence, Tuple, Callable
+import jax.numpy as jnp
+from .variables import VariableSpec
+from .constraints import eq_block, ineq_block
+
+
+def _build_affine_jacobian_matrix(
+    var_spec: VariableSpec,
+    y_blocks: Sequence[Tuple[jnp.ndarray, str]],
+) -> jnp.ndarray:
+    """Assemble a dense Jacobian from named affine variable blocks."""
+    n = var_spec.total_size
+    if len(y_blocks) == 0:
+        return jnp.zeros((0, n), dtype=jnp.float64)
+
+    m = int(jnp.asarray(y_blocks[0][0]).shape[0])
+    full = jnp.zeros((m, n), dtype=jnp.asarray(y_blocks[0][0]).dtype)
+
+    for mat, yname in y_blocks:
+        block = jnp.asarray(mat)
+        sl = var_spec.slices[yname]
+        full = full.at[:, sl].add(block)
+    return full
+
+
+def affine_eq_from_parts(
+    var_spec: VariableSpec,
+    y_blocks: Sequence[Tuple[jnp.ndarray, str]],
+    rhs_const: Optional[jnp.ndarray] = None,
+    x_blocks: Optional[Sequence[Tuple[jnp.ndarray, str]]] = None,
+    name: str = "affine_eq",
+):
+    """Build an affine equality block from named variable and parameter pieces."""
+    rhs_const = 0.0 if rhs_const is None else rhs_const
+    x_blocks = [] if x_blocks is None else list(x_blocks)
+
+    def fun(params, vars):
+        lhs = 0.0
+        for A, yname in y_blocks:
+            lhs = lhs + A @ jnp.ravel(vars[yname])
+
+        rhs = jnp.asarray(rhs_const)
+        for D, xname in x_blocks:
+            rhs = rhs + D @ jnp.ravel(params[xname])
+
+        return lhs - rhs
+
+    jac_matrix = _build_affine_jacobian_matrix(var_spec, y_blocks)
+    jac_y_fun = lambda params, y, J=jac_matrix: J
+
+    return eq_block(
+        fun,
+        var_spec,
+        name=name,
+        structure="affine",
+        jac_y_fun=jac_y_fun,
+        metadata={
+            "type": "affine_block",
+            "jac_matrix": jac_matrix,
+            "rhs_const": jnp.asarray(rhs_const),
+            "x_blocks": list(x_blocks),
+        },
+    )
+
+
+def affine_ineq_from_parts(
+    var_spec: VariableSpec,
+    y_blocks: Sequence[Tuple[jnp.ndarray, str]],
+    rhs_const: Optional[jnp.ndarray] = None,
+    x_blocks: Optional[Sequence[Tuple[jnp.ndarray, str]]] = None,
+    name: str = "affine_ineq",
+):
+    """Build an affine inequality block from named variable and parameter pieces."""
+    rhs_const = 0.0 if rhs_const is None else rhs_const
+    x_blocks = [] if x_blocks is None else list(x_blocks)
+
+    def fun(params, vars):
+        lhs = 0.0
+        for C, yname in y_blocks:
+            lhs = lhs + C @ jnp.ravel(vars[yname])
+
+        rhs = jnp.asarray(rhs_const)
+        for E, xname in x_blocks:
+            rhs = rhs + E @ jnp.ravel(params[xname])
+
+        return lhs - rhs
+
+    jac_matrix = _build_affine_jacobian_matrix(var_spec, y_blocks)
+    jac_y_fun = lambda params, y, J=jac_matrix: J
+
+    return ineq_block(
+        fun,
+        var_spec,
+        name=name,
+        structure="affine",
+        jac_y_fun=jac_y_fun,
+        metadata={
+            "type": "affine_block",
+            "jac_matrix": jac_matrix,
+            "rhs_const": jnp.asarray(rhs_const),
+            "x_blocks": list(x_blocks),
+        },
+    )
+
+
+def nonlinear_eq_from_parts(
+    var_spec: VariableSpec,
+    linear_y_terms: Optional[Sequence[Tuple[jnp.ndarray, str]]] = None,
+    nonlinear_y_terms: Optional[Sequence[Tuple[jnp.ndarray, Callable, str]]] = None,
+    rhs_const: Optional[jnp.ndarray] = None,
+    x_blocks: Optional[Sequence[Tuple[jnp.ndarray, str]]] = None,
+    x_direct_name: Optional[str] = None,
+    name: str = "nonlinear_eq",
+):
+    """Build a nonlinear equality block from linear and transformed terms."""
+    linear_y_terms = [] if linear_y_terms is None else list(linear_y_terms)
+    nonlinear_y_terms = [] if nonlinear_y_terms is None else list(nonlinear_y_terms)
+    x_blocks = [] if x_blocks is None else list(x_blocks)
+    rhs_const = 0.0 if rhs_const is None else rhs_const
+
+    def fun(params, vars):
+        lhs = 0.0
+        for A, yname in linear_y_terms:
+            lhs = lhs + A @ jnp.ravel(vars[yname])
+
+        for B, phi, yname in nonlinear_y_terms:
+            lhs = lhs + B @ phi(jnp.ravel(vars[yname]))
+
+        rhs = jnp.asarray(rhs_const)
+        for D, xname in x_blocks:
+            rhs = rhs + D @ jnp.ravel(params[xname])
+
+        if x_direct_name is not None:
+            rhs = rhs + jnp.ravel(params[x_direct_name])
+
+        return lhs - rhs
+
+    return eq_block(fun, var_spec, name=name)
+
+
+def nonlinear_ineq_from_parts(
+    var_spec: VariableSpec,
+    linear_y_terms: Optional[Sequence[Tuple[jnp.ndarray, str]]] = None,
+    nonlinear_y_terms: Optional[Sequence[Tuple[jnp.ndarray, Callable, str]]] = None,
+    rhs_const: Optional[jnp.ndarray] = None,
+    x_blocks: Optional[Sequence[Tuple[jnp.ndarray, str]]] = None,
+    name: str = "nonlinear_ineq",
+):
+    """Build a nonlinear inequality block from linear and transformed terms."""
+    linear_y_terms = [] if linear_y_terms is None else list(linear_y_terms)
+    nonlinear_y_terms = [] if nonlinear_y_terms is None else list(nonlinear_y_terms)
+    x_blocks = [] if x_blocks is None else list(x_blocks)
+    rhs_const = 0.0 if rhs_const is None else rhs_const
+
+    def fun(params, vars):
+        lhs = 0.0
+        for C, yname in linear_y_terms:
+            lhs = lhs + C @ jnp.ravel(vars[yname])
+
+        for F, phi, yname in nonlinear_y_terms:
+            lhs = lhs + F @ phi(jnp.ravel(vars[yname]))
+
+        rhs = jnp.asarray(rhs_const)
+        for E, xname in x_blocks:
+            rhs = rhs + E @ jnp.ravel(params[xname])
+
+        return lhs - rhs
+
+    return ineq_block(fun, var_spec, name=name)

jaxmodel/bounds.py

@@ -0,0 +1,15 @@
+"""Bound specifications for optimization variables."""
+
+from dataclasses import dataclass
+from typing import Optional, Callable
+import jax.numpy as jnp
+
+Array = jnp.ndarray
+
+
+@dataclass
+class BoundSpec:
+    """Hold callable lower and upper bound evaluators."""
+
+    lower_fun: Optional[Callable[[dict], Array]] = None
+    upper_fun: Optional[Callable[[dict], Array]] = None

jaxmodel/builder.py

@@ -0,0 +1,80 @@
+"""Low-level builder for constructing :class:`jaxmodel.model.JaxNLPModel`."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Callable, Optional, Sequence
+import jax.numpy as jnp
+
+from .bounds import BoundSpec
+from .constraints import ConstraintEntry, eq_block, eq_scalar, ineq_block, ineq_scalar
+from .model import JaxNLPModel
+from .parameters import ParameterSpec
+from .variables import VariableSpec
+
+
+@dataclass
+class NLPBuilder:
+    """Incrementally assemble a :class:`JaxNLPModel` from its components."""
+
+    var_spec: VariableSpec
+    parameter_spec: Optional[ParameterSpec] = None
+    dtype: jnp.dtype = jnp.float64
+    _objective_fun: Optional[Callable] = None
+    _constraints: list[ConstraintEntry] = field(default_factory=list)
+    _bounds: BoundSpec = field(default_factory=BoundSpec)
+
+    def set_objective(self, objective_fun) -> "NLPBuilder":
+        """Set the scalar objective callable or structured objective."""
+        self._objective_fun = objective_fun
+        return self
+
+    def add_constraint(self, constraint: ConstraintEntry) -> "NLPBuilder":
+        """Add one constraint entry to the model."""
+        self._constraints.append(constraint)
+        return self
+
+    def add_constraints(self, constraints: Sequence[ConstraintEntry]) -> "NLPBuilder":
+        """Add multiple constraint entries to the model."""
+        self._constraints.extend(constraints)
+        return self
+
+    def add_eq_scalar(self, fun, name: str = "eq_scalar") -> "NLPBuilder":
+        """Add a scalar equality residual."""
+        return self.add_constraint(eq_scalar(fun, self.var_spec, name=name))
+
+    def add_ineq_scalar(self, fun, name: str = "ineq_scalar") -> "NLPBuilder":
+        """Add a scalar inequality residual."""
+        return self.add_constraint(ineq_scalar(fun, self.var_spec, name=name))
+
+    def add_eq_block(self, fun, name: str = "eq_block") -> "NLPBuilder":
+        """Add a vector-valued equality block."""
+        return self.add_constraint(eq_block(fun, self.var_spec, name=name))
+
+    def add_ineq_block(self, fun, name: str = "ineq_block") -> "NLPBuilder":
+        """Add a vector-valued inequality block."""
+        return self.add_constraint(ineq_block(fun, self.var_spec, name=name))
+
+    def set_bounds(
+        self,
+        lower_fun=None,
+        upper_fun=None,
+    ) -> "NLPBuilder":
+        """Attach callable lower and upper bounds."""
+        self._bounds = BoundSpec(lower_fun=lower_fun, upper_fun=upper_fun)
+        return self
+
+    def build(self, jit_compile: bool = True) -> JaxNLPModel:
+        """Finalize and construct the :class:`JaxNLPModel`."""
+        if self._objective_fun is None:
+            raise ValueError("Objective function must be set before build().")
+
+        return JaxNLPModel(
+            var_spec=self.var_spec,
+            objective_fun=self._objective_fun,
+            constraints=self._constraints,
+            bounds=self._bounds,
+            parameter_spec=self.parameter_spec,
+            dtype=self.dtype,
+            jit_compile=jit_compile,
+        )