iota vacuum and current contributions

[ddudt]

No description provided.

[ddudt]

I left the existing iota computation alone. Technically alpha is now redundant with iota_current_num and beta is redundant with iota_vacuum_num, so there is some repeated code. But I want to be able to compute iota_vacuum and iota_current even when the iota profile is supplied. So I think there would be circular dependencies otherwise -- but let me know if you see a way to resolve that.

[unalmis]

1. We should remove these guards from iota_den, iota_den_r, iota_den_rr so that we can also compute them when an iota profile is supplied instead of a current profile.

   DESC/desc/compute/_profiles.py

   Line 1000 in c0b4441

   if profiles["current"] is None:

2. Heads up to avoid fighting the tests: Remove the current dependency on the iota_den stuff. We only used them for the above checks.

   DESC/desc/compute/_profiles.py

   Line 990 in c0b4441

   profiles=["current"],

3. We should be able to remove iota_den from this line

   DESC/tests/test_axis_limits.py

   Line 105 in c0b4441

   or (eq.current is None and ("iota_num" in name or "iota_den" in name))

[f0uriest]

1. We should remove these guards from iota_den, iota_den_r, iota_den_rr so that we can also compute them when an iota profile is supplied instead of a current profile.

   DESC/desc/compute/_profiles.py

   Line 1000 in c0b4441

   if profiles["current"] is None:

   2. Heads up to avoid fighting the tests: Remove the current dependency on the iota_den stuff. We only used them for the above checks.

      DESC/desc/compute/_profiles.py

      Line 990 in c0b4441

      profiles=["current"],

   3. We should be able to remove iota_den from this line

      DESC/tests/test_axis_limits.py

      Line 105 in c0b4441

      or (eq.current is None and ("iota_num" in name or "iota_den" in name))

I think we might not want to do your first point. Because iota_den etc are dependencies for iota, they get computed even if the iota profile is supplied, but because they are hardcoded to nan the compiler should skip that. If we remove this it may lead to redundant calculations for standard fixed iota cases

[unalmis]

I think we might not want to do your first point. Because iota_den etc are dependencies for iota, they get computed even if the iota profile is supplied, but because they are hardcoded to nan the compiler should skip that. If we remove this it may lead to redundant calculations for standard fixed iota cases

But Daniel wants to be able to compute iota_vacuum and iota_current even when the iota profile is supplied. These quantities require iota_den

[ddudt]

psi_rrr (and higher derivatives) are always 0. Can we remove all references to them from the compute functions to simplify the code?

[unalmis]

In #385 you modified the magnetic field computations to include psi_rrr. So I included psi_rrr to be consistent with all the other compute functions. I think either you or Rory mentioned in one of the meetings months ago that it could be useful to change $\rho$ so that psi_rrr is not zero.

If you decide to remove this please add psi_rrr to the zero_limits list here:

DESC/tests/test_axis_limits.py

Line 17 in c0b4441

zero_limits = {"rho", "psi", "psi_r", "e_theta", "sqrt(g)", "B_t"}

[unalmis]

I prefer names like iota_from_current and iota_in_vacuum by the way.

To compute iota_from_current and its derivatives, you should be able to just copy the compute functions for iota and replace all references to iota_num with the term alpha defined in the compute functions for iota_num. (similar with iota_num_r and alpha_r). This will also give you the correct limits.

Same thing for iota_in_vacuum except use the term beta instead.

Currently on master, the only physics constraint I used to compute these limits is that the magnetic field is physical. (This was necessary because there is an if and only if relationship between this and the limit of iota being finite).

FYI, there are two other perhaps weaker physics conditions which can make computing the limits easier. These conditions are that $\partial_{\rho} \iota$ and $\partial_{\rho} I$ are both zero at axis. I didn't explicitly enforce these and computed the limits in the more general manner because

1. whatever shortcuts they provide don't help when computing the limit of the second derivative of $\iota$.
2. didn't want to assume behavior of user-inputted quantities

[ddudt]

The new quantities are working except for the axis limit of iota from the current. I think I understand why:

I changed the references from things like 2 * jnp.pi * mu_0 * profiles["current"].compute(params["c_l"], dr=0) to instead be like 4 * jnp.pi**2 * data["I"]. This moves away from depending on the current profile as an input, so that we could compute these new quantities even when the user gives an iota profile instead. This also better matches Kaya's equations.

The axis limits then depend on the value of data["I_rr"] instead of profiles["current"].compute(params["c_l"], dr=2). The later is always well defined at the axis by the input, but the former is computed as a flux surface average of $B_\theta$. It appears that in our code this always evaluates to I_rr = 0 at the axis, even though it should generally have a finite value. So we need to implement the proper axis limit for that quantity to finish this work.

[unalmis]

The new quantities are working except for the axis limit of iota from the current...

The test was failing because iota_den was hardcoded to return jnp.nan on fixed iota equilibrium. Commiting the the suggestions here #642 (comment) resolve the issue.

The axis limits then depend on the value of data["I_rr"] ... is computed as a flux surface average of Bθ. It appears that in our code this always evaluates to I_rr = 0 at the axis, even though it should generally have a finite value. So we need to implement the proper axis limit for that quantity to finish this work.

The computation of $\partial_{\rho \rho} I = \partial_{\rho \rho} \langle B_{\theta} \rangle$ is dependent on $\partial_{\rho \rho} \iota$. All the math is done to compute $\lim_{\rho \to 0} \partial_{\rho \rho} \iota$, but we can't implement it until #587 is resolved. This means the integrand: $\partial_{\rho \rho} B_{\theta}$ evaluates as jnp.nan when the iota profile is not given as a known power series. The surface integral replaces nan integrand with zero (see here).

[unalmis]

