jaddle 0.1.0__py3-none-any.whl

jaddle/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

jaddle/glop_helpers.py

@@ -0,0 +1,112 @@
+import json
+import os
+import subprocess
+import sys
+import tempfile
+
+import numpy as np
+import scipy.sparse as sp
+from ortools.linear_solver import linear_solver_pb2
+
+from jaddle.highs_helpers import rows_to_standard_form
+
+# The OR-Tools Python bindings don't expose glop's presolver, so it runs in a
+# small C++ helper (tools/glop_presolve, built with tools/glop_presolve/build.sh).
+# Override its location with the JADDLE_GLOP_PRESOLVE environment variable.
+GLOP_PRESOLVE_BINARY = os.environ.get(
+    "JADDLE_GLOP_PRESOLVE",
+    os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        "tools",
+        "glop_presolve",
+        "build",
+        "glop_presolve",
+    ),
+)
+
+
+def glop_presolve(path: str, verbose: bool = False):
+    """
+    Reads an MPS file, relaxes integrality and presolves the LP relaxation with
+    OR-Tools' glop presolver (MainLpPreprocessor, with PDLP's settings: no
+    dualization, no implied-free preprocessor, no scaling).
+
+    Returns: (LP, offset, info) where ``offset`` is the constant objective term
+    outside ``c`` in minimisation form (min c^T x + offset; for a maximisation
+    model the original objective is ``-(c^T x + offset)``, as in
+    ``sciopt_helpers.scip_to_standard_form_sparse``) and ``info`` is the
+    helper's report (status, original/presolved sizes, read/presolve seconds).
+
+    Raises RuntimeError if glop's presolve already decided the problem (status
+    other than INIT, e.g. PRIMAL_INFEASIBLE), since no LP is left to solve.
+    """
+    if not os.path.exists(GLOP_PRESOLVE_BINARY):
+        raise FileNotFoundError(
+            f"glop_presolve helper not found at {GLOP_PRESOLVE_BINARY}; "
+            "build it with tools/glop_presolve/build.sh"
+        )
+    with tempfile.TemporaryDirectory() as tmp:
+        out_path = os.path.join(tmp, "presolved.pb")
+        cmd = [GLOP_PRESOLVE_BINARY, path, out_path]
+        if verbose:
+            cmd.insert(1, "--verbose")
+        proc = subprocess.run(cmd, capture_output=True, text=True)
+        if proc.returncode != 0:
+            raise RuntimeError(
+                f"glop_presolve failed on {path} (exit {proc.returncode}): "
+                f"{proc.stderr.strip()}"
+            )
+        lines = proc.stdout.strip().splitlines()
+        if verbose:
+            print("\n".join(lines[:-1]))
+            sys.stdout.flush()
+        info = json.loads(lines[-1])
+        if info["status"] != "INIT":
+            raise RuntimeError(
+                f"glop presolve decided {path} outright (status {info['status']})"
+            )
+        model = linear_solver_pb2.MPModelProto()
+        with open(out_path, "rb") as f:
+            model.ParseFromString(f.read())
+
+    lp, offset = mpmodel_to_standard_form_sparse(
+        model, scaling_factor=info["objective_scaling_factor"]
+    )
+    return lp, offset, info
+
+
+def mpmodel_to_standard_form_sparse(model, scaling_factor: float = 1.0):
+    """
+    Converts an MPModelProto (linear constraints only) to Jaddle's standard
+    form. Maximisation is negated into minimisation. ``scaling_factor`` is the
+    glop objective scaling factor (true objective = factor * (c^T x + offset)),
+    which MPModelProto can't store; it is folded into ``c`` and the offset.
+
+    Returns: (LP, offset).
+    """
+    if len(model.general_constraint) or model.HasField("quadratic_objective"):
+        raise ValueError("model has non-linear parts; cannot export to a Jaddle LP")
+
+    sense = -1.0 if model.maximize else 1.0
+    scale = sense * scaling_factor
+    variables = model.variable
+    c = scale * np.array([v.objective_coefficient for v in variables], dtype=np.float64)
+    lower_bounds = np.array([v.lower_bound for v in variables], dtype=np.float64)
+    upper_bounds = np.array([v.upper_bound for v in variables], dtype=np.float64)
+    offset = scale * model.objective_offset
+
+    cons = model.constraint
+    row_nnz = np.array([len(k.var_index) for k in cons], dtype=np.int64)
+    indptr = np.concatenate([[0], np.cumsum(row_nnz)])
+    indices = np.fromiter(
+        (j for k in cons for j in k.var_index), dtype=np.int64, count=indptr[-1]
+    )
+    data = np.fromiter(
+        (a for k in cons for a in k.coefficient), dtype=np.float64, count=indptr[-1]
+    )
+    A = sp.csr_matrix((data, indices, indptr), shape=(len(cons), len(variables)))
+    row_lower = np.array([k.lower_bound for k in cons], dtype=np.float64)
+    row_upper = np.array([k.upper_bound for k in cons], dtype=np.float64)
+
+    lp = rows_to_standard_form(c, A, row_lower, row_upper, lower_bounds, upper_bounds)
+    return lp, float(offset)

