From 9bb76a69a3017f1564d168e44b29b9c048fef337 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 26 Sep 2024 13:32:43 -0400
Subject: [PATCH 1/9] fix: remove unused deps

---
 Project.toml                  | 14 +-------------
 src/DiffEqCallbacks.jl        | 25 +++++++++++++------------
 src/iterative_and_periodic.jl |  2 +-
 3 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/Project.toml b/Project.toml
index 770e6937..b4de0da7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,41 +1,29 @@
 name = "DiffEqCallbacks"
 uuid = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "3.9.1"
+version = "4.0.0"
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
-NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
-[weakdeps]
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
-OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
-Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
-
 [compat]
 Aqua = "0.8"
 DataInterpolations = "4.6"
 DataStructures = "0.18.13"
 DiffEqBase = "6.154"
-ForwardDiff = "0.10.36"
 Functors = "0.4"
 LinearAlgebra = "1.10"
 Markdown = "1.10"
-NonlinearSolve = "3.7.2"
 ODEProblemLibrary = "0.1.5"
 OrdinaryDiffEq = "6.88"
-OrdinaryDiffEqCore = "1"
-Parameters = "0.12"
 QuadGK = "2.9"
 RecipesBase = "1.3.4"
 RecursiveArrayTools = "3.9"
diff --git a/src/DiffEqCallbacks.jl b/src/DiffEqCallbacks.jl
index 84967de1..85033c80 100644
--- a/src/DiffEqCallbacks.jl
+++ b/src/DiffEqCallbacks.jl
@@ -1,17 +1,18 @@
 module DiffEqCallbacks
 
-using DiffEqBase, RecursiveArrayTools, DataStructures, RecipesBase, LinearAlgebra,
-      StaticArraysCore, NonlinearSolve, ForwardDiff, Functors
-
-import Base.Iterators
-
-using Markdown
-
-using Parameters: @unpack
-
-import SciMLBase
-
-using DiffEqBase: get_tstops, get_tstops_array, get_tstops_max
+using DataStructures: DataStructures, BinaryMaxHeap, BinaryMinHeap
+using DiffEqBase: DiffEqBase, get_tstops, get_tstops_array, get_tstops_max
+using Functors: fmap
+using LinearAlgebra: LinearAlgebra, adjoint, axpy!, copyto!
+using Markdown: @doc_str
+using RecipesBase: @recipe
+using RecursiveArrayTools: RecursiveArrayTools, DiffEqArray, copyat_or_push!
+using SciMLBase: SciMLBase, CallbackSet, DiscreteCallback, NonlinearFunction,
+                 NonlinearLeastSquaresProblem, NonlinearProblem, RODEProblem,
+                 ReturnCode, SDEProblem, add_tstop!, check_error, get_du,
+                 get_proposed_dt, get_tmp_cache, init, reinit!,
+                 set_proposed_dt!, solve!, terminate!, u_modified!
+using StaticArraysCore: StaticArraysCore
 
 include("functor_helpers.jl")
 include("autoabstol.jl")
diff --git a/src/iterative_and_periodic.jl b/src/iterative_and_periodic.jl
index 9e66de88..a64c3e48 100644
--- a/src/iterative_and_periodic.jl
+++ b/src/iterative_and_periodic.jl
@@ -89,7 +89,7 @@ function (S::PeriodicCallbackAffect)(integrator)
 end
 
 function add_next_tstop!(integrator, S)
-    @unpack Δt, t0, index = S
+    (; Δt, t0, index) = S
 
     # Schedule next call to `f` using `add_tstops!`, but be careful not to keep integrating forever
     tnew = t0[] + (index[] + 1) * Δt

From c83fe3eda16d6b40756019864ff84900bc44444e Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 26 Sep 2024 13:39:18 -0400
Subject: [PATCH 2/9] fix: update dependencies and remove unused ones

---
 Project.toml | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index b4de0da7..a675a449 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "DiffEqCallbacks"
 uuid = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "4.0.0"
+version = "3.10.0"
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -16,12 +16,14 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [compat]
 Aqua = "0.8"
-DataInterpolations = "4.6"
+DataInterpolations = "4.6, 5, 6"
 DataStructures = "0.18.13"
 DiffEqBase = "6.154"
+ForwardDiff = "0.10.36"
 Functors = "0.4"
 LinearAlgebra = "1.10"
 Markdown = "1.10"
+NonlinearSolve = "3.7.2"
 ODEProblemLibrary = "0.1.5"
 OrdinaryDiffEq = "6.88"
 QuadGK = "2.9"
@@ -40,6 +42,7 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -52,4 +55,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "DataInterpolations", "OrdinaryDiffEq", "ODEProblemLibrary", "Test", "QuadGK", "SciMLSensitivity", "StaticArrays", "Tracker", "Zygote", "NonlinearSolve"]
+test = ["Aqua", "DataInterpolations", "ForwardDiff", "NonlinearSolve", "ODEProblemLibrary", "OrdinaryDiffEq", "QuadGK", "SciMLSensitivity", "StaticArrays", "Test", "Tracker", "Zygote"]

From 4312cf03b618b75fdd09db2e0f4e706000c7b93b Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 26 Sep 2024 17:10:51 -0400
Subject: [PATCH 3/9] feat: single factorize manifold projection based on
 Hairer III

---
 .JuliaFormatter.toml   |   3 +-
 Project.toml           |   8 +-
 src/DiffEqCallbacks.jl |   6 +-
 src/manifold.jl        | 353 +++++++++++++++++++++++++++++++----------
 test/manifold_tests.jl |  40 ++---
 5 files changed, 301 insertions(+), 109 deletions(-)

diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml
index 9c793591..320e0c07 100644
--- a/.JuliaFormatter.toml
+++ b/.JuliaFormatter.toml
@@ -1,2 +1,3 @@
 style = "sciml"
-format_markdown = true
\ No newline at end of file
+format_markdown = true
+annotate_untyped_fields_with_any = false
diff --git a/Project.toml b/Project.toml
index a675a449..51ab1d47 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,8 +4,10 @@ authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
 version = "3.10.0"
 
 [deps]
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
@@ -15,10 +17,13 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [compat]
+ADTypes = "1.9.0"
 Aqua = "0.8"
+ConcreteStructs = "0.2.3"
 DataInterpolations = "4.6, 5, 6"
 DataStructures = "0.18.13"
 DiffEqBase = "6.154"
+DifferentiationInterface = "0.6.1"
 ForwardDiff = "0.10.36"
 Functors = "0.4"
 LinearAlgebra = "1.10"
@@ -40,6 +45,7 @@ Zygote = "0.6.69"
 julia = "1.10"
 
 [extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -55,4 +61,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "DataInterpolations", "ForwardDiff", "NonlinearSolve", "ODEProblemLibrary", "OrdinaryDiffEq", "QuadGK", "SciMLSensitivity", "StaticArrays", "Test", "Tracker", "Zygote"]
+test = ["ADTypes", "Aqua", "DataInterpolations", "ForwardDiff", "NonlinearSolve", "ODEProblemLibrary", "OrdinaryDiffEq", "QuadGK", "SciMLSensitivity", "StaticArrays", "Test", "Tracker", "Zygote"]
diff --git a/src/DiffEqCallbacks.jl b/src/DiffEqCallbacks.jl
index 85033c80..e483795f 100644
--- a/src/DiffEqCallbacks.jl
+++ b/src/DiffEqCallbacks.jl
@@ -1,9 +1,11 @@
 module DiffEqCallbacks
 
+using ConcreteStructs: @concrete
 using DataStructures: DataStructures, BinaryMaxHeap, BinaryMinHeap
 using DiffEqBase: DiffEqBase, get_tstops, get_tstops_array, get_tstops_max
+using DifferentiationInterface: DifferentiationInterface, Constant
 using Functors: fmap
-using LinearAlgebra: LinearAlgebra, adjoint, axpy!, copyto!
+using LinearAlgebra: LinearAlgebra, adjoint, axpy!, copyto!, mul!, ldiv!
 using Markdown: @doc_str
 using RecipesBase: @recipe
 using RecursiveArrayTools: RecursiveArrayTools, DiffEqArray, copyat_or_push!
@@ -14,6 +16,8 @@ using SciMLBase: SciMLBase, CallbackSet, DiscreteCallback, NonlinearFunction,
                  set_proposed_dt!, solve!, terminate!, u_modified!
 using StaticArraysCore: StaticArraysCore
 
+const DI = DifferentiationInterface
+
 include("functor_helpers.jl")
 include("autoabstol.jl")
 include("manifold.jl")
