Fix bug in boozer coordinate computation for Omnigenous magnetic field #1166

unalmis · 2024-08-06T05:29:44Z

Move test compute everything to own file and exclude source grid quantities. from commit 8a8d9de.
Remove xyz basis file in favor of correctness check to catch bug Quantities are computed wrong unless returned data is in cylindrical basis vectors #1088
Debug zeta_B computation. zeta_B computation is sometimes wrong #1163 has been reproduced here (and on master by Dario). This issue was resolved after fixing the computation of zeta_B and theta_B for omnigenous fields.

tests/test_compute_everything.py

github-actions · 2024-08-06T06:49:51Z

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +5.26 +/- 7.49     | +2.77e-02 +/- 3.94e-02 |  5.54e-01 +/- 3.5e-02  |  5.26e-01 +/- 1.7e-02  |
 test_build_transform_fft_midres         |     +0.83 +/- 5.68     | +5.23e-03 +/- 3.57e-02 |  6.33e-01 +/- 2.5e-02  |  6.28e-01 +/- 2.5e-02  |
 test_build_transform_fft_highres        |     -0.15 +/- 3.48     | -1.66e-03 +/- 3.79e-02 |  1.09e+00 +/- 2.8e-02  |  1.09e+00 +/- 2.6e-02  |
 test_equilibrium_init_lowres            |     +3.89 +/- 7.34     | +1.50e-01 +/- 2.83e-01 |  4.01e+00 +/- 1.6e-01  |  3.85e+00 +/- 2.3e-01  |
 test_equilibrium_init_medres            |     +2.17 +/- 5.53     | +9.59e-02 +/- 2.44e-01 |  4.51e+00 +/- 1.7e-01  |  4.41e+00 +/- 1.8e-01  |
 test_equilibrium_init_highres           |     +4.14 +/- 3.41     | +2.34e-01 +/- 1.93e-01 |  5.89e+00 +/- 1.4e-01  |  5.66e+00 +/- 1.3e-01  |
 test_objective_compile_dshape_current   |     +0.40 +/- 2.20     | +1.59e-02 +/- 8.80e-02 |  4.02e+00 +/- 8.1e-02  |  4.00e+00 +/- 3.5e-02  |
 test_objective_compile_atf              |     -1.27 +/- 2.57     | -1.09e-01 +/- 2.21e-01 |  8.49e+00 +/- 1.0e-01  |  8.60e+00 +/- 2.0e-01  |
 test_objective_compute_dshape_current   |     -0.35 +/- 3.51     | -4.37e-06 +/- 4.42e-05 |  1.26e-03 +/- 2.4e-05  |  1.26e-03 +/- 3.7e-05  |
 test_objective_compute_atf              |     -3.75 +/- 6.09     | -1.65e-04 +/- 2.68e-04 |  4.23e-03 +/- 1.3e-04  |  4.39e-03 +/- 2.3e-04  |
 test_objective_jac_dshape_current       |     -7.47 +/- 9.23     | -3.01e-03 +/- 3.72e-03 |  3.73e-02 +/- 3.2e-03  |  4.03e-02 +/- 2.0e-03  |
 test_objective_jac_atf                  |     -7.57 +/- 3.10     | -1.50e-01 +/- 6.16e-02 |  1.84e+00 +/- 3.1e-02  |  1.99e+00 +/- 5.3e-02  |
 test_perturb_1                          |     -2.70 +/- 1.96     | -3.88e-01 +/- 2.82e-01 |  1.40e+01 +/- 1.1e-01  |  1.44e+01 +/- 2.6e-01  |
 test_perturb_2                          |     -0.71 +/- 2.75     | -1.40e-01 +/- 5.40e-01 |  1.95e+01 +/- 2.3e-01  |  1.96e+01 +/- 4.9e-01  |
 test_proximal_jac_atf                   |     -2.04 +/- 1.26     | -1.54e-01 +/- 9.53e-02 |  7.38e+00 +/- 3.5e-02  |  7.54e+00 +/- 8.9e-02  |
 test_proximal_freeb_compute             |     -1.89 +/- 1.15     | -3.43e-03 +/- 2.09e-03 |  1.78e-01 +/- 1.2e-03  |  1.81e-01 +/- 1.7e-03  |
 test_proximal_freeb_jac                 |     -0.03 +/- 1.33     | -2.29e-03 +/- 9.91e-02 |  7.46e+00 +/- 5.4e-02  |  7.47e+00 +/- 8.3e-02  |
 test_solve_fixed_iter                   |     -0.66 +/- 5.69     | -1.24e-01 +/- 1.07e+00 |  1.87e+01 +/- 7.1e-01  |  1.88e+01 +/- 8.1e-01  |

