Features/1900 image transformations beyond fft #2048

markak47 · 2025-11-28T09:25:23Z

DUE DILIGENCE

General

The title of the PR is suitable to appear in the Release Notes.

Implementation

Unit tests cover all split configurations
(split=None and all valid split axes for 2D and 3D).
Unit tests are currently executed with float32.
MPS testing (1 MPI process, 1 GPU) was not performed locally.
No dedicated benchmarks were added for this functionality.
Performance is maintained compared to existing ndimage operations.
Documentation and example usage were added where needed.

DESCRIPTION

This PR adds support for N-dimensional affine transformations to the
heat.ndimage module.

The new implementation supports:

2D and 3D affine transformations (2×3 and 3×4 matrices)
backward-warp semantics consistent with SciPy and PyTorch
nearest-neighbor and bilinear interpolation
multiple padding modes (constant, nearest, wrap, reflect)
correct behavior for distributed arrays across all valid split configurations

Example scripts are included to demonstrate typical usage on image and
volume data.

Issues resolved: #1900

CHANGES PROPOSED

Added affine_transform function to heat.ndimage
Implemented helper routines for input normalization, grid construction,
and sampling
Added comprehensive unit tests covering interpolation modes, padding
modes, and split configurations
Added an example demonstrating affine transformations on MRI data

TYPE OF CHANGE

New feature (non-breaking change which adds functionality)

MEMORY REQUIREMENTS

No significant additional memory overhead beyond temporary coordinate
grids required for affine sampling. Memory usage is comparable to existing
ndimage operations in both distributed and non-distributed modes.

PERFORMANCE

Performance is comparable to existing ndimage operations. No performance
regression was observed for standard use cases.

DOES THIS CHANGE MODIFY THE BEHAVIOUR OF OTHER FUNCTIONS?

no

ND)

ClaudiaComito · 2025-12-02T10:13:22Z

@markak47 thank you so much for the work and examples, this looks very nice already.

Before we go on with optimizing the implementation, a few tips on integrating this into the Heat codebase:

unless @lolacaro disagrees, we should stick to the scipy scheme, i.e. the module should be called ndimage, not transform. The function's API should adhere to scipy.ndimage.affine_transform apart from potential arguments related to distribution
check out our other modules to see what the structure of the files and dirs should be like, e.g. signal within core, or fft
the examples are great! if you haven't already, please incorporate them in the tests
datasets that are required for the tests should go into the heat/datasets directory.
please 🙏🏼 🙏🏼 🙏🏼 double-check that the test datasets can be shared publicly, if not remove them asap THANKS! 😬
other files that don't serve the module's purposes can be removed too

This is all a bit boring, but if we leave it for later, it'll be much more work.

Thank you so much Mark!

…ary files module name change

…sary files, module name change

Copilot

Pull request overview

This PR adds new affine transformation functionality to the Heat ndimage module, enabling 2D and 3D image transformations including rotation, scaling, translation, and shearing. The implementation uses backward warping with support for different interpolation methods (nearest neighbor and bilinear) and multiple padding modes (constant, nearest, wrap, reflect).

Key Changes:

New affine_transform function supporting 2D (2x3) and 3D (3x4) transformation matrices with backward warping semantics
Nearest neighbor and bilinear interpolation modes with customizable boundary handling
Comprehensive test suite validating transformations, interpolation, and padding behaviors

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 17 comments.

File	Description
heat/ndimage/affine.py	Core implementation of affine transformation with coordinate mapping, interpolation methods, and padding modes
heat/ndimage/tests/test_affine_nd.py	Unit tests covering 2D/3D transformations, identity, translation, rotation, scaling, interpolation, and padding modes
heat/ndimage/examples/mri_testscan.py	Example demonstrating affine transformations on MRI data with rotation, scaling, translation, and shearing

Comments suppressed due to low confidence (2)

heat/ndimage/examples/mri_testscan.py:79

Variable cx is not used.

    cx, cy = W/2, H/2

heat/ndimage/examples/mri_testscan.py:79

Variable cy is not used.

    cx, cy = W/2, H/2

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-12T09:25:14Z

heat/ndimage/affine.py

+    mode="nearest",
+    constant_value=0.0,
+    expand=False,  # kept for later; currently no-op
+):


The order parameter is not validated. Invalid values (e.g., order=5) will cause the function to silently use bilinear interpolation via the else branch at line 235. Add validation to ensure order is either 0 (nearest) or 1 (bilinear), and raise a ValueError for unsupported interpolation orders.

Suggested change

):

):

if order not in (0, 1):

raise ValueError(f"Unsupported interpolation order: {order}. Only 0 (nearest) and 1 (bilinear) are supported.")

Copilot · 2025-12-12T09:25:14Z

examples/mri_testscan.py

+            """
+            Input: 2x3 affine matrix.
+            Output: 2x3 affine matrix recentered around the image center.
+            """


The indentation of the docstring is inconsistent with the function definition. The opening triple quotes should align with the function's indentation level (8 spaces), not have extra indentation (12 spaces).