diff --git a/src/manifold.jl b/src/manifold.jl
index 891114b7..be30e9cc 100644
--- a/src/manifold.jl
+++ b/src/manifold.jl
@@ -1,20 +1,7 @@
-# wrapper for non-autonomous functions
-mutable struct NonAutonomousFunction{iip, F, autonomous}
-    f::F
-    t::Any
-end
-
-(f::NonAutonomousFunction{true, F, true})(res, u, p) where {F} = f.f(res, u, p)
-(f::NonAutonomousFunction{true, F, false})(res, u, p) where {F} = f.f(res, u, p, f.t)
-
-(f::NonAutonomousFunction{false, F, true})(u, p) where {F} = f.f(u, p)
-(f::NonAutonomousFunction{false, F, false})(u, p) where {F} = f.f(u, p, f.t)
-
-SciMLBase.isinplace(::NonAutonomousFunction{iip}) where {iip} = iip
-
 """
-    ManifoldProjection(g; nlsolve = nothing, save = true, nlls = Val(true),
-        isinplace = Val(true), autonomous = nothing, resid_prototype = nothing, kwargs...)
+    ManifoldProjection(
+        manifold; nlsolve = missing, save = true, autonomous = nothing,
+        manifold_jacobian = nothing, autodiff = nothing, kwargs...)
 
 In many cases, you may want to declare a manifold on which a solution lives. Mathematically,
 a manifold `M` is defined by a function `g` as the set of points where `g(u) = 0`. An
@@ -33,30 +20,34 @@ properties.
 
 ## Arguments
 
-  - `g`: The residual function for the manifold.
-
-      * This is an inplace function of form `g(resid, u, p)` or `g(resid, u, p, t)` which
-        writes to the residual `resid` the difference from the manifold components. Here, it
-        is assumed that `resid` is of the same shape as `u`.
-      * If `isinplace = Val(false)`, then `g` should be a function of the form `g(u, p)` or
-        `g(u, p, t)` which returns the residual.
+  - `manifold`: The residual function for the manifold. If the ODEProblem is an inplace
+    problem, then `g` must be an inplace function of form `g(resid, u, p)` or
+    `g(resid, u, p, t)` and similarly if the ODEProblem is an out-of-place problem then `g`
+    must be a function of form `g(u, p)` or `g(u, p, t)`.
 
 ## Keyword Arguments
 
-  - `nlsolve`: A nonlinear solver as defined in the
-    [NonlinearSolve.jl format](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/).
-    Defaults to a PolyAlgorithm.
+  - `nlsolve`: Defaults to a special single-factorize algorithm (denoted by `missing`) that
+    would work in most cases (See [1] for details). Alternatively, a nonlinear solver as
+    defined in the
+    [NonlinearSolve.jl format](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/)
+    can be specified.
   - `save`: Whether to do the standard saving (applied after the callback)
-  - `nlls`: If the problem is a nonlinear least squares problem. `nlls = Val(false)`
-    generates a `NonlinearProblem` which is typically faster than
-    `NonlinearLeastSquaresProblem`, but is only applicable if the residual size is same as
-    the state size.
   - `autonomous`: Whether `g` is an autonomous function of the form `g(resid, u, p)` or
-    `g(u, p)`. Specify it as `Val(::Bool)` to ensure this function call is type stable.
-  - `residual_prototype`: This needs to be specified if `nlls = Val(true)` and
-    `inplace = Val(true)` are specified together, else it is taken to be same as `u`.
+    `g(u, p)`. Specify it as `Val(::Bool)` to disable runtime branching. If `nothing`,
+    we attempt to infer it from the signature of `g`.
+  - `residual_prototype`: The size of the manifold residual. If it is not specified,
+    we assume it to be same as `u` in the inplace case. Else we run a single evaluation of
+    `manifold` to determine the size.
   - `kwargs`: All other keyword arguments are passed to
-    [NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/).
+    [NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/) if
+    `nlsolve` is not `missing`.
+  - `autodiff`: The autodifferentiation algorithm to use to compute the Jacobian if
+    `manifold_jacobian` is not specified. This must be specified if `manifold_jacobian` is
+    not specified and `nlsolve` is `missing`. If `nlsolve` is not `missing`, then
+    `autodiff` is ignored.
+  - `manifold_jacobian`: The Jacobian of the manifold (wrt the state). This has the same
+    signature as `manifold` and the first argument is the Jacobian if inplace.
 
 ### Saveat Warning
 
@@ -72,78 +63,276 @@ order of the integrator even with the ManifoldProjection.
 
 ## References
 
-Ernst Hairer, Christian Lubich, Gerhard Wanner. Geometric Numerical Integration:
+[1] Ernst Hairer, Christian Lubich, Gerhard Wanner. Geometric Numerical Integration:
 Structure-Preserving Algorithms for Ordinary Differential Equations. Berlin ;
 New York :Springer, 2002.
 """
-mutable struct ManifoldProjection{iip, nlls, autonomous, F, NL, R, K}
-    g::F
+@concrete mutable struct ManifoldProjection
+    manifold
+    manifold_jacobian
+    autodiff
     nlcache::Any
-    nlsolve::NL
-    resid_prototype::R
-    kwargs::K
-
-    function ManifoldProjection{iip, nlls, autonomous}(
-            g, nlsolve, resid_prototype, kwargs) where {iip, nlls, autonomous}
-        # replace residual function if it is time-dependent
-        _g = NonAutonomousFunction{iip, typeof(g), autonomous}(g, 0)
-        return new{iip, nlls, autonomous, typeof(_g), typeof(nlsolve),
-            typeof(resid_prototype), typeof(kwargs)}(
-            _g, nothing, nlsolve, resid_prototype, kwargs)
+    nlsolve
+    kwargs
+    autonomous
+end
+
+function ManifoldProjection(
+        manifold; nlsolve = missing, save = true, autonomous = nothing,
+        manifold_jacobian = nothing, autodiff = nothing, kwargs...)
+    affect! = ManifoldProjection(
+        manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous)
+    return DiscreteCallback(
+        Returns(true), affect!; initialize = initialize_manifold_projection,
+        save_positions = (false, save))
+end
+
+function ManifoldProjection(
+        manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous)
+    if autonomous isa Val{true} || autonomous isa Val{false}
+        wrapped_manifold = TypedNonAutonomousFunction{SciMLBase._unwrap_val(autonomous)}(
+            manifold, nothing)
+        wrapped_manifold_jacobian = if manifold_jacobian === nothing
+            nothing
+        else
+            TypedNonAutonomousFunction{SciMLBase._unwrap_val(autonomous)}(
+                manifold_jacobian, nothing)
+        end
+        autonomous = SciMLBase._unwrap_val(autonomous)
+    else
+        _autonomous = autonomous === nothing ? false : autonomous
+        wrapped_manifold = UntypedNonAutonomousFunction(_autonomous, manifold, nothing)
+        wrapped_manifold_jacobian = if manifold_jacobian === nothing
+            nothing
+        else
+            UntypedNonAutonomousFunction(_autonomous, manifold_jacobian, nothing)
+        end
     end
+    return ManifoldProjection(wrapped_manifold, wrapped_manifold_jacobian,
+        autodiff, nothing, nlsolve, kwargs, autonomous)
 end
 
-# Now make `affect!` for this:
-function (p::ManifoldProjection{
-        iip, nlls, autonomous, NL})(integrator) where {iip, nlls,
-        autonomous, NL}
+function (proj::ManifoldProjection)(integrator)
     # update current time if residual function is time-dependent
-    !autonomous && (p.g.t = integrator.t)
+    proj.manifold.t = integrator.t
+    proj.manifold_jacobian !== nothing && (proj.manifold_jacobian.t = integrator.t)
 
-    # solve the nonlinear problem
-    reinit!(p.nlcache, integrator.u; p = integrator.p)
-    sol = solve!(p.nlcache)
+    SciMLBase.reinit!(proj.nlcache, integrator.u; integrator.p)
 
-    if !SciMLBase.successful_retcode(sol)
-        SciMLBase.terminate!(integrator, sol.retcode)
+    if proj.nlsolve === missing
+        _, u, retcode = SciMLBase.solve!(proj.nlcache)
+    else
+        sol = SciMLBase.solve!(proj.nlcache)
+        (; u, retcode) = sol
+    end
+
+    if !SciMLBase.successful_retcode(retcode)
+        SciMLBase.terminate!(integrator, retcode)
         return
     end
 
-    copyto!(integrator.u, sol.u)
+    SciMLBase.copyto!(integrator.u, u)
 end
 
-function Manifold_initialize(cb, u, t, integrator)
-    return Manifold_initialize(cb.affect!, u, t, integrator)
+function initialize_manifold_projection(cb, u, t, integrator)
+    return initialize_manifold_projection(cb.affect!, u, t, integrator)
 end
-function Manifold_initialize(
-        affect!::ManifoldProjection{iip, nlls}, u, t, integrator) where {iip, nlls}
-    nlfunc = NonlinearFunction{iip}(affect!.g; affect!.resid_prototype)
-    nlprob = if nlls
-        NonlinearLeastSquaresProblem(nlfunc, u, integrator.p)
+function initialize_manifold_projection(affect!::ManifoldProjection, u, t, integrator)
+    if affect!.autonomous === nothing
+        autonomous = maximum(SciMLBase.numargs(affect!.manifold.f)) ==
+                     2 + SciMLBase.isinplace(integrator.f)
+        affect!.manifold.autonomous = autonomous
+        affect!.manifold_jacobian !== nothing &&
+            (affect!.manifold_jacobian.autonomous = autonomous)
+    end
+
+    if affect!.nlsolve === missing
+        affect!.manifold.t = t
+        affect!.manifold_jacobian !== nothing && (affect!.manifold_jacobian.t = t)
+        cache = init_manifold_projection(
+            Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff,
+            affect!.manifold_jacobian, u, integrator.p; affect!.kwargs...)
     else
-        NonlinearProblem(nlfunc, u, integrator.p)
+        # nlfunc = NonlinearFunction{iip}(affect!.g; affect!.resid_prototype)
+        # nlprob = NonlinearProblem(nlfunc, u, integrator.p)
+        # affect!.nlcache = init(nlprob, affect!.nlsolve; affect!.kwargs...)
+        error("Not Implemented")
     end
