Is it possible to improve a code related with iterative closest point (ICP) algorithm?

Question

I am trying to write a code related with ICP algorithm. The purpose of the code is matching a target data with a reference data by translational operation in XYZ axes (3-degree-of-freedom (DOF) operations). (It is better to have rotational operation (6-DOF in total), but it is not required.)

The code is shown below. Reference and target data (3D point cloud) are simpler than the actual data handled.

The flow of this code is as follows:

1. Find the nearest neighbor of each point in the reference data from the target data and calculate the distance between them.
2. The sum of the squares of the distances calculated for all the reference data is obtained.
3. The target data is shifted by a small distance in the X direction, and the same calculation as above is performed.
4. Repeating the small-distance shift gives the small-distance shift dependence of the sum of squares. The small-distance shift with the smallest sum of squares is the optimal target data position for the X-axis.
5. Do the same for the Y and Z axes.
6. If necessary, go back to step-1 and repeat the operations.

The problem of this code is too time consuming. I think the cause of the slow speed is using a For loop, but I couldn't think of a better code than this.

If you have a better idea, I would be very grateful if you could let me know.

Thank you in advance.

reference = {{2.2, 2.7, 2.1}, {2.3, 1.9, 2.1}, {1.8, 2.3, 2.3}};   (* Reference data. *)
data1 = {{1.6, 2.4, 1.8}, {0.4, 3.0, 1.2}, {2.1, 0.4, 2.4}, {1.5, 1.8, 2.0}, {2.9, 3.0, 2.1}, {2.2, 2.8, 1.7}, {1.0, 2.3, 1.8}, {1.8,0.9, 2.6}};    (* Target data (data to be shifted to match the reference data). *)


data1v1 = Table[{0.0, 0.0, 0.0}, {Length[data1]}];   (* Define data1v1 as an empty list. *)
totalsqdist= Table[{0.0, 0.0}, {21(*Number of steps in the following For-loops.*)}];    (* Define totalsqdist as an empty list. *)


