diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl index feb8d84d31..0ede9c1d12 100644 --- a/moment_kinetics/src/nonlinear_solvers.jl +++ b/moment_kinetics/src/nonlinear_solvers.jl @@ -1402,18 +1402,20 @@ MGS-GMRES' in Zou (2023) [https://doi.org/10.1016/j.amc.2023.127869]. residual_func!, left_preconditioner, right_preconditioner, x, residual0, rhs_delta, Jv_scale_factor, inv_Jv_scale_factor, solver_type) - # To start with we use 'w' as a buffer to make a copy of residual0 to which we can apply - # the left-preconditioner. parallel_map(solver_type, (delta_x) -> delta_x, v, delta_x) - left_preconditioner(residual0) # This function transforms the data stored in 'v' from δx to ≈J.δx # If initial δx is all-zero, we can skip a right-preconditioner evaluation because it # would just transform all-zero to all-zero. approximate_Jacobian_vector_product!(v, initial_delta_x_is_zero) + # To start with we use 'w' as a buffer to make a copy of residual0 to which we can + # apply the left-preconditioner. + parallel_map(solver_type, (residual_0) -> residual_0, w, residual0) + left_preconditioner(w) + # Now we actually set 'w' as the first Krylov vector, and normalise it. - parallel_map(solver_type, (residual0, v) -> -residual0 - v, w, residual0, v) + parallel_map(solver_type, (residual0, v) -> -residual0 - v, w, w, v) beta = distributed_norm(solver_type, w, norm_params...) parallel_map(solver_type, (w,beta) -> w/beta, select_from_V(V, 1), w, beta) if serial_solve diff --git a/moment_kinetics/test/nonlinear_solver_tests.jl b/moment_kinetics/test/nonlinear_solver_tests.jl index 745eb9f331..9d3daa10e3 100644 --- a/moment_kinetics/test/nonlinear_solver_tests.jl +++ b/moment_kinetics/test/nonlinear_solver_tests.jl @@ -11,6 +11,8 @@ using moment_kinetics.nonlinear_solvers using moment_kinetics.type_definitions: mk_float, mk_int using MPI +using LinearAlgebra +using SparseArrays function linear_test() println(" - linear test") @@ -148,13 +150,11 @@ end function nonlinear_test() println(" - non-linear test") - @testset "non-linear test" for (coord_names, serial_solve) ∈ (((:z,), false), ((:vpa,), true)) - # Test represents constant-coefficient diffusion, in 1D steady state, with a - # central finite-difference discretisation of the second derivative. - # - # Note, need to use newton_solve!() here even though it is a linear problem, - # because the inexact Jacobian-vector product we use in linear_solve!() means - # linear_solve!() on its own does not converge to the correct answer. + @testset "non-linear test, serial_solve=$serial_solve" for (coord_names, serial_solve) ∈ (((:z,), false), ((:vpa,), true)) + # Test represents diffusion with a coefficient proportional to the variable to the + # power 5/2 (similar to collisional, parallel, thermal diffusion in Braginskii), + # in 1D steady state, with a central finite-difference discretisation of the + # second derivative. n = 16 restart = 10 @@ -285,11 +285,207 @@ function nonlinear_test() end end +function nonlinear_preconditioning_test() + println(" - non-linear test") + @testset "non-linear test, precon=$precon" for precon ∈ (:none, :left, :right) + # Test represents diffusion with a coefficient proportional to the variable to the + # power 5/2 (similar to collisional, parallel, thermal diffusion in Braginskii), + # in 1D steady state, with a central finite-difference discretisation of the + # second derivative. + # + # Just do this test with 'serial solve' because it makes implementing a dummy + # preconditioner easier. This test checks the basic structure of the + # implementation (e.g. the places where the preconditioner calls are made in the + # nonlinear solver), not the parallelisation of any particular preconditioner. + + coord_names = (:vpa,) + serial_solve = true + n = 16 + restart = 10 + max_restarts = 0 + atol = 1.0e-10 + + irank_z, nrank_z, comm_sub_z, irank_r, nrank_r, comm_sub_r = + setup_distributed_memory_MPI(1, 1, 1, 1) + + setup_loop_ranges!(block_rank[], block_size[]; s=1, sn=0, r=1, z=n, vperp=1, vpa=1, + vzeta=1, vr=1, vz=1) + + z = collect(0:n-1) ./ (n-1) + b = @. - z * (1.0 - z) + + if serial_solve + coord_comm = MPI.COMM_NULL + else + coord_comm = comm_sub_z + end + the_coord = coordinate("foo", n, n, n, 1, 1, 1, 0, Cint(0), Cint(0), 1.0, + zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_int, 0), + zeros(mk_int, 0), zeros(mk_int, 0), zeros(mk_int, 0), + zeros(mk_int, 0, 0), "", "", "", "", false, nothing, + zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0), + zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0), + zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0), + zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0), + zeros(mk_int, 0), zeros(mk_float, 0), zeros(mk_float, 0), + zeros(mk_float, 0), zeros(mk_float, 0), + zeros(mk_int, 0), zeros(mk_int, 0), zeros(mk_float, 0, 0), + zeros(mk_float, 0, 0), zeros(mk_float, 0), + zeros(mk_float, 0), coord_comm, 1:n, 1:n, + zeros(mk_float, 0), zeros(mk_float, 0), "", + zeros(mk_float, 0), false, zeros(mk_float, 0, 0, 0), + zeros(mk_float, 0, 0), zeros(mk_float, 0), + zeros(mk_float, 0)) + coords = NamedTuple(c => the_coord for c ∈ coord_names) + + function rhs_func!(residual, x; krylov=false) + if serial_solve + i = 1 + D = abs(x[i])^2.5 + residual[i] = D * (- 2.0 * x[i] + x[i+1]) - b[i] + for i ∈ 2:n-1 + D = abs(x[i])^2.5 + residual[i] = D * (x[i-1] - 2.0 * x[i] + x[i+1]) - b[i] + end + i = n + D = abs(x[i])^2.5 + residual[i] = D * (x[i-1] - 2.0 * x[i]) - b[i] + else + @begin_anyzv_region() + @anyzv_serial_region begin + i = 1 + D = abs(x[i])^2.5 + residual[i] = D * (- 2.0 * x[i] + x[i+1]) - b[i] + for i ∈ 2:n-1 + D = abs(x[i])^2.5 + residual[i] = D * (x[i-1] - 2.0 * x[i] + x[i+1]) - b[i] + end + i = n + D = abs(x[i])^2.5 + residual[i] = D * (x[i-1] - 2.0 * x[i]) - b[i] + end + end + return nothing + end + + if serial_solve + x = allocate_float(:vpa=>n) + residual = allocate_float(:vpa=>n) + delta_x = allocate_float(:vpa=>n) + rhs_delta = allocate_float(:vpa=>n) + v = allocate_float(:vpa=>n) + w = allocate_float(:vpa=>n) + else + x = allocate_shared_float(:z=>n) + residual = allocate_shared_float(:z=>n) + delta_x = allocate_shared_float(:z=>n) + rhs_delta = allocate_shared_float(:z=>n) + v = allocate_shared_float(:z=>n) + w = allocate_shared_float(:z=>n) + end + + if serial_solve + x .= 1.0 + residual .= 0.0 + delta_x .= 0.0 + rhs_delta .= 0.0 + v .= 0.0 + w .= 0.0 + else + @begin_serial_region() + @serial_region begin + x .= 1.0 + residual .= 0.0 + delta_x .= 0.0 + rhs_delta .= 0.0 + v .= 0.0 + w .= 0.0 + end + end + + function get_diffusion_matrix(x) + # This is not a particularly efficient way to construct the preconditioner + # matrix, but this is only a test so simplicity is more important. + + M = zeros(n,n) + + # Dirichlet boundary conditions, so x at end points does not change. + M[1,1] = 1.0 + M[end,end] = 1.0 + + for i ∈ 2:n-1 + D = abs(x[i])^2.5 + M[i,i-1] = D + M[i,i] = -2.0 * D + M[i,i+1] = D + end + + return sparse(M) + end + + diffusion_matrix = get_diffusion_matrix(x) + diffusion_lu = lu(diffusion_matrix) + + function recalculate_preconditioner() + diffusion_matrix = get_diffusion_matrix(x) + lu!(diffusion_lu, diffusion_matrix) + end + + function preconditioner(x) + ldiv!(diffusion_lu, x) + return nothing + end + + if precon === :none + left_preconditioner = identity + right_preconditioner = identity + recalculate_preconditioner = nothing + elseif precon === :left + left_preconditioner = preconditioner + right_preconditioner = identity + recalculate_preconditioner = recalculate_preconditioner + elseif precon === :right + left_preconditioner = identity + right_preconditioner = preconditioner + recalculate_preconditioner = recalculate_preconditioner + else + error("Unrecognised value for precon=$precon.") + end + + nl_solver_params = setup_nonlinear_solve( + (rtol=0.0, atol=atol, linear_restart=restart, + linear_max_restarts=max_restarts, nonlinear_max_iterations=100, + linear_rtol=1.0e-3, linear_atol=1.0, preconditioner_update_interval=5, + total_its_soft_limit=50, adi_precon_iterations=1), + coords; serial_solve=serial_solve, anyzv_region=!serial_solve) + + if !serial_solve + @begin_r_anyzv_region() + end + newton_solve!(x, rhs_func!, residual, delta_x, rhs_delta, v, w, nl_solver_params; + coords, left_preconditioner=left_preconditioner, + right_preconditioner=right_preconditioner, + recalculate_preconditioner=recalculate_preconditioner) + + rhs_func!(residual, x) + + if serial_solve + @test isapprox(residual, zeros(n); atol=4.0*atol) + else + @begin_serial_region() + @serial_region begin + @test isapprox(residual, zeros(n); atol=4.0*atol) + end + end + end +end + function runtests() @testset "non-linear solvers" begin println("non-linear solver tests") linear_test() nonlinear_test() + nonlinear_preconditioning_test() end end