-    affect!.nlcache = init(nlprob, affect!.nlsolve; affect!.kwargs...)
+    affect!.nlcache = cache
     u_modified!(integrator, false)
 end
 
-function ManifoldProjection(g; nlsolve = nothing, save = true, nlls = Val(true),
-        isinplace = Val(true), autonomous = nothing, resid_prototype = nothing,
-        kwargs...)
-    _nlls = SciMLBase._unwrap_val(nlls)
-    iip = SciMLBase._unwrap_val(isinplace)
-    if autonomous === nothing
+export ManifoldProjection
+
+# wrapper for non-autonomous functions
+@concrete mutable struct TypedNonAutonomousFunction{autonomous}
+    f
+    t::Any
+end
+
+(f::TypedNonAutonomousFunction{false})(res, u, p) = f.f(res, u, p, f.t)
+(f::TypedNonAutonomousFunction{true})(res, u, p) = f.f(res, u, p)
+
+(f::TypedNonAutonomousFunction{false})(u, p) = f.f(u, p, f.t)
+(f::TypedNonAutonomousFunction{true})(u, p) = f.f(u, p)
+
+@concrete mutable struct UntypedNonAutonomousFunction
+    autonomous::Bool
+    f
+    t::Any
+end
+
+function (f::UntypedNonAutonomousFunction)(res, u, p)
+    return f.autonomous ? f.f(res, u, p) : f.f(res, u, p, f.t)
+end
+(f::UntypedNonAutonomousFunction)(u, p) = f.autonomous ? f.f(u, p) : f.f(u, p, f.t)
+
+# This is the algorithm described in Hairer III.
+@concrete mutable struct SingleFactorizeManifoldProjectionCache{iip}
+    manifold
+    p
+    abstol
+    maxiters::Int
+
+    ũ
+    JJᵀfact::Any  # LU might fail and we might end up doing QR
+    u_cache
+    λ_cache
+    gu_cache
+
+    first_call::Bool
+    J
+    JJᵀ
+    manifold_jacobian
+    autodiff
+    di_extras
+end
+
+function SciMLBase.reinit!(
+        cache::SingleFactorizeManifoldProjectionCache{iip}, u; p = cache.p) where {iip}
+    if !cache.first_call || (cache.ũ !== u || cache.p !== p)
+        if cache.manifold_jacobian !== nothing
+            if iip
+                cache.manifold_jacobian(cache.J, u, p)
+            else
+                cache.J = cache.manifold_jacobian(u, p)
+            end
+        else
+            if iip
+                DI.jacobian!(cache.manifold, cache.gu_cache, cache.J,
+                    cache.di_extras, cache.autodiff, u, Constant(p))
+            else
+                DI.jacobian!(cache.manifold, cache.J, cache.di_extras,
+                    cache.autodiff, u, Constant(p))
+            end
+        end
+        mul!(cache.JJᵀ, cache.J, cache.J')
+        cache.JJᵀfact = safe_factorize!(cache.JJᵀ)
+    end
+    cache.first_call = false
+    cache.ũ = u
+    cache.p = p
+end
+
+default_abstol(::Type{T}) where {T} = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)
+
+function init_manifold_projection(::Val{iip}, manifold, autodiff, manifold_jacobian, ũ,
+        p; abstol = default_abstol(eltype(ũ)), maxiters = 1000,
+        resid_prototype = nothing) where {iip}
+    if iip
+        if resid_prototype !== nothing
+            gu = similar(resid_prototype)
+            λ = similar(resid_prototype)
+        else
+            @warn "`resid_prototype` not provided for in-place problem. Assuming size of \
+                   output is the same as input. This might be incorrect." maxlog=1
+            gu = similar(ũ)
+            λ = similar(ũ)
+        end
+    else
+        gu = nothing
+        λ = manifold(ũ, p)
+    end
+
+    if manifold_jacobian !== nothing
+        if iip
+            J = similar(ũ, promote_type(eltype(gu), eltype(ũ)), (length(gu), length(ũ)))
+            manifold_jacobian(J, ũ, p)
+        else
+            J = manifold_jacobian(ũ, p)
+        end
+        di_extras = nothing
+    elseif autodiff !== nothing
         if iip
-            autonomous = maximum(SciMLBase.numargs(g)) == 3
+            di_extras = DI.prepare_jacobian(manifold, gu, autodiff, ũ, Constant(p))
+            J = DI.jacobian(manifold, gu, di_extras, autodiff, ũ, Constant(p))
         else
-            autonomous = maximum(SciMLBase.numargs(g)) == 2
+            di_extras = DI.prepare_jacobian(manifold, autodiff, ũ, Constant(p))
+            J = DI.jacobian(manifold, di_extras, autodiff, ũ, Constant(p))
         end
+    else
+        error("`autodiff` is set to `nothing` and analytic manifold jacobian is not \
+               provided.")
     end
-    affect! = ManifoldProjection{iip, _nlls, SciMLBase._unwrap_val(autonomous)}(
-        g, nlsolve, resid_prototype, kwargs)
-    condition = (u, t, integrator) -> true
-    return DiscreteCallback(condition, affect!; initialize = Manifold_initialize,
-        save_positions = (false, save))
+    JJᵀ = J * J'
+    JJᵀfact = safe_factorize!(JJᵀ)
+
+    return SingleFactorizeManifoldProjectionCache{iip}(
+        manifold, p, abstol, maxiters, ũ, JJᵀfact, similar(ũ), λ, gu,
+        true, J, JJᵀ, manifold_jacobian, autodiff, di_extras)
 end
 
-export ManifoldProjection
+function SciMLBase.solve!(cache::SingleFactorizeManifoldProjectionCache{iip}) where {iip}
+    fill!(cache.λ_cache, false)
+    ũ = cache.ũ
+    gu = cache.gu_cache
+
+    internal_solve_failed = true
+
+    if cache.gu_cache !== nothing
+        cache.manifold(gu, ũ, cache.p)
+    else
+        gu = cache.manifold(ũ, cache.p)
+    end
+
+    for _ in 1:(cache.maxiters)
+        if maximum(abs, gu) < cache.abstol
+            internal_solve_failed = false
+            break
+        end
+
+        δλ = cache.JJᵀfact \ gu
+        @. cache.λ_cache -= δλ
+
+        mul!(vec(cache.u_cache), cache.J', vec(cache.λ_cache))
+        cache.u_cache += ũ
+        if cache.gu_cache !== nothing
+            cache.manifold(gu, cache.u_cache, cache.p)
+        else
+            gu = cache.manifold(cache.u_cache, cache.p)
+        end
+    end
+
+    return (cache.λ_cache, cache.u_cache,
+        ifelse(internal_solve_failed, ReturnCode.ConvergenceFailure, ReturnCode.Success))
+end
+
+function safe_factorize!(A::AbstractMatrix)
+    if issquare(A)
+        fact = LinearAlgebra.lu(A; check = false)
+        fact_sucessful(fact) && return fact
+    elseif size(A, 1) > size(A, 2)
+        fact = LinearAlgebra.qr(A)
+        fact_sucessful(fact) && return fact
+    end
+    return LinearAlgebra.qr!(A, LinearAlgebra.ColumnNorm())
+end
+
+function fact_sucessful(F::LinearAlgebra.QRCompactWY)
+    m, n = size(F)
+    U = view(F.factors, 1:min(m, n), 1:n)
+    return all(!iszero, Iterators.reverse(@view U[diagind(U)]))
+end
+function fact_sucessful(F::FT) where {FT}
+    return hasmethod(LinearAlgebra.issuccess, (FT,)) ? LinearAlgebra.issuccess(F) : true
+end
diff --git a/test/manifold_tests.jl b/test/manifold_tests.jl
index 67e8a994..41d16575 100644
--- a/test/manifold_tests.jl
+++ b/test/manifold_tests.jl
@@ -1,4 +1,5 @@
 using OrdinaryDiffEq, Test, DiffEqBase, DiffEqCallbacks, RecursiveArrayTools, NonlinearSolve
+using ForwardDiff, ADTypes
 
 u0 = ones(2, 2)
 f = function (du, u, p, t)
@@ -14,39 +15,30 @@ end
 
 g_t(resid, u, p, t) = g(resid, u, p)
 
-function isautonomous(::ManifoldProjection{
-        iip, nlls, autonomous, NL}) where {iip, nlls, autonomous, NL}
-    return autonomous
-end
-
 sol = solve(prob, Vern7())
 @test !(sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2)
 
 # autodiff=true
 @inferred ManifoldProjection(g; autonomous = Val(false), resid_prototype = zeros(2))
-cb = ManifoldProjection(g; resid_prototype = zeros(2))
-@test isautonomous(cb.affect!)
+cb = ManifoldProjection(g; resid_prototype = zeros(2), autodiff = AutoForwardDiff())
 solve(prob, Vern7(), callback = cb)
 @time sol = solve(prob, Vern7(), callback = cb)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
 
-cb_t = ManifoldProjection(g_t; resid_prototype = zeros(2))
-@test !isautonomous(cb_t.affect!)
+cb_t = ManifoldProjection(g_t; resid_prototype = zeros(2), autodiff = AutoForwardDiff())
 solve(prob, Vern7(), callback = cb_t)
 @time sol_t = solve(prob, Vern7(), callback = cb_t)
 @test sol_t.u == sol.u && sol_t.t == sol.t
 
 # autodiff=false
 cb_false = ManifoldProjection(
-    g; nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2))
-@test isautonomous(cb_false.affect!)
+    g; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2))
 solve(prob, Vern7(), callback = cb_false)
 sol = solve(prob, Vern7(), callback = cb_false)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
 
 cb_t_false = ManifoldProjection(g_t,
-    nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2))
-@test !isautonomous(cb_t_false.affect!)
+    nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2))
 solve(prob, Vern7(), callback = cb_t_false)
 sol_t = solve(prob, Vern7(), callback = cb_t_false)
 @test sol_t.u == sol.u && sol_t.t == sol.t
