Align Motion rigid-step conventions with Affine

[azrabano23]

Addresses #812.

Description

Motion and Affine interpreted the same Euler angles and translation as different geometric transforms. Pinning down the two parametrizations on an identity-affine test volume:

• Affine (via Spatial): right-handed rotations about the world axes, R = Rz @ Ry @ Rx, pivoting on the image center, translation in mm along RAS+, applied push-forward (translation=(8, 0, 0) moves the content +8 along the first axis).
• Motion's rigid step built its affine_grid matrix directly from the raw parameters, which meant: the components acted on the reversed grid axes ((x, y, z) = (k, j, i)), the content moved in the inverse direction (affine_grid matrices map output to input coordinates), and each translation component was normalized by the extent of the wrong axis (component 0 divided by the first-axis size but applied along the last axis). Concretely, on a (32, 48, 64) volume, translation=(8, 0, 0) moved the content by −15.75 voxels along the last axis. Rotations were also only rigid for cubic volumes: for non-cubic shapes the per-axis normalization turned them into anisotropic shears, even though the transform documents rigid-body motion.

This PR rebuilds the affine_grid matrices by conjugating the inverse of the intended voxel-space rigid transform into normalized grid coordinates (axis reversal + per-axis (size − 1) / 2 scaling, align_corners=True). Motion's degrees and translation now follow the same conventions as Affine(center="image"), interpreted in voxel space:

• degrees: right-handed Euler angles about the tensor axes (i, j, k), composed as R = R_k @ R_j @ R_i, pivoting on the image center.
• translation: voxels along (i, j, k), positive direction = increasing index.

For an image whose affine is the identity, the rigid step now matches Affine(degrees=..., translation=..., center="image") to float32 precision (the parity test asserts this on a non-cubic volume for pure translations, pure rotations about each axis, and a combined transform). Rotations are now truly rigid for non-cubic volumes, and integer-voxel translations shift by exactly that many voxels (asserted against torch.roll).

The docstrings now state the shared convention explicitly, including the one remaining, deliberate difference: Motion operates in voxel space and ignores the image affine, while Affine is world-space. For anisotropic spacings or non-RAS orientations the same numbers therefore still describe different world-space motions.

Question for the maintainers: is voxel-space the intended contract for v2's Motion (the previous docstring suggested so), or would you prefer Motion to interpret its parameters in world space using the image affine, fully matching Affine? The second option is a straightforward follow-up on top of this PR (conjugate through the image affine instead of only the normalization), at the cost of making the k-space corruption depend on image metadata. If voxel-space is the intent, this PR resolves the issue and it can be closed.

Notes on compatibility: for the default symmetric usage the sampling distribution is unchanged (axis reversal and sign flips are symmetries of the parameter ranges), but outputs for a fixed seed or for deterministic parameters change, since the previous mapping from parameters to geometry was incorrect. The degrees/translation docstrings also previously claimed a scalar samples from $\mathcal{U}(-d, d)$; they now describe the actual shared value/range/distribution convention (a scalar is deterministic), matching Spatial.

Checklist

• I have read the CONTRIBUTING docs and have a developer setup ready
• Changes are
  • Non-breaking (would not break existing functionality)
  • Breaking (would cause existing functionality to change)
• Tests added or modified to cover the changes
• In-line docstrings updated
• Documentation updated
• This pull request is ready to be reviewed

[romainVala]

Hi
thanks for the PR contribution I do think it is important.
Just out of curiosity, why do you need this coherence between affine and Motion ? do you have a specific application in mind ?

About voxel versus mm space for the motion parameter, I do not know if one is preferred. I would say it is an arbitrary choice ...

[azrabano23]

Hi Romain — thanks for taking a look!

The application is augmentation calibration for medical-imaging classifiers. I work on pelvic-MRI models where we sweep motion severity (degrees/translation) to roughly match the artifact levels we actually see in patients, and to test how robust the model is as severity increases. That whole exercise assumes the parameters mean the same geometry everywhere: before this change, a sweep over translation=(8,0,0) in Motion was actually moving content ~−16 voxels along a different axis than the same sweep in Affine, and on non-cubic volumes the "rigid" step included a shear that no real scanner motion produces — so severity curves weren't comparable across transforms, and weren't physically meaningful on their own. I've been bitten by this class of bug before (train/inference preprocessing parity), which is why the PR leans on a parity test rather than just a spot fix.

On voxel vs mm: agreed it's ultimately a convention choice. I kept Motion in voxel space to stay minimal and close to its current behavior — the important part to me was that within a chosen convention, the parameters are self-consistent and documented. If you and Fernando would rather have Motion take world-space mm like Affine, I'm happy to do that as a follow-up on top of this.

[romainVala]

thanks for the precision, I still do not fully understand why you need Affine and Motion to perform similarly since you are using them independently. (May be you want to constrain the motion, and therefore you'll need to fully understand the parameters )

your said :" and on non-cubic volumes the "rigid" step included a shear that no real scanner motion produces "

this look like a real problem also I do not see why it happend (this does not happen after your fix ?)

[azrabano23]

Hi Romain, two parts.

Why parity, if they are used independently: it is about calibration, and your parenthetical is essentially it. I sweep degrees / translation for Motion and for Affine to match the artifact severity we see in patient scans, and to plot robustness against severity. Those sweeps are only comparable, and only physically meaningful on their own, if degrees=10 means a 10° rigid rotation in both transforms. Before this change it did not for Motion, so I could not constrain the motion in a way I could reason about.

Where the shear came from: Motion builds the matrix that affine_grid consumes, and affine_grid works in normalized coordinates where every axis is mapped to [−1, 1] separately, i.e. scaled by (size − 1) / 2 per axis. The old code placed the raw rotation matrix R directly in that space. Mapped back to voxel space, the transform actually applied is S R S⁻¹ with S = diag((size − 1) / 2). For a cubic volume S is a multiple of the identity and this is still R. For a non-cubic volume it is no longer orthogonal. On a (32, 48, 64) volume, a 10° rotation about the first axis has singular values [1.053, 1.000, 0.950] instead of [1, 1, 1], and the two in-plane axes end up 95.8° apart instead of 90°: a small shear plus anisotropic scaling that grows with the aspect ratio and the angle.

Yes, the fix removes it. The PR conjugates the intended voxel-space rigid transform into normalized coordinates (S⁻¹ R S on the rotation, each translation component divided by its own axis extent), so the map applied in voxel space is exactly rigid. The parity test runs on a non-cubic volume precisely to catch this: after the fix the rigid step matches Affine(center="image") to float32 precision for rotations about each axis, and integer-voxel translations match torch.roll exactly.