Optimise 1d Gauss-Legendre interpolation

Precalculate some quantities to minimise the amount of work needed to
evaluate the Lagrange polynomials, and hopefully to make the loops
vectorise better.

moment_kinetics/src/chebyshev.jl

@@ -465,7 +465,8 @@ function chebyshev_spectral_derivative!(df,f)
     end
 end
 
-function single_element_interpolate!(result, newgrid, f, imin, imax, coord, chebyshev::chebyshev_base_info)
+function single_element_interpolate!(result, newgrid, f, imin, imax, ielement, coord,
+                                     chebyshev::chebyshev_base_info)
     # Temporary buffer to store Chebyshev coefficients
     cheby_f = chebyshev.df
 

moment_kinetics/src/gauss_legendre.jl

@@ -29,7 +29,7 @@ using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float
 import ..calculus: elementwise_derivative!, mass_matrix_solve!
 import ..interpolation: single_element_interpolate!
-using ..lagrange_polynomials: lagrange_poly
+using ..lagrange_polynomials: lagrange_poly_optimised
 using ..moment_kinetics_structs: weak_discretization_info
 
 
@@ -80,6 +80,10 @@ struct gausslegendre_base_info
     Y30::Array{mk_float,3}
     # local nonlinear diffusion matrix Y31
     Y31::Array{mk_float,3}
+    # 'Other' nodes where the j'th Lagrange polynomial (which is 1 at x[j]) is equal to 0
+    other_nodes::Array{mk_float,3}
+    # One over the denominators of the Lagrange polynomials
+    one_over_denominator::Array{mk_float,2}
 end
 
 struct gausslegendre_info{TSparse, TLU} <: weak_discretization_info
@@ -189,8 +193,29 @@ function setup_gausslegendre_pseudospectral_lobatto(coord; collision_operator_di
         Y30 = allocate_float(0, 0, 0)
         Y31 = allocate_float(0, 0, 0)
     end
+
+    # Precompute some values for Lagrange polynomial evaluation
+    other_nodes = allocate_float(coord.ngrid-1, coord.ngrid, coord.nelement_local)
+    one_over_denominator = allocate_float(coord.ngrid, coord.nelement_local)
+    for ielement ∈ 1:coord.nelement_local
+        if ielement == 1
+            imin = coord.imin[ielement]
+        else
+            imin = coord.imin[ielement] - 1
+        end
+        imax = coord.imax[ielement]
+        grid = coord.grid[imin:imax]
+        for j ∈ 1:coord.ngrid
+            @views other_nodes[1:j-1,j,ielement] .= grid[1:j-1]
+            @views other_nodes[j:end,j,ielement] .= grid[j+1:end]
+
+            one_over_denominator[j,ielement] = 1.0 / prod(grid[j] - n for n ∈ @view other_nodes[:,j,ielement])
+        end
+    end
+
     return gausslegendre_base_info(Dmat,M0,M1,M2,S0,S1,
-            K0,K1,K2,P0,P1,P2,D0,Y00,Y01,Y10,Y11,Y20,Y21,Y30,Y31)
+            K0,K1,K2,P0,P1,P2,D0,Y00,Y01,Y10,Y11,Y20,Y21,Y30,Y31,
+            other_nodes, one_over_denominator)
 end
 
 function setup_gausslegendre_pseudospectral_radau(coord; collision_operator_dim=true)
@@ -261,8 +286,29 @@ function setup_gausslegendre_pseudospectral_radau(coord; collision_operator_dim=
         Y30 = allocate_float(0, 0, 0)
         Y31 = allocate_float(0, 0, 0)
     end
+
+    # Precompute some values for Lagrange polynomial evaluation
+    other_nodes = allocate_float(coord.ngrid-1, coord.ngrid, coord.nelement_local)
+    one_over_denominator = allocate_float(coord.ngrid, coord.nelement_local)
+    for ielement ∈ 1:coord.nelement_local
+        if ielement == 1
+            imin = coord.imin[ielement]
+        else
+            imin = coord.imin[ielement] - 1
+        end
+        imax = coord.imax[ielement]
+        grid = coord.grid[imin:imax]
+        for j ∈ 1:coord.ngrid
+            @views other_nodes[1:j-1,j,ielement] .= grid[1:j-1]
+            @views other_nodes[j:end,j,ielement] .= grid[j+1:end]
+
+            one_over_denominator[j,ielement] = 1.0 / prod(grid[j] - n for n ∈ @view other_nodes[:,j,ielement])
+        end
+    end
+
     return gausslegendre_base_info(Dmat,M0,M1,M2,S0,S1,
-            K0,K1,K2,P0,P1,P2,D0,Y00,Y01,Y10,Y11,Y20,Y21,Y30,Y31)
+            K0,K1,K2,P0,P1,P2,D0,Y00,Y01,Y10,Y11,Y20,Y21,Y30,Y31,
+            other_nodes, one_over_denominator)
 end 
 
 function elementwise_derivative!(coord, ff, gausslegendre::gausslegendre_info)
