Fix JAX error with `FourierCurrentPotentialField` in Flux objectives #1002

dpanici · 2024-04-19T17:21:44Z

On master, an optimization with a FourierCurrentPotentialField and the QuadraticFlux objective would fail. This happens because in QuadraticFlux.compute, field.compute_magnetic_field is called. If the field needs transforms to evaluate, then these transforms will be created on the fly if they are not provided, resulting in an error.

This PR fixes that by adding jitable=True to the .compute call inside .compute_magnetic_field

It also adds transform as an argument to MagneticField.compute_magnetic_field methods, in preparation for when #1079 is resolved and objectives can pre-compute and pass in transform objects for all magnetic fields easily.

This PR kicks the can down the road of making get_transforms work with MagneticField.compute_magnetic_field objects (to allow us to pre-compute the transforms used for magnetic field computation #1079 ) and instead opts for the simple fix.

…CurrentField to avoid issues when jitting the objective

github-actions · 2024-04-19T19:15:47Z

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +7.06 +/- 8.20     | +3.58e-02 +/- 4.16e-02 |  5.43e-01 +/- 4.1e-02  |  5.07e-01 +/- 6.4e-03  |
 test_build_transform_fft_midres         |     +6.67 +/- 3.29     | +3.96e-02 +/- 1.95e-02 |  6.33e-01 +/- 1.7e-02  |  5.93e-01 +/- 9.9e-03  |
 test_build_transform_fft_highres        |     +3.78 +/- 3.51     | +3.73e-02 +/- 3.46e-02 |  1.02e+00 +/- 1.5e-02  |  9.87e-01 +/- 3.1e-02  |
 test_equilibrium_init_lowres            |     +1.75 +/- 2.66     | +6.43e-02 +/- 9.76e-02 |  3.73e+00 +/- 9.7e-02  |  3.67e+00 +/- 9.8e-03  |
 test_equilibrium_init_medres            |     +1.82 +/- 2.17     | +7.53e-02 +/- 8.98e-02 |  4.22e+00 +/- 8.9e-02  |  4.14e+00 +/- 1.0e-02  |
 test_equilibrium_init_highres           |     +1.21 +/- 1.88     | +6.71e-02 +/- 1.04e-01 |  5.60e+00 +/- 8.5e-02  |  5.54e+00 +/- 6.0e-02  |
 test_objective_compile_dshape_current   |     +0.64 +/- 5.96     | +2.43e-02 +/- 2.26e-01 |  3.81e+00 +/- 2.2e-01  |  3.79e+00 +/- 1.9e-02  |
 test_objective_compile_atf              |     +1.09 +/- 3.08     | +8.93e-02 +/- 2.54e-01 |  8.32e+00 +/- 1.5e-01  |  8.23e+00 +/- 2.0e-01  |
 test_objective_compute_dshape_current   |     -1.81 +/- 4.61     | -2.31e-05 +/- 5.87e-05 |  1.25e-03 +/- 2.9e-05  |  1.27e-03 +/- 5.1e-05  |
 test_objective_compute_atf              |     +0.42 +/- 6.38     | +1.76e-05 +/- 2.70e-04 |  4.25e-03 +/- 2.3e-04  |  4.23e-03 +/- 1.5e-04  |
 test_objective_jac_dshape_current       |     -1.67 +/- 11.51    | -6.06e-04 +/- 4.18e-03 |  3.57e-02 +/- 2.3e-03  |  3.64e-02 +/- 3.5e-03  |
 test_objective_jac_atf                  |     +3.58 +/- 2.49     | +6.67e-02 +/- 4.63e-02 |  1.93e+00 +/- 3.4e-02  |  1.86e+00 +/- 3.2e-02  |
 test_perturb_1                          |     +0.51 +/- 0.68     | +6.71e-02 +/- 8.89e-02 |  1.31e+01 +/- 7.4e-02  |  1.30e+01 +/- 4.9e-02  |
 test_perturb_2                          |     +0.74 +/- 1.23     | +1.33e-01 +/- 2.22e-01 |  1.81e+01 +/- 1.3e-01  |  1.80e+01 +/- 1.8e-01  |
 test_proximal_jac_atf                   |     -0.61 +/- 1.44     | -4.48e-02 +/- 1.05e-01 |  7.28e+00 +/- 6.0e-02  |  7.32e+00 +/- 8.7e-02  |
 test_proximal_freeb_compute             |     -1.67 +/- 0.74     | -3.00e-03 +/- 1.32e-03 |  1.76e-01 +/- 1.0e-03  |  1.79e-01 +/- 8.2e-04  |
 test_proximal_freeb_jac                 |     -0.12 +/- 1.09     | -8.66e-03 +/- 7.98e-02 |  7.31e+00 +/- 5.9e-02  |  7.32e+00 +/- 5.4e-02  |
 test_solve_fixed_iter                   |     -0.38 +/- 9.08     | -5.59e-02 +/- 1.34e+00 |  1.47e+01 +/- 9.9e-01  |  1.48e+01 +/- 9.0e-01  |

codecov · 2024-04-19T20:29:54Z

Codecov Report

Attention: Patch coverage is 78.94737% with 4 lines in your changes missing coverage. Please review.

Project coverage is 94.97%. Comparing base (0a6b995) to head (63e1cd9).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1002      +/-   ##
==========================================
- Coverage   94.98%   94.97%   -0.01%     
==========================================
  Files          87       87              
  Lines       21749    21762      +13     
==========================================
+ Hits        20658    20669      +11     
- Misses       1091     1093       +2

Files	Coverage Δ
desc/geometry/core.py	`95.91% <100.00%> (ø)`
desc/objectives/_coils.py	`99.12% <ø> (ø)`
desc/coils.py	`96.88% <85.71%> (+0.02%)`	⬆️
desc/magnetic_fields/_current_potential.py	`98.83% <75.00%> (-0.58%)`	⬇️
desc/magnetic_fields/_core.py	`96.40% <71.42%> (-0.27%)`	⬇️

... and 1 file with indirect coverage changes

f0uriest

I think the reason coils worked while the current potential doesnt is that by default the coil classes create jitable transforms in their compute method, since they're only 1d its not really worth it to try to do the fft.

I'd recommend adding the new classes to the existing get_transforms logic, I think most of it should "just work" assuming attributes are named correctly. (might need some logic somewhere for tree-like coilsets etc, but do-able.

f0uriest · 2024-04-25T02:14:06Z

desc/objectives/_coils.py

@@ -753,13 +753,28 @@ def build(self, use_jit=True, verbose=1):
            Bplasma = compute_B_plasma(
                eq, eval_grid, self._source_grid, normal_only=True
            )
+        field = self._field
+        if hasattr(field, "Phi_mn"):


would it be better to get this all to work with the existing desc.compute.utils.get_transforms helper? That would work for coils etc as well.

dpanici · 2024-04-30T17:51:33Z

I think the reason coils worked while the current potential doesnt is that by default the coil classes create jitable transforms in their compute method, since they're only 1d its not really worth it to try to do the fft.

I'd recommend adding the new classes to the existing get_transforms logic, I think most of it should "just work" assuming attributes are named correctly. (might need some logic somewhere for tree-like coilsets etc, but do-able.

The problem is less the classes and more that not every magnetic field class needs the same keys to calculate the magnetic field. Coils need "I" and "x", the current potential fields need "K", toroidal fields dont need any keys as they just have their own compute magnetic field function.

dpanici · 2024-05-01T19:37:24Z

add case for coils and 1D transforms (by calling get_transforms in build always asking for "x" (if is a curve class or CoilSet) and maybe also "K" if it is a current potential field.

Also add in the current potential field compute, a jitable=True (as if there is a current potential field in a SumMagneticField it might not work correctly...)

tests/test_optimizer.py

…mpute for current potential field

dpanici · 2024-06-10T17:48:31Z

@f0uriest I don't exactly remember the issue with the current method/how putting the logic in get_transforms would improve it.

We still would need to call get_transforms with the correct grid based off of what the underlying object is (which would still require a check in the build method to see if it is a coil or a field), and to ask for the correct keys ("x" and maybe "K" as well if it is a current potential field).

I think I also found in this PR a fix to something causing a JAX error in a few places when current potential fields are being optimized, so if possible I'd like to get this one in. Unless there is a better case for putting more logic into get_transforms, which I think should wait until we make "B" a data index quantity for MagneticField objects and then can go through that route to get all the proper transforms etc based off of the specific MagneticField parameterization

desc/coils.py

desc/objectives/_coils.py

kianorr · 2024-06-17T16:24:19Z

tests/test_optimizer.py

+
+
+@pytest.mark.unit
+def test_tor_flux_with_surface_current_field():


Could these tests be one test with an inner test() function or is it better to keep them separate?

desc/coils.py

f0uriest · 2024-06-18T21:52:50Z

desc/magnetic_fields/_current_potential.py

@@ -644,7 +647,23 @@ def _compute_magnetic_field_from_CurrentPotentialField(
    # compute surface current, and store grid quantities
    # needed for integration in class
    # TODO: does this have to be xyz, or can it be computed in rpz as well?
-    data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params)
+    if not params and not transforms:


why this conditional?

To avoid calling field.compute inside a compute function of an objective, as there is some logic there that jit does not like

dpanici · 2024-06-19T19:35:19Z

Equinox for runtime error checking in jitted functions

dpanici · 2024-06-19T19:39:03Z

@f0uriest check this again for placing the code in a util function

rahulgaur104 · 2024-07-02T18:44:03Z

Merge after increasing the coverage.

tests/test_optimizer.py

f0uriest · 2024-07-03T01:26:31Z

tests/test_optimizer.py

+
+
+@pytest.mark.unit
+def test_quad_flux_with_surface_current_field():


Is the plan to remove/combine this test with #1025 ?

Not sure, as that one will add coil tests while this one is using surface current fields. If one test on that branch involves a sum magnetic field with coils and a surface current field then we can remove this in favor of that one

ddudt · 2024-07-04T15:51:59Z

desc/objectives/_coils.py

@@ -1272,7 +1272,10 @@ def compute(self, field_params, constants=None):

        # B_ext is not pre-computed because field is not fixed
        B_ext = constants["field"].compute_magnetic_field(
-            x, source_grid=constants["field_grid"], basis="rpz", params=field_params


No real change to this file, right? You just like to make the code longer to bother me?

I blame my IDE

but yea no change here just formatting

- doesn't pass due to jax bug, maybe PR #1002 will fix it?

dpanici added 3 commits April 19, 2024 11:54

add transforms to build method of QuadraticFlux if field is a Surface…

ee27ed5

…CurrentField to avoid issues when jitting the objective

add test for surface current field

a6716bb

add for toroidal flux objective:

46e72c4

add transforms to rest of magnetic fields

74ddc51

f0uriest requested changes Apr 25, 2024

View reviewed changes

dpanici marked this pull request as draft April 25, 2024 15:12

Merge branch 'master' into dp/quadratic-flux-transform

de4c909

Merge branch 'master' into dp/quadratic-flux-transform

e00fb85

github-actions bot reviewed May 31, 2024

View reviewed changes

tests/test_optimizer.py Outdated Show resolved Hide resolved

tests/test_optimizer.py Outdated Show resolved Hide resolved

dpanici and others added 3 commits June 10, 2024 13:23

fix name of fixparameters

27ba594

Merge branch 'master' into dp/quadratic-flux-transform

757e8b2

fix jax error by using _compute directly instead of going through .co…

6ebd349

…mpute for current potential field

dpanici marked this pull request as ready for review June 10, 2024 17:48

dpanici requested review from f0uriest, ddudt, kianorr and YigitElma June 10, 2024 17:48

Merge branch 'master' into dp/quadratic-flux-transform

d84c514

kianorr reviewed Jun 17, 2024

View reviewed changes

dpanici and others added 2 commits June 18, 2024 11:45

change compute to _compute

9334e09

Merge branch 'master' into dp/quadratic-flux-transform

48e3abb

f0uriest reviewed Jun 18, 2024

View reviewed changes

Merge branch 'master' into dp/quadratic-flux-transform

962db59

add jitable to surface compute

3a95922

dpanici mentioned this pull request Jul 2, 2024

Add jitable to surface compute to fix JAX error #1092

Closed

dpanici requested review from f0uriest, ddudt and kianorr and removed request for f0uriest, ddudt and YigitElma July 2, 2024 14:36

dpanici changed the title ~~Add transform pre-computation for FourierCurrentPotentialField in Flux objectives~~ Fix JAX error with FourierCurrentPotentialField in Flux objectives Jul 2, 2024

dpanici added the EZ-review label Jul 2, 2024

rahulgaur104 previously approved these changes Jul 2, 2024

View reviewed changes

kianorr previously approved these changes Jul 2, 2024

View reviewed changes

dpanici mentioned this pull request Jul 2, 2024

compute quantities along a fieldline #1024

Merged

dpanici and others added 2 commits July 2, 2024 19:52

Merge branch 'master' into dp/quadratic-flux-transform

e384a17

increase coverage of tests

07d110b

dpanici dismissed stale reviews from kianorr and rahulgaur104 via 07d110b July 3, 2024 00:02

dpanici requested review from rahulgaur104 and kianorr July 3, 2024 00:02

f0uriest reviewed Jul 3, 2024

View reviewed changes

add comment

6a25959

f0uriest approved these changes Jul 3, 2024

View reviewed changes

Merge branch 'master' into dp/quadratic-flux-transform

63e1cd9

dpanici requested a review from daniel-dudt July 4, 2024 15:33

ddudt approved these changes Jul 4, 2024

View reviewed changes

kianorr added a commit that referenced this pull request Jul 4, 2024

add QuadraticFlux obj

680848b

- doesn't pass due to jax bug, maybe PR #1002 will fix it?

kianorr mentioned this pull request Jul 4, 2024

Ko/test coil opt #1099

Merged

dpanici merged commit f021ce2 into master Jul 5, 2024
17 of 18 checks passed

dpanici deleted the dp/quadratic-flux-transform branch July 5, 2024 13:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix JAX error with `FourierCurrentPotentialField` in Flux objectives #1002

Fix JAX error with `FourierCurrentPotentialField` in Flux objectives #1002

dpanici commented Apr 19, 2024 •

edited

Loading

github-actions bot commented Apr 19, 2024 •

edited

Loading

codecov bot commented Apr 19, 2024 •

edited

Loading

f0uriest left a comment

f0uriest Apr 25, 2024

dpanici commented Apr 30, 2024

dpanici commented May 1, 2024 •

edited

Loading

dpanici commented Jun 10, 2024

kianorr Jun 17, 2024

f0uriest Jun 18, 2024

dpanici Jun 20, 2024

dpanici commented Jun 19, 2024

dpanici commented Jun 19, 2024

rahulgaur104 commented Jul 2, 2024

f0uriest Jul 3, 2024

dpanici Jul 3, 2024

ddudt Jul 4, 2024

dpanici Jul 5, 2024

dpanici Jul 5, 2024



		@pytest.mark.unit
		def test_tor_flux_with_surface_current_field():



		@pytest.mark.unit
		def test_quad_flux_with_surface_current_field():

Fix JAX error with FourierCurrentPotentialField in Flux objectives #1002

Fix JAX error with FourierCurrentPotentialField in Flux objectives #1002

Conversation

dpanici commented Apr 19, 2024 • edited Loading

github-actions bot commented Apr 19, 2024 • edited Loading

codecov bot commented Apr 19, 2024 • edited Loading

Codecov Report

f0uriest left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dpanici commented Apr 30, 2024

dpanici commented May 1, 2024 • edited Loading

dpanici commented Jun 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dpanici commented Jun 19, 2024

dpanici commented Jun 19, 2024

rahulgaur104 commented Jul 2, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Fix JAX error with `FourierCurrentPotentialField` in Flux objectives #1002

Fix JAX error with `FourierCurrentPotentialField` in Flux objectives #1002

dpanici commented Apr 19, 2024 •

edited

Loading

github-actions bot commented Apr 19, 2024 •

edited

Loading

codecov bot commented Apr 19, 2024 •

edited

Loading

dpanici commented May 1, 2024 •

edited

Loading