Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Ability to Write Boozmn.nc style output files #680

Draft
wants to merge 100 commits into
base: master
Choose a base branch
from

Conversation

dpanici
Copy link
Collaborator

@dpanici dpanici commented Sep 19, 2023

Merge in #1211 first, THEN merge in #1197 , then modify this PR to use the vectorized boozer transform from #1197, then it is ready to merge

Add make_boozmn_output utility function to vmec_utils, which allows DESC to write boozmn.nc style output files containing Boozer harmonics. I/O based heavily off of the hidden symmetries version of booz_xform

  • Add test of DESC-computed quantities versus boozmn file written - very good comparison
  • Add test of our boozmn.nc against an actual C++ version of booz_xform output file at same resolution/nsurfs and equilibrium - comparison not great, need to debug still (check order of m,n possibly)
  • improve speed by computing all data at once
  • make the prefactor and ["B"] basis evaluated at the necessary rho,theta_B,zeta_B nodes a name in data index (so dont need to recompute the sum of basis modes each time)
  • add non-symmetric test and fix for non-symmetric
  • add documentation (in form of notebook or a page) on our Boozer transform setup and how to use and the meaning of our "|B|_mn_B" and how to get into the familiar double angle form
  • update algorithm to use vectorized boozer transform now that Vectorize boozer transform over multiple surfaces #1197 is merged

future work (or work if someone else wants to do on this PR go ahead)

  • add capacity to save info at any surface, not just the half grid? though would need to update docstrings for everything saved... and unsure of who would use that as anyone using the .nc likely just wants it in the usual format with the half grid - pushing to a future PR
  • cannot precompute quantities that assume a single surface is passed in, i.e. theta_B, zeta_B which need nu which needs w_Boozer which needs B_theta_mn.... which itself is found with transforms["B"].fit("B_theta"), so that quantity cannot be computed using a volumetric grid right now (maybe, can FourierZernike fit then evaluate the fourier series? and pass that into the compute fxn)
  • precomputing transforms["B"], or possibly using same transform and modifying the basis matrices by multiplying by zernike_radial(rho_i)/zernike_radial(rho_{i-1}) (unsure if this second part is really possible, for now I jitted the transform build and it works reasonably well)
  • Faster Zernike evaluation #720 could help speed this up as a nontrivial amount of time is spent in zernike_to_fourier - pushing to a future PR

I will note that when comparing against the Cpp or Fortran version, DESC will always be more accurate as the other versions must interpolate the Fourier coefficients to get the Rmnc etc on the half grid, while we can exactly calculate what they are on the half grid because we know quantities for every rho in [0,1]

timing info

The time it takes to save the HELIOTRON example to a wout file with ns=100 is ~40s, then running the C++ version of booz_xform with M_Booz=N_Booz=20 on a single thread (remember it is multithreaded so this) is very quick (10s) and on 8 threads is 2s

That is compared to running the make_boozmn_output function in DESC, which takes 90s (50s on gpu though and the util function uses numpy so a lot of memory movement is happening). So there is definitely room for improvement here, on the CPU and on GPU, but at least on GPU it is comparable speed.

@dpanici
Copy link
Collaborator Author

dpanici commented Sep 21, 2023

make a norm name in data index

@dpanici dpanici changed the title first commit, machinery working but need to figure out good way to fi… Add Ability to Write Boozmn.nc style output files Sep 21, 2023
@dpanici
Copy link
Collaborator Author

dpanici commented Sep 27, 2023

notes to self on Booz_xform

modenumbers are defined here:

  • if m=0, they only include n from 0 -> nboz
  • else if m is not 0, they include all n modes from -nboz -> nboz
  • this means that they include the (0,0) mode even for their sin terms, but they skip the (0,-nboz -_ nboz) modes. This may explain the discrepancy in modes I see

@dpanici dpanici marked this pull request as draft October 11, 2023 19:33
Copy link
Contributor

