# How to vectorize rotation of 3D points?

**URL:** <https://discuss.pytorch.org/t/how-to-vectorize-rotation-of-3d-points/69476>\
**Category:** Uncategorized\
**Created:** [February 12, 2020, 11:43am UTC](https://discuss.pytorch.org/t/how-to-vectorize-rotation-of-3d-points/69476 "2020-02-12T11:43:20Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![lui.bonati](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lui.bonati/32/24092_2.png) [@lui.bonati](https://discuss.pytorch.org/u/lui.bonati)\
**Post date:** [February 12, 2020, 11:43am UTC](https://discuss.pytorch.org/t/how-to-vectorize-rotation-of-3d-points/69476/1 "2020-02-12T11:43:20Z")

</div>

Dear all,  
I have a matrix (B, 3N), where the rows are the elements in the batch and the columns correspond to N points in a 3D space, like (x1,y1,z1,x2,y2,z2,…,xN,yN,zN). I can reshape it to be (B,N,3).

Now I want to apply a rotation to the dataset, projecting each 3D point on three versors e\_x, e\_y, e\_z. Each of them is a matrix of shape (B,3) since the reference frame is different for each element in the batch.

I have implemented this with a for-loop over the batches and each atom, in the following way:

```auto
# x.shape => (B,N,3)
# e_x.shape => (B,3) {also e_y and e_x}

x_new = torch.zeros((B,N,3))

for b in range(B): # loop over batch
   for j in range(N): # loop over atoms
      x_new[b,j,0]=torch.dot(x[b,j,:], e_x[b,:]) # x component of atom j in the new reference frame 
      x_new[b,j,1]=torch.dot(x[b,j,:], e_y[b,:]) # y component
      x_new[b,j,2]=torch.dot(x[b,j,:], e_z[b,:]) # z component

```

but of course it is really slow. Do you have any suggestion about how to vectorize this calculation avoiding the for loops? Should I concatenate e\_x, e\_y, e\_z in a batch of matrices and then apply it to the data?

Thank you very much,  
Luigi

---

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [February 12, 2020, 12:57pm UTC](https://discuss.pytorch.org/t/how-to-vectorize-rotation-of-3d-points/69476/3 "2020-02-12T12:57:20Z")

</div>

Soz, version 2 just to avoid doing a for loop over each dimension:

```auto
import torch
# x.shape => (B,N,3)
# e_x.shape => (B,3) {also e_y and e_x}
B=2
N=4
x=torch.rand(B,N,3)
x_custom=x.unsqueeze(3).expand(B,N,3,3)
e_x=torch.rand(B,3)
e_y=torch.rand(B,3)
e_z=torch.rand(B,3)

x_new = torch.zeros((B,N,3))

for b in range(B): # loop over batch
   for j in range(N): # loop over atoms
      x_new[b,j,0]=torch.dot(x[b,j,:], e_x[b,:]) # x component of atom j in the new reference frame
      x_new[b,j,1]=torch.dot(x[b,j,:], e_y[b,:]) # y component
      x_new[b,j,2]=torch.dot(x[b,j,:], e_z[b,:]) # z component

einsum = torch.einsum('bijk,bjk->bik',x_custom,torch.stack([e_x,e_y,e_z],dim=2))

print(torch.norm((x_new-einsum)))

```