Coverage for src/dictk/correlation.py: 98%
82 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-09 23:57 +0000
1"""Spatial- and Fourier-domain cross-correlation criteria between a kernel and a search area."""
3from enum import Enum
5import numpy as np
8class WindowingMethod(Enum):
9 """Tapering window `window()` can apply before an FFT.
11 - HANN: tapers all the way to exactly 0 at both ends.
12 - HAMMING: stops short, around 0.08, trading a little residual
13 discontinuity for a narrower main lobe in the transformed signal.
14 """
16 HANN = "hann"
17 HAMMING = "hamming"
20def window(
21 *, arr: np.ndarray, method: WindowingMethod = WindowingMethod.HANN
22) -> np.ndarray:
23 r"""Taper `arr`'s edges toward zero with a 2D Hann or Hamming window.
25 An FFT implicitly treats an array as one period of an
26 infinitely-repeating signal. If the content doesn't tile seamlessly --
27 the general case, since nothing arranges `arr`'s edges to match up --
28 that discontinuity leaks energy across many frequencies rather than the
29 few the underlying content actually has, an effect called **spectral
30 leakage**. In a correlation surface, leakage broadens and can shift the
31 peak.
33 This counters that by tapering `arr`'s edges toward zero before it's
34 transformed, so the (still discontinuous, but now near-zero) seam
35 contributes far less energy. The 2D window is the outer product of a 1D
36 window with itself along each axis:
38 $$w_{\mathrm{Hann}}(n) = 0.5 \left(1 - \cos\left(\frac{2\pi n}{N - 1}\right)\right)$$
40 $$w_{\mathrm{Hamming}}(n) = 0.54 - 0.46 \cos\left(\frac{2\pi n}{N - 1}\right)$$
42 for $n = 0, \ldots, N-1$ across a window of length $N$.
44 See Harris FJ. "[On the use of windows for harmonic analysis with
45 the discrete Fourier
46 transform](https://www.cs.cmu.edu/afs/cs/user/bhiksha/WWW/courses/dsp/spring2013/WWW/schedule/readings/windows_comparison2_harris.pdf)."
47 *Proceedings of the IEEE* 1978;66(1):51-83. A U.S. government work,
48 not protected by U.S. copyright.
50 Args:
51 arr: A 2D array to window.
52 method: Which window to apply. Default `WindowingMethod.HANN`.
54 Returns:
55 A 2D float64 array the same shape as `arr`, with `arr` multiplied
56 elementwise by the 2D window.
58 Raises:
59 ValueError: If `arr` is not 2D.
60 """
61 if arr.ndim != 2:
62 raise ValueError(f"arr must be 2D, got shape {arr.shape}")
64 match method:
65 case WindowingMethod.HANN:
66 win_func = np.hanning
67 case WindowingMethod.HAMMING:
68 win_func = np.hamming
69 case _:
70 raise ValueError(f"Unsupported windowing method: {method}")
72 rows, cols = arr.shape
73 window_2d = np.outer(win_func(rows), win_func(cols))
74 return arr.astype(np.float64) * window_2d
77def _prepare(
78 *, kernel: np.ndarray, search: np.ndarray
79) -> tuple[np.ndarray, np.ndarray]:
80 """Validate `kernel`/`search` and cast both to float64.
82 Args:
83 kernel: The fixed template subimage.
84 search: The larger subimage to slide `kernel` across.
86 Returns:
87 `(kernel, search)`, both cast to float64.
89 Raises:
90 ValueError: If either array is not 2D, or `search` is smaller than
91 `kernel` in either dimension.
92 """
93 if kernel.ndim != 2:
94 raise ValueError(f"kernel must be 2D, got shape {kernel.shape}")
95 if search.ndim != 2:
96 raise ValueError(f"search must be 2D, got shape {search.shape}")
97 if search.shape[0] < kernel.shape[0] or search.shape[1] < kernel.shape[1]:
98 raise ValueError(
99 f"search shape {search.shape} must be >= kernel shape {kernel.shape} "
100 "in both dimensions"
101 )
102 return kernel.astype(np.float64), search.astype(np.float64)
105def _safe_divide(*, numerator: np.ndarray, denominator: np.ndarray) -> np.ndarray:
106 """Elementwise `numerator / denominator`, substituting 0 wherever `denominator` is 0.
108 Args:
109 numerator: Values to divide.
110 denominator: Values to divide by; must be broadcastable against
111 `numerator` and non-negative (as with a square root of an
112 energy term).
114 Returns:
115 `numerator / denominator`, with 0 wherever `denominator` is 0
116 (avoiding a division-by-zero warning, not just its result).
117 """
118 safe_denominator = np.where(denominator == 0, 1.0, denominator)
119 return np.where(denominator == 0, 0.0, numerator / safe_denominator)
122def _windows(*, search: np.ndarray, kernel_shape: tuple[int, int]) -> np.ndarray:
123 """Return every `kernel_shape`-sized window of `search`, as one strided view.
125 Args:
126 search: The 2D array to slide a window across.
127 kernel_shape: The `(height, width)` of each window.
129 Returns:
130 A 4D array of shape `(out_height, out_width, *kernel_shape)`, where
131 `out_height = search.shape[0] - kernel_shape[0] + 1` and likewise for
132 `out_width`. Entry `[dy, dx]` is the `kernel_shape`-sized window of
133 `search` with its own top-left corner at offset `(dx, dy)`.
134 """
135 return np.lib.stride_tricks.sliding_window_view(search, kernel_shape)
138def cc(*, kernel: np.ndarray, search: np.ndarray) -> np.ndarray:
139 r"""Cross-correlation (CC) surface of `kernel` slid over `search`.
141 At every valid position, computes $C_{\mathrm{CC}} = \sum f_i g_i$, where
142 $f$ is `kernel` and $g$ is the same-sized window of `search` at that
143 position. Robust to neither brightness nor contrast differences between
144 `kernel` and `search` — a uniform offset or scaling of either changes
145 every value.
147 See Pan B, Xie H, Wang Z. "[Equivalence of digital image correlation
148 criteria for pattern
149 matching](https://opg.optica.org/ao/viewmedia.cfm?uri=ao-49-28-5501)."
150 *Applied Optics* 2010;49(28):5501-9.
152 Args:
153 kernel: The fixed template subimage (`f`).
154 search: The larger subimage to slide `kernel` across (`g`'s source).
156 Returns:
157 A 2D float64 array of shape
158 `(search.shape[0] - kernel.shape[0] + 1, search.shape[1] - kernel.shape[1] + 1)`.
159 Entry `[dy, dx]` is $C_{\mathrm{CC}}$ with `kernel`'s top-left corner at
160 offset `(dx, dy)` in `search`'s local frame.
162 Raises:
163 ValueError: If either array is not 2D, or `search` is smaller than
164 `kernel` in either dimension.
165 """
166 kernel, search = _prepare(kernel=kernel, search=search)
167 windows = _windows(search=search, kernel_shape=kernel.shape)
168 return (windows * kernel).sum(axis=(-2, -1))
171def ncc(*, kernel: np.ndarray, search: np.ndarray) -> np.ndarray:
172 r"""Normalized cross-correlation (NCC) surface of `kernel` slid over `search`.
174 At every valid position, computes
175 $C_{\mathrm{NCC}} = \sum f_i g_i \,/\, \sqrt{\sum f_i^2 \sum g_i^2}$, where
176 $f$ is `kernel` and $g$ is the same-sized window of `search` at that
177 position. Robust to a uniform contrast (multiplicative) difference
178 between `kernel` and `search`, since scaling either side by a positive
179 constant cancels between the numerator and denominator. Not robust to
180 brightness (additive) differences. A window with zero energy (e.g. a
181 flat, constant-valued region) contributes a value of 0 rather than
182 raising a division-by-zero error.
184 See Pan B, Xie H, Wang Z. "[Equivalence of digital image correlation
185 criteria for pattern
186 matching](https://opg.optica.org/ao/viewmedia.cfm?uri=ao-49-28-5501)."
187 *Applied Optics* 2010;49(28):5501-9.
189 Args:
190 kernel: The fixed template subimage (`f`).
191 search: The larger subimage to slide `kernel` across (`g`'s source).
193 Returns:
194 A 2D float64 array of shape
195 `(search.shape[0] - kernel.shape[0] + 1, search.shape[1] - kernel.shape[1] + 1)`.
196 Entry `[dy, dx]` is $C_{\mathrm{NCC}}$ with `kernel`'s top-left corner at
197 offset `(dx, dy)` in `search`'s local frame.
199 Raises:
200 ValueError: If either array is not 2D, or `search` is smaller than
201 `kernel` in either dimension.
202 """
203 kernel, search = _prepare(kernel=kernel, search=search)
204 windows = _windows(search=search, kernel_shape=kernel.shape)
205 numerator = (windows * kernel).sum(axis=(-2, -1))
206 kernel_energy = np.sum(kernel**2)
207 window_energy = (windows**2).sum(axis=(-2, -1))
208 denominator = np.sqrt(kernel_energy * window_energy)
209 return _safe_divide(numerator=numerator, denominator=denominator)
212def zcc(*, kernel: np.ndarray, search: np.ndarray) -> np.ndarray:
213 r"""Zero-mean cross-correlation (ZCC) surface of `kernel` slid over `search`.
215 At every valid position, computes
216 $C_{\mathrm{ZCC}} = \sum (f_i - \bar{f})(g_i - \bar{g})$, where $f$ is
217 `kernel`, $g$ is the same-sized window of `search` at that position,
218 $\bar{f}$ is `kernel`'s own mean (fixed across all positions, since the
219 kernel never moves), and $\bar{g}$ is that window's own local mean
220 (recomputed at every position, not a global `search` statistic). Robust
221 to a uniform brightness (additive) difference between `kernel` and
222 `search`, since subtracting each side's own local mean cancels any
223 constant added to that side. Not robust to contrast (multiplicative)
224 differences.
226 See Pan B, Xie H, Wang Z. "[Equivalence of digital image correlation
227 criteria for pattern
228 matching](https://opg.optica.org/ao/viewmedia.cfm?uri=ao-49-28-5501)."
229 *Applied Optics* 2010;49(28):5501-9.
231 Args:
232 kernel: The fixed template subimage (`f`).
233 search: The larger subimage to slide `kernel` across (`g`'s source).
235 Returns:
236 A 2D float64 array of shape
237 `(search.shape[0] - kernel.shape[0] + 1, search.shape[1] - kernel.shape[1] + 1)`.
238 Entry `[dy, dx]` is $C_{\mathrm{ZCC}}$ with `kernel`'s top-left corner at
239 offset `(dx, dy)` in `search`'s local frame.
241 Raises:
242 ValueError: If either array is not 2D, or `search` is smaller than
243 `kernel` in either dimension.
244 """
245 kernel, search = _prepare(kernel=kernel, search=search)
246 windows = _windows(search=search, kernel_shape=kernel.shape)
247 kernel_centered = kernel - kernel.mean()
248 windows_centered = windows - windows.mean(axis=(-2, -1), keepdims=True)
249 return (windows_centered * kernel_centered).sum(axis=(-2, -1))
252def zncc(*, kernel: np.ndarray, search: np.ndarray) -> np.ndarray:
253 r"""Zero-mean normalized cross-correlation (ZNCC) surface of `kernel` slid over `search`.
255 At every valid position, computes
256 $C_{\mathrm{ZNCC}} = \sum \bar{f}_i \bar{g}_i \,/\, \sqrt{\sum \bar{f}_i^2 \sum \bar{g}_i^2}$,
257 where $\bar{f}_i = f_i - \bar{f}$ and $\bar{g}_i = g_i - \bar{g}$ ($f$ =
258 `kernel`, $g$ = the same-sized window of `search` at that position,
259 $\bar{f}$/$\bar{g}$ their respective means -- $\bar{f}$ fixed, $\bar{g}$
260 recomputed locally per position, as in `zcc`). Robust to both brightness
261 (additive) and contrast (multiplicative) differences between `kernel`
262 and `search`, combining `zcc`'s brightness invariance with `ncc`'s
263 contrast invariance. A window with zero variance (e.g. a flat,
264 constant-valued region) contributes a value of 0 rather than raising a
265 division-by-zero error.
267 See Pan B, Xie H, Wang Z. "[Equivalence of digital image correlation
268 criteria for pattern
269 matching](https://opg.optica.org/ao/viewmedia.cfm?uri=ao-49-28-5501)."
270 *Applied Optics* 2010;49(28):5501-9.
272 Args:
273 kernel: The fixed template subimage (`f`).
274 search: The larger subimage to slide `kernel` across (`g`'s source).
276 Returns:
277 A 2D float64 array of shape
278 `(search.shape[0] - kernel.shape[0] + 1, search.shape[1] - kernel.shape[1] + 1)`.
279 Entry `[dy, dx]` is $C_{\mathrm{ZNCC}}$ with `kernel`'s top-left corner at
280 offset `(dx, dy)` in `search`'s local frame.
282 Raises:
283 ValueError: If either array is not 2D, or `search` is smaller than
284 `kernel` in either dimension.
285 """
286 kernel, search = _prepare(kernel=kernel, search=search)
287 windows = _windows(search=search, kernel_shape=kernel.shape)
288 kernel_centered = kernel - kernel.mean()
289 windows_centered = windows - windows.mean(axis=(-2, -1), keepdims=True)
290 numerator = (windows_centered * kernel_centered).sum(axis=(-2, -1))
291 kernel_energy = np.sum(kernel_centered**2)
292 window_energy = (windows_centered**2).sum(axis=(-2, -1))
293 denominator = np.sqrt(kernel_energy * window_energy)
294 return _safe_divide(numerator=numerator, denominator=denominator)
297def _window(
298 *,
299 kernel: np.ndarray,
300 search: np.ndarray,
301 windowing: WindowingMethod | None,
302) -> tuple[np.ndarray, np.ndarray]:
303 """Optionally taper `kernel` and `search` toward zero at their own edges.
305 Shared by `phase_correlation` and
306 [`dictk.translation.locate`](../translation.html#locate) -- the two
307 functions that compare `kernel` against `search` via an FFT-based
308 technique, where windowing (if used at all) must happen before that
309 comparison, not after.
311 Args:
312 kernel: The fixed template subimage.
313 search: The larger subimage `kernel` is compared against.
314 windowing: If given, both `kernel` and `search` are passed through
315 `window()` with this method. `None` leaves both untouched --
316 including their dtype, so a caller that never windows sees no
317 incidental cast either.
319 Returns:
320 `(kernel, search)`, each windowed independently (or unchanged, if
321 `windowing` is `None`).
322 """
323 if windowing is not None:
324 kernel = window(arr=kernel, method=windowing)
325 search = window(arr=search, method=windowing)
326 return kernel, search
329def _kernel_pad(
330 *,
331 kernel: np.ndarray,
332 shape: tuple[int, int],
333 centered: bool = False,
334) -> tuple[np.ndarray, int, int]:
335 """Zero-pad `kernel` up to `shape`.
337 Only ever needs `search`'s *shape*, not `search` itself -- unlike
338 `_window`, which needs the actual array to taper it, padding `kernel`
339 only ever reads how big to grow it. Called on `_window`'s own output,
340 when both are used together, so windowing always happens first: pad
341 then window would taper the zero-padding along with `kernel`'s real
342 content, not just the content itself.
344 Args:
345 kernel: The fixed template subimage, before padding.
346 shape: The `(height, width)` to pad `kernel` up to -- typically
347 `search.shape`.
348 centered: If `False` (default), all padding goes after `kernel`'s
349 own content, which stays anchored at the padded array's
350 top-left corner -- `phase_correlation` relies on this exact
351 placement for the surfaces it publishes throughout
352 Correlation Visualization, so changing this default would
353 silently shift every peak position already documented there.
354 If `True`, padding is split before/after instead (as evenly
355 as possible), centering `kernel`'s content within the padded
356 array -- what `translation.locate` needs so FFT phase
357 correlation recovers a displacement symmetrically in both
358 directions, not just up to `kernel_margin_width`/
359 `kernel_margin_height` past `search_center` in the positive
360 direction. See [Recoverable Displacement
361 Range](../getting_started/recoverable_displacement_range.html)
362 for why.
364 Returns:
365 `(kernel_padded, pad_before_height, pad_before_width)` --
366 `kernel_padded` is `shape`-shaped, and the padding actually added
367 before `kernel`'s own content in each axis (always `(0, 0)` when
368 `centered=False`) -- a caller doing its own offset arithmetic on
369 `kernel_padded`'s content needs this to know where that content
370 actually sits.
371 """
372 pad_height = shape[0] - kernel.shape[0]
373 pad_width = shape[1] - kernel.shape[1]
374 if centered:
375 before_height, before_width = pad_height // 2, pad_width // 2
376 else:
377 before_height, before_width = 0, 0
378 kernel_padded = np.pad(
379 kernel,
380 (
381 (before_height, pad_height - before_height),
382 (before_width, pad_width - before_width),
383 ),
384 )
385 return kernel_padded, before_height, before_width
388def phase_correlation(
389 *,
390 kernel: np.ndarray,
391 search: np.ndarray,
392 windowing: WindowingMethod | None = None,
393 centered: bool = False,
394) -> np.ndarray:
395 r"""Phase correlation surface of `kernel` against `search`, via FFT.
397 Unlike `cc`/`ncc`/`zcc`/`zncc`, which slide `kernel` over `search` one
398 valid window at a time, this computes the same kind of answer all at
399 once in the Fourier domain: `kernel` is zero-padded (bottom and right)
400 up to `search`'s own shape, then
402 $$
403 C_{\mathrm{phase}} = \mathcal{F}^{-1}\left(\frac{\mathcal{F}(g)\,
404 \overline{\mathcal{F}(f)}}{\left|\mathcal{F}(g)\,\overline{\mathcal{F}(f)}\right|}\right)
405 $$
407 where $f$ is the zero-padded `kernel`, $g$ is `search`, and
408 $\mathcal{F}$ is the 2D discrete Fourier transform. Dividing by the
409 cross-power spectrum's own magnitude at every frequency -- rather than
410 summing raw products like `cc` does -- is the classic Kuglin-Hines
411 *phase correlation* technique, and is robust to both brightness
412 (additive) and contrast (multiplicative) differences between `kernel`
413 and `search`, the same pair of invariances `zncc` has, though by a
414 completely different mechanism: a brightness shift only touches the
415 zero-frequency (DC) term, leaving every other frequency -- and thus the
416 peak's position -- untouched, while dividing by magnitude at every
417 frequency cancels any overall contrast scaling directly. This is *not*
418 a Fourier-domain equivalent of `zncc`'s formula -- `zncc` recomputes a
419 local mean/variance at every candidate window as it slides; this
420 normalizes once, globally, per frequency, over the whole padded
421 extent -- it just lands in the same "robust to both" category.
423 This is exactly what [`dictk.translation.locate`](../translation.html#locate)
424 computes internally via `skimage.registration.phase_cross_correlation`
425 (`normalization="phase"`), reproduced here to expose the full surface
426 for visualization -- `phase_cross_correlation` itself only returns the
427 final shift, not the array it was computed from. The two aren't
428 directly comparable value-for-value, though: `locate` centers `kernel`
429 within its own zero-padded copy before this same FFT step, while this
430 function -- for backward compatibility with every peak position
431 already published in Correlation Visualization -- leaves `kernel`'s
432 content anchored at the padded array's top-left corner instead. Its
433 raw `argmax` is always in `[0, search.shape)`, matching the same
434 offset-within-`search` convention `cc`/`ncc`/`zcc`/`zncc` use. For the
435 small, comfortably-within-bounds displacements this book's examples
436 use, the two still agree once each is interpreted in its own
437 convention -- see [Recoverable Displacement
438 Range](../getting_started/recoverable_displacement_range.html) for why the
439 conventions diverge once a displacement isn't small.
441 See Kuglin CD, Hines DC. "The phase correlation image alignment
442 method." Proceedings of IEEE International Conference on Cybernetics
443 and Society, 1975:163-165.
445 Args:
446 kernel: The fixed template subimage (`f`, before padding).
447 search: The larger subimage `kernel` is compared against (`g`).
448 windowing: If given, both `kernel` and `search` are passed through
449 `window()` with this method -- tapering their edges toward
450 zero to reduce spectral leakage -- before padding/FFT. `kernel`
451 is windowed first, then zero-padded, so the padding stays
452 outside the tapered region. Default `None` applies no
453 windowing, matching this function's original behavior exactly.
454 centered: Passed straight through to
455 [`_kernel_pad`](#_kernel_pad)'s own `centered` parameter.
456 Default `False` keeps this function's permanent, original
457 bottom-right-only padding -- backward compatible with every
458 peak position already published in Correlation Visualization,
459 as described above -- regardless of what `dictk.translation.locate`
460 does internally. `True` centers `kernel`'s content instead,
461 matching `locate`'s own convention exactly, for a caller that
462 explicitly wants this function's surface to agree with
463 `locate`'s answer past the range where the two conventions
464 diverge (see [Recoverable Displacement
465 Range](../getting_started/recoverable_displacement_range.html)).
466 The raw `argmax` convention documented below only holds for
467 the default; with `centered=True`, the true match position is
468 the raw `argmax` plus `_kernel_pad`'s own returned
469 `pad_before_height`/`pad_before_width`, wrapped modulo
470 `search`'s own shape.
472 Returns:
473 A 2D float64 array the same shape as `search` (unlike
474 `cc`/`ncc`/`zcc`/`zncc`'s smaller "valid" shape, since nothing
475 here excludes any candidate offset). Entry `[dy, dx]` is
476 $C_{\mathrm{phase}}$ with `kernel`'s top-left corner at offset
477 `(dx, dy)` in `search`'s local frame (`centered=False`) -- see
478 `centered` above for the `centered=True` convention instead.
480 Raises:
481 ValueError: If either array is not 2D, or `search` is smaller than
482 `kernel` in either dimension.
483 """
484 kernel, search = _prepare(kernel=kernel, search=search)
485 kernel, search = _window(kernel=kernel, search=search, windowing=windowing)
486 kernel_padded, _pad_before_height, _pad_before_width = _kernel_pad(
487 kernel=kernel, shape=search.shape, centered=centered
488 )
490 search_freq = np.fft.fft2(search)
491 kernel_freq = np.fft.fft2(kernel_padded)
492 image_product = search_freq * kernel_freq.conj()
493 eps = np.finfo(image_product.real.dtype).eps
494 image_product /= np.maximum(np.abs(image_product), 100 * eps)
495 return np.fft.ifft2(image_product).real