github-actions bot commented Nov 8, 2023

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.64 +/- 2.10     | +9.97e-05 +/- 3.25e-04 |  1.56e-02 +/- 2.9e-04  |  1.55e-02 +/- 1.5e-04  |
 test_build_transform_fft_midres         |     +0.16 +/- 1.35     | +1.55e-04 +/- 1.29e-03 |  9.51e-02 +/- 1.2e-03  |  9.50e-02 +/- 4.9e-04  |
 test_build_transform_fft_highres        |     +1.14 +/- 0.84     | +5.29e-03 +/- 3.88e-03 |  4.69e-01 +/- 3.3e-03  |  4.63e-01 +/- 2.1e-03  |
 test_equilibrium_init_lowres            |     +0.48 +/- 1.46     | +3.14e-03 +/- 9.60e-03 |  6.61e-01 +/- 7.5e-03  |  6.58e-01 +/- 6.0e-03  |
 test_equilibrium_init_medres            |     -0.17 +/- 0.92     | -1.74e-03 +/- 9.23e-03 |  1.00e+00 +/- 5.7e-03  |  1.01e+00 +/- 7.2e-03  |
 test_equilibrium_init_highres           |     +0.41 +/- 1.07     | +1.04e-02 +/- 2.74e-02 |  2.57e+00 +/- 1.6e-02  |  2.56e+00 +/- 2.2e-02  |
 test_objective_compile_dshape_current   |     -1.97 +/- 6.87     | -8.34e-02 +/- 2.91e-01 |  4.16e+00 +/- 2.1e-01  |  4.24e+00 +/- 2.0e-01  |
 test_objective_compile_atf              |     -3.31 +/- 4.88     | -3.18e-01 +/- 4.69e-01 |  9.29e+00 +/- 3.4e-01  |  9.61e+00 +/- 3.2e-01  |
 test_objective_compute_dshape_current   |     +0.07 +/- 2.30     | +1.56e-06 +/- 5.03e-05 |  2.19e-03 +/- 3.3e-05  |  2.19e-03 +/- 3.8e-05  |
 test_objective_compute_atf              |     -1.34 +/- 2.16     | -1.04e-04 +/- 1.67e-04 |  7.63e-03 +/- 1.0e-04  |  7.73e-03 +/- 1.3e-04  |
 test_objective_jac_dshape_current       |     -6.37 +/- 7.16     | -3.01e-03 +/- 3.38e-03 |  4.42e-02 +/- 1.5e-03  |  4.72e-02 +/- 3.0e-03  |
 test_objective_jac_atf                  |     -7.77 +/- 4.22     | -1.86e-01 +/- 1.01e-01 |  2.21e+00 +/- 5.0e-02  |  2.39e+00 +/- 8.8e-02  |
 test_perturb_1                          |     -1.10 +/- 11.89    | -9.66e-02 +/- 1.04e+00 |  8.67e+00 +/- 7.6e-01  |  8.77e+00 +/- 7.2e-01  |
 test_perturb_2                          |     -2.40 +/- 5.40     | -3.52e-01 +/- 7.91e-01 |  1.43e+01 +/- 6.2e-01  |  1.47e+01 +/- 4.9e-01  |

Copy link
Contributor

github-actions bot commented Nov 9, 2023

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.74 +/- 1.72     | +1.13e-04 +/- 2.63e-04 |  1.54e-02 +/- 2.1e-04  |  1.53e-02 +/- 1.5e-04  |
 test_build_transform_fft_midres         |     +0.54 +/- 1.12     | +5.06e-04 +/- 1.06e-03 |  9.44e-02 +/- 8.6e-04  |  9.39e-02 +/- 6.1e-04  |
 test_build_transform_fft_highres        |     +0.73 +/- 0.94     | +3.36e-03 +/- 4.34e-03 |  4.65e-01 +/- 3.9e-03  |  4.62e-01 +/- 2.0e-03  |
 test_equilibrium_init_lowres            |     +2.80 +/- 2.13     | +1.84e-02 +/- 1.40e-02 |  6.75e-01 +/- 1.3e-02  |  6.57e-01 +/- 5.4e-03  |
 test_equilibrium_init_medres            |     +0.00 +/- 2.85     | +3.23e-05 +/- 2.83e-02 |  9.94e-01 +/- 2.6e-02  |  9.94e-01 +/- 1.2e-02  |
 test_equilibrium_init_highres           |     +0.07 +/- 0.69     | +1.89e-03 +/- 1.78e-02 |  2.56e+00 +/- 1.1e-02  |  2.56e+00 +/- 1.4e-02  |
 test_objective_compile_dshape_current   |     -1.40 +/- 6.95     | -5.85e-02 +/- 2.91e-01 |  4.14e+00 +/- 2.5e-01  |  4.19e+00 +/- 1.6e-01  |
 test_objective_compile_atf              |     -0.55 +/- 8.21     | -5.14e-02 +/- 7.62e-01 |  9.22e+00 +/- 5.6e-01  |  9.27e+00 +/- 5.2e-01  |
 test_objective_compute_dshape_current   |     +0.83 +/- 2.06     | +1.80e-05 +/- 4.48e-05 |  2.19e-03 +/- 2.5e-05  |  2.17e-03 +/- 3.7e-05  |
 test_objective_compute_atf              |     -1.90 +/- 1.03     | -1.45e-04 +/- 7.82e-05 |  7.48e-03 +/- 4.9e-05  |  7.62e-03 +/- 6.1e-05  |
 test_objective_jac_dshape_current       |     +5.20 +/- 10.14    | +2.31e-03 +/- 4.51e-03 |  4.68e-02 +/- 3.6e-03  |  4.45e-02 +/- 2.8e-03  |
 test_objective_jac_atf                  |     -0.24 +/- 6.29     | -5.71e-03 +/- 1.47e-01 |  2.34e+00 +/- 5.8e-02  |  2.34e+00 +/- 1.4e-01  |
 test_perturb_1                          |     -1.35 +/- 12.01    | -1.16e-01 +/- 1.03e+00 |  8.45e+00 +/- 7.3e-01  |  8.57e+00 +/- 7.3e-01  |
 test_perturb_2                          |     -1.29 +/- 3.45     | -1.85e-01 +/- 4.95e-01 |  1.42e+01 +/- 4.1e-01  |  1.44e+01 +/- 2.8e-01  |

