Solvers¶

@dataclass(frozen=True, slots=True)
class Bounded:
  """A named constraint group with an optional lower and upper bound."""

  expr: Expr
  lo: Expr | None = None
  hi: Expr | None = None
  name: str | None = None

@dataclass(frozen=True, slots=True)
class ProblemSpec[SymbolicVars]:
  """The objective, constraint groups, and box bounds returned by a problem body.

  ``lb`` and ``ub`` have the variables’ tree structure. Scalar expression leaves broadcast over
  their variable blocks. Use ``NO_LB`` or ``NO_UB`` for an open side of one leaf.
  """

  minimize: Expr
  eq: tuple[Expr, ...] = ()
  ineq: tuple[Bounded, ...] = ()
  lb: SymbolicVars | None = None
  ub: SymbolicVars | None = None

@dataclass(frozen=True, slots=True)
class StageConstraint:
  """The constraint ``lo <= fn(...) <= hi`` on one stage.

  ``fn`` takes the state and the control for a path constraint, which holds at every stage before
  the last, and the state alone for a terminal constraint. Its further inputs are matched by name, as
  for the costs. A side is ``None``, a number, or an array of the output's size, and ``lo == hi``
  makes an equality.
  """

  fn: Function[Any, Any, Any, Any]
  lo: Any = None
  hi: Any = None

  def __post_init__(self) -> None:
    if self.lo is None and self.hi is None:
      raise ValueError("StageConstraint needs at least one of lo / hi")

@dataclass(frozen=True, slots=True)
class Solver[SV, NV, SP, NP]:
  """A compiled solver called with its parameters. The initial point and multipliers default to zero.

  ``function`` is the plain ``Function`` with the full five-group signature, for code generation
  and symbolic composition. Its four outputs are its first four
  inputs, so passing a previous result as ``warm`` warm-starts the next call.
  """

  function: SolverFunction[SV, NV, SP, NP]

  @overload
  def __call__(
    self, params: NP, /, *, x0: NV | None = None, warm: tuple[NV, NV, np.ndarray, np.ndarray] | None = None
  ) -> tuple[NV, NV, np.ndarray, np.ndarray]: ...

  @overload
  def __call__(self, params: SP, /, *, x0: SV | None = None, warm: tuple[SV, SV, Expr, Expr] | None = None) -> tuple[SV, SV, Expr, Expr]: ...

  def __call__(self, params: Any, /, *, x0: Any = None, warm: Any = None) -> Any:
    """Solve for ``params`` from ``warm``, or from ``x0`` and zero multipliers, or from zero."""
    if x0 is not None and warm is not None:
      raise TypeError("pass either x0 or warm, not both")
    parts = cast(Any, self.function.instantiate().input_tree).parts
    symbolic = parts[4].is_symbolic(params)
    if warm is None:
      warm = (_zeros(parts[0], symbolic) if x0 is None else x0, *(_zeros(part, symbolic) for part in parts[1:4]))
    init = tuple(warm)
    return self.function(*init, params)

  def stats(self) -> SolverStats:
    """The statistics of the latest numerical call."""
    return self.function.solver_stats()

class SolverLibraryError(ToolchainError):
  """Raised when a solver's shared library or C headers cannot be located or loaded.

  Compiling or calling a function that reaches the solver raises it. The message lists the
  locations searched, as ``scaly_toolchain`` reports them.
  """

class NotQuadratic(ValueError):
  """A problem rejected because a QP oracle depends nonlinearly on its variables."""

