performance tuning - Compiling Map over expression that yields a ragged array

I'm trying to speed up a function that looks in the neighborhood of each 3D point in a large dataset and finds all the points within 1 unit in each direction, x, y, z.

I've started by using Select to find the points around a specified point and then Map this function over the dataset.

First, lets make up some "data" and then define the findpoints function.

data = RandomReal[10, {10^4, 3}];

findpoints[data_, pt_]:= Block[{x = pt[[1]], y = pt[[2]], z = pt[[3]]},  
  Select[data, 
   x - 1 <= #[[1]] <= x + 1 && y - 1 <= #[[2]] <= y + 1 && 
     z - 1 <= #[[3]] <= z + 1 &]
]

Now lets map it over the random data.

Map[findpoints[data, #] &, data]; // AbsoluteTiming

This is quite slow, I gave up waiting, so I started by compiling the findpoint function.

findpointsC = 
  Compile[{{data, _Real, 2}, {pt, _Real, 1}}, 
   Block[{x = pt[[1]], y = pt[[2]], z = pt[[3]]}, 
    Select[data, 
     x - 1 <= #[[1]] <= x + 1 && y - 1 <= #[[2]] <= y + 1 && 
       z - 1 <= #[[3]] <= z + 1 &]
    ], CompilationTarget -> "C", RuntimeAttributes -> {Listable}, 

   Parallelization -> True, RuntimeOptions -> "Speed"];

output= Map[findpointsC[data, #] &, data]; // AbsoluteTiming

(*{3.7542147, Null}*)

Much faster, great.

Now I figured since compiling was so beneficial here and in other functions I found that by compiling the Map function I can gain some more speed, lets compile the whole thing.

mapfindpointsC = 
  Compile[{{data, _Real, 2}}, Map[findpointsC[data, #] &, data], 

   CompilationTarget -> "C", RuntimeAttributes -> {Listable}, 
   Parallelization -> True, RuntimeOptions -> "Speed"];

mapfindpointsC[data]

(*CompiledFunction::cflist: Nontensor object generated; proceeding with uncompiled evaluation. >>*)

However, as per the error message, in this case the output is ragged.

TensorQ[output]


(*False*)

I've tried to flatten the output from map but nothing I can think of works. Is it possible to compile this expression.

Answer

Please compare

Map[findpointsC[data, #] &, data]; // AbsoluteTiming
(* {3.567939, Null} *)

and

findpointsC[data, data] ; // AbsoluteTiming

(* {0.841954, Null} *)

And of course is

Map[findpointsC[data, #] &, data] == findpointsC[data, data] 
(* True *)

The explanation is simple: You typed RuntimeAttributes->{Listable} and Parallelization but you have not used it.

Explanation for your error message

The problem with your Map arises because Compiled functions can only work with rectangular arrays. Since every point has (most likely) a different number of points in its neighborhood, the result array is ragged, or not rectangular.

It's not enough to use Flatten because even then, you still create such a ragged array. One (slower) way out of this is to collect all points in a linear list, but then you loose the information, which point-set was created by which point. This cannot be the goal, if the solution is as simple as described above. Anyway, as you might have seen we had a topic here about Internal`Bag which can be used to collect an unknown number of elements. Please see the discussion in the thread

Internal`Bag inside Compile

Based on this, you could make an all in one compiled function which I only give for reference and not for usage. It is much slower than the Listable approach. I use the already compiled function findpointsC by using With and giving the option

CompilationOptions -> {"InlineCompiledFunctions" -> True}

I have replaced your Map with a Do loop which runs over all points and finds the nearest neighbors. To insert such a list of points into an Internal`Bag we have to Flatten it which can be done with the 3rd parameter of Internal`StuffBag. To reshape the flat list at the end into a list of points I use Partition

mapfindpointsC =
 With[{fpC = findpointsC},
  Compile[{{data, _Real, 2}},
   Module[{result = Internal`Bag[Most[{.0}]]},
    Do[Internal`StuffBag[result, fpC[data, pt], 2], {pt, data}];

    Partition[Internal`BagPart[result, All], 3]
    ], CompilationTarget -> "C", RuntimeAttributes -> {Listable}, 
   Parallelization -> True, RuntimeOptions -> "Speed", 
   CompilationOptions -> {"InlineCompiledFunctions" -> True}]
]

mapfindpointsC[data]; // AbsoluteTiming

(* {3.831499, Null} *)