Require strict SciMLTesting 2.4 QA

Project.toml

@@ -1,13 +1,13 @@
 name = "ComplementaritySolve"
 uuid = "b40a91a3-bdaf-4e1c-b965-8c278a33a8d3"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "1.0.0"
+version = "2.0.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
-FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -18,6 +18,8 @@ PATHSolver = "f5f7c340-0bb3-5c69-969a-41884d311d1b"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+SciMLPublic = "431bcebd-1456-4ced-9d72-93c2757fff0b"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
@@ -26,16 +28,20 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ChainRulesCore = "1"
 CommonSolve = "0.2.6"
 ConcreteStructs = "0.2"
-FillArrays = "1.13"
+DifferentiationInterface = "0.6, 0.7"
 ForwardDiff = "1"
 GPUArraysCore = "0.2"
+LinearAlgebra = "1.10"
 LinearSolve = "3.61, 4.2, 5.0"
 NNlib = "0.9.34"
 NonlinearSolve = "4.12"
 PATHSolver = "1.4"
 Polyester = "0.7.18"
 PrecompileTools = "1.2.1"
 SciMLBase = "2.147, 3.1"
+SciMLOperators = "1.3"
+SciMLPublic = "1"
 SimpleNonlinearSolve = "2.7"
+SparseArrays = "1.10"
 Zygote = "0.7.5"
 julia = "1.10"

docs/make.jl

