Optimization¶

@runtime_checkable
class BaseProblem(HasDimensions, HasCost, HasBounds, HasEqConstraints, HasIneqConstraints, Protocol):
  """Dense problem: cost + gradient + bounds + dense constraints."""

  ...

@runtime_checkable
class SparseBaseProblem(HasDimensions, HasCost, HasBounds, HasSparseEqConstraints, HasSparseIneqConstraints, Protocol):
  """Sparse problem: cost + gradient + bounds + sparse constraints."""

  ...

@runtime_checkable
class DenseSQPProblem(BaseProblem, HasLagrangianGrad, HasHessian, Protocol):
  """Dense problem with all fields the dense SQP solver may access."""

  ...

@runtime_checkable
class SparseSQPProblem(SparseBaseProblem, HasSparseLagrangianGrad, HasSparseHessian, Protocol):
  """Sparse problem with all fields the sparse SQP solver may access."""

  ...

@dataclass(frozen=True, kw_only=True)
class DenseProblem:
  """Nonlinear optimization problem with dense Jacobians/Hessians.

  Formulation::

      min_{x}   cost_fn(x, p)
      s.t.      eq_constraints_fn(x, p) = eq_rhs
                ineq_lb <= ineq_constraints_fn(x, p) <= ineq_ub
                x_lb <= x <= x_ub

  All functions take ``(x, p)`` where ``x`` is the decision variable (shape ``(n,)``)
  and ``p`` is a parameter vector (shape ``(nparam,)``). Derivative functions
  (gradient, Hessian, Jacobians, Lagrangian) are computed lazily on first access.

  Args:
    cost_fn: Scalar cost function ``(n, nparam) -> scalar``.
    eq_constraints_fn: Equality constraints ``(n, nparam) -> (m_eq,)``, or None.
    ineq_constraints_fn: Inequality constraints ``(n, nparam) -> (m_ineq,)``, or None.
    x_lb: Lower bounds on ``x``.
    x_ub: Upper bounds on ``x``.
    eq_rhs: Right-hand side of equality constraints.
    ineq_lb: Lower bounds on inequality constraints.
    ineq_ub: Upper bounds on inequality constraints.
  """

  cost_fn: NumericalFunction
  eq_constraints_fn: NumericalFunction | None = None
  ineq_constraints_fn: NumericalFunction | None = None
  x_lb: Tensor
  x_ub: Tensor
  eq_rhs: Tensor | None = None
  ineq_lb: Tensor | None = None
  ineq_ub: Tensor | None = None

  def __post_init__(self):
    assert self.x_lb.shape == self.x_ub.shape == (self.n,)
    assert (self.x_lb <= self.x_ub).all().item()
    _validate_signature(self.cost_fn, ((self.n,), (self.nparam,)), None, "cost_fn")
    if self.eq_constraints_fn is not None:
      assert self.eq_rhs is not None and self.eq_rhs.shape == (self.m_eq,)
      _validate_signature(self.eq_constraints_fn, ((self.n,), (self.nparam,)), (self.m_eq,), "eq_constraints_fn")
    if self.ineq_constraints_fn is not None:
      assert self.ineq_lb is not None and self.ineq_ub is not None
      assert self.ineq_lb.shape == self.ineq_ub.shape == (self.m_ineq,)
      assert (self.ineq_lb <= self.ineq_ub).all().item()
      _validate_signature(self.ineq_constraints_fn, ((self.n,), (self.nparam,)), (self.m_ineq,), "ineq_constraints_fn")

  @property
  def n(self) -> int:
    return make_more(self.cost_fn.inputs[0].shape)[0]

  @property
  def nparam(self) -> int:
    return make_more(self.cost_fn.inputs[-1].shape)[0]

  @property
  def m_eq(self) -> int:
    return make_more(self.eq_constraints_fn.outputs[0].shape)[0] if self.eq_constraints_fn is not None else 0

  @property
  def m_ineq(self) -> int:
    return make_more(self.ineq_constraints_fn.outputs[0].shape)[0] if self.ineq_constraints_fn is not None else 0

  @cached_property
  def cost_grad_fn(self) -> NumericalFunction:
    return gradient(self.cost_fn, argnum=0)

  @cached_property
  def cost_hessian_fn(self) -> NumericalFunction:
    return hessian(self.cost_fn, argnum=0)

  @cached_property
  def eq_constraints_jac_fn(self) -> NumericalFunction | None:
    return jacobian(self.eq_constraints_fn, argnum=0) if self.eq_constraints_fn is not None else None

  @cached_property
  def ineq_constraints_jac_fn(self) -> NumericalFunction | None:
    return jacobian(self.ineq_constraints_fn, argnum=0) if self.ineq_constraints_fn is not None else None

  @cached_property
  def lagrangian_grad_fn(self) -> NumericalFunction:
    eq_fn, ineq_fn = self.eq_constraints_fn, self.ineq_constraints_fn
    cost_grad_fn = self.cost_grad_fn

    def lagrangian_grad(x: Tensor, lam_bounds: Tensor, lam_eq: Tensor, lam_ineq: Tensor, p: Tensor) -> Tensor:
      grad = cost_grad_fn.fn(x, p) + lam_bounds
      if eq_fn is not None:
        grad = grad + vjp((eq_fn.fn(x, p),), (x,), (lam_eq,))[0]
      if ineq_fn is not None:
        grad = grad + vjp((ineq_fn.fn(x, p),), (x,), (lam_ineq,))[0]
      return grad

    return NumericalFunction(  # ty: ignore[no-matching-overload]
      "lagrangian_grad",
      lagrangian_grad,
      (Arg(self.n), Arg(self.n), Arg(self.m_eq), Arg(self.m_ineq), Arg(self.nparam)),
    )

  @cached_property
  def lagrangian_hessian_fn(self) -> NumericalFunction:
    @numerical_function("lagrangian_hessian", (self.n, self.n, self.m_eq, self.m_ineq, self.nparam))
    def custom_hessian(x: Tensor, lam_bounds: Tensor, lam_eq: Tensor, lam_ineq: Tensor, p: Tensor) -> Tensor:
      del lam_bounds
      custom_grad = self.cost_grad_fn.fn(x, p)
      if self.eq_constraints_fn is not None:
        custom_grad = custom_grad + vjp((self.eq_constraints_fn.fn(x, p),), (x,), (lam_eq,))[0]
      if self.ineq_constraints_fn is not None:
        custom_grad = custom_grad + vjp((self.ineq_constraints_fn.fn(x, p),), (x,), (lam_ineq,))[0]
      return dvmap(lambda u: jvp((custom_grad,), (x,), (u,))[0], in_axes=1, out_axis=1)(Tensor.eye(self.n))

    return custom_hessian

