Merge pull request #123 from mabarnes/unequally-spaced-element-bounda…

…ries

Unequally spaced element boundaries for grid optimisations

.github/workflows/test.yml

@@ -10,7 +10,7 @@ jobs:
       matrix:
         os: [ubuntu-latest, macOS-latest]
       fail-fast: false
-    timeout-minutes: 30
+    timeout-minutes: 40
 
     steps:
       - uses: actions/checkout@v2

src/chebyshev.jl

@@ -41,9 +41,25 @@ end
 
 """
 initialize chebyshev grid scaled to interval [-box_length/2, box_length/2]
-"""
-function scaled_chebyshev_grid(ngrid, nelement_global, nelement_local, n,
-			irank, box_length, imin, imax)
+we no longer pass the box_length to this function, but instead pass precomputed
+arrays element_scale and element_shift that are needed to compute the grid.
+
+ngrid -- number of points per element (including boundary points)
+nelement_local -- number of elements in the local (distributed memory MPI) grid
+n -- total number of points in the local grid (excluding duplicate points)
+element_scale -- the scale factor in the transform from the coordinates 
+                 where the element limits are -1, 1 to the coordinate where
+                 the limits are Aj = coord.grid[imin[j]-1] and Bj = coord.grid[imax[j]]
+                 element_scale = 0.5*(Bj - Aj)
+element_shift -- the centre of the element in the extended grid coordinate
+                 element_shift = 0.5*(Aj + Bj)
+imin -- the array of minimum indices of each element on the extended grid.
+        By convention, the duplicated points are not included, so for element index j > 1
+        the lower boundary point is actually imin[j] - 1
+imax -- the array of maximum indices of each element on the extended grid.
+"""
+function scaled_chebyshev_grid(ngrid, nelement_local, n,
+			element_scale, element_shift, imin, imax)
     # initialize chebyshev grid defined on [1,-1]
     # with n grid points chosen to facilitate
     # the fast Chebyshev transform (aka the discrete cosine transform)
@@ -52,27 +68,22 @@ function scaled_chebyshev_grid(ngrid, nelement_global, nelement_local, n,
     chebyshev_grid = chebyshevpoints(ngrid)
     # create array for the full grid
     grid = allocate_float(n)
-    # setup the scale factor by which the Chebyshev grid on [-1,1]
-    # is to be multiplied to account for the full domain [-L/2,L/2]
-    # and the splitting into nelement elements with ngrid grid points
-    scale_factor = 0.5*box_length/float(nelement_global)
+
     # account for the fact that the minimum index needed for the chebyshev_grid
     # within each element changes from 1 to 2 in going from the first element
     # to the remaining elements
     k = 1
     @inbounds for j ∈ 1:nelement_local
-        #wgts[imin[j]:imax[j]] .= sqrt.(1.0 .- reverse(chebyshev_grid)[k:ngrid].^2) * scale_factor
-        # amount by which to shift the centre of this element from zero
-        iel_global = j + irank*nelement_local
-        shift = box_length*((float(iel_global)-0.5)/float(nelement_global) - 0.5)
+        scale_factor = element_scale[j]
+        shift = element_shift[j]
         # reverse the order of the original chebyshev_grid (ran from [1,-1])
         # and apply the scale factor and shift
         grid[imin[j]:imax[j]] .= (reverse(chebyshev_grid)[k:ngrid] * scale_factor) .+ shift
         # after first element, increase minimum index for chebyshev_grid to 2
         # to avoid double-counting boundary element
         k = 2
     end
-    wgts = clenshaw_curtis_weights(ngrid, nelement_local, n, imin, imax, scale_factor)
+    wgts = clenshaw_curtis_weights(ngrid, nelement_local, n, imin, imax, element_scale)
     return grid, wgts
 end
 
@@ -88,11 +99,9 @@ function elementwise_derivative!(coord, ff, chebyshev::chebyshev_info)
     # check array bounds
     @boundscheck nelement == size(chebyshev.f,2) || throw(BoundsError(chebyshev.f))
     @boundscheck nelement == size(df,2) && coord.ngrid == size(df,1) || throw(BoundsError(df))
-    # note that one must multiply by 2*nelement/L to get derivative
-    # in scaled coordinate
-    scale_factor = 2.0*float(coord.nelement_global)/coord.L
-	# scale factor is (length of a single element/2)^{-1}
-
+    # note that one must multiply by a coordinate transform factor 1/element_scale[j] 
+    # for each element j to get derivative on the extended grid
+
     # variable k will be used to avoid double counting of overlapping point
     # at element boundaries (see below for further explanation)
     k = 0
@@ -112,7 +121,7 @@ function elementwise_derivative!(coord, ff, chebyshev::chebyshev_info)
         # and multiply by scaling factor needed to go
         # from Chebyshev z coordinate to actual z
         for i ∈ 1:coord.ngrid
-            df[i,j] *= scale_factor
+            df[i,j] /= coord.element_scale[j]
         end
         k = 1
     end
@@ -323,23 +332,21 @@ end
 returns wgts array containing the integration weights associated
 with all grid points for Clenshaw-Curtis quadrature
 """