Copy link
Contributor

github-actions bot commented Nov 9, 2023

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +2.27 +/- 9.81     | +3.48e-04 +/- 1.50e-03 |  1.56e-02 +/- 1.5e-03  |  1.53e-02 +/- 2.7e-04  |
 test_build_transform_fft_midres         |     +2.74 +/- 1.39     | +2.54e-03 +/- 1.29e-03 |  9.53e-02 +/- 1.2e-03  |  9.28e-02 +/- 3.1e-04  |
 test_build_transform_fft_highres        |     +0.95 +/- 1.08     | +4.36e-03 +/- 4.96e-03 |  4.62e-01 +/- 3.9e-03  |  4.58e-01 +/- 3.0e-03  |
 test_equilibrium_init_lowres            |     -0.12 +/- 1.35     | -7.75e-04 +/- 8.80e-03 |  6.49e-01 +/- 6.6e-03  |  6.50e-01 +/- 5.8e-03  |
 test_equilibrium_init_medres            |     -0.04 +/- 0.92     | -3.69e-04 +/- 9.11e-03 |  9.88e-01 +/- 5.7e-03  |  9.89e-01 +/- 7.1e-03  |
 test_equilibrium_init_highres           |     +1.01 +/- 1.18     | +2.56e-02 +/- 2.98e-02 |  2.56e+00 +/- 2.6e-02  |  2.53e+00 +/- 1.5e-02  |
 test_objective_compile_dshape_current   |     +4.79 +/- 7.87     | +1.96e-01 +/- 3.22e-01 |  4.29e+00 +/- 2.2e-01  |  4.10e+00 +/- 2.3e-01  |
 test_objective_compile_atf              |     +2.39 +/- 5.66     | +2.23e-01 +/- 5.27e-01 |  9.54e+00 +/- 4.4e-01  |  9.31e+00 +/- 2.9e-01  |
 test_objective_compute_dshape_current   |     +1.45 +/- 2.72     | +3.17e-05 +/- 5.94e-05 |  2.22e-03 +/- 3.3e-05  |  2.19e-03 +/- 5.0e-05  |
 test_objective_compute_atf              |     +0.33 +/- 1.33     | +2.55e-05 +/- 1.02e-04 |  7.71e-03 +/- 9.5e-05  |  7.69e-03 +/- 3.7e-05  |
 test_objective_jac_dshape_current       |     +1.53 +/- 12.32    | +7.10e-04 +/- 5.72e-03 |  4.71e-02 +/- 2.4e-03  |  4.64e-02 +/- 5.2e-03  |
 test_objective_jac_atf                  |     -0.09 +/- 7.71     | -2.06e-03 +/- 1.80e-01 |  2.33e+00 +/- 1.5e-01  |  2.33e+00 +/- 9.3e-02  |
 test_perturb_1                          |     -1.74 +/- 11.99    | -1.53e-01 +/- 1.05e+00 |  8.63e+00 +/- 7.9e-01  |  8.78e+00 +/- 7.0e-01  |
 test_perturb_2                          |     -1.60 +/- 5.50     | -2.31e-01 +/- 7.93e-01 |  1.42e+01 +/- 5.6e-01  |  1.44e+01 +/- 5.6e-01  |

