Merge branch 'master' into dg/sym_oop

Author: DhairyaLGandhi

Project.toml

@@ -1,11 +1,12 @@
 name = "SciMLBase"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 authors = ["Chris Rackauckas <[email protected]> and contributors"]
-version = "2.91.0"
+version = "2.95.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
@@ -55,6 +56,7 @@ SciMLBaseZygoteExt = ["Zygote", "ChainRulesCore"]
 [compat]
 ADTypes = "0.2.5,1.0.0"
 Accessors = "0.1.36"
+Adapt = "4"
 ArrayInterface = "7.6"
 ChainRules = "1.58.0"
 ChainRulesCore = "1.18"
@@ -83,7 +85,7 @@ RecipesBase = "1.3.4"
 RecursiveArrayTools = "3.27.2"
 Reexport = "1"
 RuntimeGeneratedFunctions = "0.5.12"
-SciMLOperators = "0.4.0, 1"
+SciMLOperators = "0.4, 1.3"
 SciMLStructures = "1.1"
 StableRNGs = "1.0"
 StaticArrays = "1.7"
@@ -101,7 +103,6 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
-Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 RCall = "6f49c342-dc21-5d91-9882-a32aef131414"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -113,4 +114,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "ForwardDiff", "MLStyle", "PartialFunctions", "Pkg", "Plots", "SafeTestsets", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote"]
+test = ["Aqua", "ForwardDiff", "MLStyle", "PartialFunctions", "Pkg", "SafeTestsets", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote"]

docs/src/index.md

@@ -209,7 +209,7 @@ There are too many to name here. Check out the
 ## Contributing
 
   - Please refer to the
-    [SciML ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://github.com/SciML/ColPrac/blob/master/README.md)
+    [SciML ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://github.com/SciML/ColPrac)
     for guidance on PRs, issues, and other matters relating to contributing to SciML.
 
   - See the [SciML Style Guide](https://github.com/SciML/SciMLStyle) for common coding practices and other style decisions.

ext/SciMLBaseChainRulesCoreExt.jl

@@ -116,4 +116,12 @@ function ChainRulesCore.rrule(::SciMLBase.EnsembleSolution, sim, time, converged
     out, EnsembleSolution_adjoint
 end
 
+function ChainRulesCore.rrule(::Type{SciMLBase.IntervalNonlinearProblem}, args...; kwargs...)
+    function IntervalNonlinearProblemAdjoint(ȳ)
+        (NoTangent(), ȳ.f, ȳ.tspan, ȳ.p, ȳ.kwargs, ȳ.problem_type)
+    end
+
+    SciMLBase.IntervalNonlinearProblem(args...; kwargs...), IntervalNonlinearProblemAdjoint
 end
+
+end

src/SciMLBase.jl

@@ -26,6 +26,7 @@ import Accessors: @set, @reset, @delete, @insert
 using Moshi.Data: @data
 using Moshi.Match: @match
 import StaticArraysCore
+import Adapt: adapt_structure, adapt
 
 using Reexport
 using SciMLOperators
@@ -662,6 +663,13 @@ Internal. Used for signifying the AD context comes from a Tracker.jl context.
 """
 struct TrackerOriginator <: ADOriginator end
 
+"""
+$(TYPEDEF)
+
+Internal. Used for signifying the AD context comes from a Mooncake.jl context.
+"""
+struct MooncakeOriginator <: ADOriginator end
+
 include("initialization.jl")
 include("ODE_nlsolve.jl")
 include("utils.jl")
@@ -752,6 +760,8 @@ include("integrator_interface.jl")
 include("remake.jl")
 include("callbacks.jl")
 
+include("adapt.jl")
+
 include("deprecated.jl")
 
 import PrecompileTools

src/adapt.jl

@@ -0,0 +1,21 @@
+function adapt_structure(to, prob::Union{NonlinearProblem{<:Any, <:Any, iip}, ImmutableNonlinearProblem{<:Any, <:Any, iip}}) where {iip}
+    ImmutableNonlinearProblem{iip}(NonlinearFunction{iip}(adapt(to, prob.f.f)),
+        adapt(to, prob.u0),
+        adapt(to, prob.p);
+        adapt(to, prob.kwargs)...)
+end
+
+function adapt_structure(to, prob::Union{ODEProblem{<:Any, <:Any, iip}, ImmutableODEProblem{<:Any, <:Any, iip}}) where {iip}
+    ImmutableODEProblem{iip, FullSpecialize}(adapt(to, prob.f),
+        adapt(to, prob.u0),
+        adapt(to, prob.tspan),
+        adapt(to, prob.p);
+        adapt(to, prob.kwargs)...)
+end
+
+function adapt_structure(to, f::ODEFunction{iip}) where {iip}
+    if f.mass_matrix !== I && f.initialization_data !== nothing
+        error("Adaptation to GPU failed: DAEs of ModelingToolkit currently not supported.")
+    end
+    ODEFunction{iip, FullSpecialize}(f.f, jac = f.jac, mass_matrix = f.mass_matrix)
+end

src/initialization.jl

@@ -44,7 +44,7 @@ struct OverrideInitData{
     function OverrideInitData(initprob::I, update_initprob!::J, initprobmap::K,
             initprobpmap::L, metadata::M, is_update_oop::O) where {I, J, K, L, M, O}
         @assert initprob isa
-                Union{SCCNonlinearProblem, NonlinearProblem, NonlinearLeastSquaresProblem}
+                Union{SCCNonlinearProblem, ImmutableNonlinearProblem, NonlinearProblem, NonlinearLeastSquaresProblem}
         return new{I, J, K, L, M, O}(
             initprob, update_initprob!, initprobmap, initprobpmap, metadata, is_update_oop)
     end

src/problems/linear_problems.jl

@@ -22,7 +22,7 @@ are specified via the `AbstractSciMLOperator` interface. For more details, see
 the [SciMLBase Documentation](https://docs.sciml.ai/SciMLBase/stable/).
 
 Note that matrix-free versions of LinearProblem definitions are not compatible
-with all solvers. To check a solver for compatibility, use the function xxxxx.
+with all solvers. To check a solver for compatibility, use the function `needs_concrete_A(alg::AbstractLinearAlgorithm)`.
 
 ## Problem Type
 

src/problems/nonlinear_problems.jl

@@ -579,7 +579,6 @@ When a keyword argument is `nothing`, the default behaviour of the solver is use
 * `alias_u0::Union{Bool, Nothing}`: alias the `u0` array.
 * `alias::Union{Bool, Nothing}`: sets all fields of the `NonlinearAliasSpecifier` to `alias`. 
 """
-
 struct NonlinearAliasSpecifier <: AbstractAliasSpecifier
     alias_p::Union{Bool, Nothing}
     alias_f::Union{Bool, Nothing}
@@ -596,3 +595,60 @@ struct NonlinearAliasSpecifier <: AbstractAliasSpecifier
         end
     end
 end
+
+struct ImmutableNonlinearProblem{uType, iip, P, F, K, PT} <:
+       AbstractNonlinearProblem{uType, iip}
+    f::F
+    u0::uType
+    p::P
+    problem_type::PT
+    kwargs::K
+
+    SciMLBase.@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
+        SciMLBase.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).
+"""
+function ImmutableNonlinearProblem(
+        f::AbstractNonlinearFunction, u0, p = NullParameters(); kwargs...)
+    return ImmutableNonlinearProblem{SciMLBase.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{SciMLBase.isinplace(prob)}(prob.f, prob.u0, prob.p)
+end
+
+function Base.convert(
+        ::Type{ImmutableNonlinearProblem}, prob::T) where {T <: NonlinearProblem}
+    return ImmutableNonlinearProblem{SciMLBase.isinplace(prob)}(
+        prob.f, prob.u0, prob.p, prob.problem_type; prob.kwargs...)
+end

src/problems/ode_problems.jl

@@ -547,3 +547,92 @@ struct ODEAliasSpecifier <: AbstractAliasSpecifier
         end
     end
 end
+
+struct ImmutableODEProblem{uType, tType, isinplace, P, F, K, PT} <:
+       AbstractODEProblem{uType, tType, isinplace}
+    """The ODE is `du = f(u,p,t)` for out-of-place and f(du,u,p,t) for in-place."""
+    f::F
+    """The initial condition is `u(tspan[1]) = u0`."""
+    u0::uType
+    """The solution `u(t)` will be computed for `tspan[1] ≤ t ≤ tspan[2]`."""
+    tspan::tType
+    """Constant parameters to be supplied as the second argument of `f`."""
+    p::P
+    """A callback to be applied to every solver which uses the problem."""
+    kwargs::K
+    """An internal argument for storing traits about the solving process."""
+    problem_type::PT
+    @add_kwonly function ImmutableODEProblem{iip}(f::AbstractODEFunction{iip},
+            u0, tspan, p = NullParameters(),
+            problem_type = StandardODEProblem();
+            kwargs...) where {iip}
+        _u0 = prepare_initial_state(u0)
+        _tspan = promote_tspan(tspan)
+        warn_paramtype(p)
+        new{typeof(_u0), typeof(_tspan),
+            isinplace(f), typeof(p), typeof(f),
+            typeof(kwargs),
+            typeof(problem_type)}(f,
+            _u0,
+            _tspan,
+            p,
+            kwargs,
+            problem_type)
+    end
+
+    """
+        ImmutableODEProblem{isinplace}(f,u0,tspan,p=NullParameters(),callback=CallbackSet())
+
+    Define an ODE problem with the specified function.
+    `isinplace` optionally sets whether the function is inplace or not.
+    This is determined automatically, but not inferred.
+    """
+    function ImmutableODEProblem{iip}(f,
+            u0,
+            tspan,
+            p = NullParameters();
+            kwargs...) where {iip}
+        _u0 = prepare_initial_state(u0)
+        _tspan = promote_tspan(tspan)
+        _f = ODEFunction{iip, DEFAULT_SPECIALIZATION}(f)
+        ImmutableODEProblem(_f, _u0, _tspan, p; kwargs...)
+    end
+
+    @add_kwonly function ImmutableODEProblem{iip, recompile}(f, u0, tspan,
+            p = NullParameters();
+            kwargs...) where {iip, recompile}
+        ImmutableODEProblem{iip}(ODEFunction{iip, recompile}(f), u0, tspan, p; kwargs...)
+    end
+end
+
+"""
+    ImmutableODEProblem(f::ODEFunction,u0,tspan,p=NullParameters(),callback=CallbackSet())
+
+Define an ODE problem from an [`ODEFunction`](@ref).
+"""
+function ImmutableODEProblem(f::AbstractODEFunction, u0, tspan, args...; kwargs...)
+    ImmutableODEProblem{isinplace(f)}(f, u0, tspan, args...; kwargs...)
+end
+
+function ImmutableODEProblem(f, u0, tspan, p = NullParameters(); kwargs...)
+    iip = isinplace(f, 4)
+    _u0 = prepare_initial_state(u0)
+    _tspan = promote_tspan(tspan)
+    _f = ODEFunction{iip, DEFAULT_SPECIALIZATION}(f)
+    ImmutableODEProblem(_f, _u0, _tspan, p; kwargs...)
+end
+
+staticarray_itize(x) = x
+staticarray_itize(x::Vector) = StaticArraysCore.SVector{length(x)}(x)
+staticarray_itize(x::StaticArraysCore.SizedVector) = StaticArraysCore.SVector{length(x)}(x)
+staticarray_itize(x::Matrix) = StaticArraysCore.SMatrix{size(x)...}(x)
+staticarray_itize(x::StaticArraysCore.SizedMatrix) = StaticArraysCore.SMatrix{size(x)...}(x)
+
+function Base.convert(::Type{ImmutableODEProblem}, prob::T) where {T <: ODEProblem}
+    ImmutableODEProblem(prob.f,
+        staticarray_itize(prob.u0),
+        prob.tspan,
+        staticarray_itize(prob.p),
+        prob.problem_type;
+        prob.kwargs...)
+end

src/retcodes.jl

@@ -412,7 +412,7 @@ EnumX.@enumx ReturnCode begin
 
       - For nonlinear least squares optimizations, this is given for local minima which exceed 
         the chosen tolerance, i.e. `f(x)=resid` where `||resid||>tol` so it's not considered 
-        ReturnCode.Success but it is still considered a sucessful return of the solver since
+        ReturnCode.Success but it is still considered a successful return of the solver since
         it's a valid local minima (and there no minima which achieves the tolerance).
 
     ## Properties