Creating a ODESolver from DifferentialEquations.jl

[santiagobadia]

Hi @fverdugo

I have been taking a look at DifferentialEquations.jl, in order to create ODESolvers as wrappers of this library solvers.

In our current version of the GridapTimeStepper package, I think that we can do it.

For instance, looking at the DAE solvers, what one requires is a residual, which is a method that we have.

https://docs.sciml.ai/stable/tutorials/dae_example/ https://docs.sciml.ai/stable/types/dae_types/

On the other hand, you can also provide a Jacobian, with a coefficient that comes from the solver, as we have. So, we could easily provide this method too.

https://docs.sciml.ai/stable/features/performance_overloads/

We can also provide our linear solvers (it says nonlinear too, but I cannot find an example)

https://docs.sciml.ai/stable/features/linear_nonlinear/#Linear-Solvers:-linsolve-Specification-1

We could create a solver from DifferentialEquation at every time step, from t_n to t_n+dt. I cannot see how to use the methods in the library as iterators, so they cannot be used for the whole period. Furthermore, it seems they are storing the unknown at all time steps (it can be tuned), which is not what we want.

It is unclear to me how to create a wrapper for their solvers such that we can easily provide all the options of the library in our creational methods for the wrapper.

What do you think?

[santiagobadia]

I am trying to combine the ODE solvers in DifferentialEquations.jl with the PDE solvers in the Gridap.jl library.

It is mandatory for our applications to use the inplace version of the DifferentialEquations.jl package. This way, we can pre-allocate the required arrays in our Gridap solvers.

On the other hand, even when considering ODE systems, we want to use an implicit ODE expression of the solver, since our mass matrices are not just identity matrices, and can be fairly complicated or even nonlinear.

As a result, I think that the way to go is to consider an IIP problem in DifferentialEquations.jl.

I have written the following silly linear example

tspan = (0.0,1.0)
u0 = zeros(1)
u0[1] = 1/2

function residual(res,du,u,p,t)
  res[1] = - 1.01u[1] + du[1]
end

function jacobian(jac,du,u,p,gamma,t)
  jac[1,1] = gamma - 1.01
  nothing
end

f_iip = ODEFunction{false}(residual;jac=jacobian)
prob_iip = ODEProblem{false}(f_iip,u0,tspan)
sol_iip = solve(prob_inplace, ImplicitEuler(), reltol=1e-8, abstol=1e-8)

but it does not work, because it seems the library is looking for an out-of-place version of jacobian, based on the error:

MethodError: no method matching jacobian(::Array{Float64,1}, ::DiffEqBase.NullParameters, ::Float64)
Closest candidates are:
  jacobian(::Any, ::Any, ::Any, !Matched::Any, !Matched::Any, !Matched::Any) at /DifEqWrappers.jl:24
...

In fact, even without considering the Jacobian, the same problem appears with the residual when doing this:

prob_iip = ODEProblem{false}(residual,u0,tspan)
sol_iip = solve(prob_iip, ImplicitEuler(), reltol=1e-8, abstol=1e-8)

Is it possible to do what I want in DifferentialEquations.jl? If Yes, What am I doing wrong? I could not find an example in the documentation covering this (probably my fault).

[santiagobadia]

Instead, when using the DAE machinery

f_iip = DAEFunction{false}(residual;jac=jacobian)
prob_iip = DAEProblem{false}(f_iip,u0,tspan)
sol_iip = solve(prob_iip, ImplicitEuler(), reltol=1e-8, abstol=1e-8)

it returns an error saying that ImplicitEuler is an ODE solver, not a DAE solver

[fverdugo]

We could create a solver from DifferentialEquation at every time step, from t_n to t_n+dt. I cannot see how to use the methods in the library as iterators, so they cannot be used for the whole period. Furthermore, it seems they are storing the unknown at all time steps (it can be tuned), which is not what we want.

Yes, we can always implement the function solve_step! with the DifferentialEquations machinery and use our lazy GenericODESolution to represent all time steps in a lazy way. ~By the way, we only need to create the DifferentialEquations solver once for all time steps.~ EDIT: I am not sure if this is true since the inital condition and the time span are inside the ODESolver. We can always create a new solver at each time step as you suggested.

It is unclear to me how to create a wrapper for their solvers such that we can easily provide all the options of the library in our creational methods for the wrapper.

I need to have a look at the DifferentialEquations library to see their API before to be able to answer this question. If they have a "solver object". A flexible way is that the user builds this object directly in the driver using the DifferentialEquations API and passes it to our wrapper.