Copy link
Contributor

github-actions bot commented Nov 9, 2023

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -0.13 +/- 1.07     | -1.98e-05 +/- 1.65e-04 |  1.54e-02 +/- 1.2e-04  |  1.54e-02 +/- 1.1e-04  |
 test_build_transform_fft_midres         |     -0.83 +/- 1.27     | -7.86e-04 +/- 1.20e-03 |  9.39e-02 +/- 9.1e-04  |  9.47e-02 +/- 7.9e-04  |
 test_build_transform_fft_highres        |     +0.10 +/- 0.70     | +4.55e-04 +/- 3.24e-03 |  4.66e-01 +/- 2.1e-03  |  4.66e-01 +/- 2.5e-03  |
 test_equilibrium_init_lowres            |     +0.50 +/- 1.37     | +3.26e-03 +/- 8.88e-03 |  6.53e-01 +/- 5.3e-03  |  6.50e-01 +/- 7.2e-03  |
 test_equilibrium_init_medres            |     -0.57 +/- 1.52     | -5.75e-03 +/- 1.52e-02 |  9.96e-01 +/- 1.2e-02  |  1.00e+00 +/- 9.5e-03  |
 test_equilibrium_init_highres           |     -0.26 +/- 0.75     | -6.71e-03 +/- 1.92e-02 |  2.55e+00 +/- 1.5e-02  |  2.56e+00 +/- 1.2e-02  |
 test_objective_compile_dshape_current   |     +0.15 +/- 8.35     | +6.09e-03 +/- 3.46e-01 |  4.15e+00 +/- 2.2e-01  |  4.14e+00 +/- 2.6e-01  |
 test_objective_compile_atf              |     +0.66 +/- 5.56     | +6.07e-02 +/- 5.10e-01 |  9.24e+00 +/- 4.2e-01  |  9.18e+00 +/- 2.9e-01  |
 test_objective_compute_dshape_current   |     +0.01 +/- 2.04     | +3.11e-07 +/- 4.44e-05 |  2.18e-03 +/- 3.1e-05  |  2.18e-03 +/- 3.2e-05  |
 test_objective_compute_atf              |     -1.51 +/- 0.90     | -1.16e-04 +/- 6.93e-05 |  7.55e-03 +/- 4.1e-05  |  7.67e-03 +/- 5.6e-05  |
 test_objective_jac_dshape_current       |     -5.38 +/- 11.78    | -2.56e-03 +/- 5.60e-03 |  4.50e-02 +/- 4.3e-03  |  4.75e-02 +/- 3.6e-03  |
 test_objective_jac_atf                  |     -3.50 +/- 5.85     | -8.40e-02 +/- 1.40e-01 |  2.32e+00 +/- 6.8e-02  |  2.40e+00 +/- 1.2e-01  |
 test_perturb_1                          |     +0.24 +/- 13.29    | +2.04e-02 +/- 1.14e+00 |  8.58e+00 +/- 8.3e-01  |  8.56e+00 +/- 7.8e-01  |
 test_perturb_2                          |     -0.62 +/- 5.57     | -8.88e-02 +/- 8.03e-01 |  1.43e+01 +/- 7.0e-01  |  1.44e+01 +/- 4.0e-01  |

