pose.py

import numpy as np
from numpy import ndarray
from typing import *
from numbers import Number
from ..helpers import no_warnings

from .transforms import make_affine_matrix, transform_points
from .utils import safe_inv, vector_outer


__all__ = [
    'kabasch',
    'umeyama',
    'affine_umeyama',
    'solve_pose',
    'segment_solve_pose',
    'solve_poses_sequential',
    'segment_solve_poses_sequential',
]


def kabasch(cov: ndarray) -> ndarray:
    U, _, Vh = np.linalg.svd(cov)
    Vh[..., 2, :] *= np.sign(np.linalg.det(U @ Vh))[..., None]
    R = U @ Vh
    return R


def umeyama(cov_yx: ndarray, cov_xx: Optional[ndarray] = None, cov_yy: Optional[ndarray] = None, mean_x: Optional[ndarray] = None, mean_y: Optional[ndarray] = None) -> Tuple[ndarray, ndarray, ndarray]:
    """
    Procrustes analysis to solve for scale `s`, rotation `R` and translation `t` such that `y_i ~= s R x_i + t`.

    Parameters
    ----
    - `cov_yx`: (..., 3, 3) covariance matrix between y and x points.
    - `cov_xx`: (..., 3, 3) covariance matrix of x points. If None, no scaling is solved.
    - `cov_yy`: (..., 3, 3) covariance matrix of y points. If None, no scaling is solved.
    - `mean_x`: (..., 3) mean of x points. If None, no translation is solved.
    - `mean_y`: (..., 3) mean of y points. If None, no translation is solved.

    Specifically, based on provided inputs:
    
    - To solve the rotation `R`, `cov_yx` must be given.
    - To solve the scale `s`, at least one of `cov_xx` and `cov_yy` must be given.
        - (Recommended) If both `cov_xx` and `cov_yy` are given, the scale will be solved by minimizing a symmetric cost:
            `||s R X + t - Y||_F^2 / ||Y||_F^2 + ||s R^T (Y - t)  - X||_F^2 / ||X||_F^2`
        - If only `cov_xx` is given, the scale will be solved by minimizing forward cost
            `||s R X  + t - Y||_F^2`
        - If only `cov_yy` is given, the scale will be solved by minimizing inverse cost 
            `||s R^T (Y - t)  - X||_F^2`
    - To solve the translation `t`, provide `mean_x` and `mean_y`.

    Returns
    ----
    - `s`: (...) scale factor. None if both cov_xx and cov_yy are None. 
    - `R`: (..., 3, 3) rotation matrix.
    - `t`: (..., 3) translation vector. None if mean_x or mean_y is None.
    """
    dtype = cov_yx.dtype
    R = kabasch(cov_yx)
    if cov_xx is not None and cov_yy is None:
        s = np.trace(cov_yx @ R.swapaxes(-2, -1), axis1=-2, axis2=-1) / np.maximum(np.trace(cov_xx, axis1=-2, axis2=-1), np.finfo(dtype).tiny)
    if cov_xx is None and cov_yy is not None:
        s = np.trace(cov_yy, axis1=-2, axis2=-1) / np.maximum(np.trace(cov_yx @ R.swapaxes(-2, -1), axis1=-2, axis2=-1), np.finfo(dtype).tiny)
    elif cov_xx is not None and cov_yy is not None:
        x_fnorm = np.maximum(np.trace(cov_xx, axis1=-2, axis2=-1), np.finfo(dtype).tiny)
        y_fnorm = np.maximum(np.trace(cov_yy, axis1=-2, axis2=-1), np.finfo(dtype).tiny)
        s = np.sqrt(y_fnorm / x_fnorm)
    else:
        s = None
    if mean_x is not None and mean_y is not None:
        if s is not None:
            t = mean_y - transform_points(mean_x, s[..., None, None] * R)
        else:
            t = mean_y - transform_points(mean_x, R)
    else:
        t = None
    return s, R, t


