From 728149805117acc324e652c8c092935158f26457 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Fri, 7 Aug 2026 06:31:26 -0400 Subject: [PATCH] ENH: Rewrite PCA model registration around an analytic-gradient objective RegisterModelsPCA previously maximized mean intensity sampled from the fixed distance map through ITK's LinearInterpolateImageFunction, with the optimizer estimating gradients by finite differences. The objective is now stated and minimized directly: mean distance(model -> target) + w * mean distance(target -> model) # symmetric term + lambda * sum(b_i^2) # Mahalanobis shape prior Because the PCA deformation is linear in the coefficients b, the gradient is analytic and handed to the optimizer instead of being estimated, which removes one objective evaluation per coefficient per step. The post-PCA transform is folded into the mode directions so the gradient stays exact; when that transform is not affine its Jacobian is not constant, so the analytic gradient is disabled and finite differences are used with a logged warning. - Add symmetric_weight (default 0.5) so partial target coverage is penalized, and pca_prior_weight (default 0.0, disabled) for the Mahalanobis prior - Sample the distance map and its gradient with scipy.ndimage.map_coordinates, and build the target-to-model term with a scipy.spatial.cKDTree - Replace the cached ITK interpolator and _create_itk_points with _prepare_sampling, which builds the arrays the objective is made of once - Add ContourTools.sample_mesh_faces for face-density-aware point sampling and a negative_inside option on the signed distance map - Log transform fidelity after computing the PCA transforms Tests cover the pieces that were previously unverified: the analytic gradient against finite differences, recovery of known coefficients, the symmetric term penalizing partial coverage, the prior shrinking coefficients, eigenvector scaling by standard deviation, deformation happening in the template frame, and a transform round trip. --- src/physiotwin4d/contour_tools.py | 130 +++- src/physiotwin4d/register_models_pca.py | 777 ++++++++++++++++-------- tests/test_register_models_pca.py | 273 ++++++++- 3 files changed, 876 insertions(+), 304 deletions(-) diff --git a/src/physiotwin4d/contour_tools.py b/src/physiotwin4d/contour_tools.py index 67558d1f..3d7d2c73 100644 --- a/src/physiotwin4d/contour_tools.py +++ b/src/physiotwin4d/contour_tools.py @@ -318,6 +318,54 @@ def create_labelmap_from_meshes( return labelmap_image + @staticmethod + def sample_mesh_faces(mesh: pv.DataSet, max_spacing: float) -> np.ndarray: + """Return mesh points supplemented by samples across the triangle faces. + + Rasterizing vertices alone leaves gaps between them on meshes that are + coarse relative to the voxel size, which makes a distance map built from + them ripple. Adding barycentric samples dense enough that consecutive + samples are closer than ``max_spacing`` closes those gaps. + + Args: + mesh: Source mesh; its surface is triangulated if needed. + max_spacing: Target spacing between samples, in mm. + + Returns: + (n, 3) array of sample points, starting with the mesh's own points. + """ + points = np.asarray(mesh.points, dtype=np.float64) + surface = mesh.extract_surface() if not isinstance(mesh, pv.PolyData) else mesh + surface = surface.triangulate() + if surface.faces.size == 0: + return points + faces = surface.faces.reshape(-1, 4)[:, 1:] + + vertices = np.asarray(surface.points, dtype=np.float64) + corners = vertices[faces] # (n_faces, 3, 3) + edge_lengths = np.linalg.norm( + corners - np.roll(corners, 1, axis=1), axis=2 + ).max(axis=1) + + # Group faces by how finely they need to be subdivided so each division + # level is generated as one vectorized batch. + divisions = np.maximum(1, np.ceil(edge_lengths / max(max_spacing, 1e-6))) + divisions = np.minimum(divisions, 64).astype(np.int64) + + samples = [points] + for level in np.unique(divisions): + if level < 2: + continue + selected = corners[divisions == level] + # Barycentric lattice with `level` divisions per edge. + steps = np.arange(level + 1, dtype=np.float64) / level + u, v = np.meshgrid(steps, steps, indexing="ij") + mask = (u + v) <= 1.0 + weights = np.column_stack([1.0 - u[mask] - v[mask], u[mask], v[mask]]) + samples.append(np.einsum("fca,kc->fka", selected, weights).reshape(-1, 3)) + + return np.concatenate(samples, axis=0) + def create_distance_map( self, mesh: pv.DataSet | pv.UnstructuredGrid, @@ -326,37 +374,62 @@ def create_distance_map( negative_inside: bool = True, zero_inside: bool = False, norm_to_max_distance: float = 0.0, + sample_faces: bool = True, ) -> itk.Image: - self.log_info("Computing signed distance map...") + """Compute a distance map of a mesh on the reference image's grid. - # Convert mask to binary - points = mesh.points + Args: + mesh: Mesh whose surface the distances are measured to. + reference_image: Image defining the output grid. + squared_distance: Sign-preserving square of the result. Default: False + negative_inside: Keep the signed output. Default: True + zero_inside: Clip negative values to zero before anything else. + Default: False + norm_to_max_distance: If non-zero, divide by this value and clip to + [-1, 1]. Default: 0.0 (distances stay in mm) + sample_faces: Rasterize samples across the triangle faces as well as + the vertices, so that coarse meshes do not leave gaps in the + rasterized surface. Default: True + + Returns: + ITK image of distances on the reference grid. + """ + self.log_info("Computing signed distance map...") size = reference_image.GetLargestPossibleRegion().GetSize() + if sample_faces: + points = self.sample_mesh_faces( + mesh, 0.5 * float(min(reference_image.GetSpacing())) + ) + self.log_debug( + "Distance map: %d face samples from %d mesh points", + len(points), + mesh.n_points, + ) + else: + points = np.asarray(mesh.points, dtype=np.float64) + # NumPy convention is (z, y, x); ITK GetSize() returns (x, y, z) - tmp_arr = np.zeros((size[2], size[1], size[0]), dtype=np.int32) - itk_point = itk.Point[itk.D, 3]() - point_count = 0 - for point in points: - itk_point[0] = float(point[0]) - itk_point[1] = float(point[1]) - itk_point[2] = float(point[2]) - indx = reference_image.TransformPhysicalPointToIndex(itk_point) - if ( - indx[0] < 0 - or indx[1] < 0 - or indx[2] < 0 - or indx[0] >= size[0] - or indx[1] >= size[1] - or indx[2] >= size[2] - ): - continue - tmp_arr[indx[2], indx[1], indx[0]] = 1 - point_count += 1 + tmp_arr = np.zeros((size[2], size[1], size[0]), dtype=np.uint8) + + # Bulk equivalent of TransformPhysicalPointToIndex, which rounds half up. + index_to_world = itk.array_from_matrix( + reference_image.GetDirection() + ) @ np.diag(np.asarray(reference_image.GetSpacing())) + origin = np.asarray(reference_image.GetOrigin(), dtype=np.float64) + indices = np.floor( + (points - origin) @ np.linalg.inv(index_to_world).T + 0.5 + ).astype(np.int64) + size_arr = np.array([size[0], size[1], size[2]], dtype=np.int64) + inside = np.all((indices >= 0) & (indices < size_arr), axis=1) + indices = indices[inside] + point_count = len(indices) + if point_count: + tmp_arr[indices[:, 2], indices[:, 1], indices[:, 0]] = 1 self.log_info( - "Distance map: %d/%d surface points within reference image", + "Distance map: %d/%d surface samples within reference image", point_count, len(points), ) @@ -370,8 +443,17 @@ def create_distance_map( str(size), str(reference_image.GetSpacing()), ) + elif not inside.all(): + # Distances near the dropped region are measured to whatever samples + # remain in the grid, so they are larger than the true distance. + self.log_warning( + "%d of %d surface samples fall outside the reference image; " + "distances near that boundary are overestimated.", + len(points) - point_count, + len(points), + ) - tmp_binary_image = itk.GetImageFromArray(tmp_arr.astype(np.uint8)) + tmp_binary_image = itk.GetImageFromArray(tmp_arr) tmp_binary_image.CopyInformation(reference_image) assert ( tmp_binary_image.GetLargestPossibleRegion().GetSize() diff --git a/src/physiotwin4d/register_models_pca.py b/src/physiotwin4d/register_models_pca.py index ef93d97c..c2cf68d0 100644 --- a/src/physiotwin4d/register_models_pca.py +++ b/src/physiotwin4d/register_models_pca.py @@ -8,7 +8,9 @@ import itk import numpy as np import pyvista as pv +from scipy.ndimage import map_coordinates from scipy.optimize import minimize +from scipy.spatial import cKDTree from typing_extensions import Self from .contour_tools import ContourTools @@ -17,29 +19,44 @@ class RegisterModelsPCA(PhysioTwin4DBase): - """Register PCA-based shape models to medical images using mean distance optimization. + """Register PCA-based shape models to images by minimizing a distance metric. This class implements a registration pipeline for fitting statistical shape models to patient-specific medical images: **PCA Deformable Registration** - Optimizes PCA coefficients - - Model equation: P = mean + Σ(b_i * std_i * pca_eigenvector_i) - - Maximizes mean distance at deformed model points P + - Model equation: P = template + Σ(b_i * std_i * pca_eigenvector_i) + - Minimizes the mean distance-to-target at the deformed model points P **Optimization Objective:** - Maximize the mean distance of the image sampled at model points using - ITK's LinearInterpolateImageFunction. This aligns the model with bright - regions in contrast-enhanced images (e.g., blood pool in cardiac CT). + ``fixed_distance_map`` is zero on the target surface and grows with + distance away from it (in mm), so the objective is *minimized*:: + + E(b) = (1 - w) * mean_i D(P_i(b)) # model -> target + + w * mean_j ||Q_j - P_nn(j)|| # target -> model + + lambda * Σ b_i² # Mahalanobis prior + + ``w`` is ``symmetric_weight`` and ``lambda`` is ``pca_prior_weight``. + Because the deformation is linear in ``b``, the gradient is analytic and + is supplied to the optimizer directly. + + **Coordinate frames:** + The eigenvectors are directions in the statistical model's own training + frame, so they are only valid when added to a template in that same + frame. Any rigid/affine alignment to the target must be supplied as + ``post_pca_transform``, which is applied *after* the deformation, rather + than pre-applied to ``pca_template_model``. Attributes: pca_template_model (pv.DataSet): Mean shape model pca_eigenvectors (np.ndarray): PCA eigenvectors/components (modes × n_points*3) pca_std_deviations (np.ndarray): Standard deviations per mode (modes,) - fixed_distance_map (itk.Image): Patient image providing distance data - n_points (int): Number of points in the model + fixed_distance_map (itk.Image): Distance map of the target, in mm + fixed_model (pv.DataSet): Target model, when one was supplied. Required + for the symmetric (target-to-model) term. pca_number_of_modes (int): Number of PCA modes available - pca_coefficients (np.ndarray): Optimized PCA coefficients + registered_model_pca_coefficients (np.ndarray): Optimized PCA coefficients registered_model (pv.DataSet): Final registered and deformed model post_pca_transform (itk.Transform): Transform to apply after PCA registration forward_point_transform (itk.DisplacementFieldTransform): POINT transform @@ -90,63 +107,99 @@ def __init__( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ): """Initialize the PCA-based model-to-image registration. Args: pca_template_model: PyVista model containing the mean 3D shape model - (unstructured grid or polydata) + (unstructured grid or polydata). It must be in the same frame + the PCA modes were trained in; supply any alignment to the + target as post_pca_transform instead of pre-applying it here. pca_eigenvectors: Numpy array of PCA eigenvectors/components. Shape: (modes, n_points*3) Each row is a flattened eigenmode with 3D displacements: [x1,y1,z1, x2,y2,z2, ...] pca_std_deviations: Numpy array of standard deviations per PCA mode. Shape: (modes,) These are the square roots of pca_eigenvalues - pca_number_of_modes: Number of PCA modes to use. Default: -1 (use all) + pca_number_of_modes: Number of PCA modes to use. Default: 0 (use all) pca_template_model_point_subsample: Step size for subsampling model points. Default: 4 post_pca_transform: Optional ITK transform to apply after PCA registration. Default: None - fixed_distance_map: ITK image providing the distance map. + fixed_distance_map: ITK image providing the distance map, in mm. Default: None fixed_model: PyVista model used to compute the distance map, if one isn't provided. + Also supplies the target points for the symmetric term. reference_image: ITK image providing coordinate frame for computing the distance map. + pca_prior_weight: Weight (in mm) of the Mahalanobis shape prior + ``lambda * sum(b_i**2)``. Because b is expressed in standard + deviations, this term is the squared Mahalanobis distance in + shape space and makes the fit a MAP estimate rather than a pure + data fit constrained only by the coefficient bounds. + Default: 0.0 (prior disabled). + symmetric_weight: Weight in [0, 1] of the target-to-model distance + term. 0.0 measures model-to-target only, which lets the model + satisfy the metric while covering just part of the target. + Requires fixed_model; ignored with a warning when only a + distance map is available. Default: 0.5 log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO Raises: - ValueError: If pca_eigenvector dimensions don't match model points + ValueError: If pca_eigenvector dimensions don't match model points, + if the mode counts disagree, or if neither a distance map nor a + fixed model plus reference image is provided. """ # Initialize base class with logging super().__init__(class_name="RegisterModelsPCA", log_level=log_level) # Store model data self.pca_template_model: pv.DataSet = pca_template_model - self.pca_eigenvectors: np.ndarray = pca_eigenvectors - self.pca_std_deviations: np.ndarray = pca_std_deviations + self.pca_eigenvectors: np.ndarray = np.asarray( + pca_eigenvectors, dtype=np.float64 + ) + self.pca_std_deviations: np.ndarray = np.asarray( + pca_std_deviations, dtype=np.float64 + ) + + if self.pca_eigenvectors.ndim != 2: + raise ValueError( + f"pca_eigenvectors must be 2D (modes, n_points*3), got shape " + f"{self.pca_eigenvectors.shape}" + ) + expected_size = pca_template_model.n_points * 3 + if self.pca_eigenvectors.shape[1] != expected_size: + raise ValueError( + f"Component dimension mismatch: expected {expected_size} " + f"(3 × {pca_template_model.n_points} points), got " + f"{self.pca_eigenvectors.shape[1]}" + ) + if self.pca_eigenvectors.shape[0] != self.pca_std_deviations.shape[0]: + raise ValueError( + f"Mode count mismatch: {self.pca_eigenvectors.shape[0]} eigenvectors " + f"but {self.pca_std_deviations.shape[0]} standard deviations" + ) self.post_pca_transform = post_pca_transform self._contour_tools = ContourTools() + self.fixed_model: Optional[pv.DataSet] = fixed_model self.fixed_distance_map = fixed_distance_map if ( self.fixed_distance_map is None and fixed_model is not None and reference_image is not None ): - self.fixed_model = fixed_model - self.fixed_distance_map = self._contour_tools.create_distance_map( - fixed_model, - reference_image, - squared_distance=False, - negative_inside=False, - zero_inside=True, - norm_to_max_distance=200.0, + self.fixed_distance_map = self._create_distance_map( + fixed_model, reference_image ) elif self.fixed_distance_map is not None and ( fixed_model is not None or reference_image is not None ): self.log_warning( - "Fixed model and reference image will be ignored because a distance map is provided." + "A distance map was provided, so the reference image is ignored; " + "the fixed model is retained only for the symmetric metric term." ) elif self.fixed_distance_map is None and ( fixed_model is None or reference_image is None @@ -160,9 +213,11 @@ def __init__( self.pca_number_of_modes: int = pca_number_of_modes if self.pca_number_of_modes <= 0: - self.pca_number_of_modes = len(pca_std_deviations) + self.pca_number_of_modes = len(self.pca_std_deviations) self.pca_template_model_point_subsample = pca_template_model_point_subsample + self.pca_prior_weight = pca_prior_weight + self.symmetric_weight = symmetric_weight # outputs self.registered_model_pca_coefficients: Optional[np.ndarray] = None @@ -172,17 +227,26 @@ def __init__( self.forward_point_transform: Optional[itk.DisplacementFieldTransform] = None self.inverse_point_transform: Optional[itk.DisplacementFieldTransform] = None - # Image interpolator (created when needed) - self._fixed_distance_map_interpolator: Optional[ - itk.LinearInterpolateImageFunction - ] = None + # Sampling caches, built lazily by _prepare_sampling() + self._sampling_ready: bool = False + self._analytic_gradient: bool = True self._fixed_distance_map_max_distance: float = 0.0 + self._post_pca_affine_key: Optional[itk.Transform] = None + self._post_pca_affine: Optional[tuple[np.ndarray, np.ndarray]] = None self._metric_call_count: int = 0 - # Pre-convert mean shape points to ITK format - self._pca_template_model_points_itk: Optional[list[itk.Point]] = None - self._create_itk_points() + def _create_distance_map( + self, fixed_model: pv.DataSet, reference_image: itk.Image + ) -> itk.Image: + """Build the unsigned, un-normalized (mm) distance map of the target.""" + return self._contour_tools.create_distance_map( + fixed_model, + reference_image, + squared_distance=False, + negative_inside=False, + zero_inside=True, + ) @classmethod def from_json( @@ -195,6 +259,8 @@ def from_json( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ) -> Self: """Create RegisterModelsPCA from PCA model JSON file. @@ -217,6 +283,8 @@ def from_json( for registration. If None, must be set later before registration. fixed_model: Target surface mesh to register to. Default: None reference_image: Reference image defining coordinate space. Default: None + pca_prior_weight: Weight (mm) of the Mahalanobis shape prior. Default: 0.0 + symmetric_weight: Weight of the target-to-model term. Default: 0.5 log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO @@ -288,6 +356,8 @@ def from_json( fixed_distance_map=fixed_distance_map, fixed_model=fixed_model, reference_image=reference_image, + pca_prior_weight=pca_prior_weight, + symmetric_weight=symmetric_weight, log_level=log_level, ) @@ -302,6 +372,8 @@ def from_pca_model( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ) -> Self: """Create RegisterModelsPCA from a PCA model dictionary. @@ -320,6 +392,8 @@ def from_pca_model( fixed_distance_map: ITK image providing the distance values for registration. fixed_model: Target surface mesh to register to. reference_image: Reference image defining coordinate space. + pca_prior_weight: Weight (mm) of the Mahalanobis shape prior. Default: 0.0 + symmetric_weight: Weight of the target-to-model term. Default: 0.5 log_level: Logging level. Returns: @@ -334,12 +408,6 @@ def from_pca_model( if "components" not in pca_model: raise ValueError("'components' field not found in pca_model") pca_eigenvectors = np.array(pca_model["components"], dtype=np.float64) - expected_size = pca_template_model.n_points * 3 - if pca_eigenvectors.shape[1] != expected_size: - raise ValueError( - f"Component dimension mismatch: expected {expected_size} " - f"(3 × {pca_template_model.n_points} points), got {pca_eigenvectors.shape[1]}" - ) return cls( pca_template_model=pca_template_model, pca_eigenvectors=pca_eigenvectors, @@ -350,38 +418,18 @@ def from_pca_model( fixed_distance_map=fixed_distance_map, fixed_model=fixed_model, reference_image=reference_image, + pca_prior_weight=pca_prior_weight, + symmetric_weight=symmetric_weight, log_level=log_level, ) - def _create_itk_points(self) -> None: - """Pre-convert mean shape points to ITK Point format for efficiency. - - This method creates ITK Point objects once at initialization, avoiding - repeated conversions during optimization iterations. - """ - self.log_info("Converting mean shape points to ITK format...") - - self._pca_template_model_points_itk = [] - for point in self.pca_template_model.points: - itk_point = itk.Point[itk.D, 3]() - itk_point[0] = float(point[0]) - itk_point[1] = float(point[1]) - itk_point[2] = float(point[2]) - self._pca_template_model_points_itk.append(itk_point) - - self.log_info( - f" Converted {len(self._pca_template_model_points_itk)} points to ITK format" - ) - def set_fixed_model( self, fixed_model: pv.UnstructuredGrid, reference_image: Optional[itk.Image] ) -> None: - """Set the fixed model for registration. - - If this is set, the fixed distance map will be set to None. + """Set the fixed model for registration and rebuild its distance map. Args: - fixed_model: PyVista model used to compute the distance map, if one isn't provided. + fixed_model: PyVista model used to compute the distance map. reference_image: ITK image providing coordinate frame for computing the distance map. """ if reference_image is None: @@ -389,26 +437,20 @@ def set_fixed_model( "reference_image must not be None when setting a fixed model" ) - self.fixed_distance_map = self._contour_tools.create_distance_map( - fixed_model, - reference_image, - squared_distance=False, - negative_inside=False, - zero_inside=True, - norm_to_max_distance=200.0, + self.fixed_model = fixed_model + self.fixed_distance_map = self._create_distance_map( + fixed_model, reference_image ) - self._fixed_distance_map_interpolator = None + self._sampling_ready = False def set_fixed_distance_map(self, fixed_distance_map: Optional[itk.Image]) -> None: - """Set the reference image for registration. - - If this is set, the fixed model will be set to None. + """Set the distance map used as the registration target. Args: - fixed_distance_map: ITK image providing distance data + fixed_distance_map: ITK image providing distance data, in mm """ self.fixed_distance_map = fixed_distance_map - self._fixed_distance_map_interpolator = None + self._sampling_ready = False def set_pca_template_model(self, pca_template_model: pv.UnstructuredGrid) -> None: """Set the average model for registration. @@ -418,132 +460,298 @@ def set_pca_template_model(self, pca_template_model: pv.UnstructuredGrid) -> Non (unstructured grid or polydata) """ self.pca_template_model = pca_template_model + self._sampling_ready = False + self.log_info(" Average model set successfully!") - self._pca_template_model_points_itk = None + def _affine_of_transform( + self, transform: itk.Transform + ) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Recover (matrix, offset) if ``transform`` acts affinely, else None. - self._create_itk_points() - self.log_info(" Average model set successfully!") + Probing the transform rather than querying ``GetMatrix()`` works for any + ITK transform type and correctly rejects the non-affine ones (such as a + displacement field), for which no constant Jacobian exists. + """ - def _mean_distance_metric( - self, - params: np.ndarray, - ) -> float: - """Evaluate the optimization metric (mean intensity) at model points. + def apply(vector: np.ndarray) -> np.ndarray: + point = itk.Point[itk.D, 3]() + point[0], point[1], point[2] = (float(v) for v in vector) + result = transform.TransformPoint(point) + return np.array([result[0], result[1], result[2]], dtype=np.float64) - This is the objective function to be MAXIMIZED during optimization. - Higher values indicate better alignment with bright regions. + offset = apply(np.zeros(3)) + matrix = np.column_stack( + [apply(basis) - offset for basis in np.eye(3, dtype=np.float64)] + ) + probe = np.array([0.37, -0.61, 0.83], dtype=np.float64) + scale = max(1.0, float(np.abs(matrix).max()), float(np.abs(offset).max())) + if not np.allclose(apply(probe), matrix @ probe + offset, atol=1e-9 * scale): + return None + return matrix, offset - Args: - pca_deformation: Nx3 numpy array of PCA deformation vectors to add to points. - If None, no deformation is applied. + def _get_post_pca_affine(self) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Return the cached (matrix, offset) of post_pca_transform, or None. - Returns: - Mean distance value across all points + Keyed on the transform object so that reassigning post_pca_transform + invalidates the cache. """ - pca_deformation = self._compute_pca_deformation(params) - - # Create interpolator if not already cached (inline creation) - if self._fixed_distance_map_interpolator is None: - if self.fixed_distance_map is None: - self.log_error("Distance map is not set.") - raise ValueError("Distance map must be set before registering.") - ImageType = type(self.fixed_distance_map) - self._fixed_distance_map_interpolator = itk.LinearInterpolateImageFunction[ - ImageType, itk.D - ].New() - self._fixed_distance_map_interpolator.SetInputImage(self.fixed_distance_map) - fixed_distance_map_array = itk.GetArrayFromImage(self.fixed_distance_map) - self._fixed_distance_map_max_distance = fixed_distance_map_array.max() - self.log_debug("Interpolator created") - self.log_debug( - " Max distance = %s", self._fixed_distance_map_max_distance - ) + if self.post_pca_transform is None: + return None + if self._post_pca_affine_key is not self.post_pca_transform: + self._post_pca_affine_key = self.post_pca_transform + self._post_pca_affine = self._affine_of_transform(self.post_pca_transform) + return self._post_pca_affine + + def _prepare_sampling(self) -> None: + """Build the cached arrays the objective and its gradient are made of. + + Everything that does not depend on the PCA coefficients is computed once + here: the subsampled template points, the per-mode displacement vectors + already scaled by their standard deviation and mapped through the + post-PCA transform, the distance map and its gradient, and the + physical-to-index mapping used to sample them. + """ + if self.fixed_distance_map is None: + self.log_error("Distance map is not set.") + raise ValueError("Distance map must be set before registering.") + + template_points = np.asarray(self.pca_template_model.points, dtype=np.float64) + step = max(1, self.pca_template_model_point_subsample) + self._sample_slice = slice(None, None, step) + self._sample_points = template_points[self._sample_slice] + + # (modes, m, 3) displacement per unit coefficient, in template space. + modes = self.pca_eigenvectors.reshape(self.pca_eigenvectors.shape[0], -1, 3) + self._sample_modes = ( + modes[:, self._sample_slice, :] * self.pca_std_deviations[:, None, None] + ) - self.log_debug("Evaluating params = %s", params) - self.log_debug(" Max displacement = %s", pca_deformation.max(axis=0)) + # Fold the post-PCA transform into the mode directions so the gradient + # is expressed directly in world space. A non-affine post-PCA transform + # has no constant Jacobian, so the analytic gradient is disabled. + affine = self._get_post_pca_affine() + if self.post_pca_transform is not None and affine is None: + self.log_warning( + "post_pca_transform is not affine; falling back to a " + "finite-difference gradient." + ) + self._sample_modes_world = ( + self._sample_modes if affine is None else self._sample_modes @ affine[0].T + ) + self._analytic_gradient = self.post_pca_transform is None or affine is not None + + # Physical point -> continuous index: index = affine_inv @ (p - origin). + image = self.fixed_distance_map + direction = itk.array_from_matrix(image.GetDirection()) + index_to_world = direction @ np.diag(np.asarray(image.GetSpacing())) + self._index_to_world = index_to_world + self._world_to_index = np.linalg.inv(index_to_world) + self._image_origin = np.asarray(image.GetOrigin(), dtype=np.float64) + size = image.GetLargestPossibleRegion().GetSize() + self._image_size = np.array([size[0], size[1], size[2]], dtype=np.float64) + + # Array axes are (k, j, i), so index component a is array axis 2 - a. + # The forward differences are the exact derivative of the trilinear + # interpolant used to sample the map, which keeps the objective and its + # gradient consistent; a central difference would not. + self._distance_array = np.asarray( + itk.array_view_from_image(image), dtype=np.float64 + ) + self._fixed_distance_map_max_distance = float(self._distance_array.max()) + self._distance_forward_diff = tuple( + np.diff(self._distance_array, axis=2 - axis) + if self._distance_array.shape[2 - axis] > 1 + else None + for axis in range(3) + ) - # Sample distance at each point - n_valid_points = 0 - n_invalid_points = 0 - total_distance = 0.0 - center = np.zeros(3) - point = itk.Point[itk.D, 3]() - assert self.fixed_distance_map is not None, "fixed_distance_map must be set" - assert self._pca_template_model_points_itk is not None, ( - "ITK points must be initialized" + # Target points for the symmetric term. + self._target_points: Optional[np.ndarray] = None + if self.symmetric_weight > 0.0: + if self.fixed_model is None: + self.log_warning( + "symmetric_weight is %.3g but no fixed_model is available; " + "the target-to-model term is disabled.", + self.symmetric_weight, + ) + else: + self._target_points = np.asarray( + self.fixed_model.points, dtype=np.float64 + )[self._sample_slice] + + self._sampling_ready = True + self.log_debug( + "Sampling prepared: %d model points, %s target points, max distance %.3f mm", + self._sample_points.shape[0], + "no" if self._target_points is None else str(len(self._target_points)), + self._fixed_distance_map_max_distance, ) - image_size = self.fixed_distance_map.GetBufferedRegion().GetSize() - for i, base_point in enumerate(self._pca_template_model_points_itk): - if i % self.pca_template_model_point_subsample != 0: - continue - # Start with base point - point[0] = base_point[0] - point[1] = base_point[1] - point[2] = base_point[2] + def _sample_distance( + self, points: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, int]: + """Sample the distance map and its spatial gradient at world points. - # Add PCA deformation if provided - point[0] += pca_deformation[i, 0] - point[1] += pca_deformation[i, 1] - point[2] += pca_deformation[i, 2] + Points outside the image are clamped to the grid and charged the extra + travel from the clamped location, which keeps the metric continuous and + gives the optimizer a gradient that pushes such points back inside -- + unlike a constant out-of-bounds penalty, which is flat. - if self.post_pca_transform is not None: - point = self.post_pca_transform.TransformPoint(point) + Args: + points: (n, 3) array of world-space points - # Check if point is inside image bounds + Returns: + Tuple of (distances (n,), world-space gradients (n, 3), n_outside) + """ + index = (points - self._image_origin) @ self._world_to_index.T + clamped = np.clip(index, 0.0, self._image_size - 1.0) + outside = index - clamped + is_outside = np.any(outside != 0.0, axis=1) + n_outside = int(np.count_nonzero(is_outside)) + + # map_coordinates indexes the array as (k, j, i). + array_coordinates = clamped[:, ::-1] + distances = map_coordinates( + self._distance_array, array_coordinates.T, order=1, mode="nearest" + ) - coord_index = ( - self.fixed_distance_map.TransformPhysicalPointToContinuousIndex(point) + # d/dc_a of the trilinear interpolant is the forward difference along a, + # interpolated linearly across the other two axes within the same cell. + gradient_index = np.zeros_like(clamped) + for axis in range(3): + differences = self._distance_forward_diff[axis] + if differences is None: + continue + coordinates = array_coordinates.copy() + coordinates[:, 2 - axis] = np.clip( + np.floor(clamped[:, axis]), 0.0, self._image_size[axis] - 2.0 ) - if ( - 0 <= coord_index[0] < image_size[0] - and 0 <= coord_index[1] < image_size[1] - and 0 <= coord_index[2] < image_size[2] - ): - center[0] += point[0] - center[1] += point[1] - center[2] += point[2] - distance = ( - self._fixed_distance_map_interpolator.EvaluateAtContinuousIndex( - coord_index - ) - ) - total_distance += distance - n_valid_points += 1 - else: - n_invalid_points += 1 - - if n_invalid_points >= 0.05 * n_valid_points: - self.log_warning( - "%d of %d mapped outside of image. Rejecting.", - n_invalid_points, - n_valid_points + n_invalid_points, + gradient_index[:, axis] = map_coordinates( + differences, coordinates.T, order=1, mode="nearest" ) - return self._fixed_distance_map_max_distance + # A clamped axis cannot change the sampled value, so it carries no + # gradient from the map; the out-of-bounds term supplies it instead. + gradient_index[outside != 0.0] = 0.0 + gradients = gradient_index @ self._world_to_index + + if n_outside: + outside_world = outside @ self._index_to_world.T + outside_distance = np.linalg.norm(outside_world, axis=1) + safe = np.where(outside_distance > 0.0, outside_distance, 1.0) + distances = distances + outside_distance + gradients = gradients + outside_world / safe[:, None] + + return distances, gradients, n_outside + + def _deform(self, pca_coefficients: np.ndarray) -> np.ndarray: + """Deform the subsampled template points into world space.""" + n_modes = len(pca_coefficients) + points = self._sample_points + np.tensordot( + pca_coefficients, self._sample_modes[:n_modes], axes=(0, 0) + ) + return self._apply_post_pca_transform(points) - # Compute mean distance - mean_distance = total_distance / n_valid_points - center /= n_valid_points + def _objective_and_gradient(self, params: np.ndarray) -> tuple[float, np.ndarray]: + """Evaluate the registration objective and its gradient. - log_level_int = ( - self.log_level - if isinstance(self.log_level, int) - else logging.getLevelName(self.log_level) - ) - if log_level_int <= logging.DEBUG or self._metric_call_count % 100 == 0: - self.log_info( - " Metric %d: %s -> %f", - (self._metric_call_count + 1), - center, - mean_distance, + The objective is MINIMIZED: the distance map is zero on the target + surface and grows away from it. + + Args: + params: PCA coefficients b, in units of standard deviations + + Returns: + Tuple of (objective value in mm, gradient with respect to params) + """ + if not self._sampling_ready: + self._prepare_sampling() + + n_modes = len(params) + modes_world = self._sample_modes_world[:n_modes] + points = self._deform(params) + + distances, gradients, n_outside = self._sample_distance(points) + forward_distance = float(distances.mean()) + # d/db_j of mean_i D(p_i) = mean_i grad_D(p_i) . (sigma_j * v_ij) + forward_gradient = np.einsum("ia,jia->j", gradients, modes_world) / len(points) + + weight = self.symmetric_weight if self._target_points is not None else 0.0 + reverse_distance = 0.0 + reverse_gradient = np.zeros(n_modes, dtype=np.float64) + if weight > 0.0: + assert self._target_points is not None, "target points must be set" + # Target -> model: each target point is charged its distance to the + # nearest deformed model point, so a model that covers only part of + # the target scores badly. The forward term alone cannot see this. + nearest, nearest_index = cKDTree(points).query(self._target_points) + nearest_distance = np.atleast_1d(np.asarray(nearest, dtype=np.float64)) + reverse_distance = float(nearest_distance.mean()) + safe = np.where(nearest_distance > 0.0, nearest_distance, 1.0) + direction = (points[nearest_index] - self._target_points) / safe[:, None] + # Accumulate each target's pull onto the model point it selected, + # then contract once against the modes. + pull = np.zeros_like(points) + np.add.at(pull, nearest_index, direction) + pull /= len(self._target_points) + reverse_gradient = np.einsum("ia,jia->j", pull, modes_world) + + objective = (1.0 - weight) * forward_distance + weight * reverse_distance + gradient = (1.0 - weight) * forward_gradient + weight * reverse_gradient + + prior = 0.0 + if self.pca_prior_weight > 0.0: + prior = self.pca_prior_weight * float(np.dot(params, params)) + objective += prior + gradient = gradient + 2.0 * self.pca_prior_weight * params + + if n_outside > 0.25 * len(points): + self.log_warning( + "%d of %d model points mapped outside the distance map.", + n_outside, + len(points), ) + + if self.log_level <= logging.DEBUG or self._metric_call_count % 25 == 0: self.log_info( - " Params %s", - params, + " Metric %d: %.4f mm (model->target %.4f, target->model %.4f, " + "prior %.4f, outside %d)", + self._metric_call_count + 1, + objective, + forward_distance, + reverse_distance, + prior, + n_outside, ) + self.log_debug(" Params %s", params) self._metric_call_count += 1 - return mean_distance + return objective, gradient + + def _mean_distance_metric(self, params: np.ndarray) -> float: + """Evaluate the registration objective at the given PCA coefficients. + + Args: + params: PCA coefficients b, in units of standard deviations + + Returns: + Objective value, in mm. Lower is better. + """ + return self._objective_and_gradient(np.asarray(params, dtype=np.float64))[0] + + def _apply_post_pca_transform(self, points: np.ndarray) -> np.ndarray: + """Apply post_pca_transform to an (n, 3) array of world points.""" + affine = self._get_post_pca_affine() + if affine is not None: + return np.asarray(points @ affine[0].T + affine[1], dtype=np.float64) + if self.post_pca_transform is None: + return points + transformed = np.empty_like(points) + point = itk.Point[itk.D, 3]() + for i, source in enumerate(points): + point[0], point[1], point[2] = (float(v) for v in source) + result = self.post_pca_transform.TransformPoint(point) + transformed[i] = (result[0], result[1], result[2]) + return transformed def _compute_pca_deformation(self, pca_coefficients: np.ndarray) -> np.ndarray: """Compute PCA deformation vectors for all points. @@ -552,28 +760,18 @@ def _compute_pca_deformation(self, pca_coefficients: np.ndarray) -> np.ndarray: displacement = Σ(b_i * std_i * pca_eigenvector_i) Args: - pca_coefficients: Array of PCA coefficients b_i (one per mode) - pca_number_of_modes: Number of PCA modes to use. Default: use all available modes + pca_coefficients: Array of PCA coefficients b_i. Only as many modes + as there are coefficients contribute. Returns: - Nx3 array of deformation vectors (displacement from mean shape) + Nx3 array of deformation vectors (displacement from the template) """ - # Initialize deformation to zero - deformation = np.zeros((self.pca_template_model.n_points, 3), dtype=np.float64) - - # Add contribution from each PCA mode - for i in range(self.pca_number_of_modes): - pca_eigenvector_flat = self.pca_eigenvectors[i, :] - - # Reshape to (N, 3) - pca_eigenvector_3d = pca_eigenvector_flat.reshape(-1, 3) - - # Add weighted deformation: b_i * std_i * pca_eigenvector_i - deformation += ( - pca_coefficients[i] * self.pca_std_deviations[i] * pca_eigenvector_3d - ) - - return deformation + n_modes = len(pca_coefficients) + scaled = pca_coefficients * self.pca_std_deviations[:n_modes] + deformation = np.asarray( + scaled @ self.pca_eigenvectors[:n_modes], dtype=np.float64 + ) + return deformation.reshape(-1, 3) def _optimize_pca_coefficients( self, @@ -584,12 +782,12 @@ def _optimize_pca_coefficients( ) -> tuple[np.ndarray, float]: """Optimize PCA coefficients - This method optimizes PCA mode coefficients to deform the model to better match - low values in the distance map. + Minimizes the mean distance between the deformed model and the target, + supplying the analytic gradient of the objective to the optimizer. Args: pca_number_of_modes: Number of PCA modes to use in optimization. Using fewer - modes provides smoother deformations. Default: 10 + modes provides smoother deformations. Default: 0 (use all) pca_coefficient_bounds: Bound on PCA coefficients in units of std deviations. Default: 3.0 (±3 std deviations per mode) method: Optimization method for scipy.optimize.minimize. @@ -600,52 +798,65 @@ def _optimize_pca_coefficients( Returns: Tuple of (pca_coefficients, mean_distance): - pca_coefficients: Optimized PCA coefficients - - mean_distance: Final mean distance metric value + - mean_distance: Final objective value, in mm Raises: ValueError: If number of PCA modes to use exceeds available modes """ + n_available = len(self.pca_eigenvectors) if pca_number_of_modes <= 0: - pca_number_of_modes = len(self.pca_eigenvectors) - if pca_number_of_modes > len(self.pca_eigenvectors): + pca_number_of_modes = n_available + if pca_number_of_modes > n_available: raise ValueError( - f"Number of PCA modes to use ({pca_number_of_modes}) exceeds available modes ({len(self.pca_std_deviations)})" + f"Number of PCA modes to use ({pca_number_of_modes}) exceeds " + f"available modes ({n_available})" ) self.pca_number_of_modes = pca_number_of_modes + self._prepare_sampling() + self.log_info(f"Number of PCA modes: {pca_number_of_modes}") self.log_info( f"PCA coefficient bounds: ±{pca_coefficient_bounds} std deviations" ) self.log_info(f"Optimization method: {method}") self.log_info(f"Max iterations: {max_iterations}") + self.log_info(f"Shape prior weight: {self.pca_prior_weight}") + self.log_info(f"Symmetric weight: {self.symmetric_weight}") - bounds = [] - for _ in range(pca_number_of_modes): - bounds.append((-pca_coefficient_bounds, pca_coefficient_bounds)) + bounds = [ + (-pca_coefficient_bounds, pca_coefficient_bounds) + for _ in range(pca_number_of_modes) + ] - log_level_int = ( - self.log_level - if isinstance(self.log_level, int) - else logging.getLevelName(self.log_level) - ) - disp = log_level_int <= logging.INFO + disp = self.log_level <= logging.INFO + + # The metric is in mm, so the default gradient tolerance is meaningful. + # Without an analytic gradient the finite-difference step must be large + # enough to move sample points by a useful fraction of a voxel. + options: dict = {"maxiter": max_iterations, "disp": disp, "gtol": 1e-6} + if not self._analytic_gradient: + options["eps"] = 1e-2 self.log_info("Running optimization...") result_pca = minimize( # type: ignore[call-overload] - lambda params: self._mean_distance_metric(params), - np.zeros(self.pca_number_of_modes), + self._objective_and_gradient + if self._analytic_gradient + else self._mean_distance_metric, + np.zeros(pca_number_of_modes), method=method, + jac=self._analytic_gradient, bounds=bounds, - options={"maxiter": max_iterations, "disp": disp}, + options=options, ) optimized_pca_coefficients = result_pca.x - optimized_mean_distance = result_pca.fun + optimized_mean_distance = float(result_pca.fun) self.log_info("Optimization completed!") self.log_info(f"Optimized PCA coefficients: {optimized_pca_coefficients}") - self.log_info(f"Final mean intensity: {optimized_mean_distance:.2f}") + self.log_info(f"Metric evaluations: {self._metric_call_count}") + self.log_info(f"Final mean distance: {optimized_mean_distance:.4f} mm") return optimized_pca_coefficients, optimized_mean_distance @@ -672,35 +883,12 @@ def transform_template_model(self) -> pv.DataSet: self.registered_model_pca_coefficients, ) - # Apply deformation and affine transform to each point - final_points = np.zeros((self.pca_template_model.n_points, 3), dtype=np.float64) - - n_points = self.pca_template_model.n_points - progress_interval = max(1, n_points // 10) # Report progress every 10% - - point = itk.Point[itk.D, 3]() - for i in range(n_points): - # Report progress - if i % progress_interval == 0 or i == n_points - 1: - self.log_progress(i + 1, n_points, prefix="Transforming points") - - # Start with mean shape point - point[0] = float(self.pca_template_model.points[i][0]) - point[1] = float(self.pca_template_model.points[i][1]) - point[2] = float(self.pca_template_model.points[i][2]) - - # Add PCA deformation - point[0] += self.registered_model_pca_deformation[i, 0] - point[1] += self.registered_model_pca_deformation[i, 1] - point[2] += self.registered_model_pca_deformation[i, 2] - - if self.post_pca_transform is not None: - point = self.post_pca_transform.TransformPoint(point) - - # Store result - final_points[i, 0] = point[0] - final_points[i, 1] = point[1] - final_points[i, 2] = point[2] + # Deform in the template frame, then map into target space. + deformed_points = ( + np.asarray(self.pca_template_model.points, dtype=np.float64) + + self.registered_model_pca_deformation + ) + final_points = self._apply_post_pca_transform(deformed_points) # Create new model with transformed points self.registered_model = self.pca_template_model.copy(deep=True) @@ -717,30 +905,36 @@ def transform_point( point: itk.Point, include_post_pca_transform: bool = True, ) -> itk.Point: - """Transform an arbitrary point using nearest neighbor interpolation. + """Transform an arbitrary point through the PCA deformation field. Args: point: ITK point to transform (itk.Point[itk.D, 3]) + include_post_pca_transform: Also apply post_pca_transform. Default: True Returns: Transformed ITK point + Raises: + ValueError: If compute_pca_transforms() has not been called yet + Notes: - 1) if the point is outside the image bounds, the point is not transformed. - 2) if the forward point transform is set, it is applied. - 3) if the post_pca_transform is set and enabled, it is applied. - 4) if the forward point transform is not set, no errors are raised. + This samples the *approximated* deformation field built by + compute_pca_transforms(), which is splatted and blurred, so it does + not reproduce transform_template_model() exactly; the RMS of that + difference is logged when the field is built. Points outside the + field's reference image are not displaced. Example: >>> p = itk.Point[itk.D, 3]() >>> p[0], p[1], p[2] = 10.0, 20.0, 30.0 >>> transformed_p = registrar.transform_point(p) """ - - if self.forward_point_transform is not None: - transformed_point = self.forward_point_transform.TransformPoint(point) - else: - transformed_point = point + if self.forward_point_transform is None: + self.log_error("Forward point transform is not set.") + raise ValueError( + "compute_pca_transforms() must be called before transform_point()" + ) + transformed_point = self.forward_point_transform.TransformPoint(point) if include_post_pca_transform and self.post_pca_transform is not None: transformed_point = self.post_pca_transform.TransformPoint( @@ -749,9 +943,21 @@ def transform_point( return transformed_point - def compute_pca_transforms(self, reference_image: itk.Image) -> dict: + def compute_pca_transforms( + self, reference_image: itk.Image, blur_sigma: float = 2.5 + ) -> dict: """Compute PCA transforms. + The field is built by splatting the per-point PCA displacements onto the + reference grid and blurring them, so it only approximates the exact + per-point deformation. The RMS of that approximation error, and of the + forward/inverse round trip, are both logged. + + Args: + reference_image: ITK image providing the coordinate frame for the field. + blur_sigma: Sigma for Gaussian blurring of the deformation field. + Default: 2.5 + Returns: Dictionary containing: - 'forward_point_transform': POINT transform mapping template @@ -761,17 +967,19 @@ def compute_pca_transforms(self, reference_image: itk.Image) -> dict: Note: These are point transforms, oriented opposite to image-registration - transforms; see docs/developer/transform_conventions. + transforms; see docs/developer/transform_conventions. Neither + includes post_pca_transform. """ assert self.registered_model_pca_deformation is not None, ( "PCA deformation must be computed" ) + template_points = np.asarray(self.pca_template_model.points, dtype=np.float64) template_model_pca_deformation_field_image = ( self._contour_tools.create_deformation_field( - np.array(self.pca_template_model.points), + template_points, self.registered_model_pca_deformation, reference_image=reference_image, - blur_sigma=2.5, + blur_sigma=blur_sigma, ptype=itk.D, ) ) @@ -787,11 +995,44 @@ def compute_pca_transforms(self, reference_image: itk.Image) -> dict: self.forward_point_transform ) ) + + self._log_transform_fidelity(template_points) + return { "forward_point_transform": self.forward_point_transform, "inverse_point_transform": self.inverse_point_transform, } + def _log_transform_fidelity(self, template_points: np.ndarray) -> None: + """Report how well the field reproduces the deformation and inverts.""" + assert self.forward_point_transform is not None, "forward transform must be set" + assert self.inverse_point_transform is not None, "inverse transform must be set" + assert self.registered_model_pca_deformation is not None, ( + "PCA deformation must be computed" + ) + + point = itk.Point[itk.D, 3]() + forward = np.empty_like(template_points) + round_trip = np.empty_like(template_points) + for i, source in enumerate(template_points): + point[0], point[1], point[2] = (float(v) for v in source) + mapped = self.forward_point_transform.TransformPoint(point) + forward[i] = (mapped[0], mapped[1], mapped[2]) + back = self.inverse_point_transform.TransformPoint(mapped) + round_trip[i] = (back[0], back[1], back[2]) + + expected = template_points + self.registered_model_pca_deformation + field_rms = float(np.sqrt(np.mean(np.sum((forward - expected) ** 2, axis=1)))) + inverse_rms = float( + np.sqrt(np.mean(np.sum((round_trip - template_points) ** 2, axis=1))) + ) + self.log_info( + "Deformation field RMS error: %.4f mm (approximation of the " + "per-point deformation)", + field_rms, + ) + self.log_info("Forward/inverse round-trip RMS error: %.4f mm", inverse_rms) + def register( self, pca_number_of_modes: int = 0, @@ -799,40 +1040,40 @@ def register( method: str = "L-BFGS-B", max_iterations: int = 100, ) -> dict: - """Optimize PCA coefficients to deform the model to better match - low values in the distance map. + """Optimize PCA coefficients to deform the model onto the target. Args: pca_number_of_modes: Number of PCA modes to use. Default: 0 (use all available modes) - pca_coefficient_bounds: PCA coefficient bounds (±std devs). Default: 3.0 + pca_coefficient_bounds: PCA coefficient bounds (±std devs). Default: 3.5 method: Optimization method for scipy.optimize.minimize. Default: 'L-BFGS-B' (supports bounds) max_iterations: Maximum number of optimization iterations. - Default: 50 + Default: 100 Returns: Dictionary containing: - 'registered_model': Final registered PyVista model - 'pca_coefficients': Optimized PCA coefficients - - 'mean_distance': Final mean distance metric value + - 'mean_distance': Final objective value, in mm Raises: - ValueError: If reference image is not set + ValueError: If the distance map is not set Example: >>> result = registrar.register(pca_number_of_modes=10) >>> result['registered_model'].save('registered_heart.vtk') """ if self.fixed_distance_map is None: - raise ValueError("Reference image must be set before registration") + raise ValueError("A distance map must be set before registration") if pca_number_of_modes <= 0: pca_number_of_modes = self.pca_number_of_modes - self.log_section("PCA-BASED MODEL-TO-IMAGE REGISTRATION", width=70) + self.log_section("PCA-BASED MODEL-TO-MODEL REGISTRATION", width=70) self.log_info(f"Number of points: {self.pca_template_model.n_points}") self.log_info(f"Modes to use: {pca_number_of_modes}") + self._metric_call_count = 0 self.registered_model_pca_coefficients, self.registered_model_mean_distance = ( self._optimize_pca_coefficients( pca_number_of_modes=pca_number_of_modes, diff --git a/tests/test_register_models_pca.py b/tests/test_register_models_pca.py index 66b19746..2ce2fbcd 100644 --- a/tests/test_register_models_pca.py +++ b/tests/test_register_models_pca.py @@ -8,11 +8,12 @@ import numpy as np import pytest import pyvista as pv +from scipy.optimize import approx_fprime from physiotwin4d.register_models_pca import RegisterModelsPCA -def _make_registrar() -> RegisterModelsPCA: +def _make_registrar(**kwargs: Any) -> RegisterModelsPCA: """Create a small PCA registrar with a three-point template surface.""" template_model = pv.PolyData( np.array( @@ -27,27 +28,52 @@ def _make_registrar() -> RegisterModelsPCA: pca_eigenvectors = np.zeros((1, template_model.n_points * 3), dtype=np.float64) pca_std_deviations = np.ones(1, dtype=np.float64) fixed_distance_map = itk.image_from_array(np.zeros((4, 4, 4), dtype=np.float32)) + kwargs.setdefault("symmetric_weight", 0.0) return RegisterModelsPCA( pca_template_model=template_model, pca_eigenvectors=pca_eigenvectors, pca_std_deviations=pca_std_deviations, pca_number_of_modes=1, fixed_distance_map=fixed_distance_map, + **kwargs, ) -def test_itk_template_points_are_distinct_objects() -> None: - """Cached ITK points are distinct per template vertex.""" - registrar = _make_registrar() +def _sphere_registrar( + radius: float, + modes: int = 1, + **kwargs: Any, +) -> tuple[RegisterModelsPCA, np.ndarray]: + """Build a registrar whose single mode inflates a sphere radially. + + Returns the registrar and the per-point unit-norm eigenvector it was given. + """ + template = pv.Sphere(radius=radius, theta_resolution=24, phi_resolution=24) + directions = np.asarray(template.points, dtype=np.float64) + directions /= np.linalg.norm(directions, axis=1, keepdims=True) + + eigenvectors = np.zeros((modes, template.n_points * 3), dtype=np.float64) + eigenvectors[0] = directions.reshape(-1) / np.linalg.norm(directions) + for mode in range(1, modes): + # Orthogonal filler modes: displace a disjoint slab of points along x. + filler = np.zeros((template.n_points, 3), dtype=np.float64) + filler[mode::modes, 0] = 1.0 + eigenvectors[mode] = filler.reshape(-1) / np.linalg.norm(filler) - points = registrar._pca_template_model_points_itk - assert points is not None - assert len({id(point) for point in points}) == len(points) - assert [float(points[0][0]), float(points[1][0]), float(points[2][1])] == [ - 0.0, - 1.0, - 1.0, - ] + reference_image = itk.image_from_array(np.zeros((48, 48, 48), dtype=np.float32)) + reference_image.SetSpacing([1.0, 1.0, 1.0]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + registrar = RegisterModelsPCA( + pca_template_model=template, + pca_eigenvectors=eigenvectors, + pca_std_deviations=np.full(modes, 5.0), + pca_template_model_point_subsample=1, + fixed_model=pv.Sphere(radius=radius, theta_resolution=24, phi_resolution=24), + reference_image=reference_image, + **kwargs, + ) + return registrar, eigenvectors[0].reshape(-1, 3) def test_set_fixed_model_requires_reference_image() -> None: @@ -60,6 +86,48 @@ def test_set_fixed_model_requires_reference_image() -> None: ) +def test_mode_count_mismatch_is_rejected() -> None: + """Eigenvector and standard-deviation counts must agree.""" + template_model = pv.PolyData(np.zeros((3, 3), dtype=np.float64)) + with pytest.raises(ValueError, match="Mode count mismatch"): + RegisterModelsPCA( + pca_template_model=template_model, + pca_eigenvectors=np.zeros((2, 9), dtype=np.float64), + pca_std_deviations=np.ones(3, dtype=np.float64), + fixed_distance_map=itk.image_from_array( + np.zeros((4, 4, 4), dtype=np.float32) + ), + ) + + +def test_compute_pca_deformation_scales_eigenvectors_by_std() -> None: + """Deformation is exactly sum(b_i * std_i * eigenvector_i), reshaped (N, 3).""" + template_model = pv.PolyData(np.zeros((2, 3), dtype=np.float64)) + eigenvectors = np.array( + [ + [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0, 0.0, 0.0], + ], + dtype=np.float64, + ) + std_deviations = np.array([2.0, 5.0], dtype=np.float64) + registrar = RegisterModelsPCA( + pca_template_model=template_model, + pca_eigenvectors=eigenvectors, + pca_std_deviations=std_deviations, + fixed_distance_map=itk.image_from_array(np.zeros((4, 4, 4), dtype=np.float32)), + ) + + deformation = registrar._compute_pca_deformation(np.array([1.5, -1.0])) + + # Point 0: 1.5*2*[1,0,0] + (-1)*5*[0,0,1]; point 1: 1.5*2*[0,1,0]. + assert np.allclose(deformation, [[3.0, 0.0, -5.0], [0.0, 3.0, 0.0]]) + + # A shorter coefficient vector uses only the leading modes. + leading = registrar._compute_pca_deformation(np.array([1.5])) + assert np.allclose(leading, [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0]]) + + def test_transform_template_model_applies_post_pca_transform_after_deformation() -> ( None ): @@ -78,3 +146,184 @@ def test_transform_template_model_applies_post_pca_transform_after_deformation() assert np.allclose(registered_model.points[0], [2.0, 0.0, 0.0]) assert np.allclose(registered_model.points[1], [4.0, 0.0, 0.0]) + + +def test_modes_are_deformed_in_the_template_frame_then_transformed() -> None: + """Regression: modes must be rotated with the template, not added after it. + + The registered model must equal ``A @ (template + deformation)``, never + ``A @ template + deformation``. The two differ whenever the post-PCA + transform contains a rotation, which is the case for every ICP alignment. + """ + registrar = _make_registrar() + registrar.registered_model_pca_coefficients = np.array([1.0], dtype=np.float64) + deformation = np.tile( + np.array([1.0, 0.0, 0.0], dtype=np.float64), + (registrar.pca_template_model.n_points, 1), + ) + registrar.registered_model_pca_deformation = deformation + + # 90 degrees about z, so x-displacements must come out along y. + matrix = np.array( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64 + ) + offset = np.array([10.0, -3.0, 2.0], dtype=np.float64) + transform = itk.AffineTransform[itk.D, 3].New() + transform.SetMatrix(itk.matrix_from_array(matrix)) + transform.SetTranslation(offset) + registrar.post_pca_transform = transform + + registered_model: Any = registrar.transform_template_model() + + template_points = np.asarray(registrar.pca_template_model.points, dtype=np.float64) + correct = (template_points + deformation) @ matrix.T + offset + wrong = template_points @ matrix.T + offset + deformation + + assert np.allclose(registered_model.points, correct) + assert not np.allclose(registered_model.points, wrong) + + +def test_analytic_gradient_matches_finite_differences() -> None: + """The supplied Jacobian agrees with a finite-difference gradient.""" + registrar, _ = _sphere_registrar(radius=8.0, modes=3, symmetric_weight=0.5) + registrar.pca_prior_weight = 0.1 + registrar._prepare_sampling() + + params = np.array([0.4, -0.25, 0.15], dtype=np.float64) + _, analytic = registrar._objective_and_gradient(params) + numeric = approx_fprime(params, registrar._mean_distance_metric, 1e-5) + + assert np.allclose(analytic, numeric, atol=2e-3) + + +def test_register_recovers_known_coefficients() -> None: + """Fitting to a target built from known coefficients recovers them.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + + # Target = template inflated by b = 1.0 along the single radial mode. + truth = 1.0 + target = registrar.pca_template_model.copy(deep=True) + target.points = ( + np.asarray(registrar.pca_template_model.points, dtype=np.float64) + + truth * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + registrar.set_fixed_model(cast(pv.UnstructuredGrid, target), reference_image) + + result = registrar.register(pca_number_of_modes=1, max_iterations=60) + + assert result["pca_coefficients"][0] == pytest.approx(truth, abs=0.1) + assert result["mean_distance"] < 0.5 + + +def test_symmetric_term_penalizes_partial_coverage() -> None: + """The target-to-model term sees coverage the model-to-target term misses. + + A hemisphere sitting on a full sphere scores near-perfectly one-way: every + model point lies on the target surface. Only the target-to-model term + notices that half the target has no model near it. + """ + target = pv.Sphere(radius=8.0, theta_resolution=24, phi_resolution=24) + points = np.asarray(target.points, dtype=np.float64) + hemisphere = pv.PolyData(points[points[:, 2] > 0.0]) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + registrar = RegisterModelsPCA( + pca_template_model=hemisphere, + pca_eigenvectors=np.zeros((1, hemisphere.n_points * 3), dtype=np.float64), + pca_std_deviations=np.ones(1, dtype=np.float64), + pca_template_model_point_subsample=1, + fixed_model=target, + reference_image=reference_image, + symmetric_weight=0.0, + ) + + registrar._prepare_sampling() + one_way = registrar._mean_distance_metric(np.zeros(1)) + + registrar.symmetric_weight = 0.5 + registrar._prepare_sampling() + symmetric = registrar._mean_distance_metric(np.zeros(1)) + + # Model points all lie on the target surface, so the one-way term is small. + assert one_way < 0.5 + assert symmetric > 1.0 + + +def test_prior_shrinks_coefficients() -> None: + """Raising pca_prior_weight pulls the solution toward the mean shape.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + + target = registrar.pca_template_model.copy(deep=True) + target.points = ( + np.asarray(registrar.pca_template_model.points, dtype=np.float64) + + 1.0 * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + registrar.set_fixed_model(cast(pv.UnstructuredGrid, target), reference_image) + + unregularized = registrar.register(pca_number_of_modes=1, max_iterations=60) + + registrar.pca_prior_weight = 5.0 + registrar._sampling_ready = False + regularized = registrar.register(pca_number_of_modes=1, max_iterations=60) + + assert abs(regularized["pca_coefficients"][0]) < abs( + unregularized["pca_coefficients"][0] + ) + + +def test_transform_point_requires_computed_transforms() -> None: + """transform_point raises instead of silently returning the input.""" + registrar = _make_registrar() + point = itk.Point[itk.D, 3]() + point[0], point[1], point[2] = 1.0, 2.0, 3.0 + + with pytest.raises(ValueError, match="compute_pca_transforms"): + registrar.transform_point(point) + + +def test_pca_transforms_round_trip() -> None: + """forward reproduces the deformation and inverse undoes it.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + registrar.registered_model_pca_coefficients = np.array([1.0], dtype=np.float64) + registrar.registered_model_pca_deformation = ( + 1.0 * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + transforms = registrar.compute_pca_transforms(reference_image, blur_sigma=1.5) + forward = transforms["forward_point_transform"] + inverse = transforms["inverse_point_transform"] + + template_points = np.asarray(registrar.pca_template_model.points, dtype=np.float64) + expected = template_points + registrar.registered_model_pca_deformation + + point = itk.Point[itk.D, 3]() + mapped = np.empty_like(template_points) + back = np.empty_like(template_points) + for i, source in enumerate(template_points): + point[0], point[1], point[2] = (float(v) for v in source) + forward_point = forward.TransformPoint(point) + mapped[i] = (forward_point[0], forward_point[1], forward_point[2]) + inverse_point = inverse.TransformPoint(forward_point) + back[i] = (inverse_point[0], inverse_point[1], inverse_point[2]) + + # The field is splatted and blurred, so it only approximates the deformation. + field_rms = np.sqrt(np.mean(np.sum((mapped - expected) ** 2, axis=1))) + round_trip_rms = np.sqrt(np.mean(np.sum((back - template_points) ** 2, axis=1))) + + assert field_rms < 1.5 + assert round_trip_rms < 1.0