Get best alignment of 2D coords over a reference set of points, allowing for a reflection. Useful for orienting dimensionality reduction output in 2D (UMAP, t-SNE etc.)

kabsch_best.R

kabsch_best <- function(coords, ref) {
  kabsch_res1 <- kabsch(coords, ref, ret_error = TRUE)
  kabsch_res2 <- kabsch(coords |> flipx(), ref, ret_error = TRUE)
  if (kabsch_res1$error < kabsch_res2$error) {
    coords <- kabsch_res1$coords
  } else {
    coords <- kabsch_res2$coords
  }
  coords
}

# pm is moved over qm
kabsch <- function(pm, qm, ret_error = FALSE) {
  pm_dims <- dim(pm)
  if (!all(dim(qm) == pm_dims)) {
    stop(call. = TRUE, "Point sets must have the same dimensions")
  }
  # The rotation matrix will have (ncol - 1) leading ones in the diagonal
  diag_ones <- rep(1, pm_dims[2] - 1)

  # center the points
  pm <- scale(pm, center = TRUE, scale = FALSE)
  qm <- scale(qm, center = TRUE, scale = FALSE)

  am <- crossprod(pm, qm)

  svd_res <- svd(am)
  # use the sign of the determinant to ensure a right-hand coordinate system
  d <- determinant(tcrossprod(svd_res$v, svd_res$u))$sign
  dm <- diag(c(diag_ones, d))

  # rotation matrix
  um <- svd_res$v %*% tcrossprod(dm, svd_res$u)

  # Rotate and then translate to the original centroid location of qm
  aligned_pm <- sweep(t(tcrossprod(um, pm)), 2, -attr(qm, "scaled:center"))

  if (ret_error) {
    error <- (aligned_pm - qm)^2 |> rowSums() |> mean() |> sqrt()
    return(list(coords = aligned_pm, error = error))
  }
  aligned_pm
}

flipx <- function(x) {
  x[, 1] <- -x[, 1]
  x
}