restricted deformations + refactoring

[rohitrango]

Addresses #40

[rohitrango]

@claude review please

[claude]

Claude finished @rohitrango's task in 2m 10s —— View job

---

Code Review

• Gather context (PR diff, related files)
• Understand the request
• Review changed files for bugs/quality issues
• Post review feedback

Reviewed the diff between fed6f31~1 and 38866bc (which matches the listed PR changes). Overall this is a clean refactor + feature addition (gradient/deformation restriction along axes, partial warp evaluation, downsample_size dedup). A few notes:

Nits / possible follow-ups

1. fireants/registration/helpers.py:39-41 — downsample_size: when scale <= 1 it now returns the size unchanged, whereas the old inline code (max(int(s/scale), MIN_IMG_SIZE)) would still apply for any scale, including scale < 1 (which would upsample). For scale == 1 behavior is unchanged, but if any caller ever passes scale < 1 the new helper silently stops upsampling. Worth a quick check that scale < 1 is never used in this codebase (the docstring assumption scale > 1 to downsample seems to hold based on _assert_check_scales_decreasing, but it doesn't strictly enforce >= 1). Low risk, just flagging the behavior change.

2. fireants/registration/optimizers/adam.py — gradient_restriction is applied twice in the diffeomorphic (freeform=False) branch (once right after the Adam update, once again after composing + smoothing), but only once in WarpSGD.step() (sgd.py) and only once in the freeform=True branch of WarpAdam. This is harmless today (zeroed channels stay zero through composition/smoothing since neither operation mixes the displacement-axis channel), but the asymmetry between Adam/SGD paths looks like it might be defensive copy-paste rather than intentional. Consider documenting why the second call is needed in Adam (or removing it for consistency) so future readers don't have to re-derive that it's a no-op.

3. fireants/registration/deformablemixin.py:58-79 — get_partial_warped_parameters only scales the grid (displacement) component by fraction, leaving affine untouched. That seems intentional (affine is the rigid/linear part, grid is the deformable part being partially applied), but it's not explicitly called out in the docstring — worth a one-line clarification since "partial warp" could be misread as scaling the whole transform.

4. fireants/scripts/restricted_deformations/restricted_deformations.py — this is a standalone example/utility script (not part of the library) with commented-out affine block and a hardcoded notepad/data style path layout; no issues for merging, but consider whether it should live under examples/ or scripts/ with a short README given it's the motivating use case for this PR (Deformation Restriction #40).

Looks good

• downsample_size helper cleanly removes duplicated max(int(s/scale), MIN_IMG_SIZE) logic across affine.py, greedy.py, rigid.py, syn.py.
• restrict_deformations plumbing through optimizer_params → WarpAdam/WarpSGD is consistent with existing patterns (e.g. momentum).
• imageutils.py FFT-downsample guard (if not all(x <= y ...)) correctly prevents using the FFT path for upsampling, where it isn't valid.
• evaluate/evaluate_inverse now accept precomputed moved_coords/fixed_moved_coords, which is exactly what's needed to support get_partial_warped_parameters without duplicating the warp computation — good reuse.

No blocking issues found; the changes look safe to merge as-is. Items 1–3 are minor, worth a quick look but not required for this PR.