I think we might not want to do your first point. Because iota_den etc are dependencies for iota, they get computed even if the iota profile is supplied, but because they are hardcoded to nan the compiler should skip that. If we remove this it may lead to redundant calculations for standard fixed iota cases

But Daniel wants to be able to compute iota_vacuum and iota_current even when the iota profile is supplied. These quantities require iota_den

If Rory is okay with 6e9b9e2 then the master data will need to be updated in this PR since iota_den will no longer be nan

[f0uriest]

I'm generally fine with the above, it would be good if you did some basic benchmarking to see if it slows things down appreciably, but I doubt there will be a noticeable difference for most use cases

[ddudt]

This should be working now, I just need to add tests.

Everything can be computed fine when the equilibrium has a current profile assigned. The only thing that can't be computed from an iota profile is iota current, but you can still get that data manually by doing data["iota"] - data["iota vacuum"].

[f0uriest]

Fine as is (assuming the tests pass and the new code is covered) but may want to consider implementing stuff where we now return NaN, I don't think it would be too complicated

[f0uriest]

also make sure to update the changelog

[ddudt]

Fine as is (assuming the tests pass and the new code is covered) but may want to consider implementing stuff where we now return NaN, I don't think it would be too complicated

I implemented these, but you should check my formulas. The logic is a little tricky to avoid circular dependencies. Also these quantities will probably never actually be called by the user, since "iota_num" etc. are only intermediate outputs to compute "iota", etc.

[dpanici]

wasany benchmarking done? like how long it takes to to cmpute iota here versus master?

[github-actions]

|             benchmark_name             |        dt(%)        |        dt(s)        |      t_new(s)      |      t_old(s)      | 
| ------------------------------- | ------------------- | ------------------- | ------------------ | ------------------ |
 test_build_transform_fft_lowres        | -5.57+/-7.86 | -1.66e-03+/-2.3e-03 | 2.82e-02+/-1.8e-03 | 2.99e-02+/-1.5e-03 |
+test_build_transform_fft_midres        | -15.30+/-7.98 | -2.84e-02+/-1.5e-02 | 1.57e-01+/-1.0e-02 | 1.85e-01+/-1.1e-02 |
+test_build_transform_fft_highres       | -10.10+/-4.52 | -8.56e-02+/-3.8e-02 | 7.62e-01+/-2.7e-02 | 8.47e-01+/-2.8e-02 |
 test_equilibrium_init_lowres           | -1.28+/-3.50 | -8.23e-03+/-2.2e-02 | 6.33e-01+/-1.2e-02 | 6.41e-01+/-1.9e-02 |
 test_equilibrium_init_medres           | -3.71+/-3.83 | -2.98e-02+/-3.1e-02 | 7.74e-01+/-2.1e-02 | 8.03e-01+/-2.2e-02 |
 test_equilibrium_init_highres          | -1.31+/-2.90 | -1.73e-02+/-3.8e-02 | 1.30e+00+/-2.9e-02 | 1.32e+00+/-2.5e-02 |
 test_objective_compile_heliotron       | +0.35+/-3.31 | +6.69e-02+/-6.4e-01 | 1.95e+01+/-5.5e-01 | 1.94e+01+/-3.3e-01 |
 test_objective_compile_dshape_current  | +0.13+/-4.44 | +1.00e-02+/-3.3e-01 | 7.42e+00+/-3.0e-01 | 7.41e+00+/-1.4e-01 |
 test_objective_compile_atf             | +2.23+/-3.85 | +4.61e-01+/-8.0e-01 | 2.11e+01+/-6.6e-01 | 2.06e+01+/-4.4e-01 |
 test_objective_compute_heliotron       | +0.63+/-316.88 | +3.33e-04+/-1.7e-01 | 5.35e-02+/-1.2e-01 | 5.32e-02+/-1.2e-01 |
 test_objective_compute_dshape_current  | -1.54+/-373.72 | -4.07e-04+/-9.9e-02 | 2.60e-02+/-6.9e-02 | 2.64e-02+/-7.0e-02 |
 test_objective_compute_atf             | -0.40+/-316.30 | -2.10e-04+/-1.7e-01 | 5.22e-02+/-1.2e-01 | 5.24e-02+/-1.2e-01 |
 test_objective_jac_heliotron           | +0.68+/-1.43 | +4.83e-02+/-1.0e-01 | 7.16e+00+/-1.0e-01 | 7.11e+00+/-1.7e-02 |
 test_objective_jac_dshape_current      | +5.37+/-8.50 | +6.83e-03+/-1.1e-02 | 1.34e-01+/-1.0e-02 | 1.27e-01+/-3.4e-03 |
-test_objective_jac_atf                 | +2.14+/-0.99 | +1.73e-01+/-8.0e-02 | 8.23e+00+/-5.7e-02 | 8.06e+00+/-5.7e-02 |
 test_perturb_1                         | +10.45+/-39.57 | +1.37e+00+/-5.2e+00 | 1.44e+01+/-3.7e+00 | 1.31e+01+/-3.6e+00 |
 test_perturb_2                         | +4.44+/-5.54 | +1.01e+00+/-1.3e+00 | 2.37e+01+/-8.9e-01 | 2.27e+01+/-8.9e-01 |

[dpanici]

was any benchmarking done? like how long it takes to to cmpute iota here versus master?

I did this quickly for the DSHAPE and DSHAPE_current examples, seems it is as expected, in that an equilibrium with an iota profile and an equilibrium with a current profile now take the same time to calculate iota (because the iota profile eq still computes the iota contributions from the geometry/current of the equilibrium, even though it is not necessary)

for the DSHAPE this is a slowdown of about 2X but since in optimization current profiles are almost always used this is fine. the current profile eq's iota compute speed has not changed significantly, it is now slightly slower but only by like 5%