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)