-function clenshaw_curtis_weights(ngrid, nelement_local, n, imin, imax, scale_factor)
+function clenshaw_curtis_weights(ngrid, nelement_local, n, imin, imax, element_scale)
     # create array containing the integration weights
     wgts = zeros(mk_float, n)
     # calculate the modified Chebshev moments of the first kind
     μ = chebyshevmoments(ngrid)
+    # calculate the raw weights for a normalised grid on [-1,1]
+    w = clenshawcurtisweights(μ)
     @inbounds begin
         # calculate the weights within a single element and
         # scale to account for modified domain (not [-1,1])
-        wgts[1:ngrid] = clenshawcurtisweights(μ)*scale_factor
+        wgts[1:ngrid] = w*element_scale[1]
         if nelement_local > 1
-            # account for double-counting of points at inner element boundaries
-            wgts[ngrid] *= 2
             for j ∈ 2:nelement_local
-                wgts[imin[j]:imax[j]] .= wgts[2:ngrid]
+                wgts[imin[j]-1:imax[j]] .+= w*element_scale[j]
             end
-            # remove double-counting of outer element boundary for last element
-            wgts[n] *= 0.5
         end
     end
     return wgts

src/coordinates.jl

@@ -4,6 +4,7 @@ module coordinates
 
 export define_coordinate, write_coordinate
 export equally_spaced_grid
+export set_element_boundaries
 
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_int
@@ -84,6 +85,12 @@ struct coordinate
     local_io_range::UnitRange{Int64}
     # global range to write into in output file
     global_io_range::UnitRange{Int64}
