Remove laziness to speed up computations by about 25%

Author: LemonPi

src/pytorch_kinematics/transforms/transform3d.py

@@ -212,7 +212,6 @@ def __init__(
             rot_h = torch.cat((rot, zeros), dim=-2).reshape(-1, 4, 3)
             self._matrix = torch.cat((rot_h, self._matrix[:, :, 3].reshape(-1, 4, 1)), dim=-1)
 
-        self._transforms = []  # store transforms to compose
         self._lu = None
         self.device = device
         self.dtype = self._matrix.dtype
@@ -240,13 +239,12 @@ def compose(self, *others):
         Returns:
             A new Transform3d with the stored transforms
         """
-        out = Transform3d(device=self.device, dtype=self.dtype)
-        out._matrix = self._matrix.clone()
+
+        mat = self._matrix
         for other in others:
-            if not isinstance(other, Transform3d):
-                msg = "Only possible to compose Transform3d objects; got %s"
-                raise ValueError(msg % type(other))
-        out._transforms = self._transforms + list(others)
+            mat = _broadcast_bmm(mat, other.get_matrix())
+
+        out = Transform3d(device=self.device, dtype=self.dtype, matrix=mat)
         return out
 
     def get_matrix(self):
@@ -266,11 +264,7 @@ def get_matrix(self):
         Returns:
             A transformation matrix representing the composed inputs.
         """
-        composed_matrix = self._matrix
-        for other in self._transforms:
-            other_matrix = other.get_matrix()
-            composed_matrix = _broadcast_bmm(composed_matrix, other_matrix)
-        return composed_matrix
+        return self._matrix
 
     def _get_matrix_inverse(self):
         """
@@ -311,40 +305,16 @@ def inverse(self, invert_composed: bool = False):
             transformation.
         """
 
-        tinv = Transform3d(device=self.device)
+        i_matrix = self._get_matrix_inverse()
 
-        if invert_composed:
-            # first compose then invert
-            tinv._matrix = self._invert_transformation_matrix(self.get_matrix())
-        else:
-            # self._get_matrix_inverse() implements efficient inverse
-            # of self._matrix
-            i_matrix = self._get_matrix_inverse()
-
-            # 2 cases:
-            if len(self._transforms) > 0:
-                # a) Either we have a non-empty list of transforms:
-                # Here we take self._matrix and append its inverse at the
-                # end of the reverted _transforms list. After composing
-                # the transformations with get_matrix(), this correctly
-                # right-multiplies by the inverse of self._matrix
-                # at the end of the composition.
-                tinv._transforms = [t.inverse() for t in reversed(self._transforms)]
-                last = Transform3d(device=self.device)
-                last._matrix = i_matrix
-                tinv._transforms.append(last)
-            else:
-                # b) Or there are no stored transformations
-                # we just set inverted matrix
-                tinv._matrix = i_matrix
+        tinv = Transform3d(matrix=i_matrix, device=self.device)
 
         return tinv
 
     def stack(self, *others):
         transforms = [self] + list(others)
         matrix = torch.cat([t._matrix for t in transforms], dim=0)
-        out = Transform3d()
-        out._matrix = matrix
+        out = Transform3d(matrix=matrix, device=self.device, dtype=self.dtype)
         return out
 
     def transform_points(self, points, eps: Optional[float] = None):
@@ -478,7 +448,6 @@ def clone(self):
         if self._lu is not None:
             other._lu = [elem.clone() for elem in self._lu]
         other._matrix = self._matrix.clone()
-        other._transforms = [t.clone() for t in self._transforms]
         return other
 
     def to(self, device, copy: bool = False, dtype=None):
@@ -504,7 +473,6 @@ def to(self, device, copy: bool = False, dtype=None):
         other.device = device
         other.dtype = dtype if dtype is not None else other.dtype
         other._matrix = self._matrix.to(device=device, dtype=dtype)
-        other._transforms = [t.to(device, copy=copy, dtype=dtype) for t in other._transforms]
         return other
 
     def cpu(self):