def affine_umeyama(cov_yx: ndarray, cov_xx: ndarray, cov_yy: ndarray, mean_x: ndarray, mean_y: ndarray, lam: float = 1e-2, niter: int = 8) -> Tuple[ndarray, ndarray]:
    """
    Extended Procrustes analysis to solve for affine transformation `A` and translation `t` such that `y_i ~= A x_i + t`.

    Parameters
    ----
    - `cov_yx`: (..., 3, 3) covariance matrix between y
    - `cov_xx`: (..., 3, 3) covariance matrix of x points.
    - `cov_yy`: (..., 3, 3) covariance matrix of y
    - `mean_x`: (..., 3) mean of x points.
    - `mean_y`: (..., 3) mean of y points.
    - `lam`: rigidity regularization weight.
    - `gamma`: symmetricity regularization annealing factor.
    - `niter`: number of iterations for solving.

    Returns
    ----
    - `A`: (..., 3, 3) affine transformation matrix.
    - `t`: (..., 3) translation vector.
    """
    dtype = cov_yx.dtype
    R = kabasch(cov_yx)
    tr_xx = np.maximum(np.trace(cov_xx, axis1=-2, axis2=-1), np.finfo(dtype).tiny)
    tr_yy = np.maximum(np.trace(cov_yy, axis1=-2, axis2=-1), np.finfo(dtype).tiny)
    
    cov_yx, cov_xy = cov_yx / tr_xx[..., None, None], cov_yx.swapaxes(-2, -1) / tr_yy[..., None, None]
    cov_xx, cov_yy = cov_xx / tr_xx[..., None, None], cov_yy / tr_yy[..., None, None]
    
    A, B = np.zeros_like(R), np.zeros_like(R)
    I = np.eye(cov_yx.shape[-1], dtype=dtype)
    
    def _step(A, B, R, cov_yx, cov_xy, cov_xx, cov_yy, lam, gamma):
        A = (cov_yx + lam * R + gamma * B.swapaxes(-2, -1)) @ safe_inv(cov_xx + lam * I + gamma * (B @ B.swapaxes(-2, -1)))
        B = (cov_xy + lam * R.swapaxes(-2, -1) + gamma * A.swapaxes(-2, -1)) @ safe_inv(cov_yy + lam * I + gamma * (A @ A.swapaxes(-2, -1)))
        err = np.square(A @ B - I).mean(axis=(-2, -1))
        return A, B, err
    
    not_converged = np.argwhere(np.ones(R.shape[:-2], dtype=bool))
    for i in range(niter):
        gamma_i = 1.2 ** i - 1
        non_converged_indices = tuple(not_converged.T)
        A[non_converged_indices], B[non_converged_indices], err = _step(*(x[non_converged_indices] for x in (A, B, R, cov_yx, cov_xy, cov_xx, cov_yy)), lam, gamma_i)
        not_converged = not_converged[err >= 1e-6]
        if len(not_converged) == 0:
            break
    t = mean_y - transform_points(mean_x, A)
    return A, t


def solve_pose(
    p: np.ndarray, 
    q: np.ndarray, 
    w: Optional[np.ndarray] = None, 
    *,
    mode: Literal['rigid', 'similar', 'affine'] = 'rigid',
    lam: float = 1e-2, 
    niter: int = 5
) -> np.ndarray:
    """Solve for the pose (transformation from p to q) given weighted point correspondences.
    
    Parameters
    ----
    - `p`: (..., N, 3) source points
    - `q`: (..., N, 3) target points
    - `w`: optional (..., N) weights for each point correspondence. If None, uniform weights are used.
    - `mode`: mode of transformation to apply. Can be 'rigid', 'similar', or 'affine'.
        - For 'rigid', only rotation and translation are allowed.
        - For 'similar', uniform scaling, rotation and translation are allowed.
        - For 'affine', full affine transformation is allowed. Using least squares.
    - `lam`: regularization weight for 'affine' mode.
    - `niter`: number of iterations for 'affine' mode.
    
    Returns
    ----
    - `pose`: (..., 4, 4) transformations matrix from p to q.
    """
    if w is None:
        w = np.ones(p.shape[:-1], dtype=p.dtype)
    w_sum = np.maximum(np.sum(w, axis=-1), np.finfo(p.dtype).tiny)
    p_mean = np.sum(w[..., None] * p, axis=-2) / w_sum[..., None]
    q_mean = np.sum(w[..., None] * q, axis=-2) / w_sum[..., None]
    p = p - p_mean[..., None, :]
    q = q - q_mean[..., None, :]
    pw = p * w[..., None]
    qw = q * w[..., None]
    cov_qp = np.sum(vector_outer(qw, p), axis=-3) / w_sum[..., None, None]
    if mode == 'similar' or mode == 'affine':
        cov_pp = np.sum(vector_outer(pw, p), axis=-3) / w_sum[..., None, None]
        cov_qq = np.sum(vector_outer(qw, q), axis=-3) / w_sum[..., None, None]
    
    if mode == 'rigid':
        _, R, t = umeyama(cov_qp, mean_x=p_mean, mean_y=q_mean)
        pose = make_affine_matrix(R, t)
    elif mode == 'similar':
        s, R, t = umeyama(cov_qp, cov_xx=cov_pp, cov_yy=cov_qq, mean_x=p_mean, mean_y=q_mean)
        pose = make_affine_matrix(s * R, t)
    elif mode == 'affine':
        A, t = affine_umeyama(cov_qp, cov_pp, cov_qq, p_mean, q_mean, lam=lam, niter=niter)
        pose = make_affine_matrix(A, t)
    
    return pose


