FIX: Correct axis parameters and remove duplicate coefficient scaling…

[rcjackson]

This commit fixes two critical issues in the smoothness cost and gradient functions:

1. Fixed incorrect axis parameters in second derivative calculations:

   • In both JAX and NumPy implementations, the axis parameter was incorrectly specified when computing second derivatives
   • For x-direction terms (Cx), all second derivatives should use axis=2
   • For y-direction terms (Cy), all second derivatives should use axis=1
   • For z-direction terms (Cz), all second derivatives should use axis=0
   • This ensures that second derivatives are taken in the correct spatial direction

2. Removed duplicate Cx, Cy, Cz scaling in gradient functions:

   • Both JAX and NumPy gradient functions were incorrectly multiplying by Cx, Cy, Cz coefficients twice
   • The cost function already includes these coefficients, and the gradient (via chain rule or autodiff) inherently includes them
   • Removed the extra * Cx * 2, * Cy * 2, * Cz * 2 scaling at the end
   • Also removed the /dx, /dy, /dz operations that were incorrect in JAX

3. Simplified JAX gradient function:

   • Removed manual gradient calculation code that duplicated the cost function logic
   • Now relies entirely on JAX's automatic differentiation (vjp) which is more accurate and maintainable
   • Fixed the vjp call to pass Cx, Cy, Cz as positional arguments (not keyword args)

These fixes ensure that the JAX and NumPy/SciPy engines produce consistent results when using smoothness constraints with non-zero Cx, Cy, Cz parameters.

Fixes: https://openradar.discourse.group/t/smoothness-costs-in-different-engines/666

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 68.85%. Comparing base (872a709) to head (f7c7d3f).
⚠️ Report is 6 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #150      +/-   ##
==========================================
- Coverage   69.02%   68.85%   -0.17%     
==========================================
  Files          32       32              
  Lines        5169     5141      -28     
==========================================
- Hits         3568     3540      -28     
  Misses       1601     1601              

Flag	Coverage Δ
unittests	68.85% <100.00%> (-0.17%)	⬇️

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.