Copy link
Contributor

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -1.60 +/- 1.92     | -3.65e-04 +/- 4.37e-04 |  2.24e-02 +/- 3.2e-04  |  2.28e-02 +/- 2.9e-04  |
 test_build_transform_fft_midres         |     +1.25 +/- 0.85     | +1.65e-03 +/- 1.13e-03 |  1.34e-01 +/- 8.3e-04  |  1.32e-01 +/- 7.7e-04  |
 test_build_transform_fft_highres        |     +0.55 +/- 0.86     | +3.35e-03 +/- 5.19e-03 |  6.10e-01 +/- 3.7e-03  |  6.06e-01 +/- 3.7e-03  |
 test_equilibrium_init_lowres            |     -1.47 +/- 1.22     | -1.48e-02 +/- 1.22e-02 |  9.89e-01 +/- 6.6e-03  |  1.00e+00 +/- 1.0e-02  |
 test_equilibrium_init_medres            |     -1.76 +/- 1.06     | -2.64e-02 +/- 1.59e-02 |  1.48e+00 +/- 1.3e-02  |  1.50e+00 +/- 8.6e-03  |
 test_equilibrium_init_highres           |     -0.41 +/- 0.74     | -1.47e-02 +/- 2.65e-02 |  3.56e+00 +/- 1.9e-02  |  3.57e+00 +/- 1.8e-02  |
 test_objective_compile_dshape_current   |     -1.90 +/- 6.58     | -1.22e-01 +/- 4.22e-01 |  6.29e+00 +/- 2.8e-01  |  6.42e+00 +/- 3.2e-01  |
 test_objective_compile_atf              |     -0.55 +/- 3.67     | -1.04e-01 +/- 6.89e-01 |  1.87e+01 +/- 4.4e-01  |  1.88e+01 +/- 5.3e-01  |
 test_objective_compute_dshape_current   |     -2.77 +/- 3.84     | -1.13e-04 +/- 1.56e-04 |  3.96e-03 +/- 1.3e-04  |  4.07e-03 +/- 9.2e-05  |
 test_objective_compute_atf              |     -0.89 +/- 2.60     | -1.21e-04 +/- 3.52e-04 |  1.34e-02 +/- 2.7e-04  |  1.35e-02 +/- 2.2e-04  |
 test_objective_jac_dshape_current       |     +2.84 +/- 11.42    | +4.67e-03 +/- 1.87e-02 |  1.69e-01 +/- 9.0e-03  |  1.64e-01 +/- 1.6e-02  |
 test_objective_jac_atf                  |     -3.27 +/- 2.77     | -2.79e-01 +/- 2.37e-01 |  8.26e+00 +/- 1.9e-01  |  8.54e+00 +/- 1.5e-01  |
 test_perturb_1                          |     -0.88 +/- 10.48    | -1.15e-01 +/- 1.36e+00 |  1.29e+01 +/- 8.9e-01  |  1.30e+01 +/- 1.0e+00  |
 test_perturb_2                          |     -1.54 +/- 6.16     | -3.49e-01 +/- 1.40e+00 |  2.23e+01 +/- 1.0e+00  |  2.27e+01 +/- 9.7e-01  |

@dpanici
Copy link
Collaborator Author

dpanici commented Jun 19, 2024

Maybe just throw error for asym for now

desc/compute/_omnigenity.py Show resolved Hide resolved
CHANGELOG.md Outdated Show resolved Hide resolved
],
)
def _Z_mn(params, transforms, profiles, data, **kwargs):
norm = data["Boozer transform matrix"]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure id call this a "norm"

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

renamed to transform_matrix

desc/compute/_omnigenity.py Show resolved Hide resolved
desc/vmec_utils.py Show resolved Hide resolved
lasym.long_name = "0 if the configuration is stellarator-symmetric, 1 if not"
lasym[:] = not eq.sym

# FIXME: ns?
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

still tbd?

desc/compute/_omnigenity.py Outdated Show resolved Hide resolved
description="Boozer harmonics of radial toroidal coordinate of a flux surface",
dim=1,
params=[],
transforms={"B": [[0, 0, 0]]},
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@rahulgaur104 if you want the zeta_B deriv of R you would do

transforms = {"B": [[0,0,1]]}

then have "R_mn" in your data

and in the compute, do dR_dzeta_B = transforms["B"].transforms(data["R_mn"], dz=1)

@dpanici
Copy link
Collaborator Author

dpanici commented Aug 20, 2024

add as error metric G' + iota I' = -p' V'

@dpanici dpanici added waiting for other PRs and removed hackathon Stuff to work on during hackathon labels Aug 29, 2024
Base automatically changed from dp/vmecio-asym to master September 1, 2024 15:23
@dpanici
Copy link
Collaborator Author

dpanici commented Sep 2, 2024

Everything matches well except for some reason the asymmetric (cos(mt - nz)) nu modes are the opposite sign from my boozm output compared to the Fortran/C++ output.

I currently do not understand why this is, as if I had a systematic sign error in how I save the boozmn this should crop up in the other asymmetric parts of the other quantities... but it does not, those all match well, as do the symmetric modes of nu.

sym matches fine

image image

asym are opposite sign

image image

@dpanici
Copy link
Collaborator Author

dpanici commented Sep 23, 2024

  • check Psi sign of asym eq
  • accuracy check checking |B| from fit and original |B| at inboard/outboard pts

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants