Bug Fix: weight=0 for _broadcast_input='Node' fails

[singh-jaydeep]

PR #1921 added support for coil optimization with Pytree inputs for weight, target, and bounds. It relies on broadcasting and flattening the input Pytrees to have dimensions consistent with the objective's output shape. When weight=0 for a given coil or coilset, the corresponding coils (or grid nodes) are dropped during the broadcasting process. However, this wasn't being done correctly for objectives like CoilCurvature that output a single value per grid node.

Fix was suggested in original issue, #2219. We keep track of which coils are masked (weight=0) in the attribute _coilset_tree["coilset_mask"]. This always has length num_coils. The mask _coilset_tree["objective_mask"] matches the output shape of the objective, and in the case of per-node objectives, tracks which individual grid nodes are masked. Keeping both masks makes it easier to index things correctly when building the objective.

Also added some tests to catch bugs of this form.

Addresses Issue #2219.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    1.13 %    |     4.058e+03      |     4.104e+03      |    45.84     |       32.55        |       29.94        |
  test_proximal_jac_w7x_with_eq_update   |   -1.03 %    |     6.577e+03      |     6.509e+03      |    -67.58    |       157.14       |       154.29       |
  test_proximal_freeb_jac                |   -0.14 %    |     1.334e+04      |     1.332e+04      |    -18.04    |       79.83        |       81.13        |
  test_proximal_freeb_jac_blocked        |    0.26 %    |     7.641e+03      |     7.661e+03      |    20.03     |       71.00        |       69.80        |
  test_proximal_freeb_jac_batched        |   -0.72 %    |     7.645e+03      |     7.589e+03      |    -55.23    |       70.08        |       70.81        |
  test_proximal_jac_ripple               |   -0.72 %    |     3.547e+03      |     3.521e+03      |    -25.52    |       55.48        |       55.56        |
  test_proximal_jac_ripple_bounce1d      |   -0.39 %    |     3.760e+03      |     3.745e+03      |    -14.66    |       70.54        |       69.25        |
  test_eq_solve                          |   -2.19 %    |     1.977e+03      |     1.933e+03      |    -43.34    |       59.00        |       58.27        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.32%. Comparing base (7b6ed16) to head (7ae1e1f).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2221      +/-   ##
==========================================
+ Coverage   94.30%   94.32%   +0.01%     
==========================================
  Files         101      101              
  Lines       28866    28868       +2     
==========================================
+ Hits        27223    27229       +6     
+ Misses       1643     1639       -4     

Files with missing lines	Coverage Δ
desc/objectives/_coils.py	99.47% <100.00%> (+<0.01%)	⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[YigitElma]

I will test it myself too, but it looks good. I am personally fine with the breaking change; this was a pretty new addition, and going forward, I like the new name better.

[YigitElma]

Should we keep the dict key name the same as before coilset_mask? The meaning is different but at least prevents breaking some code.

[singh-jaydeep]

I switched back to coilset_mask.

[YigitElma]

I will review it properly again later. Just wanted to mention the changelog issue