@@ -73,7 +65,8 @@ function g_unsat(resid, u, p)
     resid[2] = u[2]^2 + u[1]^2 - 20
 end
 
-cb_unsat = ManifoldProjection(g_unsat; resid_prototype = zeros(2))
+cb_unsat = ManifoldProjection(
+    g_unsat; resid_prototype = zeros(2), autodiff = AutoForwardDiff())
 sol = solve(prob, Vern7(), callback = cb_unsat)
 @test !SciMLBase.successful_retcode(sol)
 @test last(sol.t) != 100.0
@@ -86,33 +79,32 @@ end
 
 g_oop_t(u, p, t) = g_oop(u, p)
 
-prob = ODEProblem(f, u0, (0.0, 100.0))
+f_oop = function (u, p, t)
+    return stack((u[2, :], -u[1, :]))
+end
+prob = ODEProblem(f_oop, u0, (0.0, 100.0))
 
 # autodiff=true
-@inferred ManifoldProjection(g_oop; autonomous = Val(false), isinplace = Val(false))
-cb = ManifoldProjection(g_oop; isinplace = Val(false))
-@test isautonomous(cb.affect!)
+@inferred ManifoldProjection(g_oop; autonomous = Val(false))
+cb = ManifoldProjection(g_oop; autodiff = AutoForwardDiff())
 solve(prob, Vern7(), callback = cb)
 @time sol = solve(prob, Vern7(), callback = cb)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
 
-cb_t = ManifoldProjection(g_oop_t; isinplace = Val(false))
-@test !isautonomous(cb_t.affect!)
+cb_t = ManifoldProjection(g_oop_t; autodiff = AutoForwardDiff())
 solve(prob, Vern7(), callback = cb_t)
 @time sol_t = solve(prob, Vern7(), callback = cb_t)
 @test sol_t.u == sol.u && sol_t.t == sol.t
 
 # autodiff=false
 cb_false = ManifoldProjection(
-    g_oop; nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), isinplace = Val(false))
-@test isautonomous(cb_false.affect!)
+    g_oop; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), isinplace = Val(false))
 solve(prob, Vern7(), callback = cb_false)
 sol = solve(prob, Vern7(), callback = cb_false)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
 
 cb_t_false = ManifoldProjection(g_oop_t,
-    nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), isinplace = Val(false))
-@test !isautonomous(cb_t_false.affect!)
+    nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), isinplace = Val(false))
 solve(prob, Vern7(), callback = cb_t_false)
 sol_t = solve(prob, Vern7(), callback = cb_t_false)
 @test sol_t.u == sol.u && sol_t.t == sol.t

From d44a4b2c3ab0046d5a16ed60b91555afd8a29045 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 26 Sep 2024 17:36:14 -0400
Subject: [PATCH 4/9] refactor: modularize the code

---
 src/manifold.jl | 90 +++++++++++++++++++++++++++++--------------------
 1 file changed, 53 insertions(+), 37 deletions(-)

diff --git a/src/manifold.jl b/src/manifold.jl
index be30e9cc..5f095c10 100644
--- a/src/manifold.jl
+++ b/src/manifold.jl
@@ -211,21 +211,8 @@ end
 function SciMLBase.reinit!(
         cache::SingleFactorizeManifoldProjectionCache{iip}, u; p = cache.p) where {iip}
     if !cache.first_call || (cache.ũ !== u || cache.p !== p)
-        if cache.manifold_jacobian !== nothing
-            if iip
-                cache.manifold_jacobian(cache.J, u, p)
-            else
-                cache.J = cache.manifold_jacobian(u, p)
-            end
-        else
-            if iip
-                DI.jacobian!(cache.manifold, cache.gu_cache, cache.J,
-                    cache.di_extras, cache.autodiff, u, Constant(p))
-            else
-                DI.jacobian!(cache.manifold, cache.J, cache.di_extras,
-                    cache.autodiff, u, Constant(p))
-            end
-        end
+        compute_manifold_jacobian!(cache.J, cache.manifold_jacobian, cache.autodiff,
+            Val(iip), cache.manifold, cache.gu_cache, u, p, cache.di_extras)
         mul!(cache.JJᵀ, cache.J, cache.J')
         cache.JJᵀfact = safe_factorize!(cache.JJᵀ)
     end
@@ -236,7 +223,7 @@ end
 
 default_abstol(::Type{T}) where {T} = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)
 
