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

Code before:
    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)
    if is_array_api_obj(indexer) and indexer.dtype == xp.bool:
        return xpx.at(matrix)[indexer].set(value)
    return xpx.at(matrix)[indexer, ...].set(value)


def normalize_dual_quaternion(dual_quat: Array) -> Array:
    """Normalize dual quaternion."""
    xp = array_namespace(dual_quat)
    real, dual = _normalize_dual_quaternion(dual_quat[..., :4], dual_quat[..., 4:])
    return xp.concat((real, dual), axis=-1)


def _create_transformation_matrix(
    translations: Array, rotation_matrices: Array

Review comment (suggestion):
We're implicitly relying on the check https://github.com/scipy/scipy/blob/e407bc4d6ee71ec1a23fce9d95b58e28451e4d94/scipy/spatial/transform/_rotation_xp.py#L496-L500 here. Are we fine with this? We could also relax that check further such that `weights` only needs to be broadcastable to `quat` / `matrix`.

Code after:
    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)
    if weights is not None and lazy:
        # We cannot raise on negative weights because jit code needs to be
        # non-branching. We return NaN instead
        mask = xp.where(any_neg_weights, xp.nan, 1.0)
        tf = mask * tf
    return tf


def setitem(
    matrix: Array,
    indexer: Array | int | tuple | slice | EllipsisType | None,
    value: Array,
) -> Array:
    xp = array_namespace(matrix)