feat: implement NonlinearProblem and NonlinearFunction for System

Author: AayushSabharwal

src/problems/compatibility.jl

@@ -23,6 +23,14 @@ function check_time_dependent(sys::System, T)
     end
 end
 
+function check_time_independent(sys::System, T)
+    if is_time_dependent(sys)
+        throw(SystemCompatibilityError("""
+        `$T` requires a time-independent system.
+        """))
+    end
+end
+
 function check_is_dde(sys::System)
     altT = get_noise_eqs(sys) === nothing ? ODEProblem : SDEProblem
     if !is_dde(sys)

src/problems/nonlinearproblem.jl

@@ -0,0 +1,78 @@
+@fallback_iip_specialize function SciMLBase.NonlinearFunction{iip, spec}(
+        sys::System, _d = nothing, u0 = nothing, p = nothing; jac = false,
+        eval_expression = false, eval_module = @__MODULE__, sparse = false,
+        checkbounds = false, sparsity = false, analytic = nothing,
+        simplify = false, cse = true, initialization_data = nothing,
+        check_compatibility = true, kwargs...) where {iip, spec}
+    check_complete(sys, NonlinearFunction)
+    check_compatibility && check_compatible_system(NonlinearFunction, sys)
+
+    dvs = unknowns(sys)
+    ps = parameters(sys)
+    f = generate_rhs(sys, dvs, ps; expression = Val{false},
+        eval_expression, eval_module, checkbounds = checkbounds, cse,
+        kwargs...)
+
+    if spec === SciMLBase.FunctionWrapperSpecialize && iip
+        if u0 === nothing || p === nothing
+            error("u0, and p must be specified for FunctionWrapperSpecialize on NonlinearFunction.")
+        end
+        f = SciMLBase.wrapfun_iip(f, (u0, u0, p))
+    end
+
+    if jac
+        _jac = generate_jacobian(sys, dvs, ps; expression = Val{false},
+            simplify, sparse, cse, eval_expression, eval_module, checkbounds, kwargs...)
+    else
+        _jac = nothing
+    end
+
+    observedfun = ObservedFunctionCache(
+        sys; steady_state = false, eval_expression, eval_module, checkbounds, cse)
+
+    if length(dvs) == length(equations(sys))
+        resid_prototype = nothing
+    else
+        resid_prototype = calculate_resid_prototype(length(equations(sys)), u0, p)
+    end
+
+    if sparse
+        jac_prototype = similar(calculate_jacobian(sys; sparse), eltype(u0))
+    else
+        jac_prototype = nothing
+    end
+
+    NonlinearFunction{iip, spec}(f;
+        sys = sys,
+        jac = _jac,
+        observed = observedfun,
+        analytic = analytic,
+        jac_prototype,
+        resid_prototype,
+        initialization_data)
+end
+
+@fallback_iip_specialize function SciMLBase.NonlinearProblem{iip, spec}(
+        sys::System, u0map, parammap = SciMLBase.NullParameters();
+        check_length = true, check_compatibility = true, kwargs...) where {iip, spec}
+    check_complete(sys, NonlinearProblem)
+    check_compatibility && check_compatible_system(NonlinearProblem, sys)
+
+    f, u0, p = process_SciMLProblem(NonlinearFunction{iip, spec}, sys, u0map, parammap;
+        check_length, check_compatibility, kwargs...)
+
+    kwargs = process_kwargs(sys; kwargs...)
+    # Call `remake` so it runs initialization if it is trivial
+    return remake(NonlinearProblem{iip}(
+        f, u0, p, StandardNonlinearProblem(); kwargs...))
+end
+
+function check_compatible_system(
+        T::Union{Type{NonlinearFunction}, Type{NonlinearProblem}}, sys::System)
+    check_time_independent(sys, T)
+    check_not_dde(sys)
+    check_no_cost(sys, T)
+    check_no_constraints(sys, T)
+    check_no_jumps(sys, T)
+    check_no_noise(sys, T)
+end

src/systems/codegen.jl

@@ -143,7 +143,9 @@ function generate_jacobian(sys::System, dvs = unknowns(sys),
     jac = calculate_jacobian(sys; simplify, sparse, dvs)
     p = reorder_parameters(sys, ps)
     t = get_iv(sys)
-    if t !== nothing
+    if t === nothing
+        wrap_code = (identity, identity)
+    else
         wrap_code = sparse ? assert_jac_length_header(sys) : (identity, identity)
     end
     res = build_function_wrapper(sys, jac, dvs, p..., t; wrap_code, expression = Val{true},