dpanici · 2024-08-06T15:33:03Z

I can't seem to reproduce #1163 when trying on a CPU compute node on della, though it actually fails with very slim margins for some of the current density and force compoents. You are running on linux locally I assume as well?

unalmis · 2024-08-06T16:30:11Z

I can't seem to reproduce #1163 when trying on a CPU compute node on della, though it actually fails with very slim margins for some of the current density and force compoents. You are running on linux locally I assume as well?

Yes. Fedora linux. CPU, jax 0.4.30, jaxlib 0.4.28.

dpanici · 2024-08-06T16:35:13Z

I can't seem to reproduce #1163 when trying on a CPU compute node on della, though it actually fails with very slim margins for some of the current density and force compoents. You are running on linux locally I assume as well?

Yes. Fedora linux. CPU, jax 0.4.30, jaxlib 0.4.28.

wait I cannot read, yes I also get that zeta_B is off by large margins, AND it is not consistent (running it twice, one it failed and the other it did not). Using older versions of jax/jaxlib (0.4.14 for both)

dpanici · 2024-08-07T19:29:18Z

take a look at the omni field in test compute everything, see if it looks weird, try changing params and see if it is more stable

dpanici · 2024-08-07T19:29:26Z

@ddudt

codecov · 2024-08-07T22:25:24Z

Codecov Report

Attention: Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.

Project coverage is 95.46%. Comparing base (cbf9ebb) to head (a0eb85a).
Report is 1764 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/compute/_geometry.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1166      +/-   ##
==========================================
+ Coverage   95.15%   95.46%   +0.31%     
==========================================
  Files          87       87              
  Lines       22281    22243      -38     
==========================================
+ Hits        21202    21235      +33     
+ Misses       1079     1008      -71

Files with missing lines	Coverage Δ
desc/compute/_omnigenity.py	`100.00% <100.00%> (ø)`
desc/compute/data_index.py	`97.05% <ø> (ø)`
desc/plotting.py	`95.80% <100.00%> (-0.01%)`	⬇️
desc/compute/_geometry.py	`99.51% <66.66%> (+16.72%)`	⬆️

... and 6 files with indirect coverage changes

f0uriest · 2024-08-08T16:11:15Z

I do see some small "blobs" near theta=0 and theta=2pi

Not sure if this is a plotting artifact, or could be something weird with the coordinate mapping there?

ddudt · 2024-08-08T17:09:17Z

I do see some small "blobs" near theta=0 and theta=2pi

Not sure if this is a plotting artifact, or could be something weird with the coordinate mapping there?

I think that is just an artifact of plotting with tricontour, although there was a small bug I just fixed that shouldn't change much.

ddudt · 2024-08-08T17:28:17Z

I think I figured out the issue: it has to do with how "zeta_B" is an alias of "theeta_B". These two quantities are always computed simultaneously -- you can't compute one without also getting the other. So we have the compute function with name="theta_B" and aliases=["zeta_B"]. This works fine if you always ask to compute both quantities or even only ask for "theta_B". But for some reason, it looks like when you ask to only compute "zeta_B" that quantity in the compute data is being overwritten with the values for "theta_B".

I'm not sure what the solution is, but I'll let you figure that out.

unalmis · 2024-08-08T18:04:07Z

I think I figured out the issue: it has to do with how "zeta_B" is an alias of "theeta_B". These two quantities are always computed simultaneously -- you can't compute one without also getting the other. So we have the compute function with name="theta_B" and aliases=["zeta_B"]. This works fine if you always ask to compute both quantities or even only ask for "theta_B". But for some reason, it looks like when you ask to only compute "zeta_B" that quantity in the compute data is being overwritten with the values for "theta_B".

I'm not sure what the solution is, but I'll let you figure that out.

Why are these aliases?
We compute