jaddle/highs_helpers.py

@@ -0,0 +1,90 @@
+import numpy as np
+import scipy.sparse as sp
+import highspy as hspy
+import jax.numpy as jnp
+import jax.experimental.sparse as jsp
+from jaddle.jaddle_basic_types import JaddleLP, LP
+
+
+def highs_to_standard_form_sparse(lp: hspy.HighsLp):
+    """
+    Converts a HighsLp object to standard form matrices:
+        min c^T x
+        s.t. A_eq x = b_eq
+             A_ineq x <= b_ineq
+             x >= lower_bounds
+             x <= upper_bounds
+    Returns: LP (JAX-native; device-side sparse). Stages that need a scipy
+    view (scaling, polish, crossover) call ``LP.to_scipy()``.
+    """
+
+    c = np.array(lp.col_cost_, dtype=np.float64)
+    lower_bounds = np.array(lp.col_lower_, dtype=np.float64)
+    upper_bounds = np.array(lp.col_upper_, dtype=np.float64)
+
+    # Build A matrix from sparse representation
+    num_row = lp.a_matrix_.num_row_
+    num_col = lp.a_matrix_.num_col_
+    A = sp.csc_matrix(
+        (lp.a_matrix_.value_, lp.a_matrix_.index_, lp.a_matrix_.start_),
+        shape=(num_row, num_col),
+        dtype=np.float64,
+    )
+
+    row_lower = np.array(lp.row_lower_, dtype=np.float64)
+    row_upper = np.array(lp.row_upper_, dtype=np.float64)
+    return rows_to_standard_form(c, A, row_lower, row_upper, lower_bounds, upper_bounds)
+
+
+def rows_to_standard_form(c, A, row_lower, row_upper, lower_bounds, upper_bounds):
+    """
+    Converts a row-bounded LP
+        min c^T x  s.t.  row_lower <= A x <= row_upper,  lower <= x <= upper
+    (infinite bounds as +-inf) to the standard form of
+    ``highs_to_standard_form_sparse``: rows with equal finite bounds become
+    A_eq x = b_eq, every other finite row bound becomes a row of
+    A_ineq x <= b_ineq. Returns: LP.
+    """
+    num_col = A.shape[1]
+
+    # Equality constraints: use a numeric tolerance and require finite bounds
+    eps = 1e-8
+    finite_lower_all = np.isfinite(row_lower)
+    finite_upper_all = np.isfinite(row_upper)
+    eq_mask = (
+        finite_lower_all & finite_upper_all & (np.abs(row_lower - row_upper) <= eps)
+    )
+
+    A_eq = A[eq_mask, :].tocsc().astype(np.float64)
+    b_eq = row_lower[eq_mask].astype(np.float64)
+
+    # Inequality constraints (rows not treated as equalities)
+    ineq_mask = ~eq_mask
+    A_ineq_rows = A[ineq_mask, :]
+    row_lower_ineq = row_lower[ineq_mask]
+    row_upper_ineq = row_upper[ineq_mask]
+
+    finite_upper = np.isfinite(row_upper_ineq)
+    finite_lower = np.isfinite(row_lower_ineq)
+
+    # Build inequality matrices more efficiently
+    matrices = []
+    vectors = []
+
+    if finite_upper.any():
+        matrices.append(A_ineq_rows[finite_upper].tocsc().astype(np.float64))
+        vectors.append(row_upper_ineq[finite_upper].astype(np.float64))
+
+    if finite_lower.any():
+        matrices.append((-A_ineq_rows[finite_lower]).tocsc().astype(np.float64))
+        vectors.append((-row_lower_ineq[finite_lower]).astype(np.float64))
+
+    if len(matrices) > 0:
+        A_ineq = sp.vstack(matrices, format="csc")
+        b_ineq = np.concatenate(vectors)
+    else:
+        # No inequality rows: return empty (0 x num_col) matrix and empty RHS
+        A_ineq = sp.csc_matrix((0, num_col), dtype=np.float64)
+        b_ineq = np.empty(0, dtype=np.float64)
+
+    return LP(c, A_eq, b_eq, A_ineq, b_ineq, lower_bounds, upper_bounds)