def segment_solve_pose(
    p: np.ndarray, 
    q: np.ndarray, 
    w: Optional[np.ndarray] = None, 
    *,
    offsets: np.ndarray, 
    mode: Literal['rigid', 'similar', 'affine'] = 'rigid',
    lam: float = 1e-2, 
    niter: int = 5
) -> np.ndarray:
    """Solve for the pose (transformation from p to q) given weighted point correspondences.
    
    Parameters
    ----
    - `p`: (N, 3) source points
    - `q`: (N, 3) target points
    - `w`: (N,) weights for each point correspondence
    - `offsets`: (S + 1,) segment offsets. Points in each segment belong to the same rigid / affine body.
    - `mode`: mode of transformation to apply. Can be 'rigid', 'similar', or 'affine'.
        - For 'rigid', only rotation and translation are allowed.
        - For 'similar', uniform scaling, rotation and translation are allowed.
        - For 'affine', full affine transformation is allowed. Using least squares.
    - `lam`: regularization weight for 'affine' mode.
    - `niter`: number of iterations for 'affine' mode.
    
    Returns
    ----
    - `pose`: (S, 4, 4) transformations matrix from p to q.
    """
    if w is None:
        w = np.ones(p.shape[:-1], dtype=p.dtype)

    lengths = np.diff(offsets)
    w_sum = np.maximum(np.add.reduceat(w, offsets[:-1], axis=0), np.finfo(p.dtype).tiny)
    p_mean = np.add.reduceat(w[..., None] * p, offsets[:-1], axis=0) / w_sum[:, None]
    q_mean = np.add.reduceat(w[..., None] * q, offsets[:-1], axis=0) / w_sum[:, None]
    p = p - np.repeat(p_mean, lengths, axis=0)
    q = q - np.repeat(q_mean, lengths, axis=0)
    pw = p * w[..., None]
    qw = q * w[..., None]
    cov_qp = np.add.reduceat(vector_outer(qw, p), offsets[:-1], axis=0) / w_sum[:, None, None]
    if mode == 'similar' or mode == 'affine':
        cov_pp = np.add.reduceat(vector_outer(pw, p), offsets[:-1], axis=0) / w_sum[:, None, None]    
        cov_qq = np.add.reduceat(vector_outer(qw, q), offsets[:-1], axis=0) / w_sum[:, None, None]
    
    if mode == 'rigid':
        _, R, t = umeyama(cov_qp, mean_x=p_mean, mean_y=q_mean)
        pose = make_affine_matrix(R, t)
    elif mode == 'similar':
        s, R, t = umeyama(cov_qp, cov_xx=cov_pp, cov_yy=cov_qq, mean_x=p_mean, mean_y=q_mean)
        pose = make_affine_matrix(s * R, t)
    elif mode == 'affine':
        A, t = affine_umeyama(cov_qp, cov_pp, cov_qq, p_mean, q_mean, lam=lam, niter=niter)
        pose = make_affine_matrix(A, t)
    
    return pose


