Allow PyTree inputs for coil objectives

[singh-jaydeep]

Adds new input formats for objectives inheriting from _CoilObjective. In particular,

• bounds, target, and weight can accept pytrees with values for individual coils, provided the input is a prefix of the input CoilSet. weights can be zero, in which case the corresponding coil is masked.
• A new option indices accepts pytrees with the same structure, indicating whether the bound/target should be applied to given coils.
• The attribute _broadcast_input is required for coil objectives, stating if they are computed per-coil or per-grid node.

Notes

• Resolves Allow for Pytree input for coil objective targets and bounds  #1152.
• This may not be the cleanest/most general approach, and am open to any suggestions!
• When the user gives an input which is not just a simple scalar we broadcast the input to be of size dim_f and store it. This isn't very efficient, but saves us from adding new logic to the compute function in _Objective.
• The setter functions in _Objective will have issues with pytree inputs. Would it make sense to add specific setters in _CoilObjective? [resolved]
• Would require tests, etc. [done]

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.20 %    |     3.793e+03      |     3.914e+03      |    121.38    |       38.76        |       36.04        |
  test_proximal_jac_w7x_with_eq_update   |    0.32 %    |     6.638e+03      |     6.659e+03      |    21.04     |       161.52       |       162.00       |
  test_proximal_freeb_jac                |   -0.55 %    |     1.324e+04      |     1.317e+04      |    -72.96    |       85.95        |       84.04        |
  test_proximal_freeb_jac_blocked        |   -0.45 %    |     7.522e+03      |     7.488e+03      |    -33.92    |       72.79        |       73.31        |
  test_proximal_freeb_jac_batched        |    0.12 %    |     7.414e+03      |     7.423e+03      |     9.01     |       72.10        |       72.39        |
  test_proximal_jac_ripple               |    0.74 %    |     3.449e+03      |     3.475e+03      |    25.68     |       65.29        |       64.39        |
  test_proximal_jac_ripple_bounce1d      |    0.21 %    |     3.566e+03      |     3.573e+03      |     7.32     |       75.23        |       77.65        |
  test_eq_solve                          |   -0.34 %    |     2.028e+03      |     2.021e+03      |    -6.95     |       93.07        |       93.18        |

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

[codecov]

Codecov Report

❌ Patch coverage is 84.37500% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.72%. Comparing base (c42d083) to head (4ed8c18).

Files with missing lines	Patch %	Lines
desc/coils.py	6.66%	14 Missing ⚠️
desc/objectives/_coils.py	98.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1921      +/-   ##
==========================================
- Coverage   95.76%   95.72%   -0.05%     
==========================================
  Files         102      102              
  Lines       28344    28424      +80     
==========================================
+ Hits        27143    27208      +65     
- Misses       1201     1216      +15     

Files with missing lines	Coverage Δ
desc/backend.py	89.69% <ø> (ø)
desc/objectives/objective_funs.py	94.81% <100.00%> (+<0.01%)	⬆️
desc/objectives/_coils.py	99.33% <98.71%> (-0.09%)	⬇️
desc/coils.py	96.25% <6.66%> (-1.71%)	⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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

[YigitElma]

@singh-jaydeep is this ready for review?

[singh-jaydeep]

I'll have it ready by Monday. There's one thing related to memory usage I wanted to check.

[dpanici]

For tests, two examples to modify with the new input structure

• DESC/tests/test_examples.py

  Line 1831 in 587fd85

  def test_second_stage_optimization_CoilSet():

• DESC/tests/test_linear_objectives.py

  Line 984 in 587fd85

  def test_fix_subset_of_params_in_collection(DummyMixedCoilSet):

[dpanici]

It might make more sense to have this as a (maybe private) method of this objective, then when you make the target setter, you can use this as well.

[dpanici]

this is probably worth some additional description as this is I think the first time indices is being used in this specific way (usually it is with a FixedParameter subclass and so indices are the specific parts of params being fixed by the constraint).

Maybe something additional like "Boolean indicating coils in the CoilSet to be targeted by the objective. If True, includes every coil."

[singh-jaydeep]

Sorry I missed this! Will do.

[dpanici]

So in this test, the first coil of the second element of the MixedCoilSet is not included in the optimization, that just means that the gradient of the objective w.r.t the params associated with that coil is zero right?

I guess what is unclear to me is why we would expect no change in its length? It is not like we are not including it at all as an optimizable DOF, it is instead a nullspace direction of our jacobian (it has no effect on f) and thus nothing really is keeping it from changing a lot vs not at all

[singh-jaydeep]

I see, yes this just zeros out the gradient associated to that coil's parameters. Would the behavior (whether the coil length still changes, or none of its params change) then depend on the particular optimizer? It passed locally, but maybe it's not a robust test.

Maybe I just test that the explicitly targeted coils do have the length we want, and check the gradients of the objective wrt. the non-optimized coils are zero? Or is there something else which makes sense

[singh-jaydeep]

I updated the test to separately check the gradients of unoptimized coils and the lengths of optimized coils.

[dpanici]

last thing I think I want for this PR: adding an example use of this to the coil tutorial notebook. In that notebook as one of the last cells, we have an optimization of a MixedCoilSet of modular coils and saddle coils. It would be nice to say have the length penalty only target the modular coils (as it effectively does now) and then maybe add some other penalty for the saddle coils, and have this take advantage of the new API as well as show the use of indices

And also while we are changing that notebook, adding to the constraints for that last finite beta coil optimization this constraint
LinkingCurrentConsistency(eq,coilset,eq_fixed=True) is something I've been meaning to add, and this would be a good spot to change that

@YigitElma do not change the coil notebook on #1877 until this PR is in to avoid annoying conflicts

[singh-jaydeep]

I'll work today on adding to the coil tutorial notebook!

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[dpanici]

hm weird that CI cannot run the noteboook due to memory, try adding jac_chunk_size=5 to the LinkingCurrentConsistency constraint.

Also I cannot tell bc the docs did not build but did you rerun the notebook?

[singh-jaydeep]

I did rerun the notebook and it seemed to work fine (even when I run the action individually on my branch?), but this isn't something I'm too familiar with. I changed as you recommended and will see

[dpanici]

looks like maybe you had dark mode on or something? Never seen that before haha but it does make the coils difficult to see here

[singh-jaydeep]

lmao I'll fix that

[dpanici]

perhaps would be good to leave a comment here saying this is due to the coils objectives, or to add to the _Objective docstring that if _mask is an attribute, it will multiply f here

[dpanici]

my last change request I promise... I think either leave a comment where mask is used in objective funs to describe why it is there, or add a line about mask to the docstring of _Objective

[singh-jaydeep]

The "mask" option is proving to be a bit clunky. In addition to getting the normalization somewhat wrong, it doesn't work well for scalar valued loss functions unless the mask gets passed in (e.g. the mean by default includes the masked values). Fixing these things would further clutter up the objective compute functions, which I don't want to do unless masking is useful for other objectives as well. Is there a preferred way of handling this?

(Another option is if "indices" gets passed into a coilobjective, when building it changes self.things to only include the coils which are being optimized. This is cleaner in that it keeps the masking within the coil objective only. But the optimizer objects when passing in a coilset and objective with "indices", since the set of optimizable things doesn't match.)

[dpanici]

The "mask" option is proving to be a bit clunky. In addition to getting the normalization somewhat wrong, it doesn't work well for scalar valued loss functions unless the mask gets passed in (e.g. the mean by default includes the masked values). Fixing these things would further clutter up the objective compute functions, which I don't want to do unless masking is useful for other objectives as well. Is there a preferred way of handling this?

(Another option is if "indices" gets passed into a coilobjective, when building it changes self.things to only include the coils which are being optimized. This is cleaner in that it keeps the masking within the coil objective only. But the optimizer objects when passing in a coilset and objective with "indices", since the set of optimizable things doesn't match.)

Yes I see the issue...

While this would be a nice feature to have, I think that we have some eventual future plans try to change the optimization API itself to allow us to, say, apply the CoilLength objective to a sub-member of the coilset and have the optimization machinery understand that this means only target one and dont touch the ohters

I think this PR even without the mask object is useful in that it opens the door to more easily do this exact sort of function. By that I mean, already one could say provide bounds as an array and just make the array index of the coil(s) you don't care about have (-np.inf,np.inf), for example, or if providing target one could also provide a weights array that sets a tiny weight on the ones you dont care about.

I would be ok with punting on something more complicated like this indices feature, as I see that it is a bit clunky and as of right now don't have any further ideas to help with it.

@PlasmaControl/desc-dev if anyone else has ideas please chime in, but this is all I have for now

[dpanici]

@singh-jaydeep can you try fixing the notebook conflict? I know that can be annoying, I think there is a jupyter nbdiff tool you can use to make it easier to merge locally