@@ -3,7 +3,7 @@ using Documenter
 
 makedocs(;
     modules = [ComplementaritySolve],
-    checkdocs = :exports,
+    checkdocs = :public,
     sitename = "ComplementaritySolve.jl",
     format = Documenter.HTML(;
         prettyurls = get(ENV, "CI", "false") == "true",
@@ -12,6 +12,7 @@ makedocs(;
     ),
     pages = [
         "Home" => "index.md",
+        "Developer API" => "developer_api.md",
     ],
 )
 

docs/src/developer_api.md

@@ -0,0 +1,20 @@
+# Developer API
+
+The qualified APIs on this page are intended for packages implementing
+ComplementaritySolve-compatible problem containers, algorithms, or sensitivity rules.
+Application code should use the exported problem constructors and algorithms instead.
+
+```@docs
+ComplementaritySolve.AbstractComplementarityProblem
+ComplementaritySolve.AbstractLinearComplementarityProblem
+ComplementaritySolve.AbstractNonlinearComplementarityProblem
+ComplementaritySolve.AbstractComplementaritySystem
+ComplementaritySolve.AbstractComplementarityAlgorithm
+ComplementaritySolve.AbstractComplementaritySystemAlgorithm
+ComplementaritySolve.AbstractComplementaritySensitivityAlgorithm
+ComplementaritySolve.AbstractComplementaritySolution
+ComplementaritySolve.AbstractLinearComplementaritySolution
+ComplementaritySolve.isbatched
+ComplementaritySolve.__solve
+ComplementaritySolve.__solve_adjoint
+```

src/ComplementaritySolve.jl

@@ -7,17 +7,20 @@ module ComplementaritySolve
 ## Core / QOL Dependencies
 using ChainRulesCore: ChainRulesCore, NoTangent, ZeroTangent
 using GPUArraysCore: GPUArraysCore
-using SciMLBase: SciMLBase, FunctionOperator, LinearProblem, NonlinearFunction,
+using SciMLBase: SciMLBase, LinearProblem, NonlinearFunction,
     NonlinearProblem, ODEFunction, ODEProblem, ReturnCode,
     SteadyStateProblem, isinplace
+using SciMLPublic: @public
 using CommonSolve: CommonSolve
 using ConcreteStructs: ConcreteStructs, @concrete
+using DifferentiationInterface: DifferentiationInterface, AutoForwardDiff, AutoZygote,
+    jacobian
 ## Stdlibs
 using LinearAlgebra: LinearAlgebra, Diagonal, I, diagind, mul!, norm, pinv, \, /
 using SparseArrays: SparseArrays
 ## SciML Dependencies
 using LinearSolve: LinearSolve
-## SciMLOperators is used transitively via SciMLBase (FunctionOperator)
+using SciMLOperators: FunctionOperator
 using SimpleNonlinearSolve: SimpleNonlinearSolve, SimpleNewtonRaphson
 using NonlinearSolve: NonlinearSolve
 ## AD Packages (for sensitivities & PATHSolver; move to extensions)
@@ -31,7 +34,6 @@ using Polyester: Polyester, @batch
 
 import CommonSolve: init, solve, solve!
 import ChainRulesCore as CRC
-import FillArrays: AbstractFill
 
 const ∂0 = ZeroTangent()
 const ∂∅ = NoTangent()
@@ -54,7 +56,7 @@ intended to be passed as the second argument to `solve(prob, alg; kwargs...)` fo
 
 # Interface
 
-- Implement an internal `__solve(prob, alg, args...; kwargs...)` method for each
+- Implement `ComplementaritySolve.__solve(prob, alg, args...; kwargs...)` for each
     supported problem family.
 - Return an `AbstractComplementaritySolution` subtype whose `prob` and `alg` fields
     reference the original problem and algorithm.
@@ -73,7 +75,7 @@ Developer interface for algorithms that solve complementarity systems.
 
 # Interface
 
-- Implement `solve(prob::AbstractComplementaritySystem, alg; kwargs...)`.
+- Implement `ComplementaritySolve.__solve(prob, alg; kwargs...)`.
 - Return the solution object produced by the wrapped ODE or steady-state solve.
 - Forward relevant solver keywords to the continuous dynamics solve and the embedded
     complementarity solve.
@@ -127,7 +129,13 @@ export PATHSolverAlgorithm
 export NaiveLCSAlgorithm
 export LinearComplementarityAdjoint, MixedComplementarityAdjoint
 export LinearComplementaritySolution, MixedComplementaritySolution
-export solve
+
+@public AbstractComplementarityAlgorithm, AbstractComplementaritySystemAlgorithm
+@public AbstractComplementaritySensitivityAlgorithm
+@public AbstractComplementarityProblem, AbstractLinearComplementarityProblem
+@public AbstractNonlinearComplementarityProblem, AbstractComplementaritySystem
+@public AbstractComplementaritySolution, AbstractLinearComplementaritySolution
+@public isbatched, __solve, __solve_adjoint
 
 include("precompilation.jl")
 

src/algorithms/generic.jl

@@ -19,7 +19,7 @@ nonlinear residual and delegating to a NonlinearSolve.jl-compatible solver.
 using ComplementaritySolve
 
 prob = LinearComplementarityProblem([2.0 -1.0; -1.0 2.0], [-1.0, -1.0])
-sol = solve(prob, NonlinearReformulation(:smooth))
+sol = ComplementaritySolve.solve(prob, NonlinearReformulation(:smooth))
 ```
 """
 @concrete struct NonlinearReformulation{method} <: AbstractComplementarityAlgorithm

src/algorithms/lcs/naive_lcs.jl

@@ -26,7 +26,7 @@ right-hand-side evaluation.
     lcp_solver
 end
 
-function solve(
+function __solve(
         prob::LinearComplementaritySystem{sstate}, alg::NaiveLCSAlgorithm;
         ode_kwargs = (;), lcp_kwargs = (;), kwargs...
     ) where {sstate}

src/algorithms/mcp/pathsolver.jl

@@ -46,9 +46,10 @@ function __solve(
 
     function J!(n, nnz, z, col, len, row, data)
         if !iip
-            J = (n ≤ 100 ? ForwardDiff.jacobian : Zygote.jacobian)(fₚ, z)
+            backend = n ≤ 100 ? AutoForwardDiff() : AutoZygote()
+            J = jacobian(fₚ, backend, z)
         else
-            J = ForwardDiff.jacobian(fₚ, similar(z, n), z)
+            J = jacobian(fₚ, similar(z, n), AutoForwardDiff(), z)
         end
         i = 1
         for c in 1:n

src/algorithms/solve.jl

@@ -21,6 +21,36 @@ function solve(
     return __solve(prob, sensealg, solver, u0, M, q, args_...; kwargs...)
 end
 
+"""
+    solve(prob::AbstractComplementaritySystem, alg::AbstractComplementaritySystemAlgorithm; kwargs...)
+
+Solve a complementarity system with a compatible system algorithm.
+
+# Arguments
+
+- `prob`: A complementarity system implementing the
+  `AbstractComplementaritySystem` field contract.
+- `alg`: A compatible `AbstractComplementaritySystemAlgorithm` implementation.
+
+# Keyword Arguments
+
+- `kwargs...`: Solver-specific keyword arguments forwarded unchanged to the
+  developer `__solve` implementation.
+
+# Interface
+
+Packages extending this interface must implement
+`ComplementaritySolve.__solve(prob, alg; kwargs...)`. The generic `solve` function
+does not inspect concrete system fields, so the extension owns validation and the
+returned continuous-solver solution type.
+"""
+function solve(
+        prob::AbstractComplementaritySystem, alg::AbstractComplementaritySystemAlgorithm;
+        kwargs...
+    )
+    return __solve(prob, alg; kwargs...)
+end
+
 function __solver_and_args(prob, args...)
     return length(args) == 0 ? (__default_solver(prob), ()) : (first(args), args[2:end])
 end
@@ -40,8 +70,67 @@ end
 ## and works well with inplace/out of place and also works OOTB with GPUs
 __default_solver(::Union{LCP, MCP}) = NonlinearReformulation(:smooth, DEFAULT_NLSOLVER)
 
-# Algorithms should dispatch on __solve
+"""
+    __solve(prob, alg, args...; kwargs...)
+
+Developer extension hook implementing a complementarity solve.
+
+# Arguments
+
+- `prob`: A subtype of `AbstractComplementarityProblem` or
+  `AbstractComplementaritySystem`.
+- `alg`: A compatible developer algorithm subtype.
+- `args...`: Normalized solve data forwarded by `solve`; problem algorithms receive
+  the effective initial state and problem data, while system algorithms receive no
+  additional positional data.
+
+# Keyword Arguments
+
+- `kwargs...`: Algorithm-specific solve options forwarded from `solve`.
+
+# Returns
+
+- A solution satisfying the documented solution contract for the problem family, or
+  the underlying continuous-solver solution for complementarity systems.
+
+# Interface
+
+This qualified, non-exported name is the stable extension point for solver packages.
+Implementations must not mutate stored problem data when processing per-call
+overrides, and must preserve the original `prob` and `alg` in complementarity solution
+objects. Applications should call `solve`, not `__solve`, directly.
+"""
 function __solve end
+
+"""
+    __solve_adjoint(prob, sensealg, sol, Δsol, args...; kwargs...)
+
+Developer extension hook for the reverse rule of a complementarity solve.
+
+# Arguments
+
+- `prob`: The original complementarity problem.
+- `sensealg`: An `AbstractComplementaritySensitivityAlgorithm` compatible with
+  `prob` and `sol`.
+- `sol`: The primal complementarity solution.
+- `Δsol`: The incoming tangent for `sol`.
+- `args...`: The normalized primal solve data passed to `__solve`.
+
+# Keyword Arguments
+
+- `kwargs...`: The solve keyword arguments supplied to the primal solve.
+
+# Returns
+
+- Tangents in the order of the differentiable primal solve data.
+
+# Interface
+
+Implementations must return `NoTangent` or `ZeroTangent` for nondifferentiable data
+and shape-compatible tangents for differentiable data. This hook is called by the
+`ChainRulesCore.rrule` for `__solve`; applications should select a concrete sensitivity
+algorithm through `solve(...; sensealg = ...)` instead of invoking it directly.
+"""
 function __solve_adjoint end
 
 function __solve(

src/problems/complementarity_problems.jl

@@ -61,6 +61,28 @@ abstract type AbstractNonlinearComplementarityProblem{iip} <:
 AbstractComplementarityProblem{iip} end
 
 SciMLBase.isinplace(::AbstractComplementarityProblem{iip}) where {iip} = iip
+
+"""
+    isbatched(prob::AbstractLinearComplementarityProblem) -> Bool
+
+Return whether `prob` stores independent linear complementarity problems in a batch.
+
+# Arguments
+
+- `prob`: A linear complementarity problem subtype.
+
+# Returns
+
+- `true` when `M` has a batch dimension and the corresponding `q` and `u0` values
+  represent independent column-wise problems; otherwise `false`.
+
+# Interface
+
+The return value is determined by the second type parameter of
+`AbstractLinearComplementarityProblem{iip, batched}`. Subtypes must use `true` only
+when their documented storage satisfies the batched matrix/vector layout expected by
+algorithms in this package.
+"""
 isbatched(::AbstractLinearComplementarityProblem{I, B}) where {I, B} = B
 
 """
@@ -97,7 +119,7 @@ using ComplementaritySolve
 M = [2.0 -1.0; -1.0 2.0]
 q = [-1.0, -1.0]
 prob = LinearComplementarityProblem(M, q)
-sol = solve(prob, PGS())
+sol = ComplementaritySolve.solve(prob, PGS())
 ```
 """
 @concrete struct LinearComplementarityProblem{iip, batched} <:

src/problems/complementarity_systems.jl

@@ -4,8 +4,9 @@
 Developer interface for complementarity system containers.
 
 Complementarity systems couple continuous dynamics with complementarity variables.
-Generic system algorithms are expected to call `solve(prob, alg; kwargs...)` and use
-the concrete problem fields documented by the system type.
+External system algorithms implement the qualified developer hook
+`ComplementaritySolve.__solve(prob, alg; kwargs...)`; users invoke the generic
+`solve(prob, alg; kwargs...)` entry point.
 
 # Interface
 
@@ -14,6 +15,9 @@ the concrete problem fields documented by the system type.
 - `prob.tspan` must contain the integration interval.
 - `prob.p` must contain parameters passed to the controller.
 - `prob.controller` must be callable as `controller(x, λ, p, t)`.
+- A subtype and an `AbstractComplementaritySystemAlgorithm` subtype must have a
+  matching `__solve` method. That method must return the solution object produced by
+  the underlying continuous solver.
 """
 abstract type AbstractComplementaritySystem{iip} end
 

src/sensitivity/mcp.jl

@@ -56,12 +56,12 @@ end
     A₂ = ∂ϕ₊∂v₊ * ∂ϕ₋∂u₋ + ∂ϕ₋∂v₋
     if isinplace(prob)
         # Using ForwardDiff for now. We can potentially use Enzyme.jl here
-        J = ForwardDiff.jacobian((y, u) -> f(y, u, p), fᵤ, u)
+        J = jacobian((y, u) -> f(y, u, p), fᵤ, AutoForwardDiff(), u)
         A = J' * A₁ .+ A₂
     else
         if length(u) ≤ 50
             # Construct the Full Matrix
-            A = only(Zygote.jacobian(Base.Fix2(f, p), u))' * A₁ .+ A₂
+            A = jacobian(Base.Fix2(f, p), AutoZygote(), u)' * A₁ .+ A₂
         else
             # Use Matrix Free Methods
             ## NOTE: If we use SparseDiffTools here we will have to mess around with a wrapper
@@ -79,7 +79,7 @@ end
 
     if isinplace(prob)
         # Using ForwardDiff for now. We can potentially use Enzyme.jl here
-        J = ForwardDiff.jacobian((y, p) -> f(y, u, p), fᵤ, p)
+        J = jacobian((y, p) -> f(y, u, p), fᵤ, AutoForwardDiff(), p)
         ∂p = -reshape((A₁ * λ)' * J, size(p))
     else
         _, pb_f = Zygote.pullback(Base.Fix1(f, u), p)

src/utils.jl

@@ -35,7 +35,7 @@ __notangent(::ZeroTangent) = true
 __notangent(::NoTangent) = true
 __notangent(::Any) = false
 
-__unfillarray(x::AbstractFill) = collect(x)
+__unfillarray(x::AbstractArray) = ismutable(x) ? x : collect(x)
 __unfillarray(x) = x
 
 function batched_matvec(A::AA3, x::AM)