Thanks everyone for the feedback here, learned a lot from it. For benchmark I did not know about the “$” notation for variables! @mikmoore , that blog post on multidimensional algorithms was great.
My original question might have been a bit obscure on why I’m trying to do this, so to give a bit more context my objective was to do an n-dimensional extension of bilinear interpolation (Bilinear interpolation - Wikipedia). I include below a simplified version of the solution I coded in case anyone finds it useful.
For example I can evaluate a function in a regular 2d grid:
# create mock array and coordinates of regular grid
using Plots
A = ones(5,5)
coords = [collect(LinRange(0,1,5)) for i in 1:2]
for index in CartesianIndices(A)
A[index] = coords[1][index[1]] + (coords[1][index[2]])^2
end
heatmap(coords[1],coords[2],A)
savefig("grid.png")
plot!()
I want to create a function that interpolates at arbitrary coordinates for an array of arbitrary dimension (in this case the dimension is 2 and I’m just doing bilinear interpolation). Using the tips from the blog post @mikmoore linked to, I adapted some code I had like this:
function interpolate_ndimensional_array(A, coords, target_coords)
dimensions = ndims(A)
#locate lower corner
lower_i = zeros(Int, dimensions) # this will store the value of the indices just below target_coords
for j in 1:dimensions
for i in 1:(length(coords[j])-1)
if coords[j][i] <= target_coords[j] && coords[j][i+1] >= target_coords[j]
lower_i[j] = i
break
end
end
end
for i in lower_i
if i==0
return NaN
end
end
#store values of array surrounding desired point
corner_values = zeros(ntuple(_->2, dimensions)...)::Array{eltype(A),dimensions}
for index in CartesianIndices(corner_values)
original_array_index = ntuple(i->index[i]+lower_i[i]-1,dimensions)
corner_values[index] = A[CartesianIndex(original_array_index)]
end
# interpolate all pair of points across the first dimension, then repeat for
# the second dimension using the interpolated points, and so on and so forth.
for i in 1:dimensions
xval1 = coords[i][lower_i[i]]
xval2 = coords[i][lower_i[i]+1]
xtarget = target_coords[i]
# Example for a 3-D case (trilinear interpolation). In this case as we interpolate across
# each dimension, we first need to do 4 interpolations, followed by 2 and then finally 1 last interpolation.
# When i=1, it starts by interpolating all pairs of points
# along the first dimension. index1 will iterate over CartesianIndices(1,2,2), meaning
# that index1 will take the values (1,1,1), (1,1,2), (1,2,1), (1,2,2).
# Inside the for loop, we compute the indeces that correspond to points
# with a different value on the first dimenison, index2= (2,1,1), (2,1,2), (2,2,1), (2,2,2).
# The code will then interpolate corner_values(index1) and corner_values[index2] at the
# desired value of the x-coordinate, target_coords[i], using the values at the edges,
# grid_values[i,1] and grid_values[i,2]. The result of each interpolation is stored in
# corner_values[index1] to avoid further allocations.
# Next up, when i=2, index1 will iterate over CartesianIndices(1,1,2), meaning
# that index1 will take the values (1,1,1) and (1,1,2), while index2 will take the
# values (1,2,1) and (1,2,2). From the previous iteration, since we stored the interpolation
# result in corner_values[index1], we have that each of these indeces corresponds to the following:
# - corner_values[(1,1,1)...]: Interpolation of corner_values[(1,1,1)...] and corner_values[(2,1,1)...]
# - corner_values[(1,1,2)...]: Interpolation of corner_values[(1,1,2)...] and corner_values[(2,1,2)...]
# - corner_values[(1,2,1)...]: Interpolation of corner_values[(1,2,1)...] and corner_values[(2,2,1)...]
# - corner_values[(1,2,2)...]: Interpolation of corner_values[(1,2,2)...] and corner_values[(2,2,2)...]
# So, by interpolating (1,1,1) and (1,2,1) we are interpolating across the second dimension
# the values that already resulted out of interpolation on the first dimension.
# Finally, when i=3, index1 will iterate over CartesianIndices(1,1,1), meaning that it
# will only take the value (1,1,1) and index2 will be (1,1,2). Interpolating these two will give
# the final desired result, which is stored in corner_values[1,1,1]
for index1 in CartesianIndices(ntuple(j->interpolator_index1_element(j,i),dimensions))
index2 = CartesianIndex(ntuple(j->interpolator_index2_element(j,i, index1),dimensions))
yval1 = corner_values[index1]
yval2 = corner_values[index2]
corner_values[index1] = yval1 +(yval2-yval1)*(xtarget-xval1)/(xval2-xval1)
end
end
#return result of interpolation using linear indexing
return corner_values[1]
end
# utility functions used to produce tuples during "reduction"
function interpolator_index1_element(j,i)
if j<= i
return 1
end
return 2
end
function interpolator_index2_element(j,i,index1)
if j< i
return 1
elseif j==i
return 2
end
return index1[j]
end
Which lets me interpolate over the array and coordinates I created before
xvals = LinRange(0,1,100)
yvals = LinRange(0,1,100)
A_interp = zeros(100,100)
for i in eachindex(xvals)
for j in eachindex(yvals)
A_interp[i,j] = interpolate_ndimensional_array(A, coords, [xvals[i],yvals[j]])
end
end
heatmap(xvals, yvals, A_interp)
savefig("interpolated.png")
plot!()
After optimizing this with your tips the calculation of one interpolation of this 2-d array takes 124 ns. This surely could be further optimized, as the bulk of this time is spent allocating some work arrays in the function. However in my particular application this is a negligible bottleneck as I am not simply checking an array with precomputed values but instead I have an array of pre-computed simulations from which I need to compute something at each node before doing the interpolation. Those evaluations make the ~120 nanosecond allocation cost just a small price to pay.