@register_compute_fun(
    name="theta_B",
    label="(\\theta_{B},\\zeta_{B})",
    units="rad",
    units_long="radians",
    description="Boozer angular coordinates",
    dim=1,
    params=[],
    transforms={},
    profiles=[],
    coordinates="rtz",
    data=["alpha", "h"],
    aliases=["zeta_B"], # should remove this
    parameterization="desc.magnetic_fields._core.OmnigenousField",
)
....
    # solve for (theta_B,zeta_B) corresponding to (eta,alpha)
    booz = matrix @ jnp.vstack((data["alpha"], data["h"]))
    data["theta_B"] = booz[0, :]
    data["zeta_B"] = booz[1, :]

Unless booz[1, :] is the same as booz[0, :], we shouldnot mark them as aliases. Otherwise they will get overridden by one another

rahulgaur104 · 2024-08-08T18:39:53Z

I think that is just an artifact of plotting with tricontour, although there was a small bug I just fixed that shouldn't change much.

@ddudt That doesn't make sense. I plotted the same thing with regular contour and it still persists. These is something else going on with those blobs at the ends.

Here's what the target looks like now
It's actually worse now. It has nothing to do with the field that I used. You can use a better field object with sensible B_lm etc. but that doesn't solve #1143 or #1144

ddudt · 2024-08-09T04:41:33Z

I think I figured out the issue: it has to do with how "zeta_B" is an alias of "theeta_B". These two quantities are always computed simultaneously -- you can't compute one without also getting the other. So we have the compute function with name="theta_B" and aliases=["zeta_B"]. This works fine if you always ask to compute both quantities or even only ask for "theta_B". But for some reason, it looks like when you ask to only compute "zeta_B" that quantity in the compute data is being overwritten with the values for "theta_B".
I'm not sure what the solution is, but I'll let you figure that out.

Why are these aliases? We compute
@register_compute_fun(
    name="theta_B",
    label="(\\theta_{B},\\zeta_{B})",
    units="rad",
    units_long="radians",
    description="Boozer angular coordinates",
    dim=1,
    params=[],
    transforms={},
    profiles=[],
    coordinates="rtz",
    data=["alpha", "h"],
    aliases=["zeta_B"], # should remove this
    parameterization="desc.magnetic_fields._core.OmnigenousField",
)
....
    # solve for (theta_B,zeta_B) corresponding to (eta,alpha)
    booz = matrix @ jnp.vstack((data["alpha"], data["h"]))
    data["theta_B"] = booz[0, :]
    data["zeta_B"] = booz[1, :]
Unless booz[1, :] is the same as booz[0, :], we shouldnot mark them as aliases. Otherwise they will get overridden by one another

They are aliases in the sense that if you ask for one of the two quantities, you get both. The problem is with the alias logic outside of the compute function, not in the compute function itself.

I don't like this solution -- it creates redundant code and computations. Maybe we shouldn't call them aliases, but I'm sure we can think of another way to have a single compute function return multiple variables.

desc/compute/_omnigenity.py

unalmis · 2024-08-09T05:05:00Z

They are aliases in the sense that if you ask for one of the two quantities, you get both. The problem is with the alias logic outside of the compute function, not in the compute function itself.

Marking some quantity A as an alias of B means that A = B.

I don't like this solution -- it creates redundant code and computations.

The redundancy is cosmetic. The computation is only done once.

Maybe we shouldn't call them aliases, but I'm sure we can think of another way to have a single compute function return multiple variables.

This would require plumbing that is outside the scope of this PR. In particular, the register_compute_fun decorator would need to register multiple quantities with multiple descriptions, dimensions, etc. I have made a new issue for this.

There is a failing regression test now. My guess is it's due to the plotting change and not the compute function change? Could you look into this sometime

Before:

After:

ddudt · 2024-08-09T14:15:03Z

I think the following will work. It's similar to your change, but avoids needing the duplicate code:

