Fix bug with recomputing quantities on incorrect grid

[unalmis]

If you call data = eq.compute(..., data=data) the contract is that stuff already in data stays there. master violates this contract. At best it causes slow downs by recomputing things, at worst it computes quantities incorrectly.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.94%. Comparing base (b68c98c) to head (05fb63b).
Report is 1552 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1006   +/-   ##
=======================================
  Coverage   94.94%   94.94%           
=======================================
  Files          87       87           
  Lines       21880    21884    +4     
=======================================
+ Hits        20773    20777    +4     
  Misses       1107     1107           

Files with missing lines	Coverage Δ
desc/equilibrium/equilibrium.py	95.95% <100.00%> (+0.02%)	⬆️

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.46 +/- 4.04     | +2.55e-03 +/- 2.24e-02 |  5.57e-01 +/- 1.6e-02  |  5.54e-01 +/- 1.5e-02  |
 test_build_transform_fft_midres         |     +0.30 +/- 2.76     | +1.85e-03 +/- 1.74e-02 |  6.30e-01 +/- 8.2e-03  |  6.28e-01 +/- 1.5e-02  |
 test_build_transform_fft_highres        |     +0.02 +/- 1.80     | +1.70e-04 +/- 1.85e-02 |  1.03e+00 +/- 9.0e-03  |  1.03e+00 +/- 1.6e-02  |
 test_equilibrium_init_lowres            |     -4.08 +/- 4.77     | -1.81e-01 +/- 2.11e-01 |  4.25e+00 +/- 1.9e-01  |  4.43e+00 +/- 1.0e-01  |
 test_equilibrium_init_medres            |     -2.18 +/- 5.69     | -1.07e-01 +/- 2.79e-01 |  4.80e+00 +/- 1.7e-01  |  4.91e+00 +/- 2.2e-01  |
 test_equilibrium_init_highres           |     -1.31 +/- 3.62     | -8.73e-02 +/- 2.42e-01 |  6.59e+00 +/- 1.7e-01  |  6.68e+00 +/- 1.7e-01  |
 test_objective_compile_dshape_current   |     +0.42 +/- 2.63     | +1.61e-02 +/- 1.00e-01 |  3.83e+00 +/- 8.6e-02  |  3.82e+00 +/- 5.2e-02  |
 test_objective_compile_atf              |     +1.37 +/- 3.47     | +1.00e-01 +/- 2.53e-01 |  7.38e+00 +/- 1.9e-01  |  7.28e+00 +/- 1.7e-01  |
 test_objective_compute_dshape_current   |     -0.44 +/- 6.23     | -1.74e-05 +/- 2.47e-04 |  3.94e-03 +/- 2.4e-04  |  3.96e-03 +/- 7.1e-05  |
 test_objective_compute_atf              |     +4.07 +/- 2.40     | +7.21e-04 +/- 4.26e-04 |  1.84e-02 +/- 3.0e-04  |  1.77e-02 +/- 3.0e-04  |
 test_objective_jac_dshape_current       |     +3.41 +/- 8.82     | +1.42e-03 +/- 3.66e-03 |  4.29e-02 +/- 2.6e-03  |  4.15e-02 +/- 2.6e-03  |
 test_objective_jac_atf                  |     -0.29 +/- 3.37     | -5.50e-03 +/- 6.41e-02 |  1.90e+00 +/- 4.4e-02  |  1.90e+00 +/- 4.7e-02  |
 test_perturb_1                          |     +2.67 +/- 3.81     | +3.13e-01 +/- 4.45e-01 |  1.20e+01 +/- 2.3e-01  |  1.17e+01 +/- 3.8e-01  |
 test_perturb_2                          |     +2.78 +/- 3.94     | +4.62e-01 +/- 6.55e-01 |  1.71e+01 +/- 5.7e-01  |  1.66e+01 +/- 3.2e-01  |
 test_proximal_jac_atf                   |     -1.64 +/- 1.15     | -1.19e-01 +/- 8.30e-02 |  7.12e+00 +/- 7.3e-02  |  7.24e+00 +/- 3.9e-02  |
 test_proximal_freeb_compute             |     -1.81 +/- 0.68     | -2.38e-03 +/- 8.89e-04 |  1.29e-01 +/- 5.0e-04  |  1.31e-01 +/- 7.3e-04  |
 test_proximal_freeb_jac                 |     +2.77 +/- 1.00     | +1.99e-01 +/- 7.20e-02 |  7.40e+00 +/- 5.0e-02  |  7.20e+00 +/- 5.2e-02  |

[dpanici]

Is there anyway to test this with existing code or is this too specific to the field line integral stuff?

[unalmis]

If you call data = eq.compute(..., data=data) the contract is that stuff already in data stays there. master violates this contract. At best it causes slow downs by recomputing things, at worst it computes quantities incorrectly. Yes this is needed for #1003. It's related to but distinct from #719 .

[dpanici]

update master data

[f0uriest]

I think changing this line:
https://github.com/PlasmaControl/DESC/blob/rc/grids/desc/equilibrium/equilibrium.py#L844
to

        deps = list(
            set(get_data_deps(names, obj=p, has_axis=grid.axis.size) + names)
            - data.keys()
        )

would also fix this issue. If the user wants to compute x which depends on y which in turn depends on z, and they pass in y then we shouldn't need to compute z at all.

[unalmis]

I think changing this line: https://github.com/PlasmaControl/DESC/blob/rc/grids/desc/equilibrium/equilibrium.py#L844 to

        deps = list(
            set(get_data_deps(names, obj=p, has_axis=grid.axis.size) + names)
            - data.keys()
        )

would also fix this issue. If the user wants to compute x which depends on y which in turn depends on z, and they pass in y then we shouldn't need to compute z at all.

I added it, but subtracting out data.keys() as you suggest does:

If the user wants to compute x which depends on y which in turn depends on z, and they pass in y then we shouldn't need to compute y at all.

The dependencies of y may still computed redundantly if they are second order dependencies of x.

1. To solve the redundant computation issue one must store the immediate dependencies of quantities. Or come up with some clever recursive logic.

@register_compute_fun(
    name="grad(|B|^2)",
    label="\\nabla |B|^{2}",
    units="T^{2} \\cdot m^{-1}",
    units_long="Tesla squared / meters",
    description="Magnetic pressure gradient",
    dim=3,
    params=[],
    transforms={},
    profiles=[],
    coordinates="rtz",
    data=["|B|", "grad(|B|)"], # i.e. also store this.
)
def _gradB2(params, transforms, profiles, data, **kwargs):
    data["grad(|B|^2)"] = 2 * (data["|B|"] * data["grad(|B|)"].T).T
    return data

2. To solve the computing things on correct grid issue, we need to add label to each compute quantity denoting what type of grid it is best used.

[f0uriest]

Ahhh ok yeah. In desc.compute.utils.compute we do things recursively, so in that case priming data with just the first order dependencies will avoid computing any extra stuff. But in eq.compute we first get the full list of dependencies (beyond just first order). That could eventually be fixed but doesn't have to be now.