def solve_poses_sequential(
    trajectories: ndarray,
    weights: Optional[ndarray] = None,
    *,
    accum: Optional[Tuple[ndarray, ...]] = None,
    min_valid_size: int = 3,
    mode: Literal['rigid', 'similar', 'affine'] = 'rigid',
    lam: float = 1e-2,
    niter: int = 8
) -> Tuple[ndarray, Tuple[ndarray, ...], Tuple[ndarray, ndarray, ndarray, ndarray]]:
    """
    Given trajectories of points over time, sequentially solve for the poses (transformations from canonical to each frame) of each body at each frame.

    Parameters
    ----
    - `trajectories`: (T, ..., N, 3) posed points. T is number of frames. `...` is optional batch dimensions. N is number of points per group.
    - `weights`: (T, ..., N) quardratic error term weights for each point at each frame
    - `accum`: accumulated statistics from previous calls. If None, start fresh.
    - `min_valid_size`: minimum number of valid points in each frame to consider the segment / group valid.
    - `mode`: mode of transformation to apply. Can be 'rigid', 'similar', or 'affine'.
        - For 'rigid', only rotation and translation are allowed.
        - For 'similar', uniform scaling, rotation and translation are allowed. 
        - For 'affine', full affine transformation is allowed. Using least squares.
    - `lam`: rigidity regularization weight for 'affine' mode.
    - `niter`: number of iterations for 'affine' mode.

    Returns
    ----
    - `poses`: (T, ..., 4, 4) transformations from canonical to each frame.
    - `valid`: (T, ...) boolean mask indicating valid segments
    - `stats`: canonical statistics of each group,
        It is a tuple of:
        - `mu`: (..., 3) weighted mean of points
        - `cov`: (..., 3, 3) weighted covariance of points
        - `tot_w`: (...,) total weight of points
        - `nnz`: (...,) number of non-zero weight points
    - `canonical_points`: (..., N, 3) canonical points.
    - `err`: (..., N,) per-point RMS error over all time := sqrt(sum_over_time(per_point_weights * per_point_squared_error) / per_point_nnz)
        Use this to filter outliers as needed.
    - `accum`: per point accumulated statistics. Just pass it to the next call for incremental solving.
        It is a tuple of:
        - `accum_sqrtw`: (..., N,) sum of sqrt(weights)
        - `accum_sqrtwx`: (..., N, 3) sum of sqrt(weights) * x
        - `accum_sqrtwxx`: (...N, 3, 3) sum of sqrt(weights) * outer(x - mean_sqrtwx, x - mean_sqrtwx)
        - `accum_w`: (..., N,) sum of weights
        - `accum_wx`: (..., N, 3) sum of weights * x
        - `accum_wxx`: (..., N, 3, 3) sum of weights * outer(x - mean_wx, x - mean_wx)
        - `accum_nnz`: (..., N,) number of non-zero weight accumulations

    Example
    ----
    ```
    accum = None
    poses, valid = [], []
    for new_trajectories_chunk in data_stream:
        # new_trajectories_chunk: (T_chunk, N, 3)
        poses_chunk, valid_chunk, stats, canonical_points, err, accum = solve_poses(
            new_trajectories_chunk,
            accum=accum,
        )
        poses.append(poses_chunk)
        valid.append(valid_chunk)
        # `stats`, `canonical_points` and `err` are returned and updated every chunk.
    poses = np.concatenate(poses, axis=0)   # (T_all, 4, 4), poses over all frames
    valid = np.concatenate(valid, axis=0)   # (T_all,), poses' validity over all frames
    """
    dtype = trajectories.dtype
    num_frames = trajectories.shape[0]
    num_points = trajectories.shape[-2]
    batch_shape = trajectories.shape[1:-2]

    if weights is None:
        weights = np.ones((num_frames, *batch_shape, num_points), dtype=dtype)

    poses = np.zeros((num_frames, *batch_shape, 4, 4), dtype=dtype)
        
    if accum is not None:
        accum_sqrtw, accum_sqrtwx, accum_sqrtwxx, accum_w, accum_wx, accum_wxx, accum_nnz = [a.copy() for a in accum]
    else:
        accum_sqrtw, accum_sqrtwx, accum_sqrtwxx, accum_w, accum_wx, accum_wxx, accum_nnz = \
            np.zeros((*batch_shape, num_points,), dtype=dtype), \
            np.zeros((*batch_shape, num_points, 3), dtype=dtype), \
            np.zeros((*batch_shape, num_points, 3, 3), dtype=dtype), \
            np.zeros((*batch_shape, num_points,), dtype=dtype), \
            np.zeros((*batch_shape, num_points, 3), dtype=dtype), \
            np.zeros((*batch_shape, num_points, 3, 3), dtype=dtype), \
            np.zeros((*batch_shape, num_points,), dtype=dtype)
    
    for i in range(num_frames):
        # Compute weighted statistics
        mean_sqrtwx = accum_sqrtwx / np.maximum(accum_sqrtw, np.finfo(trajectories.dtype).tiny)[..., None]
        wi, yi = weights[i], trajectories[i]
        sqrtwi = np.sqrt(wi)
        w = sqrtwi * accum_sqrtw
        sum_w = np.sum(w, axis=-1) + np.finfo(dtype).tiny
        center_x = np.sum(sqrtwi[..., None] * accum_sqrtwx, axis=-2) / sum_w[..., None]
        center_y = np.sum(w[..., None] * yi, axis=-2) / sum_w[..., None]
        xc = mean_sqrtwx - center_x[..., None, :]
        yc = yi - center_y[..., None, :]
        cov_yx = np.einsum('...i,...ij,...ik->...jk', w, yc, xc) / sum_w[..., None, None]
        if mode == 'affine' or mode == 'similar':
            cov_xx = (np.einsum('...i,...ij,...ik->...jk', w, xc, xc) + np.einsum('...i,...ijk->...jk', sqrtwi, accum_sqrtwxx)) / sum_w[..., None, None]
            cov_yy = np.einsum('...i,...ij,...ik->...jk', w, yc, yc) / sum_w[..., None, None]
        
        # Solve for pose
        if mode == 'rigid':
            _, R, t = umeyama(cov_yx, mean_x=center_x, mean_y=center_y)
            poses[i] = make_affine_matrix(R, t)
        elif mode == 'similar':
            s, R, t = umeyama(cov_yx, cov_xx=cov_xx, mean_x=center_x, mean_y=center_y, niter=niter)
            poses[i] = make_affine_matrix(s * R, t)
        elif mode == 'affine':
            A, t = affine_umeyama(cov_yx, cov_xx, cov_yy, center_x, center_y, lam=lam, niter=niter)
            poses[i] = make_affine_matrix(A, t)

        xi = transform_points(yi, safe_inv(poses[i])[..., None, :, :])

        # Update accum
        old_mean_sqrtwx, old_accum_sqrtw = mean_sqrtwx.copy(), accum_sqrtw.copy()
        accum_sqrtw += sqrtwi
        accum_sqrtwx += sqrtwi[..., None] * xi
        mean_sqrtwx = accum_sqrtwx / np.maximum(accum_sqrtw, np.finfo(dtype).tiny)[..., None]
        accum_sqrtwxx += old_accum_sqrtw[..., None, None] * vector_outer(mean_sqrtwx - old_mean_sqrtwx) + sqrtwi[..., None, None] * vector_outer(xi - mean_sqrtwx)

        mean_wx = accum_wx / np.maximum(accum_w, np.finfo(dtype).tiny)[..., None]
        old_mean_wx, old_accum_w = mean_wx.copy(), accum_w.copy()
        accum_w += wi
        accum_wx += wi[..., None] * xi
        mean_wx = accum_wx / np.maximum(accum_w, np.finfo(dtype).tiny)[..., None]
        accum_wxx += old_accum_w[..., None, None] * vector_outer(mean_wx - old_mean_wx) + wi[..., None, None] * vector_outer(xi - mean_wx)
        accum_nnz += wi > 0

    tot_w = np.sum(accum_w, axis=-1)
    mu = np.sum(accum_wx, axis=-2) / np.maximum(tot_w, np.finfo(dtype).tiny)[..., None]
    mean_wx = accum_wx / np.maximum(accum_w, np.finfo(dtype).tiny)[..., None]
    sigma = np.sum(accum_wxx + accum_w[..., None, None] * vector_outer(mu[..., None, :] - mean_wx), axis=-3) / np.maximum(tot_w, np.finfo(dtype).tiny)[..., None, None]
    nnz = np.sum(accum_nnz, axis=-1)
    valid = np.sum(weights > 0, axis=-1) >= min_valid_size
    err = np.sqrt(np.trace(accum_wxx, axis1=-2, axis2=-1) / np.maximum(accum_nnz, np.finfo(dtype).tiny))

    return poses, valid, (mu, sigma, tot_w, nnz), mean_wx, err, (accum_sqrtw, accum_sqrtwx, accum_sqrtwxx, accum_w, accum_wx, accum_wxx, accum_nnz)