@register_compute_fun(
    name="theta_B",
    label="\\theta_{B}",
    units="rad",
    units_long="radians",
    description="Boozer poloidal coordinate",
    dim=1,
    params=[],
    transforms={},
    profiles=[],
    coordinates="rtz",
    data=["alpha", "h"],
    parameterization="desc.magnetic_fields._core.OmnigenousField",
    helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)",
    iota="float: Value of rotational transform on the Omnigenous surface. Default 1.0",
)
def _omni_map_theta(params, transforms, profiles, data, **kwargs):
    M, N = kwargs.get("helicity", (1, 0))
    iota = kwargs.get("iota", 1)

    # coordinate mapping matrix from (alpha,h) to (theta_B,zeta_B)
    # need a bunch of wheres to avoid division by zero causing NaN in backward pass
    # this is fine since the incorrect values get ignored later, except in OT or OH
    # where fieldlines are exactly parallel to |B| contours, but this is a degenerate
    # case of measure 0 so this kludge shouldn't affect things too much.
    mat_OP = jnp.array(
        [[N, iota / jnp.where(N == 0, 1, N)], [0, 1 / jnp.where(N == 0, 1, N)]]
    )
    mat_OT = jnp.array([[0, -1], [M, -1 / jnp.where(iota == 0, 1.0, iota)]])
    den = jnp.where((N - M * iota) == 0, 1.0, (N - M * iota))
    mat_OH = jnp.array([[N, M * iota / den], [M, M / den]])
    matrix = jnp.where(
        M == 0,
        mat_OP,
        jnp.where(
            N == 0,
            mat_OT,
            mat_OH,
        ),
    )

    # solve for (theta_B,zeta_B) corresponding to (eta,alpha)
    booz = matrix @ jnp.vstack((data["alpha"], data["h"]))
    data["theta_B"] = booz[0, :]
    data["zeta_B"] = booz[1, :]
    return data


@register_compute_fun(
    name="zeta_B",
    label="\\zeta_{B}",
    units="rad",
    units_long="radians",
    description="Boozer toroidal coordinate",
    dim=1,
    params=[],
    transforms={},
    profiles=[],
    coordinates="rtz",
    data=["theta_B"],
    parameterization="desc.magnetic_fields._core.OmnigenousField",
)
def _omni_map_zeta(params, transforms, profiles, data, **kwargs):
    # some dummy line of unused code that references the dependency theta_B
    _ = data["theta_B"]
    return data

Revert changes to tesd_data_index from #1027 https://github.com/PlasmaControl/DESC/pull/1027/files/148d08d97cf6c08b922d0883700b2758dc9b000c#r1704624731

rahulgaur104

Omnigenity gives me the same results as before. Approving!

dpanici · 2024-08-09T20:58:17Z