-function init_manifold_projection(::Val{iip}, manifold, autodiff, manifold_jacobian, ũ,
+function init_manifold_projection(IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ,
         p; abstol = default_abstol(eltype(ũ)), maxiters = 1000,
         resid_prototype = nothing) where {iip}
     if iip
@@ -254,26 +241,8 @@ function init_manifold_projection(::Val{iip}, manifold, autodiff, manifold_jacob
         λ = manifold(ũ, p)
     end
 
-    if manifold_jacobian !== nothing
-        if iip
-            J = similar(ũ, promote_type(eltype(gu), eltype(ũ)), (length(gu), length(ũ)))
-            manifold_jacobian(J, ũ, p)
-        else
-            J = manifold_jacobian(ũ, p)
-        end
-        di_extras = nothing
-    elseif autodiff !== nothing
-        if iip
-            di_extras = DI.prepare_jacobian(manifold, gu, autodiff, ũ, Constant(p))
-            J = DI.jacobian(manifold, gu, di_extras, autodiff, ũ, Constant(p))
-        else
-            di_extras = DI.prepare_jacobian(manifold, autodiff, ũ, Constant(p))
-            J = DI.jacobian(manifold, di_extras, autodiff, ũ, Constant(p))
-        end
-    else
-        error("`autodiff` is set to `nothing` and analytic manifold jacobian is not \
-               provided.")
-    end
+    J, di_extras = setup_manifold_jacobian(manifold_jacobian, autodiff, IIP, manifold,
+        gu, ũ, p)
     JJᵀ = J * J'
     JJᵀfact = safe_factorize!(JJᵀ)
 
@@ -317,9 +286,56 @@ function SciMLBase.solve!(cache::SingleFactorizeManifoldProjectionCache{iip}) wh
         ifelse(internal_solve_failed, ReturnCode.ConvergenceFailure, ReturnCode.Success))
 end
 
+function setup_manifold_jacobian(
+        manifold_jacobian::M, autodiff, ::Val{iip}, manifold, gu, ũ, p) where {M, iip}
+    if iip
+        J = similar(ũ, promote_type(eltype(gu), eltype(ũ)), (length(gu), length(ũ)))
+        manifold_jacobian(J, ũ, p)
+    else
+        J = manifold_jacobian(ũ, p)
+    end
+    return J, nothing
+end
+
+function setup_manifold_jacobian(
+        ::Nothing, autodiff, ::Val{iip}, manifold, gu, ũ, p) where {iip}
+    if iip
+        di_extras = DI.prepare_jacobian(manifold, gu, autodiff, ũ, Constant(p))
+        J = DI.jacobian(manifold, gu, di_extras, autodiff, ũ, Constant(p))
+    else
+        di_extras = DI.prepare_jacobian(manifold, autodiff, ũ, Constant(p))
+        J = DI.jacobian(manifold, di_extras, autodiff, ũ, Constant(p))
+    end
+    return J, di_extras
+end
+
+function compute_manifold_jacobian!(J, manifold_jacobian, autodiff, ::Val{iip},
+        manifold, gu, ũ, p, di_extras) where {iip}
+    if iip
+        manifold_jacobian(J, ũ, p)
+    else
+        J = manifold_jacobian(ũ, p)
+    end
+    return J
+end
+
+function compute_manifold_jacobian!(J, ::Nothing, autodiff, ::Val{iip}, manifold, gu,
+        ũ, p, di_extras) where {iip}
+    if iip
+        DI.jacobian!(manifold, gu, J, di_extras, autodiff, ũ, Constant(p))
+    else
+        DI.jacobian!(manifold, J, di_extras, autodiff, ũ, Constant(p))
+    end
+    return J
+end
+
+function setup_manifold_jacobian(::Nothing, ::Nothing, args...)
+    error("`autodiff` is set to `nothing` and analytic manifold jacobian is not provided.")
+end
+
 function safe_factorize!(A::AbstractMatrix)
     if issquare(A)
-        fact = LinearAlgebra.lu(A; check = false)
+        fact = LinearAlgebra.cholesky(A; check = false)
         fact_sucessful(fact) && return fact
     elseif size(A, 1) > size(A, 2)
         fact = LinearAlgebra.qr(A)

From 7b527688fbbea355e2e4052fb4b702d54aef6ca1 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 26 Sep 2024 18:27:00 -0400
Subject: [PATCH 5/9] feat: lagrangian multiplier based projection algorithm

---
 src/manifold.jl        | 125 ++++++++++++++++++++++++++++++++++-------
 test/manifold_tests.jl |  27 +++++++--
 2 files changed, 126 insertions(+), 26 deletions(-)

diff --git a/src/manifold.jl b/src/manifold.jl
index 5f095c10..74eef9db 100644
--- a/src/manifold.jl
+++ b/src/manifold.jl
@@ -44,8 +44,7 @@ properties.
     `nlsolve` is not `missing`.
   - `autodiff`: The autodifferentiation algorithm to use to compute the Jacobian if
     `manifold_jacobian` is not specified. This must be specified if `manifold_jacobian` is
-    not specified and `nlsolve` is `missing`. If `nlsolve` is not `missing`, then
-    `autodiff` is ignored.
+    not specified.
   - `manifold_jacobian`: The Jacobian of the manifold (wrt the state). This has the same
     signature as `manifold` and the first argument is the Jacobian if inplace.
 
@@ -118,13 +117,7 @@ function (proj::ManifoldProjection)(integrator)
     proj.manifold_jacobian !== nothing && (proj.manifold_jacobian.t = integrator.t)
 
     SciMLBase.reinit!(proj.nlcache, integrator.u; integrator.p)
-
-    if proj.nlsolve === missing
-        _, u, retcode = SciMLBase.solve!(proj.nlcache)
-    else
-        sol = SciMLBase.solve!(proj.nlcache)
-        (; u, retcode) = sol
-    end
+    _, u, retcode = SciMLBase.solve!(proj.nlcache)
 
     if !SciMLBase.successful_retcode(retcode)
         SciMLBase.terminate!(integrator, retcode)
@@ -146,17 +139,17 @@ function initialize_manifold_projection(affect!::ManifoldProjection, u, t, integ
             (affect!.manifold_jacobian.autonomous = autonomous)
     end
 
+    affect!.manifold.t = t
+    affect!.manifold_jacobian !== nothing && (affect!.manifold_jacobian.t = t)
+
     if affect!.nlsolve === missing
-        affect!.manifold.t = t
-        affect!.manifold_jacobian !== nothing && (affect!.manifold_jacobian.t = t)
         cache = init_manifold_projection(
             Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff,
             affect!.manifold_jacobian, u, integrator.p; affect!.kwargs...)
     else
-        # nlfunc = NonlinearFunction{iip}(affect!.g; affect!.resid_prototype)
-        # nlprob = NonlinearProblem(nlfunc, u, integrator.p)
-        # affect!.nlcache = init(nlprob, affect!.nlsolve; affect!.kwargs...)
-        error("Not Implemented")
+        cache = init_manifold_projection_nonlinear_problem(
+            Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff,
+            affect!.manifold_jacobian, u, integrator.p, affect!.nlsolve; affect!.kwargs...)
     end
     affect!.nlcache = cache
     u_modified!(integrator, false)
@@ -187,6 +180,97 @@ function (f::UntypedNonAutonomousFunction)(res, u, p)
 end
 (f::UntypedNonAutonomousFunction)(u, p) = f.autonomous ? f.f(u, p) : f.f(u, p, f.t)
 
+# This is solving the langrange multiplier formulation. This is more accurate but at the
+# same time significantly more expensive.
+@concrete mutable struct NonlinearSolveManifoldProjectionCache{iip}
+    manifold
+    p
+    λ
+    z
+    ũ
+    gu_cache
+    nlcache
+
+    first_call::Bool
+    J
+    manifold_jacobian
+    autodiff
+    di_extras
+end
+
+function SciMLBase.reinit!(
+        cache::NonlinearSolveManifoldProjectionCache{iip}, u; p = cache.p) where {iip}
+    if !cache.first_call || (cache.ũ !== u || cache.p !== p)
+        compute_manifold_jacobian!(cache.J, cache.manifold_jacobian, cache.autodiff,
+            Val(iip), cache.manifold, cache.gu_cache, u, p, cache.di_extras)
+    end
+    cache.first_call = false
+    cache.ũ = u
+    cache.p = p
+
+    cache.z[1:length(cache.λ)] .= false
+    cache.z[(length(cache.λ) + 1):end] .= vec(u)
+    SciMLBase.reinit!(cache.nlcache, cache.z; p = (u, cache.J, p))
+end
+
+function init_manifold_projection_nonlinear_problem(
+        IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ, p, alg;
+        resid_prototype = nothing, kwargs...) where {iip}
+    if iip
+        if resid_prototype !== nothing
+            gu = similar(resid_prototype)
+            λ = similar(resid_prototype)
+        else
+            @warn "`resid_prototype` not provided for in-place problem. Assuming size of \
+                   output is the same as input. This might be incorrect." maxlog=1
+            gu = similar(ũ)
+            λ = similar(ũ)
+        end
+    else
+        gu = nothing
+        λ = manifold(ũ, p)
+    end
+
+    J, di_extras = setup_manifold_jacobian(manifold_jacobian, autodiff, IIP, manifold,
+        gu, ũ, p)
+    z = vcat(vec(λ), vec(ũ))
+
+    nlfunc = if iip
+        let λlen = length(λ), λsz = size(λ), zsz = size(ũ)
+            @views (resid, u, ps) -> begin
+                ũ2, J2, p2 = ps
+                λ2, z2 = u[1:λlen], u[(λlen + 1):end]
+                manifold(reshape(resid[1:λlen], λsz), reshape(z2, zsz), p2)
+                resid[(λlen + 1):end] .= z2 .- vec(ũ2) .+ vec(vec(J2' * λ2))
+            end
+        end
+    else
+        let λlen = length(λ), zsz = size(ũ)
+            @views (u, ps) -> begin
+                ũ2, J2, p2 = ps
+                λ2, z2 = u[1:λlen], u[(λlen + 1):end]
+                gz = vec(manifold(reshape(z2, zsz), p2))
+                resid = z2 .- vec(ũ2) .+ vec(J2' * λ2)
+                return vcat(gz, resid)
+            end
+        end
+    end
+
+    nlprob = NonlinearProblem(NonlinearFunction{iip}(nlfunc), z, (ũ, J, p))
+    nlcache = SciMLBase.init(nlprob, alg; kwargs...)
+
+    return NonlinearSolveManifoldProjectionCache{iip}(
+        manifold, p, λ, z, ũ, gu, nlcache, true, J, manifold_jacobian, autodiff, di_extras)
+end
+
+@views function SciMLBase.solve!(cache::NonlinearSolveManifoldProjectionCache{iip}) where {iip}
+    sol = SciMLBase.solve!(cache.nlcache)
+    (; u, retcode) = sol
+    λ = reshape(u[1:length(cache.λ)], size(cache.λ))
+    ũ = reshape(u[(length(cache.λ) + 1):end], size(cache.ũ))
+    return λ, ũ, retcode
+end
+
 # This is the algorithm described in Hairer III.
 @concrete mutable struct SingleFactorizeManifoldProjectionCache{iip}
     manifold
@@ -225,7 +309,7 @@ default_abstol(::Type{T}) where {T} = real(oneunit(T)) * (eps(real(one(T))))^(4
 
 function init_manifold_projection(IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ,
         p; abstol = default_abstol(eltype(ũ)), maxiters = 1000,
-        resid_prototype = nothing) where {iip}
+        resid_prototype = nothing, kwargs...) where {iip}
     if iip
         if resid_prototype !== nothing
             gu = similar(resid_prototype)
@@ -309,6 +393,11 @@ function setup_manifold_jacobian(
     return J, di_extras
 end
 
+function setup_manifold_jacobian(
+        ::Nothing, ::Nothing, ::Val{iip}, manifold, gu, ũ, p) where {iip}
+    error("`autodiff` is set to `nothing` and analytic manifold jacobian is not provided.")
+end
+
 function compute_manifold_jacobian!(J, manifold_jacobian, autodiff, ::Val{iip},
         manifold, gu, ũ, p, di_extras) where {iip}
     if iip
@@ -329,10 +418,6 @@ function compute_manifold_jacobian!(J, ::Nothing, autodiff, ::Val{iip}, manifold
     return J
 end
 
-function setup_manifold_jacobian(::Nothing, ::Nothing, args...)
-    error("`autodiff` is set to `nothing` and analytic manifold jacobian is not provided.")
-end
-
 function safe_factorize!(A::AbstractMatrix)
     if issquare(A)
         fact = LinearAlgebra.cholesky(A; check = false)
diff --git a/test/manifold_tests.jl b/test/manifold_tests.jl
index 41d16575..40fc6633 100644
--- a/test/manifold_tests.jl
+++ b/test/manifold_tests.jl
@@ -32,20 +32,27 @@ solve(prob, Vern7(), callback = cb_t)
 
 # autodiff=false
 cb_false = ManifoldProjection(
-    g; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2))
+    g; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2),
+    autodiff = AutoFiniteDiff())
 solve(prob, Vern7(), callback = cb_false)
 sol = solve(prob, Vern7(), callback = cb_false)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
 
 cb_t_false = ManifoldProjection(g_t,
-    nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2))
+    nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2),
+    autodiff = AutoFiniteDiff())
 solve(prob, Vern7(), callback = cb_t_false)
 sol_t = solve(prob, Vern7(), callback = cb_t_false)
 @test sol_t.u == sol.u && sol_t.t == sol.t
 
 # test array partitions
+function f_ap!(du, u, p, t)
+    du[1:2] .= u[3:4]
+    du[3:4] .= u[1:2]
+end
+
 u₀ = ArrayPartition(ones(2), ones(2))
-prob = ODEProblem(f, u₀, (0.0, 100.0))
+prob = ODEProblem(f_ap!, u₀, (0.0, 100.0))
 
 sol = solve(prob, Vern7(), callback = cb)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
@@ -71,6 +78,12 @@ sol = solve(prob, Vern7(), callback = cb_unsat)
 @test !SciMLBase.successful_retcode(sol)
 @test last(sol.t) != 100.0
 
+cb_unsat = ManifoldProjection(
+    g_unsat; resid_prototype = zeros(2), autodiff = AutoForwardDiff(), nlsolve = NewtonRaphson())
+sol = solve(prob, Vern7(), callback = cb_unsat)
+@test !SciMLBase.successful_retcode(sol)
+@test last(sol.t) != 100.0
+
 # Tests for OOP Manifold Projection
 function g_oop(u, p)
     return [u[2]^2 + u[1]^2 - 2
@@ -98,20 +111,22 @@ solve(prob, Vern7(), callback = cb_t)
 
 # autodiff=false
 cb_false = ManifoldProjection(
-    g_oop; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), isinplace = Val(false))
+    g_oop; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), autodiff = AutoFiniteDiff())
 solve(prob, Vern7(), callback = cb_false)
 sol = solve(prob, Vern7(), callback = cb_false)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2
 
 cb_t_false = ManifoldProjection(g_oop_t,
-    nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), isinplace = Val(false))
+    nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), autodiff = AutoFiniteDiff())
 solve(prob, Vern7(), callback = cb_t_false)
 sol_t = solve(prob, Vern7(), callback = cb_t_false)
 @test sol_t.u == sol.u && sol_t.t == sol.t
 
 # test array partitions