@dataclass(frozen=True)
class DenseSQPFunction(FunctionBase):
  """SQP solver for dense nonlinear programs.

  Wraps a ``DenseSQPProblem`` and ``SQPSettings``, using PIQP as the QP backend.
  Callable as ``solver(x0, p) -> SQPResult``. Also supports AOT code generation
  via ``generate_module``.

  Args:
    problem: The optimization problem definition.
    settings: Solver configuration.
  """

  problem: DenseSQPProblem
  settings: SQPSettings

  def __post_init__(self) -> None:
    # Validate settings
    if not self.settings.valid:
      raise ValueError(
        f"Invalid SQPSettings: {self.settings}\n"
        f"  - 0 < tau < 1: {0.0 < self.settings.tau < 1.0}\n"
        f"  - 0 < eta < 1: {0.0 < self.settings.eta < 1.0}\n"
        f"  - 0 < rho < 1: {0.0 < self.settings.rho < 1.0}\n"
        f"  - K > 0: {self.settings.K > 0}\n"
        f"  - eps_prim > 0: {self.settings.eps_prim > 0.0}\n"
        f"  - eps_dual > 0: {self.settings.eps_dual > 0.0}\n"
        f"  - 0 < min_alpha <= 1: {0.0 < self.settings.min_alpha <= 1.0}\n"
        f"  - max_iter > 0: {self.settings.max_iter > 0}\n"
        f"  - line_search_max_iter > 0: {self.settings.line_search_max_iter > 0}"
      )

    # Validate hessian approximation
    # NOTE: check the approximation setting first (cheap) to avoid triggering
    # expensive cached_property computation of the Hessian function that isn't needed.
    if self.settings.hessian_approximation == HessianApproximation.EXACT_NO_CONSTRAINTS and self.problem.cost_hessian_fn is None:
      raise ValueError("HessianApproximation.EXACT_NO_CONSTRAINTS requires cost_hessian_fn in problem")

    if self.settings.hessian_approximation == HessianApproximation.EXACT and self.problem.lagrangian_hessian_fn is None:
      raise ValueError("HessianApproximation.EXACT requires lagrangian_hessian_fn in problem")

  ##################################################################
  # problem dimensions as properties (for convenience)
  ##################################################################

  @property
  def n(self) -> int:
    return int(self.problem.n)

  @property
  def m_eq(self) -> int:
    return int(self.problem.m_eq)

  @property
  def m_ineq(self) -> int:
    return int(self.problem.m_ineq)

  @property
  def nparam(self) -> int:
    return self.problem.nparam

  @property
  def is_sparse(self) -> bool:
    return False

  @cached_property
  def hessian_fn(self) -> NumericalFunction:
    if self.settings.hessian_approximation == HessianApproximation.EXACT:
      assert self.problem.lagrangian_hessian_fn is not None
      return self.problem.lagrangian_hessian_fn
    assert self.problem.cost_hessian_fn is not None
    return self.problem.cost_hessian_fn

  ##################################################################
  # problem dimensions as CodegenIntConstants
  ##################################################################

  @cached_property
  def n_const(self) -> CodegenIntConstant:
    return CodegenIntConstant(f"{self.name}_n", self.n)

  @cached_property
  def m_eq_const(self) -> CodegenIntConstant:
    return CodegenIntConstant(f"{self.name}_m_eq", self.m_eq)

  @cached_property
  def m_ineq_const(self) -> CodegenIntConstant:
    return CodegenIntConstant(f"{self.name}_m_ineq", self.m_ineq)

  @cached_property
  def nparam_const(self) -> CodegenIntConstant:
    return CodegenIntConstant(f"{self.name}_nparam", self.nparam)

  ##################################################################
  # auxiliary NumericalFunctions for primal/dual updates
  ##################################################################

  @cached_property
  def primal_update_fn(self) -> NumericalFunction:
    def fn(x: Tensor, dx: Tensor, alpha: Tensor) -> Tensor:
      return x + alpha * dx

    return NumericalFunction(
      f"{self.name}_primal_update",
      fn,
      (Arg(self.n), Arg(self.n), Arg(())),
    )

  @cached_property
  def bounds_dual_update_fn(self) -> NumericalFunction:
    def fn(lam_bounds: Tensor, z_bu: Tensor, z_bl: Tensor, alpha: Tensor) -> Tensor:
      return lam_bounds + alpha * (z_bu - z_bl - lam_bounds)

    return NumericalFunction(
      f"{self.name}_bounds_dual_update",
      fn,
      (Arg(self.n), Arg(self.n), Arg(self.n), Arg(())),
    )

  @cached_property
  def eq_dual_update_fn(self) -> NumericalFunction | None:
    if self.m_eq == 0:
      return None

    def fn(lam_eq: Tensor, y: Tensor, alpha: Tensor) -> Tensor:
      return lam_eq + alpha * (y - lam_eq)

    return NumericalFunction(
      f"{self.name}_eq_dual_update",
      fn,
      (Arg(self.m_eq), Arg(self.m_eq), Arg(())),
    )

  @cached_property
  def ineq_dual_update_fn(self) -> NumericalFunction | None:
    if self.m_ineq == 0:
      return None

    def fn(lam_ineq: Tensor, z_u: Tensor, z_l: Tensor, alpha: Tensor) -> Tensor:
      return lam_ineq + alpha * (z_u - z_l - lam_ineq)

    return NumericalFunction(
      f"{self.name}_ineq_dual_update",
      fn,
      (Arg(self.m_ineq), Arg(self.m_ineq), Arg(self.m_ineq), Arg(())),
    )

  ##################################################################
  # auxiliary NumericalFunctions for violation computation
  ##################################################################

  @cached_property
  def constraints_violation_fn(self) -> NumericalFunction:
    m_eq = self.m_eq
    m_ineq = self.m_ineq

    def fn(
      x: Tensor,
      x_lb: Tensor,
      x_ub: Tensor,
      eq_val: Tensor,
      eq_rhs: Tensor,
      ineq_val: Tensor,
      ineq_lb: Tensor,
      ineq_ub: Tensor,
    ) -> Tensor:
      parts = [(x_lb - x).max(), (x - x_ub).max()]
      if m_eq > 0:
        parts.append((eq_rhs - eq_val).abs().max())
      if m_ineq > 0:
        parts.append((ineq_lb - ineq_val).max())
        parts.append((ineq_val - ineq_ub).max())
      return Tensor.max(Tensor.stack(*parts))

    return NumericalFunction(  # ty: ignore[no-matching-overload]
      f"{self.name}_constraints_violation",
      fn,
      (
        Arg(self.n),
        Arg(self.n),
        Arg(self.n),
        Arg(self.m_eq),
        Arg(self.m_eq),
        Arg(self.m_ineq),
        Arg(self.m_ineq),
        Arg(self.m_ineq),
      ),
    )

  @cached_property
  def lagrangian_grad_fn(self) -> NumericalFunction:
    return self.problem.lagrangian_grad_fn

  @cached_property
  def stationarity_violation_fn(self) -> NumericalFunction:
    def fn(grad_L: Tensor) -> Tensor:
      return grad_L.abs().max()

    return NumericalFunction(
      f"{self.name}_stationarity_violation",
      fn,
      (Arg(self.n),),
    )

  @cached_property
  def complementarity_violation_fn(self) -> NumericalFunction:
    m_eq = self.m_eq
    m_ineq = self.m_ineq

    def fn(
      x: Tensor,
      x_lb: Tensor,
      x_ub: Tensor,
      lam_bounds: Tensor,
      eq_val: Tensor,
      eq_rhs: Tensor,
      lam_eq: Tensor,
      ineq_val: Tensor,
      ineq_lb: Tensor,
      ineq_ub: Tensor,
      lam_ineq: Tensor,
    ) -> Tensor:
      parts = [
        (x_lb - x).maximum(0.0).dot(lam_bounds),
        (x - x_ub).maximum(0.0).dot(lam_bounds),
      ]
      if m_eq > 0:
        parts.append((eq_rhs - eq_val).abs().dot(lam_eq))
      if m_ineq > 0:
        parts.append((ineq_lb - ineq_val).maximum(0.0).dot(lam_ineq))
        parts.append((ineq_val - ineq_ub).maximum(0.0).dot(lam_ineq))
      return Tensor.max(Tensor.stack(*parts))

    return NumericalFunction(  # ty: ignore[no-matching-overload]
      f"{self.name}_complementarity_violation",
      fn,
      (
        Arg(self.n),
        Arg(self.n),
        Arg(self.n),
        Arg(self.n),
        Arg(self.m_eq),
        Arg(self.m_eq),
        Arg(self.m_eq),
        Arg(self.m_ineq),
        Arg(self.m_ineq),
        Arg(self.m_ineq),
        Arg(self.m_ineq),
      ),
    )

  ##################################################################
  # auxiliary NumericalFunctions for QP data preparation
  ##################################################################

  @cached_property
  def qp_bounds_fn(self) -> NumericalFunction:
    def fn(x: Tensor, x_lb: Tensor, x_ub: Tensor) -> tuple[Tensor, Tensor]:
      return x_lb - x, x_ub - x

    return NumericalFunction(
      f"{self.name}_qp_bounds",
      fn,
      (Arg(self.n), Arg(self.n), Arg(self.n)),
    )

  @cached_property
  def qp_eq_rhs_fn(self) -> NumericalFunction | None:
    if self.m_eq == 0:
      return None

    def fn(eq_rhs: Tensor, eq_val: Tensor) -> Tensor:
      return eq_rhs - eq_val

    return NumericalFunction(
      f"{self.name}_qp_eq_rhs",
      fn,
      (Arg(self.m_eq), Arg(self.m_eq)),
    )

  @cached_property
  def qp_ineq_bounds_fn(self) -> NumericalFunction | None:
    if self.m_ineq == 0:
      return None

    def fn(ineq_val: Tensor, ineq_lb: Tensor, ineq_ub: Tensor) -> tuple[Tensor, Tensor]:
      return ineq_lb - ineq_val, ineq_ub - ineq_val

    return NumericalFunction(
      f"{self.name}_qp_ineq_bounds",
      fn,
      (Arg(self.m_ineq), Arg(self.m_ineq), Arg(self.m_ineq)),
    )

  ##################################################################
  # codegen helpers: constant buffer names and argument strings
  ##################################################################

  @cached_property
  def x_lb_buf_name(self) -> str:
    return f"{self.name}_x_lb"

  @cached_property
  def x_ub_buf_name(self) -> str:
    return f"{self.name}_x_ub"

  @cached_property
  def eq_rhs_buf_name(self) -> str:
    return f"{self.name}_eq_rhs"

  @cached_property
  def ineq_lb_buf_name(self) -> str:
    return f"{self.name}_ineq_lb"

  @cached_property
  def ineq_ub_buf_name(self) -> str:
    return f"{self.name}_ineq_ub"

  @cached_property
  def p_buf_name(self) -> str:
    return f"{self.name}_p"

  # Constant buffer values for template rendering
  @cached_property
  def x_lb_values(self) -> str:
    return tensor_to_values_string(self.problem.x_lb)

  @cached_property
  def x_ub_values(self) -> str:
    return tensor_to_values_string(self.problem.x_ub)

  @cached_property
  def eq_rhs_values(self) -> str:
    if self.problem.eq_rhs is None:
      return ""
    return tensor_to_values_string(self.problem.eq_rhs)

  @cached_property
  def ineq_lb_values(self) -> str:
    if self.problem.ineq_lb is None:
      return ""
    return tensor_to_values_string(self.problem.ineq_lb)

  @cached_property
  def ineq_ub_values(self) -> str:
    if self.problem.ineq_ub is None:
      return ""
    return tensor_to_values_string(self.problem.ineq_ub)

  # Argument strings for template rendering
  @property
  def x_lb_arg(self) -> str:
    return self.x_lb_buf_name

  @property
  def x_ub_arg(self) -> str:
    return self.x_ub_buf_name

  @property
  def eq_val_arg(self) -> str:
    return "ws.eq_val"

  @property
  def eq_rhs_arg(self) -> str:
    return self.eq_rhs_buf_name

  @property
  def lam_eq_arg(self) -> str:
    return "ws.lam_eq"

  @property
  def lam_eq_for_hessian_arg(self) -> str:
    return "ws.lam_eq"

  @property
  def ineq_val_arg(self) -> str:
    return "ws.ineq_val"

  @property
  def ineq_lb_arg(self) -> str:
    return self.ineq_lb_buf_name

  @property
  def ineq_ub_arg(self) -> str:
    return self.ineq_ub_buf_name

  @property
  def lam_ineq_arg(self) -> str:
    return "ws.lam_ineq"

  @property
  def lam_ineq_for_hessian_arg(self) -> str:
    return "ws.lam_ineq"

  @property
  def p_arg(self) -> str:
    return "p" if int(self.nparam) > 0 else self.p_buf_name

  ##################################################################
  # JIT compilation
  ##################################################################

  @cached_property
  def jit_source(self) -> str:
    from anvil.codegen.jit import generate_jit_source

    return generate_jit_source(self.name, [self])

  @cached_property
  def _jit_module(self):
    from anvil.codegen.jit import compile_to_shared_lib, load_jit_module
    from anvil.piqp_paths import INCLUDE_DIR, LIB_DIR

    extra_flags = [f"-I{INCLUDE_DIR}", f"-L{LIB_DIR}", "-lpiqpc", f"-Wl,-rpath,{LIB_DIR}"]
    path = compile_to_shared_lib(self.jit_source, self.name, extra_flags=extra_flags)
    return load_jit_module(path)

  @cached_property
  def _jit_ws(self) -> ctypes.c_void_p:
    init_fn = getattr(self._jit_module.lib, f"jit_{self.name}_init_ws")
    init_fn.argtypes = []
    init_fn.restype = ctypes.c_void_p
    ws = init_fn()
    deinit_fn = getattr(self._jit_module.lib, f"jit_{self.name}_deinit_ws")
    deinit_fn.argtypes = [ctypes.c_void_p]
    deinit_fn.restype = None
    weakref.finalize(self, deinit_fn, ws)
    return ctypes.c_void_p(ws)

  @cached_property
  def _jit_call_fn(self):
    fn = getattr(self._jit_module.lib, f"jit_{self.name}_call")
    fn.argtypes = [
      ctypes.c_void_p,  # ws
      ctypes.c_int,
      ctypes.c_int,
      ctypes.c_int,  # max_iter, verbose, compute_timings
      ctypes.c_double,
      ctypes.c_double,  # eps_prim, eps_dual
      ctypes.c_int,  # globalization_strategy
      ctypes.c_int,  # line_search_max_iter
      ctypes.c_double,
      ctypes.c_double,
      ctypes.c_double,
      ctypes.c_double,
      ctypes.c_double,  # tau, eta, rho, K, min_alpha
      ctypes.c_void_p,
      ctypes.c_void_p,  # initial_x, initial_lam_bounds
      ctypes.c_void_p,
      ctypes.c_void_p,
      ctypes.c_void_p,  # lam_eq, lam_ineq, p
      ctypes.POINTER(ctypes.c_int),
      ctypes.POINTER(ctypes.c_int),
      ctypes.POINTER(ctypes.c_int),  # status, iter, qp_iter
      ctypes.c_void_p,  # timings
    ]
    fn.restype = None
    return fn

  @cached_property
  def _jit_get_x(self):
    fn = getattr(self._jit_module.lib, f"jit_{self.name}_get_x")
    fn.argtypes = [ctypes.c_void_p]
    fn.restype = ctypes.c_void_p
    return fn

  @cached_property
  def _jit_get_lam_bounds(self):
    fn = getattr(self._jit_module.lib, f"jit_{self.name}_get_lam_bounds")
    fn.argtypes = [ctypes.c_void_p]
    fn.restype = ctypes.c_void_p
    return fn

  @cached_property
  def _jit_get_lam_eq(self):
    if int(self.m_eq) == 0:
      return None
    fn = getattr(self._jit_module.lib, f"jit_{self.name}_get_lam_eq")
    fn.argtypes = [ctypes.c_void_p]
    fn.restype = ctypes.c_void_p
    return fn

  @cached_property
  def _jit_get_lam_ineq(self):
    if int(self.m_ineq) == 0:
      return None
    fn = getattr(self._jit_module.lib, f"jit_{self.name}_get_lam_ineq")
    fn.argtypes = [ctypes.c_void_p]
    fn.restype = ctypes.c_void_p
    return fn

  ##################################################################
  # python calling interface helpers
  ##################################################################

  def _constraint_buffers(self) -> tuple[Tensor, Tensor, Tensor]:
    eq_rhs = self.problem.eq_rhs if self.problem.eq_rhs is not None else Tensor.zeros(self.m_eq)
    ineq_lb = self.problem.ineq_lb if self.problem.ineq_lb is not None else Tensor.zeros(self.m_ineq)
    ineq_ub = self.problem.ineq_ub if self.problem.ineq_ub is not None else Tensor.zeros(self.m_ineq)
    return eq_rhs, ineq_lb, ineq_ub

  def _parameter_buffer(self, p: FloatArray | None) -> Tensor:
    if p is None:
      if int(self.nparam) > 0:
        raise ValueError(f"Problem expects parameter vector of shape ({self.nparam},), got None")
      return Tensor.zeros(0)
    if p.shape != (int(self.nparam),):
      raise ValueError(f"Expected parameter vector shape ({self.nparam},), got {p.shape}")
    return p if isinstance(p, Tensor) else Tensor(p)

  def _init_primal_dual(
    self,
    initial_x: FloatArray | None,
    initial_lam_bounds: FloatArray | None,
    initial_lam_eq: FloatArray | None,
    initial_lam_ineq: FloatArray | None,
  ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    _to_tensor = lambda v: v if isinstance(v, Tensor) else Tensor(v)
    x = Tensor.zeros(self.n).contiguous() if initial_x is None else _to_tensor(initial_x)
    lam_bounds = Tensor.zeros(self.n).contiguous() if initial_lam_bounds is None else _to_tensor(initial_lam_bounds)
    lam_eq = Tensor.zeros(self.m_eq).contiguous() if initial_lam_eq is None else _to_tensor(initial_lam_eq)
    lam_ineq = Tensor.zeros(self.m_ineq).contiguous() if initial_lam_ineq is None else _to_tensor(initial_lam_ineq)
    return x, lam_bounds, lam_eq, lam_ineq

  def _init_qp_solver(self) -> piqp.DenseSolver:
    solver = piqp.DenseSolver()
    if self.settings.qp_settings is not None:
      solver.settings = self.settings.qp_settings
    return solver

  def _setup_qp_solver(self, qp_solver: piqp.DenseSolver) -> None:
    qp_setup_kwargs: dict = {
      "P": Tensor.eye(self.n).numpy(),
      "c": Tensor.ones(self.n).numpy(),
      "x_l": self.problem.x_lb.numpy(),
      "x_u": self.problem.x_ub.numpy(),
    }
    if self.m_eq > 0:
      qp_setup_kwargs["A"] = np.zeros((self.m_eq, self.n))
      qp_setup_kwargs["b"] = np.zeros(self.m_eq)
    if self.m_ineq > 0:
      qp_setup_kwargs["G"] = np.zeros((self.m_ineq, self.n))
      qp_setup_kwargs["h_l"] = np.zeros(self.m_ineq)
      qp_setup_kwargs["h_u"] = np.zeros(self.m_ineq)
    qp_solver.setup(**qp_setup_kwargs)

  def _evaluate_linearization(self, x: Tensor, p: Tensor) -> dict[str, Tensor]:
    return {
      "cost_val": self.problem.cost_fn.fn(x, p),
      "cost_grad": self.problem.cost_grad_fn.fn(x, p),
      "eq_val": self.problem.eq_constraints_fn.fn(x, p) if self.problem.eq_constraints_fn is not None else Tensor.zeros(self.m_eq),
      "ineq_val": self.problem.ineq_constraints_fn.fn(x, p) if self.problem.ineq_constraints_fn is not None else Tensor.zeros(self.m_ineq),
      "eq_jac": self.problem.eq_constraints_jac_fn.fn(x, p) if self.problem.eq_constraints_jac_fn is not None else Tensor.zeros(self.m_eq, self.n),
      "ineq_jac": self.problem.ineq_constraints_jac_fn.fn(x, p)
      if self.problem.ineq_constraints_jac_fn is not None
      else Tensor.zeros(self.m_ineq, self.n),
    }

  def _violations(
    self,
    x: Tensor,
    p: Tensor,
    lam_bounds: Tensor,
    lam_eq: Tensor,
    lam_ineq: Tensor,
    eq_rhs: Tensor,
    ineq_lb: Tensor,
    ineq_ub: Tensor,
    linearization: dict[str, Tensor],
  ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
    lam_eq_for_fn = lam_eq
    lam_ineq_for_fn = lam_ineq
    constraints_violation = self.constraints_violation_fn.fn(
      x,
      self.problem.x_lb,
      self.problem.x_ub,
      linearization["eq_val"],
      eq_rhs,
      linearization["ineq_val"],
      ineq_lb,
      ineq_ub,
    )
    complementarity_violation = self.complementarity_violation_fn.fn(
      x,
      self.problem.x_lb,
      self.problem.x_ub,
      lam_bounds,
      linearization["eq_val"],
      eq_rhs,
      lam_eq_for_fn,
      linearization["ineq_val"],
      ineq_lb,
      ineq_ub,
      lam_ineq_for_fn,
    )
    lagrangian_grad = self.lagrangian_grad_fn.fn(x, lam_bounds, lam_eq_for_fn, lam_ineq_for_fn, p)
    stationarity_violation = self.stationarity_violation_fn.fn(lagrangian_grad)
    return constraints_violation, complementarity_violation, stationarity_violation, lagrangian_grad, lam_eq_for_fn, lam_ineq_for_fn

  def _compute_hessian(self, x: Tensor, lam_bounds: Tensor, lam_eq: Tensor, lam_ineq: Tensor, p: Tensor) -> Tensor:
    if self.settings.hessian_approximation == HessianApproximation.EXACT:
      return self.hessian_fn.fn(x, lam_bounds, lam_eq, lam_ineq, p)
    return self.hessian_fn.fn(x, p)

  def _build_qp_update(
    self,
    x: Tensor,
    eq_rhs: Tensor,
    ineq_lb: Tensor,
    ineq_ub: Tensor,
    linearization: dict[str, Tensor],
    hessian: Tensor,
  ) -> dict:
    qp_x_l, qp_x_u = self.qp_bounds_fn.fn(x, self.problem.x_lb, self.problem.x_ub)
    qp_update_kwargs: dict = {
      "P": hessian.numpy(),
      "c": linearization["cost_grad"].numpy(),
      "x_l": qp_x_l.numpy(),
      "x_u": qp_x_u.numpy(),
    }
    if self.m_ineq > 0 and self.qp_ineq_bounds_fn is not None:
      qp_h_l, qp_h_u = self.qp_ineq_bounds_fn.fn(linearization["ineq_val"], ineq_lb, ineq_ub)
      qp_update_kwargs["G"] = linearization["ineq_jac"].numpy()
      qp_update_kwargs["h_l"] = qp_h_l.numpy()
      qp_update_kwargs["h_u"] = qp_h_u.numpy()
    if self.m_eq > 0 and self.qp_eq_rhs_fn is not None:
      qp_b = self.qp_eq_rhs_fn.fn(eq_rhs, linearization["eq_val"])
      qp_update_kwargs["A"] = linearization["eq_jac"].numpy()
      qp_update_kwargs["b"] = qp_b.numpy()
    return qp_update_kwargs

  def _build_qp_update_from_prep(
    self,
    linearization: dict[str, Tensor],
    hessian: Tensor,
    prep: tuple[Tensor, ...],
  ) -> dict:
    """Build QP update dict from pre-computed (JIT'd) prep tensors.

    prep layout: [qp_x_l, qp_x_u, (qp_b)?, (qp_h_l, qp_h_u)?]
    """
    qp_update_kwargs: dict = {
      "P": hessian.numpy(),
      "c": linearization["cost_grad"].numpy(),
      "x_l": prep[0].numpy(),
      "x_u": prep[1].numpy(),
    }
    idx = 2
    if self.m_eq > 0 and self.qp_eq_rhs_fn is not None:
      qp_update_kwargs["A"] = linearization["eq_jac"].numpy()
      qp_update_kwargs["b"] = prep[idx].numpy()
      idx += 1
    if self.m_ineq > 0 and self.qp_ineq_bounds_fn is not None:
      qp_update_kwargs["G"] = linearization["ineq_jac"].numpy()
      qp_update_kwargs["h_l"] = prep[idx].numpy()
      qp_update_kwargs["h_u"] = prep[idx + 1].numpy()
    return qp_update_kwargs

  def _update_primal_dual(
    self,
    qp_solver: piqp.DenseSolver,
    x: Tensor,
    lam_bounds: Tensor,
    lam_eq: Tensor,
    lam_ineq: Tensor,
    alpha: float,
  ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    dx = Tensor(qp_solver.result.x.copy())
    alpha_tensor = Tensor([alpha])
    x = self.primal_update_fn.fn(x, dx, alpha_tensor)
    z_bu = Tensor(qp_solver.result.z_bu.copy())
    z_bl = Tensor(qp_solver.result.z_bl.copy())
    lam_bounds = self.bounds_dual_update_fn.fn(lam_bounds, z_bu, z_bl, alpha_tensor)
    if self.m_eq > 0 and self.eq_dual_update_fn is not None:
      y = Tensor(qp_solver.result.y.copy())
      lam_eq = self.eq_dual_update_fn.fn(lam_eq, y, alpha_tensor)
    if self.m_ineq > 0 and self.ineq_dual_update_fn is not None:
      z_u = Tensor(qp_solver.result.z_u.copy())
      z_l = Tensor(qp_solver.result.z_l.copy())
      lam_ineq = self.ineq_dual_update_fn.fn(lam_ineq, z_u, z_l, alpha_tensor)
    Tensor.realize(x, lam_bounds, lam_eq, lam_ineq)
    return x, lam_bounds, lam_eq, lam_ineq

  def _realize_tensors(
    self,
    linearization: dict[str, Tensor],
    constraints_violation: Tensor,
    complementarity_violation: Tensor,
    stationarity_violation: Tensor,
    x: Tensor,
    lam_bounds: Tensor,
    lam_eq: Tensor,
    lam_ineq: Tensor,
  ) -> None:
    tensors_to_realize = [
      linearization["cost_val"],
      linearization["cost_grad"],
      constraints_violation,
      complementarity_violation,
      stationarity_violation,
      x,
      lam_bounds,
      lam_eq,
      lam_ineq,
      linearization["eq_val"],
      linearization["ineq_val"],
    ]
    for key in ("eq_jac", "ineq_jac"):
      value = linearization.get(key)
      if value is not None:
        tensors_to_realize.append(value)
    Tensor.realize(*tensors_to_realize)

  ##################################################################
  # python calling interface (__call__ implements the SQP loop)
  ##################################################################

  @override
  def __call__(
    self,
    initial_x: FloatArray | None = None,
    initial_lam_bounds: FloatArray | None = None,
    initial_lam_eq: FloatArray | None = None,
    initial_lam_ineq: FloatArray | None = None,
    p: FloatArray | None = None,
  ) -> SQPResult:
    from anvil.codegen.jit import SQP_TIMING_FIELDS

    if not self.settings.valid:
      return SQPResult(status=SQPStatus.INVALID_SETTINGS)

    n = int(self.n)
    CPTR = ctypes.c_void_p

    # prepare inputs
    x = np.zeros(n, dtype=np.float64) if initial_x is None else np.ascontiguousarray(initial_x, dtype=np.float64)
    lam_bounds = np.zeros(n, dtype=np.float64) if initial_lam_bounds is None else np.ascontiguousarray(initial_lam_bounds, dtype=np.float64)
    lam_eq = np.zeros(int(self.m_eq), dtype=np.float64) if initial_lam_eq is None else np.ascontiguousarray(initial_lam_eq, dtype=np.float64)
    lam_ineq = np.zeros(int(self.m_ineq), dtype=np.float64) if initial_lam_ineq is None else np.ascontiguousarray(initial_lam_ineq, dtype=np.float64)
    if p is None:
      if int(self.nparam) > 0:
        raise ValueError(f"Problem expects parameter vector of shape ({self.nparam},), got None")
      p_arr = np.zeros(0, dtype=np.float64)
    else:
      p_np = p.numpy() if isinstance(p, Tensor) else np.asarray(p)
      if p_np.shape != (int(self.nparam),):
        raise ValueError(f"Expected parameter vector shape ({self.nparam},), got {p_np.shape}")
      p_arr = np.ascontiguousarray(p_np, dtype=np.float64)

    # output scalars
    status = ctypes.c_int(0)
    iter_out = ctypes.c_int(0)
    qp_iter_out = ctypes.c_int(0)

    # timings
    timed = self.settings.compute_timings
    timings_arr = np.zeros(len(SQP_TIMING_FIELDS), dtype=np.float64) if timed else None

    # call JIT function
    self._jit_call_fn(
      self._jit_ws,
      ctypes.c_int(self.settings.max_iter),
      ctypes.c_int(int(self.settings.verbose)),
      ctypes.c_int(int(timed)),
      ctypes.c_double(self.settings.eps_prim),
      ctypes.c_double(self.settings.eps_dual),
      ctypes.c_int(self.settings.globalization_strategy.value),
      ctypes.c_int(self.settings.line_search_max_iter),
      ctypes.c_double(self.settings.tau),
      ctypes.c_double(self.settings.eta),
      ctypes.c_double(self.settings.rho),
      ctypes.c_double(self.settings.K),
      ctypes.c_double(self.settings.min_alpha),
      x.ctypes.data_as(CPTR),
      lam_bounds.ctypes.data_as(CPTR),
      lam_eq.ctypes.data_as(CPTR) if int(self.m_eq) > 0 else CPTR(),
      lam_ineq.ctypes.data_as(CPTR) if int(self.m_ineq) > 0 else CPTR(),
      p_arr.ctypes.data_as(CPTR) if int(self.nparam) > 0 else CPTR(),
      ctypes.byref(status),
      ctypes.byref(iter_out),
      ctypes.byref(qp_iter_out),
      timings_arr.ctypes.data_as(CPTR) if timed else CPTR(),  # ty: ignore[unresolved-attribute]
    )

    # extract results from workspace
    result_x = np.ctypeslib.as_array(ctypes.cast(self._jit_get_x(self._jit_ws), ctypes.POINTER(ctypes.c_double)), shape=(n,)).copy()
    result_lam_bounds = np.ctypeslib.as_array(ctypes.cast(self._jit_get_lam_bounds(self._jit_ws), ctypes.POINTER(ctypes.c_double)), shape=(n,)).copy()
    result_lam_eq = None
    if int(self.m_eq) > 0 and self._jit_get_lam_eq is not None:
      result_lam_eq = np.ctypeslib.as_array(
        ctypes.cast(self._jit_get_lam_eq(self._jit_ws), ctypes.POINTER(ctypes.c_double)), shape=(int(self.m_eq),)
      ).copy()
    result_lam_ineq = None
    if int(self.m_ineq) > 0 and self._jit_get_lam_ineq is not None:
      result_lam_ineq = np.ctypeslib.as_array(
        ctypes.cast(self._jit_get_lam_ineq(self._jit_ws), ctypes.POINTER(ctypes.c_double)), shape=(int(self.m_ineq),)
      ).copy()

    timings = None
    if timed and timings_arr is not None:
      timings = SQPTimings(**dict(zip(SQP_TIMING_FIELDS, timings_arr.tolist())))

    return SQPResult(
      status=SQPStatus(status.value),
      iter=iter_out.value,
      qp_iter=qp_iter_out.value,
      x=result_x,
      lam_bounds=result_lam_bounds,
      lam_eq=result_lam_eq,
      lam_ineq=result_lam_ineq,
      timings=timings,
    )

  ##################################################################
  # FunctionBase methods and properties to implement
  ##################################################################

  @override
  def dependencies(self) -> list[FunctionBase]:
    hessian_deps = (
      [self.problem.lagrangian_hessian_fn]
      if self.settings.hessian_approximation == HessianApproximation.EXACT and self.problem.lagrangian_hessian_fn is not None
      else [self.problem.cost_hessian_fn]
      if self.problem.cost_hessian_fn is not None
      else []
    )
    return [
      # Problem functions
      self.problem.cost_fn,
      self.problem.cost_grad_fn,
      *hessian_deps,
      *([self.problem.eq_constraints_fn] if self.problem.eq_constraints_fn is not None else []),
      *([self.problem.eq_constraints_jac_fn] if self.problem.eq_constraints_jac_fn is not None else []),
      *([self.problem.ineq_constraints_fn] if self.problem.ineq_constraints_fn is not None else []),
      *([self.problem.ineq_constraints_jac_fn] if self.problem.ineq_constraints_jac_fn is not None else []),
      # Auxiliary functions (always present)
      self.primal_update_fn,
      self.bounds_dual_update_fn,
      self.constraints_violation_fn,
      self.lagrangian_grad_fn,
      self.stationarity_violation_fn,
      self.complementarity_violation_fn,
      self.qp_bounds_fn,
      # Optional auxiliary functions
      *([self.eq_dual_update_fn] if self.eq_dual_update_fn is not None else []),
      *([self.ineq_dual_update_fn] if self.ineq_dual_update_fn is not None else []),
      *([self.qp_eq_rhs_fn] if self.qp_eq_rhs_fn is not None else []),
      *([self.qp_ineq_bounds_fn] if self.qp_ineq_bounds_fn is not None else []),
    ]

  @cached_property
  @override
  def codegen_constants(self) -> set[CodegenIntConstant]:
    return {self.n_const, self.nparam_const, self.m_eq_const, self.m_ineq_const}

  @cached_property
  @override
  def header_includes(self) -> set[str]:
    return {'#include "piqp.h"', "#include <cmath>"}

  @cached_property
  @override
  def header_code(self) -> str:
    return JINJA_ENV.get_template("sqp_function_header.j2").render(fn=self)

  @cached_property
  @override
  def source_includes(self) -> set[str]:
    return {"#include <cstdlib>", "#include <cstring>", "#include <utility>", "#include <chrono>"}

  @cached_property
  @override
  def source_code(self) -> str:
    return JINJA_ENV.get_template("sqp_function_source.j2").render(fn=self)