Using PreallocationTools with LinearSolve

I’m trying to solve a large 2D PDE using this example as a template. I’m effectively trying to solve some simplified (though still nonlinear) pseudo-Boussinesq Navier Stokes equations where I evolve u and temperature and solve for w using \frac{\partial u}{\partial x}+\frac{\partial w}{\partial z}=0. This becomes a system of equations that I have to solve for w at each time step (with my problem geometry I can assume variations in w are small, so divergence in u should be the only thing driving w)

I went through many iterations before realizing that the minimal non-working example is actually very straightforward. This problem is trivialized for simplicity, but it highlights the line that is giving me an error.

using DifferentialEquations, SparseArrays, LinearSolve, PreallocationTools

function ddt(dy, y, p, t)
    Nz,Nx,A_w,linsolve,w,divu = p

    ### I'm not certain I understand how get_tmp works, but I know it doesn't work without it.
    divu = get_tmp(divu,y)
    w = get_tmp(w,y)
    

    ### Ultimately, I would be computing "divu" before and then doing something with "w"
    
    ### This approach works when I comment out the line that makes "A_w_" sparse.
    # w[2:end-1,2:end-1] = reshape(solve(LinearProblem(A_w,reshape(divu,Nx*Nz))).u,(Nz,Nx))


    ### I can't make this approach work regardless of sparisty
    w[2:end-1,2:end-1] = reshape(solve(LinearSolve.set_b(linsolve,reshape(divu,Nx*Nz))).u,(Nz,Nx))


    ### Set all to zero for simplicity
    dy[:,:,:] .= 0.
end

Nx_= 10
Nz_= 10

A_w_ = Matrix(1.0I, Nx_*Nz_, Nx_*Nz_)
A_w_ = sparse(A_w_)

### Always going to use the same A matrix, so LU factorize and cache.
bw = ones(Nz_*Nx_)
linsolve_ = init(LinearProblem(A_w_, bw))
sol1 = solve(linsolve_)

w_ = DiffCache(zeros(Nz_+2, Nx_+2))
divu_ = DiffCache(zeros(Nz_, Nx_))

p = Nz_,Nx_,A_w_,linsolve_,w_,divu_

y0 = zeros(Nz_+2, Nx_+2, 2)

prob = ODEProblem(ddt, y0, (0.0, 1.0), p);
sol = solve(prob, FBDF(),save_everystep=false);"

I get an error that resembles:

First call to automatic differentiation for the Jacobian
failed. This means that the user `f` function is not compatible
with automatic differentiation. Methods to fix this include:...

As I’ve indicated in my comments, I can make the first method work if my A matrix is not sparse, but I cannot get the second approach (caching the LU decomposition of A) to work at all. If I allocate divu and w (and don’t pass them in through p), the code also works fine, but is far slower.

I think my issue has something to do with a misuse of PreallocationTools.jl (likely that I’m not applying it to linsolve_?) but I’m having a hard time understanding how to do it. It looks like LazyBufferCache might be what I need somehow, but can’t wrap my head around what it does. Any help would be great.

1 Like

Try forcing it to SparspakFactorization(), i.e. w[2:end-1,2:end-1] = reshape(solve(LinearProblem(A_w,reshape(divu,Nx*Nz)),SparspakFactorization()).u,(Nz,Nx)). If that works I can update the automated algorithm to handle this case.

Yeah for that you should probably use the GeneralLazyBufferCache.

lbc = GeneralLazyBufferCache(function (y)
    init(LinearProblem(A_w,similar(y,Nx*Nz), SparspakFactorization())
end)

and then prob = lbc[y] gives you a dual/non-dual prob, and then

w[2:end-1,2:end-1] = reshape(solve(LinearSolve.set_b(linsolve,reshape(divu,Nx*Nz))).u,(Nz,Nx))

should work. Note I didn’t try all of this code and just writing this on the fly since I’m short on time, but thought you might want the untested pointers instead of nothing.

I appreciate you taking the time. It’s giving me more avenues to try. Somewhat blindly permuting different parameters seems to have worked!

The first approach gave me the same error as before.

I wasn’t able to get the second approach to work, but “SparspakFactorization()” seems like it was the cause. Once I removed it and just defined lbc as follows, it worked.

lbc = GeneralLazyBufferCache(function (y)
    init(LinearProblem(A_w_,y))
end)
linsolve_ = lbc[bw]
sol1 = solve(linsolve_)

Thanks, Chris. Very much appreciated.

While I’ve had this code working for a while, I’ve encountered an argument error. I don’t know if it’s been doing this all along and I haven’t noticed, or if this started when I updated a couple of packages recently. Admittedly, it doesn’t seem to be causing any trouble, but I wanted to see if it was something I was doing wrong (in case it starts causing trouble later). Here’s a very simple MWE:

def_mat(M) = [diagm(1 => ones(M-1))-I [zeros(M-1,1); 1.0]]

n = 4
A = sparse(def_mat(n))

x_vec = ones(n)

lbc_ = GeneralLazyBufferCache(function (y)
    init(LinearProblem(A,y))
end)

linsolve_w_ = lbc_[x_vec]
sol1 = solve(linsolve_w_)
linsolve_w_.b = [4.0; 3.0; 2.0; 1.0]
sol2 = solve(linsolve_w_)

Which works as I would hope. However, if I comment out the last 3 lines (or otherwise end a line with “linsolve_w_” in any way), I get the following:

ArgumentError: pointer to the SparseArrays.LibSuiteSparse.cholmod_factor_struct object is null. This can happen if the object has been serialized.

Should I be structuring my GeneralLazyBufferCache function differently, or is this something within SparseArrays?