(* Finding optimum x-shift distance. *)
For[dx = -1.0, dx <= 1.0, dx = dx + 0.1,    (* Iterative shift is from -1.0 to +1.0 with 0.1-step. *)
 data1v1 = data1[[#]] + {dx, 0.0, 0.0} & /@ Range[8];   (* Iteratively shift data1 by dx in x direction. *)
 
 nearestpoints = First[Nearest[data1v1, reference[[#]]]] & /@ Range[3];   (* Find the point in "data1" that is closest to each point in "reference". *)
 dist = EuclideanDistance[reference[[#]], nearestpoints[[#]]] & /@ Range[3];   (* Each distance between "reference" and "nearestpoints" *)
 sqdist = dist^2;   (* Square "dist". *)
 
 totalsqdist[[Round[(dx + 1.)*10 + 1], 1]] = dx;
 totalsqdist[[Round[(dx + 1.)*10 + 1], 2]] = Total[sqdist]    (* Total "sqdist". *)
 ]
xshift = MinimalBy[totalsqdist, Last][[1, 1]];    (* Optimum x-shift distance, which minimizes "totalsqdist". *)


(* Finding optimum y-shift distance. *)
For[dy = -1.0, dy <= 1.0, dy = dy + 0.1,    (* Iterative shift is from -1.0 to +1.0 with 0.1-step. *)
 data1v1 = data1[[#]] + {xshift, dy, 0.0} & /@ Range[8];   (* Shift data1 by xshift in x direction. Iteratively shift data1 by dy in y direction. *)
 
 nearestpoints = First[Nearest[data1v1, reference[[#]]]] & /@ Range[3];   (* Find the point in "data1" that is closest to each point in "reference". *)
 dist = EuclideanDistance[reference[[#]], nearestpoints[[#]]] & /@ Range[3];   (* Each distance between "reference" and "nearestpoints" *)
 sqdist = dist^2;   (* Square "dist". *)
 
 totalsqdist[[Round[(dy + 1.)*10 + 1], 1]] = dy;
 totalsqdist[[Round[(dy + 1.)*10 + 1], 2]] = Total[sqdist]    (* Total "sqdist". *)
 ]
yshift = MinimalBy[totalsqdist, Last][[1, 1]];    (* Optimum y-shift distance (, which minimizes "totalsqdist"). *)


(* Finding optimum z-shift distance. *)
For[dz = -1.0, dz <= 1.0, dz = dz + 0.1,    (* Iterative shift is from -1.0 to +1.0 with 0.1-step. *)
 data1v1 = data1[[#]] + {xshift, yshift, dz} & /@ Range[8];   (* Shift data1 by xshift and yshift in x and y direction, respectively. Iteratively shift data1 by dz in y direction. *)
 
 nearestpoints = First[Nearest[data1v1, reference[[#]]]] & /@ Range[3];   (* Find the point in "data1" that is closest to each point in "reference". *)
 dist = EuclideanDistance[reference[[#]], nearestpoints[[#]]] & /@ Range[3];   (* Each distance between "reference" and "nearestpoints" *)
 sqdist = dist^2;   (* Square "dist". *)
 
 totalsqdist[[Round[(dz + 1.)*10 + 1], 1]] = dz;
 totalsqdist[[Round[(dz + 1.)*10 + 1], 2]] = Total[sqdist]    (* Total "sqdist". *)
 ]
zshift = MinimalBy[totalsqdist, Last][[1, 1]];    (* Optimum z-shift distance (, which minimizes "totalsqdist"). *)


data1shifted = data1[[#]] + {xshift, yshift, zshift} & /@ Range[8];    (* Final result. *)


pl1 = {PlotLabel -> Style["Before shift process.", 16], PlotStyle -> {{PointSize[0.03], Black}, {PointSize[0.03], Lighter[Blue, 0.5]}}, AxesLabel -> {"x", "y", "z"}, ImageSize -> Medium, PlotRange -> {{0, 4}, {0, 4}, {0, 4}}};
pl2 = {PlotLabel -> Style["After shift process.", 16], PlotStyle -> {{PointSize[0.03], Black}, {PointSize[0.03], Lighter[Blue, 0.5]}}, AxesLabel -> {"x", "y", "z"}, ImageSize -> Medium, PlotRange -> {{0, 4}, {0, 4}, {0, 4}}};
Row[{ListPointPlot3D[{Legended[reference, "Reference"], Legended[data1, "data1 (to be shifted)"]}, pl1], "        ", ListPointPlot3D[{Legended[reference, "Reference"], Legended[data1shifted, "data1shifted"]}, pl2]}]

Answer 1

Okay, after I have seen that Newton's method does not work well here, here my reformulation of OP's algorithm. This implementation finds the same shifts, but it is far more efficient for large datasets.

ClearAll[findShift];
findShift[data_, reference_, {a_, b_, n_Integer}] := 
 Module[{shift, d, h, Y, nf},
  
  (* n+1 = the number of shifting operations per dimension.*)
  (* 
  a = leftmost shift in each direction*)
  (* 
  b rightmost shift in each direction*)
  
  (* The step size.*)
  h = (b - a)/n;
  
  (*Build only one nearest function for all shifting operations.*)
  
  nf = Nearest[data -> "Distance"];
  
  d = Dimensions[data][[-1]];
  (*a vector to store the shifts*)
  shift = ConstantArray[0., d];
  Y = reference;
  Do[
   Y[[All, direction]] -= (a - h);
   shift[[direction]] = MinimalBy[
      (*You can parallize this loop if you want.*)
      Table[
       (*The trick is to shift the reference, not the data! *)
       (*This allows us to reuse the Nearest function nf.*)
       Y[[All, direction]] -= h;
       {s, Total[nf[Y]^2, 2]}
       , {s, a, b, h}
       ],
      Last
      ][[1, 1]];
   (*Apply optimal shift in current direction.*)
   
   Y[[All, direction]] += b - shift[[direction]];
   , {direction, 1, d}];
  
  (*Finally, apply the shifts to the data once.*)
  
  data + ConstantArray[shift, Length[data]]
  ]

And here is a a usage example.

reference = RandomReal[{-1, 1}, {100, 3}];
data1 = RandomReal[{-1, 1}, {100000, 3}];
data1shifted2 = findShift[data1, reference, {-1., 1., 20}]; // AbsoluteTiming // First

0.027269

OP's code requires 5.87231 seconds for the same task on my machine. So this is an improvement of factor 200.

Newton's method is not very effective for this type of problem since there are many local minima. I keep the following only as reference.

The following implementation is basically Newton's method. It assumes that towards the end of the optimization, the relation between data points and closest reference points does not change anymore. In that regime, the objective function is quadratic and Newton converges in a single step. The Newton search direction is also a good choice outside of this regime, and so we can employ Armijo line search to globalize Newton's method. (Newton's method is notorious for wild behavior outsize a basin of attraction. One of the reasons (when the Hessian is positive-definite) is that the size of the step is typically too large.)

This improves two parts of OP's implementation:

1. It updates all coordinate directions at once.

2. It builds only one NearestFunction for the reference data and applies is multiple times. This is an advantage because everytime we call Nearest, a certain space partition tree hast to be build. So while the actual data queries are very fast, there is a certain overhead for building the seach index. By calling Nearest only once, we have to pay the overhead only once, of course.

Here is the actual code:

(*We have to generate only on NearestFunction for the reference data; then we can apply it repeatedly.*)
nf = Nearest[reference];
X = data1;

(*Some hard-coded parameters for the algorithm.*)
\[Sigma] = 0.5;
\[Gamma] = 0.5;
TOL = 1. 10^-8;
maxbacktrackings = 100;
maxiterations = 100;


Y = nf[X, 1][[All, 1]];
(*Objective: sum of squares*)
F = 1/2 Total[(Y - X)^2, 2];
(*This is the derivative of the objective.*)
DF = Total[X - Y];
(*This is the Newton search direction.*)
u = -DF/Length[X];
residualsquared = u.u;


iter = 0;
While[residualsquared > TOL^2 && iter < maxiterations,
  iter++;
  (*Armijo backtracking line search "with quadratic fit".*)
 
 (*Eventually, step size t = 1 should work because we use Newton's method.*)
  t = 1.;
  Xt = X + t ConstantArray[DF, Length[X]];
  Yt = nf[Xt, 1][[All, 1]];
  Ft = 1/2 Total[(Yt - Xt)^2, 2];
  
  slope = DF.u;
  backtrackings = 0;
  While[Ft > F - t \[Sigma] slope && 
    backtrackings < maxbacktrackings,
   backtrackings++;
   tnew = -\[Sigma] t^2 slope/(Ft - F - t slope);
   t = Max[\[Gamma] t, tnew];
   Xt = X + t ConstantArray[u, Length[X]];
   Yt = nf[Xt, 1][[All, 1]];
   Ft = 1/2 Total[(Yt - Xt)^2, 2]
   ];
  
  (*Now we have found a feasible step size t. Updating data*)
  X = Xt;
  Y = Yt;
  F = Ft;
  (*This is the derivative of the objective.*)
  DF = Total[X - Y];
  (*This is the Newton search direction.*)
  u = -DF/Length[X];
  residualsquared = u.u;
  ];

The result (the optimally translated data) is now stored in X.

On the example data, this converges with two iterations.

Edit

The algorithm minimizes the sum of squared distances from each reference point to its clostests point in the dataset. If you want to minimize the sum of squared distances of each data point from its closest point in the references, the following should do:

Z = reference;
X = data1;
nf = Nearest[X];
(*Some hard-coded parameters for the algorithm.*)
\[Sigma] = 0.5;
\[Gamma] = 0.5;
TOL = 1. 10^-8;
maxbacktrackings = 20;
maxiterations = 100;

Y = nf[Z, 1][[All, 1]];
(*Objective:sum of squares*)
F = 1/2 Total[(Z - Y)^2, 2];
Print["initial fit" -> F];
(*This is the derivative of the objective.*)
DF = Total[Y - Z];
(*This is the Newton search direction.*)
u = -DF/Length[Z];
residualsquared = u.u;

iter = 0;
While[
  residualsquared > TOL^2 && iter < maxiterations,
  iter++;
  (*Armijo backtracking line search "with quadratic \
fit".*)(*Eventually,step size t=
  1 should work because we use Newton's method.*)t = 1.;
  t = 1.;
  Xt = X + t ConstantArray[u, Length[X]];
  nft = Nearest[Xt];
  Yt = nft[Z, 1][[All, 1]];
  Ft = 1/2 Total[(Z - Yt)^2, 2];
  slope = DF.u;
  backtrackings = 0;
  While[Ft > F + t \[Sigma] slope && backtrackings < maxbacktrackings,
    backtrackings++;
   tnew = -\[Sigma] t^2 slope/(Ft - F - t slope);
   t = Max[\[Gamma] t, tnew];
   Xt = X + t ConstantArray[u, Length[X]];
   nft = Nearest[Xt];
   Yt = nft[Z, 1][[All, 1]];
   Ft = 1/2 Total[(Z - Yt)^2, 2];
   ];
  (*Now we have found a feasible step size t.Updating data*)
  
  X = Xt;
  Y = Yt;
  F = Ft;
  nf = nft;
  (*This is the derivative of the objective.*)
  DF = Total[Y - Z];
  (*This is the Newton search direction.*)
  u = -DF/Length[Z];
  residualsquared = u.u;
  ];
Print["final fit" -> F];
Print["residual" -> Sqrt[residualsquared]];