Add stdlib_spatial module with Kabsch–Umeyama vector alignment algorithm

doc/specs/stdlib_spatial.md

@@ -0,0 +1,59 @@
+---
+title: spatial
+---
+
+# The `stdlib_spatial` module
+
+This module provides implementations of algorithms for spatial data processing.
+
+[TOC]
+
+## `kabsch_umeyama` - Finding optimal rotation matrix
+
+### Status
+
+Experimental
+
+### Description
+
+Compute the optimal similarity transform (Kabsch–Umeyama):
+\[
+    P \approx c \, R \, Q + t
+\]
+where:
+
+- R is an orthogonal rotation matrix,
+- c is an optional scaling factor,
+- t is a translation vector.
+
+The transformation minimizes the root-mean-square deviation (RMSD) between corresponding columns
+of P and Q, optionally using weights and with optional scaling.
+The implementation is based on the algorithm described here: [Aligning point patterns with Kabsch–Umeyama algorithm](https://zpl.fi/aligning-point-patterns-with-kabsch-umeyama-algorithm)
+
+### Syntax
+
+`call ` [[stdlib_spatial(module):kabsch_umeyama(interface)]] `(P, Q, R, t, c, rmsd [, W, scale])`
+
+### Arguments
+
+`P`: Shall be a `real` or `complex` rank-2 array. It is an `intent(in)` argument.
+
+`Q`: Shall be a rank-2 array with same kind as `P`. It is an `intent(in)` argument.
+
+`R`: Shall be a rank-2 array with same kind as `P`. For `real` kinds, the algorithm returns a proper rotation matrix, meaning `det(R) = 1`. It is an `intent(out)` argument.
+
+`t`: Shall be a rank-1 array with same kind as `P`. It is an `intent(out)` argument.
+
+`c`: Scalar value of the same type as `P`. It is an `intent(out)` argument. If `scale` is disabled `c` will be returned with a value of `1`.
+
+`rmsd`: Scalar value of `real` kind. It is an `intent(out)` argument.
+
+`W` (optional): Shall be a rank-1 array of `real` kind. It is an `intent(in)` argument. By default, `W` is an array of `1`s.
+
+`scale` (optional): Shall be a logical type. It is an `intent(in)` argument. By default, `scale = .true.`.
+
+### Example
+
+```fortran
+{!example/spatial/example_kabsch_umeyama.f90!}
+```

example/CMakeLists.txt

@@ -46,6 +46,7 @@ if (STDLIB_QUADRATURE)
 endif()
 add_subdirectory(selection)
 add_subdirectory(sorting)
+add_subdirectory(spatial)
 add_subdirectory(specialfunctions_gamma)
 if (STDLIB_SPECIALMATRICES)
   add_subdirectory(specialmatrices)

example/spatial/CMakeLists.txt

@@ -0,0 +1 @@
+ADD_EXAMPLE(kabsch_umeyama)

example/spatial/example_kabsch_umeyama.f90

@@ -0,0 +1,36 @@
+program example_kabsch_umeyama
+    use stdlib_linalg_constants, only: dp
+    use stdlib_spatial, only: kabsch_umeyama
+    implicit none
+
+    integer, parameter :: d = 2, N = 7
+    real(dp) :: P(d,N), Q(d,N), R(d, d), t(d), c, rmsd
+
+    integer :: i
+
+    ! 2x7 matrices.
+    P(1,:) = [23.0_dp, 66.0_dp, 88.0_dp, 119.0_dp, 122.0_dp, 170.0_dp, 179.0_dp]
+    P(2,:) = [178.0_dp, 173.0_dp, 187.0_dp, 202.0_dp, 229.0_dp, 232.0_dp, 199.0_dp]
+
+    Q(1,:) = [232.0_dp, 208.0_dp, 181.0_dp, 155.0_dp, 142.0_dp, 121.0_dp, 139.0_dp]
+    Q(2,:) = [38.0_dp, 32.0_dp, 31.0_dp, 45.0_dp, 33.0_dp, 59.0_dp, 69.0_dp]
+
+    call kabsch_umeyama(P, Q, R, t, c, rmsd)
+
+    print *
+    print *, "Recovered rotation R:"
+    do i = 1, d
+        print *, R(i,:)
+    end do
+
+    print *
+    print *, "Recovered scale c:", c
+
+    print *
+    print *, "Recovered translation t:"
+    print *, t
+
+    print *
+    print *, "RMSD:", rmsd
+
+end program example_kabsch_umeyama

src/CMakeLists.txt

@@ -46,6 +46,7 @@ ADD_SUBDIR(system)
 ADD_SUBDIR(stats)
 
 add_subdirectory(sparse)
+add_subdirectory(spatial)
 
 set(fppFiles
     stdlib_version.fypp
@@ -73,5 +74,6 @@ target_link_libraries(${PROJECT_NAME} PUBLIC
   ${PROJECT_NAME}_strings
   ${PROJECT_NAME}_blas ${PROJECT_NAME}_lapack ${PROJECT_NAME}_lapack_extended
   ${PROJECT_NAME}_sparse
+  ${PROJECT_NAME}_spatial
   ${OPTIONAL_LIB}
   )

src/spatial/CMakeLists.txt

@@ -0,0 +1,14 @@
+set(spatial_fppFiles
+    stdlib_spatial.fypp
+    stdlib_spatial_kabsch_umeyama.fypp
+    )
+
+set(spatial_cppFiles
+    )
+
+set(spatial_f90Files
+    )
+
+configure_stdlib_target(${PROJECT_NAME}_spatial spatial_f90Files spatial_fppFiles spatial_cppFiles)
+
+target_link_libraries(${PROJECT_NAME}_spatial PUBLIC ${PROJECT_NAME}_core ${PROJECT_NAME}_constants ${PROJECT_NAME}_lapack ${PROJECT_NAME}_linalg ${PROJECT_NAME}_intrinsics)

src/spatial/stdlib_spatial.fypp

@@ -0,0 +1,41 @@
+#:include "common.fypp"
+#:set R_KINDS_TYPES = list(zip(REAL_KINDS, REAL_TYPES, REAL_SUFFIX))
+#:set C_KINDS_TYPES = list(zip(CMPLX_KINDS, CMPLX_TYPES, CMPLX_SUFFIX))
+#:set KINDS_TYPES = R_KINDS_TYPES+C_KINDS_TYPES
+module stdlib_spatial
+    use stdlib_kinds, only: sp, dp, xdp, qp
+    use stdlib_constants
+    implicit none
+    private
+    public :: kabsch_umeyama
+
+    interface kabsch_umeyama
+        !! ([Specifications](../page/specs/stdlib_spatial.html#kabsch_umeyama))
+        !! This interface computes the optimal similarity transform (Kabsch–Umeyama):
+        !! \[
+        !!      P \approx c \, R \, Q + t
+        !! \]
+        !! The transformation minimizes the RMSD between corresponding columns
+        !! of P and Q, optionally using weights and with optional scaling.
+        #:for k, t, s in (KINDS_TYPES)
+        module subroutine kabsch_umeyama_${s}$(P, Q, R, t, c, rmsd, W, scale)
+            ${t}$, intent(in) :: P(:, :)
+            !! Target point set (d × N)
+            ${t}$, intent(in) :: Q(:, :)
+            !! Reference point set (d × N)
+            ${t}$, intent(out) :: R(:, :)
+            !! Optimal rotation matrix (d × d)
+            ${t}$, intent(out) :: t(:)
+            !! Translation vector (d)
+            ${t}$, intent(out) :: c
+            !! Scale factor
+            real(${k}$), intent(out) :: rmsd
+            !! Root-mean-square deviation
+            real(${k}$), intent(in), optional :: W(:)
+            !! Optional weights
+            logical, intent(in), optional :: scale
+            !! Enable scaling
+        end subroutine
+        #:endfor
+    end interface
+end module stdlib_spatial

src/spatial/stdlib_spatial_kabsch_umeyama.fypp

@@ -0,0 +1,194 @@
+#:include "common.fypp"
+#:set R_KINDS_TYPES = list(zip(REAL_KINDS, REAL_TYPES, REAL_SUFFIX))
+#:set C_KINDS_TYPES = list(zip(CMPLX_KINDS, CMPLX_TYPES, CMPLX_SUFFIX))
+#:set KINDS_TYPES = R_KINDS_TYPES+C_KINDS_TYPES
+submodule(stdlib_spatial) stdlib_spatial_kabsch_umeyama
+    use stdlib_linalg, only: svd, det
+    use stdlib_intrinsics, only: stdlib_sum_kahan, stdlib_dot_product_kahan, kahan_kernel
+    use stdlib_error, only: error_stop
+    use stdlib_optval, only: optval
+    use stdlib_linalg_lapack, only: gemm, gemv
+
+contains
+    #:for k, t, s in (KINDS_TYPES)
+    module subroutine kabsch_umeyama_${s}$(P, Q, R, t, c, rmsd, W, scale)
+        ${t}$, intent(in) :: P(:, :)
+        !! Target point set (d × N)
+        ${t}$, intent(in) :: Q(:, :)
+        !! Reference point set (d × N)
+        ${t}$, intent(out) :: R(:, :)
+        !! Optimal rotation matrix (d × d)
+        ${t}$, intent(out) :: t(:)
+        !! Translation vector (d)
+        ${t}$, intent(out) :: c
+        !! Scale factor
+        real(${k}$), intent(out) :: rmsd
+        !! Root-mean-square deviation
+        real(${k}$), intent(in), optional :: W(:)
+        !! Optional weights
+        logical, intent(in), optional :: scale
+        !! Enable scaling
+
+        ! Internal variables.
+        integer :: i, j, point, d, N
+        ${t}$, allocatable :: covariance(:,:), U(:,:), Vt(:,:), vec(:), tmp_N(:), tmp_d(:), c_P(:), c_Q(:)
+        real(${k}$) :: sum_w, variance_p
+        real(${k}$), allocatable :: S(:)
+        ${t}$ :: temp
+        #:if t.startswith('complex')
+        real(${k}$) :: rtemp
+        #:endif
+        logical :: scale_
+        #:if t.startswith('real')
+        logical :: reflect_
+        #:endif
+        real(${k}$) :: rmsd_err
+
+        scale_ = optval(scale, .true.)
+        ! Dimension checks
+        if(size(P,dim=1)/=size(Q,dim=1) .or. size(P,dim=1)/=size(R,dim=1) .or. size(P,dim=1)/=size(R,dim=2) &
+                    .or. size(P,dim=1)/=size(t)) then
+            call error_stop("array sizes do not match")
+        end if
+        if(size(P,dim=2)/=size(Q,dim=2)) then
+            call error_stop("array sizes do not match")
+        end if
+        if (present(W)) then
+            if (size(W) /= size(P,dim=2)) then
+                call error_stop("array sizes do not match")
+            end if
+        end if
+        d = size(P,dim=1)
+        N = size(P,dim=2)
+
+        if(present(W)) then
+            sum_w = stdlib_sum_kahan(W)
+        else
+            sum_w = real(N, kind = ${k}$)
+        end if
+        !> leave opportunity for future discussion on how to add spmd reduction needed here to reduce sum_w before division
+
+        sum_w = one_${k}$ / sum_w
+        if(sum_w<zero_${k}$) call error_stop("Invalid weights: sum of weights must be positive")
+
+        allocate(c_P(d), source=zero_${s}$)
+        allocate(c_Q(d), source=zero_${s}$)
+        allocate(tmp_N(N), source=zero_${s}$)
+
+        ! Compute centroids of P and Q
+        if(present(W)) then
+            do i = 1, d
+                tmp_N(:) = W(:) * P(i,:)
+                c_P(i) = stdlib_sum_kahan(tmp_N)
+                tmp_N(:) = W(:) * Q(i,:)
+                c_Q(i) = stdlib_sum_kahan(tmp_N)
+            end do
+        else
+            c_P = stdlib_sum_kahan(P, dim=2)
+            c_Q = stdlib_sum_kahan(Q, dim=2)
+        end if
+        c_P = c_P * sum_w
+        c_Q = c_Q * sum_w
+
+        ! Compute covariance matrix H = (P - c_P) * (Q - c_Q)^T and variance of P
+        allocate(covariance(d,d), source=zero_${s}$)
+        allocate(tmp_d(d), source=zero_${s}$)
+        variance_p = zero_${k}$
+
+        if (present(W)) then
+            do point = 1, N
+                tmp_d = P(:, point) - c_P(:)
+                tmp_N(point) = stdlib_dot_product_kahan(tmp_d, tmp_d)
+            end do
+            tmp_N(:) = W(:) * tmp_N(:)
+            variance_p = stdlib_sum_kahan(tmp_N)
+            do j = 1, d
+                do i = 1, d
+                    #:if t.startswith('complex')
+                    tmp_N(:) = W(:) * (P(i,:) - c_P(i)) * conjg(Q(j,:) - c_Q(j))
+                    covariance(i,j) = stdlib_sum_kahan(tmp_N)
+                    #:else
+                    tmp_N(:) = W(:) * (P(i,:) - c_P(i)) * (Q(j,:) - c_Q(j))
+                    covariance(i,j) = stdlib_sum_kahan(tmp_N)
+                    #:endif
+                end do
+            end do
+        else
+            ! Calculate variance by the formula (1/n)*sigma(P - c_P)^2
+            do point = 1, N
+                tmp_d = P(:, point) - c_P(:)
+                tmp_N(point) = stdlib_dot_product_kahan(tmp_d, tmp_d)
+            end do
+            variance_p = stdlib_sum_kahan(tmp_N)
+            do j = 1, d
+                do i = 1, d
+                    #:if t.startswith('complex')
+                    tmp_N(:) = (P(i,:) - c_P(i)) * conjg(Q(j,:) - c_Q(j))
+                    covariance(i,j) = stdlib_sum_kahan(tmp_N)
+                    #:else
+                    tmp_N(:) = (P(i,:) - c_P(i)) * (Q(j,:) - c_Q(j))
+                    covariance(i,j) = stdlib_sum_kahan(tmp_N)
+                    #:endif
+                end do
+            end do
+        end if
+
+        covariance = covariance * sum_w
+        variance_p = variance_p * sum_w
+
+        allocate(U(d,d), source=zero_${s}$)
+        allocate(Vt(d,d), source=zero_${s}$)
+        allocate(S(d), source=zero_${k}$)
+
+        ! SVD of covariance matrix H -> H = U * S * Vt
+        call svd(covariance, S, U, Vt)
+
+        ! Check for reflections in case of real entries.
+        #:if t.startswith('real')
+        reflect_ = det(matmul(U,Vt)) < zero_${s}$
+        if(reflect_) Vt(d,:) = -Vt(d,:)
+        #:endif
+
+        ! Optimal rotation matrix.
+        call gemm(transa='N', transb='N', m=d,n=d,k=d, alpha=one_${s}$, a=U,lda=d, b=Vt, ldb=d, beta=zero_${s}$, c=R, ldc=d)
+
+        ! Scaling factor
+        c = one_${s}$
+        if(scale_) then
+            #:if t.startswith('real')
+            if(reflect_) then
+                c = sum(S(1:d-1)) - S(d)
+            else
+                c = sum(S(1:d))
+            end if
+            #:else
+            c = sum(S(1:d))
+            #:endif
+            c = variance_p / c
+        end if
+
+        ! Translation vector t = c_P - c*R*c_Q
+        t = c_P
+        call gemv(trans='N', m=d, n=d, alpha=-c, A=R, lda=d, x=c_Q, incx=1, beta=one_${s}$, y=t, incy=1)
+
+        ! Compute RMSD
+        allocate(vec(d), source=zero_${s}$)
+        rmsd = zero_${k}$
+        rmsd_err = zero_${k}$
+        do point = 1, N
+            ! Calculate the k^th difference vector by the formula vec_k = c*R*Q_k + t - P_k
+            vec = t
+            call gemv(trans='N', m=d, n=d, alpha=c, A=R, lda=d, x=Q(:, point), incx=1, beta=one_${s}$, y=vec, incy=1)
+            vec = vec - P(:,point)
+            temp = stdlib_dot_product_kahan(vec,vec)
+            #:if t.startswith('complex')
+            rtemp = real(temp, kind=${k}$)
+            call kahan_kernel(rtemp, rmsd, rmsd_err)
+            #:else
+            call kahan_kernel(temp, rmsd, rmsd_err)
+            #:endif
+        end do
+        rmsd = sqrt(rmsd * sum_w)
+    end subroutine
+    #:endfor
+end submodule stdlib_spatial_kabsch_umeyama

test/CMakeLists.txt

@@ -50,6 +50,7 @@ endif()
 add_subdirectory(optval)
 add_subdirectory(selection)
 add_subdirectory(sorting)
+add_subdirectory(spatial)
 add_subdirectory(specialfunctions)
 if (STDLIB_STATS)
   add_subdirectory(stats)

test/spatial/CMakeLists.txt

@@ -0,0 +1,7 @@
+set(
+  fppFiles
+  "test_spatial_kabsch_umeyama.fypp"
+)
+fypp_f90("${fyppFlags}" "${fppFiles}" outFiles)
+
+ADDTEST(spatial_kabsch_umeyama)