Continue optimizing the previous code

Here are some possible ways to optimize the previous code:

1. Vectorize the calculations: Instead of using nested loops to compute the responsibility matrix, we can use vectorized operations to speed up the computation. For example, we can use broadcasting to compute the Euclidean distance between each pair of points in a matrix form. Similarly, we can use matrix multiplication to compute the weighted sums of the point clouds.

def em_for_alignment(xs: np.ndarray, ys: np.ndarray, num_iter: int = 10) -> Tuple[np.ndarray, np.ndarray]:
    """
    The em algorithm for aligning two point clouds based on affine transformation
    :param xs: a set of points with size (N, D), N is the number of samples, D is the dimension of points
    :param ys: a set of points with size (M, D), M is the number of samples, D is the dimension of points
    :param num_iter: the number of EM iterations
    :return:
        ys_new: the aligned points: ys_new = ys @ affine + translation
        responsibility: the responsibility matrix P=[p(y_m | x_n)] with size (N, M),
        whose elements indicating the correspondence between the points
    """
    # initialize the affine matrix and translation vector
    affine = np.eye(xs.shape[1])
    translation = np.zeros(xs.shape[1])

    # initialize the responsibility matrix
    responsibility = np.zeros((xs.shape[0], ys.shape[0]))

    for i in range(num_iter):
        # E-step: compute the responsibility matrix
        diff = xs[:, np.newaxis, :] - ys[np.newaxis, :, :]
        sq_dist = np.sum(diff ** 2, axis=-1)
        responsibility = np.exp(-0.5 * sq_dist) / (2 * np.pi) ** (xs.shape[1] / 2)
        responsibility /= np.sum(responsibility, axis=1, keepdims=True)

        # M-step: update the affine matrix and translation vector
        xs_weighted = responsibility.T @ xs
        ys_weighted = responsibility.T @ ys
        affine, _, _, _ = np.linalg.lstsq(xs_weighted, ys_weighted, rcond=None)
        translation = np.mean(ys, axis=0) - np.mean(xs @ affine, axis=0)

    # compute the aligned points
    ys_new = ys @ affine + translation

    return ys_new, responsibility

2. Use the Kabsch algorithm: Instead of using the weighted least squares solution to update the affine matrix, we can use the Kabsch algorithm, which is a more efficient and numerically stable method for finding the optimal rigid transformation between two point clouds. The Kabsch algorithm consists of three steps: centering the point clouds, computing the covariance matrix, and finding the optimal rotation matrix.

def em_for_alignment(xs: np.ndarray, ys: np.ndarray, num_iter: int = 10) -> Tuple[np.ndarray, np.ndarray]:
    """
    The em algorithm for aligning two point clouds based on affine transformation
    :param xs: a set of points with size (N, D), N is the number of samples, D is the dimension of points
    :param ys: a set of points with size (M, D), M is the number of samples, D is the dimension of points
    :param num_iter: the number of EM iterations
    :return:
        ys_new: the aligned points: ys_new = ys @ affine + translation
        responsibility: the responsibility matrix P=[p(y_m | x_n)] with size (N, M),
        whose elements indicating the correspondence between the points
    """
    # center the point clouds
    xs_centered = xs - np.mean(xs, axis=0)
    ys_centered = ys - np.mean(ys, axis=0)

    # initialize the affine matrix and translation vector
    affine = np.eye(xs.shape[1])
    translation = np.zeros(xs.shape[1])

    # initialize the responsibility matrix
    responsibility = np.zeros((xs.shape[0], ys.shape[0]))

    for i in range(num_iter):
        # E-step: compute the responsibility matrix
        diff = xs_centered[:, np.newaxis, :] - ys_centered[np.newaxis, :, :]
        sq_dist = np.sum(diff ** 2, axis=-1)
        responsibility = np.exp(-0.5 * sq_dist) / (2 * np.pi) ** (xs.shape[1] / 2)
        responsibility /= np.sum(responsibility, axis=1, keepdims=True)

        # M-step: update the affine matrix and translation vector
        cov = xs_centered.T @ responsibility @ ys_centered
        u, _, vh = np.linalg.svd(cov)
        r = vh.T @ u.T
        t = np.mean(ys, axis=0) - np.mean(xs @ r, axis=0)
        affine = np.hstack((r, t[:, np.newaxis]))

    # compute the aligned points
    ys_new = ys @ affine[:, :-1] + affine[:, -1]

    return ys_new, responsibility

The Kabsch algorithm is more efficient than the weighted least squares solution, especially when the point clouds are high-dimensional or noisy. However, it only works for rigid transformations, i.e., rotations and translations. If the transformation between the point clouds is not rigid, we need to use a more general method, such as the Procrustes analysis or the Iterative Closest Point (ICP) algorithm.