Suggested change

"""

Input: 2x3 affine matrix.

Output: 2x3 affine matrix recentered around the image center.

"""

"""

Input: 2x3 affine matrix.

Output: 2x3 affine matrix recentered around the image center.

"""

Copilot · 2025-12-12T09:25:15Z

heat/ndimage/affine.py

+    b = torch.tensor(M[:, ND:], device=device).reshape(ND, 1)
+
+    # Backward warp matrix
+    A_inv = torch.inverse(A)


The matrix inversion at line 199 can fail if the transformation matrix A is singular (non-invertible). This will raise an exception from torch.inverse, but without a clear error message explaining the issue. Consider adding a try-except block to catch this case and provide a more informative error message to users about why their transformation matrix is invalid.

Copilot · 2025-12-12T09:25:15Z

examples/mri_testscan.py

+    # Define transforms (all center-aware)
+    # ========================================================
+
+    cx, cy = W/2, H/2


The variables cx and cy are defined but never used (they are redefined within the recenter function at line 88). Consider removing this line to avoid confusion.

Suggested change

cx, cy = W/2, H/2

examples/mri_testscan.py

heat/ndimage/affine.py

Copilot · 2025-12-12T09:25:17Z

heat/ndimage/affine.py

+def _normalize_input(x, ND):
+    orig_shape = x.shape
+    t = x.larray
+
+    # 2D image: (H, W) or (C, H, W)
+    if ND == 2:
+        if x.ndim == 2:  # (H,W)
+            t = t.unsqueeze(0).unsqueeze(0)  # → (1,1,H,W)
+        elif x.ndim == 3:  # (C,H,W)
+            t = t.unsqueeze(0)  # → (1,C,H,W)
+
+    # 3D image: (D,H,W) or (C,D,H,W)
+    else:
+        if x.ndim == 3:  # (D,H,W)
+            t = t.unsqueeze(0).unsqueeze(0)  # → (1,1,D,H,W)
+        elif x.ndim == 4:  # (C,D,H,W)
+            t = t.unsqueeze(0)  # → (1,C,D,H,W)
+
+    return t, orig_shape


The _normalize_input function doesn't handle invalid input dimensions. For 2D transforms (ND=2), if the input has ndim != 2 and ndim != 3, the function returns the original tensor without proper normalization. Similarly for 3D. This will cause errors downstream in the sampling functions. Add an else clause to raise a ValueError with a clear message about expected dimensions.

Copilot · 2025-12-12T09:25:17Z

heat/ndimage/tests/test_affine_nd.py

+    def test_expand_rotation(self):
+        x = ht.array([
+            [1.,0.],
+            [0.,1.]
+        ])
+
+        M = [[0,-1,0],
+             [1,0,0]]
+
+        y = affine_transform(x, M, expand=True)
+
+        self.assertTrue(y.shape[0] >= 2 and y.shape[1] >= 2)


This test checks that the shape doesn't shrink when expand=True, but since expand is currently a no-op (as documented in affine.py line 178), this test doesn't actually validate the expand functionality. Consider either skipping this test with a clear comment explaining the feature is not yet implemented, or adding an assertion that the output shape equals the input shape to make the test's purpose clearer.

heat/ndimage/affine.py

heat/ndimage/tests/test_affine_nd.py

ClaudiaComito · 2025-12-12T09:41:13Z

TODOs after today's meeting:

fix pre-commit error
incorporate co-pilot suggestions (after reviewing them)
update branch
mark ready for review
move examples to dedicated directory in heat/examples
TBD: make a jupyter notebook out of examples?

next meeting Tues 15.12. 15:00 with @lolacaro @markak47 @ClaudiaComito

markak47 · 2025-12-14T12:06:30Z

**@github-copilot please re-review the PR after the latest changes.
**

github-actions · 2025-12-15T04:44:33Z

Thank you for the PR!

github-actions · 2026-01-04T12:49:32Z

Thank you for the PR!

codecov · 2026-01-04T13:15:56Z

Codecov Report

❌ Patch coverage is 54.26829% with 75 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.50%. Comparing base (1a69cc7) to head (2ebbdf0).
⚠️ Report is 127 commits behind head on main.

Files with missing lines	Patch %	Lines
heat/ndimage/affine.py	54.26%	75 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2048      +/-   ##
==========================================
- Coverage   91.96%   91.50%   -0.46%     
==========================================
  Files          88       89       +1     
  Lines       13496    13660     +164     
==========================================
+ Hits        12411    12500      +89     
- Misses       1085     1160      +75