[fverdugo]

Is it possible to do what I want in DifferentialEquations.jl? If Yes, What am I doing wrong? I could not find an example in the documentation covering this (probably my fault).

It is possible to use in-place functions (at least for the residual) for systems of ODEs. See here: https://docs.sciml.ai/latest/tutorials/ode_example/#Example-2:-Solving-Systems-of-Equations-1

In the link, they write ODEProblem instead of ODEProblem{false}. I don't know if this is the problem.

[santiagobadia]

The issue is a little bit old, please take a look at this post

https://discourse.julialang.org/t/differentialequations-jl-inplace-implicit-ode-problem/37495

It must be ODEProblem{true} where true means inplace but the tutorial does not apply to us, because there is no mass matrix. But you can consider a non-trivial mass matrix, as replied in the post. I knew this option but still, it is not all we need.

But still, we cannot provide a jacobian so I consider this is not enough. And it assumes a linear mass matrix, which is not always the case. So, up to now, I don't know how to make it work for the most general case.

I like the DAE definition in DifferentialEquations.jl, it really matches our needs, but when using this constructor combined with a ODE algorithm, e.g., ImplicitEuler(), it return an error, because we are trying to solve a DAE with an ODE algorithm. However, this is not always the case, when jacobian_t part is non-singular...

[ChrisRackauckas]

Hey, just found this. I can help you out here.

I cannot see how to use the methods in the library as iterators, so they cannot be used for the whole period

https://docs.sciml.ai/latest/basics/integrator/

Furthermore, it seems they are storing the unknown at all time steps (it can be tuned), which is not what we want.

https://docs.sciml.ai/latest/basics/common_solver_opts/#Output-Control-1 save_everystep=false,save_start=false saves nothing except the final timepoint, so you can save as much or as little as you want.

it says nonlinear too, but I cannot find an example

You probably never want to supply that, which is why we haven't updated support for that in awhile. Most standard nonlinear solver schemes are highly highly inefficient for time-stepping problems because you always want to do some kind of fancy quasi-Newton with Jacobian reuse. We'll re-add support here, but the built-in ones of quasi-Newton, Anderson acceleration, and fixed point are probably what you want anyways (and they do the proper reuse stuff).

It must be ODEProblem{true} where true means inplace but the tutorial does not apply to us, because there is no mass matrix. But you can consider a non-trivial mass matrix, as replied in the post. I knew this option but still, it is not all we need.

Yes, you probably always want to do ODEProblem{true}. I'm not sure why you were setting ODEProblem{false}, or where you even found that.

But still, we cannot provide a jacobian so I consider this is not enough. And it assumes a linear mass matrix, which is not always the case. So, up to now, I don't know how to make it work for the most general case.

Why not set the Jacobian? There's an example in the documentation: https://docs.sciml.ai/latest/features/performance_overloads/#Declaring-Explicit-Jacobians-for-DAEs-1

But still, we cannot provide a jacobian so I consider this is not enough. And it assumes a linear mass matrix, which is not always the case. So, up to now, I don't know how to make it work for the most general case.

It doesn't assume a linear mass matrix. It mentions in the documentation which solvers can handle a state-dependent mass matrix: M(u,p,t)u' = f(u,p,t), and ImplicitEuler is one of the ones in the list.

https://docs.sciml.ai/latest/solvers/dae_solve/#OrdinaryDiffEq.jl-(Mass-Matrix)-1

There's some examples of using it here:

https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/test/interface/mass_matrix_tests.jl#L36-L63

But you can consider a non-trivial mass matrix, as replied in the post. I knew this option but still, it is not all we need.

I'm curious it isn't, because it can express anything that can be expressed in residual form.

---

Path Forward

I think the easiest path forward is to first target the DAEProblem, defining the residual form and the Jacobian, and using IDA() . In this style, your example is just:

using Sundials
tspan = (0.0,1.0)
u0 = zeros(1)
u0[1] = 1/2

function residual(res,du,u,p,t)
  res[1] = - 1.01u[1] + du[1]
end

function jacobian(jac,du,u,p,gamma,t)
  jac[1,1] = gamma - 1.01
  nothing
end

f_iip = DAEFunction{true}(residual;jac=jacobian)
prob_iip = DAEProblem{true}(f_iip,u0,u0,tspan,differential_vars=[true])
sol_iip = solve(prob_iip, IDA(), reltol=1e-8, abstol=1e-8)

