Merge pull request #1031 from SciML/adapt

Add immutable problem types and adapt

Author: ChrisRackauckas

Project.toml

@@ -1,11 +1,12 @@
 name = "SciMLBase"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 authors = ["Chris Rackauckas <[email protected]> and contributors"]
-version = "2.91.1"
+version = "2.92.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"

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
@@ -752,6 +753,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/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