Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use LinearSolve precs interface #2318

Open
wants to merge 6 commits into
base: master
Choose a base branch
from

Conversation

oscardssmith
Copy link
Contributor

@oscardssmith oscardssmith commented Aug 7, 2024

This is a post OrdinaryDiffEq solver split version of #2247 (remaking was substantially easier than rebasing). It still needs a little bit of work (and for the LinearSolve side to be merged first), but it's getting pretty close.

@oscardssmith oscardssmith marked this pull request as draft August 7, 2024 21:41
@avik-pal
Copy link
Member

and for the LinearSolve side to be merged first)

Is there a link for this? I might want to do a similar cleanup of the boilerplate in NonlinearSolve

@oscardssmith oscardssmith marked this pull request as ready for review August 15, 2024 16:50
@oscardssmith
Copy link
Contributor Author

yeah it's SciML/LinearSolve.jl#514. It's not yet fully documented but there are some docs for it and it works.

@oscardssmith
Copy link
Contributor Author

I believe this now works (once SciML/SciMLBase.jl#769 is merged so remake on the Linear solver works properly)

@oscardssmith
Copy link
Contributor Author

Tests all pass locally (with deved SciMLBase). @ChrisRackauckas can you confirm that this is the API we want?

@oscardssmith oscardssmith force-pushed the os/fix-dolinsolve-2 branch 2 times, most recently from 1e1d569 to 25415c9 Compare August 24, 2024 16:09
@ChrisRackauckas
Copy link
Member

It looks right.

@oscardssmith
Copy link
Contributor Author

hmm. The Interface II test seems to be failing stochastically.

@oscardssmith
Copy link
Contributor Author

planning on merging as soon as tests pass

@oscardssmith
Copy link
Contributor Author

Ah, it's not random, but it is dependent on the order of runs. Does Krylov.jl depend on the global random number generator? IMO that's pretty bad UI.

@oscardssmith
Copy link
Contributor Author

@ChrisRackauckas This error does actually appear to be at least somewhat stochastic. Do you know where the randomness would be coming from and how to work around it? Currently, it looks like this test will fail roughly 1/60 solve calls that gets run (roughly equally likely for various solvers and precs types).

@ChrisRackauckas
Copy link
Member

Are you factorizing any uninitialized memory? The reason for the init stuff in ArrayInterface is to avoid that, so double check it is still in play.

@oscardssmith
Copy link
Contributor Author

I'm pretty sure that's not the problem (since I'm not getting undef errors, but rather just timesteps failing repeatedly)

@ChrisRackauckas
Copy link
Member

You wouldn't get undef errors from similar memory, you'd just get random non-zeroed values.

@oscardssmith
Copy link
Contributor Author

It looks like the undef J prototype that we were creating was causing all sorts of problems. Changing it to a zeromatrix seems to fix them.

@oscardssmith
Copy link
Contributor Author

oscardssmith commented Sep 4, 2024

Looks like tests are finally mostly passing

@oscardssmith oscardssmith force-pushed the os/fix-dolinsolve-2 branch 2 times, most recently from 6b3e731 to 0b0ba2e Compare September 9, 2024 16:56
@oscardssmith
Copy link
Contributor Author

The last failure is StabilizedIRK and I honestly have absolutely no idea why it's any different. @ChrisRackauckas any thoughts?

@oscardssmith
Copy link
Contributor Author

I think this is finally good to merge.

@ChrisRackauckas
Copy link
Member

Show some before and after WPDs from the stiff ODE set.

@oscardssmith
Copy link
Contributor Author

WPDs with or without precs?

@oscardssmith
Copy link
Contributor Author

also why work-precision? this shouldn't change accuracy

@ChrisRackauckas
Copy link
Member

It should take like 10 minutes to make them and eyeball if there's a change. This is a pretty large change so it's worth double checking.

@oscardssmith
Copy link
Contributor Author

Here's the WP results:

using OrdinaryDiffEq, LinearSolve, Test, IncompleteLU, ModelingToolkit, DiffEqDevTools
import AlgebraicMultigrid

# Required due to https://github.com/JuliaSmoothOptimizers/Krylov.jl/pull/477
Base.eltype(::IncompleteLU.ILUFactorization{Tv, Ti}) where {Tv, Ti} = Tv
Base.eltype(::AlgebraicMultigrid.Preconditioner) = Float64

const N = 32
const xyd_brusselator = range(0, stop = 1, length = N)
brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0
limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a

const iter = Ref(0)
function brusselator_2d_loop(du, u, p, t)
    global iter[] += 1
    A, B, alpha, dx = p
    alpha = alpha / dx^2
    @inbounds for I in CartesianIndices((N, N))
        i, j = Tuple(I)
        x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
        ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
        limit(j - 1, N)
        du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
                       4u[i, j, 1]) +
                      B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
                      brusselator_f(x, y, t)
        du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] -
                       4u[i, j, 2]) +
                      A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
    end
    nothing
end
p = (3.4, 1.0, 10.0, step(xyd_brusselator))

function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I, 1] = 22 * (y * (1 - y))^(3 / 2)
        u[I, 2] = 27 * (x * (1 - x))^(3 / 2)
    end
    u
end
u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop, u0, (0.0, 11.5), p)

du0 = copy(u0)
jac = ModelingToolkit.Symbolics.jacobian_sparsity(
    (du, u) -> brusselator_2d_loop(du, u, p,
        0.0), du0,
    u0)

prob_ode_brusselator_2d_sparse = ODEProblem(
    ODEFunction(brusselator_2d_loop,
        jac_prototype = float.(jac)),
    u0, (0.0, 11.5), p)
function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
    if newW === nothing || newW
        Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(convert(
            AbstractMatrix,
            W)))
    else
        Pl = Plprev
    end
    Pl, nothing
end
abstols = 1.0 ./ 10.0 .^ (4:11)
reltols = 1.0 ./ 10.0 .^ (1:8)
setups = [Dict(:alg=>Rosenbrock23(linsolve = KrylovJL_GMRES(precs = algebraicmultigrid), concrete_jac=true))]
test_sol = [solve(prob_ode_brusselator_2d_sparse, Rodas5P(), save_everystep = false, abstol=1e-10, reltol=1e-10)]
probs=[prob_ode_brusselator_2d_sparse]
wp = WorkPrecisionSet(probs,abstols,reltols,setups;
                             save_everystep=false,appxsol=test_sol,maxiters=Int(1e5),numruns=10)
plot(wp)

master:
image

pr:
image

@ChrisRackauckas
Copy link
Member

It does look like there are some significant regressions?

@oscardssmith oscardssmith changed the title simplify dolinsolve use LinearSolve precs interface Oct 2, 2024
@ChrisRackauckas
Copy link
Member

What's the current status here?

@oscardssmith
Copy link
Contributor Author

I need to figure out what's causing this to sometimes be slower, but it's not getting merged urgently since it's not going in until OrdinaryDiffEqv7

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants