Optimzie the implementation of StructureModule.

Author: Xreki

apps/protein_folding/helixfold/alphafold_paddle/model/r3.py

@@ -410,22 +410,28 @@ def rots_from_two_vecs(e0_unnormalized: Vecs, e1_unnormalized: Vecs) -> Rots:
 
 def rots_mul_rots(a: Rots, b: Rots) -> Rots:
     """Composition of rotations 'a' and 'b'."""
-    c0 = rots_mul_vecs(a, Vecs(b.xx, b.yx, b.zx))
-    c1 = rots_mul_vecs(a, Vecs(b.xy, b.yy, b.zy))
-    c2 = rots_mul_vecs(a, Vecs(b.xz, b.yz, b.zz))
-    return Rots(c0.x, c1.x, c2.x, c0.y, c1.y, c2.y, c0.z, c1.z, c2.z)
+    if a.shape == b.shape:$
+        return Rots(paddle.matmul(a.rotation, b.rotation))
+    else:
+        c0 = rots_mul_vecs(a, Vecs(b.xx, b.yx, b.zx))
+        c1 = rots_mul_vecs(a, Vecs(b.xy, b.yy, b.zy))
+        c2 = rots_mul_vecs(a, Vecs(b.xz, b.yz, b.zz))
+        return Rots(c0.x, c1.x, c2.x, c0.y, c1.y, c2.y, c0.z, c1.z, c2.z)
 
 
 def rots_mul_vecs(m: Rots, v: Vecs) -> Vecs:
     """Apply rotations 'm' to vectors 'v'."""
-    return Vecs(m.xx * v.x + m.xy * v.y + m.xz * v.z,
-                m.yx * v.x + m.yy * v.y + m.yz * v.z,
-                m.zx * v.x + m.zy * v.y + m.zz * v.z)
+    v_x = v.x
+    v_y = v.y
+    v_z = v.z
+    return Vecs(m.xx * v_x + m.xy * v_y + m.xz * v_z,
+                m.yx * v_x + m.yy * v_y + m.yz * v_z,
+                m.zx * v_x + m.zy * v_y + m.zz * v_z)
 
 
 def vecs_add(v1: Vecs, v2: Vecs) -> Vecs:
     """Add two vectors 'v1' and 'v2'."""
-    return Vecs(v1.x + v2.x, v1.y + v2.y, v1.z + v2.z)
+    return Vecs(v1.translation + v2.translation)
 
 
 def vecs_dot_vecs(v1: Vecs, v2: Vecs) -> paddle.Tensor: