numpy.linalg.lstsq 的向量化

Question

I have two sets of points (A1, A2, B1, B2) for which I want to calculate the affine transformation (from A1 to B1, from A2 to B2). Using numpy.linalg.lstsq this is very straightforward for a single case:

A1 = [[100   0   0]
      [  0 100   0]
      [100 100   0]
      [  0   0 100]
      [100 100 100]]

B1 = [[160   0   0]
      [  0 160   0]
      [160 160   0]
      [  0   0 160]
      [160 160 160]]

A1 = np.hstack([A1, np.ones((A1.shape[0], 1))])
B1 = np.hstack([B1, np.ones((B2.shape[0], 1))])
affine_transformation = np.linalg.lstsq(A1, B1, rcond=None)[0].transpose()

I want to vectorise this to calculate it for multiple sets of points without a loop. I would like to end up with something like this:

pts_a = np.array([A1, A2])
pts_a = np.pad(pts_a, ((0, 0), (0, 0), (0,1)), mode='constant', constant_values=1)
pts_b = np.array([B1, B2])
pts_b = np.pad(pts_b, ((0, 0), (0, 0), (0,1)), mode='constant', constant_values=1)
affine_transformation = np.linalg.lstsq(pts_a, pts_b, rcond=None)[0].transpose()

Any help is much appreciated!

Answer 1

I sorted it by using this:

 affine_transformation = np.asarray([np.linalg.lstsq(pts_a[i,:,:], pts_b[i,:,:], rcond=None)[0].transpose() for i in range(pts_a.shape[0])])

Is there a more efficient way?