or in the integrator interface:

using Sundials
tspan = (0.0,1.0)
u0 = zeros(1)
u0[1] = 1/2

function residual(res,du,u,p,t)
  res[1] = - 1.01u[1] + du[1]
end

function jacobian(jac,du,u,p,gamma,t)
  jac[1,1] = gamma - 1.01
  nothing
end

f_iip = DAEFunction{true}(residual;jac=jacobian)
prob_iip = DAEProblem{true}(f_iip,u0,u0,tspan,differential_vars=[true])
integ = init(prob_iip, IDA(), reltol=1e-8, abstol=1e-8)
step!(integ)
step!(integ)

# Show of using integrators as iterators
using Base.Iterators
for i in take(integ,100)
      @show integ.u
end

The fully implicit form will get you working, and we can then discuss the mass matrix forms to open up new solver possibilities, but this can at least get a version 1 working. The downside with using the Sundials solver here is it doesn't have the full flexibility yet, but we can get it to what you need quite easily. Do you have an example of the residual form I could test on?

[santiagobadia]

Dear @ChrisRackauckas

Thanks a lot for you reply! And sorry for the huge delay answering.

I have followed your suggestions, and created a driver in a branch in which I have used Sundials with Gridap

https://github.com/gridap/GridapODEs.jl/blob/diffeqwrapper/test/TransientFEsTests/DiffEqsWrapperTests.jl

If you have some time, you can take a look and let me know.

At the end of this file, I have listed the main issues I see in the current solution, attached below:

• Issue 1: How do I initialize the jacobian (sparse CSR matrix)? Dense matrices cannot be used in FE codes.

• Issue 2: No control over the (non)linear solver, we would like to be able to provide certainly linear and possibly nonlinear solvers. Let us assume that we just want to consider a fixed point algorithm and we consider an implicit time integration of a nonlinear PDE. Our solvers are efficient since they re-use cache among nonlinear steps (e.g., symbolic factorization, arrays for numerical factorization) and for the transient case in our library, we can also reuse all this between time steps. Could we attain something like this using DifferentialEquations/Sundials?

• Issue 3: Iterator-like version as in GridapODEs. HOWEVER, is this computation lazy? I don't think so, so we need to store all time steps, e.g., for computing time-dependent functionals (drag or lift in CFD, etc), which is going to incur in a lot of memory consumption.

It would be great if you can suggest any idea or recommendation.

Cheers

--Santi

[ChrisRackauckas]

Issue 1: How do I initialize the jacobian (sparse CSR matrix)? Dense matrices cannot be used in FE codes.

f_iip = DAEFunction{true}(residual;jac=jacobian,jac_prototype=mysparsematrix)

See https://docs.sciml.ai/latest/tutorials/advanced_ode_example/ for more tutorials on that stuff.

Issue 2: No control over the (non)linear solver, we would like to be able to provide certainly linear and possibly nonlinear solvers. Let us assume that we just want to consider a fixed point algorithm and we consider an implicit time integration of a nonlinear PDE.

https://docs.sciml.ai/latest/tutorials/advanced_ode_example/ shows swapping out linear solvers. https://docs.sciml.ai/latest/features/linear_nonlinear/ is all of the extra details.

Our solvers are efficient since they re-use cache among nonlinear steps (e.g., symbolic factorization, arrays for numerical factorization) and for the transient case in our library, we can also reuse all this between time steps. Could we attain something like this using DifferentialEquations/Sundials?

All of the DiffEq and Sundials solvers already do that.

Issue 3: Iterator-like version as in GridapODEs. HOWEVER, is this computation lazy? I don't think so, so we need to store all time steps, e.g., for computing time-dependent functionals (drag or lift in CFD, etc), which is going to incur in a lot of memory consumption.

Yes, it's lazy. It doesn't store anything if you do save_everystep=false,save_start=false.

[santiagobadia]

I have worked a little bit on the driver that makes use of Sundials in this file

I have found a couple of issues 1) related to jac_prototype (when I pass my own sparse matrix I get an internal error). If I don't pass the Jacobian prototype, 2) the nonlinear solver does not converge if the dimension of the ODE system is > 1, while the problem I am solving is linear (?). I guess I am doing something wrong.

To be investigated.

[kanav99]

Hi @santiagobadia, I am from SciML. I was having a look, do you have an example of the following issues? Any code snippets which I can change to reproduce the error?

[santiagobadia]

Hi @kanav99

Please take a look at this test

