Skip to content

Use compact [1,-2,1] kernel for BendingEnergyLoss second derivatives#8918

Open
vishnukannaujia wants to merge 2 commits into
Project-MONAI:devfrom
vishnukannaujia:5939-bending-energy-compact-second-derivative
Open

Use compact [1,-2,1] kernel for BendingEnergyLoss second derivatives#8918
vishnukannaujia wants to merge 2 commits into
Project-MONAI:devfrom
vishnukannaujia:5939-bending-energy-compact-second-derivative

Conversation

@vishnukannaujia

Copy link
Copy Markdown

Fixes #5939.

Summary

BendingEnergyLoss previously computed second-order derivatives by applying the central first-order helper spatial_gradient twice, which yields a wide [1, 0, -2, 0, 1] / 4 stencil for pure second derivatives that spans four voxels per axis. That stencil is less accurate than the standard compact one and forced the validation "all spatial dims > 4".

This PR replaces the second-order computation with compact central stencils evaluated directly on pred:

  • Pure d^2/dx_i^2: x[i+1] - 2 * x[i] + x[i-1] (the standard [1, -2, 1] kernel).
  • Mixed d^2/(dx_i dx_j): (x[i+1,j+1] - x[i+1,j-1] - x[i-1,j+1] + x[i-1,j-1]) / 4 (compact 4-point central scheme).

Both span three voxels per axis, so the spatial-size validation is relaxed from > 4 to > 2, matching DiffusionLoss. The public API (__init__(normalize, reduction), forward(pred)) and normalize semantics are unchanged. The existing spatial_gradient helper is left untouched because DiffusionLoss still uses it.

Why TEST_CASES expected values do not change

For f(x) = x^2, the analytical second derivative is the constant 2. Both the previous central-of-central stencil (f[i+2] - 2*f[i] + f[i-2]) / 4 and the new compact [1, -2, 1] stencil f[i+1] - 2*f[i] + f[i-1] are exact for quadratics, so both return 2 at every interior voxel. Mixed-partial test inputs are constant in at least one of the two axes, so both stencils return 0 for mixed terms on these cases. Squared and reduced by mean, the existing TEST_CASES expected values (0.0, 4.0, 100.0) are therefore invariant under this change.

What does change in the tests:

  • test_ill_shape is updated to trigger on shape 2 (was 4) so it still exercises the spatial-size guard.
  • A new TEST_CASES row covers shape (1, 3, 3, 3, 3) of ones (previously rejected by the > 4 guard) → expected 0.0, validating the relaxed guard.

Reference for the compact mixed-partial scheme: Pavel Holoborodko's finite-difference notes cited in the original issue.

The previous implementation differentiated the central first-order helper
twice, yielding a wide [1, 0, -2, 0, 1] / 4 stencil for pure second
derivatives that spans four voxels per axis. This is less accurate than
the standard compact [1, -2, 1] stencil and forced the validation
"all spatial dims > 4".

Switch to compact stencils computed directly from pred:
- Pure d^2/dx_i^2 uses x[i+1] - 2 * x[i] + x[i-1].
- Mixed d^2/(dx_i dx_j) uses
  (x[i+1,j+1] - x[i+1,j-1] - x[i-1,j+1] + x[i-1,j-1]) / 4.

Both span three voxels per axis, so the spatial-size validation relaxes
to "> 2", matching DiffusionLoss. The public API
(__init__(normalize, reduction), forward(pred)) and normalize semantics
are unchanged.

For quadratic inputs both old and new stencils return the exact analytic
second derivative (= 2 for f(x) = x^2), so the existing TEST_CASES
expected values (0.0, 4.0, 100.0) are invariant under this change; only
test_ill_shape is updated to trigger on shape 2 instead of 4, and a new
TEST_CASES row exercises the relaxed guard on shape (1, 3, 3, 3, 3).

Fixes Project-MONAI#5939.

Signed-off-by: Vishnu Kannaujia <vishnu.kannaujia@gmail.com>
@coderabbitai

coderabbitai Bot commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: f587e8f9-8bcd-46f9-9b37-4c6d7fee1158

📥 Commits

Reviewing files that changed from the base of the PR and between 3dcbac4 and 51e60da.

📒 Files selected for processing (1)
  • monai/losses/deform.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • monai/losses/deform.py

📝 Walkthrough

Walkthrough

spatial_gradient_squared is added to compute pure second-order spatial derivatives using the compact [1, -2, 1] stencil and mixed partial derivatives via forward-difference cross-term formula, with outputs sliced to [1:-1] on differentiated dimensions. BendingEnergyLoss.forward is updated to call this helper instead of chaining two first-order spatial_gradient calls, accumulating energy from pure and twice-mixed derivative terms with revised normalization. Minimum valid spatial size drops from > 4 to > 2. Tests are adjusted accordingly.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed Title clearly and concisely summarizes the main change: adopting compact [1,-2,1] kernel for computing second derivatives in BendingEnergyLoss.
Description check ✅ Passed Description comprehensively covers the rationale, technical implementation, API stability, and test validation. Includes proper issue reference and detailed justification for expected value invariance.
Linked Issues check ✅ Passed Changes directly address #5939 objectives: replaces numerically unstable double-application of first-order gradient with compact [1,-2,1] kernel for pure derivatives and compact schemes for mixed partials, improving accuracy and stability.
Out of Scope Changes check ✅ Passed All changes are scoped to BendingEnergyLoss implementation and its tests. New spatial_gradient_squared helper, validation relaxation from >4 to >2, and test adjustments directly support the stated objectives.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@monai/losses/deform.py`:
- Line 150: The `torch.tensor(0)` initialization at line 150 omits the device
parameter, causing a CPU tensor to be added to GPU tensors in the loop,
inconsistent with line 147's approach using `torch.tensor(pred.shape,
device=pred.device)`. Additionally, torch.tensor(0) defaults to int64 dtype
which may be incorrect. Update the initialization at line 150 to include the
device parameter set to pred.device to match the pattern used earlier in the
function. Also apply the same fix to the DiffusionLoss class around line 237
which has the identical pattern for consistency across the file.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d7c6cee0-9dec-4f81-b50d-50cf4b5fc43b

📥 Commits

Reviewing files that changed from the base of the PR and between eccefc5 and 3dcbac4.

📒 Files selected for processing (2)
  • monai/losses/deform.py
  • tests/losses/deform/test_bending_energy.py

Comment thread monai/losses/deform.py Outdated
The compact pure-derivative stencil added in this branch
(`x[i+1] - 2 * x[i] + x[i-1]`) intentionally avoids the `/2.0` factor of the
old first-order helper. That means an integer-dtype `pred` (e.g. the
arange-squared inputs in ``test_shape_5``: a 1-d input with no mixed terms
to force float promotion via the `/4.0` factor) now stays Long all the way
through the accumulator, and `torch.mean(energy)` raises:

    RuntimeError: mean(): could not infer output dtype. Input dtype must be
    either a floating point or complex dtype. Got: Long

Initialize the accumulator explicitly as a float and bind it to
`pred.device` so it broadcasts against integer-dtype inputs and matches
GPU tensors. Also matches the device pattern already used a few lines
above for `spatial_dims`.

`DiffusionLoss` is left unchanged: its first-order helper still divides
by 2.0, so its accumulator is always promoted to float by the first
addition; updating it here would expand the scope of Project-MONAI#5939.

Signed-off-by: Vishnu Kannaujia <vishnu.kannaujia@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

BendingEnergyLoss

1 participant