-
Notifications
You must be signed in to change notification settings - Fork 225
Add stdlib_spatial module with Kabsch–Umeyama vector alignment algorithm #1119
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
ea6b295
655af91
60818bb
400b85d
5208695
25df8d2
33859b6
c7b3227
78f35a4
a9c489c
fb08933
88e76fb
e73a93f
e3a1cc2
27b68ca
c895ed5
d722366
b9b9832
37b3cba
723661b
3a7221f
51a193d
7554914
12f1513
fb4f1a2
fe7d96e
81a228c
a9f0fc1
42db95b
fa159bd
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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!} | ||
| ``` |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| ADD_EXAMPLE(kabsch_umeyama) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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}_linalg ${PROJECT_NAME}_intrinsics) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,201 @@ | ||
| #: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 | ||
|
|
||
| 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(:) | ||
| logical :: scale_ | ||
| #:if t.startswith('real') | ||
| logical :: reflect_ | ||
| #:endif | ||
| real(${k}$) :: rmsd_err | ||
|
|
||
|
|
||
| ! 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) | ||
| scale_ = .true. | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| if(present(scale)) scale_ = scale | ||
|
|
||
| sum_w = one_${k}$ / N | ||
| if(present(W)) sum_w = one_${k}$ / stdlib_sum_kahan(W) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. it would be better to compute 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 | ||
| do i = 1, d | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| c_P(i) = stdlib_sum_kahan(P(i, :)) | ||
| c_Q(i) = stdlib_sum_kahan(Q(i, :)) | ||
| end do | ||
| 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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The loops in this section should be revisited. This algorithm is usually applied for N >> d, this means that the best approach in terms of memory access is to handle, for each point in |
||
| 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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
| 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}$ | ||
| #:endif | ||
|
|
||
| #:if t.startswith('real') | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. this fypp macro can be merged in the previous #:if |
||
| if(reflect_) Vt(d,:) = -Vt(d,:) | ||
| #:endif | ||
|
|
||
| ! Optimal rotation matrix. | ||
| do i = 1,d | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same comment as before. I'm worried that using |
||
| do j = 1,d | ||
| #:if t.startswith('complex') | ||
| R(i,j) = stdlib_dot_product_kahan(conjg(U(i,:)), Vt(:, j)) | ||
| #:else | ||
| R(i,j) = stdlib_dot_product_kahan(U(i,:), Vt(:, j)) | ||
| #:endif | ||
| end do | ||
| end do | ||
|
Comment on lines
+135
to
+156
|
||
|
|
||
| ! Scaling factor | ||
| if(scale_) then | ||
| #:if t.startswith('real') | ||
| if(reflect_) then | ||
| c = variance_p / (sum(S(1:d-1)) - S(d)) | ||
| else | ||
| c = variance_p / (sum(S(1:d))) | ||
| end if | ||
| #:else | ||
| c = variance_p / (sum(S(1:d))) | ||
| #:endif | ||
| else | ||
| c = one_${s}$ | ||
| end if | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. a small suggestion to make this section easier to follow (includding aligning the fypp macro with the inner scope. ! 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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This operation might be better of being writtent using |
||
| do i = 1, d | ||
| #:if t.startswith('complex') | ||
| t(i) = c_P(i) - c * stdlib_dot_product_kahan(conjg(R(i,1:d)), c_Q(1:d)) | ||
| #:else | ||
| t(i) = c_P(i) - c * stdlib_dot_product_kahan(R(i,1:d), c_Q(1:d)) | ||
| #:endif | ||
| end do | ||
|
|
||
| ! 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 | ||
| do i = 1, d | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same comment as before, |
||
| #:if t.startswith('complex') | ||
| vec(i) = c * stdlib_dot_product_kahan(conjg(R(i,1:d)), Q(1:d,point)) | ||
| #:else | ||
| vec(i) = c * stdlib_dot_product_kahan(R(i,1:d), Q(1:d,point)) | ||
| #:endif | ||
| end do | ||
| vec(1:d) = vec(1:d) + t(1:d) - P(1:d,point) | ||
| call kahan_kernel(real(stdlib_dot_product_kahan(vec,vec), kind=${k}$), rmsd, rmsd_err) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It is preferable to avoid passing expressions as arguments to procedures. This induces the creation of temporal variables. Better to declare the required variables, assign the expression and then pass the argument. |
||
| end do | ||
| rmsd = sqrt(rmsd * sum_w) | ||
| end subroutine | ||
| #:endfor | ||
| end submodule stdlib_spatial_kabsch_umeyama | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Wis declared as the same type as the point arrays (${t}$). For complex variants this implies complex weights, which then flow intosum_w/centroid/covariance computations viastdlib_sum_kahan/stdlib_dot_product_kahanand require implicit complex→real conversions. If weights are intended (as in the issue description) to be real and non-negative, consider typingWasreal(${k}$)for both real and complex point sets and validatesum(W) > 0.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Updated.