Repository: scipy/scipy
File: scipy/spatial/transform/_rigid_transform.py (Python)

Code before:
        >>> t = rng.random((4, 3))
        >>> r = R.random(4, rng=rng)
        >>> tf = Tf.from_components(t, r)
        >>> tf.mean().as_matrix()
        array([[ 0.61593485, -0.74508342,  0.25588075,  0.66999034],
               [-0.59353615, -0.65246765, -0.47116962,  0.25481794],
               [ 0.51801458,  0.13833531, -0.84411151,  0.52429339],
               [0., 0., 0., 1.]])
        >>> Tf.from_components(t.mean(axis=0), r.mean()).as_matrix()
        array([[ 0.61593485, -0.74508342,  0.25588075,  0.66999034],
               [-0.59353615, -0.65246765, -0.47116962,  0.25481794],
               [ 0.51801458,  0.13833531, -0.84411151,  0.52429339],
               [0., 0., 0., 1.]])
        """
        t, r = self.as_components()
        r_mean = r.mean(weights=weights)
        xp = array_namespace(t)
        if weights is None:
            t_mean = xp.mean(t, axis=tuple(range(t.ndim - 1)))
        else:
            axis = tuple(range(t.ndim - 1))
            norm = xp.sum(weights[..., None], axis=axis)
            wsum = xp.sum(t * weights[..., None], axis=axis)
            t_mean = wsum / norm
        mean = RigidTransform.from_components(t_mean, r_mean)
        return mean

    @xp_capabilities(
        skip_backends=[("dask.array", "missing linalg.cross/det functions")]
    )
    def as_matrix(self) -> Array:
        """Return a copy of the matrix representation of the transform.

        4x4 rigid transformation matrices are of the form::

            [       tx]
            [   R   ty]
            [       tz]
            [ 0 0 0  1]

        where ``R`` is a 3x3 orthonormal rotation matrix and
        ``t = [tx, ty, tz]`` is a 3x1 translation vector.

        Returns
        -------
        matrix : numpy.ndarray, shape (4, 4) or (N, 4, 4)
            A single transformation matrix or a stack of transformation
            matrices.

        Examples
        --------

Review comment (suggestion):
If we return the full matrix from the backend, we could use `_from_raw_matrix` to skip additional normalization steps on lazy backends.

Code after:
        >>> t = rng.random((4, 3))
        >>> r = R.random(4, rng=rng)
        >>> tf = Tf.from_components(t, r)
        >>> tf.mean().as_matrix()
        array([[ 0.61593485, -0.74508342,  0.25588075,  0.66999034],
               [-0.59353615, -0.65246765, -0.47116962,  0.25481794],
               [ 0.51801458,  0.13833531, -0.84411151,  0.52429339],
               [0., 0., 0., 1.]])
        >>> Tf.from_components(t.mean(axis=0), r.mean()).as_matrix()
        array([[ 0.61593485, -0.74508342,  0.25588075,  0.66999034],
               [-0.59353615, -0.65246765, -0.47116962,  0.25481794],
               [ 0.51801458,  0.13833531, -0.84411151,  0.52429339],
               [0., 0., 0., 1.]])
        """
        mean = self._backend.mean(self._matrix, weights=weights)
        return RigidTransform._from_raw_matrix(mean, xp=self._xp,
                                               backend=self._backend)

    @xp_capabilities(
        skip_backends=[("dask.array", "missing linalg.cross/det functions")]
    )
    def as_matrix(self) -> Array:
        """Return a copy of the matrix representation of the transform.

        4x4 rigid transformation matrices are of the form::

            [       tx]
            [   R   ty]
            [       tz]
            [ 0 0 0  1]

        where ``R`` is a 3x3 orthonormal rotation matrix and
        ``t = [tx, ty, tz]`` is a 3x1 translation vector.

        Returns
        -------
        matrix : numpy.ndarray, shape (..., 4, 4)
            Transformation matrices with the same leading dimensions as the transform.

        Examples
        --------
        >>> from scipy.spatial.transform import RigidTransform as Tf
        >>> from scipy.spatial.transform import Rotation as R
        >>> import numpy as np

        A transformation matrix is a 4x4 matrix formed from a 3x3 rotation
        matrix and a 3x1 translation vector:

        >>> t = np.array([2, 3, 4])
        >>> r = R.from_matrix([[0, 0, 1],
        ...                    [1, 0, 0],