Implicit nonlinear odes with variable mass matrix

Hi all. I’m quite new to Julia, moving over from Python and am mostly attracted by the more advanced differential equations solvers that let me incorporate mass matrices. I’m trying to use variable mass matrices in some multiple shooting code I’ve written, but am having trouble. I’m trying to solve an implicit, nonlinear ODE of the form A(y,z)y’ = R(y,z)

My full code is probably too much to sift through, so I’ve tried to cut it down to the bare minimum needed to show the problem. I’m just using some arbitrary differential equation to show what I’m hoping to achieve.

function K(z)
    return 0.2*z*exp(-z^2/(2*10^2))

function update_func(A,y,p,z)
    a,K = p
    A[1,1] = a
    A[1,2] = 0

    A[2,1] = 0
    A[2,2] = K(z)

function R(dy,y,p,z)
    y1,y2 = y
    a,K = p
    dy[1] = y2 - K(z)
    dy[2] = y1+a*z

nsub = 8 # number of subdomains
s = ones(nsub+1,2) # Initial guesses for y1 and y2
p = (0.1,K)
### For multiple shooting, we need to integrate over a range of subdomains
zrange = range(0.1, stop=1.5, length=nsub+1)

for i in nsub
M = DiffEqArrayOperator(ones(2,2),update_func=update_func) ### See note
sol = solve(ODEProblem(ODEFunction(R, mass_matrix=M), s[i,:], (zranges[i],zranges[i+1]), p),RadauIIA5())

Ideally, I would like to further generalize K to be K(u,y).

I must admit that I can’t quite figure out what DiffEqArrayOperator is doing here, but this isn’t working as intended. From my testing/comparing simple examples to python code, it appears that the mass matrix A is being updated only at the start of the loop, and so in each looping I’m solving the more static equations:
i=1 → A(z=0.1)y’ = R(y,z)
i=2 → A(z=0.275)y’ = R(y,z)
i=3 → A(z=0.45)y’ = R(y,z)

However, this isn’t what I want – I want to be solving A(z)y’ = R(y,z). When i=1 in the loop, z should vary from 0.1 to 0.275 and I’d like A(z) to reflect that.

Any help on how this can be achieved would be greatly appreciated. Thank you!

Easiest thing is to just use a DAEProblem: