feat: add ImmutableNonlinearProblem

Author: avik-pal

lib/NonlinearSolveBase/Project.toml

@@ -16,10 +16,12 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [weakdeps]
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [extensions]
+NonlinearSolveBaseDiffEqBaseExt = "DiffEqBase"
 NonlinearSolveBaseForwardDiffExt = "ForwardDiff"
 NonlinearSolveBaseSparseArraysExt = "SparseArrays"
 

lib/NonlinearSolveBase/ext/NonlinearSolveBaseDiffEqBaseExt.jl

@@ -0,0 +1,14 @@
+module NonlinearSolveBaseDiffEqBaseExt
+
+using DiffEqBase: DiffEqBase
+using NonlinearSolveBase: ImmutableNonlinearProblem
+
+function DiffEqBase.get_concrete_problem(
+        prob::ImmutableNonlinearProblem, isadapt; kwargs...)
+    u0 = DiffEqBase.get_concrete_u0(prob, isadapt, nothing, kwargs)
+    u0 = DiffEqBase.promote_u0(u0, prob.p, nothing)
+    p = DiffEqBase.get_concrete_p(prob, kwargs)
+    return SciMLBase.remake(prob; u0, p)
+end
+
+end

lib/NonlinearSolveBase/src/NonlinearSolveBase.jl

@@ -7,7 +7,9 @@ using FastClosures: @closure
 using LinearAlgebra: norm
 using Markdown: @doc_str
 using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition
-using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator
+using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinearProblem,
+                 AbstractNonlinearFunction, @add_kwonly, StandardNonlinearProblem,
+                 NullParameters, NonlinearProblem, isinplace
 using StaticArraysCore: StaticArray
 
 include("public.jl")
@@ -25,4 +27,6 @@ export RelTerminationMode, AbsTerminationMode, NormTerminationMode, RelNormTermi
        AbsNormTerminationMode, RelNormSafeTerminationMode, AbsNormSafeTerminationMode,
        RelNormSafeNormTerminationMode, AbsNormSafeNormTerminationMode
 
+export ImmutableNonlinearProblem
+
 end

lib/NonlinearSolveBase/src/immutable_problem.jl

@@ -1 +1,57 @@
+struct ImmutableNonlinearProblem{uType, iip, P, F, K, PT} <:
+       AbstractNonlinearProblem{uType, iip}
+    f::F
+    u0::uType
+    p::P
+    problem_type::PT
+    kwargs::K
 
+    @add_kwonly function ImmutableNonlinearProblem{iip}(
+            f::AbstractNonlinearFunction{iip}, u0, p = NullParameters(),
+            problem_type = StandardNonlinearProblem(); kwargs...) where {iip}
+        if haskey(kwargs, :p)
+            error("`p` specified as a keyword argument `p = $(kwargs[:p])` to \
+                   `NonlinearProblem`. This is not supported.")
+        end
+        warn_paramtype(p)
+        return new{
+            typeof(u0), iip, typeof(p), typeof(f), typeof(kwargs), typeof(problem_type)}(
+            f, u0, p, problem_type, kwargs)
+    end
+
+    """
+    Define a steady state problem using the given function.
+    `isinplace` optionally sets whether the function is inplace or not.
+    This is determined automatically, but not inferred.
+    """
+    function ImmutableNonlinearProblem{iip}(
+            f, u0, p = NullParameters(); kwargs...) where {iip}
+        return ImmutableNonlinearProblem{iip}(NonlinearFunction{iip}(f), u0, p; kwargs...)
+    end
+end
+
+"""
+Define a nonlinear problem using an instance of
+[`AbstractNonlinearFunction`](@ref AbstractNonlinearFunction).
+"""
+function ImmutableNonlinearProblem(
+        f::AbstractNonlinearFunction, u0, p = NullParameters(); kwargs...)
+    return ImmutableNonlinearProblem{isinplace(f)}(f, u0, p; kwargs...)
+end
+
+function ImmutableNonlinearProblem(f, u0, p = NullParameters(); kwargs...)
+    return ImmutableNonlinearProblem(NonlinearFunction(f), u0, p; kwargs...)
+end
+
+"""
+Define a ImmutableNonlinearProblem problem from SteadyStateProblem
+"""
+function ImmutableNonlinearProblem(prob::AbstractNonlinearProblem)
+    return ImmutableNonlinearProblem{isinplace(prob)}(prob.f, prob.u0, prob.p)
+end
+
+function Base.convert(
+        ::Type{ImmutableNonlinearProblem}, prob::T) where {T <: NonlinearProblem}
+    return ImmutableNonlinearProblem{isinplace(prob)}(
+        prob.f, prob.u0, prob.p, prob.problem_type; prob.kwargs...)
+end