warpfield.register

  1import warnings
  2import gc
  3import pathlib
  4import os
  5from typing import List, Union, Callable
  6
  7import numpy as np
  8import scipy.signal
  9import cupy as cp
 10import cupyx
 11import cupyx.scipy.ndimage
 12from pydantic import BaseModel, ValidationError
 13from tqdm.auto import tqdm
 14import h5py
 15
 16from .warp import warp_volume
 17from .utils import create_rgb_video, mips_callback
 18from .ndimage import (
 19    accumarray,
 20    dogfilter,
 21    gausskernel_sheared,
 22    infill_nans,
 23    ndwindow,
 24    periodic_smooth_decomposition_nd_rfft,
 25    sliding_block,
 26    upsampled_dft_rfftn,
 27    soften_edges,
 28)
 29
 30_ArrayType = Union[np.ndarray, cp.ndarray]
 31
 32class WarpMap:
 33    """Represents a 3D displacement field
 34
 35    Args:
 36        warp_field (numpy.array): the displacement field data (3-x-y-z)
 37        block_size (3-element list or numpy.array):
 38        block_stride (3-element list or numpy.array):
 39        ref_shape (tuple): shape of the reference volume
 40        mov_shape (tuple): shape of the moving volume
 41    """
 42
 43    def __init__(self, warp_field, block_size, block_stride, ref_shape, mov_shape):
 44        self.warp_field = cp.array(warp_field, dtype="float32")
 45        self.block_size = cp.array(block_size, dtype="float32")
 46        self.block_stride = cp.array(block_stride, dtype="float32")
 47        self.ref_shape = ref_shape
 48        self.mov_shape = mov_shape
 49
 50    def warp(self, vol, out=None):
 51        """Apply the warp to a volume. Can be thought of as pulling the moving volume to the fixed volume space.
 52
 53        Args:
 54            vol (cupy.array): the volume to be warped
 55
 56        Returns:
 57            cupy.array: warped volume
 58        """
 59        if np.any(vol.shape != np.array(self.mov_shape)):
 60            warnings.warn(f"Volume shape {vol.shape} does not match the expected shape {self.mov_shape}.")
 61        if out is None:
 62            out = cp.zeros(self.ref_shape, dtype="float32", order="C")
 63        vol_out = warp_volume(
 64            vol, self.warp_field, self.block_stride, cp.array(-self.block_size / self.block_stride / 2), out=out
 65        )
 66        return vol_out
 67
 68    def apply(self, *args, **kwargs):
 69        """Alias of warp method"""
 70        return self.warp(*args, **kwargs)
 71
 72    def fit_affine(self, target=None):
 73        """Fit affine transformation and return new fitted WarpMap
 74
 75        Args:
 76            target (dict): dict with keys "blocks_shape", "block_size", and "block_stride"
 77
 78        Returns:
 79            WarpMap:
 80            numpy.array: affine tranformation coefficients
 81        """
 82        if target is None:
 83            warp_field_shape = self.warp_field.shape
 84            block_size = self.block_size
 85            block_stride = self.block_stride
 86        else:
 87            warp_field_shape = target["warp_field_shape"]
 88            block_size = cp.array(target["block_size"]).astype("float32")
 89            block_stride = cp.array(target["block_stride"]).astype("float32")
 90
 91        ix = cp.indices(self.warp_field.shape[1:]).reshape(3, -1).T
 92        ix = ix * self.block_stride + self.block_size / 2
 93        M = cp.zeros(self.warp_field.shape[1:])
 94        #M[1:-1, 1:-1, 1:-1] = 1
 95        M[:,:,:] = 1
 96        ixg = cp.where(M.flatten() > 0)[0]
 97        a = cp.hstack([ix[ixg], cp.ones((len(ixg), 1))])
 98        b = ix[ixg] + self.warp_field.reshape(3, -1).T[ixg]
 99        coeff = cp.linalg.lstsq(a, b, rcond=None)[0]
