python-microscopy · zacsimile · Feb 14, 2025 · Feb 14, 2025 · Feb 16, 2025 · Feb 17, 2025
diff --git a/PYME/Analysis/points/coordinate_tools.py b/PYME/Analysis/points/coordinate_tools.py
@@ -276,3 +276,109 @@ def distance_to_image_mask(mask, points):
 
     return distances
 
+def unit_quaternion_to_rotation_matrix(q):
+    """
+    Convert a unit quaternion to a rotation matrix. Helper function for 
+    absolute_orientation(). Section 3E of reference.
+
+    Parameters
+    ----------
+    q : np.typing.ArrayLike
+        (4,) vector representing the unit quaternion q0 + iqx + jqy + kqz
+
+    References
+    ----------
+    Berthold K. P. Horn, "Closed-form solution of absolute orientation using unit 
+    quaternions," J. Opt. Soc. Am. A 4, 629-642 (1987) 
+
+    """
+    q0, qx, qy, qz = q
+    q02, qx2, qy2, qz2 = q*q
+
+    return np.array([[q02+qx2-qy2-qz2, 2*(qx*qy-q0*qz), 2*(qx*qz+q0*qy)],
+                    [2*(qy*qx+q0*qz), q02-qx2+qy2-qz2, 2*(qy*qz-q0*qx)],
+                    [2*(qz*qx-q0*qy), 2*(qz*qy+q0*qz), q02-qx2-qy2+qz2]])
+
+def absolute_orientation(reference, target, weights_reference=None, weights_target=None):
+    """
+    Compute the rotation and shift from target to reference point cloud.
+
+    Parameters
+    ----------
+    reference : np.typing.ArrayLike
+        (3, N) matrix of points
+    target : np.typing.ArrayLike
+        (3, N) matrix of points
+    weights_reference : np.typing.ArrayLike
+        (3, N) matrix of weights. Should sum to 1 along axis 1.
+    weights_target : np.typing.ArrayLike
+        (3, N) matrix of weights. Should sum to 1 along axis 1.
+
+    References
+    ----------
+    Berthold K. P. Horn, "Closed-form solution of absolute orientation using unit 
+    quaternions," J. Opt. Soc. Am. A 4, 629-642 (1987) 
+
+    """
+    # Map target coordinates to reference.
+    # Note we do not use any scaling
+
+    if weights_reference is None:
+        weights_reference = np.ones(reference.shape, dtype=target.dtype)
+    if weights_target is None:
+        weights_target = np.ones(target.shape, dtype=target.dtype)
+
+    # Move to operating w.r.t centroid
+    reference_cent = (weights_reference*reference).sum(1)/weights_reference.sum(1)
+    target_cent = (weights_target*target).sum(1)/weights_target.sum(1)
+    r = reference - reference_cent[:,None]
+    t = target - target_cent[:,None]
+
+    # Compute pairwise sums
+    w00 = 1/np.sqrt((1/weights_target[0,:])**2+(1/weights_reference[0,:])**2)
+    S_txrx = (w00*t[0,:]*r[0,:]).sum()
+    w01 = 1/np.sqrt((1/weights_target[0,:])**2+(1/weights_reference[1,:])**2)
+    S_txry = (w01*t[0,:]*r[1,:]).sum()
+    w02 = 1/np.sqrt((1/weights_target[0,:])**2+(1/weights_reference[2,:])**2)
+    S_txrz = (w02*t[0,:]*r[2,:]).sum()
+    w10 = 1/np.sqrt((1/weights_target[1,:])**2+(1/weights_reference[0,:])**2)
+    S_tyrx = (w10*t[1,:]*r[0,:]).sum()
+    w11 = 1/np.sqrt((1/weights_target[1,:])**2+(1/weights_reference[1,:])**2)
+    S_tyry = (w11*t[1,:]*r[1,:]).sum()
+    w12 = 1/np.sqrt((1/weights_target[1,:])**2+(1/weights_reference[2,:])**2)
+    S_tyrz = (w12*t[1,:]*r[2,:]).sum()
+    w20 = 1/np.sqrt((1/weights_target[2,:])**2+(1/weights_reference[0,:])**2)
+    S_tzrx = (w20*t[2,:]*r[0,:]).sum()
+    w21 = 1/np.sqrt((1/weights_target[2,:])**2+(1/weights_reference[1,:])**2)
+    S_tzry = (w21*t[2,:]*r[1,:]).sum()
+    w22 = 1/np.sqrt((1/weights_target[2,:])**2+(1/weights_reference[2,:])**2)
+    S_tzrz = (w22*t[2,:]*r[2,:]).sum()
+
+    # Compute quaternion matrix
+    N = np.array([[(S_txrx+S_tyry+S_tzrz), S_tyrz-S_tzry, S_tzrx-S_txrz, S_txry-S_tyrx],
+                [S_tyrz-S_tzry, (S_txrx-S_tyry-S_tzrz), S_txry+S_tyrx, S_tzrx+S_txrz],
+                [S_tzrx-S_txrz, S_txry+S_tyrx, (-S_txrx+S_tyry-S_tzrz), S_tyrz+S_tzry],
+                [S_txry-S_tyrx, S_tzrx+S_txrz, S_tyrz+S_tzry, (-S_txrx-S_tyry+S_tzrz)]])
+
+    # Get the eigenvalues/eigenvectors. They are sorted high to low, and we only
+    # need the largest one (vec[:,0]), so we don't have to check the values.
+    _, vec = np.linalg.eig(N)
+
+    # Convert quaternion to rotation
+    rotm = unit_quaternion_to_rotation_matrix(vec[:,0])
+
+    # Apply rotation
+    target_rotm = np.matmul(rotm, target)
+
+    # Compute and apply shift
+    # shift0 = reference_cent - (weights_target*target_rotm).sum(1)/weights_target.sum(1)
+    shift = reference_cent - np.dot(rotm, target_cent)
+    # print(f"shif0: {shift0} shift: {shift}")
+    target_rotm += shift[:,None]
+
+    res = np.sum((target_rotm - reference)**2)
+
+    print(rotm)
+    print(shift)
+
+    return target_rotm, rotm, shift, res
diff --git a/PYME/recipes/pointcloud.py b/PYME/recipes/pointcloud.py
@@ -40,7 +40,7 @@ def execute(self, namespace):
 
 @register_module('LocalPointDensity')
 class LocalPointDensity(ModuleBase):