+f_ap(u, p, t) = ArrayPartition(u[3:4], u[1:2])
+
 u₀ = ArrayPartition(ones(2), ones(2))
-prob = ODEProblem(f, u₀, (0.0, 100.0))
+prob = ODEProblem(f_ap, u₀, (0.0, 100.0))
 
 sol = solve(prob, Vern7(), callback = cb)
 @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2

From cd67a515d65d6433f08e668fbb6fa78c86532fe1 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 13:10:45 -0400
Subject: [PATCH 6/9] fix: update GeneralDomain to the new ManifoldProjection
 API

---
 docs/Project.toml      |   2 +
 docs/src/projection.md |   5 +-
 src/domain.jl          | 132 ++++++++++++++++++++++++-----------------
 src/manifold.jl        |  41 ++++++-------
 test/domain_tests.jl   |  14 ++++-
 5 files changed, 112 insertions(+), 82 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index 759ff0db..70dbd1e3 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -6,6 +7,7 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [compat]
+ADTypes = "1.9.0"
 DiffEqCallbacks = "3"
 Documenter = "1"
 OrdinaryDiffEq = "6.88"
diff --git a/docs/src/projection.md b/docs/src/projection.md
index 31b08bc2..53289c50 100644
--- a/docs/src/projection.md
+++ b/docs/src/projection.md
@@ -12,7 +12,7 @@ ManifoldProjection
 Here we solve the harmonic oscillator:
 
 ```@example manifold
-using OrdinaryDiffEq, DiffEqCallbacks, Plots
+using OrdinaryDiffEq, DiffEqCallbacks, Plots, ADTypes
 
 u0 = ones(2)
 function f(du, u, p, t)
@@ -28,14 +28,13 @@ to conserve the sum of squares:
 ```@example manifold
 function g(resid, u, p, t)
     resid[1] = u[2]^2 + u[1]^2 - 2
-    resid[2] = 0
 end
 ```
 
 To build the callback, we just call
 
 ```@example manifold