jaddle/jaddle_basic_types.py

@@ -0,0 +1,272 @@
+import jax.numpy as jnp
+from typing import NamedTuple
+import optax
+from typing import Any, Callable, NamedTuple, Sequence, Union
+import jax
+import numpy as np
+
+ScheduleLike = Union[float, Callable[[jnp.ndarray], jnp.ndarray]]
+
+
+# %%
+# Basic Types
+class SaddleState(NamedTuple):
+    primal: jnp.ndarray
+    dual_ineq: jnp.ndarray
+    dual_eq: jnp.ndarray
+
+
+class JaddleCP:
+    def __init__(
+        self,
+        num_variables,
+        objective,
+        constraints_eq,
+        constraints_ineq,
+        lower_bounds,
+        upper_bounds,
+        dual_bound=None,
+    ):
+        self.num_variables = num_variables
+        self.objective = objective
+        self.constraints_eq = constraints_eq
+        self.constraints_ineq = constraints_ineq
+        self.lower_bounds = lower_bounds
+        self.upper_bounds = upper_bounds
+        self.dual_bound = dual_bound
+
+    def initial_primal_solution(self):
+        return jnp.zeros(self.num_variables)
+
+    def num_eq_constraints(self):
+        return len(self.constraints_eq(self.initial_primal_solution()))
+
+    def num_ineq_constraints(self):
+        return len(self.constraints_ineq(self.initial_primal_solution()))
+
+    def num_constraints(self):
+        return self.num_eq_constraints() + self.num_ineq_constraints()
+
+    def ineq_slack(self, x):
+        return jnp.max(jnp.maximum(self.constraints_ineq(x), 0.0))
+
+    def eq_slack(self, x):
+        return jnp.max(jnp.abs(self.constraints_eq(x)))
+
+    def complementarity_slack(self, x, dual_ineq):
+        return dual_ineq * (self.constraints_ineq(x))
+
+    def initial_solution(self):
+        return SaddleState(
+            primal=jnp.zeros(self.num_variables),
+            dual_ineq=jnp.zeros(self.num_ineq_constraints()),
+            dual_eq=jnp.zeros(self.num_eq_constraints()),
+        )
+
+
+class LP:
+    def __init__(
+        self,
+        c,
+        A_eq,
+        b_eq,
+        A_ineq,
+        b_ineq,
+        lower_bounds,
+        upper_bounds,
+    ):
+        self.c = c
+        self.A_eq = A_eq
+        self.b_eq = b_eq
+        self.A_ineq = A_ineq
+        self.b_ineq = b_ineq
+        self.lower_bounds = lower_bounds
+        self.upper_bounds = upper_bounds
+
+    def objective(self, x):
+        return self.c @ x
+
+    def num_variables(self):
+        return len(self.c)
+
+    def num_eq_constraints(self):
+        return self.A_eq.shape[0]
+
+    def num_ineq_constraints(self):
+        return self.A_ineq.shape[0]
+
+    def num_constraints(self):
+        return self.A_eq.shape[0] + self.A_ineq.shape[0]
+
+    def ineq_slack(self, x):
+        return jnp.max(jnp.maximum(self.A_ineq @ x - self.b_ineq, 0.0), initial=0.0)
+
+    def eq_slack(self, x):
+        return jnp.max(jnp.abs(self.A_eq @ x - self.b_eq), initial=0.0)
+
+    def diff_eq_slack(self, x):
+        return self.A_eq @ x - self.b_eq
+
+    def complementarity_slack(self, x, dual_ineq):
+        return (dual_ineq * (self.A_ineq @ x - self.b_ineq)).sum()
+
+
+class JaddleLP:
+    def __init__(
+        self,
+        c,
+        A_eq,
+        b_eq,
+        A_ineq,
+        b_ineq,
+        lower_bounds,
+        upper_bounds,
+    ):
+        self.c = c
+        self.A_eq = A_eq
+        self.b_eq = b_eq
+        self.A_ineq = A_ineq
+        self.b_ineq = b_ineq
+        self.lower_bounds = lower_bounds
+        self.upper_bounds = upper_bounds
+        # Fused [A_eq; A_ineq] for 2-matvec gradient computation.
+        #
+        # Both A and Aᵀ are kept as BCOO. BCSR was benchmarked here (Aᵀ
+        # materialised as its own row-major CSR, not a lazy `.T`) and was MUCH
+        # slower than BCOO on this workload — do not switch back to BCSR without
+        # re-benchmarking; BCOO's matvec wins for these matrices/precision/device.
+        #
+        # Aᵀ is stored as an explicit transposed BCOO (column-swapped indices)
+        # rather than the lazy `self.A.T`, so `Aᵀ @ y` runs its own matvec
+        # instead of the transposed-operand code path.
+        import jax.experimental.sparse as _jsp
+
+        self.n_eq = A_eq.shape[0]
+        A_bcoo = _jsp.bcoo_concatenate([A_eq, A_ineq], dimension=0)
+        A_T_bcoo = _jsp.BCOO(
+            (A_bcoo.data, A_bcoo.indices[:, ::-1]),
+            shape=A_bcoo.shape[::-1],
+        )
+
+        # Column-swapping A's indices destroys row-major sort order, leaving A_T
+        # with indices_sorted=False/unique_indices=False. JAX's BCOO matvec has a
+        # faster kernel for sorted+unique indices (and the unsorted path can't
+        # skip duplicate accumulation), so sort A_T once at build time. The LP
+        # has no duplicate (row, col) entries, so unique_indices is safe to set.
+        A_bcoo = A_bcoo.sort_indices()
+        A_bcoo.unique_indices = True
+
+        A_T_bcoo = A_T_bcoo.sort_indices()
+        A_T_bcoo.unique_indices = True
+
+        self.A = A_bcoo
+        self.A_T = A_T_bcoo
+        self.b = jnp.concatenate([b_eq, b_ineq])
+
+    @classmethod
+    def from_scipy(cls, c, A_eq, b_eq, A_ineq, b_ineq, lower_bounds, upper_bounds):
+        """Build a ``JaddleLP`` directly from scipy sparse blocks + numpy vectors.
+
+        Mirrors ``jaddle_linear.to_jaddle_sparse``: resolves the active precision
+        profile (float64 under x64, else float32; float16 lives only in JAX since
+        scipy cannot hold it), converts each constraint block to a sorted,
+        deduplicated BCOO, and casts the data array to the profile dtype. This is
+        the JAX-native entry point — callers (e.g. ``highs_to_standard_form_sparse``)
+        no longer route through a scipy ``LP``.
+        """
+        import jax.experimental.sparse as _jsp
+        from jaddle.jaddle_optimisers import jaddle_dtype
+
+        float_dtype = jaddle_dtype()
+        if float_dtype == jnp.float64 and not jax.config.jax_enable_x64:
+            float_dtype = jnp.float32
+        np_float = np.float64 if float_dtype == jnp.float64 else np.float32
+
+        def _to_bcoo(A):
+            A = A.astype(np_float).sorted_indices()
+            A.sum_duplicates()
+            bcoo = _jsp.BCOO.from_scipy_sparse(A.tocoo()).sort_indices()
+            return _jsp.BCOO(
+                (bcoo.data.astype(float_dtype), bcoo.indices), shape=bcoo.shape
+            )
+
+        return cls(
+            jnp.asarray(c, dtype=float_dtype),
+            _to_bcoo(A_eq),
+            jnp.asarray(b_eq, dtype=float_dtype),
+            _to_bcoo(A_ineq),
+            jnp.asarray(b_ineq, dtype=float_dtype),
+            jnp.asarray(lower_bounds, dtype=float_dtype),
+            jnp.asarray(upper_bounds, dtype=float_dtype),
+        )
+
+    def objective(self, x):
+        return self.c @ x
+
+    def num_variables(self):
+        return len(self.c)
+
+    def num_eq_constraints(self):
+        return self.A_eq.shape[0]
+
+    def num_ineq_constraints(self):
+        return self.A_ineq.shape[0]
+
+    def num_constraints(self):
+        return self.A_eq.shape[0] + self.A_ineq.shape[0]
+
+    def ineq_slack(self, x):
+        return jnp.max(jnp.maximum(self.A_ineq @ x - self.b_ineq, 0.0), initial=0.0)
+
+    def eq_slack(self, x):
+        return jnp.max(jnp.abs(self.A_eq @ x - self.b_eq), initial=0.0)
+
+    def diff_eq_slack(self, x):
+        return self.A_eq @ x - self.b_eq
+
+    def complementarity_slack(self, x, dual_ineq):
+        return (dual_ineq * (self.A_ineq @ x - self.b_ineq)).sum()
+
+    def initial_solution(self):
+        """Default start point: the box-projected zero primal with zero duals
+        (the PDLP default)."""
+        return SaddleState(
+            primal=optax.projections.projection_box(
+                jnp.zeros(self.num_variables()), self.lower_bounds, self.upper_bounds
+            ),
+            dual_ineq=jnp.zeros(self.num_ineq_constraints()),
+            dual_eq=jnp.zeros(self.num_eq_constraints()),
+        )
+
+    def to_scipy(self) -> "LP":
+        """Materialise a scipy-backed ``LP`` (CSC matrices, numpy vectors).
+
+        ``JaddleLP`` is the device-side representation the saddle solver iterates
+        on; the scaling, primal/dual polish, crossover and eq-projection LU stages
+        run host-side on scipy sparse (sparse direct solve, ``lsmr``, boolean
+        row-slicing — none of which JAX provides). Those stages call this to get
+        the scipy view they need; the solver never does.
+        """
+        import scipy.sparse as sp
+
+        def _to_csc(bcoo):
+            data = np.asarray(bcoo.data)
+            idx = np.asarray(bcoo.indices)
+            return sp.csc_matrix(
+                (data, (idx[:, 0], idx[:, 1])),
+                shape=bcoo.shape,
+                dtype=np.float64,
+            )
+
+        return LP(
+            np.asarray(self.c, dtype=np.float64),
+            _to_csc(self.A_eq),
+            np.asarray(self.b_eq, dtype=np.float64),
+            _to_csc(self.A_ineq),
+            np.asarray(self.b_ineq, dtype=np.float64),
+            np.asarray(self.lower_bounds, dtype=np.float64),
+            np.asarray(self.upper_bounds, dtype=np.float64),
+        )
+
+
+# %%