https://github.com/gridap/GridapODEs.jl/blob/diffeqwrapper/test/TransientFEsTests/SimpleDiffEqsTest.jl

When I pass jac_prototype it gets stuck.

[kanav99]

Yes I was having a look at that file. What I wanted to know was what you pass as jac_prototype that it gets stuck. If I could get a script that I can run to reproduce the error, that would be very helpful

[santiagobadia]

I pass a SparseArrays.SparseMatrixCSC{Float64,Int64}, which is J in the code.

[kanav99]

I was trying DImplicitEuler instead of IDA, I am not sure if IDA is compatible with Sparse Matrices (cc: @ChrisRackauckas ). There where some changes needed for OrdinaryDiffEq (https://github.com/SciML/OrdinaryDiffEq.jl/pull/1231) but we need the jacobian function to have 5 (J, du, u, p, t) instead of 6 arguments (gamma is extra)

[santiagobadia]

Let us assume that I have a nonlinear problem A(t,u,u') = 0, and I can provide ∂A/∂u and ∂A/∂(u') or ∂A/∂u + γ*∂A/∂(u'), where γ represents d(Δₜ u)/du, with Δₜ being the discretisation of the time derivative in the time integrator. Without γ, I am not sure how to combine the partial derivative of the operator A wrt u and its time derivative u'. Should I assume a particular value for γ based on the solver I pick? E.g., assume that I am going to use Implicit Euler and take γ=1. It is even worse that that, because γ should depend on the time step size, so I cannot see how to build the Jacobian if γ is extra.

[ChrisRackauckas]

I was trying DImplicitEuler instead of IDA, I am not sure if IDA is compatible with Sparse Matrices (cc: @ChrisRackauckas )

It should be. https://github.com/SciML/Sundials.jl/blob/master/test/common_interface/jacobians.jl#L41-L92

There where some changes needed for OrdinaryDiffEq (SciML/OrdinaryDiffEq.jl#1231) but we need the jacobian function to have 5 (J, du, u, p, t) instead of 6 arguments (gamma is extra)

No, not for DAEs. For DAEs the definition of gamma is required, as defined in https://diffeq.sciml.ai/stable/features/performance_overloads/#Declaring-Explicit-Jacobians-for-DAEs-1 and as @santiagobadia explains it is necessary for the definition of a DAE's Jacobian.

Step one for testing though, just try IDA(linear_solver=:GMRES) and make sure that works (since that's matrix-free). Then try to optimize it and get sparse direct.

[ChrisRackauckas]

I just noticed we only had a test on CVODE with sparse (https://github.com/SciML/Sundials.jl/blob/master/test/common_interface/jacobians.jl#L27) so indeed that could be the issue. But when you choose a sparse matrix, you also have to switch the linear solver in Sundials. Right now Super LU isn't built (https://github.com/SciML/Sundials.jl/issues/249), so you need to use IDA(linear_solver=:KLU). That's more optimized for heterogeneous sparse matrices, so it's not the optimal one for PDEs, but let me get things going on the Super LU artifact building and see if we can get that shipping in the binary to wrap.

[ChrisRackauckas]

BTW @santiagobadia, is there any way for this to get a definition as a DAE in mass matrix form? That would open up a lot more opportunities to test.

[santiagobadia]

Yes, I will do it and let you know.

[kanav99]

Yes, I was confusing between ODEs and DAEs sorry. I just fixed this and I now get a dimension mismatch. The u0 is of size 1 and the jacobian is of size 4x4, why is that?

[santiagobadia]

@kanav99 are working on this file?

I cannot reproduce the error you mention. Everything is dim correct. If you play with n (number cells per dim in 2D square heat eq) you can change the dim of the system.

@ChrisRackauckas I provide the mass matrix now in the same file. But I guess that this is only for explicit solvers. I also provide the stiffness matrix, i.e., the Jacobian of the steady part, in case it can be used for implict methods (e.g., ImplicitEuler).

[ChrisRackauckas]

I provide the mass matrix now in the same file. But I guess that this is only for explicit solvers.

Mass matrices can only be used with implicit methods, e.g. https://diffeq.sciml.ai/stable/solvers/dae_solve/#OrdinaryDiffEq.jl-(Mass-Matrix)-1 . @kanav99 can probably take it from here.

[kanav99]

@santiagobadia I pulled the latest commit and it works now, thanks! And yes, mass matrices work for Implicit methods. I see that IDA() works well now.

[santiagobadia]

Ok, @kanav99

I have cleaned things up. I would say that with the methods I provide you could use play with DifferentialEquations.jl solvers. You can see the different methods being used here:

https://github.com/gridap/GridapODEs.jl/blob/ee3d0b64fb1323ddb16d279bc3e4bcdd47f9a714/test/DiffEqsWrappersTests/DiffEqsTests.jl#L71-L94

You can extract the full jacobian J, the mass matrix M, what I call the stiffmes matrix K s.t. J = M*gamma + K, and a method to allocate these arrays.

If you play with this and can add examples to the tests with some other solvers (right now only DAE Sundials solver), it would be great. Let me know if you miss methods or they should be changed somehow.

[kanav99]

I am trying to find more examples, but for now these are the methods I am trying -

# DABDF2 only works non adaptive for now, some issue in the error estimator
sol_iip = OrdinaryDiffEq.solve(prob_iip, DABDF2(), reltol = 1e-8, abstol = 1e-8, adaptive = false, dt = 0.1)
@show sol_iip.u

# DImplicitEuler works adaptive too
sol_iip = OrdinaryDiffEq.solve(prob_iip, DImplicitEuler(), reltol = 1e-8, abstol = 1e-8)
@show sol_iip.u

[ChrisRackauckas]

What about mass matrix form?

Also, IDA?

[kanav99]

IDA Fails for n >= 3 (this just increases the state size), we get a [IDAS ERROR] IDACalcIC Newton/Linesearch algorithm failed to converge..

For mass matrix form, I think we need the ODEFunction instead of DAEFunction.

[jd-lara]

@ChrisRackauckas FWIW I make tests will all the currently supported Linear Solvers and IDA in https://github.com/NREL-SIIP/PowerSimulationsDynamics.jl/blob/master/test/test_sundials.jl and works well.

[ChrisRackauckas]

I'll take a quick look at the Sundials thing because that's a bit weird, but let's revive the discussion on the mass matrices. We have res!, but f(du,u,p,t) = res!(...) - M*u would be a very inefficient way to do it. Is there no way to direct get the RHS @santiagobadia ?

[ChrisRackauckas]

The only issue in the other example is that when you use a sparse matrix you have to use a linear solver for sparse matrices.

# To explore the Sundials solver options, e.g., BE with fixed time step dtd
f_iip = DAEFunction{true}(res!; jac = jac!, jac_prototype=J)
# jac_prototype is the way to pass my pre-allocated jacobian matrix
prob_iip = DAEProblem{true}(f_iip, u0, u0, tspan, differential_vars = trues(length(u0)))

sol_iip = Sundials.solve(prob_iip, IDA(linear_solver=:KLU), reltol = 1e-8, abstol = 1e-8)
@show sol_iip.u

works fine. Sundials does give you the ability to swap out linear and nonlinear solvers. We haven't fully exposed that yet, but it's not too difficult to do. @jd-lara was working on https://github.com/SciML/Sundials.jl/pull/281 to get SuperLUMT which is a better linear solver than KLU for this case.

[fverdugo]

Hi @ChrisRackauckas !

By looking at the interfaces at src/ODETools/ODEOperators.jl it seems that we dont have (yet) methods to get the RHS directly. However, we have the OperatorType trait to signal which class of ode we have in order to do optimizations.

I believe we need to enrich the interfaces for the cases Affine and Constant, to get M, K, f (following notation in src/ODETools/ODEOperators.jl) directly.

[ChrisRackauckas]

Great. Yes notice I showed how to get Sundials working with the sparse matrices, that should be decent for now. Given you are interested, I'll make sure we get our ops people building the SuperLUMT version for Sundials. But having a Mu' = f version of the DAEs would allow us to use more pure Julia tools. I'd like to try and get a fully GPU-accelerated solve with RadauIIA5 once we can and highlight it as a cross-organization result, so let me know when you have it ready. No rush of course, and we plan to work on more native Julia solvers (BDF) anyways.

[fverdugo]

Hi @ChrisRackauckas,

I just realized that functions for M and K are in fact jacobian_t and jacobian respectively. The missing function is for the rhs f, but one can always compute it as minus residual for u=0 and u_t=0. Of course this will integrate some terms more than a dedicated function for the rhs, but you can give a try for the moment.

[oriolcg]

I just realized that functions for M and K are in fact jacobian_t and jacobian respectively.

Just a minor remark here, jacobian_t does not return M, but γ*M with γ the corresponding coefficient of the ODE solver.

[fverdugo]

Yes, good point @oriolcg! One needs to call jacobian_t with dut_u=1.0 for getting M.

[ChrisRackauckas]

Under some assumptions this works:

using Test
using Gridap
using GridapODEs
using GridapODEs.ODETools
using GridapODEs.TransientFETools
using GridapODEs.DiffEqWrappers

# using DifferentialEquations
using Sundials
using Gridap.Algebra: NewtonRaphsonSolver
using Base.Iterators

# FE problem (heat eq) using Gridap
function fe_problem(u, n)

  f(t) = x -> ∂t(u)(x, t) - Δ(u(t))(x)

  domain = (0, 1, 0, 1)
  partition = (n, n)
  model = CartesianDiscreteModel(domain, partition)

  order = 1

  reffe = ReferenceFE(lagrangian,Float64,order)
  V0 = FESpace(
    model,
    reffe,
    conformity = :H1,
    dirichlet_tags = "boundary",
  )
  U = TransientTrialFESpace(V0, u)

  Ω = Triangulation(model)
  degree = 2 * order
  dΩ = Measure(Ω, degree)

  a(u, v) = ∫( ∇(v) ⋅ ∇(u) )dΩ
  b(v, t) = ∫( v * f(t) )dΩ
  m(u, v) = ∫( v * u )dΩ

  res(t, u, ut, v) = a(u, v) + m(ut, v) - b(v, t)
  jac(t, u, ut, du, v) = a(du, v)
  jac_t(t, u, ut, dut, v) = m(dut, v)

  op = TransientFEOperator(res, jac, jac_t, U, V0)

  U0 = U(0.0)
  uh0 = interpolate_everywhere(u(0.0), U0)
  u0 = get_free_values(uh0)

  return op, u0

end

# Solving the heat equation using GridapODEs and DiffEqs
tspan = (0.0, 1.0)

u(x, t) = t
u(t) = x -> u(x, t)

# ISSUE 1: When I choose n > 2, even though the problem that we will solve is
# linear, the Sundials solvers seems to have convergence issues in the nonlinear
# solver (?). Ut returns errors
# [IDAS ERROR]  IDACalcIC Newton/Linesearch algorithm failed to converge.
# ISSUE 2: When I pass `jac_prototype` the code gets stuck

n = 3 # cells per dim (2D)
op, u0 = fe_problem(u,n)

# Some checks
res!, jac!, mass!, stif! = diffeq_wrappers(op)
J = prototype_jacobian(op,u0)
r = copy(u0)
θ = 1.0; t0 = 0.0; tF = 1.0; dt = 0.1; tθ = 1.0; dtθ = dt*θ
res!(r, u0, u0, nothing, tθ)
jac!(J, u0, u0, nothing, (1 / dtθ), tθ)

K = prototype_jacobian(op,u0)
M = prototype_jacobian(op,u0)
stif!(K, u0, u0, nothing, tθ)
mass!(M, u0, u0, nothing, tθ)
# Here you have the mass matrix M

#=
function f(du,u,p,t)
  du .= stif!(K, zero(u0), u, nothing, t) * u
end
=#

function f(du,u,p,t)
  res!(du, zero(u), u, nothing, t)
  du .*= -1
end

M_op = DiffEqArrayOperator(M; update_func = (M,u,p,t) -> mass!(M, zero(u), u, nothing, t))
f_ode = ODEFunction{true}(f; mass_matrix = M_op, jac = (J,u,p,t)->stif!(J, zero(u), u, nothing, tθ), tgrad = (grad,u,p,t)->(grad.=0))
prob = ODEProblem(f_ode,u0,tspan)
sol = solve(prob,RadauIIA5(), reltol = 1e-8, abstol = 1e-8)

@test (1/dtθ)*M+K ≈ J

# To explore the Sundials solver options, e.g., BE with fixed time step dtd
f_iip = DAEFunction{true}(res!; jac = jac!)#, jac_prototype=J)
# jac_prototype is the way to pass my pre-allocated jacobian matrix
prob_iip = DAEProblem{true}(f_iip, u0, u0, tspan, differential_vars = [true,true,true,true])
# When I pass `jac_prototype` the code get stuck here:
sol_iip = Sundials.solve(prob_iip, IDA(), reltol = 1e-8, abstol = 1e-8)
# @show sol_iip.u

plot(sol)
plot!(sol_iip)

savefig("plot.png")

I do get some oscillations at zero though, probably because the assumptions aren't satisfied.