Flag	Coverage Δ
unit	`91.50% <54.26%> (-0.46%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

github-actions · 2026-01-09T05:01:50Z

Thank you for the PR!

ClaudiaComito

Post-meeting comments for @markak47 , more later!

ClaudiaComito · 2026-01-09T08:52:52Z

heat/ndimage/affine.py

+    if x.split is not None:
+        x_full = x.resplit(None)
+        y_full = _affine_transform_local(x_full, M, order, mode, constant_value, expand)


This is bad because x might be a huge DNDarray that doesn't fit in the RAM of the single process

if

M is 2-D and

x is 3-D and

x.split = 0

example: huge stack of images, (1M, 1024, 1024) distributed along the stacking dimension

then _affine_transform_local can be applied to the local slices and resplit(None) is not necessary

same with 3-D M and 4-D x

if x.is_distributed() and x.split != 0: raise NotImplementedError("Affine transforms along the distributed axis not supported yet")

ClaudiaComito · 2026-01-09T08:54:08Z

heat/ndimage/affine.py

+    ht.DNDarray
+        Transformed array.
+    """
+    M = np.asarray(M)


for more information, see https://pre-commit.ci

github-actions · 2026-01-16T09:00:09Z

Thank you for the PR!

markak47 added 2 commits November 16, 2025 17:53

affine for ND still to fix 4D

af28e1e

Fix Normalize shape (npx expo start -c

6798efb

ND)

github-actions bot added the features label Nov 28, 2025

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from 505a7e0 to 9bdd8d2 Compare November 30, 2025 17:34

JuanPedroGHM added this to Roadmap Dec 2, 2025

github-project-automation bot moved this to Todo in Roadmap Dec 2, 2025

JuanPedroGHM added this to the 1.8.0 milestone Dec 2, 2025

markak47 added 3 commits December 5, 2025 15:48

changed nii.gz into heat datasets with liscence file removed unnecess…

cb9b523

…ary files module name change

changed nii.gz into heat datasets with liscence file ,removed unneces…

4bbd12e

…sary files, module name change

changed nii.gz into heat datasets with liscence file ,removed unneces…

188bc5c

…sary files, module name change

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from 1e81333 to 188bc5c Compare December 5, 2025 14:52

ClaudiaComito requested a review from Copilot December 12, 2025 09:17

Copilot started reviewing on behalf of ClaudiaComito December 12, 2025 09:17 View session

Copilot AI reviewed Dec 12, 2025

View reviewed changes

added documentation

d986169

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from 8a83636 to d986169 Compare December 14, 2025 11:47

added documentation to mri scan

2cda26b

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from b0619f6 to 2cda26b Compare December 14, 2025 11:49

removed duplicate file

970c25c

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from c0b55a5 to 970c25c Compare December 14, 2025 11:53

markak47 self-assigned this Dec 14, 2025

markak47 added 2 commits December 14, 2025 13:01

chore: trigger copilot review

681ba56

chore: trigger copilot review

2a29fda

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from 345ac2c to 2a29fda Compare December 14, 2025 12:02

helmholtz-analytics deleted a comment from Copilot AI Dec 14, 2025

updated test cases to also test for all valid split

20ec2d9

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from 176c194 to 20ec2d9 Compare December 14, 2025 12:43

examples

74f37fa

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from a1731ee to 74f37fa Compare December 14, 2025 13:24

ClaudiaComito marked this pull request as ready for review December 15, 2025 04:33

fixed mpirun + new test_cube.py example file

144a884

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from e3d66ad to 144a884 Compare January 4, 2026 12:45

ClaudiaComito moved this from Todo to In Progress in Roadmap Jan 9, 2026

ClaudiaComito requested changes Jan 9, 2026

View reviewed changes

MPI-safe + raise error, TODO Halo exchange for big data sets

7a6fd1c

markak47 force-pushed the features/1900-Image_transformations_beyond_FFT branch from 43a3ebe to 7a6fd1c Compare January 16, 2026 08:54

[pre-commit.ci] auto fixes from pre-commit.com hooks

2ebbdf0

for more information, see https://pre-commit.ci

-):
+):
+    if order not in (0, 1):
+        raise ValueError(f"Unsupported interpolation order: {order}. Only 0 (nearest) and 1 (bilinear) are supported.")

Features/1900 image transformations beyond fft #2048

Are you sure you want to change the base?

Features/1900 image transformations beyond fft #2048

Uh oh!

Conversation

markak47 commented Nov 28, 2025 • edited by ClaudiaComito Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ClaudiaComito commented Dec 2, 2025 • edited by markak47 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ClaudiaComito commented Dec 12, 2025 • edited by markak47 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

markak47 commented Dec 14, 2025

Uh oh!

github-actions bot commented Dec 15, 2025

Uh oh!

github-actions bot commented Jan 4, 2026

Uh oh!

codecov bot commented Jan 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions bot commented Jan 9, 2026

Uh oh!

ClaudiaComito left a comment

Choose a reason for hiding this comment

Uh oh!

ClaudiaComito Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

ClaudiaComito Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ClaudiaComito Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ClaudiaComito Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jan 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

markak47 commented Nov 28, 2025 •

edited by ClaudiaComito

Loading

ClaudiaComito commented Dec 2, 2025 •

edited by markak47

Loading

ClaudiaComito commented Dec 12, 2025 •

edited by markak47

Loading

codecov bot commented Jan 4, 2026 •

edited

Loading

ClaudiaComito Jan 9, 2026 •

edited

Loading

ClaudiaComito Jan 9, 2026 •

edited

Loading