+    # scale for each element
+    element_scale::Array{mk_float,1}
+    # shift for each element
+    element_shift::Array{mk_float,1}
+    # option used to set up element spacing
+    element_spacing_option::String
 end
 
 """
@@ -105,10 +112,14 @@ function define_coordinate(input, parallel_io::Bool=false)
     # obtain (local) index mapping from the grid within each element
     # to the full grid
     imin, imax = elemental_to_full_grid_map(input.ngrid, input.nelement_local)
+    # initialise the data used to construct the grid
+    # boundaries for each element
+    element_boundaries = set_element_boundaries(input.nelement_global, input.L, input.element_spacing_option)
+    # shift and scale factors for each local element
+    element_scale, element_shift = set_element_scale_and_shift(input.nelement_global, input.nelement_local, input.irank, element_boundaries)
     # initialize the grid and the integration weights associated with the grid
     # also obtain the Chebyshev theta grid and spacing if chosen as discretization option
-    grid, wgts, uniform_grid = init_grid(input.ngrid, input.nelement_global,
-        input.nelement_local, n_global, n_local, input.irank, input.L,
+    grid, wgts, uniform_grid = init_grid(input.ngrid, input.nelement_local, n_global, n_local, input.irank, input.L, element_scale, element_shift,
         imin, imax, igrid, input.discretization, input.name)
     # calculate the widths of the cells between neighboring grid points
     cell_width = grid_spacing(grid, n_local)
@@ -126,7 +137,6 @@ function define_coordinate(input, parallel_io::Bool=false)
     # endpoints, so only two pieces of information must be shared
     send_buffer = allocate_float(1)
     receive_buffer = allocate_float(1)
-
     # Add some ranges to support parallel file io
     if !parallel_io
         # No parallel io, just write everything
@@ -149,7 +159,7 @@ function define_coordinate(input, parallel_io::Bool=false)
         cell_width, igrid, ielement, imin, imax, input.discretization, input.fd_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch), copy(scratch),
         scratch_2d, copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
-        local_io_range, global_io_range)
+        local_io_range, global_io_range, element_scale, element_shift, input.element_spacing_option)
 
     if input.discretization == "chebyshev_pseudospectral" && coord.n > 1
         # create arrays needed for explicit Chebyshev pseudospectral treatment in this
@@ -167,10 +177,56 @@ function define_coordinate(input, parallel_io::Bool=false)
     return coord, spectral
 end
 
+function set_element_boundaries(nelement_global, L, element_spacing_option)
+    # set global element boundaries
+    element_boundaries = allocate_float(nelement_global+1)
+    if element_spacing_option == "sqrt" && nelement_global > 3
+        # number of boundaries of sqrt grid
+        nsqrt = floor(mk_int,(nelement_global)/2) + 1
+        if nelement_global%2 > 0 # odd
+            if nsqrt < 3
+                fac = 2.0/3.0
+            else
+                fac = 1.0/( 3.0/2.0 - 0.5*((nsqrt-2)/(nsqrt-1))^2)
+            end
+        else
+            fac = 1.0
+        end
+
+        for j in 1:nsqrt
+            element_boundaries[j] = -(L/2.0) + fac*(L/2.0)*((j-1)/(nsqrt-1))^2
+        end
+        for j in 1:nsqrt
+            element_boundaries[(nelement_global+1)+ 1 - j] = (L/2.0) - fac*(L/2.0)*((j-1)/(nsqrt-1))^2
+        end
+
+    elseif element_spacing_option == "uniform" || (element_spacing_option == "sqrt" && nelement_global < 4) # uniform spacing 
+        for j in 1:nelement_global+1
+            element_boundaries[j] = L*((j-1)/(nelement_global) - 0.5)
+        end
+    else 
+        println("ERROR: element_spacing_option: ",element_spacing_option, " not supported")
+    end
+    return element_boundaries
+end
+
+function set_element_scale_and_shift(nelement_global, nelement_local, irank, element_boundaries)
+    element_scale = allocate_float(nelement_local)
+    element_shift = allocate_float(nelement_local)
+
+    for j in 1:nelement_local
+        iel_global = j + irank*nelement_local
+        upper_boundary = element_boundaries[iel_global+1]
+        lower_boundary = element_boundaries[iel_global]
+        element_scale[j] = 0.5*(upper_boundary-lower_boundary)
+        element_shift[j] = 0.5*(upper_boundary+lower_boundary)
+    end
+    return element_scale, element_shift
+end
 """
 setup a grid with n_global grid points on the interval [-L/2,L/2]
 """
-function init_grid(ngrid, nelement_global, nelement_local, n_global, n_local, irank, L,
+function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_scale, element_shift,
                    imin, imax, igrid, discretization, name)
     uniform_grid = equally_spaced_grid(n_global, n_local, irank, L)
     uniform_grid_shifted = equally_spaced_grid_shifted(n_global, n_local, irank, L)
@@ -187,7 +243,7 @@ function init_grid(ngrid, nelement_global, nelement_local, n_global, n_local, ir
     elseif discretization == "chebyshev_pseudospectral"
         if name == "vperp"
             # initialize chebyshev grid defined on [-L/2,L/2]
-            grid, wgts = scaled_chebyshev_grid(ngrid, nelement_global, nelement_local, n_local, irank, L, imin, imax)
+            grid, wgts = scaled_chebyshev_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax)
             grid .= grid .+ L/2.0 # shift to [0,L] appropriate to vperp variable
             wgts = 2.0 .* wgts .* grid # to include 2 vperp in jacobian of integral
                                         # see note above on normalisation
@@ -198,7 +254,7 @@ function init_grid(ngrid, nelement_global, nelement_local, n_global, n_local, ir
             # needed to obtain Chebyshev spectral coefficients
             # 'wgts' are the integration weights attached to each grid points
             # that are those associated with Clenshaw-Curtis quadrature
-            grid, wgts = scaled_chebyshev_grid(ngrid, nelement_global, nelement_local, n_local, irank, L, imin, imax)
+            grid, wgts = scaled_chebyshev_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax)
         end
     elseif discretization == "finite_difference"
         if name == "vperp"

src/file_io.jl

@@ -521,6 +521,10 @@ function define_io_coordinate!(parent, coord, coord_name, description, parallel_
         write_single_value!(group, "bc", coord.bc; parallel_io=parallel_io,
                             description="boundary condition for $coord_name")
 
+        # write the element spacing option for the coordinate
+        write_single_value!(group, "element_spacing_option", coord.element_spacing_option; parallel_io=parallel_io,
+                            description="element_spacing_option for $coord_name")
+
         return group
     end
 

src/input_structs.jl

@@ -127,6 +127,8 @@ mutable struct grid_input_mutable
     bc::String
     # mutable struct containing advection speed options
     advection::advection_input_mutable
+    # string option determining boundary spacing
+    element_spacing_option::String
 end
 
 """
@@ -156,6 +158,8 @@ struct grid_input
     advection::advection_input
     # MPI communicator
     comm::MPI.Comm
+    # string option determining boundary spacing
+    element_spacing_option::String
 end
 
 """

src/load_data.jl

@@ -213,6 +213,7 @@ function load_coordinate_data(fid, name; printout=false)
     discretization = load_variable(coord_group, "discretization")
     fd_option = load_variable(coord_group, "fd_option")
     bc = load_variable(coord_group, "bc")
+    element_spacing_option = load_variable(coord_group, "element_spacing_option")
 
     nelement_local = nothing
     if n_local == 1 && ngrid == 1
@@ -225,11 +226,10 @@ function load_coordinate_data(fid, name; printout=false)
     else
         nelement_global = (n_global-1) ÷ (ngrid-1)
     end
-
     # Define input to create coordinate struct
     input = grid_input(name, ngrid, nelement_global, nelement_local, nrank, irank, L,
                        discretization, fd_option, bc, advection_input("", 0.0, 0.0, 0.0),
-                       MPI.COMM_NULL)
+                       MPI.COMM_NULL, element_spacing_option)
 
     coord, spectral = define_coordinate(input)