def segment_solve_poses_sequential(
    trajectories: ndarray,
    weights: Optional[ndarray] = None,
    offsets: ndarray = None,
    *,
    accum: Optional[Tuple[ndarray, ...]] = None,
    min_valid_size: int = 3,
    mode: Literal['rigid', 'similar', 'affine'] = 'rigid',
    lam: float = 1e-2,
    niter: int = 8
) -> Tuple[ndarray, Tuple[ndarray, ...], Tuple[ndarray, ndarray, ndarray, ndarray]]:
    """
    Segment array mode for `solve_poses_sequential`.

    Parameters
    ----
    - `trajectories`: (T, N, 3) posed points.
    - `weights`: (T, N) quardratic error term weights for each point at each frame
    - `offsets`: (S + 1,) segment offsets. Points in each segment belong to the same rigid / affine body.
    - `accum`: accumulated statistics from previous calls. If None, start fresh.
    - `min_valid_size`: minimum number of valid points in each frame to consider the segment / group valid.
    - `mode`: mode of transformation to apply. Can be 'rigid', 'similar', or 'affine'.
        - For 'rigid', only rotation and translation are allowed.
        - For 'similar', uniform scaling, rotation and translation are allowed. 
        - For 'affine', full affine transformation is allowed. Using least squares.
    - `lam`: rigidity regularization weight for 'affine' mode.
    - `niter`: number of iterations for 'affine' mode.

    Returns
    ----
    - `poses`: (T, S, 4, 4) transformations from canonical to each frame.
    - `valid`: (T, S) boolean mask indicating valid segments
    - `stats`: canonical statistics of each group,
        It is a tuple of:
        - `mu`: (S, 3) weighted mean of points
        - `cov`: (S, 3, 3) weighted covariance of points
        - `tot_w`: (S,) total weight of points
        - `nnz`: (S,) number of non-zero weight points
    - `canonical_points`: (N, 3) canonical points.
    - `err`: (N,) per-point RMS error over all time := sqrt(sum_over_time(per_point_weights * per_point_squared_error) / per_point_nnz)
        Use this to filter outliers as needed.
    - `accum`: per point accumulated statistics. Just pass it to the next call for incremental solving.
        It is a tuple of:
        - `accum_sqrtw`: (N,) sum of sqrt(weights)
        - `accum_sqrtwx`: (N, 3) sum of sqrt(weights) * x
        - `accum_sqrtwxx`: (N, 3, 3) sum of sqrt(weights) * outer(x - mean_sqrtwx, x - mean_sqrtwx)
        - `accum_w`: (N,) sum of weights
        - `accum_wx`: (N, 3) sum of weights * x
        - `accum_wxx`: (N, 3, 3) sum of weights * outer(x - mean_wx, x - mean_wx)
        - `accum_nnz`: (N,) number of non-zero weight accumulations
    """
    dtype = trajectories.dtype
    num_frames = trajectories.shape[0]
    num_points = trajectories.shape[1]

    if weights is None:
        weights = np.ones((num_frames, num_points), dtype=dtype)

    num_segments = len(offsets) - 1
    lengths = np.diff(offsets)
    poses = np.zeros((num_frames, num_segments, 4, 4), dtype=dtype)
        
    if accum is not None:
        accum_sqrtw, accum_sqrtwx, accum_sqrtwxx, accum_w, accum_wx, accum_wxx, accum_nnz = [a.copy() for a in accum]
    else:
        accum_sqrtw, accum_sqrtwx, accum_sqrtwxx, accum_w, accum_wx, accum_wxx, accum_nnz = \
            np.zeros((num_points,), dtype=dtype), \
            np.zeros((num_points, 3), dtype=dtype), \
            np.zeros((num_points, 3, 3), dtype=dtype), \
            np.zeros((num_points,), dtype=dtype), \
            np.zeros((num_points, 3), dtype=dtype), \
            np.zeros((num_points, 3, 3), dtype=dtype), \
            np.zeros((num_points,), dtype=dtype)
    
    for i in range(num_frames):
        # Compute weighted statistics
        mean_sqrtwx = accum_sqrtwx / np.maximum(accum_sqrtw, np.finfo(trajectories.dtype).tiny)[..., None]
        wi, yi = weights[i], trajectories[i]
        sqrtwi = np.sqrt(wi)
        w = sqrtwi * accum_sqrtw
        sum_w = np.add.reduceat(w, offsets[:-1], axis=0) + np.finfo(dtype).tiny
        center_x = np.add.reduceat(sqrtwi[:, None] * accum_sqrtwx, offsets[:-1], axis=0) / sum_w[:, None]
        center_y = np.add.reduceat(w[:, None] * yi, offsets[:-1], axis=0) / sum_w[:, None]
        center_x_broadcast = np.repeat(center_x, lengths, axis=0)
        center_y_broadcast = np.repeat(center_y, lengths, axis=0)
        xc = mean_sqrtwx - center_x_broadcast
        yc = yi - center_y_broadcast
        cov_yx = np.add.reduceat(w[:, None, None] * vector_outer(yc, xc), offsets[:-1], axis=0) / sum_w[:, None, None]
        if mode == 'affine' or mode == 'similar':
            cov_xx = np.add.reduceat(sqrtwi[:, None, None] * accum_sqrtwxx + w[:, None, None] * vector_outer(xc), offsets[:-1], axis=0) / sum_w[:, None, None]
            cov_yy = np.add.reduceat(w[:, None, None] * vector_outer(yc), offsets[:-1], axis=0) / sum_w[:, None, None]
        
        # Solve for pose
        if mode == 'rigid':
            _, R, t = umeyama(cov_yx, mean_x=center_x, mean_y=center_y)
            poses[i] = make_affine_matrix(R, t)
        elif mode == 'similar':
            s, R, t = umeyama(cov_yx, cov_xx=cov_xx, mean_x=center_x, mean_y=center_y, niter=niter)
            poses[i] = make_affine_matrix(s * R, t)
        elif mode == 'affine':
            A, t = affine_umeyama(cov_yx, cov_xx, cov_yy, center_x, center_y, lam=lam, niter=niter)
            poses[i] = make_affine_matrix(A, t)

        xi = transform_points(yi, np.repeat(safe_inv(poses[i]), lengths, axis=0))

        # Update accum
        old_mean_sqrtwx, old_accum_sqrtw = mean_sqrtwx.copy(), accum_sqrtw.copy()
        accum_sqrtw += sqrtwi
        accum_sqrtwx += sqrtwi[..., None] * xi
        mean_sqrtwx = accum_sqrtwx / np.maximum(accum_sqrtw, np.finfo(dtype).tiny)[..., None]
        accum_sqrtwxx += old_accum_sqrtw[..., None, None] * vector_outer(mean_sqrtwx - old_mean_sqrtwx) + sqrtwi[..., None, None] * vector_outer(xi - mean_sqrtwx)

        mean_wx = accum_wx / np.maximum(accum_w, np.finfo(dtype).tiny)[..., None]
        old_mean_wx, old_accum_w = mean_wx.copy(), accum_w.copy()
        accum_w += wi
        accum_wx += wi[..., None] * xi
        mean_wx = accum_wx / np.maximum(accum_w, np.finfo(dtype).tiny)[..., None]
        accum_wxx += old_accum_w[..., None, None] * vector_outer(mean_wx - old_mean_wx) + wi[..., None, None] * vector_outer(xi - mean_wx)
        accum_nnz += wi > 0

    tot_w = np.add.reduceat(accum_w, offsets[:-1], axis=0)
    mu = np.add.reduceat(accum_wx, offsets[:-1], axis=0) / np.maximum(tot_w, np.finfo(dtype).tiny)[:, None]
    mean_wx = accum_wx / np.maximum(accum_w, np.finfo(dtype).tiny)[:, None]
    mu_broadcast = np.repeat(mu, lengths, axis=0)
    sigma = np.add.reduceat(accum_wxx + accum_w[:, None, None] * vector_outer(mu_broadcast - mean_wx), offsets[:-1], axis=0) / np.maximum(tot_w, np.finfo(dtype).tiny)[:, None, None]
    nnz = np.add.reduceat(accum_nnz, offsets[:-1], axis=0)
    valid = np.add.reduceat(weights > 0, offsets[:-1], axis=1) >= min_valid_size
    err = np.sqrt(np.trace(accum_wxx, axis1=-2, axis2=-1) / np.maximum(accum_nnz, np.finfo(dtype).tiny))

    return poses, valid, (mu, sigma, tot_w, nnz), mean_wx, err, (accum_sqrtw, accum_sqrtwx, accum_sqrtwxx, accum_w, accum_wx, accum_wxx, accum_nnz)