fix adi preconditioner

mabarnes · Oct 26, 2024 · f7b625d · f7b625d
1 parent c2a16ce
commit f7b625d
Showing 1 changed file with 3 additions and 8 deletions.
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1131,9 +1131,7 @@ global_rank[] == 0 && println("recalculating precon")
                             end
                         end
 
-                        v_range = collect((iz - 1)*v_size + 1: iz*v_size)
-                        push!(v_range, z.n*v_size + iz)
-                        adi_info.v_solve_explicit_matrices[v_solve_counter] = sparse(@view(explicit_J[v_range,:]))
+                        adi_info.v_solve_explicit_matrices[v_solve_counter] = sparse(@view(explicit_J[adi_info.v_solve_global_inds[v_solve_counter],:]))
                     end
                     @boundscheck v_solve_counter == adi_info.v_solve_nsolve || error("v_solve_counter($v_solve_counter) != v_solve_nsolve($(adi_info.v_solve_nsolve))")
 
@@ -1181,8 +1179,7 @@ global_rank[] == 0 && println("recalculating precon")
                             end
                         end
 
-                        z_range = collect((ivperp - 1)*vpa.n + ivpa:v_size:(z.n - 1)*v_size + (ivperp - 1)*vpa.n + ivpa - 1)
-                        adi_info.z_solve_explicit_matrices[z_solve_counter] = sparse(@view(explicit_J[z_range,:]))
+                        adi_info.z_solve_explicit_matrices[z_solve_counter] = sparse(@view(explicit_J[adi_info.z_solve_global_inds[z_solve_counter],:]))
                     end
                     @serial_region begin
                         # Do the solve for ppar on the rank-0 process, which has the
@@ -1218,8 +1215,7 @@ global_rank[] == 0 && println("recalculating precon")
                             end
                         end
 
-                        z_range = z.n*v_size + 1:z.n*v_size + z.n
-                        adi_info.v_solve_explicit_matrices[v_solve_counter] = sparse(@view(explicit_J[z_range,:]))
+                        adi_info.z_solve_explicit_matrices[z_solve_counter] = sparse(@view(explicit_J[adi_info.z_solve_global_inds[z_solve_counter],:]))
                     end
                     @boundscheck z_solve_counter == adi_info.z_solve_nsolve || error("z_solve_counter($z_solve_counter) != z_solve_nsolve($(adi_info.z_solve_nsolve))")
                 end
@@ -1309,7 +1305,6 @@ global_rank[] == 0 && println("recalculating precon")
                             # terms are being added to the left-hand-side preconditioner
                             # matrix, but here the 'explicit matrix' terms are added on
                             # the right-hand-side.
-                            println("check sizes ", size(z_solve_buffer, " ", size(z_solve_explicit_matrices[isolve]), " ", size(this_intermediate_buffer)))
                             mul!(z_solve_buffer, z_solve_explicit_matrices[isolve],
                                  this_intermediate_buffer, -1.0, 1.0)
                             ldiv!(z_solve_implicit_lus[isolve], z_solve_buffer)