@@ -308,11 +354,23 @@ function elementwise_derivative!(coord, ff, adv_fac, spectral::gausslegendre_inf
     return elementwise_derivative!(coord, ff, spectral)
 end
 
-function single_element_interpolate!(result, newgrid, f, imin, imax, coord, gausslegendre::gausslegendre_base_info)
-    for i ∈ 1:length(newgrid)
-        result[i] = 0.0
-        for j ∈ 1:coord.ngrid
-            @views result[i] += f[j] * lagrange_poly(j, coord.grid[imin:imax], newgrid[i])
+function single_element_interpolate!(result, newgrid, f, imin, imax, ielement, coord,
+                                     gausslegendre::gausslegendre_base_info)
+    n_new = length(newgrid)
+
+    i = 1
+    other_nodes = @view gausslegendre.other_nodes[:,i,ielement]
+    one_over_denominator = gausslegendre.one_over_denominator[i,ielement]
+    this_f = f[i]
+    for j ∈ 1:n_new
+        result[j] = this_f * lagrange_poly_optimised(other_nodes, one_over_denominator, newgrid[j])
+    end
+    for i ∈ 2:coord.ngrid
+        other_nodes = @view gausslegendre.other_nodes[:,i,ielement]
+        one_over_denominator = gausslegendre.one_over_denominator[i,ielement]
+        this_f = f[i]
+        for j ∈ 1:n_new
+            result[j] += this_f * lagrange_poly_optimised(other_nodes, one_over_denominator, newgrid[j])
         end
     end
 

moment_kinetics/src/interpolation.jl

@@ -99,7 +99,7 @@ function interpolate_to_grid_1d!(result, newgrid, f, coord, spectral)
         kmin = kstart[1]
         kmax = kstart[2] - 1
         @views single_element_interpolate!(result[kmin:kmax], newgrid[kmin:kmax],
-                                           f[imin:imax], imin, imax, coord,
+                                           f[imin:imax], imin, imax, 1, coord,
                                            first_element_spectral)
     end
     @inbounds for j ∈ 2:nelement
@@ -109,7 +109,7 @@ function interpolate_to_grid_1d!(result, newgrid, f, coord, spectral)
             imin = coord.imin[j] - 1
             imax = coord.imax[j]
             @views single_element_interpolate!(result[kmin:kmax], newgrid[kmin:kmax],
-                                               f[imin:imax], imin, imax, coord,
+                                               f[imin:imax], imin, imax, j, coord,
                                                spectral.lobatto)
         end
     end

moment_kinetics/src/lagrange_polynomials.jl

@@ -8,7 +8,7 @@ their being scattered (and possibly duplicated) in other modules.
 """
 module lagrange_polynomials
 
-export lagrange_poly
+export lagrange_poly, lagrange_poly_optimised
 
 """
 Lagrange polynomial
@@ -33,4 +33,21 @@ function lagrange_poly(j,x_nodes,x)
     return poly
 end
 
+"""
+    lagrange_poly_optimised(other_nodes, one_over_denominator, x)
+
+Optimised version of Lagrange polynomial calculation, making use of pre-calculated quantities.
+
+`other_nodes` is a vector of the grid points in this element where this Lagrange
+polynomial is zero (the other nodes than the one where it is 1).
+
+`one_over_denominator` is `1/prod(x0 - n for n ∈ other_nodes)` where `x0` is the grid
+point where this Lagrange polynomial is 1.
+
+`x` is the point to evaluate the Lagrange polynomial at.
+"""
+function lagrange_poly_optimised(other_nodes, one_over_denominator, x)
+    return prod(x - n for n ∈ other_nodes) * one_over_denominator
+end
+
 end