-cb = ManifoldProjection(g)
+cb = ManifoldProjection(g; autodiff = AutoForwardDiff(), resid_prototype = zeros(1))
 ```
 
 Using this callback, the Runge-Kutta method `Vern7` conserves energy. Note that the
diff --git a/src/domain.jl b/src/domain.jl
index d9879886..581a18f8 100644
--- a/src/domain.jl
+++ b/src/domain.jl
@@ -2,33 +2,35 @@
 
 abstract type AbstractDomainAffect{T, S, uType} end
 
+(f::AbstractDomainAffect)(integrator) = affect!(integrator, f)
+
 struct PositiveDomainAffect{T, S, uType} <: AbstractDomainAffect{T, S, uType}
     abstol::T
     scalefactor::S
     u::uType
 end
 
-struct GeneralDomainAffect{autonomous, F, T, S, uType} <: AbstractDomainAffect{T, S, uType}
+struct GeneralDomainAffect{F <: AbstractNonAutonomousFunction, T, S, uType, A} <:
+       AbstractDomainAffect{T, S, uType}
     g::F
     abstol::T
     scalefactor::S
     u::uType
     resid::uType
+    autonomous::A
+end
 
-    function GeneralDomainAffect{autonomous}(g::F, abstol::T, scalefactor::S, u::uType,
-            resid::uType) where {autonomous, F, T, S, uType
-    }
-        new{autonomous, F, T, S, uType}(g, abstol, scalefactor, u, resid)
+function initialize_general_domain_affect(cb, u, t, integrator)
+    return initialize_general_domain_affect(cb.affect!, u, t, integrator)
+end
+function initialize_general_domain_affect(affect!::GeneralDomainAffect, u, t, integrator)
+    if affect!.autonomous === nothing
+        autonomous = maximum(SciMLBase.numargs(affect!.g.f)) ==
+                     2 + SciMLBase.isinplace(integrator.f)
+        affect!.g.autonomous = autonomous
     end
 end
 
-# definitions of callback functions
-
-# Workaround since it is not possible to add methods to an abstract type:
-# https://github.com/JuliaLang/julia/issues/14919
-(f::PositiveDomainAffect)(integrator) = affect!(integrator, f)
-(f::GeneralDomainAffect)(integrator) = affect!(integrator, f)
-
 # general method definitions for domain callbacks
 
 """
@@ -41,6 +43,8 @@ function affect!(integrator, f::AbstractDomainAffect{T, S, uType}) where {T, S,
         throw(ArgumentError("domain callback can only be applied to adaptive algorithms"))
     end
 
+    iip = Val(SciMLBase.isinplace(integrator.f))
+
     # define array of next time step, absolute tolerance, and scale factor
     if uType <: Nothing
         if integrator.u isa Union{Number, StaticArraysCore.SArray}
@@ -55,7 +59,7 @@ function affect!(integrator, f::AbstractDomainAffect{T, S, uType}) where {T, S,
     scalefactor = S <: Nothing ? 1 // 2 : f.scalefactor
 
     # setup callback and save additional arguments for checking next time step
-    args = setup(f, integrator)
+    args = setup(f, integrator, iip)
 
     # obtain proposed next time step
     dt = get_proposed_dt(integrator)
@@ -80,7 +84,7 @@ function affect!(integrator, f::AbstractDomainAffect{T, S, uType}) where {T, S,
         end
 
         # check whether time step is accepted
-        isaccepted(u, p, t, abstol, f, args...) && break
+        isaccepted(u, p, t, abstol, f, iip, args...) && break
 
         # reduce time step
         dtcache = dt
@@ -120,12 +124,12 @@ was modified.
 modify_u!(integrator, ::AbstractDomainAffect) = false
 
 """
-    setup(f::AbstractDomainAffect, integrator)
+    setup(f::AbstractDomainAffect, integrator, ::Val{iip}) where {iip}
 
 Setup callback `f` and return an arbitrary tuple whose elements are used as additional
 arguments in checking whether time step is accepted.
 """
-setup(::AbstractDomainAffect, integrator) = ()
+setup(::AbstractDomainAffect, integrator, ::Val{iip}) where {iip} = ()
 
 """
     isaccepted(u, abstol, f::AbstractDomainAffect, args...)
@@ -133,7 +137,7 @@ setup(::AbstractDomainAffect, integrator) = ()
 Return whether `u` is an acceptable state vector at the next time point given absolute
 tolerance `abstol`, callback `f`, and other optional arguments.
 """
-isaccepted(u, p, t, tolerance, ::AbstractDomainAffect, args...) = true
+isaccepted(u, p, t, tolerance, ::AbstractDomainAffect, ::Val{iip}, args...) where {iip} = true
 
 # specific method definitions for positive domain callback
 
@@ -175,27 +179,30 @@ function _set_neg_zero!(integrator, u::StaticArraysCore.SArray)
 end
 
 # state vector is accepted if its entries are greater than -abstol
-isaccepted(u, p, t, abstol::Number, ::PositiveDomainAffect) = all(ui -> ui > -abstol, u)
-function isaccepted(u, p, t, abstol, ::PositiveDomainAffect)
+function isaccepted(u, p, t, abstol::Number, ::PositiveDomainAffect, ::Val{iip}) where {iip}
+    return all(ui -> ui > -abstol, u)
+end
+function isaccepted(u, p, t, abstol, ::PositiveDomainAffect, ::Val{iip}) where {iip}
     length(u) == length(abstol) ||
         throw(DimensionMismatch("numbers of states and tolerances do not match"))
-    all(ui > -tol for (ui, tol) in zip(u, abstol))
+    return all(ui > -tol for (ui, tol) in zip(u, abstol))
 end
 
 # specific method definitions for general domain callback
 
 # create array of residuals
-function setup(f::GeneralDomainAffect, integrator)
-    f.resid isa Nothing ? (similar(integrator.u),) : (f.resid,)
+setup(f::GeneralDomainAffect, integrator, ::Val{false}) = (nothing,)
+function setup(f::GeneralDomainAffect, integrator, ::Val{true})
+    return f.resid === nothing ? (similar(integrator.u),) : (f.resid,)
 end
 
-function isaccepted(u, p, t, abstol, f::GeneralDomainAffect{autonomous, F, T, S, uType},
-        resid) where {autonomous, F, T, S, uType}
+function isaccepted(u, p, t, abstol, f::GeneralDomainAffect, ::Val{iip}, resid) where {iip}
     # calculate residuals
-    if autonomous
+    f.g.t = t
+    if iip
         f.g(resid, u, p)
     else
-        f.g(resid, u, p, t)
+        resid = f.g(u, p)
     end
 
     # accept time step if residuals are smaller than the tolerance
@@ -214,26 +221,32 @@ end
 """
     GeneralDomain(
         g, u = nothing; save = true, abstol = nothing, scalefactor = nothing,
-        autonomous = maximum(SciMLBase.numargs(g)) == 3, nlsolve_kwargs = (;
-            abstol = 10 * eps()), kwargs...)
+        autonomous = nothing, domain_jacobian = nothing,
+        nlsolve_kwargs = (; abstol = 10 * eps()), kwargs...)
 
 A `GeneralDomain` callback in DiffEqCallbacks.jl generalizes the concept of
-a `PositiveDomain` callback to arbitrary domains. Domains are specified by
-in-place functions `g(resid, u, p)` or `g(resid, u, p, t)` that calculate residuals of a
-state vector `u` at time `t` relative to that domain, with `p` the parameters of the
-corresponding integrator. As for `PositiveDomain`, steps are accepted if residuals
-of the extrapolated values at the next time step are below
-a certain tolerance. Moreover, this callback is automatically coupled with a
-`ManifoldProjection` that keeps all calculated state vectors close to the desired
-domain, but in contrast to a `PositiveDomain` callback the nonlinear solver in a
-`ManifoldProjection` cannot guarantee that all state vectors of the solution are
-actually inside the domain. Thus, a `PositiveDomain` callback should generally be
-preferred.
+a `PositiveDomain` callback to arbitrary domains.
+
+Domains are specified by
+  - in-place functions `g(resid, u, p)` or `g(resid, u, p, t)` if the corresponding
+    ODEProblem is an inplace problem, or
+  - out-of-place functions `g(u, p)` or `g(u, p, t)` if the corresponding ODEProblem is
+    an out-of-place problem.
+
+The function calculates residuals of a state vector `u` at time `t` relative to that domain,
+with `p` the parameters of the corresponding integrator.
+
+As for `PositiveDomain`, steps are accepted if residuals of the extrapolated values at the
+next time step are below a certain tolerance. Moreover, this callback is automatically
+coupled with a `ManifoldProjection` that keeps all calculated state vectors close to the
+desired domain, but in contrast to a `PositiveDomain` callback the nonlinear solver in a
+`ManifoldProjection` cannot guarantee that all state vectors of the solution are actually
+inside the domain. Thus, a `PositiveDomain` callback should generally be preferred.
 
 ## Arguments
 
-  - `g`: the implicit definition of the domain as a function `g(resid, u, p)` or
-    `g(resid, u, p, t)` which is zero when the value is in the domain.
+  - `g`: the implicit definition of the domain as a function as described above which is
+    zero when the value is in the domain.
   - `u`: A prototype of the state vector of the integrator. A copy of it is saved and
     extrapolated values are written to it. If it is not specified,
     every application of the callback allocates a new copy of the state vector.
@@ -248,9 +261,13 @@ preferred.
     specified, time steps are halved.
   - `autonomous`: Whether `g` is an autonomous function of the form `g(resid, u, p)`.
     If it is not specified, it is determined automatically.
-  - `kwargs`: All other keyword arguments are passed to `ManifoldProjection`.
+  - `kwargs`: All other keyword arguments are passed to [`ManifoldProjection`](@ref).
   - `nlsolve_kwargs`: All keyword arguments are passed to the nonlinear solver in
     `ManifoldProjection`. The default is `(; abstol = 10 * eps())`.
+  - `domain_jacobian`: The Jacobian of the domain (wrt the state). This has the same
+    signature as `g` and the first argument is the Jacobian if inplace. This corresponds to
+    the `manifold_jacobian` argument of [`ManifoldProjection`](@ref). Note that passing
+    a `manifold_jacobian` is not supported for `GeneralDomain` and results in an error.
 
 ## References
 
@@ -260,20 +277,27 @@ Non-negative solutions of ODEs. Applied Mathematics and Computation 170
 """
 function GeneralDomain(
         g, u = nothing; save = true, abstol = nothing, scalefactor = nothing,
-        autonomous = maximum(SciMLBase.numargs(g)) == 3, nlsolve_kwargs = (;
-            abstol = 10 * eps()), kwargs...)
-    _autonomous = SciMLBase._unwrap_val(autonomous)
-    if u isa Nothing
-        affect! = GeneralDomainAffect{_autonomous}(g, abstol, scalefactor, nothing, nothing)
+        autonomous = nothing, domain_jacobian = nothing, manifold_jacobian = missing,
+        nlsolve_kwargs = (; abstol = 10 * eps()), kwargs...)
+    if manifold_jacobian !== missing
+        throw(ArgumentError("`manifold_jacobian` is not supported for `GeneralDomain`. \
+                             Use `domain_jacobian` instead."))
+    end
+    manifold_projection = ManifoldProjection(
+        g; save = false, autonomous, manifold_jacobian = domain_jacobian,
+        kwargs..., nlsolve_kwargs...)
+    domain = wrap_autonomous_function(autonomous, g)
+    domain_jacobian = wrap_autonomous_function(autonomous, domain_jacobian)
+    affect! = if u === nothing
+        GeneralDomainAffect(domain, abstol, scalefactor, nothing, nothing, autonomous)
     else
-        affect! = GeneralDomainAffect{_autonomous}(g, abstol, scalefactor, deepcopy(u),
-            deepcopy(u))
+        GeneralDomainAffect(
+            domain, abstol, scalefactor, deepcopy(u), deepcopy(u), autonomous)
     end
-    condition = (u, t, integrator) -> true
-    CallbackSet(
-        ManifoldProjection(
-            g; save = false, autonomous, isinplace = Val(true), kwargs..., nlsolve_kwargs...),
-        DiscreteCallback(condition, affect!; save_positions = (false, save)))
+    domain_cb = DiscreteCallback(
+        Returns(true), affect!; initialize = initialize_general_domain_affect,
+        save_positions = (false, save))
+    return CallbackSet(manifold_projection, domain_cb)
 end
 
 @doc doc"""
diff --git a/src/manifold.jl b/src/manifold.jl
index 74eef9db..9d427d40 100644
--- a/src/manifold.jl
+++ b/src/manifold.jl
@@ -31,7 +31,8 @@ properties.
     would work in most cases (See [1] for details). Alternatively, a nonlinear solver as
     defined in the
     [NonlinearSolve.jl format](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/)
-    can be specified.
+    can be specified. Additionally if NonlinearSolve.jl is loaded and `nothing` is specified
+    a polyalgorithm is used.
   - `save`: Whether to do the standard saving (applied after the callback)
   - `autonomous`: Whether `g` is an autonomous function of the form `g(resid, u, p)` or
     `g(u, p)`. Specify it as `Val(::Bool)` to disable runtime branching. If `nothing`,
@@ -88,25 +89,8 @@ end
 
 function ManifoldProjection(
         manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous)
-    if autonomous isa Val{true} || autonomous isa Val{false}
-        wrapped_manifold = TypedNonAutonomousFunction{SciMLBase._unwrap_val(autonomous)}(
-            manifold, nothing)
-        wrapped_manifold_jacobian = if manifold_jacobian === nothing
-            nothing
-        else
-            TypedNonAutonomousFunction{SciMLBase._unwrap_val(autonomous)}(
-                manifold_jacobian, nothing)
-        end
-        autonomous = SciMLBase._unwrap_val(autonomous)
-    else
-        _autonomous = autonomous === nothing ? false : autonomous
-        wrapped_manifold = UntypedNonAutonomousFunction(_autonomous, manifold, nothing)
-        wrapped_manifold_jacobian = if manifold_jacobian === nothing
-            nothing
-        else
-            UntypedNonAutonomousFunction(_autonomous, manifold_jacobian, nothing)
-        end
-    end
+    wrapped_manifold = wrap_autonomous_function(autonomous, manifold)
+    wrapped_manifold_jacobian = wrap_autonomous_function(autonomous, manifold_jacobian)
     return ManifoldProjection(wrapped_manifold, wrapped_manifold_jacobian,
         autodiff, nothing, nlsolve, kwargs, autonomous)
 end
@@ -158,7 +142,20 @@ end
 export ManifoldProjection
 
 # wrapper for non-autonomous functions
-@concrete mutable struct TypedNonAutonomousFunction{autonomous}
+function wrap_autonomous_function(autonomous::Union{Val{true}, Val{false}}, g)
+    g === nothing && return nothing
+    return TypedNonAutonomousFunction{SciMLBase._unwrap_val(autonomous)}(g, nothing)
+end
+function wrap_autonomous_function(autonomous::Union{Bool, Nothing}, g)
+    g === nothing && return nothing
+    autonomous = autonomous === nothing ? false : autonomous
+    return UntypedNonAutonomousFunction(autonomous, g, nothing)
+end
+
+abstract type AbstractNonAutonomousFunction end
+
+@concrete mutable struct TypedNonAutonomousFunction{autonomous} <:
+                         AbstractNonAutonomousFunction
     f
     t::Any
 end
@@ -169,7 +166,7 @@ end
 (f::TypedNonAutonomousFunction{false})(u, p) = f.f(u, p, f.t)
 (f::TypedNonAutonomousFunction{true})(u, p) = f.f(u, p)
 
-@concrete mutable struct UntypedNonAutonomousFunction
+@concrete mutable struct UntypedNonAutonomousFunction <: AbstractNonAutonomousFunction
     autonomous::Bool
     f
     t::Any
diff --git a/test/domain_tests.jl b/test/domain_tests.jl
index 9aaee626..0130b246 100644
--- a/test/domain_tests.jl
+++ b/test/domain_tests.jl
@@ -1,4 +1,4 @@
-using DiffEqCallbacks, OrdinaryDiffEq, Test
+using DiffEqCallbacks, OrdinaryDiffEq, Test, ADTypes, NonlinearSolve
 
 # Non-negative ODE examples
 #
@@ -39,7 +39,11 @@ naive_sol_absval = solve(prob_absval, BS3())
 function g(resid, u, p)
     resid[1] = u[1] < 0 ? -u[1] : 0
 end
-general_sol_absval = solve(prob_absval, BS3(); callback = GeneralDomain(g, [1.0]),
+general_sol_absval = solve(
+    prob_absval, BS3();
+    callback = GeneralDomain(g, [1.0];
+        autodiff = AutoForwardDiff(),
+        nlsolve=NewtonRaphson(; autodiff = AutoForwardDiff())),
     save_everystep = false)
 @test all(x -> x[1] ≥ 0, general_sol_absval.u)
 @test general_sol_absval.errors[:l∞] < 9.9e-5
@@ -49,7 +53,11 @@ general_sol_absval = solve(prob_absval, BS3(); callback = GeneralDomain(g, [1.0]
 # test "non-autonomous" function
 g_t(resid, u, p, t) = g(resid, u, p)
 
-general_t_sol_absval = solve(prob_absval, BS3(); callback = GeneralDomain(g_t, [1.0]),
+general_t_sol_absval = solve(
+    prob_absval, BS3();
+    callback = GeneralDomain(g_t, [1.0];
+        autodiff = AutoForwardDiff(),
+        nlsolve=NewtonRaphson(; autodiff = AutoForwardDiff())),
     save_everystep = false)
 @test general_sol_absval.t ≈ general_t_sol_absval.t
 @test general_sol_absval.u ≈ general_t_sol_absval.u

From 25edaa4f18ec34f34df21b47275ebf77885ac7ba Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 13:25:13 -0400
Subject: [PATCH 7/9] fix: update minimum compats

---
 .github/workflows/Downgrade.yml |  1 -
 Project.toml                    | 18 +++++++++---------
 src/manifold.jl                 |  8 ++++----
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml
index cf17c8d8..48d3b873 100644
--- a/.github/workflows/Downgrade.yml
+++ b/.github/workflows/Downgrade.yml
@@ -22,7 +22,6 @@ jobs:
         with:
           version: ${{ matrix.version }}
       - uses: julia-actions/julia-downgrade-compat@v1
-#        if: ${{ matrix.version == '1.6' }}
         with:
           skip: Pkg,TOML
       - uses: julia-actions/julia-buildpkg@v1
diff --git a/Project.toml b/Project.toml
index 51ab1d47..9b425732 100644
--- a/Project.toml
+++ b/Project.toml
@@ -22,25 +22,25 @@ Aqua = "0.8"
 ConcreteStructs = "0.2.3"
 DataInterpolations = "4.6, 5, 6"
 DataStructures = "0.18.13"
-DiffEqBase = "6.154"
+DiffEqBase = "6.155.3"
 DifferentiationInterface = "0.6.1"
 ForwardDiff = "0.10.36"
 Functors = "0.4"
 LinearAlgebra = "1.10"
 Markdown = "1.10"
-NonlinearSolve = "3.7.2"
-ODEProblemLibrary = "0.1.5"
+NonlinearSolve = "3.14"
+ODEProblemLibrary = "0.1.8"
 OrdinaryDiffEq = "6.88"
 QuadGK = "2.9"
 RecipesBase = "1.3.4"
-RecursiveArrayTools = "3.9"
-SciMLBase = "2.26"
-SciMLSensitivity = "7.56"
-StaticArrays = "1.8"
+RecursiveArrayTools = "3.27"
+SciMLBase = "2.54"
+SciMLSensitivity = "7.68"
+StaticArrays = "1.9.7"
 StaticArraysCore = "1.4"
 Sundials = "4.19.2"
-Test = "1"
-Tracker = "0.2.26"
+Test = "1.10"
+Tracker = "0.2.35"
 Zygote = "0.6.69"
 julia = "1.10"
 
diff --git a/src/manifold.jl b/src/manifold.jl
index 9d427d40..6b593da7 100644
--- a/src/manifold.jl
+++ b/src/manifold.jl
@@ -418,19 +418,19 @@ end
 function safe_factorize!(A::AbstractMatrix)
     if issquare(A)
         fact = LinearAlgebra.cholesky(A; check = false)
-        fact_sucessful(fact) && return fact
+        fact_successful(fact) && return fact
     elseif size(A, 1) > size(A, 2)
         fact = LinearAlgebra.qr(A)
-        fact_sucessful(fact) && return fact
+        fact_successful(fact) && return fact
     end
     return LinearAlgebra.qr!(A, LinearAlgebra.ColumnNorm())
 end
 
-function fact_sucessful(F::LinearAlgebra.QRCompactWY)
+function fact_successful(F::LinearAlgebra.QRCompactWY)
     m, n = size(F)
     U = view(F.factors, 1:min(m, n), 1:n)
     return all(!iszero, Iterators.reverse(@view U[diagind(U)]))
 end
-function fact_sucessful(F::FT) where {FT}
+function fact_successful(F::FT) where {FT}
     return hasmethod(LinearAlgebra.issuccess, (FT,)) ? LinearAlgebra.issuccess(F) : true
 end

From 014ba51317b65acdfb326c72cfc8f1bd2c137145 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 13:34:44 -0400
Subject: [PATCH 8/9] chore: run formatter

---
 src/iterative_and_periodic.jl | 3 ++-
 test/domain_tests.jl          | 4 ++--
 test/preset_time.jl           | 3 ++-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/iterative_and_periodic.jl b/src/iterative_and_periodic.jl
index a64c3e48..05166174 100644
--- a/src/iterative_and_periodic.jl
+++ b/src/iterative_and_periodic.jl
@@ -177,7 +177,8 @@ end
     # Checking for floating point equality is OK here as `DifferentialEquations.jl`
     # sets the time exactly to the final time in the last iteration
     return integrator.t == last(integrator.sol.prob.tspan) ||
-           (hasfield(typeof(integrator), :iter) && (integrator.iter == integrator.opts.maxiters))
+           (hasfield(typeof(integrator), :iter) &&
+            (integrator.iter == integrator.opts.maxiters))
 end
 
 export PeriodicCallback
diff --git a/test/domain_tests.jl b/test/domain_tests.jl
index 0130b246..6777aa07 100644
--- a/test/domain_tests.jl
+++ b/test/domain_tests.jl
@@ -43,7 +43,7 @@ general_sol_absval = solve(
     prob_absval, BS3();
     callback = GeneralDomain(g, [1.0];
         autodiff = AutoForwardDiff(),
-        nlsolve=NewtonRaphson(; autodiff = AutoForwardDiff())),
+        nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff())),
     save_everystep = false)
 @test all(x -> x[1] ≥ 0, general_sol_absval.u)
 @test general_sol_absval.errors[:l∞] < 9.9e-5
@@ -57,7 +57,7 @@ general_t_sol_absval = solve(
     prob_absval, BS3();
     callback = GeneralDomain(g_t, [1.0];
         autodiff = AutoForwardDiff(),
-        nlsolve=NewtonRaphson(; autodiff = AutoForwardDiff())),
+        nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff())),
     save_everystep = false)
 @test general_sol_absval.t ≈ general_t_sol_absval.t
 @test general_sol_absval.u ≈ general_t_sol_absval.u
diff --git a/test/preset_time.jl b/test/preset_time.jl
index bf9b4a23..be082135 100644
--- a/test/preset_time.jl
+++ b/test/preset_time.jl
@@ -91,4 +91,5 @@ sol2 = solve(prob2, Tsit5(), callback = cb1)
 @test sol2(48.0 + eps(48.0)) ≈ [10.0]
 
 _some_test_func(integrator) = u_modified!(integrator, false)
-@inferred PresetTimeCallback(collect(range(0, 10, 100)), _some_test_func, save_positions=(false, false))
+@inferred PresetTimeCallback(
+    collect(range(0, 10, 100)), _some_test_func, save_positions = (false, false))

From 06b4bd85246c787765879d1af27ffdcd74ed38e3 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 15:17:59 -0400
Subject: [PATCH 9/9] chore!: bump version to 4

---
 Project.toml      | 2 +-
 docs/Project.toml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 9b425732..407671e7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "DiffEqCallbacks"
 uuid = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "3.10.0"
+version = "4.0.0"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
diff --git a/docs/Project.toml b/docs/Project.toml
index 70dbd1e3..8688ff38 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -8,7 +8,7 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [compat]
 ADTypes = "1.9.0"
-DiffEqCallbacks = "3"
+DiffEqCallbacks = "4"
 Documenter = "1"
 OrdinaryDiffEq = "6.88"
 Plots = "1.36"