-    """
+    r"""
     Estimate the local density around a localization by fitting a scaling function to the number of
     Neigbours vs distance. The expected scaling function for a uniform density is used ($N \propto r^2$
     for 2D, $N\propto r^3$ for 3D.
@@ -508,3 +508,171 @@ def execute(self, namespace):
             avg_log_prob[mask] = np.mean(log_prob[mask])
         out.addColumn(self.label_key + '_avg_log_prob', avg_log_prob)
         namespace[self.output_labeled] = out
+
+
+@register_module('IterativeClosestPoint')
+class IterativeClosestPoint(ModuleBase):
+    """ Use iterative closest point algorithm to register target point cloud to
+    reference point cloud. Useful for accounting for unknown translations and
+    rotations introduced in between, e.g., PAINT imager washes.
+
+    References
+    ----------
+    [1] https://www.cs.princeton.edu/courses/archive/fall18/cos526/notes/cos526_f18_lecture10_acquisition_registration.pdf
+    [2] Zhang, Z. Iterative point matching for registration of free-form curves and surfaces. Int J Comput Vision 13, 119–152 (1994).
+    [3] Berthold K. P. Horn, "Closed-form solution of absolute orientation using unit quaternions," J. Opt. Soc. Am. A 4, 629-642 (1987) 
+
+    Parameters
+    ----------
+    reference: PYME.IO.tabular
+        The point cloud to which we will register the other point cloud.
+    to_register: PYME.IO.tabular
+        This point cloud will be translated and rotated to match the 
+        reference pointcloud.
+    max_iters : int
+        Maximum number of iterations to use to register points
+    distance_threshold: float
+        Maximum distance between points at which they are considered close enough
+        for registration. If set to -1, this will be automatically calculated.
+    max_points : Int
+        Maximum number of points to use per iteration for registration. If set to
+        -1, all points passing `distance_threshold` will be used.
+    sigma_x : str
+        Key for localization precision along x dimension
+    sigma_y : str
+        Key for localization precision along y dimension
+    sigma_z : str
+        Key for localization precision along z dimension
+
+    Returns
+    -------
+    output : PYME.IO.tabular
+        The to_register point cloud mapped onto the reference point cloud.
+    """
+
+    reference = Input('reference')
+    to_register = Input('to_register')
+    output = Output('registered')
+    max_iters = Int(5)
+    distance_threshold = Float(-1)
+    max_points = Int(1000)
+    sigma_x = CStr('error_x')
+    sigma_y = CStr('error_y')
+    sigma_z = CStr('error_z')
+
+    def execute(self, namespace):
+        from scipy.spatial import KDTree
+        from PYME.Analysis.points.coordinate_tools import absolute_orientation
+        from PYME.IO import MetaDataHandler
+
+        reference = namespace[self.reference]
+        target = namespace[self.to_register]
+
+        reference_pts0 = np.vstack([reference['x'], reference['y'], reference['z']]).T
+        reference_tree = KDTree(reference_pts0)
+
+        target_pts = np.vstack([target['x'], target['y'], target['z']])
+        # print(reference_pts0.shape, target_pts.shape)
+
+        # print("reference x:", reference['x'][:10])
+        # print("referencepts0 x:", reference_pts0[:10,0])
+
+        rot_tot = []
+        shift_tot = []
+
+        for k in range(self.max_iters):
+            # Get the 1 nearest neighbor of each target point in reference
+            dist, idxs_reference = reference_tree.query(target_pts.T)
+
+            # print("idxs_reference: ", idxs_reference[:10])
+
+            # Reject pairs further apart than self.distance_threshold 
+            if self.distance_threshold  < 0:
+                # Calculate the distance threshold as median + mad
+                # See [2] for other threshold options
+                median = np.median(dist)
+                mad = np.median(np.abs(dist-median))
+                distance_threshold = median + mad
+                if distance_threshold <= 0:
+                    # all points are close
+                    break
+            else:
+                distance_threshold = self.distance_threshold
+            idxs_dist = np.flatnonzero(dist < distance_threshold)
+
+            # Crop down to self.max_points
+            if (self.max_points > 0) and (len(idxs_dist) > self.max_points):
+                idxs_dist = np.random.choice(idxs_dist, size=self.max_points)
+            # Now grab the points on which to iterate
+            idxs_reference = idxs_reference[idxs_dist]
+            # print("idxs_reference: ", idxs_reference[:10])
+            idxs_target = np.arange(target_pts.shape[1])[idxs_dist]
+
+            reference_pts = np.vstack([reference['x'][idxs_reference], 
+                                       reference['y'][idxs_reference], 
+                                       reference['z'][idxs_reference]])
+            try:
+                reference_weights = np.vstack([1/reference[self.sigma_x][idxs_reference], 
+                                               1/reference[self.sigma_y][idxs_reference], 
+                                               1/reference[self.sigma_z][idxs_reference]])
+
+                # The error has to be less than the localization precision of the dataset
+                rescmp = ((1/reference_weights)**2).sum()
+            except KeyError:
+                reference_weights = None
+                # No error? Then we should be able to register the points exactly.
+                rescmp = 1  # TODO: What if target_weights finds the sigma keys?
+
+            target_pts_sm = target_pts[:, idxs_target]
+
+            try:
+                target_weights = np.vstack([1/target[self.sigma_x][idxs_target], 
+                                            1/target[self.sigma_y][idxs_target], 
+                                            1/target[self.sigma_z][idxs_target]])
+            except KeyError:
+                target_weights = None
+
+            # print(reference_pts.shape, target_pts_sm.shape, reference_weights, target_weights)
+
+            # print("reference_pts x: ", reference_pts[0,:10])
+            # print("target_pts_sm x: ", target_pts_sm[0,:10])
+
+            _, rotm, shift, res = absolute_orientation(reference_pts, 
+                                                       target_pts_sm, 
+                                                       reference_weights, 
+                                                       target_weights)
+
+            # Keep track
+            rot_tot.append(rotm)
+            shift_tot.append(shift)
+
+            logger.debug(f"Iteration {k} res: {res} rscmp: {rescmp}")
+            # print(f"Iteration {k} res: {res} rscmp: {rescmp}")
+
+            if res <= rescmp:
+                # Residuals less than the sum of the error on the reference data set
+                break
+
+            # Update the full set of points to see which moved closer
+            target_pts = np.matmul(rotm, target_pts) + shift[:,None]
+
+        # Create mapping strings
+        xstr, ystr, zstr = "x", "y", "z"
+        for r, s in zip(rot_tot, shift_tot):
+            xstrp=f"{r[0,0]}*({xstr})+{r[0,1]}*({ystr})+{r[0,2]}*({zstr})+{s[0]}"
+            ystrp=f"{r[1,0]}*({xstr})+{r[1,1]}*({ystr})+{r[1,2]}*({zstr})+{s[1]}"
+            zstrp=f"{r[2,0]}*({xstr})+{r[2,1]}*({ystr})+{r[2,2]}*({zstr})+{s[2]}"
+            xstr, ystr, zstr = xstrp, ystrp, zstrp
+
+        out = tabular.MappingFilter(target, x=xstr, y=ystr, z=zstr)
+        # out.addColumn('xp', target_pts[0,...])
+        # out.addColumn('yp', target_pts[1,...])
+        # out.addColumn('zp', target_pts[2,...])
+
+        try:
+            out.mdh = MetaDataHandler.DictMDHandler(target.mdh)
+        except AttributeError:
+            pass
+
+        namespace[self.output] = out
+
diff --git a/tests/PYME/Analysis/points/test_coordinate_tools.py b/tests/PYME/Analysis/points/test_coordinate_tools.py
@@ -137,3 +137,44 @@ def test_simple_distance_to_image_mask():
 
     distances = coordinate_tools.distance_to_image_mask(mask, points)
     np.testing.assert_array_equal(distances, np.arange(size) - 0.5 * size)
+
+def test_unit_quaternion_to_rotation_matrix():
+    mat = coordinate_tools.unit_quaternion_to_rotation_matrix(np.array([0,0,0,0]))
+    np.testing.assert_array_equal(mat, np.zeros((3,3)))
+    mat = coordinate_tools.unit_quaternion_to_rotation_matrix(np.array([1,0,0,0]))
+    np.testing.assert_array_equal(mat, np.eye(3))
+    mat = coordinate_tools.unit_quaternion_to_rotation_matrix(np.array([0,1,0,0]))
+    np.testing.assert_array_equal(mat, np.diag([1,-1,-1]))
+    mat = coordinate_tools.unit_quaternion_to_rotation_matrix(np.array([0,0,1,0]))
+    np.testing.assert_array_equal(mat, np.diag([-1,1,-1]))
+    mat = coordinate_tools.unit_quaternion_to_rotation_matrix(np.array([0,0,0,1]))
+    np.testing.assert_array_equal(mat, np.diag([-1,-1,1]))
+
+def test_absolute_orientation():
+    from PYME.simulation import locify
+    def round_box(p, w, r):
+        w = np.array(w)
+        q = np.abs(p) - w[:,None]
+        return np.linalg.norm(np.maximum(q,0.0),axis=0) + np.minimum(np.maximum(q[0,:],np.maximum(q[1,:],q[2,:])),0.0) - r
+
+    shift = [0.5,0,0]
+    rot = np.pi/4
+    cube0 = locify.points_from_sdf(lambda x: round_box(x, [0.5,0.5,0.5], 0), r_max=1, centre=(0,0,0), dx_min=0.2, p=1.0)
+    cube1 = cube0.copy()
+    cube1 += np.array(shift)[:,None]
+    cube1 = np.dot(np.array([[np.cos(rot), 0, np.sin(rot)], [0, 1, 0], [-np.sin(rot), 0, np.cos(rot)]]), cube1)
+
+    target_rotm, rotm, shift, res = coordinate_tools.absolute_orientation(cube0, cube1)
+
+    np.testing.assert_allclose(cube0, target_rotm)
+    assert res < 1e-6
+
+    weights_reference = np.random.rand(*cube0.shape)
+    weights_target = np.random.rand(*cube1.shape)
+
+    target_rotm, rotm, shift, res = coordinate_tools.absolute_orientation(cube0, 
+                                                                          cube1, 
+                                                                          weights_reference=weights_reference, 
+                                                                          weights_target=weights_target)
+
+    np.testing.assert_array_less(np.abs(cube0-target_rotm), 1/weights_reference)