desc/plotting.py

        grid_compute = _get_grid(**grid_kwargs)
    if grid_plot is None:
        grid_kwargs = {
            "rho": rho,
            "theta": 91,
            "zeta": 91,
            "NFP": thing.NFP,
-            "endpoint": True,
+            "endpoint": eq_switch,


why can't there be an endpoint in the plot grid for the Omnigenous field?

The old code used grid_compute for the Omnigenous field, which also had endpoint=False so this doesn't change that. The only thing this changes is to correctly plot on grid_plot instead of the grid_compute.

In theory I think it should be able to use endpoint=True, but it was giving a worse looking plot for the test. I suspect it's an issue with tricontour but I'm not sure.

dpanici · 2024-08-09T20:59:17Z

If this resolves any of the issues on boozer stuff, make sure to link them to this

Cherry pick changes from commit #8a8d9de to this commit

88d6b8f

unalmis changed the title ~~Cherry pick changes from commit #8a8d9de to this commit~~ Test compute everything Aug 6, 2024

unalmis commented Aug 6, 2024

View reviewed changes

tests/test_compute_everything.py Outdated Show resolved Hide resolved

unalmis assigned ddudt Aug 7, 2024

unalmis linked an issue Aug 7, 2024 that may be closed by this pull request

zeta_B computation is sometimes wrong #1163

Closed

Add test to ensure basis conversion in #1088 is resolved

db10448

Merge branch 'master' into test_compute

c34e180

fix bug in plot_boozer_surface for OmnigenousField

c53a134

Fix bug in boozer coordinate computation for Omnigenous magnetic field

85f28a9

unalmis added the bug fix Something was fixed label Aug 8, 2024

unalmis changed the title ~~Test compute everything~~ Fix bug in boozer coordinate computation for Omnigenous magnetic field Aug 8, 2024

Merge branch 'master' into test_compute

48dd52e

unalmis requested review from dpanici, f0uriest and ddudt August 8, 2024 18:22

This was referenced Aug 8, 2024

Misaligned target and eq fields in omnigenity objective #1143

Open

Overlapping target field in omnigenity #1144

Open

unalmis requested a review from rahulgaur104 August 8, 2024 18:32

f0uriest previously approved these changes Aug 9, 2024

View reviewed changes

ddudt requested changes Aug 9, 2024

View reviewed changes

desc/compute/_omnigenity.py Show resolved Hide resolved

daniel-dudt added 2 commits August 9, 2024 08:51

repair plotting test

bee9f54

simplify omnigenity compute funs

32e39f1

daniel-dudt dismissed f0uriest’s stale review via 32e39f1 August 9, 2024 14:54

unalmis added 2 commits August 9, 2024 13:54

Update test_data_index for changes in #1000 and

1573297

Revert changes to tesd_data_index from #1027 https://github.com/PlasmaControl/DESC/pull/1027/files/148d08d97cf6c08b922d0883700b2758dc9b000c#r1704624731

Cosmetic change to reduce duplicate code

8ea83c6

unalmis linked an issue Aug 9, 2024 that may be closed by this pull request

Have a single compute function register multiple variables #1171

Open

2 tasks

unalmis requested a review from ddudt August 9, 2024 18:03

rahulgaur104 approved these changes Aug 9, 2024

View reviewed changes

ddudt approved these changes Aug 9, 2024

View reviewed changes

Merge branch 'master' into test_compute

a0eb85a

f0uriest approved these changes Aug 9, 2024

View reviewed changes

dpanici reviewed Aug 9, 2024

View reviewed changes

dpanici approved these changes Aug 9, 2024

View reviewed changes

f0uriest merged commit b622eae into master Aug 9, 2024
17 of 18 checks passed

f0uriest deleted the test_compute branch August 9, 2024 23:15

ddudt mentioned this pull request Aug 10, 2024

Have a single compute function register multiple variables #1171

Open

2 tasks

unalmis mentioned this pull request Aug 14, 2024

Reverse-Mode Differentiable Ideal-Ballooning Stability Solver #1170

Merged

8 tasks

unalmis self-assigned this Aug 17, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix bug in boozer coordinate computation for Omnigenous magnetic field #1166

Fix bug in boozer coordinate computation for Omnigenous magnetic field #1166

unalmis commented Aug 6, 2024 •

edited

Loading

github-actions bot commented Aug 6, 2024 •

edited

Loading

dpanici commented Aug 6, 2024

unalmis commented Aug 6, 2024

dpanici commented Aug 6, 2024

dpanici commented Aug 7, 2024

dpanici commented Aug 7, 2024

codecov bot commented Aug 7, 2024 •

edited

Loading

f0uriest commented Aug 8, 2024

ddudt commented Aug 8, 2024

ddudt commented Aug 8, 2024

unalmis commented Aug 8, 2024 •

edited

Loading

rahulgaur104 commented Aug 8, 2024 •

edited

Loading

ddudt commented Aug 9, 2024

unalmis commented Aug 9, 2024 •

edited

Loading

ddudt commented Aug 9, 2024

rahulgaur104 left a comment

dpanici Aug 9, 2024

ddudt Aug 9, 2024 •

edited

Loading

dpanici commented Aug 9, 2024

Fix bug in boozer coordinate computation for Omnigenous magnetic field #1166

Fix bug in boozer coordinate computation for Omnigenous magnetic field #1166

Conversation

unalmis commented Aug 6, 2024 • edited Loading

github-actions bot commented Aug 6, 2024 • edited Loading

dpanici commented Aug 6, 2024

unalmis commented Aug 6, 2024

dpanici commented Aug 6, 2024

dpanici commented Aug 7, 2024

dpanici commented Aug 7, 2024

codecov bot commented Aug 7, 2024 • edited Loading

Codecov Report

f0uriest commented Aug 8, 2024

ddudt commented Aug 8, 2024

ddudt commented Aug 8, 2024

unalmis commented Aug 8, 2024 • edited Loading

rahulgaur104 commented Aug 8, 2024 • edited Loading

ddudt commented Aug 9, 2024

unalmis commented Aug 9, 2024 • edited Loading

ddudt commented Aug 9, 2024

rahulgaur104 left a comment

Choose a reason for hiding this comment

dpanici Aug 9, 2024

Choose a reason for hiding this comment

ddudt Aug 9, 2024 • edited Loading

Choose a reason for hiding this comment

dpanici commented Aug 9, 2024

unalmis commented Aug 6, 2024 •

edited

Loading

github-actions bot commented Aug 6, 2024 •

edited

Loading

codecov bot commented Aug 7, 2024 •

edited

Loading

unalmis commented Aug 8, 2024 •

edited

Loading

rahulgaur104 commented Aug 8, 2024 •

edited

Loading

unalmis commented Aug 9, 2024 •

edited

Loading

ddudt Aug 9, 2024 •

edited

Loading