Coverage for src/dictk/grid.py: 100%
55 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-09 23:57 +0000
1"""A rectangular collection of points, and batch point tracking across it."""
3from collections.abc import Sequence
4from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
5from enum import Enum
6from functools import partial
8import numpy as np
10from dictk import translation
11from dictk.correlation import WindowingMethod
12from dictk.image import PixelCoordinate, SubpixelCoordinate
15class Executor(Enum):
16 """Pool type `locate`'s `max_workers` parameter runs on.
18 - THREAD: shares memory with the caller, no pickling. Pays a small
19 scheduling cost on every task, which does not shrink as task count
20 grows.
21 - PROCESS: separate memory per worker, needs pickling. Pays a large
22 fixed cost once, spawning workers and importing dependencies in
23 each -- which then amortizes as task count grows.
25 See [Parallelization](../getting_started/parallelization.html) for
26 measured trade-offs between the two.
27 """
29 THREAD = "thread"
30 PROCESS = "process"
33def _locate_worker(
34 args: tuple[PixelCoordinate, PixelCoordinate],
35 *,
36 reference_image: np.ndarray,
37 current_image: np.ndarray,
38 kernel_margin_width: int,
39 kernel_margin_height: int,
40 search_margin_width: int,
41 search_margin_height: int,
42 windowing: WindowingMethod | None,
43) -> PixelCoordinate:
44 """One point's worth of `translation.locate`, as a single positional
45 argument -- `Executor.map` (thread or process) always calls its
46 target positionally, one item per iterable, so a keyword-only
47 signature is not an option for the function being mapped over.
48 Module-level on purpose: `ProcessPoolExecutor` needs a real,
49 importable function to hand to spawned workers, not a closure.
50 """
51 reference_point, search_center = args
52 return translation.locate(
53 reference_image=reference_image,
54 current_image=current_image,
55 reference_point=reference_point,
56 search_center=search_center,
57 kernel_margin_width=kernel_margin_width,
58 kernel_margin_height=kernel_margin_height,
59 search_margin_width=search_margin_width,
60 search_margin_height=search_margin_height,
61 windowing=windowing,
62 )
65def _locate_subpixel_worker(
66 args: tuple[PixelCoordinate, PixelCoordinate],
67 *,
68 reference_image: np.ndarray,
69 current_image: np.ndarray,
70 kernel_margin_width: int,
71 kernel_margin_height: int,
72 search_margin_width: int,
73 search_margin_height: int,
74 windowing: WindowingMethod | None,
75 upsample_factor: int,
76) -> SubpixelCoordinate:
77 """One point's worth of `translation.locate_subpixel` -- see
78 `_locate_worker`'s own docstring for why this is a module-level
79 function taking a single positional argument, not a closure."""
80 reference_point, search_center = args
81 return translation.locate_subpixel(
82 reference_image=reference_image,
83 current_image=current_image,
84 reference_point=reference_point,
85 search_center=search_center,
86 kernel_margin_width=kernel_margin_width,
87 kernel_margin_height=kernel_margin_height,
88 search_margin_width=search_margin_width,
89 search_margin_height=search_margin_height,
90 windowing=windowing,
91 upsample_factor=upsample_factor,
92 )
95def generate(
96 *,
97 origin: PixelCoordinate,
98 count_x: int,
99 count_y: int,
100 spacing_x: int,
101 spacing_y: int,
102) -> list[PixelCoordinate]:
103 """Generate a rectangular collection of points spanning x and y.
105 Points are returned in row-major order (top-left to bottom-right: all
106 of row 0 first, then row 1, and so on) -- point `i` sits at the same
107 index a later batch tracking call (see `locate`) returns its found
108 position at.
110 `count_x` and `count_y` need not be equal, and `spacing_x` and
111 `spacing_y` need not be equal -- this is a general rectangular
112 collection, not a square or uniformly-spaced one.
114 `count_x` or `count_y` of exactly `1` is allowed here -- a single
115 row or column of points is still a meaningful point-tracking grid on
116 its own. That's looser than [`elements`](#elements)'s own
117 requirement of `>= 2` along each axis, since forming even 1 finite
118 element needs 2 points per axis. The two functions check their own,
119 different preconditions independently -- `generate` doesn't know or
120 care about elements.
122 Args:
123 origin: Position of the top-left point, in the source image's
124 pixel reference frame.
125 count_x: Number of points along x. Must be >= 1.
126 count_y: Number of points along y. Must be >= 1.
127 spacing_x: Pixel spacing between adjacent points along x.
128 spacing_y: Pixel spacing between adjacent points along y.
130 Returns:
131 A list of `count_x * count_y` `PixelCoordinate`s, in row-major
132 order.
134 Raises:
135 ValueError: If `count_x` or `count_y` is less than 1.
136 """
137 if count_x < 1:
138 raise ValueError(f"count_x {count_x} must be >= 1")
139 if count_y < 1:
140 raise ValueError(f"count_y {count_y} must be >= 1")
142 return [
143 PixelCoordinate(x=origin.x + i * spacing_x, y=origin.y + j * spacing_y)
144 for j in range(count_y)
145 for i in range(count_x)
146 ]
149def elements(*, count_x: int, count_y: int) -> list[tuple[int, int, int, int]]:
150 """Q4 connectivity for the regular `count_x` x `count_y` lattice `generate` produces.
152 Each element is one of `generate`'s point grid's unit cells, its 4
153 corner nodes given as indices into that same points list -- so
154 `[points[i] for i in elements(...)[0]]` are one element's 4 corner
155 `PixelCoordinate`s, ready to hand to
156 [`dictk.element.gauss_point_green_lagrange_strains`](../api/dictk/element.html#gauss_point_green_lagrange_strains)
157 or
158 [`dictk.element.gauss_point_log_strains`](../api/dictk/element.html#gauss_point_log_strains).
160 Each 4-tuple is `(top_left, top_right, bottom_right, bottom_left)`
161 point indices -- the same $N_1$..$N_4$ corner order those functions'
162 `shape_functions` convention expects (see [Shape
163 Functions](../getting_started/finite_element_method.html#shape-functions)).
164 "Top"/"bottom" here use `generate`'s own image-pixel convention (y
165 increasing downward, origin at the top-left point) -- not the
166 math-style y-up convention `finite_element_method.md`'s figures use,
167 where the same 4 points would be called `N1`..`N4`'s "bottom-left,
168 bottom-right, top-right, top-left" instead. Only the labels differ
169 between the two: the actual point order (start at one corner, then
170 +x, then +x and +y together, then +y) is identical either way, and
171 that -- not which direction is called "up" -- is what makes it match
172 `shape_functions`' expected winding.
174 Args:
175 count_x: Number of points along x in the source `generate` call.
176 Must be >= 2 (at least 2 points make 1 element along x).
177 count_y: Number of points along y in the source `generate` call.
178 Must be >= 2.
180 Returns:
181 A list of `(count_x - 1) * (count_y - 1)` 4-tuples, in row-major
182 order (all of element row 0 first, then row 1, and so on) --
183 matching `generate`'s own point ordering.
185 Raises:
186 ValueError: If `count_x` or `count_y` is less than 2.
187 """
188 if count_x < 2:
189 raise ValueError(f"count_x {count_x} must be >= 2")
190 if count_y < 2:
191 raise ValueError(f"count_y {count_y} must be >= 2")
193 return [
194 (
195 j * count_x + i, # top_left
196 j * count_x + i + 1, # top_right
197 (j + 1) * count_x + i + 1, # bottom_right
198 (j + 1) * count_x + i, # bottom_left
199 )
200 for j in range(count_y - 1)
201 for i in range(count_x - 1)
202 ]
205def locate(
206 *,
207 reference_image: np.ndarray,
208 current_image: np.ndarray,
209 reference_points: Sequence[PixelCoordinate],
210 search_centers: Sequence[PixelCoordinate] | None = None,
211 kernel_margin_width: int,
212 kernel_margin_height: int,
213 search_margin_width: int,
214 search_margin_height: int,
215 windowing: WindowingMethod | None = None,
216 max_workers: int | None = None,
217 executor: Executor = Executor.THREAD,
218) -> list[PixelCoordinate]:
219 """Batch version of `dictk.translation.locate`: track many points at once.
221 Given a collection of `reference_points` (e.g. from `generate`), finds
222 each one's position in `current_image` by calling
223 [`dictk.translation.locate`](./translation.html#locate) once per point.
224 Returns a flat list, index-aligned with `reference_points` -- the same
225 row-major order `generate` produces, so point `i`'s found position is
226 at index `i` of the result.
228 Args:
229 reference_image: The reference (undeformed) 2D grayscale image.
230 current_image: The current (deformed) 2D grayscale image.
231 reference_points: Each point's fixed, known position, in
232 `reference_image`'s pixel reference frame.
233 search_centers: Where to center each point's search area, in
234 `current_image`'s pixel reference frame -- a guess of roughly
235 where each `reference_points` entry ended up. If `None`
236 (default), each point's own `reference_points` entry is used
237 as its own search center, the same reasonable default
238 `translation.locate` documents for the single-point case. If
239 given, must be the same length as `reference_points`.
240 kernel_margin_width: Half each kernel's width, in pixels. Must be
241 >= 1.
242 kernel_margin_height: Half each kernel's height, in pixels. Must
243 be >= 1.
244 search_margin_width: Half each search area's width, in pixels.
245 Must be greater than `kernel_margin_width`.
246 search_margin_height: Half each search area's height, in pixels.
247 Must be greater than `kernel_margin_height`.
248 windowing: Passed straight through to each per-point
249 [`dictk.translation.locate`](./translation.html#locate) call
250 -- see its own `windowing` parameter. Default `None` applies
251 no windowing to any point, matching this function's original
252 behavior exactly.
253 max_workers: If given, points are tracked concurrently across
254 this many workers instead of one at a time. Default `None`
255 stays sequential -- a plain loop, no pool, no overhead --
256 matching this function's original behavior exactly. See
257 [Parallelization](../getting_started/parallelization.html)
258 before setting this: at this book's own teaching scale
259 (small kernels and search areas), sequential outperforms
260 both pool types at any point count up to 1,000,000, measured.
261 Concurrency only pays for itself once each point's own
262 correlation is large enough, or point count is large enough
263 to amortize `Executor.PROCESS`'s fixed startup cost -- see
264 that page for the measured trade-offs.
265 executor: Which pool type `max_workers` runs on. Ignored if
266 `max_workers` is `None`. Default `Executor.THREAD`.
268 Returns:
269 Each point's location, in `current_image`'s pixel reference
270 frame, in the same order as `reference_points`.
272 Raises:
273 ValueError: If `search_centers` is given and its length doesn't
274 match `reference_points`, `max_workers` is given and less
275 than 1, or (from the underlying per-point call) if the
276 margin arguments are invalid.
277 """
278 if search_centers is None:
279 search_centers = reference_points
280 elif len(search_centers) != len(reference_points):
281 raise ValueError(
282 f"search_centers length ({len(search_centers)}) must match "
283 f"reference_points length ({len(reference_points)})"
284 )
285 if max_workers is not None and max_workers < 1:
286 raise ValueError(f"max_workers {max_workers} must be >= 1")
288 if max_workers is None:
289 return [
290 translation.locate(
291 reference_image=reference_image,
292 current_image=current_image,
293 reference_point=reference_point,
294 search_center=search_center,
295 kernel_margin_width=kernel_margin_width,
296 kernel_margin_height=kernel_margin_height,
297 search_margin_width=search_margin_width,
298 search_margin_height=search_margin_height,
299 windowing=windowing,
300 )
301 for reference_point, search_center in zip(reference_points, search_centers)
302 ]
304 worker = partial(
305 _locate_worker,
306 reference_image=reference_image,
307 current_image=current_image,
308 kernel_margin_width=kernel_margin_width,
309 kernel_margin_height=kernel_margin_height,
310 search_margin_width=search_margin_width,
311 search_margin_height=search_margin_height,
312 windowing=windowing,
313 )
314 executor_cls = (
315 ThreadPoolExecutor if executor is Executor.THREAD else ProcessPoolExecutor
316 )
317 with executor_cls(max_workers=max_workers) as pool:
318 return list(pool.map(worker, zip(reference_points, search_centers)))
321def locate_subpixel(
322 *,
323 reference_image: np.ndarray,
324 current_image: np.ndarray,
325 reference_points: Sequence[PixelCoordinate],
326 search_centers: Sequence[PixelCoordinate] | None = None,
327 kernel_margin_width: int,
328 kernel_margin_height: int,
329 search_margin_width: int,
330 search_margin_height: int,
331 windowing: WindowingMethod | None = None,
332 upsample_factor: int = 100,
333 max_workers: int | None = None,
334 executor: Executor = Executor.THREAD,
335) -> list[SubpixelCoordinate]:
336 """Batch version of `dictk.translation.locate_subpixel`: track many
337 points at once, with subpixel refinement.
339 Same structure as [`locate`](#locate) -- see its own docstring for
340 the full explanation of `search_centers`, `windowing`,
341 `max_workers`, and `executor`. The one behavioral difference beyond
342 the added `upsample_factor` parameter: each point's found position
343 is a [`SubpixelCoordinate`](../image.html#SubpixelCoordinate)
344 (generally fractional), not a `PixelCoordinate`, matching
345 `translation.locate_subpixel`'s own return type.
347 Args:
348 reference_image: The reference (undeformed) 2D grayscale image.
349 current_image: The current (deformed) 2D grayscale image.
350 reference_points: Each point's fixed, known position, in
351 `reference_image`'s pixel reference frame.
352 search_centers: Where to center each point's search area -- see
353 `locate`'s own docstring. If `None` (default), each point's
354 own `reference_points` entry is used.
355 kernel_margin_width: Half each kernel's width, in pixels. Must
356 be >= 1.
357 kernel_margin_height: Half each kernel's height, in pixels. Must
358 be >= 1.
359 search_margin_width: Half each search area's width, in pixels.
360 Must be greater than `kernel_margin_width`.
361 search_margin_height: Half each search area's height, in pixels.
362 Must be greater than `kernel_margin_height`.
363 windowing: Passed straight through to each per-point
364 [`dictk.translation.locate_subpixel`](../translation.html#locate_subpixel)
365 call. Default `None` applies no windowing.
366 upsample_factor: Passed straight through to each per-point
367 `locate_subpixel` call. Default `100`, matching that
368 function's own default. Must be >= 1.
369 max_workers: If given, points are tracked concurrently across
370 this many workers instead of one at a time -- see `locate`'s
371 own docstring for the measured trade-offs.
372 executor: Which pool type `max_workers` runs on. Ignored if
373 `max_workers` is `None`. Default `Executor.THREAD`.
375 Returns:
376 Each point's location, in `current_image`'s pixel reference
377 frame, generally fractional, in the same order as
378 `reference_points`.
380 Raises:
381 ValueError: If `search_centers` is given and its length doesn't
382 match `reference_points`, `max_workers` is given and less
383 than 1, or (from the underlying per-point call) if the
384 margin or `upsample_factor` arguments are invalid.
385 """
386 if search_centers is None:
387 search_centers = reference_points
388 elif len(search_centers) != len(reference_points):
389 raise ValueError(
390 f"search_centers length ({len(search_centers)}) must match "
391 f"reference_points length ({len(reference_points)})"
392 )
393 if max_workers is not None and max_workers < 1:
394 raise ValueError(f"max_workers {max_workers} must be >= 1")
396 if max_workers is None:
397 return [
398 translation.locate_subpixel(
399 reference_image=reference_image,
400 current_image=current_image,
401 reference_point=reference_point,
402 search_center=search_center,
403 kernel_margin_width=kernel_margin_width,
404 kernel_margin_height=kernel_margin_height,
405 search_margin_width=search_margin_width,
406 search_margin_height=search_margin_height,
407 windowing=windowing,
408 upsample_factor=upsample_factor,
409 )
410 for reference_point, search_center in zip(reference_points, search_centers)
411 ]
413 worker = partial(
414 _locate_subpixel_worker,
415 reference_image=reference_image,
416 current_image=current_image,
417 kernel_margin_width=kernel_margin_width,
418 kernel_margin_height=kernel_margin_height,
419 search_margin_width=search_margin_width,
420 search_margin_height=search_margin_height,
421 windowing=windowing,
422 upsample_factor=upsample_factor,
423 )
424 executor_cls = (
425 ThreadPoolExecutor if executor is Executor.THREAD else ProcessPoolExecutor
426 )
427 with executor_cls(max_workers=max_workers) as pool:
428 return list(pool.map(worker, zip(reference_points, search_centers)))