100        ix_out = cp.indices(warp_field_shape[1:]).reshape(3, -1).T * block_stride + block_size / 2
101        linfit = ((ix_out @ (coeff[:3] - cp.eye(3))) + coeff[3]).T.reshape(warp_field_shape)
102        return WarpMap(linfit, block_size, block_stride, self.ref_shape, self.mov_shape), coeff
103
104    def median_filter(self):
105        """Apply median filter to the displacement field
106
107        Returns:
108            WarpMap: new WarpMap with median filtered displacement field
109        """
110        warp_field = cupyx.scipy.ndimage.median_filter(self.warp_field, size=[1, 3, 3, 3], mode="nearest")
111        return WarpMap(warp_field, self.block_size, self.block_stride, self.ref_shape, self.mov_shape)
112
113    def resize_to(self, target):
114        """Resize to target WarpMap, using linear interpolation
115
116        Args:
117            target (WarpMap or WarpMapper): target to resize to
118                or a dict with keys "shape", "block_size", and "block_stride"
119
120        Returns:
121            WarpMap: resized WarpMap
122        """
123        if isinstance(target, WarpMap):
124            t_sh, t_bsz, t_bst = target.warp_field.shape[1:], target.block_size, target.block_stride
125        elif isinstance(target, WarpMapper):
126            t_sh, t_bsz, t_bst = target.blocks_shape[:3], cp.array(target.block_size), cp.array(target.block_stride)
127        elif isinstance(target, dict):
128            t_sh, t_bsz, t_bst = (
129                target["warp_field_shape"][1:],
130                cp.array(target["block_size"]),
131                cp.array(target["block_stride"]),
132            )
133        else:
134            raise ValueError("target must be a WarpMap, WarpMapper, or dict")
135        ix = cp.array(cp.indices(t_sh).reshape(3, -1))
136        # ix = (ix + 0.5) / cp.array(self.block_size / t_bsz)[:, None] - 0.5
137        ix = (ix * t_bst[:, None] + (t_bsz - self.block_size)[:, None] / 2) / self.block_stride[:, None]
138        dm_r = cp.array(
139            [
140                cupyx.scipy.ndimage.map_coordinates(cp.array(self.warp_field[i]), ix, mode="nearest", order=1).reshape(
141                    t_sh
142                )
143                for i in range(3)
144            ]
145        )
146        return WarpMap(dm_r, t_bsz, t_bst, self.ref_shape, self.mov_shape)
147
148    def chain(self, target):
149        """Chain displacement maps
150
151        Args:
152            target (WarpMap): WarpMap to be added to existing map
153
154        Returns:
155            WarpMap: new WarpMap with chained displacement field
156        """
157        indices = cp.indices(target.warp_field.shape[1:])
158        warp_field = self.warp_field.copy()
159        warp_field += target.warp_field
160        return WarpMap(warp_field, target.block_size, target.block_stride, self.ref_shape, self.mov_shape)
161
162    def invert(self, **kwargs):
163        """alias for invert_fast method"""
164        return self.invert_fast(**kwargs)
165
166    def invert_fast(self, sigma=0.5, truncate=20):
167        """Invert the displacement field using accumulation and Gaussian basis interpolation.
168
169        Args:
170            sigma (float): standard deviation for Gaussian basis interpolation
171            truncate (float): truncate parameter for Gaussian basis interpolation
172
173        Returns:
174            WarpMap: inverted WarpMap
175        """
176        warp_field = self.warp_field.get()
177        target_coords = np.indices(warp_field.shape[1:]) + warp_field / self.block_stride[:, None, None, None].get()
178        wf_shape = np.ceil(np.array(self.mov_shape) / self.block_stride.get() + 1).astype("int")
179        num_coords = accumarray(target_coords, wf_shape)
180        inv_field = np.zeros((3, *wf_shape), dtype=warp_field.dtype)
181        for i in range(3):
182            inv_field[i] = -accumarray(target_coords, wf_shape, weights=warp_field[i].ravel())
183            with np.errstate(invalid="ignore"):
184                inv_field[i] /= num_coords
185            inv_field[i][num_coords == 0] = np.nan
186            inv_field[i] = infill_nans(inv_field[i], sigma=sigma, truncate=truncate)
187        return WarpMap(inv_field, self.block_size, self.block_stride, self.mov_shape, self.ref_shape)
188
189    def push_coordinates(self, coords, negative_shifts=False):
190        """Push voxel coordinates from fixed to moving space.
191
192        Args:
193            coords (numpy.array): 3D *voxel* coordinates to be warped (3-by-n array)
194
195        Returns:
196            numpy.array: transformed voxel coordinates
197        """
198        assert coords.shape[0] == 3
199        coords = cp.array(coords, dtype="float32")
200        # coords_blocked = coords / self.block_size[:, None] - 0.5
201        coords_blocked = coords / self.block_stride[:, None] - (self.block_size / (2 * self.block_stride))[:, None]
202        warp_field = self.warp_field.copy()
203        shifts = cp.zeros_like(coords)
204        for idim in range(3):
205            shifts[idim] = cupyx.scipy.ndimage.map_coordinates(
206                warp_field[idim], coords_blocked, order=1, mode="nearest"
207            )
208        if negative_shifts:
209            shifts = -shifts
210        return coords + shifts
211
212    def pull_coordinates(self, coords):
213        """Pull voxel coordinates through the warp field. Involves inversion, followed by pushing coordinates.
214
215        Args:
216            coords (numpy.array): 3D *voxel* coordinates to be warped (3-by-n array)
217
218        Returns:
219            numpy.array: transformed voxel coordinates
220        """
221        return self.invert().push_coordinates(coords, negative_shifts=True)
222
223    def jacobian_det(self, units_per_voxel=[1, 1, 1], edge_order=1):
224        """
225        Compute det J = det(∇φ) for φ(x)=x+u(x), using np.indices for the identity grid.
226
227        Args:
228            edge_order : passed to np.gradient (1 or 2)
229
230        Returns:
231            detJ: cp.ndarray of shape spatial
232        """
233        scaling = cp.array(units_per_voxel, dtype="float32") * self.block_stride
234        coords = cp.indices(self.warp_field.shape[1:], dtype="float32") * scaling[:, None, None, None]
235        phi = coords + self.warp_field
236        J = cp.empty(self.warp_field.shape[1:] + (3, 3), dtype="float32")
237        for i in range(3):
238            grads = cp.gradient(phi[i], *scaling, edge_order=edge_order)
239            for j in range(3):
240                J[..., i, j] = grads[j]
241        return cp.linalg.det(J)
242
243    def as_ants_image(self, voxel_size_um=1):
244        """Convert to ANTsImage
245
246        Args:
247            voxel_size_um (scalar or array): voxel size (default is 1)
248
249        Returns:
250            ants.core.ants_image.ANTsImage:
251        """
252        try:
253            import ants
254        except ImportError:
255            raise ImportError("ANTs is not installed. Please install it using 'pip install ants'")
256
257        ants_image = ants.from_numpy(
258            self.warp_field.get().transpose(1, 2, 3, 0),
259            origin=list((self.block_size.get() - 1) / 2 * voxel_size_um),
260            spacing=list(self.block_stride.get() * voxel_size_um),
261            has_components=True,
262        )
263        return ants_image
264
265    def __repr__(self):
266        """String representation of the WarpMap object."""
267        info = (
268            f"WarpMap("
269            f"warp_field_shape={self.warp_field.shape}, "
270            f"block_size={self.block_size.get()}, "
271            f"block_stride={self.block_stride.get()}, "
272            f"transformation: {str(self.mov_shape)} --> {str(self.ref_shape)}"
273        )
274        return info
275
276    def to_h5(self, h5_path, group="warp_map", compression="gzip", overwrite=True):
277        """
278        Save this WarpMap to an HDF5 file.
279
280        Args:
281            h5_path (str or os.PathLike): Path to the HDF5 file.
282            group (str): Group path inside the HDF5 file to store the WarpMap (created if missing).
283            compression (str or None): Dataset compression (e.g., 'gzip', None).
284            overwrite (bool): If True, overwrite existing datasets/attrs inside the group.
285        """
286        with h5py.File(h5_path, "a") as f:
287            if overwrite and (group not in (None, "", "/")) and (group in f):
288                del f[group]
289            if (not overwrite) and (group in f):
290                raise ValueError(f"Group '{group}' already exists in {h5_path}. Set 'overwrite=True' to overwrite it.")
291            grp = f.require_group(group) if group not in (None, "", "/") else f
292            grp.create_dataset("warp_field", data=self.warp_field.get(), compression=compression)
293            grp.create_dataset("block_size", data=self.block_size.get())
294            grp.create_dataset("block_stride", data=self.block_stride.get())
295            grp.create_dataset("ref_shape", data=np.array(self.ref_shape, dtype="int64"))
296            grp.create_dataset("mov_shape", data=np.array(self.mov_shape, dtype="int64"))
297            grp.attrs["class"] = "WarpMap"
298
299    @classmethod
300    def from_h5(cls, h5_path, group="warp_map"):
301        """
302        Load a WarpMap from an HDF5 file.
303
304        Args:
305            h5_path (str or os.PathLike): Path to the HDF5 file.
306            group (str): Group path inside the HDF5 file where the WarpMap is stored.
307
308        Returns:
309            WarpMap: The loaded WarpMap object.
310        """
311        with h5py.File(h5_path, "r") as f:
312            grp = f[group]
313            warp_field = grp["warp_field"][:]
314            block_size = grp["block_size"][:]
315            block_stride = grp["block_stride"][:]
316            ref_shape = tuple(grp["ref_shape"][:].tolist())
317            mov_shape = tuple(grp["mov_shape"][:].tolist())
318        return cls(warp_field, block_size, block_stride, ref_shape, mov_shape)
319
320
321class WarpMapper:
322    """Class that estimates warp field using cross-correlation, based on a piece-wise rigid model.
323
324    Args:
325        ref_vol (numpy.array): The reference volume
326        block_size (3-element list or numpy.array): shape of blocks, whose rigid displacement is estimated
327        block_stride (3-element list or numpy.array): stride (usually identical to block_size)
328        proj_method (str or callable): Projection method
329    """
330
331    def __init__(
332        self, ref_vol, block_size, block_stride=None, proj_method=None, subpixel=4, epsilon=1e-6, tukey_alpha=0.5
333    ):
334        if np.any(block_size > np.array(ref_vol.shape)):
335            raise ValueError(f"Block size (currently: {block_size}) must be smaller than the volume shape ({np.array(ref_vol.shape)}).")
336        self.proj_method = proj_method
337        self.plan_rev = [None, None, None]
338        self.subpixel = subpixel
339        self.epsilon = epsilon
340        self.tukey_alpha = tukey_alpha
341        self.update_reference(ref_vol, block_size, block_stride)
342        self.ref_shape = np.array(ref_vol.shape)
343
344    def update_reference(self, ref_vol, block_size, block_stride=None):
345        ft = lambda arr: cp.fft.rfftn(arr, axes=(-2, -1))
346        block_size = np.array(block_size)
347        block_stride = block_size if block_stride is None else np.array(block_stride)
348        ref_blocks = sliding_block(cp.array(ref_vol), block_size=block_size, block_stride=block_stride)
349        self.blocks_shape = ref_blocks.shape
350        ref_blocks_proj = [self.proj_method(ref_blocks, axis=iax) for iax in [-3, -2, -1]]
351        if self.tukey_alpha < 1:
352            ref_blocks_proj = [
353                ref_blocks_proj[i]
354                * cp.array(
355                    ndwindow(
356                        [1, 1, 1, *ref_blocks_proj[i].shape[-2:]], lambda n: scipy.signal.windows.tukey(n, alpha=0.5)
357                    )
358                ).astype("float32")
359                for i in range(3)
360            ]
361        self.plan_fwd = [
362            cupyx.scipy.fft.get_fft_plan(ref_blocks_proj[i], axes=(-2, -1), value_type="R2C") for i in range(3)
363        ]
364        self.ref_blocks_proj_ft_conj = [
365            cupyx.scipy.fft.rfftn(ref_blocks_proj[i], axes=(-2, -1), plan=self.plan_fwd[i]).conj() for i in range(3)
366        ]
367        self.block_size = block_size
368        self.block_stride = block_stride
369
370    def get_displacement(self, vol, smooth_func=None):
371        """Estimate the displacement of vol with the reference volume, via piece-wise rigid cross-correlation with the pre-saved blocks.
372
373        Args:
374            vol (numpy.array): Input volume
375            smooth_func (callable): Smoothing function to be applied to the cross-correlation volume
376
377        Returns:
378            WarpMap
379        """
380        vol_blocks = sliding_block(vol, block_size=self.block_size, block_stride=self.block_stride)
381        vol_blocks_proj = [self.proj_method(vol_blocks, axis=iax) for iax in [-3, -2, -1]]
382        del vol_blocks
383
384        disp_field = []
385        for i in range(3):
386            R = (
387                cupyx.scipy.fft.rfftn(vol_blocks_proj[i], axes=(-2, -1), plan=self.plan_fwd[i])
388                * self.ref_blocks_proj_ft_conj[i]
389            )
390            if self.plan_rev[i] is None:
391                self.plan_rev[i] = cupyx.scipy.fft.get_fft_plan(R, axes=(-2, -1), value_type="C2R")
392            xcorr_proj = cp.fft.fftshift(cupyx.scipy.fft.irfftn(R, axes=(-2, -1), plan=self.plan_rev[i]), axes=(-2, -1))
393            if smooth_func is not None:
394                xcorr_proj = smooth_func(xcorr_proj, self.block_size)
395            xcorr_proj[..., xcorr_proj.shape[-2] // 2, xcorr_proj.shape[-1] // 2] += self.epsilon
396
397            max_ix = cp.array(cp.unravel_index(cp.argmax(xcorr_proj, axis=(-2, -1)), xcorr_proj.shape[-2:]))
398            max_ix = max_ix - cp.array(xcorr_proj.shape[-2:])[:, None, None, None] // 2
399            del xcorr_proj
400            i0, j0 = max_ix.reshape(2, -1)
401            shifts = upsampled_dft_rfftn(
402                R.reshape(-1, *R.shape[-2:]),
403                upsampled_region_size=int(self.subpixel * 2 + 1),
404                upsample_factor=self.subpixel,
405                axis_offsets=(i0, j0),
406            )
407            del R
408            max_sub = cp.array(cp.unravel_index(cp.argmax(shifts, axis=(-2, -1)), shifts.shape[-2:]))
409            max_sub = (
410                max_sub.reshape(max_ix.shape) - cp.array(shifts.shape[-2:])[:, None, None, None] // 2
411            ) / self.subpixel
412            del shifts
413            disp_field.append(max_ix + max_sub)
414
415        disp_field = cp.array(disp_field)
416        disp_field = (
417            cp.array(
418                [
419                    disp_field[1, 0] + disp_field[2, 0],
420                    disp_field[0, 0] + disp_field[2, 1],
421                    disp_field[0, 1] + disp_field[1, 1],
422                ]
423            ).astype("float32")
424            / 2
425        )
426        return WarpMap(disp_field, self.block_size, self.block_stride, self.ref_shape, vol.shape)
427
428
429class RegistrationPyramid:
430    """A class for performing multi-resolution registration.
431
432    Args:
433        ref_vol (numpy.array): Reference volume
434        settings (pandas.DataFrame): Settings for each level of the pyramid.
435            IMPORTANT: the block sizea in the last level cannot be larger than the block_size in any previous level.
436        reg_mask (numpy.array): Mask for registration
437        clip_thresh (float): Threshold for clipping the reference volume
438    """
439
440    def __init__(self, ref_vol, recipe, reg_mask=1):
441        recipe.model_validate(recipe.model_dump())
442        self.recipe = recipe
443        self.reg_mask = cp.array(reg_mask, dtype="float32", copy=False, order="C")
444        self.mappers = []
445        ref_vol = cp.array(ref_vol, dtype="float32", copy=False, order="C")
446        self.ref_shape = ref_vol.shape
447        if self.recipe.pre_filter is not None:
448            ref_vol = self.recipe.pre_filter(ref_vol, reg_mask=self.reg_mask)
449        self.mapper_ix = []
450        for i in range(len(recipe.levels)):
451            if recipe.levels[i].repeats < 1:
452                continue
453            block_size = np.array(recipe.levels[i].block_size)
454            tmp = np.r_[ref_vol.shape] // -block_size
455            block_size[block_size < 0] = tmp[block_size < 0]
456            if isinstance(recipe.levels[i].block_stride, (int, float)):
457                block_stride = (block_size * recipe.levels[i].block_stride).astype("int")
458            else:
459                block_stride = np.array(recipe.levels[i].block_stride)
460            self.mappers.append(
461                WarpMapper(
462                    ref_vol,
463                    block_size,
464                    block_stride=block_stride,
465                    proj_method=recipe.levels[i].project,
466                    tukey_alpha=recipe.levels[i].tukey_ref,
467                )
468            )
469            self.mapper_ix.append(i)
470        assert len(self.mappers) > 0, "At least one level of registration is required"
471
472    def register_single(self, vol, callback=None, verbose=False):
473        """Register a single volume to the reference volume.
474
475        Args:
476            vol (array_like): Volume to be registered (numpy or cupy array)
477            callback (function): Callback function to be called after each level of registration
478
479        Returns:
480            - vol (array_like): Registered volume (numpy or cupy array, depending on input)
481            - warp_map (WarpMap): Displacement field
482            - callback_output (list): List of outputs from the callback function
483        """
484        was_numpy = isinstance(vol, np.ndarray)
485        vol = cp.array(vol, "float32", copy=False, order="C")
486        offsets = (cp.array(vol.shape) - cp.array(self.ref_shape)) / 2
487        warp_map = WarpMap(offsets[:, None, None, None], cp.ones(3), cp.ones(3), self.ref_shape, vol.shape)
488        warp_map = warp_map.resize_to(self.mappers[-1])
489        callback_output = []
490        vol_tmp0 = self.recipe.pre_filter(vol, reg_mask=self.reg_mask) if self.recipe.pre_filter is not None else vol
491        vol_tmp = cp.zeros(self.ref_shape, dtype="float32", order="C")
492        warp_map.warp(vol_tmp0, out=vol_tmp)
493        min_block_stride = np.min([mapper.block_stride for mapper in self.mappers], axis=0)
494        if callback is not None:
495            callback_output.append(callback(vol_tmp))
496
497        if np.any(self.mappers[-1].block_stride > min_block_stride[0]):
498            warnings.warn(
499                "The block stride (in voxels) in the last level should not be larger than the block stride in any previous level (along any axis)."
500            )
501        for k, mapper in enumerate(tqdm(self.mappers, desc=f"Levels", disable=not verbose)):
502            for _ in tqdm(
503                range(self.recipe.levels[self.mapper_ix[k]].repeats), leave=False, desc=f"Repeats", disable=not verbose
504            ):
505                wm = mapper.get_displacement(
506                    vol_tmp, smooth_func=self.recipe.levels[self.mapper_ix[k]].smooth  # * self.reg_mask,
507                )
508                wm.warp_field *= self.recipe.levels[self.mapper_ix[k]].update_rate
509                if self.recipe.levels[self.mapper_ix[k]].median_filter:
510                    wm = wm.median_filter()
511                if self.recipe.levels[self.mapper_ix[k]].affine:
512                    if (np.array(mapper.blocks_shape[:3]) < 2).sum() > 1:
513                        raise ValueError(
514                            f"Affine fit needs at least two axes with at least 2 blocks! Volume shape: {self.ref_shape}; block size: {mapper.block_size}"
515                        )
516                    wm, _ = wm.fit_affine(
517                        target=dict(
518                            warp_field_shape=(3, *self.mappers[-1].blocks_shape[:3]),
519                            block_size=self.mappers[-1].block_size,
520                            block_stride=self.mappers[-1].block_stride,
521                        )
522                    )
523                else:
524                    wm = wm.resize_to(self.mappers[-1])
525
526                warp_map = warp_map.chain(wm)
527                warp_map.warp(vol_tmp0, out=vol_tmp)
528                if callback is not None:
529                    # callback_output.append(callback(warp_map.unwarp(vol)))
530                    callback_output.append(callback(vol_tmp))
531        warp_map.warp(vol, out=vol_tmp)
532        if was_numpy:
533            vol_tmp = vol_tmp.get()
534        return vol_tmp, warp_map, callback_output
535
536
537def register_volumes(ref, vol, recipe, reg_mask=1, callback=None, verbose=True, video_path=None, vmax=None):
538    """Register a volume to a reference volume using a registration pyramid.
539
540    Args:
541        ref (numpy.array or cupy.array): Reference volume
542        vol (numpy.array or cupy.array): Volume to be registered
543        recipe (Recipe): Registration recipe
544        reg_mask (numpy.array): Mask to be multiplied with the reference volume. Default is 1 (no mask)
545        callback (function): Callback function to be called on the volume after each iteration. Default is None.
546            Can be used to monitor and optimize registration. Example: `callback = lambda vol: vol.mean(1).get()`
547            (note that `vol` is a 3D cupy array. Use `.get()` to turn the output into a numpy array and save GPU memory).
548            Callback outputs for each registration step will be returned as a list.
549        verbose (bool): If True, show progress bars. Default is True
550        video_path (str): Save a video of the registration process, using callback outputs. The callback has to return 2D frames. Default is None.
551        vmax (float): Maximum pixel value (to scale video brightness). If none, set to 99.9 percentile of pixel values.
552
553    Returns:
554        - numpy.array or cupy.array (depending on vol input): Registered volume
555        - WarpMap: Displacement field
556        - list: List of outputs from the callback function
557    """
558    recipe.model_validate(recipe.model_dump())
559    reg = RegistrationPyramid(ref, recipe, reg_mask=reg_mask)
560    registered_vol, warp_map, cbout = reg.register_single(vol, callback=callback, verbose=verbose)
561    del reg
562    gc.collect()
563    cp.fft.config.get_plan_cache().clear()
564
565    if video_path is not None:
566        try:
567            assert cbout[0].ndim == 2, "Callback output must be a 2D array"
568            ref = callback(recipe.pre_filter(ref))
569            vmax = np.percentile(ref, 99.9).item() if vmax is None else vmax
570            create_rgb_video(video_path, ref / vmax, np.array(cbout) / vmax, fps=10)
571        except (ValueError, AssertionError) as e:
572            warnings.warn(f"Video generation failed with error: {e}")
573    return registered_vol, warp_map, cbout
574
575
576class Projector(BaseModel):
577    """A class to apply a 2D projection and filters to a volume block
578
579    Parameters:
580        max: if True, apply a max filter to the volume block. Default is True
581        normalize: if True, normalize projections by the L2 norm (to get correlations, not covariances). Default is False
582        dog: if True, apply a DoG filter to the volume block. Default is True
583        low: the lower sigma value for the DoG filter. Default is 0.5
584        high: the higher sigma value for the DoG filter. Default is 10.0
585        tukey_env: if True, apply a Tukey window to the output. Default is False
586        gauss_env: if True, apply a Gaussian window to the output. Default is False
587    """
588
589    max: bool = True
590    normalize: Union[bool, float] = False
591    dog: bool = True
592    low: Union[Union[int, float], List[Union[int, float]]] = 0.5
593    high: Union[Union[int, float], List[Union[int, float]]] = 10.0
594    periodic_smooth: bool = False
595
596    def __call__(self, vol_blocks, axis):
597        """Apply a 2D projection and filters to a volume block
598        Args:
599            vol_blocks (cupy.array): Blocked volume to be projected (6D dataset, with the first 3 dimensions being blocks and the last 3 dimensions being voxels)
600            axis (int): Axis along which to project
601        Returns:
602            cupy.array: Projected volume block (5D dataset, with the first 3 dimensions being blocks and the last 2 dimensions being 2D projections)
603        """
604        if self.max:
605            out = vol_blocks.max(axis)
606        else:
607            out = vol_blocks.mean(axis)
608        if self.periodic_smooth:
609            out = periodic_smooth_decomposition_nd_rfft(out)
610        low = np.delete(np.r_[1,1,1] * self.low, axis)
611        high = np.delete(np.r_[1,1,1] * self.high, axis)
612        if self.dog:
613            out = dogfilter(out, [0, 0, 0, *low], [0, 0, 0, *high], mode="reflect")
614        elif not np.all(np.array(self.low) == 0):
615            out = cupyx.scipy.ndimage.gaussian_filter(out, [0, 0, 0, *low], mode="reflect", truncate=5.0)
616        if self.normalize > 0:
617            out /= cp.sqrt(cp.sum(out**2, axis=(-2, -1), keepdims=True)) ** self.normalize + 1e-9
618        return out
619
620
621class Smoother(BaseModel):
622    """Smooth blocks with a Gaussian kernel
623    Args:
624        sigmas (list): [sigma0, sigma1, sigma2]. If None, no smoothing is applied.
625        truncate (float): truncate parameter for gaussian kernel. Default is 5.
626        shear (float): shear parameter for gaussian kernel. Default is None.
627        long_range_ratio (float): long range ratio for double gaussian kernel. Default is None.
628    """
629
630    sigmas: Union[float, List[float]] = [1.0, 1.0, 1.0]
631    shear: Union[float, None] = None
632    long_range_ratio: Union[float, None] = 0.05
633
634    def __call__(self, xcorr_proj, block_size=None):
635        """Apply a Gaussian filter to the cross-correlation data
636        Args:
637            xcorr_proj (cupy.array): cross-correlation data (5D array, with the first 3 dimensions being the blocks and the last 2 dimensions being the 2D projection)
638            block_size (list): shape of blocks, whose rigid displacement is estimated
639        Returns:
640            cupy.array: smoothed cross-correlation volume
641        """
642        truncate = 4.0
643        if self.sigmas is None:
644            return xcorr_proj
645        if self.shear is not None:
646            shear_blocks = self.shear * (block_size[1] / block_size[0])
647            gw = gausskernel_sheared(self.sigma[:2], shear_blocks, truncate=truncate)
648            gw = cp.array(gw[:, :, None, None, None])
649            xcorr_proj = cupyx.scipy.ndimage.convolve(xcorr_proj, gw, mode="constant")
650            xcorr_proj = cupyx.scipy.ndimage.gaussian_filter1d(
651                xcorr_proj, self.sigmas[2], axis=2, mode="constant", truncate=truncate
652            )
653        else:  # shear is None:
654            xcorr_proj = cupyx.scipy.ndimage.gaussian_filter(
655                xcorr_proj, [*self.sigmas, 0, 0], mode="constant", truncate=truncate
656            )
657        if self.long_range_ratio is not None:
658            xcorr_proj *= 1 - self.long_range_ratio
659            xcorr_proj += (
660                cupyx.scipy.ndimage.gaussian_filter(
661                    xcorr_proj, [*np.array(self.sigmas) * 5, 0, 0], mode="constant", truncate=truncate
662                )
663                * self.long_range_ratio
664            )
665        return xcorr_proj
666
667
668class RegFilter(BaseModel):
669    """A class to apply a filter to the volume before registration
670
671    Parameters:
672        clip_thresh: threshold for clipping the reference volume. Default is 0
673        dog: if True, apply a DoG filter to the volume. Default is True
674        low: the lower sigma value for the DoG filter. Default is 0.5
675        high: the higher sigma value for the DoG filter. Default is 10.0
676    """
677
678    clip_thresh: float = 0
679    dog: bool = True
680    low: float = 0.5
681    high: float = 10.0
682    soft_edge: Union[Union[int, float], List[Union[int, float]]] = 0.0
683
684    def __call__(self, vol, reg_mask=None):
685        """Apply the filter to the volume
686        Args:
687            vol (cupy or numpy array): 3D volume to be filtered
688            reg_mask (array): Mask for registration
689        Returns:
690            cupy.ndarray: Filtered volume
691        """
692        vol = cp.clip(cp.array(vol, "float32", copy=False) - self.clip_thresh, 0, None)
693        if np.any(np.array(self.soft_edge) > 0):
694            vol = soften_edges(vol, soft_edge=self.soft_edge, copy=False)
695        if reg_mask is not None:
696            vol *= cp.array(reg_mask, dtype="float32", copy=False)
697        if self.dog:
698            vol = dogfilter(vol, self.low, self.high, mode="reflect")
699        return vol
700
701
702class LevelConfig(BaseModel):
703    """Configuration for each level of the registration pyramid
704
705    Args:
706        block_size (list): shape of blocks, whose rigid displacement is estimated
707        block_stride (list): stride (usually identical to block_size)
708        repeats (int): number of iterations for this level (deisable level by setting repeats to 0)
709        smooth (Smoother or None): Smoother object
710        project (Projector, callable or None): Projector object. The callable should take a volume block and an axis as input and return a projected volume block.
711        tukey_ref (float): if not None, apply a Tukey window to the reference volume (alpha = tukey_ref). Default is 0.5
712        affine (bool): if True, apply affine transformation to the displacement field
713        median_filter (bool): if True, apply median filter to the displacement field
714        update_rate (float): update rate for the displacement field. Default is 1.0. Can be lowered to dampen oscillations.
715    """
716
717    block_size: Union[List[int]]
718    block_stride: Union[List[int], float] = 1.0
719    project: Union[Projector, Callable[[_ArrayType, int], _ArrayType]] = Projector()
720    tukey_ref: Union[float, None] = 0.5
721    smooth: Union[Smoother, None] = Smoother()
722    affine: bool = False
723    median_filter: bool = True
724    update_rate: float = 1.0
725    repeats: int = 5
726
727
728class Recipe(BaseModel):
729    """Configuration for the registration recipe. Recipe is initialized with a single affine level.
730
731    Args:
732        reg_filter (RegFilter, callable or None): Filter to be applied to the reference volume
733        levels (list): List of LevelConfig objects
734    """
735
736    pre_filter: Union[RegFilter, Callable[[_ArrayType], _ArrayType], None] = RegFilter()
737    levels: List[LevelConfig] = [
738        LevelConfig(block_size=[-1, -1, -1], repeats=3),  # translation level
739        LevelConfig(  # affine level
740            block_size=[-2, -2, -2],
741            block_stride=0.5,
742            repeats=10,
743            affine=True,
744            median_filter=False,
745            smooth=Smoother(sigmas=[0.5, 0.5, 0.5]),
746        ),
747    ]
748
749    def add_level(self, block_size, **kwargs):
750        """Add a level to the registration recipe
751
752        Args:
753            block_size (list): shape of blocks, whose rigid displacement is estimated
754            **kwargs: additional arguments for LevelConfig
755        """
756        if isinstance(block_size, (int, float)):
757            block_size = [block_size] * 3
758        if len(block_size) != 3:
759            raise ValueError("block_size must be a list of 3 integers")
760        self.levels.append(LevelConfig(block_size=block_size, **kwargs))
761
762    def insert_level(self, index, block_size, **kwargs):
763        """Insert a level to the registration recipe
764
765        Args:
766            index (int): A number specifying in which position to insert the level
767            block_size (list): shape of blocks, whose rigid displacement is estimated
768            **kwargs: additional arguments for LevelConfig
769        """
770        if isinstance(block_size, (int, float)):
771            block_size = [block_size] * 3
772        if len(block_size) != 3:
773            raise ValueError("block_size must be a list of 3 integers")
774        self.levels.insert(index, LevelConfig(block_size=block_size, **kwargs))
775
776    @classmethod
777    def from_yaml(cls, yaml_path):
778        """Load a recipe from a YAML file
779
780        Args:
781            yaml_path (str): path to the YAML file
782
783        Returns:
784            Recipe: Recipe object
785        """
786        import yaml
787
788        this_file_dir = pathlib.Path(__file__).resolve().parent
789        if os.path.isfile(yaml_path):
790            yaml_path = yaml_path
791        else:
792            yaml_path = os.path.join(this_file_dir, "recipes", yaml_path)
793
794        with open(yaml_path, "r") as f:
795            data = yaml.safe_load(f)
796
797        return cls.model_validate(data)
798
799    def to_yaml(self, yaml_path):
800        """Save the recipe to a YAML file
801
802        Args:
803            yaml_path (str): path to the YAML file
804        """
805        import yaml
806
807        with open(yaml_path, "w") as f:
808            yaml.dump(self.model_dump(), f)
809        print(f"Recipe saved to {yaml_path}")