From f573df680f1c05a69c30c362a34d03909ba3cef7 Mon Sep 17 00:00:00 2001
From: Elia Cereda <eliacereda@gmail.com>
Date: Wed, 21 Feb 2024 18:41:00 +0100
Subject: [PATCH] Fix affines.decompose44() behavior in presence of NANs

---
 transforms3d/affines.py | 30 ++++++++++++------------------
 transforms3d/utils.py   | 24 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/transforms3d/affines.py b/transforms3d/affines.py
index a93d375..380f4c2 100644
--- a/transforms3d/affines.py
+++ b/transforms3d/affines.py
@@ -5,7 +5,7 @@
 import numpy as np
 
 from .shears import striu2mat
-from .utils import diag_2d, transpose_2d, unpack
+from .utils import diag_2d, transpose_2d, unpack, vdot_batch
 
 
 def decompose44(A44):
@@ -49,18 +49,12 @@ def decompose44(A44):
 
     Examples
     --------
-    >>> T = [20, 30, 40] # translations
-    >>> R = [[0, -1, 0], [1, 0, 0], [0, 0, 1]] # rotation matrix
-    >>> Z = [2.0, 3.0, 4.0] # zooms
-    >>> S = [0.2, 0.1, 0.3] # shears
+    >>> T = [[20, 30, 40], [50, 60, 70]]
+    >>> R = [[[0, -1, 0], [1, 0, 0], [0, 0, 1]],  [[1,  0, 0], [0, 0, -1], [0, 1, 0]]]
+    >>> Z = [[2.0, 3.0, 4.0], [-1.0, 2.0, 4.0]]
+    >>> S = None # FIXME: Not implemented [[0.2, 0.1, 0.3], [0.4, 0.5, 0.6]]
     >>> # Now we make an affine matrix
-    >>> A = np.eye(4)
-    >>> Smat = np.array([[1, S[0], S[1]],
-    ...                  [0,    1, S[2]],
-    ...                  [0,    0,    1]])
-    >>> RZS = np.dot(R, np.dot(np.diag(Z), Smat))
-    >>> A[:3,:3] = RZS
-    >>> A[:-1,-1] = T # set translations
+    >>> A = compose(T, R, Z, S)
     >>> Tdash, Rdash, Zdash, Sdash = decompose44(A)
     >>> np.allclose(T, Tdash)
     True
@@ -68,7 +62,7 @@ def decompose44(A44):
     True
     >>> np.allclose(Z, Zdash)
     True
-    >>> np.allclose(S, Sdash)
+    >>> # FIXME: Not implemented np.allclose(S, Sdash)
     True
 
     Notes
@@ -134,8 +128,8 @@ def decompose44(A44):
     M0 /= sx[..., np.newaxis]
 
     # orthogonalize M1 with respect to M0
-    sx_sxy = np.vdot(M0, M1)
-    M1 -= sx_sxy * M0
+    sx_sxy = vdot_batch(M0, M1)
+    M1 -= sx_sxy[..., np.newaxis] * M0
 
     # extract y scale and normalize
     sy = np.sqrt(np.sum(M1**2, axis=-1))
@@ -143,9 +137,9 @@ def decompose44(A44):
     sxy = sx_sxy / sx
 
     # orthogonalize M2 with respect to M0 and M1
-    sx_sxz = np.vdot(M0, M2)
-    sy_syz = np.vdot(M1, M2)
-    M2 -= (sx_sxz * M0 + sy_syz * M1)
+    sx_sxz = vdot_batch(M0, M2)
+    sy_syz = vdot_batch(M1, M2)
+    M2 -= (sx_sxz[..., np.newaxis] * M0 + sy_syz[..., np.newaxis] * M1)
 
     # extract z scale and normalize
     sz = np.sqrt(np.sum(M2**2, axis=-1))
diff --git a/transforms3d/utils.py b/transforms3d/utils.py
index e91e6ac..7300938 100644
--- a/transforms3d/utils.py
+++ b/transforms3d/utils.py
@@ -237,3 +237,27 @@ def transpose_2d(array):
     (128, 3, 2)
     '''
     return np.transpose(array, (*range(array.ndim - 2), -1, -2))
+
+def vdot_batch(a, b):
+    ''' Compute the vector dot product on batches of 1D arrays
+
+    Parameters
+    ----------
+    a, b : array-like batches of 1D arrays with shape (..., N)
+
+    Returns
+    -------
+    out : np.ndarray with shape (...)
+       sum-product of the last dimension from arrays a, b
+
+    Examples
+    --------
+    >>> a = np.array([[0, 1, 2], [2, 3, 4]])
+    >>> b = np.array([[2, 2, 2], [3, 3, 3]])
+    >>> c = vdot_batch(a, b)
+    >>> c.shape
+    (2,)
+    >>> c
+    array([ 6, 27])
+    '''
+    return np.einsum('...i,...i', a, b)