diff --git a/src/physiotwin4d/contour_tools.py b/src/physiotwin4d/contour_tools.py index 67558d1..3d7d2c7 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 ef93d97..c2cf68d 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 66b1974..2ce2fbc 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