class SolverBackend(Protocol):
  """The full solver plugin protocol: packaging metadata plus the C wrapper
  template. Solves always run through the generated C wrapper, and plugins ship
  no Python solve path."""

  name: str
  kind: str  # "qp" | "nlp" - which descriptor family the plugin solves
  protocol_version: int
  lib_stem: str  # shared library stem: lib<stem>.{dylib,so}
  link_flags: tuple[str, ...]
  header: str  # C header path relative to include_dir(), e.g. "piqp/piqp.h"

  def lib_dir(self) -> Path: ...

  def include_dir(self) -> Path: ...

  def prepare_options(self, options: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
    """Validate options, supply defaults, and split compilation choices from runtime tuning."""
    ...

  def render_wrapper(self, fun: ConcreteFunction, ctx: SolverWrapperCtx) -> list[str]:
    """Emit the C wrapper for one solver ``ConcreteFunction`` (see docs/dev/solver_plugins.md).

    Must define ``static void <ctx.raw_symbol>(...)`` with the descriptor's
    ``in*``/``out*`` signature followed by ``double* w`` and one option-array pointer,
    drive the solver's
    C API with the oracle kernels (``ctx.raw_symbol_of``), and fill
    ``ctx.stats_symbol`` on every call.
    """
    ...

class NlpSolverBackend(SolverBackend, Protocol):
  """Solver backend protocol for NLP plugins, including Hessian layout and stage oracles."""

  hess_triangle: Literal["lower", "upper"]
  stage_oracles: bool = False

@dataclass(frozen=True, slots=True)
class SolverWrapperCtx:
  """Codegen kit handed to a plugin's ``render_wrapper`` hook.

  ``symbol`` is the solver's mangled C identifier (prefix for any static the
  template declares), ``raw_symbol`` the function the template must define,
  ``stats_symbol`` the ``scaly_solver_stats`` static it must fill (declared by
  core, one per solver). ``raw_symbol_of`` resolves the C symbol of an oracle
  / derivative ConcreteFunction from the descriptor.
  """

  symbol: str
  raw_symbol: str
  stats_symbol: str
  function_symbols: Mapping[ConcreteFunction, str] = field(default_factory=dict)

  def raw_symbol_of(self, fun: ConcreteFunction | ExternalOracle) -> str:
    from scaly.solvers.model import ExternalOracle

    return fun.raw_symbol if isinstance(fun, ExternalOracle) else f"{self.function_symbols.get(fun, c_ident(fun.name))}_raw"

@dataclass(frozen=True)
class StageLayout:
  """State/control sizes and original constraint rows for stages ``0..N``.

  Variables are ordered ``[x_0, u_0, ..., x_N, u_N]``. ``dynamics[k]`` contains
  the ``nx[k+1]`` equality rows into stage ``k+1``, in next-state order.
  ``eq[k]`` and ``ineq[k]`` contain the remaining constraints of stage ``k``.
  All row indices refer to the stacked ``[equalities; inequalities]`` vector.
  The number of local constraints is the length of the corresponding tuple.
  Next-state Jacobian blocks are structurally diagonal, without a guarantee
  about their coefficients.
  """

  nx: tuple[int, ...]
  nu: tuple[int, ...]
  dynamics: tuple[tuple[int, ...], ...]
  eq: tuple[tuple[int, ...], ...]
  ineq: tuple[tuple[int, ...], ...]

  @property
  def N(self) -> int:
    """The number of transitions, one less than the number of stages."""
    return len(self.nx) - 1

@dataclass(frozen=True)
class StageOracles:
  """Mapped dense derivatives in recorded stage order, with row-major stage blocks.

  ``grad`` returns the stacked objective gradient. ``jac`` returns the dynamics,
  initial, path and terminal Jacobians, each in local ``[x, u]`` order.
  ``hess`` takes the objective scale and stacked constraint multipliers after the
  ordinary inputs. It returns stage and terminal Lagrangian blocks, each a dense
  Hessian followed by its gradient row. A stage excludes the incoming transition's
  linear ``-xnext`` term. Box bounds are handled by the plugin.
  """

  grad: ConcreteFunction
  jac: ConcreteFunction
  hess: ConcreteFunction

@dataclass(frozen=True, slots=True)
class ExternalOracle:
  """A C-ABI oracle supplied by a plugin consumer instead of an Scaly graph.

  ``source`` must define ``raw_symbol`` with the same flat-buffer convention
  as generated Scaly kernels: one ``const double*`` per input, one ``double*``
  per output, and a trailing ``double*`` workspace argument. ``workspace_size``
  declares the number of doubles available through that final argument.
  """

  name: str
  raw_symbol: str
  source: str
  input_signature: tuple[tuple[str, tuple[int, ...]], ...]
  output_signature: tuple[tuple[str, tuple[int, ...]], ...]
  workspace_size: int = 0

  def __post_init__(self) -> None:
    if self.workspace_size < 0:
      raise ValueError(f"external oracle workspace_size must be nonnegative, got {self.workspace_size}")

@dataclass(frozen=True, slots=True)
class SolverStats:
  """The statistics of one numerical solve, as returned by ``Solver.stats()``.

  The fields mirror the generated ``scaly_solver_stats`` struct, whose layout and meaning are
  described on the generated-interface page of *How it works*. Diagnostics
  a backend does not report are zero.
  """

  version: int
  status: ScalySolveStatus
  native_status: int
  iter: int
  obj: float
  t_total: float
  t_fe: float
  t_solver: float
  t_qp: float
  t_globalization: float
  t_glue: float
  n_eval_f: int
  n_eval_grad_f: int
  n_eval_g: int
  n_eval_jac_g: int
  n_eval_h: int
  _pad0: int = 0
  primal_viol: float = 0.0
  step_inf: float = 0.0
  alpha: float = 0.0
  merit_penalty: float = 0.0
  backtracks: int = 0
  qp_iter: int = 0

  @classmethod
  def from_c(cls, value: CSolverStats) -> SolverStats:
    return cls(**{name: ScalySolveStatus(raw) if name == "status" else raw for name, _ in STATS_FIELDS if (raw := getattr(value, name)) is not None})

  def to_solver_status(self) -> SolverStatus:
    """Summarize these statistics. ``ok`` holds for ``OK`` and ``ACCEPTABLE``."""
    counts = {name: getattr(self, name) for name, _ in STATS_FIELDS if name.startswith("n_eval_")}
    return SolverStatus(
      code=int(self.status),
      name=self.status.name,
      iter=self.iter,
      stats=counts,
      _ok=self.status in (ScalySolveStatus.OK, ScalySolveStatus.ACCEPTABLE),
    )

@dataclass(frozen=True, slots=True)
class SolverStatus:
  """A short solve summary from ``SolverStats.to_solver_status()``: the status, iterations and oracle evaluation counts."""

  code: int
  name: str
  iter: int = 0
  stats: dict[str, int] | None = None
  _ok: bool | None = None

  @property
  def ok(self) -> bool:
    # code is the scaly status enum: OK == 0, ACCEPTABLE == 1.
    return self.code in (0, 1) if self._ok is None else self._ok

class ScalySolveStatus(enum.IntEnum):
  """Backend-neutral solve outcome, as reported in the generated statistics struct.

  Each backend maps its own native status onto these, so calling code does not have to know
  which solver ran. ``OK`` and ``ACCEPTABLE`` are the successful ones.
  """

  OK = 0
  ACCEPTABLE = 1
  MAX_ITER = 2
  PRIMAL_INFEASIBLE = 3
  DUAL_INFEASIBLE = 4
  NUMERICS = 5
  USER_STOP = 6
  ERROR = 7