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

Code before:
    if is_lazy_array(n):
        # Lazy execution. We compute all special cases and the general case
        result = from_exp_coords(as_exp_coords(matrix) * n)
        identity = xp.eye(4, dtype=matrix.dtype, device=device)
        result = xp.where(n == 0, identity, result)
        result = xp.where(n == -1, inv(matrix), result)
        result = xp.where(n == 1, matrix, result)
        return result
    if n == 0:
        identity = xp.eye(4, dtype=matrix.dtype, device=device)
        identity = xpx.atleast_nd(identity, ndim=matrix.ndim, xp=xp)
        return xp.tile(identity, (*matrix.shape[:-2], 1, 1))
    elif n == -1:
        return inv(matrix)
    elif n == 1:
        return matrix
    return from_exp_coords(as_exp_coords(matrix) * n)


def mean(matrix: Array, weights: ArrayLike | None = None) -> Array:
    xp = array_namespace(matrix)
    if matrix.shape[0] == 0:
        raise ValueError("Mean of an empty rotation set is undefined.")

    quats = quat_from_matrix(matrix[..., :3, :3])
    if weights is None:
        quats_mean = quat_mean(quats)
    else:
        quats_mean = quat_mean(quats, weights=weights)
    r_mean = quat_as_matrix(quats_mean)

    t = matrix[..., :3, 3]
    axis = tuple(range(t.ndim - 1))
    if weights is None:
        t_mean = xp.mean(t, axis=axis)
    else:
        norm = xp.sum(weights[..., None], axis=axis)
        wsum = xp.sum(t * weights[..., None], axis=axis)
        t_mean = wsum / norm

    return _create_transformation_matrix(t_mean, r_mean)


def setitem(
    matrix: Array,
    indexer: Array | int | tuple | slice | EllipsisType | None,
    value: Array,
) -> Array:
    xp = array_namespace(matrix)
    if isinstance(indexer, EllipsisType):
        return xpx.at(matrix)[indexer].set(value)

Review comment (performance):
More food for thoughts on optimizations we could do: At the moment, every call to `from_matrix` is checking if the matrix is a valid rotation matrix, and normalizes the matrix for lazy frameworks. In cases like this where we call the function from within our own backend, the matrix is guaranteed to be valid and there is no need to normalize again. We could have a `from_matrix_unsafe` function that implements only the conversion from a valid matrix to a quaternion, and then only normalize inputs from outside the backend.

Not relevant for this PR, but it could significantly improve performance.

Code after:
    if is_lazy_array(n):
        # Lazy execution. We compute all special cases and the general case
        result = from_exp_coords(as_exp_coords(matrix) * n)
        identity = xp.eye(4, dtype=matrix.dtype, device=device)
        result = xp.where(n == 0, identity, result)
        result = xp.where(n == -1, inv(matrix), result)
        result = xp.where(n == 1, matrix, result)
        return result
    if n == 0:
        identity = xp.eye(4, dtype=matrix.dtype, device=device)
        identity = xpx.atleast_nd(identity, ndim=matrix.ndim, xp=xp)
        return xp.tile(identity, (*matrix.shape[:-2], 1, 1))
    elif n == -1:
        return inv(matrix)
    elif n == 1:
        return matrix
    return from_exp_coords(as_exp_coords(matrix) * n)


def mean(matrix: Array, weights: ArrayLike | None = None) -> Array:
    xp = array_namespace(matrix)
    if matrix.shape[0] == 0:
        raise ValueError("Mean of an empty rotation set is undefined.")

    lazy = is_lazy_array(matrix)
    quats = quat_from_matrix(matrix[..., :3, :3])
    if weights is None:
        quats_mean = quat_mean(quats)
    else:
        neg_weights = weights < 0
        any_neg_weights = xp.any(neg_weights)
        if not lazy and any_neg_weights:
            raise ValueError("`weights` must be non-negative.")
        if weights.shape != matrix.shape[:-2]:
            raise ValueError(
                f"Expected `weights` to match transform shape, got shape "
                f"{weights.shape} for {matrix.shape[:-2]} transformations."
            )
        quats_mean = quat_mean(quats, weights=weights)
    r_mean = quat_as_matrix(quats_mean)

    t = matrix[..., :3, 3]
    axis = tuple(range(t.ndim - 1))
    if weights is None:
        t_mean = xp.mean(t, axis=axis)
    else:
        norm = xp.sum(weights[..., None], axis=axis)
        wsum = xp.sum(t * weights[..., None], axis=axis)
        t_mean = wsum / norm

    tf = _create_transformation_matrix(t_mean, r_mean)