Class Frame

modelscope/models/science/unifold/modules/frame.py:183–459

class Frame:

    def __init__(
        self,
        rotation: Optional[Rotation],
        translation: Optional[torch.Tensor],
    ):
        if rotation is None and translation is None:
            rotation = Rotation.identity((0, ))
            translation = zero_translation((0, ))
        elif translation is None:
            translation = zero_translation(rotation.shape, rotation.dtype,
                                           rotation.device,
                                           rotation.requires_grad)

        elif rotation is None:
            rotation = Rotation.identity(
                translation.shape[:-1],
                translation.dtype,
                translation.device,
                translation.requires_grad,
            )

        if (rotation.shape != translation.shape[:-1]) or (rotation.device
                                                          !=  # noqa W504
                                                          translation.device):
            raise ValueError('RotationMatrix and translation incompatible')

        self._r = rotation
        self._t = translation

    @staticmethod
    def identity(
        shape: Iterable[int],
        dtype: Optional[torch.dtype] = torch.float,
        device: Optional[torch.device] = torch.device('cpu'),
        requires_grad: bool = False,
    ) -> Frame:
        return Frame(
            Rotation.identity(shape, dtype, device, requires_grad),
            zero_translation(shape, dtype, device, requires_grad),
        )

    def __getitem__(
        self,
        index: Any,
    ) -> Frame:
        if type(index) != tuple:
            index = (index, )

        return Frame(
            self._r[index],
            self._t[index + (slice(None), )],
        )

    def __mul__(
        self,
        right: torch.Tensor,