diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 7abd4cb6d8..ad00eeef5d 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -11,7 +11,7 @@ on: inputs: debug_enabled: type: boolean - description: 'Run the build with tmate debugging enabled (https://github.com/marketplace/actions/debugging-with-tmate)' + description: '(https://github.com/marketplace/actions/debugging-with-tmate)' required: false default: false @@ -44,7 +44,10 @@ jobs: pwd lscpu cd tests/benchmarks - python -m pytest benchmark_cpu_small.py --benchmark-save='Latest_Commit' -vv --durations=0 --benchmark-save-data + python -m pytest benchmark_cpu_small.py -vv \ + --benchmark-save='Latest_Commit' \ + --durations=0 \ + --benchmark-save-data - name: Checkout current master uses: actions/checkout@v4 with: @@ -57,7 +60,10 @@ jobs: pwd lscpu cd tests/benchmarks - python -m pytest benchmark_cpu_small.py --benchmark-save='master' -vv --durations=0 --benchmark-save-data + python -m pytest benchmark_cpu_small.py -vv \ + --benchmark-save='master' \ + --durations=0 \ + --benchmark-save-data - name: put benchmark results in same folder run: | pwd @@ -88,6 +94,3 @@ jobs: with: name: benchmark_artifact path: tests/benchmarks/.benchmarks - - name: Setup tmate session - if: ${{ failure() }} - uses: mxschmitt/action-tmate@v3 diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index d588563e44..c751a37045 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -12,7 +12,15 @@ jobs: strategy: fail-fast: false matrix: - jax-version: [0.3.0, 0.3.1, 0.3.2, 0.3.3, 0.3.4, 0.3.5, 0.3.6, 0.3.7, 0.3.8, 0.3.9, 0.3.10, 0.3.11, 0.3.12, 0.3.13, 0.3.14, 0.3.15, 0.3.16, 0.3.17, 0.3.19, 0.3.20, 0.3.21, 0.3.22, 0.3.23, 0.3.24, 0.3.25, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.5, 0.4.6, 0.4.7, 0.4.8, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14] + jax-version: [0.3.0, 0.3.1, 0.3.2, 0.3.3, 0.3.4, 0.3.5, + 0.3.6, 0.3.7, 0.3.8, 0.3.9, 0.3.10, 0.3.11, + 0.3.12, 0.3.13, 0.3.14, 0.3.15, 0.3.16, 0.3.17, + 0.3.19, 0.3.20, 0.3.21, 0.3.22, 0.3.23, 0.3.24, + 0.3.25, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.5, + 0.4.6, 0.4.7, 0.4.8, 0.4.9, 0.4.10, 0.4.11, + 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, + 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, + 0.4.24, 0.4.25] group: [1, 2] steps: - uses: actions/checkout@v4 @@ -36,4 +44,10 @@ jobs: run: | pwd lscpu - python -m pytest -m unit --durations=0 --mpl --maxfail=1 --splits 3 --group ${{ matrix.group }} --splitting-algorithm least_duration + python -m pytest -m unit \ + --durations=0 \ + --mpl \ + --maxfail=1 \ + --splits 3 \ + --group ${{ matrix.group }} \ + --splitting-algorithm least_duration diff --git a/.github/workflows/nbtests.yml b/.github/workflows/nbtests.yml index 89c194cc3a..1a8f52e204 100644 --- a/.github/workflows/nbtests.yml +++ b/.github/workflows/nbtests.yml @@ -21,6 +21,7 @@ jobs: strategy: matrix: python-version: ['3.10'] + group: [1, 2] steps: - uses: actions/checkout@v4 @@ -37,4 +38,8 @@ jobs: pwd lscpu export PYTHONPATH=$(pwd) - pytest -v --nbmake "./docs/notebooks" --nbmake-timeout=1000 --ignore=./docs/notebooks/zernike_eval.ipynb --ignore=./docs/notebooks/DESC_Solve_from_pyQSC.ipynb + pytest -v --nbmake "./docs/notebooks" \ + --nbmake-timeout=1000 \ + --ignore=./docs/notebooks/zernike_eval.ipynb \ + --splits 2 \ + --group ${{ matrix.group }} \ diff --git a/.github/workflows/regression_test.yml b/.github/workflows/regression_test.yml index f9961e2ebd..d9ef1072d6 100644 --- a/.github/workflows/regression_test.yml +++ b/.github/workflows/regression_test.yml @@ -42,7 +42,18 @@ jobs: run: | pwd lscpu - python -m pytest -v -m regression --durations=0 --cov-report xml:cov.xml --cov-config=setup.cfg --cov=desc/ --mpl --mpl-results-path=mpl_results.html --mpl-generate-summary=html --splits 6 --group ${{ matrix.group }} --splitting-algorithm least_duration --db ./prof.db + python -m pytest -v -m regression \ + --durations=0 \ + --cov-report xml:cov.xml \ + --cov-config=setup.cfg \ + --cov=desc/ \ + --mpl \ + --mpl-results-path=mpl_results.html \ + --mpl-generate-summary=html \ + --splits 6 \ + --group ${{ matrix.group }} \ + --splitting-algorithm least_duration \ + --db ./prof.db - name: save coverage file and plot comparison results if: always() uses: actions/upload-artifact@v4 diff --git a/.github/workflows/scheduled.yml b/.github/workflows/scheduled.yml index fb4878ae58..a584db5cb8 100644 --- a/.github/workflows/scheduled.yml +++ b/.github/workflows/scheduled.yml @@ -11,7 +11,11 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - combos: [{group: 1, python_version: '3.8'}, {group: 2, python_version: '3.9'}, {group: 3, python_version: '3.10'}, {group: 4, python_version: '3.11'}, {group: 5, python_version: '3.12'}] + combos: [{group: 1, python_version: '3.8'}, + {group: 2, python_version: '3.9'}, + {group: 3, python_version: '3.10'}, + {group: 4, python_version: '3.11'}, + {group: 5, python_version: '3.12'}] steps: - uses: actions/checkout@v4 @@ -31,4 +35,8 @@ jobs: run: | pwd lscpu - python -m pytest -v -m unit --durations=0 --splits 5 --group ${{ matrix.combos.group }} --splitting-algorithm least_duration + python -m pytest -v -m unit \ + --durations=0 \ + --splits 5 \ + --group ${{ matrix.combos.group }} \ + --splitting-algorithm least_duration diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 19f4eda907..57e5881a05 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -20,7 +20,10 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - combos: [{group: 1, python_version: '3.9'}, {group: 2, python_version: '3.10'}, {group: 3, python_version: '3.11'}, {group: 4, python_version: '3.12'}] + combos: [{group: 1, python_version: '3.9'}, + {group: 2, python_version: '3.10'}, + {group: 3, python_version: '3.11'}, + {group: 4, python_version: '3.12'}] steps: - uses: actions/checkout@v4 @@ -41,7 +44,18 @@ jobs: run: | pwd lscpu - python -m pytest -v -m unit --durations=0 --cov-report xml:cov.xml --cov-config=setup.cfg --cov=desc/ --mpl --mpl-results-path=mpl_results.html --mpl-generate-summary=html --splits 4 --group ${{ matrix.combos.group }} --splitting-algorithm least_duration --db ./prof.db + python -m pytest -v -m unit \ + --durations=0 \ + --cov-report xml:cov.xml \ + --cov-config=setup.cfg \ + --cov=desc/ \ + --mpl \ + --mpl-results-path=mpl_results.html \ + --mpl-generate-summary=html \ + --splits 4 \ + --group ${{ matrix.combos.group }} \ + --splitting-algorithm least_duration \ + --db ./prof.db - name: save coverage file and plot comparison results if: always() uses: actions/upload-artifact@v4 diff --git a/.test_durations b/.test_durations index d8f7060fcb..7331618ec5 100644 --- a/.test_durations +++ b/.test_durations @@ -1,475 +1,642 @@ { - "tests/test_axis_limits.py::TestAxisLimits::test_axis_limit_api": 4.548233151435852, - "tests/test_axis_limits.py::TestAxisLimits::test_limit_continuity": 71.97202825546265, - "tests/test_axis_limits.py::TestAxisLimits::test_magnetic_field_is_physical": 25.841275453567505, - "tests/test_backend.py::test_put": 0.3379814624786377, - "tests/test_backend.py::test_sign": 0.3835097551345825, - "tests/test_backend.py::test_vmap": 0.3952150344848633, - "tests/test_basis.py::TestBasis::test_polyder": 0.21661221981048584, - "tests/test_basis.py::TestBasis::test_polyval": 0.2100837230682373, - "tests/test_basis.py::TestBasis::test_zernike_coeffs": 0.2981588840484619, - "tests/test_basis.py::TestBasis::test_polyval_exact": 0.2100837230682373, - "tests/test_basis.py::TestBasis::test_powers": 0.2231835126876831, - "tests/test_basis.py::TestBasis::test_chebyshev": 0.3672276735305786, - "tests/test_basis.py::TestBasis::test_zernike_radial": 1.423325777053833, - "tests/test_basis.py::TestBasis::test_fourier": 0.39130139350891113, - "tests/test_basis.py::TestBasis::test_power_series": 0.23145699501037598, - "tests/test_basis.py::TestBasis::test_double_fourier": 0.4569425582885742, - "tests/test_basis.py::TestBasis::test_change_resolution": 1.3561828136444092, - "tests/test_basis.py::TestBasis::test_repr": 2.206360936164856, - "tests/test_basis.py::TestBasis::test_zernike_indexing": 0.20653140544891357, - "tests/test_basis.py::TestBasis::test_derivative_not_in_basis_zeros": 0.6462588310241699, - "tests/test_basis.py::TestBasis::test_basis_resolutions_assert_integers": 0.24122905731201172, + "tests/test_axis_limits.py::TestAxisLimits::test_axis_limit_api": 10.306883258999733, + "tests/test_axis_limits.py::TestAxisLimits::test_limit_continuity": 27.196473163996416, + "tests/test_axis_limits.py::TestAxisLimits::test_magnetic_field_is_physical": 6.181816910004272, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[<1/|B|>]": 15.636104067001725, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[<|B|>]": 13.036220199002855, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[<|B|^2>]": 13.710468647001107, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[<|B|^2>_r]": 14.318662270001369, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0]": 10.275923050998244, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_r]": 10.9314324520019, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_rr]": 11.45641257600073, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_rt]": 11.51512692599863, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_rz]": 11.840702208995936, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_t]": 10.802180848000717, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_tt]": 11.650608847001422, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_tz]": 11.722970514001645, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_z]": 10.839945408999483, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[B0_zz]": 11.56075947899808, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[J^rho]": 14.653282455998124, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e^rho_r]": 11.195839553998667, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e^rho_t]": 11.16860465800346, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e^rho_z]": 11.113929876999464, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e^zeta_r]": 11.045268964004208, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e^zeta_t]": 11.145609544000763, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e^zeta_z]": 11.20895503599968, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[e_theta/sqrt(g)]": 10.351756926000235, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[grad(B)]": 15.32959635700172, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[grad(|B|)]": 15.452872551999462, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota current]": 12.528394308003044, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota vacuum]": 11.84887967600298, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota]": 12.618590641999617, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_den]": 10.120928225998796, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_den_r]": 10.911997907001933, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_den_rr]": 11.443894683998224, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_num current]": 10.042572182002914, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_num vacuum]": 10.309513904994674, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_num_r current]": 10.377449766998325, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_num_r vacuum]": 10.650603052003135, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_num_rr]": 11.84881359800056, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[iota_r]": 12.44367556600264, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[n_rho]": 10.437926040001912, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[n_zeta]": 10.570785170999443, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[trapped fraction]": 13.358987843999785, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[|e_rho x e_theta|_r]": 9.826671683000313, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[|e_rho x e_theta|_rr]": 10.230985248996149, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[|e_theta x e_zeta|_r]": 10.383659332001116, + "tests/test_axis_limits.py::test_reverse_mode_ad_axis[|e_theta x e_zeta|_rr]": 10.250626319997536, + "tests/test_backend.py::test_put": 0.5784192619976238, + "tests/test_backend.py::test_sign": 0.4722795660018164, + "tests/test_backend.py::test_vmap": 0.3483646079985192, + "tests/test_basis.py::TestBasis::test_basis_resolutions_assert_integers": 0.27035030200568144, + "tests/test_basis.py::TestBasis::test_change_resolution": 0.2740671179999481, + "tests/test_basis.py::TestBasis::test_chebyshev": 0.33105514000271796, + "tests/test_basis.py::TestBasis::test_derivative_not_in_basis_zeros": 0.5189595209994877, + "tests/test_basis.py::TestBasis::test_double_fourier": 0.36261462599941297, + "tests/test_basis.py::TestBasis::test_fourier": 0.3442120290019375, + "tests/test_basis.py::TestBasis::test_jacobi_poly_single": 0.43588841199380113, + "tests/test_basis.py::TestBasis::test_polyder": 0.3267408269966836, + "tests/test_basis.py::TestBasis::test_polyval": 0.33155919999990147, + "tests/test_basis.py::TestBasis::test_polyval_exact": 5.259696545999759, + "tests/test_basis.py::TestBasis::test_power_series": 0.34384299799785367, + "tests/test_basis.py::TestBasis::test_powers": 0.4021272939971823, + "tests/test_basis.py::TestBasis::test_repr": 0.279552201005572, + "tests/test_basis.py::TestBasis::test_zernike_coeffs": 0.29283089700038545, + "tests/test_basis.py::TestBasis::test_zernike_indexing": 0.2757554160016298, + "tests/test_basis.py::TestBasis::test_zernike_radial": 2.7187563000006776, "tests/test_basis.py::test_jacobi_jvp": 2.096624255180359, - "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_analytic": 3.5495134592056274, - "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_Kim": 13.277099013328552, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_second_pass": 1.2246217727661133, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_2_3": 2.7978817224502563, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_4_5": 10.157285571098328, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_tokamak_benchmark": 8.322012901306152, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QA": 9.539125442504883, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QH": 9.056997179985046, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_normalization": 41.65271031856537, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_resolution": 95.99948453903198, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_iota": 421.759339094162, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_2_3": 2.421199860997149, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_4_5": 4.598695638000208, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_second_pass": 1.9836588429971016, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QA": 8.149377257002925, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QH": 1.6454791360010859, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_tokamak_benchmark": 6.201658457001031, + "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_Kim": 8.299000467999576, + "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_analytic": 2.636759080000047, + "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_normalization": 34.827340141000604, "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_current": 387.01775670051575, - "tests/test_bootstrap.py::test_bootstrap_objective_build": 27.412319898605347, - "tests/test_coils.py::TestCoil::test_biot_savart_all_coils": 4.966217041015625, - "tests/test_coils.py::TestCoil::test_properties": 5.437066316604614, - "tests/test_coils.py::TestCoil::test_SumMagneticField_with_Coil": 3.111448287963867, - "tests/test_coils.py::TestCoilSet::test_linspaced_linear": 2.3614059686660767, - "tests/test_coils.py::TestCoilSet::test_linspaced_angular": 6.477421283721924, - "tests/test_coils.py::TestCoilSet::test_from_symmetry": 10.293273448944092, - "tests/test_coils.py::TestCoilSet::test_properties": 5.437066316604614, - "tests/test_coils.py::TestCoilSet::test_dunder_methods": 1.1951448917388916, - "tests/test_coils.py::test_load_and_save_makegrid_coils": 13.879701614379883, - "tests/test_coils.py::test_save_and_load_makegrid_coils_rotated": 37.21220946311951, - "tests/test_coils.py::test_save_and_load_makegrid_coils_rotated_int_grid": 13.974516749382019, - "tests/test_coils.py::test_save_makegrid_coils_assert_NFP": 0.4364185333251953, - "tests/test_coils.py::test_load_makegrid_coils_header_asserts": 0.41977834701538086, - "tests/test_coils.py::test_repr": 2.206360936164856, - "tests/test_compute_funs.py::test_total_volume": 5.2514989376068115, - "tests/test_compute_funs.py::test_enclosed_volumes": 5.10689377784729, - "tests/test_compute_funs.py::test_enclosed_areas": 9.250897884368896, - "tests/test_compute_funs.py::test_surface_areas": 5.958480000495911, - "tests/test_compute_funs.py::test_surface_areas_2": 5.958480000495911, - "tests/test_compute_funs.py::test_elongation": 8.7400062084198, - "tests/test_compute_funs.py::test_magnetic_field_derivatives": 33.20082712173462, - "tests/test_compute_funs.py::test_metric_derivatives": 5.997546434402466, - "tests/test_compute_funs.py::test_magnetic_pressure_gradient": 12.57770299911499, - "tests/test_compute_funs.py::test_currents": 10.046160459518433, - "tests/test_compute_funs.py::test_BdotgradB": 9.498456478118896, - "tests/test_compute_funs.py::test_boozer_transform": 8.135252833366394, - "tests/test_compute_funs.py::test_compute_grad_p_volume_avg": 3.611231803894043, - "tests/test_compute_funs.py::test_compare_quantities_to_vmec": 23.354363560676575, - "tests/test_compute_funs.py::test_compute_everything": 416.05528950691223, - "tests/test_compute_funs.py::test_compute_averages": 14.273451089859009, - "tests/test_compute_funs.py::test_covariant_basis_vectors": 35.97991585731506, - "tests/test_data_index.py::TestDataIndex::test_data_index_deps": 12.303489208221436, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals": 44.133379101753235, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals_transform": 1.8601322174072266, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_vector_functions": 4.124471187591553, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_area": 5.989154577255249, - "tests/test_compute_utils.py::TestComputeUtils::test_line_length": 13.470165729522705, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_identity_op": 11.656662225723267, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_homomorphism": 13.759432792663574, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals_against_shortcut": 1.070227861404419, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_against_shortcut": 9.63156795501709, - "tests/test_compute_utils.py::TestComputeUtils::test_symmetry_surface_average_1": 12.463624119758606, - "tests/test_compute_utils.py::TestComputeUtils::test_symmetry_surface_average_2": 6.5595245361328125, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_variance": 1.9017333984375, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_min_max": 2.2540769577026367, - "tests/test_configuration.py::TestConstructor::test_defaults": 2.065831422805786, - "tests/test_configuration.py::TestConstructor::test_supplied_objects": 6.169172406196594, - "tests/test_configuration.py::TestConstructor::test_dict": 3.6929678916931152, - "tests/test_configuration.py::TestConstructor::test_asserts": 1.7440376281738281, - "tests/test_configuration.py::TestConstructor::test_supplied_coeffs": 1.8669953346252441, - "tests/test_configuration.py::TestInitialGuess::test_default_set": 4.367171764373779, - "tests/test_configuration.py::TestInitialGuess::test_errors": 2.6561394929885864, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_other": 2.640138626098633, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface": 3.513801336288452, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface2": 3.513801336288452, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_points": 5.633792042732239, - "tests/test_configuration.py::TestInitialGuess::test_NFP_error": 1.977283000946045, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_file": 4.338591694831848, - "tests/test_configuration.py::TestGetSurfaces::test_get_rho_surface": 2.8503174781799316, - "tests/test_configuration.py::TestGetSurfaces::test_get_zeta_surface": 2.5673550367355347, - "tests/test_configuration.py::TestGetSurfaces::test_get_theta_surface": 1.638696551322937, - "tests/test_configuration.py::TestGetSurfaces::test_asserts": 1.7440376281738281, - "tests/test_configuration.py::test_magnetic_axis": 5.0468586683273315, - "tests/test_configuration.py::test_is_nested": 4.586352825164795, - "tests/test_configuration.py::test_is_nested_theta": 4.808726787567139, - "tests/test_configuration.py::test_get_profile": 8.75450313091278, - "tests/test_configuration.py::test_kinetic_errors": 4.169016122817993, - "tests/test_constrain_current.py::TestConstrainCurrent::test_iota_to_current_and_back": 33.03442454338074, - "tests/test_constrain_current.py::TestConstrainCurrent::test_current_to_iota_and_back": 24.11527442932129, - "tests/test_curves.py::TestRZCurve::test_length": 2.5587551593780518, - "tests/test_curves.py::TestRZCurve::test_curvature": 2.93274188041687, - "tests/test_curves.py::TestRZCurve::test_torsion": 2.4702093601226807, - "tests/test_curves.py::TestRZCurve::test_frenet": 2.836435317993164, - "tests/test_curves.py::TestRZCurve::test_coords": 1.7616546154022217, - "tests/test_curves.py::TestRZCurve::test_misc": 2.985080599784851, - "tests/test_curves.py::TestRZCurve::test_asserts": 1.7440376281738281, - "tests/test_curves.py::TestRZCurve::test_to_FourierXYZCurve": 6.9260677099227905, - "tests/test_curves.py::TestRZCurve::test_to_SplineXYZCurve": 8.358957767486572, - "tests/test_curves.py::TestRZCurve::test_from_input_file": 7.244992971420288, - "tests/test_curves.py::TestFourierXYZCurve::test_length": 2.5587551593780518, - "tests/test_curves.py::TestFourierXYZCurve::test_curvature": 2.93274188041687, - "tests/test_curves.py::TestFourierXYZCurve::test_torsion": 2.4702093601226807, - "tests/test_curves.py::TestFourierXYZCurve::test_frenet": 2.836435317993164, - "tests/test_curves.py::TestFourierXYZCurve::test_coords": 1.7616546154022217, - "tests/test_curves.py::TestFourierXYZCurve::test_to_FourierXYZCurve": 6.9260677099227905, - "tests/test_curves.py::TestFourierXYZCurve::test_misc": 2.985080599784851, - "tests/test_curves.py::TestPlanarCurve::test_length": 2.5587551593780518, - "tests/test_curves.py::TestPlanarCurve::test_curvature": 2.93274188041687, - "tests/test_curves.py::TestPlanarCurve::test_torsion": 2.4702093601226807, - "tests/test_curves.py::TestPlanarCurve::test_frenet": 2.836435317993164, - "tests/test_curves.py::TestPlanarCurve::test_coords": 1.7616546154022217, - "tests/test_curves.py::TestPlanarCurve::test_misc": 2.985080599784851, - "tests/test_curves.py::TestPlanarCurve::test_asserts": 1.7440376281738281, - "tests/test_curves.py::TestSplineXYZCurve::test_length": 2.5587551593780518, - "tests/test_curves.py::TestSplineXYZCurve::test_coords": 1.7616546154022217, - "tests/test_curves.py::TestSplineXYZCurve::test_curvature": 2.93274188041687, - "tests/test_curves.py::TestSplineXYZCurve::test_torsion": 2.4702093601226807, - "tests/test_curves.py::TestSplineXYZCurve::test_to_SplineXYZCurve": 8.358957767486572, - "tests/test_curves.py::TestSplineXYZCurve::test_asserts_and_errors": 1.7440376281738281, - "tests/test_curves.py::TestSplineXYZCurve::test_misc": 2.985080599784851, - "tests/test_curves.py::TestSplineXYZCurve::test_compute_ndarray_error": 2.4868667125701904, - "tests/test_derivatives.py::TestDerivative::test_finite_diff_vec": 1.5384825468063354, - "tests/test_derivatives.py::TestDerivative::test_finite_diff_scalar": 1.6321053504943848, - "tests/test_derivatives.py::TestDerivative::test_auto_diff": 1.6796185970306396, - "tests/test_derivatives.py::TestDerivative::test_compare_AD_FD": 1.5109567642211914, - "tests/test_derivatives.py::TestDerivative::test_fd_hessian": 1.321077823638916, - "tests/test_derivatives.py::TestDerivative::test_block_jacobian": 1.6304693222045898, - "tests/test_derivatives.py::TestDerivative::test_jac_looped": 1.59993577003479, - "tests/test_derivatives.py::TestJVP::test_autodiff_jvp": 1.6057462692260742, - "tests/test_derivatives.py::TestJVP::test_finitediff_jvp": 1.5164278745651245, - "tests/test_derivatives.py::TestJVP::test_autodiff_jvp2": 1.6057462692260742, - "tests/test_derivatives.py::TestJVP::test_finitediff_jvp2": 1.5164278745651245, - "tests/test_derivatives.py::TestJVP::test_autodiff_jvp3": 1.9495824575424194, - "tests/test_derivatives.py::TestJVP::test_finitediff_jvp3": 1.406630516052246, - "tests/test_derivatives.py::TestJVP::test_vjp": 49.77897334098816, - "tests/test_equilibrium.py::test_compute_geometry": 2.961799144744873, - "tests/test_equilibrium.py::test_compute_theta_coords": 9.089770913124084, - "tests/test_equilibrium.py::test_compute_flux_coords": 78.98275232315063, - "tests/test_equilibrium.py::test_map_coordinates": 78.44930052757263, - "tests/test_equilibrium.py::test_map_coordinates2": 78.44930052757263, - "tests/test_equilibrium.py::test_to_sfl": 39.330466985702515, - "tests/test_equilibrium.py::test_continuation_resolution": 268.36252522468567, - "tests/test_equilibrium.py::test_grid_resolution_warning": 534.3893562555313, - "tests/test_equilibrium.py::test_eq_change_grid_resolution": 2.5595715045928955, - "tests/test_equilibrium.py::test_eq_change_symmetry": 5.290525197982788, - "tests/test_equilibrium.py::test_resolution": 8.523245215415955, - "tests/test_equilibrium.py::test_equilibrium_from_near_axis": 30.988573789596558, - "tests/test_equilibrium.py::test_poincare_solve_not_implemented": 2.5252435207366943, - "tests/test_equilibrium.py::test_equilibriafamily_constructor": 6.004562854766846, - "tests/test_equilibrium.py::test_change_NFP": 13.311393737792969, - "tests/test_equilibrium.py::test_error_when_ndarray_or_integer_passed": 2.2681093215942383, - "tests/test_examples.py::test_SOLOVEV_vacuum": 3.47597336769104, + "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_iota": 421.759339094162, + "tests/test_bootstrap.py::test_bootstrap_objective_build": 33.5478269860032, + "tests/test_bootstrap.py::test_bootstrap_optimization_comparison_qa": 368.391193716001, + "tests/test_coils.py::TestCoil::test_SumMagneticField_with_Coil": 1.2285185000037018, + "tests/test_coils.py::TestCoil::test_adding_MagneticField_with_Coil_or_CoilSet": 0.9661065750005946, + "tests/test_coils.py::TestCoil::test_biot_savart_all_coils": 3.9202290960020036, + "tests/test_coils.py::TestCoil::test_converting_coil_types": 1.5931856890019844, + "tests/test_coils.py::TestCoil::test_properties": 0.3970797659967502, + "tests/test_coils.py::TestCoilSet::test_coilset_convert": 4.40795374499794, + "tests/test_coils.py::TestCoilSet::test_dunder_methods": 1.0025817760033533, + "tests/test_coils.py::TestCoilSet::test_from_symmetry": 4.516506672000105, + "tests/test_coils.py::TestCoilSet::test_linspaced_angular": 4.781867071003944, + "tests/test_coils.py::TestCoilSet::test_linspaced_linear": 2.043203753994021, + "tests/test_coils.py::TestCoilSet::test_properties": 2.5864213140012, + "tests/test_coils.py::TestCoilSet::test_symmetry_magnetic_field": 15.481033231997571, + "tests/test_coils.py::test_load_and_save_makegrid_coils": 4.837137193000672, + "tests/test_coils.py::test_load_makegrid_coils_header_asserts": 0.4405547799942724, + "tests/test_coils.py::test_repr": 1.5094817869976396, + "tests/test_coils.py::test_save_and_load_makegrid_coils_rotated": 9.566098835999583, + "tests/test_coils.py::test_save_and_load_makegrid_coils_rotated_int_grid": 8.960116868998739, + "tests/test_coils.py::test_save_makegrid_coils_assert_NFP": 0.8933817200013436, + "tests/test_compat.py::test_flip_helicity_axisym": 15.645121720001043, + "tests/test_compat.py::test_flip_helicity_current": 27.13188577099936, + "tests/test_compat.py::test_flip_helicity_iota": 25.78969065599813, + "tests/test_compat.py::test_rescale": 26.37215985399962, + "tests/test_compute_funs.py::test_BdotgradB": 13.458342846999585, + "tests/test_compute_funs.py::test_aliases": 15.90729377200114, + "tests/test_compute_funs.py::test_boozer_transform": 9.761044317005144, + "tests/test_compute_funs.py::test_compare_quantities_to_vmec": 23.47456439400048, + "tests/test_compute_funs.py::test_compute_averages": 23.204242305000662, + "tests/test_compute_funs.py::test_compute_everything": 41.52274995600237, + "tests/test_compute_funs.py::test_contravariant_basis_vectors": 92.40470739599914, + "tests/test_compute_funs.py::test_covariant_basis_vectors": 19.219057816000714, + "tests/test_compute_funs.py::test_elongation": 26.32173658100146, + "tests/test_compute_funs.py::test_enclosed_areas": 18.132758002000628, + "tests/test_compute_funs.py::test_enclosed_volumes": 18.970259471003374, + "tests/test_compute_funs.py::test_iota_components": 17.655027994002012, + "tests/test_compute_funs.py::test_magnetic_field_derivatives": 18.572860783006035, + "tests/test_compute_funs.py::test_magnetic_pressure_gradient": 13.666031557000679, + "tests/test_compute_funs.py::test_metric_derivatives": 9.299652367000817, + "tests/test_compute_funs.py::test_surface_areas": 14.468549634002557, + "tests/test_compute_funs.py::test_surface_areas_2": 13.126867309001682, + "tests/test_compute_funs.py::test_surface_equilibrium_geometry": 28.360896346999652, + "tests/test_compute_funs.py::test_total_volume": 18.43933450300392, + "tests/test_compute_utils.py::TestComputeUtils::test_line_length": 4.704729233999387, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_area": 2.8100511410011677, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_against_shortcut": 3.83928267799638, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_homomorphism": 1.4975819430037518, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_identity_op": 1.1260491000030015, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_vector_functions": 1.8768331540049985, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals": 19.622801577002974, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals_against_shortcut": 1.0201668540030369, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals_transform": 1.2729154780026875, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_min_max": 1.3701877329986019, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_variance": 1.3257562700018752, + "tests/test_compute_utils.py::TestComputeUtils::test_symmetry_surface_average_1": 5.496868898997491, + "tests/test_compute_utils.py::TestComputeUtils::test_symmetry_surface_average_2": 10.967256519998045, + "tests/test_compute_utils.py::test_rotation_matrix": 4.061818482001399, + "tests/test_configuration.py::TestConstructor::test_asserts": 1.4049688549966959, + "tests/test_configuration.py::TestConstructor::test_defaults": 10.572905429999082, + "tests/test_configuration.py::TestConstructor::test_dict": 8.3489596670006, + "tests/test_configuration.py::TestConstructor::test_supplied_coeffs": 0.9374606740020681, + "tests/test_configuration.py::TestConstructor::test_supplied_objects": 22.871558008999273, + "tests/test_configuration.py::TestGetSurfaces::test_asserts": 1.1521193240005232, + "tests/test_configuration.py::TestGetSurfaces::test_get_rho_surface": 15.097044327005278, + "tests/test_configuration.py::TestGetSurfaces::test_get_theta_surface": 1.1302085750030528, + "tests/test_configuration.py::TestGetSurfaces::test_get_zeta_surface": 4.384280577996833, + "tests/test_configuration.py::TestInitialGuess::test_NFP_error": 0.9750648899971566, + "tests/test_configuration.py::TestInitialGuess::test_default_set": 11.891957347001153, + "tests/test_configuration.py::TestInitialGuess::test_errors": 1.3177527570078382, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_coordinate_mapping": 20.587510880999616, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_coordinate_mapping_no_sym": 18.74857710800279, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_file": 11.965596940997784, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_other": 12.467711595996661, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_points": 13.174946844002989, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface": 5.232776146003744, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface2": 5.747947469997598, + "tests/test_configuration.py::test_get_profile": 18.315085360998637, + "tests/test_configuration.py::test_is_nested": 17.02455229900079, + "tests/test_configuration.py::test_is_nested_theta": 16.069276984999306, + "tests/test_configuration.py::test_kinetic_errors": 14.796659756997542, + "tests/test_configuration.py::test_magnetic_axis": 7.3246193419945484, + "tests/test_constrain_current.py::TestConstrainCurrent::test_current_to_iota_and_back": 13.809269275003317, + "tests/test_constrain_current.py::TestConstrainCurrent::test_iota_to_current_and_back": 20.610450854997907, + "tests/test_curves.py::TestFourierXYZCurve::test_asserts": 0.9961255910020554, + "tests/test_curves.py::TestFourierXYZCurve::test_coords": 1.144117637999443, + "tests/test_curves.py::TestFourierXYZCurve::test_curvature": 1.1052798269993218, + "tests/test_curves.py::TestFourierXYZCurve::test_frenet": 1.4736654719999933, + "tests/test_curves.py::TestFourierXYZCurve::test_length": 3.3653678620030405, + "tests/test_curves.py::TestFourierXYZCurve::test_misc": 1.2293614089976472, + "tests/test_curves.py::TestFourierXYZCurve::test_to_FourierXYZCurve": 3.8908353169972543, + "tests/test_curves.py::TestFourierXYZCurve::test_torsion": 1.166831981998257, + "tests/test_curves.py::TestPlanarCurve::test_asserts": 0.9783399230000214, + "tests/test_curves.py::TestPlanarCurve::test_coords": 1.2120063109978219, + "tests/test_curves.py::TestPlanarCurve::test_curvature": 1.1603255669979262, + "tests/test_curves.py::TestPlanarCurve::test_frenet": 1.406318638997618, + "tests/test_curves.py::TestPlanarCurve::test_length": 1.9474872020000475, + "tests/test_curves.py::TestPlanarCurve::test_misc": 1.2922219290012436, + "tests/test_curves.py::TestPlanarCurve::test_rotation": 2.723928325005545, + "tests/test_curves.py::TestPlanarCurve::test_torsion": 1.159477840999898, + "tests/test_curves.py::TestRZCurve::test_asserts": 0.9437689430014871, + "tests/test_curves.py::TestRZCurve::test_coords": 1.2269516010019288, + "tests/test_curves.py::TestRZCurve::test_curvature": 1.1763031689988566, + "tests/test_curves.py::TestRZCurve::test_frenet": 1.42519166999773, + "tests/test_curves.py::TestRZCurve::test_from_input_file": 32.137807690000045, + "tests/test_curves.py::TestRZCurve::test_length": 3.607348298002762, + "tests/test_curves.py::TestRZCurve::test_misc": 1.6830140780002694, + "tests/test_curves.py::TestRZCurve::test_to_FourierXYZCurve": 3.318559372997697, + "tests/test_curves.py::TestRZCurve::test_to_SplineXYZCurve": 3.0842853100002685, + "tests/test_curves.py::TestRZCurve::test_torsion": 1.1583341479999945, + "tests/test_curves.py::TestSplineXYZCurve::test_asserts_and_errors": 1.0821766250010114, + "tests/test_curves.py::TestSplineXYZCurve::test_compute_ndarray_error": 1.034954820999701, + "tests/test_curves.py::TestSplineXYZCurve::test_coords": 2.233524409999518, + "tests/test_curves.py::TestSplineXYZCurve::test_curvature": 2.0053053730080137, + "tests/test_curves.py::TestSplineXYZCurve::test_length": 21.65546567400088, + "tests/test_curves.py::TestSplineXYZCurve::test_misc": 1.0600306280030054, + "tests/test_curves.py::TestSplineXYZCurve::test_to_SplineXYZCurve": 3.3972826739991433, + "tests/test_curves.py::TestSplineXYZCurve::test_torsion": 1.9396493819986063, + "tests/test_data_index.py::TestDataIndex::test_data_index_deps": 3.654910685996583, + "tests/test_derivatives.py::TestDerivative::test_auto_diff": 1.3889989999988757, + "tests/test_derivatives.py::TestDerivative::test_block_jacobian": 1.2823353329949896, + "tests/test_derivatives.py::TestDerivative::test_compare_AD_FD": 1.009793486999115, + "tests/test_derivatives.py::TestDerivative::test_fd_hessian": 0.9568148079997627, + "tests/test_derivatives.py::TestDerivative::test_finite_diff_scalar": 0.95879512099782, + "tests/test_derivatives.py::TestDerivative::test_finite_diff_vec": 0.9704352070002642, + "tests/test_derivatives.py::TestDerivative::test_jac_looped": 1.1326828239980387, + "tests/test_derivatives.py::TestJVP::test_autodiff_jvp": 1.4069660029999795, + "tests/test_derivatives.py::TestJVP::test_autodiff_jvp2": 1.2605693780024012, + "tests/test_derivatives.py::TestJVP::test_autodiff_jvp3": 1.2917560200003209, + "tests/test_derivatives.py::TestJVP::test_finitediff_jvp": 1.1071212779970665, + "tests/test_derivatives.py::TestJVP::test_finitediff_jvp2": 1.0138074659989798, + "tests/test_derivatives.py::TestJVP::test_finitediff_jvp3": 0.9896067409972602, + "tests/test_derivatives.py::TestJVP::test_vjp": 1.2450714299993706, + "tests/test_equilibrium.py::test_backward_compatible_load_and_resolve": 36.62463790200127, + "tests/test_equilibrium.py::test_change_NFP": 15.607664676004788, + "tests/test_equilibrium.py::test_compute_theta_coords": 9.970856725998601, + "tests/test_equilibrium.py::test_continuation_resolution": 136.30576520200702, + "tests/test_equilibrium.py::test_eq_change_symmetry": 15.52943572300137, + "tests/test_equilibrium.py::test_equilibriafamily_constructor": 37.23395097699904, + "tests/test_equilibrium.py::test_equilibrium_from_near_axis": 28.39438554100343, + "tests/test_equilibrium.py::test_equilibrium_unused_kwargs": 12.015114074001758, + "tests/test_equilibrium.py::test_error_when_ndarray_or_integer_passed": 3.127018835002673, + "tests/test_equilibrium.py::test_grid_resolution_warning": 15.262868713001808, + "tests/test_equilibrium.py::test_map_coordinates": 22.34353607900266, + "tests/test_equilibrium.py::test_map_coordinates_derivative": 47.83286731299813, + "tests/test_equilibrium.py::test_poincare_solve_not_implemented": 10.693831380001939, + "tests/test_equilibrium.py::test_resolution": 22.519933166000556, + "tests/test_equilibrium.py::test_to_sfl": 44.09785354499763, + "tests/test_examples.py::TestGetExample::test_example_get_boundary": 8.60165646599853, + "tests/test_examples.py::TestGetExample::test_example_get_current": 2.388091346998408, + "tests/test_examples.py::TestGetExample::test_example_get_eq": 2.860374152995064, + "tests/test_examples.py::TestGetExample::test_example_get_eqf": 2.6269075269992754, + "tests/test_examples.py::TestGetExample::test_example_get_iota": 2.4244833209995704, + "tests/test_examples.py::TestGetExample::test_example_get_pressure": 2.87473690400293, + "tests/test_examples.py::TestGetExample::test_missing_example": 1.9945495780011697, + "tests/test_examples.py::test_1d_optimization": 101.42265871397103, + "tests/test_examples.py::test_1d_optimization_old": 126.5154206950101, + "tests/test_examples.py::test_ATF_results": 722.2917343610025, + "tests/test_examples.py::test_DSHAPE_current_results": 126.09806620600284, + "tests/test_examples.py::test_DSHAPE_results": 148.24703417899946, + "tests/test_examples.py::test_ESTELL_results": 1321.8716867980256, + "tests/test_examples.py::test_HELIOTRON_results": 518.3004652379896, + "tests/test_examples.py::test_HELIOTRON_vac_results": 315.85949871699995, + "tests/test_examples.py::test_NAE_QIC_solve": 237.0647383928299, + "tests/test_examples.py::test_NAE_QSC_solve": 238.0105233192444, "tests/test_examples.py::test_SOLOVEV_results": 9.616402626037598, - "tests/test_examples.py::test_DSHAPE_results": 13.075568437576294, - "tests/test_examples.py::test_DSHAPE_current_results": 13.897011280059814, - "tests/test_examples.py::test_HELIOTRON_results": 14.499746799468994, - "tests/test_examples.py::test_HELIOTRON_vac_results": 18.51157259941101, - "tests/test_examples.py::test_HELIOTRON_vac2_results": 65.61901330947876, - "tests/test_examples.py::test_force_balance_grids": 278.0637106895447, - "tests/test_examples.py::test_solve_bounds": 44.96602940559387, - "tests/test_examples.py::test_1d_optimization": 76.29036378860474, - "tests/test_examples.py::test_1d_optimization_old": 76.29036378860474, - "tests/test_examples.py::test_qh_optimization": 900.0, - "tests/test_examples.py::test_ATF_results": 1421.3091161251068, - "tests/test_examples.py::test_ESTELL_results": 5973.843339681625, - "tests/test_examples.py::test_simsopt_QH_comparison": 1908.741860628128, - "tests/test_examples.py::test_NAE_QSC_solve": 294.0105233192444, - "tests/test_examples.py::test_NAE_QIC_solve": 498.0647383928299, - "tests/test_examples.py::TestGetExample::test_missing_example": 1.8569834232330322, - "tests/test_examples.py::TestGetExample::test_example_get_eq": 2.801750421524048, - "tests/test_examples.py::TestGetExample::test_example_get_eqf": 2.801750421524048, - "tests/test_examples.py::TestGetExample::test_example_get_boundary": 6.7636377811431885, - "tests/test_examples.py::TestGetExample::test_example_get_pressure": 4.8439977169036865, - "tests/test_examples.py::TestGetExample::test_example_get_iota": 6.1666529178619385, - "tests/test_examples.py::TestGetExample::test_example_get_current": 5.687736749649048, - "tests/test_geometry.py::test_rotation_matrix": 1.7277274131774902, - "tests/test_geometry.py::test_xyz2rpz": 1.6413753032684326, - "tests/test_geometry.py::test_rpz2xyz": 1.5908217430114746, - "tests/test_grid.py::TestGrid::test_custom_grid": 1.6421220302581787, - "tests/test_grid.py::TestGrid::test_linear_grid": 1.502720594406128, - "tests/test_grid.py::TestGrid::test_linear_grid_spacing_consistency": 1.2147817611694336, - "tests/test_grid.py::TestGrid::test_linear_grid_symmetric_nodes_consistency": 1.502720594406128, - "tests/test_grid.py::TestGrid::test_linear_grid_spacing_two_nodes": 1.2147817611694336, - "tests/test_grid.py::TestGrid::test_spacing_when_duplicate_node_is_removed": 1.6561402082443237, - "tests/test_grid.py::TestGrid::test_node_spacing_non_sym": 2.0402016639709473, - "tests/test_grid.py::TestGrid::test_symmetry_spacing_basic": 1.5400123596191406, - "tests/test_grid.py::TestGrid::test_node_spacing_sym": 2.327052116394043, - "tests/test_grid.py::TestGrid::test_concentric_grid": 1.492321491241455, - "tests/test_grid.py::TestGrid::test_concentric_grid_high_res": 1.5603344440460205, - "tests/test_grid.py::TestGrid::test_quadrature_grid": 1.500448226928711, - "tests/test_grid.py::TestGrid::test_quad_grid_volume_integration": 2.609942317008972, - "tests/test_grid.py::TestGrid::test_repr": 2.206360936164856, - "tests/test_grid.py::TestGrid::test_change_resolution": 1.3561828136444092, - "tests/test_grid.py::TestGrid::test_enforce_symmetry": 3.3707268238067627, - "tests/test_grid.py::TestGrid::test_symmetry_volume_integral": 1.6521550416946411, - "tests/test_grid.py::TestGrid::test_compress_expand_inverse_op": 0.8788290023803711, - "tests/test_grid.py::test_find_most_rational_surfaces": 2.3906614780426025, - "tests/test_grid.py::test_find_least_rational_surfaces": 5.716690540313721, - "tests/test_input_output.py::test_vmec_input": 1.593571662902832, - "tests/test_input_output.py::test_write_desc_input_Nones": 1.6852741241455078, + "tests/test_examples.py::test_freeb_axisym": 165.2098564740445, + "tests/test_examples.py::test_freeb_vacuum": 666.5655182629998, + "tests/test_examples.py::test_multiobject_optimization_al": 121.51872604500022, + "tests/test_examples.py::test_multiobject_optimization_prox": 181.0960068869972, + "tests/test_examples.py::test_non_eq_optimization": 45.560406255997805, + "tests/test_examples.py::test_only_non_eq_optimization": 11.020795070999156, + "tests/test_examples.py::test_qh_optimization": 1103.3707119170285, + "tests/test_examples.py::test_simsopt_QH_comparison": 688.8956553020107, + "tests/test_examples.py::test_solve_bounds": 43.388052278984105, + "tests/test_geometry.py::test_rotation_matrix": 3.0645941510010744, + "tests/test_geometry.py::test_rpz2xyz": 2.1571697639992635, + "tests/test_geometry.py::test_xyz2rpz": 2.182286054998258, + "tests/test_grid.py::TestGrid::test_change_resolution": 1.542892319997918, + "tests/test_grid.py::TestGrid::test_compress_expand_inverse_op": 1.5772511959985422, + "tests/test_grid.py::TestGrid::test_concentric_grid": 1.5311077970000042, + "tests/test_grid.py::TestGrid::test_concentric_grid_high_res": 1.693702033997397, + "tests/test_grid.py::TestGrid::test_custom_grid": 2.8138667390048795, + "tests/test_grid.py::TestGrid::test_enforce_symmetry": 1.6266411370015703, + "tests/test_grid.py::TestGrid::test_linear_grid": 1.5034972640023625, + "tests/test_grid.py::TestGrid::test_linear_grid_spacing_consistency": 1.5571907359990291, + "tests/test_grid.py::TestGrid::test_linear_grid_spacing_two_nodes": 1.5128135260019917, + "tests/test_grid.py::TestGrid::test_linear_grid_symmetric_nodes_consistency": 1.5305054739983461, + "tests/test_grid.py::TestGrid::test_node_spacing_non_sym": 1.6383949640003266, + "tests/test_grid.py::TestGrid::test_node_spacing_sym": 1.644689387998369, + "tests/test_grid.py::TestGrid::test_quad_grid_volume_integration": 14.33542403300089, + "tests/test_grid.py::TestGrid::test_quadrature_grid": 1.5233443490033096, + "tests/test_grid.py::TestGrid::test_repr": 1.529734384996118, + "tests/test_grid.py::TestGrid::test_spacing_when_duplicate_node_is_removed": 1.5191250130010303, + "tests/test_grid.py::TestGrid::test_symmetry_spacing_basic": 1.5363636809997843, + "tests/test_grid.py::TestGrid::test_symmetry_volume_integral": 1.6720038189996558, + "tests/test_grid.py::test_custom_jitable_grid_indexing": 14.675376830004097, + "tests/test_grid.py::test_custom_jitable_grid_weights": 3.068923705002817, + "tests/test_grid.py::test_find_least_rational_surfaces": 8.115539762999106, + "tests/test_grid.py::test_find_most_rational_surfaces": 5.016010666000511, + "tests/test_input_output.py::TestInputReader::test_min_input": 1.5275491909997072, + "tests/test_input_output.py::TestInputReader::test_no_input_file": 1.5882188270006736, + "tests/test_input_output.py::TestInputReader::test_node_pattern_warning": 1.9307869609983754, + "tests/test_input_output.py::TestInputReader::test_nonexistant_input_file": 1.51432798600581, + "tests/test_input_output.py::TestInputReader::test_np_environ": 1.53816907700093, + "tests/test_input_output.py::TestInputReader::test_quiet_verbose": 1.5910946509975474, + "tests/test_input_output.py::TestInputReader::test_vacuum_objective_with_iota_yields_current": 1.6110654220065044, + "tests/test_input_output.py::test_ascii_io": 30.808994811002776, + "tests/test_input_output.py::test_copy": 6.101126101002592, + "tests/test_input_output.py::test_desc_output_to_input": 16.195360117995733, "tests/test_input_output.py::test_descout_to_input": 3.763477087020874, - "tests/test_input_output.py::test_near_axis_input_files": 10.827251434326172, - "tests/test_input_output.py::test_vmec_input_surface_threshold": 1.593571662902832, - "tests/test_input_output.py::TestInputReader::test_no_input_file": 1.4909934997558594, - "tests/test_input_output.py::TestInputReader::test_nonexistant_input_file": 1.5385918617248535, - "tests/test_input_output.py::TestInputReader::test_min_input": 1.5834696292877197, - "tests/test_input_output.py::TestInputReader::test_np_environ": 1.4599570035934448, - "tests/test_input_output.py::TestInputReader::test_quiet_verbose": 1.506763219833374, - "tests/test_input_output.py::TestInputReader::test_vacuum_objective_with_iota_yields_current": 1.6538304090499878, - "tests/test_input_output.py::test_writer_given_filename": 1.5545578002929688, - "tests/test_input_output.py::test_writer_given_file": 1.5664550065994263, - "tests/test_input_output.py::test_writer_close_on_delete": 1.4603271484375, - "tests/test_input_output.py::test_writer_write_dict": 1.510363221168518, - "tests/test_input_output.py::test_writer_write_list": 1.5396735668182373, - "tests/test_input_output.py::test_writer_write_obj": 1.5099352598190308, - "tests/test_input_output.py::test_reader_given_filename": 1.5020062923431396, - "tests/test_input_output.py::test_reader_given_file": 1.5020062923431396, - "tests/test_input_output.py::test_reader_read_obj": 1.4568850994110107, - "tests/test_input_output.py::test_pickle_io": 2.3418238162994385, - "tests/test_input_output.py::test_ascii_io": 12.55675745010376, - "tests/test_input_output.py::test_copy": 2.303689479827881, - "tests/test_input_output.py::test_save_none": 2.2719184160232544, - "tests/test_input_output.py::test_load_eq_without_current": 3.3425700664520264, - "tests/test_input_output.py::test_io_SplineMagneticField": 3.267856240272522, - "tests/test_interpolate.py::TestInterp1D::test_interp1d": 5.272557854652405, - "tests/test_interpolate.py::TestInterp1D::test_interp1d_extrap_periodic": 5.272557854652405, - "tests/test_interpolate.py::TestInterp1D::test_interp1d_monotonic": 5.272557854652405, - "tests/test_interpolate.py::TestInterp2D::test_interp2d": 6.148774027824402, - "tests/test_interpolate.py::TestInterp3D::test_interp3d": 8.586695075035095, - "tests/test_linear_objectives.py::test_LambdaGauge_sym": 2.644054412841797, - "tests/test_linear_objectives.py::test_LambdaGauge_asym": 2.5627944469451904, - "tests/test_linear_objectives.py::test_bc_on_interior_surfaces": 34.068824768066406, - "tests/test_linear_objectives.py::test_constrain_bdry_with_only_one_mode": 2.1073665618896484, - "tests/test_linear_objectives.py::test_constrain_asserts": 96.83970987796783, - "tests/test_linear_objectives.py::test_fixed_mode_solve": 41.12145459651947, - "tests/test_linear_objectives.py::test_fixed_modes_solve": 54.25764226913452, - "tests/test_linear_objectives.py::test_fixed_axis_and_theta_SFL_solve": 29.61748719215393, - "tests/test_linear_objectives.py::test_factorize_linear_constraints_asserts": 6.690843820571899, - "tests/test_linear_objectives.py::test_build_init": 11.40738296508789, - "tests/test_linear_objectives.py::test_kinetic_constraints": 4.064720392227173, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes": 38.067320704460144, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target": 38.067320704460144, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_axis": 38.067320704460144, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target_axis": 38.067320704460144, - "tests/test_linear_objectives.py::test_FixBoundary_with_single_weight": 4.974437475204468, - "tests/test_linear_objectives.py::test_FixBoundary_passed_target_no_passed_modes_error": 2.0891618728637695, - "tests/test_linear_objectives.py::test_FixAxis_passed_target_no_passed_modes_error": 1.9397294521331787, - "tests/test_linear_objectives.py::test_FixMode_passed_target_no_passed_modes_error": 2.0636954307556152, - "tests/test_linear_objectives.py::test_FixSumModes_passed_target_too_long": 1.830045461654663, - "tests/test_linear_objectives.py::test_FixMode_False_or_None_modes": 1.8329894542694092, - "tests/test_linear_objectives.py::test_FixSumModes_False_or_None_modes": 1.6984522342681885, - "tests/test_linear_objectives.py::test_FixAxis_util_correct_objectives": 1.736714482307434, - "tests/test_linear_objectives.py::test_FixNAE_util_correct_objectives": 2.9634666442871094, - "tests/test_magnetic_fields.py::TestMagneticFields::test_basic_fields": 2.0924978256225586, - "tests/test_magnetic_fields.py::TestMagneticFields::test_scalar_field": 2.300666332244873, - "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field": 12.266662836074829, - "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field_axisym": 12.266662836074829, - "tests/test_magnetic_fields.py::TestMagneticFields::test_field_line_integrate": 4.013031244277954, - "tests/test_magnetic_fields.py::TestMagneticFields::test_Bnormal_calculation": 4.293668150901794, - "tests/test_magnetic_fields.py::TestMagneticFields::test_Bnormal_save_and_load_DSHAPE": 17.31321120262146, - "tests/test_magnetic_fields.py::TestMagneticFields::test_Bnormal_save_and_load_HELIOTRON": 22.930991649627686, - "tests/test_objective_funs.py::TestObjectiveFunction::test_generic": 7.045788288116455, - "tests/test_objective_funs.py::TestObjectiveFunction::test_objective_from_user": 3.0326939821243286, - "tests/test_objective_funs.py::TestObjectiveFunction::test_volume": 5.788979768753052, - "tests/test_objective_funs.py::TestObjectiveFunction::test_aspect_ratio": 5.5107598304748535, - "tests/test_objective_funs.py::TestObjectiveFunction::test_elongation": 8.7400062084198, - "tests/test_objective_funs.py::TestObjectiveFunction::test_energy": 5.575179576873779, - "tests/test_objective_funs.py::TestObjectiveFunction::test_target_iota": 4.241585612297058, - "tests/test_objective_funs.py::TestObjectiveFunction::test_target_shear": 9.692948698997498, - "tests/test_objective_funs.py::TestObjectiveFunction::test_toroidal_current": 6.84657621383667, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qa_boozer": 16.396090507507324, - "tests/test_objective_funs.py::TestObjectiveFunction::test_jax_compile_boozer": 22.362881541252136, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qh_boozer": 24.601386547088623, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_twoterm": 28.74255645275116, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_tripleproduct": 13.618701815605164, - "tests/test_objective_funs.py::TestObjectiveFunction::test_isodynamicity": 6.120284557342529, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_boozer_grids": 5.078736186027527, - "tests/test_objective_funs.py::TestObjectiveFunction::test_mercier_stability": 55.6435444355011, - "tests/test_objective_funs.py::TestObjectiveFunction::test_magnetic_well": 15.848889827728271, - "tests/test_objective_funs.py::test_derivative_modes": 202.97485494613647, - "tests/test_objective_funs.py::test_rejit": 2.860276937484741, - "tests/test_objective_funs.py::test_generic_compute": 7.045788288116455, - "tests/test_objective_funs.py::test_getter_setter": 2.431445837020874, - "tests/test_objective_funs.py::test_bounds_format": 2.1949069499969482, - "tests/test_objective_funs.py::test_target_profiles": 11.689410924911499, - "tests/test_objective_funs.py::test_plasma_vessel_distance": 19.831424236297607, - "tests/test_objective_funs.py::test_mean_curvature": 17.688599348068237, - "tests/test_objective_funs.py::test_principal_curvature": 11.401215553283691, - "tests/test_objective_funs.py::test_field_scale_length": 58.9319269657135, - "tests/test_objective_funs.py::test_profile_objective_print": 8.955283880233765, - "tests/test_objective_funs.py::test_plasma_vessel_distance_print": 11.333454370498657, - "tests/test_objective_funs.py::test_rebuild": 100.36731946468353, - "tests/test_objective_funs.py::test_jvp_scaled": 6.008568048477173, - "tests/test_objective_funs.py::test_vjp": 49.77897334098816, - "tests/test_objective_funs.py::test_objective_target_bounds": 5.067176938056946, - "tests/test_objective_funs.py::test_softmax_and_softmin": 2.521897792816162, - "tests/test_optimizer.py::TestFmin::test_convex_full_hess_dogleg": 5.1186981201171875, - "tests/test_optimizer.py::TestFmin::test_convex_full_hess_subspace": 6.205466270446777, - "tests/test_optimizer.py::TestFmin::test_convex_full_hess_exact": 6.87457013130188, - "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_dogleg": 5.980814456939697, - "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_subspace": 6.666441917419434, - "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_exact": 6.87121319770813, - "tests/test_optimizer.py::TestSGD::test_sgd_convex": 3.96983540058136, - "tests/test_optimizer.py::TestLSQTR::test_lsqtr_exact": 6.493544340133667, - "tests/test_optimizer.py::test_no_iterations": 5.529320478439331, - "tests/test_optimizer.py::test_overstepping": 444.7169290781021, - "tests/test_optimizer.py::test_maxiter_1_and_0_solve": 66.5764102935791, - "tests/test_optimizer.py::test_scipy_fail_message": 40.96630001068115, - "tests/test_optimizer.py::test_not_implemented_error": 1.752349853515625, - "tests/test_optimizer.py::test_wrappers": 36.56393492221832, - "tests/test_optimizer.py::test_all_optimizers": 100.39988565444946, - "tests/test_optimizer.py::test_scipy_constrained_solve": 586.6609632968903, - "tests/test_optimizer.py::test_solve_with_x_scale": 54.16777300834656, - "tests/test_optimizer.py::test_bounded_optimization": 11.421349048614502, - "tests/test_optimizer.py::test_auglag": 66.71971952915192, - "tests/test_optimizer.py::test_constrained_AL_lsq": 78.91495037078857, - "tests/test_optimizer.py::test_constrained_AL_scalar": 356.7853684425354, - "tests/test_perturbations.py::test_perturbation_orders": 2968.4424810409546, - "tests/test_perturbations.py::test_perturb_with_float_without_error": 11.654300451278687, - "tests/test_perturbations.py::test_optimal_perturb": 103.26844382286072, - "tests/test_perturbations.py::test_perturb_axis": 46.70373606681824, - "tests/test_plotting.py::test_kwarg_warning": 2.9493203163146973, - "tests/test_plotting.py::test_kwarg_future_warning": 10.476717233657837, - "tests/test_plotting.py::test_1d_p": 4.144308567047119, - "tests/test_plotting.py::test_1d_fsa_consistency": 55.39805483818054, - "tests/test_plotting.py::test_1d_dpdr": 4.301635503768921, - "tests/test_plotting.py::test_1d_iota": 7.030443429946899, - "tests/test_plotting.py::test_1d_iota_radial": 7.030443429946899, - "tests/test_plotting.py::test_1d_logpsi": 3.6051418781280518, - "tests/test_plotting.py::test_2d_logF": 13.238114833831787, - "tests/test_plotting.py::test_2d_g_tz": 5.49422812461853, - "tests/test_plotting.py::test_2d_g_rz": 8.476567149162292, - "tests/test_plotting.py::test_2d_lambda": 4.4203550815582275, - "tests/test_plotting.py::test_3d_B": 7.013176918029785, - "tests/test_plotting.py::test_3d_J": 10.500922203063965, - "tests/test_plotting.py::test_3d_tz": 10.512905597686768, - "tests/test_plotting.py::test_3d_rz": 7.675700902938843, - "tests/test_plotting.py::test_3d_rt": 7.302988767623901, - "tests/test_plotting.py::test_plot_fsa_axis_limit": 34.448630571365356, - "tests/test_plotting.py::test_fsa_I": 10.042901039123535, - "tests/test_plotting.py::test_fsa_G": 12.94868779182434, - "tests/test_plotting.py::test_fsa_F_normalized": 15.39680826663971, - "tests/test_plotting.py::test_section_J": 13.633682370185852, - "tests/test_plotting.py::test_section_Z": 8.474269151687622, - "tests/test_plotting.py::test_section_R": 9.828718423843384, - "tests/test_plotting.py::test_section_F": 13.354544162750244, - "tests/test_plotting.py::test_section_F_normalized_vac": 15.453603267669678, - "tests/test_plotting.py::test_section_logF": 10.222967624664307, - "tests/test_plotting.py::test_plot_surfaces": 17.72052574157715, - "tests/test_plotting.py::test_plot_surfaces_no_theta": 11.31156075000763, - "tests/test_plotting.py::test_plot_boundary": 9.28933835029602, - "tests/test_plotting.py::test_plot_boundaries": 17.290244579315186, - "tests/test_plotting.py::test_plot_comparison": 24.554563283920288, - "tests/test_plotting.py::test_plot_comparison_no_theta": 11.492709875106812, - "tests/test_plotting.py::test_plot_con_basis": 5.39559268951416, - "tests/test_plotting.py::test_plot_cov_basis": 4.859344482421875, - "tests/test_plotting.py::test_plot_magnetic_tension": 7.140878438949585, - "tests/test_plotting.py::test_plot_magnetic_pressure": 7.344560861587524, - "tests/test_plotting.py::test_plot_gradpsi": 4.3144437074661255, - "tests/test_plotting.py::test_plot_normF_2d": 11.558874249458313, - "tests/test_plotting.py::test_plot_normF_section": 14.094926834106445, - "tests/test_plotting.py::test_plot_coefficients": 4.083799958229065, - "tests/test_plotting.py::test_plot_logo": 4.796693801879883, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_linear": 2.3446993827819824, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_quad": 2.4298858642578125, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_jacobi": 2.456006646156311, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb1": 2.446089506149292, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb2": 2.553113579750061, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_ocs": 2.441372871398926, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_powerseries": 2.3069249391555786, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierseries": 2.7168502807617188, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_doublefourierseries": 11.475860476493835, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierzernike": 39.484002351760864, - "tests/test_plotting.py::TestPlotFieldLines::test_find_idx": 2.351112127304077, - "tests/test_plotting.py::TestPlotFieldLines::test_plot_field_line": 13.026328086853027, - "tests/test_plotting.py::TestPlotFieldLines::test_plot_field_lines": 13.026328086853027, - "tests/test_plotting.py::test_plot_boozer_modes": 61.49421417713165, - "tests/test_plotting.py::test_plot_boozer_surface": 17.859545946121216, - "tests/test_plotting.py::test_plot_qs_error": 190.67440521717072, - "tests/test_plotting.py::test_plot_coils": 8.402854442596436, - "tests/test_plotting.py::test_plot_b_mag": 17.19118618965149, - "tests/test_plotting.py::test_plot_surfaces_HELIOTRON": 166.8380891084671, - "tests/test_profiles.py::TestProfiles::test_same_result": 149.00100541114807, - "tests/test_profiles.py::TestProfiles::test_close_values": 14.516500115394592, - "tests/test_profiles.py::TestProfiles::test_PowerSeriesProfile_even_sym": 6.23789656162262, - "tests/test_profiles.py::TestProfiles::test_SplineProfile_methods": 8.58802855014801, - "tests/test_profiles.py::TestProfiles::test_repr": 2.206360936164856, - "tests/test_profiles.py::TestProfiles::test_get_set": 3.9549028873443604, - "tests/test_profiles.py::TestProfiles::test_auto_sym": 2.459155559539795, - "tests/test_profiles.py::TestProfiles::test_sum_profiles": 7.875421047210693, - "tests/test_profiles.py::TestProfiles::test_product_profiles": 6.731966733932495, - "tests/test_profiles.py::TestProfiles::test_product_profiles_derivative": 6.731966733932495, - "tests/test_profiles.py::TestProfiles::test_scaled_profiles": 2.645738363265991, - "tests/test_profiles.py::TestProfiles::test_profile_errors": 5.228930950164795, - "tests/test_profiles.py::TestProfiles::test_default_profiles": 5.27898907661438, - "tests/test_profiles.py::TestProfiles::test_solve_with_combined": 57.84936344623566, - "tests/test_profiles.py::TestProfiles::test_kinetic_pressure": 6.029354810714722, - "tests/test_random.py::test_random_pressure": 9.484513282775879, - "tests/test_random.py::test_random_surface": 15.788919448852539, - "tests/test_stability_funs.py::test_mercier_vacuum": 9.476405262947083, - "tests/test_stability_funs.py::test_compute_d_shear": 22.89810848236084, - "tests/test_stability_funs.py::test_compute_d_current": 30.48818826675415, - "tests/test_stability_funs.py::test_compute_d_well": 37.768325090408325, - "tests/test_stability_funs.py::test_compute_d_geodesic": 34.81776785850525, - "tests/test_stability_funs.py::test_compute_d_mercier": 34.14472937583923, - "tests/test_stability_funs.py::test_compute_magnetic_well": 31.520386695861816, - "tests/test_stability_funs.py::test_mercier_print": 14.523894309997559, - "tests/test_stability_funs.py::test_magwell_print": 16.67170536518097, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_area": 6.229729056358337, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_compute_ndarray_error": 2.4868667125701904, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_normal": 3.5685609579086304, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_misc": 2.985080599784851, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_input_file": 7.244992971420288, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_near_axis": 3.170508623123169, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_curvature": 2.93274188041687, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_area": 6.229729056358337, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_normal": 3.5685609579086304, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_misc": 2.985080599784851, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_curvature": 2.93274188041687, - "tests/test_surfaces.py::test_surface_orientation": 6.0597487688064575, - "tests/test_transform.py::TestTransform::test_eq": 2.918691635131836, - "tests/test_transform.py::TestTransform::test_transform_order_error": 2.4013243913650513, - "tests/test_transform.py::TestTransform::test_profile": 2.5705788135528564, - "tests/test_transform.py::TestTransform::test_surface": 3.5995852947235107, - "tests/test_transform.py::TestTransform::test_volume_chebyshev": 3.6701056957244873, - "tests/test_transform.py::TestTransform::test_volume_zernike": 3.659238815307617, - "tests/test_transform.py::TestTransform::test_set_grid": 3.2190189361572266, - "tests/test_transform.py::TestTransform::test_set_basis": 2.969343900680542, - "tests/test_transform.py::TestTransform::test_fft": 3.400732159614563, - "tests/test_transform.py::TestTransform::test_direct_fft_equal": 11.848833560943604, - "tests/test_transform.py::TestTransform::test_project": 7.40645694732666, - "tests/test_transform.py::TestTransform::test_fft_warnings": 5.197953701019287, - "tests/test_transform.py::TestTransform::test_direct2_warnings": 2.8224503993988037, - "tests/test_transform.py::TestTransform::test_fit_direct1": 4.676076889038086, - "tests/test_transform.py::TestTransform::test_fit_direct2": 3.7976577281951904, - "tests/test_transform.py::TestTransform::test_fit_fft": 4.245433330535889, - "tests/test_transform.py::TestTransform::test_empty_grid": 3.6825473308563232, - "tests/test_transform.py::TestTransform::test_Z_projection": 7.416264057159424, - "tests/test_transform.py::test_transform_pytree": 5.119142532348633, - "tests/test_transform.py::test_NFP_warning": 4.844725966453552, - "tests/test_utils.py::test_isalmostequal": 3.151520013809204, - "tests/test_utils.py::test_islinspaced": 3.082701086997986, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd": 2.498995542526245, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd_sin_series": 2.498995542526245, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev": 2.3764783143997192, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev_sin_sym": 2.3764783143997192, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identities_inverse": 3.0618289709091187, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_linear_transform": 4.089492082595825, - "tests/test_vmec.py::TestVMECIO::test_fourier_to_zernike": 2.991957187652588, - "tests/test_vmec.py::TestVMECIO::test_zernike_to_fourier": 3.233712911605835, - "tests/test_vmec.py::test_vmec_load_profiles": 49.678714752197266, - "tests/test_vmec.py::test_load_then_save": 97.55081272125244, - "tests/test_vmec.py::test_axis_surf_after_load": 44.66222405433655, - "tests/test_vmec.py::test_vmec_save_asym": 60.512279748916626, - "tests/test_vmec.py::test_vmec_save_1": 2.5306615829467773, - "tests/test_vmec.py::test_vmec_save_2": 2.9303231239318848, - "tests/test_vmec.py::test_plot_vmec_comparison": 11.675038933753967, - "tests/test_vmec.py::test_vmec_boundary_subspace": 2.914440870285034 + "tests/test_input_output.py::test_io_SplineMagneticField": 3.2883595250023063, + "tests/test_input_output.py::test_load_eq_without_current": 3.6407579550032096, + "tests/test_input_output.py::test_near_axis_input_files": 6.113512972002354, + "tests/test_input_output.py::test_pickle_io": 3.23799625600077, + "tests/test_input_output.py::test_reader_given_file": 1.5546266630008176, + "tests/test_input_output.py::test_reader_given_filename": 1.55079097899943, + "tests/test_input_output.py::test_reader_read_obj": 1.5273862780013587, + "tests/test_input_output.py::test_save_after_load": 12.613824988999113, + "tests/test_input_output.py::test_save_none": 12.252932658997452, + "tests/test_input_output.py::test_vmec_input": 1.7186922430046252, + "tests/test_input_output.py::test_vmec_input_surface_threshold": 4.158838191000541, + "tests/test_input_output.py::test_write_desc_input_Nones": 1.5303001239990408, + "tests/test_input_output.py::test_writer_close_on_delete": 1.5522544029990968, + "tests/test_input_output.py::test_writer_given_file": 1.6992341809964273, + "tests/test_input_output.py::test_writer_given_filename": 1.557395614003326, + "tests/test_input_output.py::test_writer_write_dict": 1.5580429099973117, + "tests/test_input_output.py::test_writer_write_list": 1.5497696940001333, + "tests/test_input_output.py::test_writer_write_obj": 1.5774160399996617, + "tests/test_linear_objectives.py::test_FixAxis_passed_target_no_passed_modes_error": 13.526608840002154, + "tests/test_linear_objectives.py::test_FixAxis_util_correct_objectives": 13.284155260000261, + "tests/test_linear_objectives.py::test_FixBoundary_passed_target_no_passed_modes_error": 12.836838494000403, + "tests/test_linear_objectives.py::test_FixBoundary_with_single_weight": 15.53305328699571, + "tests/test_linear_objectives.py::test_FixMode_False_or_None_modes": 12.976503283000056, + "tests/test_linear_objectives.py::test_FixMode_passed_target_no_passed_modes_error": 12.958678402996156, + "tests/test_linear_objectives.py::test_FixNAE_util_correct_objectives": 13.359161319996929, + "tests/test_linear_objectives.py::test_FixSumModes_False_or_None_modes": 13.134607334002794, + "tests/test_linear_objectives.py::test_FixSumModes_passed_target_too_long": 13.388925283998105, + "tests/test_linear_objectives.py::test_LambdaGauge_asym": 17.750804241997685, + "tests/test_linear_objectives.py::test_LambdaGauge_sym": 3.426120066000294, + "tests/test_linear_objectives.py::test_bc_on_interior_surfaces": 39.445212638001976, + "tests/test_linear_objectives.py::test_build_init": 19.491527664002206, + "tests/test_linear_objectives.py::test_constrain_asserts": 19.33108093199553, + "tests/test_linear_objectives.py::test_constrain_bdry_with_only_one_mode": 13.012020947997371, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes": 21.17463499099904, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target": 20.90506877000007, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target_axis": 23.282829511001182, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_axis": 21.98586076799984, + "tests/test_linear_objectives.py::test_factorize_linear_constraints_asserts": 18.34482362300332, + "tests/test_linear_objectives.py::test_fixed_axis_and_theta_SFL_solve": 35.77478713099845, + "tests/test_linear_objectives.py::test_fixed_mode_solve": 46.75125164797646, + "tests/test_linear_objectives.py::test_fixed_modes_solve": 42.279382451000856, + "tests/test_linear_objectives.py::test_kinetic_constraints": 16.506718756998453, + "tests/test_magnetic_fields.py::TestMagneticFields::test_Bnormal_calculation": 4.837572751002881, + "tests/test_magnetic_fields.py::TestMagneticFields::test_Bnormal_save_and_load_DSHAPE": 7.48490697899615, + "tests/test_magnetic_fields.py::TestMagneticFields::test_Bnormal_save_and_load_HELIOTRON": 16.140960685999744, + "tests/test_magnetic_fields.py::TestMagneticFields::test_basic_fields": 3.0971555670003, + "tests/test_magnetic_fields.py::TestMagneticFields::test_change_Phi_basis_fourier_current_field": 2.167428494998603, + "tests/test_magnetic_fields.py::TestMagneticFields::test_combined_fields": 2.3817777869953716, + "tests/test_magnetic_fields.py::TestMagneticFields::test_current_potential_field": 7.338101552006265, + "tests/test_magnetic_fields.py::TestMagneticFields::test_field_line_integrate": 2.5038658199991914, + "tests/test_magnetic_fields.py::TestMagneticFields::test_field_line_integrate_bounds": 2.516783448001661, + "tests/test_magnetic_fields.py::TestMagneticFields::test_fourier_current_potential_field": 10.437403561001702, + "tests/test_magnetic_fields.py::TestMagneticFields::test_fourier_current_potential_field_asserts": 1.7217379279973102, + "tests/test_magnetic_fields.py::TestMagneticFields::test_fourier_current_potential_field_symmetry": 1.8540061119929305, + "tests/test_magnetic_fields.py::TestMagneticFields::test_init_Phi_mn_fourier_current_field": 2.801191710997955, + "tests/test_magnetic_fields.py::TestMagneticFields::test_io_fourier_current_field": 1.8819897070025036, + "tests/test_magnetic_fields.py::TestMagneticFields::test_mgrid_io": 6.422455993000767, + "tests/test_magnetic_fields.py::TestMagneticFields::test_scalar_field": 2.2889806749990385, + "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field": 8.430371877999278, + "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field_axisym": 3.1920034830000077, + "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field_jit": 5.132572205999168, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_bootstrap": 10.938637015991844, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_boundary_error": 28.82819817197742, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_generic_profile": 8.87665598798776, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_generic_scalar": 15.481857894017594, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_generic_volume": 9.822074825002346, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_mercier": 21.103813255031127, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[AspectRatio]": 4.525908992975019, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[BScaleLength]": 6.885055861988803, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[CurrentDensity]": 7.6574104170140345, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[Elongation]": 4.1126805539825, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[Energy]": 4.3215977270156145, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[ForceBalanceAnisotropic]": 5.083369603002211, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[ForceBalance]": 5.81725922197802, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[GoodCoordinates]": 4.3965783509775065, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[HelicalForceBalance]": 4.42610470901127, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[Isodynamicity]": 6.47041044497746, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[MagneticWell]": 5.655285602988442, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[MeanCurvature]": 4.5763303140120115, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[Pressure]": 3.3859170679643285, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[PrincipalCurvature]": 3.563800739997532, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[QuasisymmetryBoozer]": 9.03863914401154, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[QuasisymmetryTripleProduct]": 4.531183171027806, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[QuasisymmetryTwoTerm]": 4.539271248009754, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[RadialForceBalance]": 4.92413304198999, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[RotationalTransform]": 3.10291648801649, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[Shear]": 4.324439108022489, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[ToroidalCurrent]": 4.580668631970184, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_others[Volume]": 3.898456078983145, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_plasma_vessel": 21.156958234030753, + "tests/test_objective_funs.py::TestComputeScalarResolution::test_compute_scalar_resolution_vacuum_boundary_error": 5.402465846011182, + "tests/test_objective_funs.py::TestObjectiveFunction::test_aspect_ratio": 3.3635769089996757, + "tests/test_objective_funs.py::TestObjectiveFunction::test_boundary_error_biest": 21.398838486995373, + "tests/test_objective_funs.py::TestObjectiveFunction::test_boundary_error_biestsc": 36.988582465000945, + "tests/test_objective_funs.py::TestObjectiveFunction::test_boundary_error_nestor": 19.785346114000276, + "tests/test_objective_funs.py::TestObjectiveFunction::test_boundary_error_vacuum": 17.436402035003994, + "tests/test_objective_funs.py::TestObjectiveFunction::test_elongation": 8.438965100001951, + "tests/test_objective_funs.py::TestObjectiveFunction::test_energy": 3.352779314001964, + "tests/test_objective_funs.py::TestObjectiveFunction::test_field_scale_length": 37.18068048100031, + "tests/test_objective_funs.py::TestObjectiveFunction::test_generic": 16.024178745003155, + "tests/test_objective_funs.py::TestObjectiveFunction::test_isodynamicity": 3.724591848000273, + "tests/test_objective_funs.py::TestObjectiveFunction::test_jax_compile_boozer": 25.236037753005803, + "tests/test_objective_funs.py::TestObjectiveFunction::test_linear_objective_from_user": 2.844101209000655, + "tests/test_objective_funs.py::TestObjectiveFunction::test_magnetic_well": 3.698493749005138, + "tests/test_objective_funs.py::TestObjectiveFunction::test_mean_curvature": 8.964764428004855, + "tests/test_objective_funs.py::TestObjectiveFunction::test_mercier_stability": 4.598655678000796, + "tests/test_objective_funs.py::TestObjectiveFunction::test_objective_from_user": 2.555016718000843, + "tests/test_objective_funs.py::TestObjectiveFunction::test_plasma_vessel_distance": 10.042417220996867, + "tests/test_objective_funs.py::TestObjectiveFunction::test_pressure": 2.6660704920032003, + "tests/test_objective_funs.py::TestObjectiveFunction::test_principal_curvature": 7.934462338998856, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qa_boozer": 15.994433098003356, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qh_boozer": 13.213613214003999, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_boozer_grids": 3.10034963799626, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_tripleproduct": 5.860502887000621, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_twoterm": 18.462982232002105, + "tests/test_objective_funs.py::TestObjectiveFunction::test_target_iota": 3.0118673049983045, + "tests/test_objective_funs.py::TestObjectiveFunction::test_target_max_iota": 4.955381063999084, + "tests/test_objective_funs.py::TestObjectiveFunction::test_target_mean_iota": 17.385818001999723, + "tests/test_objective_funs.py::TestObjectiveFunction::test_target_min_iota": 4.972477066003194, + "tests/test_objective_funs.py::TestObjectiveFunction::test_target_shear": 4.190596986998571, + "tests/test_objective_funs.py::TestObjectiveFunction::test_toroidal_current": 3.322180211998784, + "tests/test_objective_funs.py::TestObjectiveFunction::test_volume": 3.512971648997336, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[AspectRatio]": 5.052777159002289, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[BScaleLength]": 5.441058575997886, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[CurrentDensity]": 3.7936064269997587, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[Elongation]": 3.3977684090023104, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[Energy]": 3.777643142002489, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[ForceBalance]": 4.325833215996681, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[GoodCoordinates]": 3.230575322999357, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[HelicalForceBalance]": 3.682147666000674, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[Isodynamicity]": 3.772865343002195, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[MagneticWell]": 4.302671537006972, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[MeanCurvature]": 3.3762891870028398, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[MercierStability]": 4.685184687998117, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[Pressure]": 2.84491098700164, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[PrincipalCurvature]": 3.1221340559968667, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[QuasisymmetryBoozer]": 10.952548792000016, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[QuasisymmetryTripleProduct]": 4.4712163519943715, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[QuasisymmetryTwoTerm]": 3.822364287003438, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[RadialForceBalance]": 4.109406987001421, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[RotationalTransform]": 3.133018026001082, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[Shear]": 3.763396199003182, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[ToroidalCurrent]": 3.189751669000543, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad[Volume]": 2.934465738002473, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad_anisotropy": 12.26007259799735, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad_bootstrap": 12.399535646003642, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad_boundary_error": 29.14739112600364, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad_plasma_vessel": 24.57573320800293, + "tests/test_objective_funs.py::TestObjectiveNaNGrad::test_objective_no_nangrad_vacuum_boundary_error": 4.006381560000591, + "tests/test_objective_funs.py::test_asymmetric_normalization": 17.90830730599555, + "tests/test_objective_funs.py::test_boundary_error_print": 31.719327356000576, + "tests/test_objective_funs.py::test_bounds_format": 12.815263700002106, + "tests/test_objective_funs.py::test_derivative_modes": 107.23479169298662, + "tests/test_objective_funs.py::test_getter_setter": 13.223397965000913, + "tests/test_objective_funs.py::test_jvp_scaled": 18.01336596600231, + "tests/test_objective_funs.py::test_loss_function_asserts": 12.929548473995965, + "tests/test_objective_funs.py::test_objective_fun_things": 28.585008473000926, + "tests/test_objective_funs.py::test_objective_target_bounds": 30.379466548998607, + "tests/test_objective_funs.py::test_plasma_vessel_distance_print": 20.774819416001264, + "tests/test_objective_funs.py::test_profile_objective_print": 21.027601255002082, + "tests/test_objective_funs.py::test_softmax_and_softmin": 3.491235953006253, + "tests/test_objective_funs.py::test_target_profiles": 28.2688175459989, + "tests/test_objective_funs.py::test_vjp": 24.924405622998165, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_lsq[lsq-auglag]": 15.054546579001908, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_lsq[lsq-exact]": 5.327012796999043, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_lsq[scipy-dogbox]": 4.609646222997981, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_lsq[scipy-lm]": 4.167987420998543, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_lsq[scipy-trf]": 4.473903854999662, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[fmin-auglag-bfgs]": 4.634653371998866, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[fmin-auglag]": 30.172641609002312, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[fmintr-bfgs]": 4.24471804699715, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[fmintr]": 7.447562761000881, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-CG]": 4.465172360003635, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-Newton-CG]": 4.494747399996413, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-SLSQP]": 4.2800934679980855, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-bfgs]": 4.672026657997776, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-dogleg]": 4.400435430998186, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-trust-constr]": 4.812981226001284, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-trust-exact]": 4.167372934000014, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-trust-krylov]": 4.438994068001193, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-trust-ncg]": 4.809485538004083, + "tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[sgd]": 4.846750221000548, + "tests/test_optimizer.py::TestFmin::test_convex_full_hess_dogleg": 5.363728240005003, + "tests/test_optimizer.py::TestFmin::test_convex_full_hess_exact": 3.4878661910006485, + "tests/test_optimizer.py::TestFmin::test_convex_full_hess_subspace": 3.0701052249933127, + "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_dogleg": 4.552741784002137, + "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_exact": 3.872639163000713, + "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_subspace": 3.162206826000329, + "tests/test_optimizer.py::TestLSQTR::test_lsqtr_exact": 7.8507093689986505, + "tests/test_optimizer.py::TestSGD::test_sgd_convex": 4.2401257020028424, + "tests/test_optimizer.py::test_LinearConstraint_jacobian": 54.30764804998762, + "tests/test_optimizer.py::test_auglag": 20.970893525001884, + "tests/test_optimizer.py::test_bounded_optimization": 11.128514849002386, + "tests/test_optimizer.py::test_constrained_AL_lsq": 148.80931450499338, + "tests/test_optimizer.py::test_constrained_AL_scalar": 100.66257204298745, + "tests/test_optimizer.py::test_maxiter_1_and_0_solve": 60.678130401996896, + "tests/test_optimizer.py::test_no_iterations": 4.720541481005057, + "tests/test_optimizer.py::test_not_implemented_error": 2.2516019210015656, + "tests/test_optimizer.py::test_optimize_multiple_things_different_order": 24.858137269999133, + "tests/test_optimizer.py::test_optimize_with_single_constraint": 20.4751513600022, + "tests/test_optimizer.py::test_overstepping": 254.93381108899484, + "tests/test_optimizer.py::test_proximal_jacobian": 277.7108911679825, + "tests/test_optimizer.py::test_scipy_constrained_solve": 120.41281288999016, + "tests/test_optimizer.py::test_scipy_fail_message": 51.639403586043045, + "tests/test_optimizer.py::test_solve_with_x_scale": 47.20604335801909, + "tests/test_optimizer.py::test_wrappers": 32.09230667100201, + "tests/test_perturbations.py::test_optimal_perturb": 45.21181147899915, + "tests/test_perturbations.py::test_perturb_axis": 41.52180520800175, + "tests/test_perturbations.py::test_perturb_with_float_without_error": 23.925306844997976, + "tests/test_perturbations.py::test_perturbation_orders": 255.70213129601325, + "tests/test_plotting.py::TestPlot1D::test_1d_elongation": 7.39161468698876, + "tests/test_plotting.py::TestPlot1D::test_1d_iota": 5.816118283997639, + "tests/test_plotting.py::TestPlot1D::test_1d_iota_radial": 6.645610381980077, + "tests/test_plotting.py::TestPlot1D::test_1d_logpsi": 0.6680448819970479, + "tests/test_plotting.py::TestPlot1D::test_1d_p": 1.8197079929959727, + "tests/test_plotting.py::TestPlot1D::test_plot_1d_curve": 0.9738203239976428, + "tests/test_plotting.py::TestPlot1D::test_plot_1d_surface": 1.4176946110092103, + "tests/test_plotting.py::TestPlot2D::test_2d_g_rz": 0.9800739470083499, + "tests/test_plotting.py::TestPlot2D::test_2d_g_tz": 4.02319536198047, + "tests/test_plotting.py::TestPlot2D::test_2d_logF": 11.32121155099594, + "tests/test_plotting.py::TestPlot2D::test_2d_plot_Bn": 17.361508290006896, + "tests/test_plotting.py::TestPlot2D::test_plot_2d_surface": 1.4505801339983009, + "tests/test_plotting.py::TestPlot2D::test_plot_con_basis": 1.4606957460055128, + "tests/test_plotting.py::TestPlot2D::test_plot_cov_basis": 0.6839415080030449, + "tests/test_plotting.py::TestPlot2D::test_plot_normF_2d": 9.59870039501402, + "tests/test_plotting.py::TestPlot3D::test_3d_plot_Bn": 12.080953058000887, + "tests/test_plotting.py::TestPlot3D::test_3d_rt": 0.6598505790025229, + "tests/test_plotting.py::TestPlot3D::test_3d_rz": 5.69653678398754, + "tests/test_plotting.py::TestPlot3D::test_3d_tz": 8.93025886900432, + "tests/test_plotting.py::TestPlot3D::test_plot_3d_surface": 1.828075653000269, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_doublefourierseries": 4.297180719993776, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierseries": 0.8649517310113879, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierzernike": 7.051884510976379, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_powerseries": 1.2842499969992787, + "tests/test_plotting.py::TestPlotBoozerModes::test_plot_boozer_modes": 12.516068537996034, + "tests/test_plotting.py::TestPlotBoozerModes::test_plot_boozer_modes_breaking_only": 4.926561071988544, + "tests/test_plotting.py::TestPlotBoozerModes::test_plot_boozer_modes_max": 4.42618233000394, + "tests/test_plotting.py::TestPlotBoozerModes::test_plot_boozer_modes_no_norm": 6.169192904999363, + "tests/test_plotting.py::TestPlotBoozerSurface::test_plot_boozer_surface": 10.770511821014225, + "tests/test_plotting.py::TestPlotBoozerSurface::test_plot_omnigenous_field": 2.5080141939979512, + "tests/test_plotting.py::TestPlotBoundary::test_plot_boundaries": 12.157784113995149, + "tests/test_plotting.py::TestPlotBoundary::test_plot_boundary": 8.667275105995941, + "tests/test_plotting.py::TestPlotBoundary::test_plot_boundary_surface": 3.225397364003584, + "tests/test_plotting.py::TestPlotComparison::test_plot_comparison": 68.2459130659845, + "tests/test_plotting.py::TestPlotComparison::test_plot_comparison_no_theta": 2.8547526549809845, + "tests/test_plotting.py::TestPlotFSA::test_fsa_F_normalized": 7.336087165997014, + "tests/test_plotting.py::TestPlotFSA::test_fsa_I": 9.353267086000415, + "tests/test_plotting.py::TestPlotFSA::test_plot_fsa_axis_limit": 15.69285398699867, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb1": 0.6533244179881876, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb2": 0.6581470770033775, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_jacobi": 0.6626311550062383, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_linear": 1.060141088993987, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_ocs": 0.6494206589850364, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_quad": 0.6860662049875828, + "tests/test_plotting.py::TestPlotSection::test_plot_normF_section": 8.501898501985124, + "tests/test_plotting.py::TestPlotSection::test_section_F": 2.231009934999747, + "tests/test_plotting.py::TestPlotSection::test_section_J": 11.778287986991927, + "tests/test_plotting.py::TestPlotSection::test_section_logF": 2.1756981370126596, + "tests/test_plotting.py::TestPlotSurfaces::test_plot_surfaces": 27.21949966799002, + "tests/test_plotting.py::TestPlotSurfaces::test_plot_surfaces_HELIOTRON": 23.4548534630012, + "tests/test_plotting.py::TestPlotSurfaces::test_plot_surfaces_no_theta": 1.3530529010022292, + "tests/test_plotting.py::test_kwarg_future_warning": 0.5322974350128789, + "tests/test_plotting.py::test_kwarg_warning": 15.721355953995953, + "tests/test_plotting.py::test_plot_b_mag": 17.60216506497818, + "tests/test_plotting.py::test_plot_coefficients": 2.7744905070139794, + "tests/test_plotting.py::test_plot_coils": 3.1842678650136804, + "tests/test_plotting.py::test_plot_logo": 2.637325112009421, + "tests/test_plotting.py::test_plot_qs_error": 22.63121982998564, + "tests/test_profiles.py::TestProfiles::test_PowerSeriesProfile_even_sym": 2.9862372129973664, + "tests/test_profiles.py::TestProfiles::test_SplineProfile_methods": 6.018229676999908, + "tests/test_profiles.py::TestProfiles::test_auto_sym": 2.5092655509979522, + "tests/test_profiles.py::TestProfiles::test_close_values": 16.05634071500026, + "tests/test_profiles.py::TestProfiles::test_default_profiles": 5.469991549001861, + "tests/test_profiles.py::TestProfiles::test_get_set": 3.3157052079986897, + "tests/test_profiles.py::TestProfiles::test_kinetic_pressure": 18.91112590099874, + "tests/test_profiles.py::TestProfiles::test_product_profiles": 5.031340598998213, + "tests/test_profiles.py::TestProfiles::test_product_profiles_derivative": 3.007900715005235, + "tests/test_profiles.py::TestProfiles::test_profile_errors": 2.9162500700040255, + "tests/test_profiles.py::TestProfiles::test_repr": 3.2363957870002196, + "tests/test_profiles.py::TestProfiles::test_same_result": 93.83305923800799, + "tests/test_profiles.py::TestProfiles::test_scaled_profiles": 2.6538628470043477, + "tests/test_profiles.py::TestProfiles::test_solve_with_combined": 54.50625906395726, + "tests/test_profiles.py::TestProfiles::test_sum_profiles": 6.9174393029970815, + "tests/test_random.py::test_random_pressure": 9.319775469000888, + "tests/test_random.py::test_random_surface": 25.388372931003687, + "tests/test_singularities.py::test_biest_interpolators": 4.716703420996055, + "tests/test_singularities.py::test_singular_integral_greens_id": 32.9272484800058, + "tests/test_singularities.py::test_singular_integral_vac_estell": 15.568489260003844, + "tests/test_stability_funs.py::test_compute_d_current": 24.52692030599792, + "tests/test_stability_funs.py::test_compute_d_geodesic": 27.397302482997475, + "tests/test_stability_funs.py::test_compute_d_mercier": 28.462580421000894, + "tests/test_stability_funs.py::test_compute_d_shear": 22.696304549997876, + "tests/test_stability_funs.py::test_compute_d_well": 25.56624591199943, + "tests/test_stability_funs.py::test_compute_magnetic_well": 24.079023827998753, + "tests/test_stability_funs.py::test_magwell_print": 15.778229646999534, + "tests/test_stability_funs.py::test_mercier_print": 23.06626790899827, + "tests/test_stability_funs.py::test_mercier_vacuum": 15.945780925998406, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_area": 5.243621954003174, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_compute_ndarray_error": 2.473588169996219, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_constant_offset_surface_circle": 8.832547789003002, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_constant_offset_surface_rot_ellipse": 21.54778232699755, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_curvature": 2.7695129120002093, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_input_file": 9.124532824000198, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_near_axis": 2.764549903997249, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_misc": 3.076208914997551, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_normal": 3.3924093739988166, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_position": 2.728874722997716, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_surface_from_values": 12.313251707000745, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_area": 16.957669101997453, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_curvature": 6.525485035999736, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_misc": 3.0236848810018273, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_normal": 8.12666688700483, + "tests/test_surfaces.py::test_surface_orientation": 17.786888768998324, + "tests/test_transform.py::TestTransform::test_Z_projection": 8.725720057998842, + "tests/test_transform.py::TestTransform::test_direct2_warnings": 2.876706176004518, + "tests/test_transform.py::TestTransform::test_direct_fft_equal": 15.199172624998027, + "tests/test_transform.py::TestTransform::test_empty_grid": 5.249251693003316, + "tests/test_transform.py::TestTransform::test_eq": 7.471260242004064, + "tests/test_transform.py::TestTransform::test_fft": 3.653772793997632, + "tests/test_transform.py::TestTransform::test_fft_warnings": 9.022689112996886, + "tests/test_transform.py::TestTransform::test_fit_direct1": 7.881664825999906, + "tests/test_transform.py::TestTransform::test_fit_direct2": 3.270915659995808, + "tests/test_transform.py::TestTransform::test_fit_fft": 6.026062988006743, + "tests/test_transform.py::TestTransform::test_profile": 2.6665852229998563, + "tests/test_transform.py::TestTransform::test_project": 14.427902427003573, + "tests/test_transform.py::TestTransform::test_set_basis": 7.858471010000358, + "tests/test_transform.py::TestTransform::test_set_grid": 8.912943017003272, + "tests/test_transform.py::TestTransform::test_surface": 3.587806878000265, + "tests/test_transform.py::TestTransform::test_transform_order_error": 2.790843536997272, + "tests/test_transform.py::TestTransform::test_volume_chebyshev": 3.4869363360012358, + "tests/test_transform.py::TestTransform::test_volume_zernike": 2.9709311130027345, + "tests/test_transform.py::test_NFP_warning": 12.250644918003673, + "tests/test_transform.py::test_transform_pytree": 9.803377583004476, + "tests/test_utils.py::test_isalmostequal": 3.218680423004116, + "tests/test_utils.py::test_islinspaced": 2.581942224998784, + "tests/test_vmec.py::TestVMECIO::test_fourier_to_zernike": 5.898660079994443, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identities_inverse": 3.000647447999654, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd": 2.842231549999269, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd_sin_series": 2.839034000004176, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev": 2.633703292998689, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev_sin_sym": 2.8674020860016753, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_linear_transform": 3.621225193997816, + "tests/test_vmec.py::TestVMECIO::test_zernike_to_fourier": 2.6461375140024757, + "tests/test_vmec.py::test_axis_surf_after_load": 41.88807087100213, + "tests/test_vmec.py::test_load_then_save": 70.19107655300104, + "tests/test_vmec.py::test_plot_vmec_comparison": 24.87666376099878, + "tests/test_vmec.py::test_vmec_boundary_subspace": 4.572544892005681, + "tests/test_vmec.py::test_vmec_load_profiles": 53.93892134500129, + "tests/test_vmec.py::test_vmec_save_1": 65.66046185197774, + "tests/test_vmec.py::test_vmec_save_2": 2.5563492279907223, + "tests/test_vmec.py::test_vmec_save_asym": 41.46463073199993, + "tests/test_vmec.py::test_vmec_save_kinetic": 49.06321475199729, + "tests/test_vmec.py::test_write_vmec_input": 74.78046189900488 } diff --git a/desc/plotting.py b/desc/plotting.py index 7c10a059a2..88c5b77241 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -11,8 +11,6 @@ from matplotlib import cycler, rcParams from mpl_toolkits.axes_grid1 import make_axes_locatable from pylatexenc.latex2text import LatexNodes2Text -from scipy.integrate import solve_ivp -from scipy.interpolate import Rbf from termcolor import colored from desc.backend import sign @@ -3531,481 +3529,3 @@ def plot_logo(save_path=None, **kwargs): fig.savefig(save_path, facecolor=fig.get_facecolor(), edgecolor="none") return fig, ax - - -def plot_field_lines_sfl( - eq, - rho, - seed_thetas=0, - phi_start=0, - phi_end=2 * np.pi, - dphi=1e-2, - ax=None, - return_data=False, - **kwargs, -): - r"""Plots field lines on specified flux surface. - - Traces field lines at specified initial vartheta (:math:`\\vartheta`) seed - locations, then plots them. - Field lines traced by first finding the corresponding straight field line (SFL) - coordinates :math:`(\\rho,\\vartheta,\\phi)` for each field line, then - converting those to the computational :math:`(\\rho,\\theta,\\phi)` coordinates, - then finally computing from those the toroidal :math:`(R,\\phi,Z)` coordinates of - each field line. - - The SFL angle coordinates are found with the SFL relation: - - :math:`\\vartheta = \\iota \\phi + \\vartheta_0` - - Parameters - ---------- - eq : Equilibrium - Object from which to plot. - rho : float - Flux surface to trace field lines at. - seed_thetas : float or array-like of floats - Poloidal positions at which to seed magnetic field lines. - If array-like, will plot multiple field lines. - phi_start: float - Toroidal angle to integrate field line from, in radians. Default is 0. - phi_end: float - Toroidal angle to integrate field line until, in radians. Default is 2*pi. - dphi: float - spacing in phi to sample field lines along, in radians. Default is 1e-2. - ax : matplotlib AxesSubplot, optional - Axis to plot on. - if True, return the data plotted as well as fig,ax - return_data : bool - if True, return the data plotted as well as fig,ax - **kwargs : dict, optional - Specify properties of the figure, axis, and plot appearance e.g.:: - - plot_X(figsize=(4,6)) - - Valid keyword arguments are: - - * ``figsize``: tuple of length 2, the size of the figure (to be passed to - matplotlib) - - Returns - ------- - fig : matplotlib.figure.Figure - Figure being plotted to. - ax : matplotlib.axes.Axes or ndarray of Axes - Axes being plotted to. - plot_data : dict - dictionary of the data plotted, only returned if ``return_data=True`` - - Examples - -------- - .. image:: ../../_static/images/plotting/DSHAPE_field_lines_plot.png - - .. code-block:: python - - from desc.plotting import plot_field_lines_sfl - import desc.examples - import numpy as np - eq = desc.examples.get("DSHAPE") - seed_thetas=np.linspace(0, 2 * np.pi, 3,endpoint=False) - fig, ax, _ = plot_field_lines_sfl( - eq, rho=1,seed_thetas=seed_thetas , phi_end=2 * np.pi - ) - - """ - # TODO: can this be removed now? - if rho == 0: - raise NotImplementedError( - "Currently does not support field line tracing of the magnetic axis, " - + "please input 0 < rho <= 1" - ) - - fig, ax = _format_ax(ax, is3d=True, figsize=kwargs.get("figsize", None)) - - # check how many field lines to plot - if seed_thetas is list: - n_lines = len(seed_thetas) - elif isinstance(seed_thetas, np.ndarray): - n_lines = seed_thetas.size - else: - n_lines = 1 - - phi0 = phi_start - N_pts = int((phi_end - phi0) / dphi) - - grid_single_rho = Grid( - nodes=np.array([[rho, 0, 0]]) - ) # grid to get the iota value at the specified rho surface - iota = eq.compute("iota", grid=grid_single_rho)["iota"][0] - - varthetas = [] - phi = np.linspace(phi0, phi_end, N_pts) - if n_lines > 1: - for i in range(n_lines): - varthetas.append( - seed_thetas[i] + iota * phi - ) # list of varthetas corresponding to the field line - else: - varthetas.append( - seed_thetas + iota * phi - ) # list of varthetas corresponding to the field line - theta_coords = ( - [] - ) # list of nodes in (rho,theta,phi) corresponding to each (rho,vartheta,phi) - print( - ( - "Calculating field line (rho,theta,zeta) coordinates corresponding " - + "to sfl coordinates", - ) - ) - for vartheta_list in varthetas: - rhos = rho * np.ones_like(vartheta_list) - sfl_coords = np.vstack((rhos, vartheta_list, phi)).T - theta_coords.append(eq.compute_theta_coords(sfl_coords)) - - # calculate R,phi,Z of nodes in grid - # only need to do this after finding the grid corresponding to - # desired rho, vartheta, phi - print( - "Calculating field line (R,phi,Z) coordinates corresponding to " - + "(rho,theta,zeta) coordinates" - ) - field_line_coords = {"R": [], "Z": [], "phi": [], "seed_thetas": seed_thetas} - for coords in theta_coords: - grid = Grid(nodes=coords) - toroidal_coords = eq.compute(["R", "Z"], grid=grid) - field_line_coords["R"].append(toroidal_coords["R"]) - field_line_coords["Z"].append(toroidal_coords["Z"]) - field_line_coords["phi"].append(phi) - - for i in range(n_lines): - xline = np.asarray(field_line_coords["R"][i]) * np.cos( - field_line_coords["phi"][i] - ) - yline = np.asarray(field_line_coords["R"][i]) * np.sin( - field_line_coords["phi"][i] - ) - - ax.plot(xline, yline, field_line_coords["Z"][i], linewidth=2) - - ax.set_xlabel(_AXIS_LABELS_XYZ[0]) - ax.set_ylabel(_AXIS_LABELS_XYZ[1]) - ax.set_zlabel(_AXIS_LABELS_XYZ[2]) - ax.set_title( - "%d Magnetic Field Lines Traced On $\\rho=%1.2f$ Surface" % (n_lines, rho) - ) - _set_tight_layout(fig) - - # need this stuff to make all the axes equal, ax.axis('equal') doesn't work for 3d - x_limits = ax.get_xlim3d() - y_limits = ax.get_ylim3d() - z_limits = ax.get_zlim3d() - - x_range = abs(x_limits[1] - x_limits[0]) - x_middle = np.mean(x_limits) - y_range = abs(y_limits[1] - y_limits[0]) - y_middle = np.mean(y_limits) - z_range = abs(z_limits[1] - z_limits[0]) - z_middle = np.mean(z_limits) - - # The plot bounding box is a sphere in the sense of the infinity - # norm, hence I call half the max range the plot radius. - plot_radius = 0.5 * max([x_range, y_range, z_range]) - - ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius]) - ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) - ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) - - plot_data = field_line_coords - - if return_data: - return fig, ax, plot_data - - return fig, ax - - -def plot_field_lines_real_space( - eq, - rho, - seed_thetas=0, - phi_end=2 * np.pi, - grid=None, - ax=None, - B_interp=None, - return_B_interp=False, - **kwargs, -): - r"""Traces and plots field lines on a flux surface at specified seed locations. - - Field lines integrated by first fitting the magnetic field with radial basis - functions (RBF) in R,Z,phi, then integrating the field line from phi=0 up to the - specified phi angle, by solving: - - :math:`\\frac{dR}{d\\phi} = \\frac{RB_R}{B_{\\phi}}` - - :math:`\\frac{dZ}{d\\phi} = \\frac{RB_Z}{B_{\\phi}}` - - :math:`B_R = \\mathbf{B} \\cdot \\hat{\\mathbf{R}}` - - :math:`B_Z = \\mathbf{B} \\cdot \\hat{\\mathbf{Z}}` - - :math:`B_{\\phi} = \\mathbf{B} \\cdot \\hat{\\mathbf{\\phi}}` - - Parameters - ---------- - eq : Equilibrium - object from which to plot - rho : float - flux surface to trace field lines at - seed_thetas : float or array-like of floats - theta positions at which to seed magnetic field lines, if array-like, will plot - multiple field lines - phi_end: float - phi to integrate field line until, in radians. Default is 2*pi - grid : Grid, optional - grid of rho, theta, zeta coordinates used to evaluate magnetic field at, which - is then interpolated with RBF - ax : matplotlib AxesSubplot, optional - axis to plot on - B_interp : dict of scipy.interpolate.rbf.Rbf or equivalent interpolators, optional - if not None, uses the passed-in interpolation objects instead of fitting the - magnetic field with Rbf's. Useful if have already ran plot_field_lines once and - want to change the seed thetas or how far to integrate in phi. Dict should have - the following keys: ['B_R'], ['B_Z'], and ['B_phi'], corresponding to the - interpolating object for each cylindrical component of the magnetic field. - return_B_interp: bool, default False - If true, in addition to returning the fig, axis and field line coordinates, - will also return the dictionary of interpolating radial basis functions - interpolating the magnetic field in (R,phi,Z) - - Returns - ------- - fig : matplotlib.figure.Figure - figure being plotted to - ax : matplotlib.axes.Axes or ndarray of Axes - axes being plotted to - field_line_coords : dict - dict containing the R,phi,Z coordinates of each field line traced. Dictionary - entries are lists corresponding to the field lines for each seed_theta given. - Also contains the scipy IVP solutions for info on how each line was integrated. - B_interp : dict, only returned if return_B_interp is True - dict of scipy.interpolate.rbf.Rbf or equivalent call signature interpolators, - which interpolate the cylindrical components of magnetic field in (R,phi,Z). - Dict has the following keys: ['B_R'], ['B_Z'], and ['B_phi'], corresponding to - the interpolating object for each cylindrical component of the magnetic field, - and the interpolators have call signature B(R,phi,Z) = interpolator(R,phi,Z) - - Notes - ----- - Use plot_field_lines_sfl if plotting from a solved equilibrium, as that is faster - and more accurate than real space interpolation - - """ - nfp = 1 - if grid is None: - grid_kwargs = {"M": 30, "N": 30, "L": 20, "NFP": nfp, "axis": False} - grid = _get_grid(**grid_kwargs) - - fig, ax = _format_ax(ax, is3d=True, figsize=kwargs.get("figsize", None)) - - # check how many field lines to plot - if seed_thetas is list: - n_lines = len(seed_thetas) - elif isinstance(seed_thetas, np.ndarray): - n_lines = seed_thetas.size - else: - n_lines = 1 - phi0 = kwargs.get("phi0", 0) - - # calculate toroidal coordinates - toroidal_coords = eq.compute("phi", grid=grid) - Rs = toroidal_coords["R"] - Zs = toroidal_coords["Z"] - phis = toroidal_coords["phi"] - - # calculate cylindrical B - magnetic_field = eq.compute("B", grid=grid) - BR = magnetic_field["B_R"] - BZ = magnetic_field["B_Z"] - Bphi = magnetic_field["B_phi"] - - if B_interp is None: # must fit RBfs to interpolate B field in R,phi,Z - print("Fitting magnetic field with radial basis functions in R,phi,Z") - BRi = Rbf(Rs, Zs, phis, BR) - BZi = Rbf(Rs, Zs, phis, BZ) - Bphii = Rbf(Rs, Zs, phis, Bphi) - B_interp = {"B_R": BRi, "B_Z": BZi, "B_phi": Bphii} - - field_line_coords = { - "Rs": [], - "Zs": [], - "phis": [], - "IVP solutions": [], - "seed_thetas": seed_thetas, - } - if n_lines > 1: - for theta in seed_thetas: - field_line_Rs, field_line_phis, field_line_Zs, sol = _field_line_Rbf( - rho, theta, phi_end, grid, Rs, Zs, B_interp, phi0 - ) - field_line_coords["Rs"].append(field_line_Rs) - field_line_coords["Zs"].append(field_line_Zs) - field_line_coords["phis"].append(field_line_phis) - field_line_coords["IVP solutions"].append(sol) - - else: - field_line_Rs, field_line_phis, field_line_Zs, sol = _field_line_Rbf( - rho, seed_thetas, phi_end, grid, Rs, Zs, B_interp, phi0 - ) - field_line_coords["Rs"].append(field_line_Rs) - field_line_coords["Zs"].append(field_line_Zs) - field_line_coords["phis"].append(field_line_phis) - field_line_coords["IVP solutions"].append(sol) - - for i, solution in enumerate(field_line_coords["IVP solutions"]): - if not solution.success: - print( - "Integration from seed theta %1.2f radians was not successful!" - % seed_thetas[i] - ) - - for i in range(n_lines): - xline = np.asarray(field_line_coords["Rs"][i]) * np.cos( - field_line_coords["phis"][i] - ) - yline = np.asarray(field_line_coords["Rs"][i]) * np.sin( - field_line_coords["phis"][i] - ) - - ax.plot(xline, yline, field_line_coords["Zs"][i], linewidth=2) - - ax.set_xlabel(_AXIS_LABELS_XYZ[0]) - ax.set_ylabel(_AXIS_LABELS_XYZ[1]) - ax.set_zlabel(_AXIS_LABELS_XYZ[2]) - ax.set_title( - "%d Magnetic Field Lines Traced On $\\rho=%1.2f$ Surface" % (n_lines, rho) - ) - _set_tight_layout(fig) - - # need this stuff to make all the axes equal, ax.axis('equal') doesn't work for 3d - x_limits = ax.get_xlim3d() - y_limits = ax.get_ylim3d() - z_limits = ax.get_zlim3d() - - x_range = abs(x_limits[1] - x_limits[0]) - x_middle = np.mean(x_limits) - y_range = abs(y_limits[1] - y_limits[0]) - y_middle = np.mean(y_limits) - z_range = abs(z_limits[1] - z_limits[0]) - z_middle = np.mean(z_limits) - - # The plot bounding box is a sphere in the sense of the infinity - # norm, hence I call half the max range the plot radius. - plot_radius = 0.5 * max([x_range, y_range, z_range]) - - ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius]) - ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) - ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) - - if return_B_interp: - return ( - fig, - ax, - field_line_coords, - B_interp, - ) - else: - return fig, ax - - -def _find_idx(rho0, theta0, phi0, grid): - """Finds the index of the node closest to the given rho0, theta0, phi0. - - Parameters - ---------- - rho0 : float - rho to find closest grid point to. - theta0 : float - theta to find closest grid point to. - phi0 : float - phi to find closest grid point to. - grid : Grid - grid to find closest point on - - Returns - ------- - idx_pt : int - index of the grid node closest to the given point. - - """ - rhos = grid.nodes[:, 0] - thetas = grid.nodes[:, 1] - phis = grid.nodes[:, 2] - - if theta0 < 0: - theta0 = 2 * np.pi + theta0 - if theta0 > 2 * np.pi: - theta0 = np.mod(theta0, 2 * np.pi) - if phi0 < 0: - phi0 = 2 * np.pi + phi0 - if phi0 > 2 * np.pi: - phi0 = np.mod(phi0, 2 * np.pi) - - bool1 = np.logical_and( - np.abs(rhos - rho0) == np.min(np.abs(rhos - rho0)), - np.abs(thetas - theta0) == np.min(np.abs(thetas - theta0)), - ) - bool2 = np.logical_and(bool1, np.abs(phis - phi0) == np.min(np.abs(phis - phi0))) - idx_pt = np.where(bool2)[0][0] - return idx_pt - - -def _field_line_Rbf(rho, theta0, phi_end, grid, Rs, Zs, B_interp, phi0=0): - """Integrate along interpolated field lines. - - Takes the initial poloidal angle you want to seed a field line at (at phi=0), - and integrates along the field line to the specified phi_end. returns fR,fZ,fPhi, - the R,Z,Phi coordinates of the field line trajectory. - - """ - fR = [] - fZ = [] - fPhi = [] - idx0 = _find_idx(rho, theta0, phi0, grid) - curr_R = Rs[idx0] - curr_Z = Zs[idx0] - fR.append(curr_R) - fZ.append(curr_Z) - fPhi.append(phi0) - - # integrate field lines in Phi - print( - "Integrating Magnetic Field Line Equation from seed theta = %f radians" % theta0 - ) - y0 = [fR[0], fZ[0]] - - def rhs(phi, y): - """RHS of magnetic field line equation.""" - dRdphi = ( - y[0] - * B_interp["B_R"](y[0], y[1], np.mod(phi, 2 * np.pi)) - / B_interp["B_phi"](y[0], y[1], np.mod(phi, 2 * np.pi)) - ) - dZdphi = ( - y[0] - * B_interp["B_Z"](y[0], y[1], np.mod(phi, 2 * np.pi)) - / B_interp["B_phi"](y[0], y[1], np.mod(phi, 2 * np.pi)) - ) - return [dRdphi, dZdphi] - - n_tries = 1 - max_step = 0.01 - sol = solve_ivp(rhs, [0, phi_end], y0, max_step=max_step) - while not sol.success and n_tries < 4: - max_step = 0.5 * max_step - n_tries += 1 - sol = solve_ivp(rhs, [0, phi_end], y0, max_step=max_step) - fR = sol.y[0, :] - fZ = sol.y[1, :] - fPhi = sol.t - return fR, fPhi, fZ, sol diff --git a/docs/api.rst b/docs/api.rst index 08a8ede7a1..8d730601b0 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -258,7 +258,6 @@ Plotting desc.plotting.plot_coefficients desc.plotting.plot_coils desc.plotting.plot_comparison - desc.plotting.plot_field_lines_sfl desc.plotting.plot_fsa desc.plotting.plot_grid desc.plotting.plot_logo diff --git a/tests/baseline/test_1d_dpdr.png b/tests/baseline/test_1d_dpdr.png deleted file mode 100644 index d949bc1d93..0000000000 Binary files a/tests/baseline/test_1d_dpdr.png and /dev/null differ diff --git a/tests/baseline/test_1d_elongation.png b/tests/baseline/test_1d_elongation.png index 4670a13118..ccf04b533c 100644 Binary files a/tests/baseline/test_1d_elongation.png and b/tests/baseline/test_1d_elongation.png differ diff --git a/tests/baseline/test_1d_iota.png b/tests/baseline/test_1d_iota.png index 83cf2ec3cf..7d40e1c4e1 100644 Binary files a/tests/baseline/test_1d_iota.png and b/tests/baseline/test_1d_iota.png differ diff --git a/tests/baseline/test_1d_iota_radial.png b/tests/baseline/test_1d_iota_radial.png index df42538d4d..c910e0cba1 100644 Binary files a/tests/baseline/test_1d_iota_radial.png and b/tests/baseline/test_1d_iota_radial.png differ diff --git a/tests/baseline/test_1d_logpsi.png b/tests/baseline/test_1d_logpsi.png index 6fbb627e9d..562f4895b6 100644 Binary files a/tests/baseline/test_1d_logpsi.png and b/tests/baseline/test_1d_logpsi.png differ diff --git a/tests/baseline/test_1d_p.png b/tests/baseline/test_1d_p.png index cfc2704eb1..36db721bfc 100644 Binary files a/tests/baseline/test_1d_p.png and b/tests/baseline/test_1d_p.png differ diff --git a/tests/baseline/test_2d_g_rz.png b/tests/baseline/test_2d_g_rz.png index 397d9533f6..a8c3609ef6 100644 Binary files a/tests/baseline/test_2d_g_rz.png and b/tests/baseline/test_2d_g_rz.png differ diff --git a/tests/baseline/test_2d_g_tz.png b/tests/baseline/test_2d_g_tz.png index dce5e4d4fa..e58a3fb603 100644 Binary files a/tests/baseline/test_2d_g_tz.png and b/tests/baseline/test_2d_g_tz.png differ diff --git a/tests/baseline/test_2d_lambda.png b/tests/baseline/test_2d_lambda.png deleted file mode 100644 index b0e71d807e..0000000000 Binary files a/tests/baseline/test_2d_lambda.png and /dev/null differ diff --git a/tests/baseline/test_2d_logF.png b/tests/baseline/test_2d_logF.png index 7a757124c6..12bf243827 100644 Binary files a/tests/baseline/test_2d_logF.png and b/tests/baseline/test_2d_logF.png differ diff --git a/tests/baseline/test_2d_plot_Bn.png b/tests/baseline/test_2d_plot_Bn.png index 971207f19b..7e3ad74c80 100644 Binary files a/tests/baseline/test_2d_plot_Bn.png and b/tests/baseline/test_2d_plot_Bn.png differ diff --git a/tests/baseline/test_fsa_F_normalized.png b/tests/baseline/test_fsa_F_normalized.png index b8d1223945..170263d9e1 100644 Binary files a/tests/baseline/test_fsa_F_normalized.png and b/tests/baseline/test_fsa_F_normalized.png differ diff --git a/tests/baseline/test_fsa_G.png b/tests/baseline/test_fsa_G.png deleted file mode 100644 index 99876534b1..0000000000 Binary files a/tests/baseline/test_fsa_G.png and /dev/null differ diff --git a/tests/baseline/test_plot_boozer_modes.png b/tests/baseline/test_plot_boozer_modes.png index 8f00d3e75d..b72e308394 100644 Binary files a/tests/baseline/test_plot_boozer_modes.png and b/tests/baseline/test_plot_boozer_modes.png differ diff --git a/tests/baseline/test_plot_boozer_modes_breaking_only.png b/tests/baseline/test_plot_boozer_modes_breaking_only.png index 5953b525b7..424ba1f669 100644 Binary files a/tests/baseline/test_plot_boozer_modes_breaking_only.png and b/tests/baseline/test_plot_boozer_modes_breaking_only.png differ diff --git a/tests/baseline/test_plot_boozer_modes_max.png b/tests/baseline/test_plot_boozer_modes_max.png index 4be79a0afa..ac59a48ed3 100644 Binary files a/tests/baseline/test_plot_boozer_modes_max.png and b/tests/baseline/test_plot_boozer_modes_max.png differ diff --git a/tests/baseline/test_plot_boozer_modes_no_norm.png b/tests/baseline/test_plot_boozer_modes_no_norm.png index 69ed46ffa0..b2fe1e29b5 100644 Binary files a/tests/baseline/test_plot_boozer_modes_no_norm.png and b/tests/baseline/test_plot_boozer_modes_no_norm.png differ diff --git a/tests/baseline/test_plot_boundary.png b/tests/baseline/test_plot_boundary.png index 6db322b6bf..8f2dd27692 100644 Binary files a/tests/baseline/test_plot_boundary.png and b/tests/baseline/test_plot_boundary.png differ diff --git a/tests/baseline/test_plot_coefficients.png b/tests/baseline/test_plot_coefficients.png index 43f37c1fec..7852b9fbfd 100644 Binary files a/tests/baseline/test_plot_coefficients.png and b/tests/baseline/test_plot_coefficients.png differ diff --git a/tests/baseline/test_plot_comparison.png b/tests/baseline/test_plot_comparison.png index 99089581ba..9f7a7f85d5 100644 Binary files a/tests/baseline/test_plot_comparison.png and b/tests/baseline/test_plot_comparison.png differ diff --git a/tests/baseline/test_plot_comparison_no_theta.png b/tests/baseline/test_plot_comparison_no_theta.png index 8320e44410..2d1cd7a7ae 100644 Binary files a/tests/baseline/test_plot_comparison_no_theta.png and b/tests/baseline/test_plot_comparison_no_theta.png differ diff --git a/tests/baseline/test_plot_cov_basis.png b/tests/baseline/test_plot_cov_basis.png index cd1865690f..68ed3a5906 100644 Binary files a/tests/baseline/test_plot_cov_basis.png and b/tests/baseline/test_plot_cov_basis.png differ diff --git a/tests/baseline/test_plot_field_line.png b/tests/baseline/test_plot_field_line.png deleted file mode 100644 index ed563dd490..0000000000 Binary files a/tests/baseline/test_plot_field_line.png and /dev/null differ diff --git a/tests/baseline/test_plot_field_lines.png b/tests/baseline/test_plot_field_lines.png deleted file mode 100644 index 5ee63467a0..0000000000 Binary files a/tests/baseline/test_plot_field_lines.png and /dev/null differ diff --git a/tests/baseline/test_plot_gradpsi.png b/tests/baseline/test_plot_gradpsi.png deleted file mode 100644 index 00d4592727..0000000000 Binary files a/tests/baseline/test_plot_gradpsi.png and /dev/null differ diff --git a/tests/baseline/test_plot_magnetic_pressure.png b/tests/baseline/test_plot_magnetic_pressure.png deleted file mode 100644 index 8655097a28..0000000000 Binary files a/tests/baseline/test_plot_magnetic_pressure.png and /dev/null differ diff --git a/tests/baseline/test_plot_magnetic_tension.png b/tests/baseline/test_plot_magnetic_tension.png deleted file mode 100644 index 34cdb08e8e..0000000000 Binary files a/tests/baseline/test_plot_magnetic_tension.png and /dev/null differ diff --git a/tests/baseline/test_plot_normF_2d.png b/tests/baseline/test_plot_normF_2d.png index 3b0899c477..9552700b53 100644 Binary files a/tests/baseline/test_plot_normF_2d.png and b/tests/baseline/test_plot_normF_2d.png differ diff --git a/tests/baseline/test_plot_normF_section.png b/tests/baseline/test_plot_normF_section.png index 3cf477acb5..ead4cf980f 100644 Binary files a/tests/baseline/test_plot_normF_section.png and b/tests/baseline/test_plot_normF_section.png differ diff --git a/tests/baseline/test_plot_omnigenous_field.png b/tests/baseline/test_plot_omnigenous_field.png index 03b3c5bc52..a4d36172c4 100644 Binary files a/tests/baseline/test_plot_omnigenous_field.png and b/tests/baseline/test_plot_omnigenous_field.png differ diff --git a/tests/baseline/test_plot_qs_error.png b/tests/baseline/test_plot_qs_error.png index 53b495dfa2..c90e59fb80 100644 Binary files a/tests/baseline/test_plot_qs_error.png and b/tests/baseline/test_plot_qs_error.png differ diff --git a/tests/baseline/test_plot_surfaces.png b/tests/baseline/test_plot_surfaces.png index e13108e153..bdae244b07 100644 Binary files a/tests/baseline/test_plot_surfaces.png and b/tests/baseline/test_plot_surfaces.png differ diff --git a/tests/baseline/test_plot_surfaces_HELIOTRON.png b/tests/baseline/test_plot_surfaces_HELIOTRON.png index 7ebfc750fa..1b2c8c9dbd 100644 Binary files a/tests/baseline/test_plot_surfaces_HELIOTRON.png and b/tests/baseline/test_plot_surfaces_HELIOTRON.png differ diff --git a/tests/baseline/test_plot_surfaces_no_theta.png b/tests/baseline/test_plot_surfaces_no_theta.png index 65a713fe4f..f2f3e9c6bd 100644 Binary files a/tests/baseline/test_plot_surfaces_no_theta.png and b/tests/baseline/test_plot_surfaces_no_theta.png differ diff --git a/tests/baseline/test_section_F.png b/tests/baseline/test_section_F.png index 1ce56fb8b2..573206b0b4 100644 Binary files a/tests/baseline/test_section_F.png and b/tests/baseline/test_section_F.png differ diff --git a/tests/baseline/test_section_F_normalized_vac.png b/tests/baseline/test_section_F_normalized_vac.png deleted file mode 100644 index 352761a194..0000000000 Binary files a/tests/baseline/test_section_F_normalized_vac.png and /dev/null differ diff --git a/tests/baseline/test_section_J.png b/tests/baseline/test_section_J.png index 8d687d3ebc..afc8cfce37 100644 Binary files a/tests/baseline/test_section_J.png and b/tests/baseline/test_section_J.png differ diff --git a/tests/baseline/test_section_R.png b/tests/baseline/test_section_R.png deleted file mode 100644 index 6478780739..0000000000 Binary files a/tests/baseline/test_section_R.png and /dev/null differ diff --git a/tests/baseline/test_section_Z.png b/tests/baseline/test_section_Z.png deleted file mode 100644 index 5ef0f93b98..0000000000 Binary files a/tests/baseline/test_section_Z.png and /dev/null differ diff --git a/tests/baseline/test_section_logF.png b/tests/baseline/test_section_logF.png index 05821dd748..02456ea0a9 100644 Binary files a/tests/baseline/test_section_logF.png and b/tests/baseline/test_section_logF.png differ diff --git a/tests/conftest.py b/tests/conftest.py index e000574706..cf2b1d020c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ import os import h5py +import jax import numpy as np import pytest from netCDF4 import Dataset @@ -12,6 +13,12 @@ from desc.vmec import VMECIO +@pytest.fixture(scope="class", autouse=True) +def clear_caches_before(): + """Automatically run before each test to clear caches and reduce OOM issues.""" + jax.clear_caches() + + @pytest.fixture(scope="session") def TmpDir(tmpdir_factory): """Create a temporary directory to store testing files.""" @@ -19,38 +26,6 @@ def TmpDir(tmpdir_factory): return dir_path -@pytest.fixture(scope="session") -def SOLOVEV_vac(tmpdir_factory): - """Run SOLOVEV vacuum example.""" - input_path = ".//tests//inputs//SOLOVEV_vac" - output_dir = tmpdir_factory.mktemp("result") - desc_h5_path = output_dir.join("SOLOVEV_vac_out.h5") - desc_nc_path = output_dir.join("SOLOVEV_vac_out.nc") - vmec_nc_path = ".//tests//inputs//wout_SOLOVEV_vac.nc" - booz_nc_path = output_dir.join("SOLOVEV_vac_bx.nc") - - cwd = os.path.dirname(__file__) - exec_dir = os.path.join(cwd, "..") - input_filename = os.path.join(exec_dir, input_path) - - print("Running SOLOVEV vacuum test.") - print("exec_dir=", exec_dir) - print("cwd=", cwd) - - args = ["-o", str(desc_h5_path), input_filename, "--numpy", "-vv"] - with pytest.warns(UserWarning, match="Left handed coordinates"): - main(args) - - SOLOVEV_vac_out = { - "input_path": input_path, - "desc_h5_path": desc_h5_path, - "desc_nc_path": desc_nc_path, - "vmec_nc_path": vmec_nc_path, - "booz_nc_path": booz_nc_path, - } - return SOLOVEV_vac_out - - @pytest.fixture(scope="session") def SOLOVEV(tmpdir_factory): """Run SOLOVEV example.""" @@ -175,24 +150,6 @@ def HELIOTRON(tmpdir_factory): return HELIOTRON_out -@pytest.fixture(scope="session") -def HELIOTRON_ex(tmpdir_factory): - """Saved HELIOTRON fixed rotational transform example.""" - input_path = ".//tests//inputs//HELIOTRON" - output_dir = tmpdir_factory.mktemp("result") - desc_h5_path = ".//desc//examples//HELIOTRON_output.h5" - vmec_nc_path = ".//tests//inputs//wout_HELIOTRON.nc" - booz_nc_path = output_dir.join("HELIOTRON_bx.nc") - - HELIOTRON_out = { - "input_path": input_path, - "desc_h5_path": desc_h5_path, - "vmec_nc_path": vmec_nc_path, - "booz_nc_path": booz_nc_path, - } - return HELIOTRON_out - - @pytest.fixture(scope="session") def HELIOTRON_vac(tmpdir_factory): """Run HELIOTRON vacuum (vacuum) example.""" @@ -225,37 +182,6 @@ def HELIOTRON_vac(tmpdir_factory): return HELIOTRON_vacuum_out -@pytest.fixture(scope="session") -def HELIOTRON_vac2(tmpdir_factory): - """Run HELIOTRON vacuum (fixed current) example.""" - input_path = ".//tests//inputs//HELIOTRON_vacuum2" - output_dir = tmpdir_factory.mktemp("result") - desc_h5_path = output_dir.join("HELIOTRON_vacuum2_out.h5") - desc_nc_path = output_dir.join("HELIOTRON_vacuum2_out.nc") - vmec_nc_path = ".//tests//inputs//wout_HELIOTRON_vacuum2.nc" - booz_nc_path = output_dir.join("HELIOTRON_vacuum2_bx.nc") - - cwd = os.path.dirname(__file__) - exec_dir = os.path.join(cwd, "..") - input_filename = os.path.join(exec_dir, input_path) - - print("Running HELIOTRON vacuum (fixed current) test.") - print("exec_dir=", exec_dir) - print("cwd=", cwd) - - args = ["-o", str(desc_h5_path), input_filename, "-vv"] - main(args) - - HELIOTRON_vacuum2_out = { - "input_path": input_path, - "desc_h5_path": desc_h5_path, - "desc_nc_path": desc_nc_path, - "vmec_nc_path": vmec_nc_path, - "booz_nc_path": booz_nc_path, - } - return HELIOTRON_vacuum2_out - - @pytest.fixture(scope="session") def DummyStellarator(tmpdir_factory): """Create and save a dummy stellarator configuration for testing.""" diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index 77f81f31be..1316192c8c 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -309,9 +309,13 @@ def test_limit_continuity(self): # same as 'weaker_tolerance | zero_limit', but works on Python 3.8 (PEP 584) kwargs = dict(weaker_tolerance, **zero_map) # fixed iota - assert_is_continuous(get("W7-X"), kwargs=kwargs) + eq = get("W7-X") + eq.change_resolution(4, 4, 4, 8, 8, 8) + assert_is_continuous(eq, kwargs=kwargs) # fixed current - assert_is_continuous(get("NCSX"), kwargs=kwargs) + eq = get("NCSX") + eq.change_resolution(4, 4, 4, 8, 8, 8) + assert_is_continuous(eq, kwargs=kwargs) @pytest.mark.unit def test_magnetic_field_is_physical(self): @@ -344,8 +348,12 @@ def test(eq): np.testing.assert_allclose(B[:, 1], B[0, 1]) np.testing.assert_allclose(B[:, 2], B[0, 2]) - test(get("W7-X")) - test(get("NCSX")) + eq = get("W7-X") + eq.change_resolution(4, 4, 4, 8, 8, 8) + test(eq) + eq = get("NCSX") + eq.change_resolution(4, 4, 4, 8, 8, 8) + test(eq) def _reverse_mode_unsafe_names(): @@ -385,10 +393,10 @@ def get_source(name): unsafe_names.append(name) unsafe_names = sorted(unsafe_names) - print("Unsafe names: ", unsafe_names) return unsafe_names +@pytest.mark.unit @pytest.mark.parametrize("name", _reverse_mode_unsafe_names()) def test_reverse_mode_ad_axis(name): """Asserts that the rho=0 axis limits are reverse mode differentiable.""" @@ -396,7 +404,7 @@ def test_reverse_mode_ad_axis(name): grid = LinearGrid(rho=0.0, M=2, N=2, NFP=eq.NFP, sym=eq.sym) eq.change_resolution(2, 2, 2, 4, 4, 4) - obj = ObjectiveFunction(GenericObjective(name, eq, grid=grid)) + obj = ObjectiveFunction(GenericObjective(name, eq, grid=grid), use_jit=False) obj.build(verbose=0) g = obj.grad(obj.x()) assert not np.any(np.isnan(g)) diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index f4b7b1c7e8..34ffe08e22 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -1,5 +1,7 @@ """Tests for bootstrap current functions.""" +import warnings + import matplotlib.pyplot as plt import numpy as np import pytest @@ -1196,54 +1198,6 @@ def test_BootstrapRedlConsistency_normalization(self): # Results are not perfectly identical because ln(Lambda) is not quite invariant. np.testing.assert_allclose(results, expected, rtol=2e-3) - @pytest.mark.unit - @pytest.mark.solve - def test_BootstrapRedlConsistency_resolution(self, DSHAPE_current): - """Confirm that the objective function is ~independent of grid resolution.""" - helicity = (1, 0) - - eq = desc.io.load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - - eq.electron_density = PowerSeriesProfile( - 2.0e19 * np.array([1, -0.85]), modes=[0, 4] - ) - eq.electron_temperature = PowerSeriesProfile( - 1e3 * np.array([1.02, -3, 3, -1]), modes=[0, 2, 4, 6] - ) - eq.ion_temperature = PowerSeriesProfile( - 1.1e3 * np.array([1.02, -3, 3, -1]), modes=[0, 2, 4, 6] - ) - eq.atomic_number = 1.4 - - def test(grid_type, kwargs, L, M, N): - grid = grid_type(L=L, M=M, N=N, NFP=eq.NFP, **kwargs) - obj = ObjectiveFunction( - BootstrapRedlConsistency(eq=eq, grid=grid, helicity=helicity) - ) - obj.build() - scalar_objective = obj.compute_scalar(obj.x(eq)) - print(f"grid_type:{grid_type} L:{L} M:{M} N:{N} obj:{scalar_objective}") - return scalar_objective - - results = [] - - # Loop over grid types. For LinearGrid we need to drop the - # point at rho=0 to avoid a divide-by-0 - grid_types = [LinearGrid, QuadratureGrid] - kwargss = [{"axis": False}, {}] - - # Loop over grid resolutions: - Ls = [150, 300, 150, 150] - Ms = [10, 10, 20, 10] - Ns = [0, 0, 0, 2] - - for grid_type, kwargs in zip(grid_types, kwargss): - for L, M, N in zip(Ls, Ms, Ns): - results.append(test(grid_type, kwargs, L, M, N)) - - results = np.array(results) - np.testing.assert_allclose(results, np.mean(results), rtol=0.04) - @pytest.mark.regression def test_bootstrap_consistency_iota(self, TmpDir): """Try optimizing for bootstrap consistency in axisymmetry, at fixed shape. @@ -1492,14 +1446,18 @@ def test_bootstrap_objective_build(): L=3, M=3, N=3, NFP=2, electron_temperature=1e3, electron_density=1e25 ) # density too high - with pytest.warns(UserWarning): - BootstrapRedlConsistency(eq=eq).build() + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + BootstrapRedlConsistency(eq=eq).build() eq = Equilibrium( L=3, M=3, N=3, NFP=2, electron_temperature=1e5, electron_density=1e21 ) # electron temperature too high - with pytest.warns(UserWarning): - BootstrapRedlConsistency(eq=eq).build() + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + BootstrapRedlConsistency(eq=eq).build() eq = Equilibrium( L=3, M=3, @@ -1510,8 +1468,10 @@ def test_bootstrap_objective_build(): ion_temperature=1e5, ) # ion temperature too high - with pytest.warns(UserWarning): - BootstrapRedlConsistency(eq=eq).build() + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + BootstrapRedlConsistency(eq=eq).build() eq = Equilibrium( L=3, M=3, @@ -1522,8 +1482,10 @@ def test_bootstrap_objective_build(): ion_temperature=1e3, ) # density too low - with pytest.warns(UserWarning): - BootstrapRedlConsistency(eq=eq).build() + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + BootstrapRedlConsistency(eq=eq).build() eq = Equilibrium( L=3, M=3, @@ -1534,8 +1496,10 @@ def test_bootstrap_objective_build(): ion_temperature=1e3, ) # electron temperature too low - with pytest.warns(UserWarning): - BootstrapRedlConsistency(eq=eq).build() + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + BootstrapRedlConsistency(eq=eq).build() eq = Equilibrium( L=3, M=3, @@ -1546,8 +1510,10 @@ def test_bootstrap_objective_build(): ion_temperature=1, ) # ion temperature too low - with pytest.warns(UserWarning): - BootstrapRedlConsistency(eq=eq).build() + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + BootstrapRedlConsistency(eq=eq).build() eq = Equilibrium( L=4, diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 91cbcf66e8..4ff9e7482e 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -4,12 +4,11 @@ import numpy as np import pytest -from scipy.io import netcdf_file from scipy.signal import convolve2d from desc.coils import FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, SplineXYZCoil from desc.compute import data_index, rpz2xyz_vec -from desc.equilibrium import EquilibriaFamily, Equilibrium +from desc.equilibrium import Equilibrium from desc.examples import get from desc.geometry import ( FourierPlanarCurve, @@ -18,7 +17,7 @@ FourierXYZCurve, ZernikeRZToroidalSection, ) -from desc.grid import LinearGrid, QuadratureGrid +from desc.grid import LinearGrid from desc.io import load from desc.magnetic_fields import ( CurrentPotentialField, @@ -1085,25 +1084,6 @@ def test_magnetic_pressure_gradient(DummyStellarator): ) -@pytest.mark.unit -@pytest.mark.solve -def test_currents(DSHAPE_current): - """Test that different methods for computing I and G agree.""" - eq = EquilibriaFamily.load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - - grid_full = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) - grid_sym = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=True) - - data_booz = eq.compute("|B|_mn", grid=grid_full, M_booz=eq.M, N_booz=eq.N) - data_full = eq.compute(["I", "G"], grid=grid_full) - data_sym = eq.compute(["I", "G"], grid=grid_sym) - - np.testing.assert_allclose(data_full["I"].mean(), data_booz["I"], atol=1e-16) - np.testing.assert_allclose(data_sym["I"].mean(), data_booz["I"], atol=1e-16) - np.testing.assert_allclose(data_full["G"].mean(), data_booz["G"], atol=1e-16) - np.testing.assert_allclose(data_sym["G"].mean(), data_booz["G"], atol=1e-16) - - @pytest.mark.slow @pytest.mark.unit def test_BdotgradB(DummyStellarator): @@ -1134,10 +1114,10 @@ def test_partial_derivative(name): @pytest.mark.unit @pytest.mark.solve -def test_boozer_transform(DSHAPE_current): +def test_boozer_transform(): """Test that Boozer coordinate transform agrees with BOOZ_XFORM.""" # TODO: add test with stellarator example - eq = EquilibriaFamily.load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] + eq = get("DSHAPE_CURRENT") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) data = eq.compute("|B|_mn", grid=grid, M_booz=eq.M, N_booz=eq.N) booz_xform = np.array( @@ -1166,47 +1146,6 @@ def test_boozer_transform(DSHAPE_current): ) -@pytest.mark.unit -def test_compute_grad_p_volume_avg(): - """Test calculation of volume averaged pressure gradient.""" - eq = Equilibrium() # default pressure profile is 0 pressure - pres_grad_vol_avg = eq.compute("<|grad(p)|>_vol")["<|grad(p)|>_vol"] - np.testing.assert_allclose(pres_grad_vol_avg, 0) - - -@pytest.mark.unit -def test_compare_quantities_to_vmec(): - """Compare several computed quantities to vmec.""" - wout_file = ".//tests//inputs//wout_DSHAPE.nc" - desc_file = ".//tests//inputs//DSHAPE_output_saved_without_current.h5" - - fid = netcdf_file(wout_file, mmap=False) - ns = fid.variables["ns"][()] - J_dot_B_vmec = fid.variables["jdotb"][()] - volavgB = fid.variables["volavgB"][()] - betatotal = fid.variables["betatotal"][()] - fid.close() - - with pytest.warns(RuntimeWarning, match="Save attribute '_current'"): - eq = EquilibriaFamily.load(desc_file)[-1] - - # Compare 0D quantities: - grid = QuadratureGrid(eq.L, M=eq.M, N=eq.N, NFP=eq.NFP) - data = eq.compute("_vol", grid=grid) - data = eq.compute("<|B|>_rms", grid=grid, data=data) - - np.testing.assert_allclose(volavgB, data["<|B|>_rms"], rtol=1e-7) - np.testing.assert_allclose(betatotal, data["_vol"], rtol=1e-5) - - # Compare radial profile quantities: - s = np.linspace(0, 1, ns) - rho = np.sqrt(s) - grid = LinearGrid(rho=rho, M=eq.M, N=eq.N, NFP=eq.NFP) - data = eq.compute("", grid=grid) - J_dot_B_desc = grid.compress(data[""]) - np.testing.assert_allclose(J_dot_B_desc, J_dot_B_vmec, rtol=0.005) - - @pytest.mark.unit def test_compute_everything(): """Test that the computations on this branch agree with those on master. @@ -1657,8 +1596,7 @@ def test_contravariant_basis_vectors(): @pytest.mark.unit -@pytest.mark.solve -def test_iota_components(HELIOTRON_vac): +def test_iota_components(): """Test that iota components are computed correctly.""" # axisymmetric, so all rotational transform should be from the current eq_i = get("DSHAPE") # iota profile assigned @@ -1672,7 +1610,7 @@ def test_iota_components(HELIOTRON_vac): np.testing.assert_allclose(data_c["iota vacuum"], 0) # vacuum stellarator, so all rotational transform should be from the external field - eq = load(load_from=str(HELIOTRON_vac["desc_h5_path"]), file_format="hdf5")[-1] + eq = get("ESTELL") grid = LinearGrid(L=100, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, axis=True) data = eq.compute(["iota", "iota current", "iota vacuum"], grid) np.testing.assert_allclose(data["iota"], data["iota vacuum"]) diff --git a/tests/test_compute_utils.py b/tests/test_compute_utils.py index b8a69c20fa..3136828620 100644 --- a/tests/test_compute_utils.py +++ b/tests/test_compute_utils.py @@ -108,6 +108,7 @@ def test_b_theta(surface_label, grid, eq): ) eq = get("W7-X") + eq.change_resolution(3, 3, 3, 6, 6, 6) lg = LinearGrid(L=L, M=M, N=N, NFP=eq.NFP, endpoint=False) lg_endpoint = LinearGrid(L=L, M=M, N=N, NFP=eq.NFP, endpoint=True) cg_sym = ConcentricGrid(L=L, M=M, N=N, NFP=eq.NFP, sym=True) @@ -321,6 +322,7 @@ def test(grid): def test_surface_averages_identity_op(self): """Test flux surface averages of surface functions are identity operations.""" eq = get("W7-X") + eq.change_resolution(3, 3, 3, 6, 6, 6) grid = ConcentricGrid(L=L, M=M, N=N, NFP=eq.NFP, sym=eq.sym) data = eq.compute(["p", "sqrt(g)"], grid=grid) pressure_average = surface_averages(grid, data["p"], data["sqrt(g)"]) @@ -333,6 +335,7 @@ def test_surface_averages_homomorphism(self): Meaning average(a + b) = average(a) + average(b). """ eq = get("W7-X") + eq.change_resolution(3, 3, 3, 6, 6, 6) grid = ConcentricGrid(L=L, M=M, N=N, NFP=eq.NFP, sym=eq.sym) data = eq.compute(["|B|", "|B|_t", "sqrt(g)"], grid=grid) a = surface_averages(grid, data["|B|"], data["sqrt(g)"]) @@ -388,6 +391,7 @@ def test_surface_averages_against_shortcut(self): # test on grids with a single rho surface eq = get("W7-X") + eq.change_resolution(3, 3, 3, 6, 6, 6) rho = np.array((1 - 1e-4) * np.random.default_rng().random() + 1e-4) grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym) data = eq.compute(["|B|", "sqrt(g)"], grid=grid) diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 380aeb68bb..fed3c02356 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -5,6 +5,7 @@ import numpy as np import pytest +import desc.examples from desc.backend import put from desc.equilibrium import EquilibriaFamily, Equilibrium from desc.equilibrium.initial_guess import _initial_guess_surface @@ -410,11 +411,10 @@ def test_NFP_error(self): _ = Equilibrium(surface=surface3, axis=axis, NFP=2) @pytest.mark.unit - @pytest.mark.solve - def test_guess_from_file(self, DSHAPE_current): + def test_guess_from_file(self): """Test setting initial guess from saved equilibrium file.""" - path = DSHAPE_current["desc_h5_path"] - eq1 = Equilibrium(M=13, sym=True, spectral_indexing="fringe") + path = "tests//inputs//iotest_HELIOTRON.h5" + eq1 = Equilibrium(L=9, M=14, N=3, sym=True, spectral_indexing="ansi") eq1.set_initial_guess(path) eq2 = EquilibriaFamily.load(path)[-1] @@ -491,10 +491,9 @@ def test_asserts(self): @pytest.mark.unit -@pytest.mark.solve -def test_magnetic_axis(HELIOTRON_vac): +def test_magnetic_axis(): """Test that Configuration.axis returns the true axis location.""" - eq = EquilibriaFamily.load(load_from=str(HELIOTRON_vac["desc_h5_path"]))[-1] + eq = desc.examples.get("HELIOTRON") axis = eq.axis grid = LinearGrid(N=3 * eq.N_grid, NFP=eq.NFP, rho=np.array(0.0)) @@ -546,10 +545,9 @@ def test_is_nested_theta(): @pytest.mark.unit -@pytest.mark.solve -def test_get_profile(DSHAPE_current): +def test_get_profile(): """Test getting/setting iota and current profiles.""" - eq = EquilibriaFamily.load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] + eq = desc.examples.get("DSHAPE_CURRENT") current0 = eq.current current1 = eq.get_profile("current", kind="power_series") current2 = eq.get_profile("current", kind="spline") diff --git a/tests/test_constrain_current.py b/tests/test_constrain_current.py index c5b367b69f..0678f28a2f 100644 --- a/tests/test_constrain_current.py +++ b/tests/test_constrain_current.py @@ -43,8 +43,7 @@ class TestConstrainCurrent: """ @pytest.mark.unit - @pytest.mark.solve - def test_iota_to_current_and_back(self, DSHAPE): + def test_iota_to_current_and_back(self): """Test we can recover the rotational transform in simple cases. Given ``iota``, compute ``I(iota)`` as defined in equation 6. @@ -65,8 +64,7 @@ def test_iota_to_current_and_back(self, DSHAPE): while names[-1] + "r" in data_index["desc.equilibrium.equilibrium.Equilibrium"]: names += [names[-1] + "r"] - def test(stellarator, grid_type): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] + def test(eq, grid_type): kwargs = {"L": eq.L_grid, "M": eq.M_grid, "N": eq.N_grid, "NFP": eq.NFP} if grid_type != QuadratureGrid: kwargs["sym"] = eq.sym @@ -119,14 +117,15 @@ def test(stellarator, grid_type): err_msg=name, ) + eq = desc.examples.get("DSHAPE") + eq.change_resolution(3, 3, 0, 6, 6, 0) # Only makes sense to test on configurations with fixed iota profiles. - test(DSHAPE, QuadratureGrid) - test(DSHAPE, ConcentricGrid) - test(DSHAPE, LinearGrid) + test(eq, QuadratureGrid) + test(eq, ConcentricGrid) + test(eq, LinearGrid) @pytest.mark.unit - @pytest.mark.solve - def test_current_to_iota_and_back(self, HELIOTRON_vac): + def test_current_to_iota_and_back(self): """Test we can recover the enclosed net toroidal current in simple cases. Given ``I``, compute ``iota(I)`` as defined in equation 18. @@ -147,8 +146,7 @@ def test_current_to_iota_and_back(self, HELIOTRON_vac): while names[-1] + "r" in data_index["desc.equilibrium.equilibrium.Equilibrium"]: names += [names[-1] + "r"] - def test(stellarator, grid_type): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] + def test(eq, grid_type): kwargs = {"L": eq.L_grid, "M": eq.M_grid, "N": eq.N_grid, "NFP": eq.NFP} if grid_type != QuadratureGrid: kwargs["sym"] = eq.sym @@ -192,6 +190,8 @@ def test(stellarator, grid_type): ) # Only makes sense to test on configurations with fixed current profiles. - test(HELIOTRON_vac, QuadratureGrid) - test(HELIOTRON_vac, ConcentricGrid) - test(HELIOTRON_vac, LinearGrid) + eq = desc.examples.get("ESTELL") + eq.change_resolution(3, 3, 3, 6, 6, 6) + test(eq, QuadratureGrid) + test(eq, ConcentricGrid) + test(eq, LinearGrid) diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index ab9cb39d71..f63f0593eb 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -6,7 +6,6 @@ import numpy as np import pytest -from netCDF4 import Dataset from desc.__main__ import main from desc.backend import sign @@ -22,42 +21,10 @@ @pytest.mark.unit -@pytest.mark.solve -def test_compute_geometry(DSHAPE_current): - """Test computation of plasma geometric values.""" - - def test(stellarator): - # VMEC values - file = Dataset(str(stellarator["vmec_nc_path"]), mode="r") - V_vmec = float(file.variables["volume_p"][-1]) - R0_vmec = float(file.variables["Rmajor_p"][-1]) - a_vmec = float(file.variables["Aminor_p"][-1]) - ar_vmec = float(file.variables["aspect"][-1]) - file.close() - - # DESC values - eq = EquilibriaFamily.load(load_from=str(stellarator["desc_h5_path"]))[-1] - data = eq.compute("R0/a") - V_desc = data["V"] - R0_desc = data["R0"] - a_desc = data["a"] - ar_desc = data["R0/a"] - - assert abs(V_vmec - V_desc) < 5e-3 - assert abs(R0_vmec - R0_desc) < 5e-3 - assert abs(a_vmec - a_desc) < 5e-3 - assert abs(ar_vmec - ar_desc) < 5e-3 - - test(DSHAPE_current) - - -@pytest.mark.slow -@pytest.mark.unit -@pytest.mark.solve -def test_compute_theta_coords(DSHAPE_current): +def test_compute_theta_coords(): """Test root finding for theta(theta*,lambda(theta)).""" - eq = EquilibriaFamily.load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - + eq = get("DSHAPE_CURRENT") + eq.change_resolution(3, 3, 0, 6, 6, 0) rho = np.linspace(0.01, 0.99, 200) theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) zeta = np.linspace(0, 2 * np.pi, 200, endpoint=False) @@ -79,37 +46,9 @@ def test_compute_theta_coords(DSHAPE_current): @pytest.mark.unit def test_map_coordinates(): - """Test root finding for (rho,theta,zeta) from (R,phi,Z).""" - eq = get("DSHAPE") - - inbasis = ["alpha", "phi", "rho"] - outbasis = ["rho", "theta_PEST", "zeta"] - - rho = np.linspace(0.01, 0.99, 20) - theta = np.linspace(0, np.pi, 20, endpoint=False) - zeta = np.linspace(0, np.pi, 20, endpoint=False) - - grid = Grid(np.vstack([rho, theta, zeta]).T, sort=False) - in_data = eq.compute(inbasis, grid=grid) - in_coords = np.stack([in_data[k] for k in inbasis], axis=-1) - out_data = eq.compute(outbasis, grid=grid) - out_coords = np.stack([out_data[k] for k in outbasis], axis=-1) - - out = eq.map_coordinates( - in_coords, - inbasis, - outbasis, - period=(2 * np.pi, 2 * np.pi, np.inf), - maxiter=40, - ) - np.testing.assert_allclose(out, out_coords, rtol=1e-4, atol=1e-4) - - -@pytest.mark.unit -def test_map_coordinates2(): """Test root finding for (rho,theta,zeta) for common use cases.""" eq = get("W7-X") - + eq.change_resolution(3, 3, 3, 6, 6, 6) n = 100 # finding coordinates along a single field line coords = np.array([np.ones(n), np.zeros(n), np.linspace(0, 10 * np.pi, n)]).T @@ -143,7 +82,7 @@ def test_map_coordinates2(): def test_map_coordinates_derivative(): """Test root finding for (rho,theta,zeta) from (R,phi,Z).""" eq = get("DSHAPE") - + eq.change_resolution(3, 3, 0, 6, 6, 0) inbasis = ["alpha", "phi", "rho"] outbasis = ["rho", "theta_PEST", "zeta"] @@ -201,17 +140,16 @@ def bar(L_lmn): @pytest.mark.slow @pytest.mark.unit -@pytest.mark.solve -def test_to_sfl(DSHAPE_current): +def test_to_sfl(): """Test converting an equilibrium to straight field line coordinates.""" - eq = EquilibriaFamily.load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - + eq = get("DSHAPE_CURRENT") + eq.change_resolution(6, 6, 0, 12, 12, 0) Rr1, Zr1, Rv1, Zv1 = compute_coords(eq) Rr2, Zr2, Rv2, Zv2 = compute_coords(eq.to_sfl()) rho_err, theta_err = area_difference(Rr1, Rr2, Zr1, Zr2, Rv1, Rv2, Zv1, Zv2) - np.testing.assert_allclose(rho_err, 0, atol=2.5e-4) - np.testing.assert_allclose(theta_err, 0, atol=1e-4) + np.testing.assert_allclose(rho_err, 0, atol=1e-2) + np.testing.assert_allclose(theta_err, 0, atol=2e-4) @pytest.mark.slow @@ -237,26 +175,24 @@ def test_grid_resolution_warning(): eq = Equilibrium(L=3, M=3, N=3) eqN = eq.copy() eqN.change_resolution(N=1, N_grid=0) - with pytest.warns(Warning): - eqN.solve(ftol=1e-2, maxiter=2) + # if we first raise warnings to errors then check for error we can avoid + # actually running the full solve + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + eqN.solve(ftol=1e-2, maxiter=2) eqM = eq.copy() eqM.change_resolution(M=eq.M, M_grid=eq.M - 1) - with pytest.warns(Warning): - eqM.solve(ftol=1e-2, maxiter=2) + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + eqM.solve(ftol=1e-2, maxiter=2) eqL = eq.copy() eqL.change_resolution(L=eq.L, L_grid=eq.L - 1) - with pytest.warns(Warning): - eqL.solve(ftol=1e-2, maxiter=2) - - -@pytest.mark.unit -def test_eq_change_grid_resolution(): - """Test changing equilibrium grid resolution.""" - eq = Equilibrium(L=2, M=2, N=2) - eq.change_resolution(L_grid=10, M_grid=10, N_grid=10) - assert eq.L_grid == 10 - assert eq.M_grid == 10 - assert eq.N_grid == 10 + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + eqL.solve(ftol=1e-2, maxiter=2) @pytest.mark.unit @@ -316,34 +252,6 @@ def test_resolution(): assert eq1.resolution == eq2.resolution -@pytest.mark.unit -def test_symmetry(): - """Test changing equilibrium symmetry.""" - M = 6 - N = 3 - surface = get("W7-X").surface.change_resolution(M=M, N=N) - eq_sym1 = Equilibrium(M=M, N=N, surface=surface, sym=True) - eq_asym1 = Equilibrium(M=M, N=N, surface=surface, sym=False) - - eq_sym2 = eq_asym1.copy() - eq_asym2 = eq_sym1.copy() - - eq_sym2.change_resolution(sym=True) - eq_asym2.change_resolution(sym=False) - - np.testing.assert_allclose(eq_sym1.R_lmn, eq_sym2.R_lmn) - np.testing.assert_allclose(eq_sym1.Z_lmn, eq_sym2.Z_lmn) - np.testing.assert_allclose(eq_sym1.L_lmn, eq_sym2.L_lmn) - np.testing.assert_allclose(eq_sym1.Rb_lmn, eq_sym2.Rb_lmn) - np.testing.assert_allclose(eq_sym1.Zb_lmn, eq_sym2.Zb_lmn) - - np.testing.assert_allclose(eq_asym1.R_lmn, eq_asym2.R_lmn) - np.testing.assert_allclose(eq_asym1.Z_lmn, eq_asym2.Z_lmn) - np.testing.assert_allclose(eq_asym1.L_lmn, eq_asym2.L_lmn) - np.testing.assert_allclose(eq_asym1.Rb_lmn, eq_asym2.Rb_lmn) - np.testing.assert_allclose(eq_asym1.Zb_lmn, eq_asym2.Zb_lmn) - - @pytest.mark.unit def test_equilibrium_from_near_axis(): """Test loading a solution from pyQSC/pyQIC.""" @@ -416,6 +324,7 @@ def test_change_NFP(): with warnings.catch_warnings(): warnings.simplefilter("error") eq = get("HELIOTRON") + eq.change_resolution(3, 3, 1, 6, 6, 2) eq.change_resolution(NFP=4) obj = get_equilibrium_objective(eq=eq) obj.build() @@ -452,7 +361,7 @@ def test_backward_compatible_load_and_resolve(): eq.change_resolution(4, 4, 4, 4, 4, 4) f_obj = ForceBalance(eq=eq) - obj = ObjectiveFunction(f_obj) + obj = ObjectiveFunction(f_obj, use_jit=False) eq.solve(maxiter=1, objective=obj) diff --git a/tests/test_examples.py b/tests/test_examples.py index 9acdbe597c..3f93c58921 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -11,7 +11,7 @@ from desc.backend import jnp from desc.coils import FourierRZCoil -from desc.continuation import _solve_axisym, solve_continuation_automatic +from desc.continuation import solve_continuation_automatic from desc.equilibrium import EquilibriaFamily, Equilibrium from desc.examples import get from desc.geometry import FourierRZToroidalSurface @@ -38,7 +38,6 @@ ForceBalance, ForceBalanceAnisotropic, GenericObjective, - HelicalForceBalance, LinearObjectiveFromUser, MeanCurvature, ObjectiveFunction, @@ -47,7 +46,6 @@ PrincipalCurvature, QuasisymmetryBoozer, QuasisymmetryTwoTerm, - RadialForceBalance, VacuumBoundaryError, Volume, get_fixed_boundary_constraints, @@ -60,26 +58,6 @@ from .utils import area_difference_desc, area_difference_vmec -@pytest.mark.regression -@pytest.mark.solve -def test_SOLOVEV_vacuum(SOLOVEV_vac): - """Tests that the SOLOVEV vacuum example gives no rotational transform.""" - eq = EquilibriaFamily.load(load_from=str(SOLOVEV_vac["desc_h5_path"]))[-1] - data = eq.compute("|J|") - - np.testing.assert_allclose(data["iota"], 0, atol=1e-16) - np.testing.assert_allclose(data["|J|"], 0, atol=3e-3) - - # test that solving with the continuation method works correctly - # when eq resolution is lower than the mres_step - eq.change_resolution(L=3, M=3) - eqf = _solve_axisym(eq, mres_step=6) - assert len(eqf) == 1 - assert eqf[-1].L == eq.L - assert eqf[-1].M == eq.M - assert eqf[-1].N == eq.N - - @pytest.mark.regression @pytest.mark.solve def test_SOLOVEV_results(SOLOVEV): @@ -154,66 +132,6 @@ def test_HELIOTRON_vac_results(HELIOTRON_vac): np.testing.assert_allclose(curr(np.linspace(0, 1, 20)), 0, atol=1e-8) -@pytest.mark.regression -@pytest.mark.solve -def test_HELIOTRON_vac2_results(HELIOTRON_vac, HELIOTRON_vac2): - """Tests that the 2 methods for solving vacuum give the same results.""" - eq1 = EquilibriaFamily.load(load_from=str(HELIOTRON_vac["desc_h5_path"]))[-1] - eq2 = EquilibriaFamily.load(load_from=str(HELIOTRON_vac2["desc_h5_path"]))[-1] - rho_err, theta_err = area_difference_desc(eq1, eq2) - np.testing.assert_allclose(rho_err[:, 4:], 0, atol=1e-2) - np.testing.assert_allclose(theta_err, 0, atol=1e-4) - curr1 = eq1.get_profile("current") - curr2 = eq2.get_profile("current") - iota1 = eq1.get_profile("iota") - iota2 = eq2.get_profile("iota") - np.testing.assert_allclose(curr1(np.linspace(0, 1, 20)), 0, atol=1e-8) - np.testing.assert_allclose(curr2(np.linspace(0, 1, 20)), 0, atol=1e-8) - np.testing.assert_allclose(iota1.params, iota2.params, rtol=1e-1, atol=1e-1) - - -@pytest.mark.regression -@pytest.mark.solve -def test_force_balance_grids(): - """Compares radial & helical force balance on same vs different grids.""" - # When ConcentricGrid had a rotation option, RadialForceBalance, HelicalForceBalance - # defaulted to cos, sin rotation, respectively. - # This test has been kept to increase code coverage. - - def test(iota=False): - if iota: - eq1 = Equilibrium(iota=PowerSeriesProfile(0), sym=True) - eq2 = Equilibrium(iota=PowerSeriesProfile(0), sym=True) - else: - eq1 = Equilibrium(current=PowerSeriesProfile(0), sym=True) - eq2 = Equilibrium(current=PowerSeriesProfile(0), sym=True) - - res = 3 - eq1.change_resolution(L=res, M=res) - eq1.L_grid = res - eq1.M_grid = res - eq2.change_resolution(L=res, M=res) - eq2.L_grid = res - eq2.M_grid = res - - # force balances on the same grids - obj1 = ObjectiveFunction(ForceBalance(eq=eq1)) - eq1.solve(objective=obj1) - - # force balances on different grids - obj2 = ObjectiveFunction( - (RadialForceBalance(eq=eq2), HelicalForceBalance(eq=eq2)) - ) - eq2.solve(objective=obj2) - - np.testing.assert_allclose(eq1.R_lmn, eq2.R_lmn, atol=5e-4) - np.testing.assert_allclose(eq1.Z_lmn, eq2.Z_lmn, atol=5e-4) - np.testing.assert_allclose(eq1.L_lmn, eq2.L_lmn, atol=2e-3) - - test(iota=True) - test(iota=False) - - @pytest.mark.regression @pytest.mark.solve def test_solve_bounds(): @@ -396,38 +314,6 @@ def test_ATF_results(tmpdir_factory): np.testing.assert_allclose(theta_err, 0, atol=5e-4) -@pytest.mark.regression -@pytest.mark.solve -def test_ESTELL_results(tmpdir_factory): - """Test automatic continuation method with ESTELL.""" - output_dir = tmpdir_factory.mktemp("result") - eq0 = get("ESTELL") - eq = Equilibrium( - Psi=eq0.Psi, - NFP=eq0.NFP, - L=eq0.L, - M=eq0.M, - N=eq0.N, - L_grid=eq0.L_grid, - M_grid=eq0.M_grid, - N_grid=eq0.N_grid, - pressure=eq0.pressure, - current=eq0.current, - surface=eq0.get_surface_at(rho=1), - sym=eq0.sym, - spectral_indexing=eq0.spectral_indexing, - ) - eqf = EquilibriaFamily.solve_continuation_automatic( - eq, - verbose=2, - checkpoint_path=output_dir.join("ESTELL.h5"), - ) - eqf = load(output_dir.join("ESTELL.h5")) - rho_err, theta_err = area_difference_desc(eq0, eqf[-1]) - np.testing.assert_allclose(rho_err[:, 4:], 0, atol=5e-2) - np.testing.assert_allclose(theta_err, 0, atol=1e-4) - - @pytest.mark.regression @pytest.mark.optimize def test_simsopt_QH_comparison(): @@ -735,7 +621,7 @@ def test_NAE_QIC_solve(): @pytest.mark.unit @pytest.mark.optimize -def test_multiobject_optimization(): +def test_multiobject_optimization_al(): """Test for optimizing multiple objects at once.""" eq = Equilibrium(L=4, M=4, N=0, iota=2) surf = FourierRZToroidalSurface( diff --git a/tests/test_input_output.py b/tests/test_input_output.py index 3e1ad9dbc2..0d2fe6648c 100644 --- a/tests/test_input_output.py +++ b/tests/test_input_output.py @@ -8,6 +8,7 @@ import numpy as np import pytest +import desc.examples from desc.basis import FourierZernikeBasis from desc.equilibrium import Equilibrium from desc.grid import LinearGrid @@ -478,11 +479,11 @@ def test_reader_read_obj(reader_test_file): @pytest.mark.unit @pytest.mark.solve -def test_pickle_io(DSHAPE_current, tmpdir_factory): +def test_pickle_io(tmpdir_factory): """Test saving and loading equilibrium in pickle format.""" tmpdir = tmpdir_factory.mktemp("desc_inputs") tmp_path = tmpdir.join("solovev_test.pkl") - eqf = load(load_from=str(DSHAPE_current["desc_h5_path"])) + eqf = desc.examples.get("DSHAPE_CURRENT", "all") eqf.save(tmp_path, file_format="pickle") peqf = load(tmp_path, file_format="pickle") assert equals(eqf, peqf) @@ -490,11 +491,11 @@ def test_pickle_io(DSHAPE_current, tmpdir_factory): @pytest.mark.unit @pytest.mark.solve -def test_ascii_io(DSHAPE_current, tmpdir_factory): +def test_ascii_io(tmpdir_factory): """Test saving and loading equilibrium in ASCII format.""" tmpdir = tmpdir_factory.mktemp("desc_inputs") tmp_path = tmpdir.join("solovev_test.txt") - eq1 = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] + eq1 = desc.examples.get("DSHAPE_CURRENT") eq1.iota = eq1.get_profile("iota", grid=LinearGrid(30, 16, 0)).to_powerseries( sym=True ) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index af4bea4437..cecbf2bbe4 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -40,8 +40,8 @@ FixSumModesR, FixSumModesZ, FixThetaSFL, + GenericObjective, ObjectiveFunction, - QuasisymmetryTwoTerm, get_equilibrium_objective, get_fixed_axis_constraints, get_fixed_boundary_constraints, @@ -105,7 +105,8 @@ def test_LambdaGauge_asym(): np.testing.assert_allclose(lam, 0, atol=1e-15) -@pytest.mark.unit +@pytest.mark.regression +@pytest.mark.solve def test_bc_on_interior_surfaces(): """Test applying boundary conditions on internal surface.""" surf = FourierRZToroidalSurface(rho=0.5) @@ -408,15 +409,15 @@ def test_correct_indexing_passed_modes(): """Test Indexing when passing in specified modes, related to gh issue #380.""" n = 1 eq = desc.examples.get("W7-X") - grid = LinearGrid( - M=eq.M, N=eq.N, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True - ) + eq.change_resolution(3, 3, 3, 6, 6, 6) + eq.surface = eq.get_surface_at(1.0) objective = ObjectiveFunction( ( - QuasisymmetryTwoTerm(eq=eq, weight=1e-2, helicity=(1, -eq.NFP), grid=grid), - AspectRatio(eq=eq, target=8, weight=1e2), + # just need dummy objective for factorizing constraints + GenericObjective("0", eq=eq), ), + use_jit=False, ) objective.build() @@ -438,7 +439,7 @@ def test_correct_indexing_passed_modes(): BoundaryZSelfConsistency(eq=eq), FixPressure(eq=eq), ) - constraint = ObjectiveFunction(constraints) + constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( @@ -461,15 +462,12 @@ def test_correct_indexing_passed_modes_and_passed_target(): """Test Indexing when passing in specified modes, related to gh issue #380.""" n = 1 eq = desc.examples.get("W7-X") - grid = LinearGrid( - M=eq.M, N=eq.N, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True - ) + eq.change_resolution(3, 3, 3, 6, 6, 6) + eq.surface = eq.get_surface_at(1.0) objective = ObjectiveFunction( - ( - QuasisymmetryTwoTerm(eq=eq, weight=1e-2, helicity=(1, -eq.NFP), grid=grid), - AspectRatio(eq=eq, target=8, weight=1e2), - ), + (GenericObjective("0", eq=eq),), + use_jit=False, ) objective.build() @@ -484,6 +482,7 @@ def test_correct_indexing_passed_modes_and_passed_target(): idxs = [] for mode in R_modes: idxs.append(eq.surface.R_basis.get_idx(*mode)) + idxs = np.array(idxs) target_R = eq.surface.R_lmn[idxs] Z_modes = eq.surface.Z_basis.modes[ @@ -492,6 +491,7 @@ def test_correct_indexing_passed_modes_and_passed_target(): idxs = [] for mode in Z_modes: idxs.append(eq.surface.Z_basis.get_idx(*mode)) + idxs = np.array(idxs) target_Z = eq.surface.Z_lmn[idxs] constraints = ( @@ -501,7 +501,7 @@ def test_correct_indexing_passed_modes_and_passed_target(): BoundaryZSelfConsistency(eq=eq), FixPressure(eq=eq), ) - constraint = ObjectiveFunction(constraints) + constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( @@ -524,15 +524,13 @@ def test_correct_indexing_passed_modes_axis(): """Test Indexing when passing in specified axis modes, related to gh issue #380.""" n = 1 eq = desc.examples.get("W7-X") - grid = LinearGrid( - M=eq.M, N=eq.N, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True - ) + eq.change_resolution(3, 3, 3, 6, 6, 6) + eq.surface = eq.get_surface_at(1.0) + eq.axis = eq.get_axis() objective = ObjectiveFunction( - ( - QuasisymmetryTwoTerm(eq=eq, weight=1e-2, helicity=(1, -eq.NFP), grid=grid), - AspectRatio(eq=eq, target=8, weight=1e2), - ), + (GenericObjective("0", eq=eq),), + use_jit=False, ) objective.build() @@ -562,7 +560,7 @@ def test_correct_indexing_passed_modes_axis(): FixSumModesZ(eq=eq, modes=np.array([[3, 3, -3], [4, 4, -4]]), normalize=False), FixPressure(eq=eq), ) - constraint = ObjectiveFunction(constraints) + constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( @@ -586,16 +584,13 @@ def test_correct_indexing_passed_modes_and_passed_target_axis(): n = 1 eq = desc.examples.get("W7-X") - - grid = LinearGrid( - M=eq.M, N=eq.N, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True - ) + eq.change_resolution(4, 4, 4, 8, 8, 8) + eq.surface = eq.get_surface_at(1.0) + eq.axis = eq.get_axis() objective = ObjectiveFunction( - ( - QuasisymmetryTwoTerm(eq=eq, weight=1e-2, helicity=(1, -eq.NFP), grid=grid), - AspectRatio(eq=eq, target=8, weight=1e2), - ), + (GenericObjective("0", eq=eq),), + use_jit=False, ) objective.build() @@ -609,6 +604,7 @@ def test_correct_indexing_passed_modes_and_passed_target_axis(): idxs = [] for mode in R_modes: idxs.append(eq.axis.R_basis.get_idx(*mode)) + idxs = np.array(idxs) target_R = eq.axis.R_n[idxs] Z_modes = eq.axis.Z_basis.modes[np.max(np.abs(eq.axis.Z_basis.modes), 1) > n + 1, :] @@ -616,6 +612,7 @@ def test_correct_indexing_passed_modes_and_passed_target_axis(): idxs = [] for mode in Z_modes: idxs.append(eq.axis.Z_basis.get_idx(*mode)) + idxs = np.array(idxs) target_Z = eq.axis.Z_n[idxs] constraints = ( @@ -691,7 +688,7 @@ def test_correct_indexing_passed_modes_and_passed_target_axis(): normalize=False, ), ) - constraint = ObjectiveFunction(constraints) + constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 9bae1b2151..ae5b51c35f 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -427,6 +427,7 @@ def test_fourier_current_potential_field_symmetry(self): with pytest.raises(ValueError): field.Phi_mn = np.ones((basis.num_modes + 1,)) + @pytest.mark.unit def test_io_fourier_current_field(self, tmpdir_factory): """Test that i/o works for FourierCurrentPotentialField.""" surface = FourierRZToroidalSurface( @@ -519,6 +520,7 @@ def test_fourier_current_potential_field_asserts(self): NFP=10, ) + @pytest.mark.unit def test_change_Phi_basis_fourier_current_field(self): """Test that change_Phi_resolution works for FourierCurrentPotentialField.""" surface = FourierRZToroidalSurface( @@ -575,6 +577,7 @@ def test_change_Phi_basis_fourier_current_field(self): np.testing.assert_allclose(field.Phi_basis.modes, basis.modes) assert field.Phi_basis.sym == "sin" + @pytest.mark.unit def test_init_Phi_mn_fourier_current_field(self): """Test initial Phi_mn size is correct for FourierCurrentPotentialField.""" surface = FourierRZToroidalSurface( diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 3741a3f4f5..d8cf475014 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -6,6 +6,8 @@ This module primarily tests the constructing/building/calling methods. """ +import warnings + import numpy as np import pytest from scipy.constants import elementary_charge, mu_0 @@ -31,14 +33,11 @@ CoilCurvature, CoilLength, CoilTorsion, - CurrentDensity, Elongation, Energy, ForceBalance, ForceBalanceAnisotropic, GenericObjective, - GoodCoordinates, - HelicalForceBalance, Isodynamicity, LinearObjectiveFromUser, MagneticWell, @@ -53,7 +52,6 @@ QuasisymmetryBoozer, QuasisymmetryTripleProduct, QuasisymmetryTwoTerm, - RadialForceBalance, RotationalTransform, Shear, ToroidalCurrent, @@ -62,7 +60,6 @@ ) from desc.objectives._free_boundary import BoundaryErrorNESTOR from desc.objectives.normalization import compute_scaling_factors -from desc.objectives.objective_funs import _Objective from desc.objectives.utils import softmax, softmin from desc.profiles import FourierZernikeProfile, PowerSeriesProfile from desc.vmec_utils import ptolemy_linear_transform @@ -577,6 +574,171 @@ def test(eq): test(get("DSHAPE")) test(get("HELIOTRON")) + @pytest.mark.unit + def test_plasma_vessel_distance(self): + """Test calculation of min distance from plasma to vessel.""" + R0 = 10.0 + a_p = 1.0 + a_s = 2.0 + # default eq has R0=10, a=1 + eq = Equilibrium(M=3, N=2) + # surface with same R0, a=2, so true d=1 for all pts + surface = FourierRZToroidalSurface( + R_lmn=[R0, a_s], Z_lmn=[-a_s], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]] + ) + # For equally spaced grids, should get true d=1 + surf_grid = LinearGrid(M=5, N=6) + plas_grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, plasma_grid=plas_grid, surface_grid=surf_grid, surface=surface + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + np.testing.assert_allclose(d, a_s - a_p) + + # for unequal M, should have error of order M_spacing*a_p + surf_grid = LinearGrid(M=5, N=6) + plas_grid = LinearGrid(M=10, N=6) + obj = PlasmaVesselDistance( + eq=eq, + plasma_grid=plas_grid, + surface_grid=surf_grid, + surface=surface, + surface_fixed=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert abs(d.min() - (a_s - a_p)) < 1e-14 + assert abs(d.max() - (a_s - a_p)) < surf_grid.spacing[0, 1] * a_p + + # for unequal N, should have error of order N_spacing*R0 + surf_grid = LinearGrid(M=5, N=6) + plas_grid = LinearGrid(M=5, N=12) + obj = PlasmaVesselDistance( + eq=eq, plasma_grid=plas_grid, surface_grid=surf_grid, surface=surface + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert abs(d.min() - (a_s - a_p)) < 1e-14 + assert abs(d.max() - (a_s - a_p)) < surf_grid.spacing[0, 2] * R0 + # ensure that it works (dimension-wise) when compute_scaled is called + _ = obj.compute_scaled(*obj.xs(eq, surface)) + + grid = LinearGrid(L=3, M=3, N=3) + eq = Equilibrium() + surf = FourierRZToroidalSurface() + obj = PlasmaVesselDistance( + surface=surf, surface_grid=grid, plasma_grid=grid, eq=eq + ) + with pytest.raises(UserWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + obj.build() + + # test softmin, should give value less than true minimum + surf_grid = LinearGrid(M=5, N=6) + plas_grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + plasma_grid=plas_grid, + surface_grid=surf_grid, + surface=surface, + use_softmin=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert np.all(np.abs(d) < a_s - a_p) + + # for large enough alpha, should be same as actual min + obj = PlasmaVesselDistance( + eq=eq, + plasma_grid=plas_grid, + surface_grid=surf_grid, + surface=surface, + use_softmin=True, + alpha=100, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + np.testing.assert_allclose(d, a_s - a_p) + + @pytest.mark.unit + def test_mean_curvature(self): + """Test for mean curvature objective function.""" + # torus should have mean curvature negative everywhere + eq = Equilibrium() + obj = MeanCurvature(eq=eq) + obj.build() + H = obj.compute_unscaled(*obj.xs(eq)) + assert np.all(H <= 0) + + # more shaped case like NCSX should have some positive curvature + eq = get("NCSX") + obj = MeanCurvature(eq=eq) + obj.build() + H = obj.compute_unscaled(*obj.xs(eq)) + assert np.any(H > 0) + + # check using the surface + obj = MeanCurvature(eq=eq.surface) + obj.build() + H = obj.compute_unscaled(*obj.xs(eq.surface)) + assert np.any(H > 0) + + @pytest.mark.unit + def test_principal_curvature(self): + """Test for principal curvature objective function.""" + eq1 = get("DSHAPE") + eq2 = get("NCSX") + obj1 = PrincipalCurvature(eq=eq1, normalize=False) + obj1.build() + K1 = obj1.compute_unscaled(*obj1.xs(eq1)) + obj2 = PrincipalCurvature(eq=eq2, normalize=False) + obj2.build() + K2 = obj2.compute_unscaled(*obj2.xs(eq2)) + + # simple test: NCSX should have higher mean absolute curvature than DSHAPE + assert K1.mean() < K2.mean() + + # same test but using the surface directly + obj1 = PrincipalCurvature(eq=eq1.surface, normalize=False) + obj1.build() + K1 = obj1.compute_unscaled(*obj1.xs(eq1.surface)) + obj2 = PrincipalCurvature(eq=eq2.surface, normalize=False) + obj2.build() + K2 = obj2.compute_unscaled(*obj2.xs(eq2.surface)) + + # simple test: NCSX should have higher mean absolute curvature than DSHAPE + assert K1.mean() < K2.mean() + + @pytest.mark.unit + def test_field_scale_length(self): + """Test for B field scale length objective function.""" + surf1 = FourierRZToroidalSurface( + R_lmn=[5, 1], Z_lmn=[-1], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]], NFP=1 + ) + surf2 = FourierRZToroidalSurface( + R_lmn=[10, 2], + Z_lmn=[-2], + modes_R=[[0, 0], [1, 0]], + modes_Z=[[-1, 0]], + NFP=1, + ) + eq1 = Equilibrium(L=2, M=2, N=0, surface=surf1) + eq2 = Equilibrium(L=2, M=2, N=0, surface=surf2) + eq1.solve() + eq2.solve() + + obj1 = BScaleLength(eq=eq1, normalize=False) + obj2 = BScaleLength(eq=eq2, normalize=False) + obj1.build() + obj2.build() + + L1 = obj1.compute_unscaled(*obj1.xs(eq1)) + L2 = obj2.compute_unscaled(*obj2.xs(eq2)) + + np.testing.assert_array_less(L1, L2) + @pytest.mark.unit def test_coil_length(self): """Tests coil length.""" @@ -655,7 +817,7 @@ def test(coil, grid=None): test(nested_coils, grid=nested_grids) -@pytest.mark.unit +@pytest.mark.regression def test_derivative_modes(): """Test equality of derivatives using batched, looped methods.""" eq = Equilibrium(M=2, N=1, L=2) @@ -711,69 +873,6 @@ def test_derivative_modes(): np.testing.assert_allclose(H1, H3, atol=1e-10) -@pytest.mark.unit -@pytest.mark.xfail -def test_rejit(): - """Test that updating attributes and recompiling correctly updates.""" - - class DummyObjective(_Objective): - def __init__(self, y, eq=None, target=0, weight=1, name="dummy"): - self.y = y - super().__init__(things=eq, target=target, weight=weight, name=name) - - def build(self, use_jit=True, verbose=1): - self._dim_f = 1 - super().build(use_jit, verbose) - - def compute(self, params, constants=None): - return 200 + self.target * self.weight - self.y * params["R_lmn"] ** 3 - - eq = Equilibrium() - obj = DummyObjective(3, eq=eq) - obj.build() - assert obj.compute_unscaled({"R_lmn": 4}) == 8 - assert obj.compute_scaled_error({"R_lmn": 4}) == 8 - obj.target = 1 - obj.weight = 2 - assert obj.compute({"R_lmn": 4}) == 10 # compute method is not JIT compiled - assert ( - obj.compute_scaled_error({"R_lmn": 4}) == 8 - ) # only compute_scaled is JIT compiled - obj.jit() - assert obj.compute({"R_lmn": 4}) == 10 - assert obj.compute_scaled_error({"R_lmn": 4}) == 18 - - objFun = ObjectiveFunction(obj) - objFun.build() - x = objFun.x(eq) - - f = objFun.compute_scaled_error(x) - J = objFun.jac_scaled(x) - np.testing.assert_allclose(f, [-5598, 402, 396]) - np.testing.assert_allclose(J[:, eq.x_idx["R_lmn"]], np.diag([-1800, 0, -18])) - objFun.objectives[0].target = 3 - objFun.objectives[0].weight = 4 - objFun.objectives[0].y = 2 - np.testing.assert_allclose(objFun.compute_scaled_error(x), f) - np.testing.assert_allclose(objFun.jac_scaled(x), J) - objFun.jit() - np.testing.assert_allclose(objFun.compute_scaled_error(x), [-7164, 836, 828]) - np.testing.assert_allclose(objFun.jac_scaled(x), J * 4 / 3) - - -@pytest.mark.unit -def test_generic_compute(): - """Test for GH issue #388.""" - eq = Equilibrium() - obj = ObjectiveFunction(AspectRatio(target=2, weight=1, eq=eq)) - obj.build() - a1 = obj.compute_scalar(obj.x(eq)) - obj = ObjectiveFunction(GenericObjective("R0/a", target=2, weight=1, eq=eq)) - obj.build() - a2 = obj.compute_scalar(obj.x(eq)) - assert np.allclose(a1, a2) - - @pytest.mark.unit def test_getter_setter(): """Test getter and setter methods of Objectives.""" @@ -901,167 +1000,6 @@ def test_target_profiles(): ) -@pytest.mark.unit -def test_plasma_vessel_distance(): - """Test calculation of min distance from plasma to vessel.""" - R0 = 10.0 - a_p = 1.0 - a_s = 2.0 - # default eq has R0=10, a=1 - eq = Equilibrium(M=3, N=2) - # surface with same R0, a=2, so true d=1 for all pts - surface = FourierRZToroidalSurface( - R_lmn=[R0, a_s], Z_lmn=[-a_s], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]] - ) - # For equally spaced grids, should get true d=1 - surf_grid = LinearGrid(M=5, N=6) - plas_grid = LinearGrid(M=5, N=6) - obj = PlasmaVesselDistance( - eq=eq, plasma_grid=plas_grid, surface_grid=surf_grid, surface=surface - ) - obj.build() - d = obj.compute_unscaled(*obj.xs(eq, surface)) - np.testing.assert_allclose(d, a_s - a_p) - - # for unequal M, should have error of order M_spacing*a_p - surf_grid = LinearGrid(M=5, N=6) - plas_grid = LinearGrid(M=10, N=6) - obj = PlasmaVesselDistance( - eq=eq, - plasma_grid=plas_grid, - surface_grid=surf_grid, - surface=surface, - surface_fixed=True, - ) - obj.build() - d = obj.compute_unscaled(*obj.xs(eq, surface)) - assert abs(d.min() - (a_s - a_p)) < 1e-14 - assert abs(d.max() - (a_s - a_p)) < surf_grid.spacing[0, 1] * a_p - - # for unequal N, should have error of order N_spacing*R0 - surf_grid = LinearGrid(M=5, N=6) - plas_grid = LinearGrid(M=5, N=12) - obj = PlasmaVesselDistance( - eq=eq, plasma_grid=plas_grid, surface_grid=surf_grid, surface=surface - ) - obj.build() - d = obj.compute_unscaled(*obj.xs(eq, surface)) - assert abs(d.min() - (a_s - a_p)) < 1e-14 - assert abs(d.max() - (a_s - a_p)) < surf_grid.spacing[0, 2] * R0 - # ensure that it works (dimension-wise) when compute_scaled is called - _ = obj.compute_scaled(*obj.xs(eq, surface)) - - grid = LinearGrid(L=3, M=3, N=3) - eq = Equilibrium() - surf = FourierRZToroidalSurface() - obj = PlasmaVesselDistance(surface=surf, surface_grid=grid, plasma_grid=grid, eq=eq) - with pytest.warns(UserWarning): - obj.build() - - # test softmin, should give value less than true minimum - surf_grid = LinearGrid(M=5, N=6) - plas_grid = LinearGrid(M=5, N=6) - obj = PlasmaVesselDistance( - eq=eq, - plasma_grid=plas_grid, - surface_grid=surf_grid, - surface=surface, - use_softmin=True, - ) - obj.build() - d = obj.compute_unscaled(*obj.xs(eq, surface)) - assert np.all(np.abs(d) < a_s - a_p) - - # for large enough alpha, should be same as actual min - obj = PlasmaVesselDistance( - eq=eq, - plasma_grid=plas_grid, - surface_grid=surf_grid, - surface=surface, - use_softmin=True, - alpha=100, - ) - obj.build() - d = obj.compute_unscaled(*obj.xs(eq, surface)) - np.testing.assert_allclose(d, a_s - a_p) - - -@pytest.mark.unit -def test_mean_curvature(): - """Test for mean curvature objective function.""" - # torus should have mean curvature negative everywhere - eq = Equilibrium() - obj = MeanCurvature(eq=eq) - obj.build() - H = obj.compute_unscaled(*obj.xs(eq)) - assert np.all(H <= 0) - - # more shaped case like NCSX should have some positive curvature - eq = get("NCSX") - obj = MeanCurvature(eq=eq) - obj.build() - H = obj.compute_unscaled(*obj.xs(eq)) - assert np.any(H > 0) - - # check using the surface - obj = MeanCurvature(eq=eq.surface) - obj.build() - H = obj.compute_unscaled(*obj.xs(eq.surface)) - assert np.any(H > 0) - - -@pytest.mark.unit -def test_principal_curvature(): - """Test for principal curvature objective function.""" - eq1 = get("DSHAPE") - eq2 = get("NCSX") - obj1 = PrincipalCurvature(eq=eq1, normalize=False) - obj1.build() - K1 = obj1.compute_unscaled(*obj1.xs(eq1)) - obj2 = PrincipalCurvature(eq=eq2, normalize=False) - obj2.build() - K2 = obj2.compute_unscaled(*obj2.xs(eq2)) - - # simple test: NCSX should have higher mean absolute curvature than DSHAPE - assert K1.mean() < K2.mean() - - # same test but using the surface directly - obj1 = PrincipalCurvature(eq=eq1.surface, normalize=False) - obj1.build() - K1 = obj1.compute_unscaled(*obj1.xs(eq1.surface)) - obj2 = PrincipalCurvature(eq=eq2.surface, normalize=False) - obj2.build() - K2 = obj2.compute_unscaled(*obj2.xs(eq2.surface)) - - # simple test: NCSX should have higher mean absolute curvature than DSHAPE - assert K1.mean() < K2.mean() - - -@pytest.mark.unit -def test_field_scale_length(): - """Test for B field scale length objective function.""" - surf1 = FourierRZToroidalSurface( - R_lmn=[5, 1], Z_lmn=[-1], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]], NFP=1 - ) - surf2 = FourierRZToroidalSurface( - R_lmn=[10, 2], Z_lmn=[-2], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]], NFP=1 - ) - eq1 = Equilibrium(L=2, M=2, N=0, surface=surf1) - eq2 = Equilibrium(L=2, M=2, N=0, surface=surf2) - eq1.solve() - eq2.solve() - - obj1 = BScaleLength(eq=eq1, normalize=False) - obj2 = BScaleLength(eq=eq2, normalize=False) - obj1.build() - obj2.build() - - L1 = obj1.compute_unscaled(*obj1.xs(eq1)) - L2 = obj2.compute_unscaled(*obj2.xs(eq2)) - - np.testing.assert_array_less(L1, L2) - - @pytest.mark.unit def test_profile_objective_print(capsys): """Test that the profile objectives print correctly.""" @@ -1516,7 +1454,7 @@ def test_jvp_scaled(): weight = 3 target = 5 objective = ObjectiveFunction( - Volume(target=target, normalize=True, weight=weight, eq=eq) + Volume(target=target, normalize=True, weight=weight, eq=eq), use_jit=False ) objective.build() x = objective.x(eq) @@ -1552,7 +1490,7 @@ def test_vjp(): weight = 3 target = 5 objective = ObjectiveFunction( - ForceBalance(target=target, normalize=True, weight=weight, eq=eq) + ForceBalance(target=target, normalize=True, weight=weight, eq=eq), use_jit=False ) objective.build() x = objective.x(eq) @@ -1575,7 +1513,7 @@ def test_objective_target_bounds(): asp = AspectRatio(bounds=(2, 3), normalize=False, weight=3, eq=eq) fbl = ForceBalance(normalize=True, bounds=(-1, 2), weight=5, eq=eq) - objective = ObjectiveFunction((vol, asp, fbl)) + objective = ObjectiveFunction((vol, asp, fbl), use_jit=False) objective.build() target = objective.target_scaled @@ -1608,7 +1546,9 @@ def test_objective_target_bounds(): eq = Equilibrium(L=8, M=2, N=2, iota=PowerSeriesProfile(0.42)) - con = ObjectiveFunction(RotationalTransform(eq=eq, bounds=(0.41, 0.43))) + con = ObjectiveFunction( + RotationalTransform(eq=eq, bounds=(0.41, 0.43)), use_jit=False + ) con.build() np.testing.assert_allclose(con.compute_scaled_error(con.x(eq)), 0) @@ -1663,463 +1603,417 @@ def test_loss_function_asserts(): RotationalTransform(eq=eq, loss_function=fun) -@pytest.mark.unit -@pytest.mark.slow -def test_compute_scalar_resolution(): # noqa: C901 +class TestComputeScalarResolution: """Test that compute_scalar values are roughly independent of grid resolution.""" + + # get a list of all the objectives + objectives = [ + getattr(desc.objectives, obj) + for obj in dir(desc.objectives) + if obj[0].isupper() + and (not obj.startswith("Fix")) + and (obj != "ObjectiveFunction") + and ("SelfConsistency" not in obj) + ] + specials = [ + # these require special logic + PlasmaVesselDistance, + BootstrapRedlConsistency, + BoundaryError, + VacuumBoundaryError, + GenericObjective, + Omnigenity, + CoilLength, + CoilTorsion, + CoilCurvature, + # need to avoid blowup near the axis + MercierStability, + # don't test these since they depend on what user wants + LinearObjectiveFromUser, + ObjectiveFromUser, + ] + other_objectives = list(set(objectives) - set(specials)) + eq = get("HELIOTRON") res_array = np.array([2, 2.5, 3]) - # BootstrapRedlConsistency - # this is already covered in tests/test_bootstrap.py - # by TestBootstrapObjectives.test_BootstrapRedlConsistency_resolution - - # CurrentDensity - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(CurrentDensity(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=3e-2) - - # Energy - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(Energy(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # ForceBalance - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(ForceBalance(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # ForceBalanceAnisotropic - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(ForceBalanceAnisotropic(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # HelicalForceBalance - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(HelicalForceBalance(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-1) - - # RadialForceBalance - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(RadialForceBalance(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-1) - - # GenericObjective - # scalar - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = QuadratureGrid( - L=int(eq.L * res), M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP - ) - obj = ObjectiveFunction(GenericObjective("_vol", eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - # radial profile - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = LinearGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - axis=False, + @pytest.mark.regression + def test_compute_scalar_resolution_plasma_vessel(self): + """PlasmaVesselDistance.""" + f = np.zeros_like(self.res_array, dtype=float) + surface = FourierRZToroidalSurface( + R_lmn=[10, 1.5], Z_lmn=[-1.5], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]] ) - obj = ObjectiveFunction(GenericObjective("", eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - # volume quantity - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(GenericObjective("sqrt(g)", eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # AspectRatio - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = QuadratureGrid( - L=int(eq.L * res), M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP - ) - obj = ObjectiveFunction(AspectRatio(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # BScaleLength - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP, sym=eq.sym) - obj = ObjectiveFunction(BScaleLength(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # Elongation - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = QuadratureGrid( - L=int(eq.L * res), M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP + for i, res in enumerate(self.res_array): + grid = LinearGrid( + M=int(self.eq.M * res), N=int(self.eq.N * res), NFP=self.eq.NFP + ) + obj = ObjectiveFunction( + PlasmaVesselDistance( + surface=surface, eq=self.eq, surface_grid=grid, plasma_grid=grid + ), + use_jit=False, + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_bootstrap(self): + """BootstrapRedlConsistency.""" + eq = self.eq.copy() + eq.electron_density = PowerSeriesProfile([1e19, 0, -1e19]) + eq.electron_temperature = PowerSeriesProfile([1e3, 0, -1e3]) + eq.ion_temperature = PowerSeriesProfile([1e3, 0, -1e3]) + eq.atomic_number = 1.0 + + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + grid = LinearGrid( + M=int(self.eq.M * res), N=int(self.eq.N * res), NFP=self.eq.NFP + ) + obj = ObjectiveFunction( + BootstrapRedlConsistency(eq=eq, grid=grid), use_jit=False + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_boundary_error(self): + """BoundaryError.""" + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + + pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) + iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) + surf = FourierRZToroidalSurface( + R_lmn=[4.0, 1.0], + modes_R=[[0, 0], [1, 0]], + Z_lmn=[-1.0], + modes_Z=[[-1, 0]], + NFP=1, ) - obj = ObjectiveFunction(Elongation(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # MeanCurvature - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP, sym=eq.sym) - obj = ObjectiveFunction(MeanCurvature(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) + eq = Equilibrium(M=6, N=0, Psi=1.0, surface=surf, pressure=pres, iota=iota) - # PlasmaVesselDistance - f = np.zeros_like(res_array, dtype=float) - surface = FourierRZToroidalSurface( - R_lmn=[10, 1.5], Z_lmn=[-1.5], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]] - ) - for i, res in enumerate(res_array): - grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP) - obj = ObjectiveFunction( - PlasmaVesselDistance( - surface=surface, eq=eq, surface_grid=grid, plasma_grid=grid - ), - ) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq, surface)) - np.testing.assert_allclose(f, f[-1], rtol=5e-2) - - # PrincipalCurvature - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP, sym=eq.sym) - obj = ObjectiveFunction(PrincipalCurvature(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # Volume - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = QuadratureGrid( - L=int(eq.L * res), M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP - ) - obj = ObjectiveFunction(Volume(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # RotationalTransform - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = LinearGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - axis=False, - ) - obj = ObjectiveFunction(RotationalTransform(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # ToroidalCurrent - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = LinearGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - axis=False, - ) - obj = ObjectiveFunction(ToroidalCurrent(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=4e-2) - - # Isodynamicity - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction(Isodynamicity(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # QuasisymmetryBoozer - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate((res_array + 4) * 2): - grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP) - obj = ObjectiveFunction( - QuasisymmetryBoozer(eq=eq, helicity=(1, -eq.NFP), grid=grid) + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + eq.change_resolution( + L_grid=int(eq.L * res), M_grid=int(eq.M * res), N_grid=int(eq.N * res) + ) + obj = ObjectiveFunction(BoundaryError(self.eq, ext_field), use_jit=False) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_vacuum_boundary_error(self): + """VacuumBoundaryError.""" + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + + pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) + iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) + surf = FourierRZToroidalSurface( + R_lmn=[4.0, 1.0], + modes_R=[[0, 0], [1, 0]], + Z_lmn=[-1.0], + modes_Z=[[-1, 0]], + NFP=1, ) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # QuasisymmetryTripleProduct - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), - NFP=eq.NFP, - sym=eq.sym, + eq = Equilibrium(M=6, N=0, Psi=1.0, surface=surf, pressure=pres, iota=iota) + + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + eq.change_resolution( + L_grid=int(eq.L * res), M_grid=int(eq.M * res), N_grid=int(eq.N * res) + ) + obj = ObjectiveFunction( + VacuumBoundaryError(self.eq, ext_field), use_jit=False + ) + with pytest.warns(UserWarning): + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_generic_scalar(self): + """Generic objective with scalar qty.""" + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + grid = QuadratureGrid( + L=int(self.eq.L * res), + M=int(self.eq.M * res), + N=int(self.eq.N * res), + NFP=self.eq.NFP, + ) + obj = ObjectiveFunction( + GenericObjective("_vol", eq=self.eq, grid=grid), use_jit=False + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=1e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_generic_profile(self): + """Generic objective with profile qty.""" + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + grid = LinearGrid( + L=int(self.eq.L * res), + M=int(self.eq.M * res), + N=int(self.eq.N * res), + NFP=self.eq.NFP, + sym=self.eq.sym, + axis=False, + ) + obj = ObjectiveFunction( + GenericObjective("", eq=self.eq, grid=grid), use_jit=False + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=2e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_generic_volume(self): + """Generic objective with volume qty.""" + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + grid = ConcentricGrid( + L=int(self.eq.L * res), + M=int(self.eq.M * res), + N=int(self.eq.N * res), + NFP=self.eq.NFP, + sym=self.eq.sym, + ) + obj = ObjectiveFunction( + GenericObjective("sqrt(g)", eq=self.eq, grid=grid), use_jit=False + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=2e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_mercier(self): + """Mercier stability.""" + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + rho = np.linspace(0.2, 1, int(self.eq.L * res)) + grid = LinearGrid( + rho=rho, + M=int(self.eq.M * res), + N=int(self.eq.N * res), + NFP=self.eq.NFP, + sym=self.eq.sym, + ) + obj = ObjectiveFunction( + MercierStability(eq=self.eq, grid=grid), use_jit=False + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=2e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_omnigenity(self): + """Omnigenity.""" + surf = FourierRZToroidalSurface.from_qp_model( + major_radius=1, + aspect_ratio=20, + elongation=6, + mirror_ratio=0.2, + torsion=0.1, + NFP=1, + sym=True, ) - obj = ObjectiveFunction(QuasisymmetryTripleProduct(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # QuasisymmetryTwoTerm - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - grid = ConcentricGrid( - L=int(eq.L * res), - M=int(eq.M * res), - N=int(eq.N * res), + eq = Equilibrium(Psi=6e-3, M=4, N=4, surface=surf) + eq, _ = eq.solve(objective="force", verbose=3) + field = OmnigenousField( + L_B=0, + M_B=2, + L_x=0, + M_x=0, + N_x=0, NFP=eq.NFP, - sym=eq.sym, - ) - obj = ObjectiveFunction( - QuasisymmetryTwoTerm(eq=eq, helicity=(1, -eq.NFP), grid=grid), - ) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=2e-2) - - # MagneticWell - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - rho = np.linspace(0.2, 1, int(eq.L * res)) - grid = LinearGrid( - rho=rho, M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP, sym=eq.sym - ) - obj = ObjectiveFunction(MagneticWell(eq=eq, grid=grid, target=0)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # MercierStability - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array): - rho = np.linspace(0.2, 1, int(eq.L * res)) - grid = LinearGrid( - rho=rho, M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP, sym=eq.sym + helicity=(0, eq.NFP), + B_lm=np.array([0.8, 1.2]), ) - obj = ObjectiveFunction(MercierStability(eq=eq, grid=grid)) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq)) - np.testing.assert_allclose(f, f[-1], rtol=1e-2) - - # TODO: add test for PlasmaVesselDistance - - # Omnigenity - # (needs an approximately QP equilibrium and field) - surf = FourierRZToroidalSurface.from_qp_model( - major_radius=1, - aspect_ratio=20, - elongation=6, - mirror_ratio=0.2, - torsion=0.1, - NFP=1, - sym=True, - ) - eq = Equilibrium(Psi=6e-3, M=4, N=4, surface=surf) - eq, _ = eq.solve(objective="force", verbose=3) - field = OmnigenousField( - L_B=0, - M_B=2, - L_x=0, - M_x=0, - N_x=0, - NFP=eq.NFP, - helicity=(0, eq.NFP), - B_lm=np.array([0.8, 1.2]), + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array + 0.5): # omnigenity needs higher res + grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP) + obj = ObjectiveFunction( + Omnigenity(eq=eq, field=field, eq_grid=grid, field_grid=grid) + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x(eq, field)) + np.testing.assert_allclose(f, f[-1], rtol=1e-3) + + @pytest.mark.regression + @pytest.mark.parametrize( + "objective", sorted(other_objectives, key=lambda x: str(x.__name__)) ) - f = np.zeros_like(res_array, dtype=float) - for i, res in enumerate(res_array + 0.5): # omnigenity needs higher resolution - grid = LinearGrid(M=int(eq.M * res), N=int(eq.N * res), NFP=eq.NFP) - obj = ObjectiveFunction( - Omnigenity(eq=eq, field=field, eq_grid=grid, field_grid=grid) - ) - obj.build(verbose=0) - f[i] = obj.compute_scalar(obj.x(eq, field)) - np.testing.assert_allclose(f, f[-1], rtol=1e-3) + def test_compute_scalar_resolution_others(self, objective): + """All other objectives.""" + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + # just change eq resolution and let objective pick the right grid type + self.eq.change_resolution( + L_grid=int(self.eq.L * res), + M_grid=int(self.eq.M * res), + N_grid=int(self.eq.N * res), + ) + obj = ObjectiveFunction(objective(eq=self.eq), use_jit=False) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + + @pytest.mark.regression + @pytest.mark.parametrize("objective", [CoilLength, CoilTorsion, CoilCurvature]) + def test_compute_scalar_resolution_coils(self, objective): + """Coil objectives.""" + coil = FourierXYZCoil() + coilset = CoilSet.linspaced_angular(coil) + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + obj = ObjectiveFunction( + objective(coilset, grid=LinearGrid(N=int(5 + 3 * res))), use_jit=False + ) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=1e-2, atol=1e-12) -@pytest.mark.unit -def test_objective_no_nangrad(): + +class TestObjectiveNaNGrad: """Make sure reverse mode AD works correctly for all objectives.""" - # these need some special logic - eq = Equilibrium(L=2, M=2, N=2) - surf = FourierRZToroidalSurface() - obj = ObjectiveFunction(PlasmaVesselDistance(eq, surf)) - obj.build() - g = obj.grad(obj.x(eq, surf)) - assert not np.any(np.isnan(g)), "plasma vessel distance" - eq = Equilibrium(L=2, M=2, N=2, anisotropy=FourierZernikeProfile()) - obj = ObjectiveFunction(ForceBalanceAnisotropic(eq)) - obj.build() - g = obj.grad(obj.x(eq)) - assert not np.any(np.isnan(g)), "anisotropic" + # get a list of all the objectives + objectives = [ + getattr(desc.objectives, obj) + for obj in dir(desc.objectives) + if obj[0].isupper() + and (not obj.startswith("Fix")) + and (obj != "ObjectiveFunction") + and ("SelfConsistency" not in obj) + ] + specials = [ + # these require special logic + PlasmaVesselDistance, + ForceBalanceAnisotropic, + BootstrapRedlConsistency, + BoundaryError, + VacuumBoundaryError, + CoilLength, + CoilCurvature, + CoilTorsion, + # we don't test these since they depend too much on what exactly the user wants + GenericObjective, + LinearObjectiveFromUser, + ObjectiveFromUser, + # TODO: add Omnigenity objective (see GH issue #943) + Omnigenity, + ] + other_objectives = list(set(objectives) - set(specials)) - eq = Equilibrium( - L=2, - M=2, - N=2, - electron_density=PowerSeriesProfile([1e19, 0, -1e19]), - electron_temperature=PowerSeriesProfile([1e3, 0, -1e3]), - current=PowerSeriesProfile([1, 0, -1]), - ) - obj = ObjectiveFunction(BootstrapRedlConsistency(eq)) - obj.build() - g = obj.grad(obj.x(eq)) - assert not np.any(np.isnan(g)), "redl bootstrap" + @pytest.mark.unit + def test_objective_no_nangrad_plasma_vessel(self): + """PlasmaVesselDistance.""" + eq = Equilibrium(L=2, M=2, N=2) + surf = FourierRZToroidalSurface() + obj = ObjectiveFunction(PlasmaVesselDistance(eq, surf), use_jit=False) + obj.build() + g = obj.grad(obj.x(eq, surf)) + assert not np.any(np.isnan(g)), "plasma vessel distance" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + @pytest.mark.unit + def test_objective_no_nangrad_anisotropy(self): + """ForceBalanceAnisotropic.""" + eq = Equilibrium(L=2, M=2, N=2, anisotropy=FourierZernikeProfile()) + obj = ObjectiveFunction(ForceBalanceAnisotropic(eq), use_jit=False) + obj.build() + g = obj.grad(obj.x(eq)) + assert not np.any(np.isnan(g)), "anisotropic" - pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) - iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) - surf = FourierRZToroidalSurface( - R_lmn=[4.0, 1.0], - modes_R=[[0, 0], [1, 0]], - Z_lmn=[-1.0], - modes_Z=[[-1, 0]], - NFP=1, - ) + @pytest.mark.unit + def test_objective_no_nangrad_bootstrap(self): + """BootstrapRedlConsistency.""" + eq = Equilibrium( + L=2, + M=2, + N=2, + electron_density=PowerSeriesProfile([1e19, 0, -1e19]), + electron_temperature=PowerSeriesProfile([1e3, 0, -1e3]), + current=PowerSeriesProfile([1, 0, -1]), + ) + obj = ObjectiveFunction(BootstrapRedlConsistency(eq), use_jit=False) + obj.build() + g = obj.grad(obj.x(eq)) + assert not np.any(np.isnan(g)), "redl bootstrap" - eq = Equilibrium(M=10, N=0, Psi=1.0, surface=surf, pressure=pres, iota=iota) - obj = ObjectiveFunction(BoundaryError(eq, ext_field)) - obj.build() - g = obj.grad(obj.x()) - assert not np.any(np.isnan(g)), "boundary error" + @pytest.mark.unit + def test_objective_no_nangrad_boundary_error(self): + """BoundaryError.""" + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + + pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) + iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) + surf = FourierRZToroidalSurface( + R_lmn=[4.0, 1.0], + modes_R=[[0, 0], [1, 0]], + Z_lmn=[-1.0], + modes_Z=[[-1, 0]], + NFP=1, + ) - obj = ObjectiveFunction(VacuumBoundaryError(eq, ext_field)) - with pytest.warns(UserWarning): + eq = Equilibrium(M=6, N=0, Psi=1.0, surface=surf, pressure=pres, iota=iota) + obj = ObjectiveFunction(BoundaryError(eq, ext_field), use_jit=False) obj.build() - g = obj.grad(obj.x()) - assert not np.any(np.isnan(g)), "vacuum boundary error" + g = obj.grad(obj.x(eq, ext_field)) + assert not np.any(np.isnan(g)), "boundary error" - # TODO: add Omnigenity objective (see GH issue #943) + @pytest.mark.unit + def test_objective_no_nangrad_vacuum_boundary_error(self): + """VacuumBoundaryError.""" + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + + pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) + iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) + surf = FourierRZToroidalSurface( + R_lmn=[4.0, 1.0], + modes_R=[[0, 0], [1, 0]], + Z_lmn=[-1.0], + modes_Z=[[-1, 0]], + NFP=1, + ) - # these only need basic equilibrium - eq = Equilibrium(L=2, M=2, N=2) - objectives = [ - CurrentDensity, - Energy, - ForceBalance, - HelicalForceBalance, - RadialForceBalance, - AspectRatio, - BScaleLength, - Elongation, - GoodCoordinates, - MeanCurvature, - PrincipalCurvature, - Volume, - Pressure, - RotationalTransform, - Shear, - ToroidalCurrent, - Isodynamicity, - QuasisymmetryBoozer, - QuasisymmetryTripleProduct, - QuasisymmetryTwoTerm, - MagneticWell, - MercierStability, - ] + eq = Equilibrium(M=6, N=0, Psi=1.0, surface=surf, pressure=pres, iota=iota) - for objective in objectives: - obj = ObjectiveFunction(objective(eq)) + obj = ObjectiveFunction(VacuumBoundaryError(eq, ext_field), use_jit=False) + with pytest.warns(UserWarning): + obj.build() + g = obj.grad(obj.x(eq, ext_field)) + assert not np.any(np.isnan(g)), "vacuum boundary error" + + @pytest.mark.unit + @pytest.mark.parametrize( + "objective", sorted(other_objectives, key=lambda x: str(x.__name__)) + ) + def test_objective_no_nangrad(self, objective): + """Generic test for other objectives.""" + eq = Equilibrium(L=2, M=2, N=2) + obj = ObjectiveFunction(objective(eq), use_jit=False) obj.build() g = obj.grad(obj.x(eq)) assert not np.any(np.isnan(g)), str(objective) + @pytest.mark.unit + @pytest.mark.parametrize("objective", [CoilLength, CoilTorsion, CoilCurvature]) + def test_objective_no_nangrad_coils(self, objective): + """Coil objectives.""" + coil = FourierXYZCoil() + coilset = CoilSet.linspaced_angular(coil) + obj = ObjectiveFunction(objective(coilset), use_jit=False) + obj.build(verbose=0) + g = obj.grad(obj.x()) + assert not np.any(np.isnan(g)), str(objective) + @pytest.mark.unit def test_asymmetric_normalization(): @@ -2142,13 +2036,3 @@ def test_asymmetric_normalization(): assert np.all(np.isfinite(val)) for val in scales_eq.values(): assert np.all(np.isfinite(val)) - - -@pytest.mark.unit -def test_force_balance_axis_error(): - """Test that ForceBalance objective is not NaN if the grid contains axis.""" - eq = get("SOLOVEV") - grid = LinearGrid(L=2, M=2, N=2, axis=True) - obj = ForceBalance(eq, grid=grid) - obj.build() - assert not np.any(np.isnan(obj.compute_unscaled(*obj.xs(eq)))) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index c066b642f2..64f71ee2b3 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -319,7 +319,7 @@ def test_no_iterations(): np.testing.assert_allclose(x0, out2["x"]) -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.slow @pytest.mark.optimize def test_overstepping(): @@ -351,6 +351,7 @@ def compute(self, params, constants=None): return x eq = desc.examples.get("DSHAPE") + eq.change_resolution(2, 2, 0, 4, 4, 0) np.random.seed(0) objective = ObjectiveFunction(DummyObjective(things=eq), use_jit=False) @@ -384,7 +385,7 @@ def compute(self, params, constants=None): objective=objective, constraints=constraints, optimizer=optimizer, - maxiter=50, + maxiter=5, verbose=3, gtol=-1, # disable gradient stopping ftol=-1, # disable function stopping @@ -392,8 +393,8 @@ def compute(self, params, constants=None): copy=True, options={ "initial_trust_radius": 0.5, - "perturb_options": {"verbose": 0}, - "solve_options": {"verbose": 0}, + "perturb_options": {"verbose": 0, "order": 1}, + "solve_options": {"verbose": 0, "maxiter": 2}, }, ) @@ -405,13 +406,14 @@ def compute(self, params, constants=None): np.testing.assert_allclose(x0, x1, rtol=1e-14, atol=1e-14) -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.slow @pytest.mark.solve def test_maxiter_1_and_0_solve(): """Test that solves with maxiter 1 and 0 terminate correctly.""" # correctly meaning they terminate, instead of looping infinitely eq = desc.examples.get("SOLOVEV") + eq.change_resolution(2, 2, 0, 4, 4, 0) constraints = ( FixBoundaryR(eq=eq), FixBoundaryZ(eq=eq), @@ -420,7 +422,7 @@ def test_maxiter_1_and_0_solve(): FixPsi(eq=eq), ) objectives = ForceBalance(eq=eq) - obj = ObjectiveFunction(objectives) + obj = ObjectiveFunction(objectives, use_jit=False) for opt in ["lsq-exact", "fmintr-bfgs"]: eq, result = eq.solve( maxiter=1, constraints=constraints, objective=obj, optimizer=opt, verbose=3 @@ -433,7 +435,7 @@ def test_maxiter_1_and_0_solve(): assert result["nit"] == 0 -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.slow @pytest.mark.solve def test_scipy_fail_message(): @@ -447,7 +449,7 @@ def test_scipy_fail_message(): FixPsi(eq=eq), ) objectives = ForceBalance(eq=eq) - obj = ObjectiveFunction(objectives) + obj = ObjectiveFunction(objectives, use_jit=False) # should fail on maxiter, and should NOT throw an error for opt in ["scipy-trf"]: @@ -464,7 +466,7 @@ def test_scipy_fail_message(): assert "Maximum number of iterations has been exceeded" in result["message"] eq.set_initial_guess() objectives = Energy(eq=eq) - obj = ObjectiveFunction(objectives) + obj = ObjectiveFunction(objectives, use_jit=False) for opt in ["scipy-trust-exact"]: eq, result = eq.solve( maxiter=3, @@ -490,6 +492,7 @@ def test_not_implemented_error(): def test_wrappers(): """Tests for using wrapped objectives.""" eq = desc.examples.get("SOLOVEV") + eq.change_resolution(2, 2, 0, 4, 4, 0) con = ( FixBoundaryR(eq=eq), FixBoundaryZ(eq=eq), @@ -559,38 +562,50 @@ def test_wrappers(): np.testing.assert_allclose(ob.weights, con[0].weight) -@pytest.mark.unit -@pytest.mark.slow -def test_all_optimizers(): - """Just tests that the optimizers run without error, eg tests for the wrappers.""" +class TestAllOptimizers: + """Tests all optimizers run without error, eg tests for wrappers.""" + eqf = desc.examples.get("SOLOVEV") + eqf.change_resolution(3, 3, 0, 6, 6, 0) eqe = eqf.copy() fobj = ObjectiveFunction(ForceBalance(eq=eqf)) eobj = ObjectiveFunction(Energy(eq=eqe)) - fobj.build() - eobj.build() econ = get_fixed_boundary_constraints(eq=eqe) fcon = get_fixed_boundary_constraints(eq=eqf) - options = {"sgd": {"alpha": 1e-4}} - - for opt in optimizers: - print("TESTING ", opt) - if optimizers[opt]["scalar"]: - obj = eobj - eq = eqe - con = econ - else: - obj = fobj - eq = eqf - con = fcon - eq.solve( - objective=obj, - constraints=con, + + scalar_methods = [opt for opt in optimizers if optimizers[opt]["scalar"]] + lsq_methods = [opt for opt in optimizers if not optimizers[opt]["scalar"]] + + @pytest.mark.unit + @pytest.mark.parametrize("opt", scalar_methods) + def test_all_optimizers_scalar(self, opt): + """Test all scalar methods.""" + if not self.eobj.built: + self.eobj.build() + + self.eqe.solve( + objective=self.eobj, + constraints=self.econ, + optimizer=opt, + copy=True, + verbose=3, + maxiter=5, + ) + + @pytest.mark.unit + @pytest.mark.parametrize("opt", lsq_methods) + def test_all_optimizers_lsq(self, opt): + """Test all least squares methods.""" + if not self.fobj.built: + self.fobj.build() + + self.eqf.solve( + objective=self.fobj, + constraints=self.fcon, optimizer=opt, copy=True, verbose=3, maxiter=5, - options=options.get(opt, None), ) @@ -661,7 +676,7 @@ def test_scipy_constrained_solve(): assert eq2.is_nested() -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.solve def test_solve_with_x_scale(): """Make sure we can manually specify x_scale when solving/optimizing.""" @@ -970,60 +985,6 @@ def test_constrained_AL_scalar(): np.testing.assert_array_less(-Dwell, ctol) -@pytest.mark.slow -@pytest.mark.unit -@pytest.mark.optimize -def test_proximal_with_PlasmaVesselDistance(): - """Tests that the proximal projection works with fixed surface distance obj.""" - eq = desc.examples.get("SOLOVEV") - - constraints = ( - ForceBalance(eq=eq), - FixPressure(eq=eq), # fix pressure profile - FixIota(eq=eq), # fix rotational transform profile - FixPsi(eq=eq), # fix total toroidal magnetic flux - ) - # circular surface - a = 2 - R0 = 4 - surf = FourierRZToroidalSurface( - R_lmn=[R0, a], - Z_lmn=[0.0, -a], - modes_R=np.array([[0, 0], [1, 0]]), - modes_Z=np.array([[0, 0], [-1, 0]]), - sym=True, - NFP=eq.NFP, - ) - - grid = LinearGrid(M=eq.M, N=0, NFP=eq.NFP) - obj = PlasmaVesselDistance( - surface=surf, eq=eq, target=0.5, plasma_grid=grid, surface_fixed=True - ) - objective = ObjectiveFunction((obj,)) - - optimizer = Optimizer("proximal-lsq-exact") - eq, result = optimizer.optimize( - eq, - objective, - constraints, - verbose=3, - maxiter=3, - ) - - # make sure it also works if proximal is given multiple objects in things - obj = PlasmaVesselDistance( - surface=surf, eq=eq, target=0.5, plasma_grid=grid, surface_fixed=False - ) - objective = ObjectiveFunction((obj,)) - (eq, surf), result = optimizer.optimize( - (eq, surf), - objective, - constraints, - verbose=3, - maxiter=3, - ) - - @pytest.mark.unit @pytest.mark.optimize def test_optimize_multiple_things_different_order(): @@ -1125,16 +1086,19 @@ def test_optimize_with_single_constraint(): """Tests that Optimizer.optimize prints afterwards with a single constraint.""" eq = Equilibrium() optimizer = Optimizer("lsq-exact") - objectective = ObjectiveFunction(GenericObjective("|B|", eq)) + objectective = ObjectiveFunction(GenericObjective("|B|", eq), use_jit=False) constraints = FixParameter( # Psi is not constrained eq, ["R_lmn", "Z_lmn", "L_lmn", "Rb_lmn", "Zb_lmn", "p_l", "c_l"] ) # test depends on verbose > 0 - optimizer.optimize(eq, objective=objectective, constraints=constraints, verbose=2) + optimizer.optimize( + eq, objective=objectective, constraints=constraints, verbose=2, maxiter=1 + ) -@pytest.mark.unit +@pytest.mark.slow +@pytest.mark.regression def test_proximal_jacobian(): """Test that JVPs and manual concatenation give the same result as full jac.""" eq = desc.examples.get("HELIOTRON") @@ -1142,9 +1106,9 @@ def test_proximal_jacobian(): eq1 = eq.copy() eq2 = eq.copy() eq3 = eq.copy() - con1 = ObjectiveFunction(ForceBalance(eq1)) - con2 = ObjectiveFunction(ForceBalance(eq2)) - con3 = ObjectiveFunction(ForceBalance(eq3)) + con1 = ObjectiveFunction(ForceBalance(eq1), use_jit=False) + con2 = ObjectiveFunction(ForceBalance(eq2), use_jit=False) + con3 = ObjectiveFunction(ForceBalance(eq3), use_jit=False) obj1 = ObjectiveFunction( ( QuasisymmetryTripleProduct(eq1, deriv_mode="fwd"), @@ -1152,6 +1116,7 @@ def test_proximal_jacobian(): Volume(eq1, deriv_mode="fwd"), ), deriv_mode="batched", + use_jit=False, ) obj2 = ObjectiveFunction( ( @@ -1160,6 +1125,7 @@ def test_proximal_jacobian(): Volume(eq2, deriv_mode="fwd"), ), deriv_mode="looped", + use_jit=False, ) obj3 = ObjectiveFunction( ( @@ -1168,10 +1134,13 @@ def test_proximal_jacobian(): Volume(eq3, deriv_mode="rev"), ), deriv_mode="blocked", + use_jit=False, ) - prox1 = ProximalProjection(obj1, con1, eq1) - prox2 = ProximalProjection(obj2, con2, eq2) - prox3 = ProximalProjection(obj3, con3, eq3) + perturb_options = {"order": 1} + solve_options = {"maxiter": 1} + prox1 = ProximalProjection(obj1, con1, eq1, perturb_options, solve_options) + prox2 = ProximalProjection(obj2, con2, eq2, perturb_options, solve_options) + prox3 = ProximalProjection(obj3, con3, eq3, perturb_options, solve_options) prox1.build() prox2.build() prox3.build() @@ -1244,7 +1213,8 @@ def test_proximal_jacobian(): np.testing.assert_allclose(jac_unscaled, jac3, rtol=1e-12, atol=1e-12) -@pytest.mark.unit +@pytest.mark.slow +@pytest.mark.regression def test_LinearConstraint_jacobian(): """Test that JVPs and manual concatenation give the same result as full jac.""" eq = desc.examples.get("HELIOTRON") @@ -1253,9 +1223,15 @@ def test_LinearConstraint_jacobian(): eq2 = eq.copy() eq3 = eq.copy() - obj1 = ObjectiveFunction(ForceBalance(eq1, deriv_mode="auto"), deriv_mode="batched") - obj2 = ObjectiveFunction(ForceBalance(eq2, deriv_mode="fwd"), deriv_mode="looped") - obj3 = ObjectiveFunction(ForceBalance(eq3, deriv_mode="rev"), deriv_mode="blocked") + obj1 = ObjectiveFunction( + ForceBalance(eq1, deriv_mode="auto"), deriv_mode="batched", use_jit=False + ) + obj2 = ObjectiveFunction( + ForceBalance(eq2, deriv_mode="fwd"), deriv_mode="looped", use_jit=False + ) + obj3 = ObjectiveFunction( + ForceBalance(eq3, deriv_mode="rev"), deriv_mode="blocked", use_jit=False + ) con1 = ObjectiveFunction(get_fixed_boundary_constraints(eq1)) con2 = ObjectiveFunction(get_fixed_boundary_constraints(eq2)) diff --git a/tests/test_perturbations.py b/tests/test_perturbations.py index aef2f82388..822a3bde29 100644 --- a/tests/test_perturbations.py +++ b/tests/test_perturbations.py @@ -4,7 +4,7 @@ import pytest import desc.examples -from desc.equilibrium import EquilibriaFamily, Equilibrium +from desc.equilibrium import Equilibrium from desc.geometry import FourierRZCurve from desc.grid import ConcentricGrid, QuadratureGrid from desc.objectives import ( @@ -19,10 +19,9 @@ @pytest.mark.regression @pytest.mark.slow -@pytest.mark.solve -def test_perturbation_orders(SOLOVEV): +def test_perturbation_orders(): """Test that higher-order perturbations are more accurate.""" - eq = EquilibriaFamily.load(load_from=str(SOLOVEV["desc_h5_path"]))[-1] + eq = desc.examples.get("SOLOVEV") objective = get_equilibrium_objective(eq=eq) constraints = get_fixed_boundary_constraints(eq=eq) @@ -129,13 +128,16 @@ def test_optimal_perturb(): # particular solution. Here we do a simple test to ensure the interior and boundary # agree eq1 = desc.examples.get("DSHAPE") + eq1.change_resolution(3, 3, 0, 6, 6, 0) eq1.change_resolution(N=1, N_grid=5) + eq1.surface = eq1.get_surface_at(1.0) objective = ObjectiveFunction( ToroidalCurrent( eq=eq1, grid=QuadratureGrid(eq1.L, eq1.M, eq1.N), target=0, weight=1 - ) + ), + use_jit=False, ) - constraint = ObjectiveFunction(ForceBalance(eq=eq1, target=0)) + constraint = ObjectiveFunction(ForceBalance(eq=eq1, target=0), use_jit=False) objective.build() constraint.build() diff --git a/tests/test_plotting.py b/tests/test_plotting.py index a10a1280c8..aeb784f090 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -1,5 +1,7 @@ """Regression tests for plotting functions, by comparing to saved baseline images.""" +import warnings + import matplotlib.pyplot as plt import numpy as np import pytest @@ -20,7 +22,6 @@ from desc.io import load from desc.magnetic_fields import OmnigenousField, ToroidalMagneticField from desc.plotting import ( - _find_idx, plot_1d, plot_2d, plot_3d, @@ -32,7 +33,6 @@ plot_coefficients, plot_coils, plot_comparison, - plot_field_lines_sfl, plot_fsa, plot_grid, plot_logo, @@ -60,581 +60,489 @@ def test_kwarg_warning(DummyStellarator): def test_kwarg_future_warning(DummyStellarator): """Test that passing in deprecated kwargs throws a warning.""" eq = load(load_from=str(DummyStellarator["output_path"])) - with pytest.warns(FutureWarning): - fig, ax = plot_surfaces(eq, zeta=2) + with pytest.raises(FutureWarning): + with warnings.catch_warnings(): + warnings.simplefilter("error") + fig, ax = plot_surfaces(eq, zeta=2) return None -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_1d_p(): - """Test plotting 1d pressure profile.""" - eq = get("SOLOVEV") - fig, ax, data = plot_1d(eq, "p", figsize=(4, 4), return_data=True) - assert "p" in data.keys() - return fig +class TestPlot1D: + """Tests for plot_1d.""" + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_1d_p(self): + """Test plotting 1d pressure profile.""" + eq = get("SOLOVEV") + fig, ax, data = plot_1d(eq, "p", figsize=(4, 4), return_data=True) + assert "p" in data.keys() + return fig -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_1d_elongation(): - """Test plotting 1d elongation as a function of toroidal angle.""" - eq = get("precise_QA") - grid = LinearGrid(N=20, NFP=eq.NFP) - fig, ax, data = plot_1d( - eq, "a_major/a_minor", grid=grid, figsize=(4, 4), return_data=True - ) - assert "a_major/a_minor" in data.keys() - return fig - - -@pytest.mark.unit -def test_1d_fsa_consistency(): - """Test that plot_1d uses 2d grid to compute quantities with surface averages.""" - eq = get("W7-X") - - def test(name, with_sqrt_g=True, grid=None): - _, ax_0 = plot_1d(eq, name, grid=grid) - # 100 rho points is plot_1d default - _, ax_1 = plot_fsa( - eq, - name, - with_sqrt_g=with_sqrt_g, - rho=100 if grid is None else grid.nodes[:, 0], - ) - np.testing.assert_allclose( - ax_0.lines[0].get_xydata(), ax_1.lines[0].get_xydata() + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_1d_elongation(self): + """Test plotting 1d elongation as a function of toroidal angle.""" + eq = get("precise_QA") + grid = LinearGrid(N=20, NFP=eq.NFP) + fig, ax, data = plot_1d( + eq, "a_major/a_minor", grid=grid, figsize=(4, 4), return_data=True ) + assert "a_major/a_minor" in data.keys() + return fig - lg = LinearGrid(rho=np.linspace(0, 1, 30)) - test("magnetic well") - test("magnetic well", grid=lg) - test("current", with_sqrt_g=False) - test("current", with_sqrt_g=False, grid=lg) - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_1d_dpdr(): - """Test plotting 1d pressure derivative.""" - eq = get("DSHAPE_current") - fig, ax, data = plot_1d(eq, "p_r", figsize=(4, 4), return_data=True) - assert "p_r" in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_1d_iota(): - """Test plotting 1d rotational transform.""" - eq = get("DSHAPE_current") - grid = LinearGrid(rho=0.5, theta=100, zeta=0.0) - fig, ax, data = plot_1d(eq, "iota", grid=grid, figsize=(4, 4), return_data=True) - assert "theta" in data.keys() - assert "iota" in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_1d_iota_radial(): - """Test plotting 1d rotational transform.""" - eq = get("DSHAPE_current") - fig, ax, data = plot_1d(eq, "iota", figsize=(4, 4), return_data=True) - assert "rho" in data.keys() - assert "iota" in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_1d_logpsi(): - """Test plotting 1d flux function with log scale.""" - eq = get("DSHAPE_current") - fig, ax, data = plot_1d(eq, "psi", log=True, figsize=(4, 4), return_data=True) - ax.set_ylim([1e-5, 1e0]) - assert "rho" in data.keys() - assert "psi" in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=10) -def test_2d_logF(DSHAPE_current): - """Test plotting 2d force error with log scale.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=100, theta=100, zeta=0.0) - fig, ax, data = plot_2d( - eq, "|F|", log=True, grid=grid, figsize=(4, 4), return_data=True - ) - assert "|F|" in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_2d_g_tz(DSHAPE_current): - """Test plotting 2d metric coefficients vs theta/zeta.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=0.5, theta=100, zeta=100) - fig, ax, data = plot_2d(eq, "sqrt(g)", grid=grid, figsize=(4, 4), return_data=True) - assert "theta" in data.keys() - assert "zeta" in data.keys() - - assert "sqrt(g)" in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_2d_g_rz(DSHAPE_current): - """Test plotting 2d metric coefficients vs rho/zeta.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=100, theta=0.0, zeta=100) - fig, ax, data = plot_2d(eq, "sqrt(g)", grid=grid, figsize=(4, 4), return_data=True) - assert "rho" in data.keys() - assert "zeta" in data.keys() - assert "sqrt(g)" in data.keys() - - return fig - - -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_2d_lambda(DSHAPE_current): - """Test plotting lambda on 2d grid.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d(eq, "lambda", figsize=(4, 4), return_data=True) - assert "lambda" in data.keys() - - return fig - - -@pytest.mark.unit -@pytest.mark.solve -def test_3d_B(DSHAPE_current): - """Test 3d plot of toroidal field.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, data = plot_3d(eq, "B^zeta", return_data=True) - assert "X" in data.keys() - assert "Y" in data.keys() - assert "Z" in data.keys() - - assert "B^zeta" in data.keys() - - return fig - - -@pytest.mark.unit -@pytest.mark.solve -def test_3d_J(DSHAPE_current): - """Test 3d plotting of poloidal current.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=1.0, theta=100, zeta=100) - fig = plot_3d(eq, "J^theta", grid=grid) - return fig - - -@pytest.mark.unit -@pytest.mark.solve -def test_3d_tz(DSHAPE_current): - """Test 3d plot of force on interior surface.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=0.5, theta=100, zeta=100) - fig = plot_3d(eq, "|F|", log=True, grid=grid) - return fig - - -@pytest.mark.unit -@pytest.mark.solve -def test_3d_rz(DSHAPE_current): - """Test 3d plotting of pressure on toroidal cross section.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=100, theta=0.0, zeta=100) - fig = plot_3d(eq, "p", grid=grid) - return fig - - -@pytest.mark.unit -@pytest.mark.solve -def test_3d_rt(DSHAPE_current): - """Test 3d plotting of flux on poloidal ribbon.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - grid = LinearGrid(rho=100, theta=100, zeta=0.0) - fig = plot_3d(eq, "psi", grid=grid) - return fig - + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_1d_iota(self): + """Test plotting 1d rotational transform.""" + eq = get("DSHAPE_current") + grid = LinearGrid(rho=0.5, theta=100, zeta=0.0) + fig, ax, data = plot_1d(eq, "iota", grid=grid, figsize=(4, 4), return_data=True) + assert "theta" in data.keys() + assert "iota" in data.keys() + return fig -@pytest.mark.unit -def test_plot_fsa_axis_limit(): - """Test magnetic axis limit of flux surface average is plotted.""" - eq = get("W7-X") - rho = np.linspace(0, 1, 10) - grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=rho) - assert grid.axis.size - - name = "J*B" - assert ( - "<" + name + ">" in data_index["desc.equilibrium.equilibrium.Equilibrium"] - ), "Test with a different quantity." - # should forward computation to compute function - _, _, plot_data = plot_fsa( - eq=eq, - name=name, - rho=rho, - M=eq.M_grid, - N=eq.N_grid, - with_sqrt_g=True, - return_data=True, - ) - desired = grid.compress( - eq.compute(names="<" + name + ">", grid=grid)["<" + name + ">"] - ) - np.testing.assert_allclose(plot_data["<" + name + ">"], desired, equal_nan=False) - - name = "B0" - assert ( - "<" + name + ">" not in data_index["desc.equilibrium.equilibrium.Equilibrium"] - ), "Test with a different quantity." - # should automatically compute axis limit - _, _, plot_data = plot_fsa( - eq=eq, - name=name, - rho=rho, - M=eq.M_grid, - N=eq.N_grid, - with_sqrt_g=True, - return_data=True, - ) - data = eq.compute(names=[name, "sqrt(g)", "sqrt(g)_r"], grid=grid) - desired = surface_averages( - grid=grid, - q=data[name], - sqrt_g=grid.replace_at_axis(data["sqrt(g)"], data["sqrt(g)_r"], copy=True), - expand_out=False, - ) - np.testing.assert_allclose( - plot_data["<" + name + ">_fsa"], desired, equal_nan=False - ) + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_1d_iota_radial(self): + """Test plotting 1d rotational transform.""" + eq = get("DSHAPE_current") + fig, ax, data = plot_1d(eq, "iota", figsize=(4, 4), return_data=True) + assert "rho" in data.keys() + assert "iota" in data.keys() + return fig - name = "|B|" - assert ( - "<" + name + ">" in data_index["desc.equilibrium.equilibrium.Equilibrium"] - ), "Test with a different quantity." - _, _, plot_data = plot_fsa( - eq=eq, - name=name, - rho=rho, - M=eq.M_grid, - N=eq.N_grid, - with_sqrt_g=False, # Test that does not compute data_index["<|B|>"] - return_data=True, - ) - data = eq.compute(names=name, grid=grid) - desired = surface_averages(grid=grid, q=data[name], expand_out=False) - np.testing.assert_allclose( - plot_data["<" + name + ">_fsa"], desired, equal_nan=False - ) + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_1d_logpsi(self): + """Test plotting 1d flux function with log scale.""" + eq = get("DSHAPE_current") + fig, ax, data = plot_1d(eq, "psi", log=True, figsize=(4, 4), return_data=True) + ax.set_ylim([1e-5, 1e0]) + assert "rho" in data.keys() + assert "psi" in data.keys() + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_1d_curve(self): + """Test plot_1d function for Curve objects.""" + curve = FourierXYZCurve([0, 10, 1]) + fig, ax = plot_1d(curve, "curvature") + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_fsa_I(DSHAPE_current): - """Test plotting of flux surface average toroidal current.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_fsa(eq, "B_theta", with_sqrt_g=False, return_data=True) - assert "rho" in data.keys() - assert "_fsa" in data.keys() - assert "normalization" in data.keys() - assert data["normalization"] == 1 + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_1d_surface(self): + """Test plot_1d function for Surface objects.""" + surf = FourierRZToroidalSurface() + fig, ax = plot_1d(surf, "curvature_H_rho", grid=LinearGrid(M=50)) + return fig - return fig +class TestPlot2D: + """Tests for plot_2d.""" -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_fsa_G(DSHAPE_current): - """Test plotting of flux surface average poloidal current.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_fsa(eq, "B_zeta", with_sqrt_g=False, log=True) - ax.set_ylim([1e-1, 1e0]) - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=10) + def test_2d_logF(self): + """Test plotting 2d force error with log scale.""" + eq = get("DSHAPE_CURRENT") + grid = LinearGrid(rho=100, theta=100, zeta=0.0) + fig, ax, data = plot_2d( + eq, "|F|", log=True, grid=grid, figsize=(4, 4), return_data=True + ) + assert "|F|" in data.keys() + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_2d_g_tz(self): + """Test plotting 2d metric coefficients vs theta/zeta.""" + eq = get("DSHAPE_CURRENT") + grid = LinearGrid(rho=0.5, theta=100, zeta=100) + fig, ax, data = plot_2d( + eq, "sqrt(g)", grid=grid, figsize=(4, 4), return_data=True + ) + assert "theta" in data.keys() + assert "zeta" in data.keys() -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_fsa_F_normalized(DSHAPE_current): - """Test plotting flux surface average normalized force error on log scale.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_fsa(eq, "|F|", log=True, norm_F=True, norm_name="<|grad(p)|>_vol") - ax.set_ylim([1e-6, 1e-3]) - return fig + assert "sqrt(g)" in data.keys() + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_2d_g_rz(self): + """Test plotting 2d metric coefficients vs rho/zeta.""" + eq = get("DSHAPE_CURRENT") + grid = LinearGrid(rho=100, theta=0.0, zeta=100) + fig, ax, data = plot_2d( + eq, "sqrt(g)", grid=grid, figsize=(4, 4), return_data=True + ) + assert "rho" in data.keys() + assert "zeta" in data.keys() + assert "sqrt(g)" in data.keys() -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_section_J(DSHAPE_current): - """Test plotting poincare section of radial current.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_section(eq, "J^rho", return_data=True) - assert "R" in data.keys() - assert "Z" in data.keys() - assert "J^rho" in data.keys() - assert "normalization" in data.keys() - assert data["normalization"] == 1 + return fig - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_con_basis(self): + """Test 2d plot of R component of e^rho.""" + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_2d( + eq, "e^rho", component="R", figsize=(4, 4), return_data=True + ) + for string in ["e^rho", "normalization", "theta", "zeta"]: + assert string in data.keys() + assert data["normalization"] == 1 + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=24) -def test_section_Z(DSHAPE_current): - """Test plotting poincare section of Z coordinate.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_section(eq, "Z") - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_cov_basis(self): + """Test 2d plot of norm of e_rho.""" + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_2d(eq, "e_rho", figsize=(4, 4), return_data=True) + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_normF_2d(self): + """Test 2d plot of normalized force.""" + grid = LinearGrid(rho=np.array(0.8), M=20, N=2) + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_2d( + eq, "|F|", norm_F=True, figsize=(4, 4), return_data=True, grid=grid + ) + for string in ["|F|", "theta", "zeta"]: + assert string in data.keys() + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_section_R(DSHAPE_current): - """Test plotting poincare section of R coordinate.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_section(eq, "R") - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_2d_surface(self): + """Test plot_2d function for Surface objects.""" + surf = FourierRZToroidalSurface() + fig, ax = plot_2d(surf, "curvature_H_rho") + return fig + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + @pytest.mark.unit + def test_2d_plot_Bn(self): + """Test 2d plotting of Bn on equilibrium surface.""" + eq = get("HELIOTRON") + fig, _ = plot_2d( + eq, + "B*n", + field=ToroidalMagneticField(1, 1), + field_grid=LinearGrid(M=10, N=10), + grid=LinearGrid(M=30, N=30, NFP=eq.NFP, endpoint=True), + ) + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_section_F(DSHAPE_current): - """Test plotting poincare section of radial force.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_section(eq, "F_rho") - return fig +class TestPlot3D: + """Tests for plot_3d.""" -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_section_F_normalized_vac(DSHAPE_current): - """Test plotting poincare section of normalized vacuum force error.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[1] - fig, ax = plot_section(eq, "|F|", norm_F=True) - return fig + @pytest.mark.unit + def test_3d_tz(self): + """Test 3d plot of force on interior surface.""" + eq = get("DSHAPE_CURRENT") + grid = LinearGrid(rho=0.5, theta=100, zeta=100) + fig, data = plot_3d(eq, "|F|", log=True, grid=grid, return_data=True) + assert "X" in data.keys() + assert "Y" in data.keys() + assert "Z" in data.keys() + + assert "|F|" in data.keys() + return fig + @pytest.mark.unit + def test_3d_rz(self): + """Test 3d plotting of pressure on toroidal cross section.""" + eq = get("DSHAPE_CURRENT") + grid = LinearGrid(rho=100, theta=0.0, zeta=100) + fig = plot_3d(eq, "p", grid=grid) + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=50) -def test_section_logF(DSHAPE_current): - """Test plotting poincare section of force magnitude on log scale.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_section(eq, "|F|", log=True) - return fig + @pytest.mark.unit + def test_3d_rt(self): + """Test 3d plotting of flux on poloidal ribbon.""" + eq = get("DSHAPE_CURRENT") + grid = LinearGrid(rho=100, theta=100, zeta=0.0) + fig = plot_3d(eq, "psi", grid=grid) + return fig + @pytest.mark.unit + def test_plot_3d_surface(self): + """Test 3d plotting of surface object.""" + surf = FourierRZToroidalSurface() + fig = plot_3d(surf, "curvature_H_rho") + return fig -@pytest.mark.slow -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_surfaces(DSHAPE_current): - """Test plotting flux surfaces.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_surfaces(eq, return_data=True) - for string in [ - "rho_R_coords", - "rho_Z_coords", - "vartheta_R_coords", - "vartheta_Z_coords", - ]: - assert string in data.keys() + @pytest.mark.unit + def test_3d_plot_Bn(self): + """Test 3d plotting of Bn on equilibrium surface.""" + eq = get("precise_QA") + eq.change_resolution(M=4, N=4, L=4, M_grid=8, N_grid=8, L_grid=8) + fig = plot_3d( + eq, + "B*n", + field=ToroidalMagneticField(1, 1), + grid=LinearGrid(M=30, N=30, NFP=1, endpoint=True), + ) + return fig - return fig +class TestPlotFSA: + """Tests for plot_fsa.""" -@pytest.mark.slow -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_surfaces_no_theta(DSHAPE_current): - """Test plotting flux surfaces without theta contours.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_surfaces(eq, theta=False, return_data=True) - for string in ["rho_R_coords", "rho_Z_coords"]: - assert string in data.keys() + @pytest.mark.unit + def test_plot_fsa_axis_limit(self): + """Test magnetic axis limit of flux surface average is plotted.""" + eq = get("W7-X") + rho = np.linspace(0, 1, 10) + grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=rho) + assert grid.axis.size + + name = "J*B" + assert ( + "<" + name + ">" in data_index["desc.equilibrium.equilibrium.Equilibrium"] + ), "Test with a different quantity." + # should forward computation to compute function + _, _, plot_data = plot_fsa( + eq=eq, + name=name, + rho=rho, + M=eq.M_grid, + N=eq.N_grid, + with_sqrt_g=True, + return_data=True, + ) + desired = grid.compress( + eq.compute(names="<" + name + ">", grid=grid)["<" + name + ">"] + ) + np.testing.assert_allclose( + plot_data["<" + name + ">"], desired, equal_nan=False + ) - return fig + name = "B0" + assert ( + "<" + name + ">" + not in data_index["desc.equilibrium.equilibrium.Equilibrium"] + ), "Test with a different quantity." + # should automatically compute axis limit + _, _, plot_data = plot_fsa( + eq=eq, + name=name, + rho=rho, + M=eq.M_grid, + N=eq.N_grid, + with_sqrt_g=True, + return_data=True, + ) + data = eq.compute(names=[name, "sqrt(g)", "sqrt(g)_r"], grid=grid) + desired = surface_averages( + grid=grid, + q=data[name], + sqrt_g=grid.replace_at_axis(data["sqrt(g)"], data["sqrt(g)_r"], copy=True), + expand_out=False, + ) + np.testing.assert_allclose( + plot_data["<" + name + ">_fsa"], desired, equal_nan=False + ) + name = "|B|" + assert ( + "<" + name + ">" in data_index["desc.equilibrium.equilibrium.Equilibrium"] + ), "Test with a different quantity." + _, _, plot_data = plot_fsa( + eq=eq, + name=name, + rho=rho, + M=eq.M_grid, + N=eq.N_grid, + with_sqrt_g=False, # Test that does not compute data_index["<|B|>"] + return_data=True, + ) + data = eq.compute(names=name, grid=grid) + desired = surface_averages(grid=grid, q=data[name], expand_out=False) + np.testing.assert_allclose( + plot_data["<" + name + ">_fsa"], desired, equal_nan=False + ) -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_boundary(): - """Test plotting boundary.""" - eq = get("W7-X") - fig, ax, data = plot_boundary(eq, plot_axis=True, return_data=True) - assert "R" in data.keys() - assert "Z" in data.keys() + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_fsa_I(self): + """Test plotting of flux surface average toroidal current.""" + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_fsa(eq, "B_theta", with_sqrt_g=False, return_data=True) + assert "rho" in data.keys() + assert "_fsa" in data.keys() + assert "normalization" in data.keys() + assert data["normalization"] == 1 - return fig + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_fsa_F_normalized(self): + """Test plotting flux surface average normalized force error on log scale.""" + eq = get("DSHAPE_CURRENT") + fig, ax = plot_fsa( + eq, "|F|", log=True, norm_F=True, norm_name="<|grad(p)|>_vol" + ) + ax.set_ylim([1e-6, 1e-3]) + return fig -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_boundaries(): - """Test plotting boundaries.""" - eq1 = get("SOLOVEV") - eq2 = get("DSHAPE") - eq3 = get("W7-X") - fig, ax, data = plot_boundaries((eq1, eq2, eq3), return_data=True) - assert "R" in data.keys() - assert "Z" in data.keys() - assert len(data["R"]) == 3 - assert len(data["Z"]) == 3 - return fig +class TestPlotSection: + """Tests for plot_section.""" + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_section_J(self): + """Test plotting poincare section of radial current.""" + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_section(eq, "J^rho", return_data=True) + assert "R" in data.keys() + assert "Z" in data.keys() + assert "J^rho" in data.keys() + assert "normalization" in data.keys() + assert data["normalization"] == 1 -@pytest.mark.slow -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_comparison(DSHAPE_current): - """Test plotting comparison of flux surfaces.""" - eqf = load(load_from=str(DSHAPE_current["desc_h5_path"])) - fig, ax, data = plot_comparison(eqf, return_data=True) - for string in [ - "rho_R_coords", - "rho_Z_coords", - "vartheta_R_coords", - "vartheta_Z_coords", - ]: - assert string in data.keys() - assert len(data[string]) == len(eqf) + return fig - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_section_F(self): + """Test plotting poincare section of radial force.""" + eq = get("DSHAPE_CURRENT") + fig, ax = plot_section(eq, "F_rho") + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=50) + def test_section_logF(self): + """Test plotting poincare section of force magnitude on log scale.""" + eq = get("DSHAPE_CURRENT") + fig, ax = plot_section(eq, "|F|", log=True) + return fig -@pytest.mark.slow -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_comparison_no_theta(DSHAPE_current): - """Test plotting comparison of flux surfaces without theta contours.""" - eqf = load(load_from=str(DSHAPE_current["desc_h5_path"])) - fig, ax = plot_comparison(eqf, theta=0) - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_normF_section(self): + """Test Poincare section plot of normalized force on log scale.""" + eq = get("DSHAPE_CURRENT") + fig, ax = plot_section(eq, "|F|", norm_F=True, log=True) + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_con_basis(DSHAPE_current): - """Test 2d plot of R component of e^rho.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d( - eq, "e^rho", component="R", figsize=(4, 4), return_data=True - ) - for string in ["e^rho", "normalization", "theta", "zeta"]: - assert string in data.keys() - assert data["normalization"] == 1 +class TestPlotSurfaces: + """Tests for plot_surfaces.""" - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_surfaces(self): + """Test plotting flux surfaces.""" + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_surfaces(eq, return_data=True) + for string in [ + "rho_R_coords", + "rho_Z_coords", + "vartheta_R_coords", + "vartheta_Z_coords", + ]: + assert string in data.keys() + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_cov_basis(DSHAPE_current): - """Test 2d plot of norm of e_rho.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d(eq, "e_rho", figsize=(4, 4), return_data=True) - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_surfaces_no_theta(self): + """Test plotting flux surfaces without theta contours.""" + eq = get("DSHAPE_CURRENT") + fig, ax, data = plot_surfaces(eq, theta=False, return_data=True) + for string in ["rho_R_coords", "rho_Z_coords"]: + assert string in data.keys() + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_magnetic_tension(DSHAPE_current): - """Test 2d plot of magnetic tension.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d(eq, "|(B*grad)B|", figsize=(4, 4), return_data=True) - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_surfaces_HELIOTRON(self): + """Test plot surfaces of equilibrium for correctness of vartheta lines.""" + fig, ax = plot_surfaces(get("HELIOTRON")) + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_magnetic_pressure(DSHAPE_current): - """Test 2d plot of magnetic pressure.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d(eq, "|grad(|B|^2)|/2mu0", figsize=(4, 4), return_data=True) - return fig +class TestPlotBoundary: + """Tests for plot_boundary and plot_boundaries.""" + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_boundary(self): + """Test plotting boundary.""" + eq = get("W7-X") + fig, ax, data = plot_boundary(eq, plot_axis=True, return_data=True) + assert "R" in data.keys() + assert "Z" in data.keys() -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_gradpsi(DSHAPE_current): - """Test 2d plot of norm of grad(rho).""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d(eq, "|grad(rho)|", figsize=(4, 4), return_data=True) - return fig + return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_boundary_surface(self): + """Test plot_boundary function for Surface objects.""" + surf = FourierRZToroidalSurface() + fig, ax = plot_boundary(surf) + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_normF_2d(DSHAPE_current): - """Test 2d plot of normalized force.""" - grid = LinearGrid(rho=np.array(0.8), M=20, N=2) - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_2d( - eq, "|F|", norm_F=True, figsize=(4, 4), return_data=True, grid=grid - ) - for string in ["|F|", "theta", "zeta"]: - assert string in data.keys() - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_boundaries(self): + """Test plotting boundaries.""" + eq1 = get("SOLOVEV") + eq2 = get("DSHAPE") + eq3 = get("W7-X") + fig, ax, data = plot_boundaries((eq1, eq2, eq3), return_data=True) + assert "R" in data.keys() + assert "Z" in data.keys() + assert len(data["R"]) == 3 + assert len(data["Z"]) == 3 + return fig -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_normF_section(DSHAPE_current): - """Test Poincare section plot of normalized force on log scale.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_section(eq, "|F|", norm_F=True, log=True) - return fig +class TestPlotComparison: + """Tests for plot_comparison.""" -@pytest.mark.unit -@pytest.mark.solve -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_coefficients(DSHAPE_current): - """Test scatter plot of spectral coefficients.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_coefficients(eq, color="b", marker="o") - ax[0, 0].set_ylim([1e-8, 1e1]) - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_comparison(self): + """Test plotting comparison of flux surfaces.""" + eqf = get("DSHAPE_CURRENT", "all") + fig, ax, data = plot_comparison(eqf, return_data=True) + for string in [ + "rho_R_coords", + "rho_Z_coords", + "vartheta_R_coords", + "vartheta_Z_coords", + ]: + assert string in data.keys() + assert len(data[string]) == len(eqf) + return fig -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_logo(): - """Test plotting the DESC logo.""" - fig, ax = plot_logo() - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_comparison_no_theta(self): + """Test plotting comparison of flux surfaces without theta contours.""" + eqf = get("DSHAPE_CURRENT", "all") + fig, ax = plot_comparison(eqf, theta=0) + return fig class TestPlotGrid: @@ -737,123 +645,118 @@ def test_plot_basis_fourierzernike(self): return fig -class TestPlotFieldLines: - """Tests for plotting field lines.""" +class TestPlotBoozerModes: + """Tests for plot_boozer_modes.""" @pytest.mark.unit - def test_find_idx(self): - """Test finding the index of the node closest to a given point.""" - # pick the first grid node point, add epsilon to it, check it returns idx of 0 - grid = LinearGrid(L=1, M=2, N=2, axis=False) - epsilon = np.finfo(float).eps - test_point = grid.nodes[0, :] + epsilon - idx = _find_idx(*test_point, grid=grid) - assert idx == 0 - - @pytest.mark.unit - @pytest.mark.solve @pytest.mark.slow - @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) - def test_plot_field_line(self, DSHAPE_current): - """Test plotting single field line over 1 transit.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax, data = plot_field_lines_sfl( - eq, rho=1, seed_thetas=0, phi_end=2 * np.pi, return_data=True + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_boozer_modes(self): + """Test plotting boozer spectrum.""" + eq = get("WISTELL-A") + fig, ax, data = plot_boozer_modes( + eq, + M_booz=eq.M, + N_booz=eq.N, + num_modes=7, + return_data=True, + norm=True, + rho=5, ) - for string in ["R", "Z", "phi", "seed_thetas"]: + ax.set_ylim([1e-6, 5e0]) + for string in ["|B|_mn", "B modes", "rho"]: assert string in data.keys() return fig @pytest.mark.unit - @pytest.mark.solve @pytest.mark.slow - @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_3d) - def test_plot_field_lines(self, DSHAPE_current): - """Test plotting multiple field lines over 1 transit.""" - eq = load(load_from=str(DSHAPE_current["desc_h5_path"]))[-1] - fig, ax = plot_field_lines_sfl( - eq, rho=1, seed_thetas=np.linspace(0, 2 * np.pi, 4), phi_end=2 * np.pi + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_boozer_modes_breaking_only(self): + """Test plotting symmetry breaking boozer spectrum.""" + eq = get("WISTELL-A") + fig, ax = plot_boozer_modes( + eq, + M_booz=eq.M, + N_booz=eq.N, + helicity=(1, -eq.NFP), + norm=True, + num_modes=7, + rho=5, ) + ax.set_ylim([1e-6, 5e0]) return fig + @pytest.mark.unit + @pytest.mark.slow + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_boozer_modes_max(self): + """Test plotting symmetry breaking boozer spectrum.""" + eq = get("WISTELL-A") + fig, ax = plot_boozer_modes( + eq, + M_booz=eq.M, + N_booz=eq.N, + helicity=(1, -eq.NFP), + max_only=True, + label="WISTELL-A", + color="r", + norm=True, + rho=5, + ) + ax.set_ylim([1e-6, 5e0]) + return fig -@pytest.mark.unit -@pytest.mark.slow -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_boozer_modes(): - """Test plotting boozer spectrum.""" - eq = get("WISTELL-A") - fig, ax, data = plot_boozer_modes( - eq, M_booz=eq.M, N_booz=eq.N, num_modes=7, return_data=True, norm=True - ) - ax.set_ylim([1e-6, 5e0]) - for string in ["|B|_mn", "B modes", "rho"]: - assert string in data.keys() - return fig - - -@pytest.mark.unit -@pytest.mark.slow -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_boozer_modes_breaking_only(): - """Test plotting symmetry breaking boozer spectrum.""" - eq = get("WISTELL-A") - fig, ax = plot_boozer_modes( - eq, - M_booz=eq.M, - N_booz=eq.N, - helicity=(1, -eq.NFP), - norm=True, - num_modes=7, - ) - ax.set_ylim([1e-6, 5e0]) - return fig - - -@pytest.mark.unit -@pytest.mark.slow -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_boozer_modes_max(): - """Test plotting symmetry breaking boozer spectrum.""" - eq = get("WISTELL-A") - fig, ax = plot_boozer_modes( - eq, - M_booz=eq.M, - N_booz=eq.N, - helicity=(1, -eq.NFP), - max_only=True, - label="WISTELL-A", - color="r", - norm=True, - ) - ax.set_ylim([1e-6, 5e0]) - return fig + @pytest.mark.unit + @pytest.mark.slow + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) + def test_plot_boozer_modes_no_norm(self): + """Test plotting boozer spectrum without B0 and norm.""" + eq = get("ESTELL") + fig, ax = plot_boozer_modes( + eq, + M_booz=eq.M, + N_booz=eq.N, + num_modes=7, + B0=False, + log=False, + rho=5, + ) + return fig -@pytest.mark.unit -@pytest.mark.slow -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_boozer_modes_no_norm(): - """Test plotting boozer spectrum without B0 and norm.""" - eq = get("ESTELL") - fig, ax = plot_boozer_modes( - eq, M_booz=eq.M, N_booz=eq.N, num_modes=7, B0=False, log=False - ) - return fig +class TestPlotBoozerSurface: + """Tests for plot_boozer_surface.""" + @pytest.mark.unit + @pytest.mark.slow + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_boozer_surface(self): + """Test plotting B in boozer coordinates.""" + eq = get("WISTELL-A") + fig, ax, data = plot_boozer_surface( + eq, M_booz=eq.M, N_booz=eq.N, return_data=True, rho=0.5, fieldlines=4 + ) + for string in ["|B|", "theta_B", "zeta_B"]: + assert string in data.keys() + return fig -@pytest.mark.unit -@pytest.mark.slow -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_boozer_surface(): - """Test plotting B in boozer coordinates.""" - eq = get("WISTELL-A") - fig, ax, data = plot_boozer_surface( - eq, M_booz=eq.M, N_booz=eq.N, return_data=True, rho=0.5, fieldlines=4 - ) - for string in ["|B|", "theta_B", "zeta_B"]: - assert string in data.keys() - return fig + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + def test_plot_omnigenous_field(self): + """Test plot omnigenous magnetic field.""" + field = OmnigenousField( + L_B=0, + M_B=4, + L_x=0, + M_x=1, + N_x=1, + NFP=4, + helicity=(1, 4), + B_lm=np.array([0.8, 0.9, 1.1, 1.2]), + x_lmn=np.array([0, -np.pi / 8, 0, np.pi / 8, 0, np.pi / 4]), + ) + fig, ax = plot_boozer_surface(field, iota=0.6, fieldlines=4) + return fig @pytest.mark.unit @@ -863,7 +766,13 @@ def test_plot_qs_error(): """Test plotting qs error metrics.""" eq = get("WISTELL-A") fig, ax, data = plot_qs_error( - eq, helicity=(1, -eq.NFP), M_booz=eq.M, N_booz=eq.N, log=True, return_data=True + eq, + helicity=(1, -eq.NFP), + M_booz=eq.M, + N_booz=eq.N, + log=True, + return_data=True, + rho=5, ) ax.set_ylim([1e-3, 2e-1]) for string in ["rho", "f_T", "f_B", "f_C"]: @@ -938,99 +847,17 @@ def test_plot_b_mag(): @pytest.mark.unit @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_surfaces_HELIOTRON(): - """Test plot surfaces of equilibrium for correctness of vartheta lines.""" - fig, ax = plot_surfaces(get("HELIOTRON")) - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_1d_curve(): - """Test plot_1d function for Curve objects.""" - curve = FourierXYZCurve([0, 10, 1]) - fig, ax = plot_1d(curve, "curvature") - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_2d_surface(): - """Test plot_2d function for Surface objects.""" - surf = FourierRZToroidalSurface() - fig, ax = plot_2d(surf, "curvature_H_rho") - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_1d_surface(): - """Test plot_1d function for Surface objects.""" - surf = FourierRZToroidalSurface() - fig, ax = plot_1d(surf, "curvature_H_rho", grid=LinearGrid(M=50)) - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_plot_boundary_surface(): - """Test plot_boundary function for Surface objects.""" - surf = FourierRZToroidalSurface() - fig, ax = plot_boundary(surf) - return fig - - -@pytest.mark.unit -def test_plot_3d_surface(): - """Test 3d plotting of surface object.""" - surf = FourierRZToroidalSurface() - fig = plot_3d(surf, "curvature_H_rho") +def test_plot_coefficients(): + """Test scatter plot of spectral coefficients.""" + eq = get("DSHAPE_CURRENT") + fig, ax = plot_coefficients(eq, color="b", marker="o") + ax[0, 0].set_ylim([1e-8, 1e1]) return fig @pytest.mark.unit @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -def test_plot_omnigenous_field(): - """Test plot omnigenous magnetic field.""" - field = OmnigenousField( - L_B=0, - M_B=4, - L_x=0, - M_x=1, - N_x=1, - NFP=4, - helicity=(1, 4), - B_lm=np.array([0.8, 0.9, 1.1, 1.2]), - x_lmn=np.array([0, -np.pi / 8, 0, np.pi / 8, 0, np.pi / 4]), - ) - fig, ax = plot_boozer_surface(field, iota=0.6, fieldlines=4) - return fig - - -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) -@pytest.mark.unit -def test_2d_plot_Bn(): - """Test 2d plotting of Bn on equilibrium surface.""" - eq = get("HELIOTRON") - fig, _ = plot_2d( - eq, - "B*n", - field=ToroidalMagneticField(1, 1), - field_grid=LinearGrid(M=10, N=10), - grid=LinearGrid(M=30, N=30, NFP=eq.NFP, endpoint=True), - ) - return fig - - -@pytest.mark.unit -def test_3d_plot_Bn(): - """Test 3d plotting of Bn on equilibrium surface.""" - eq = get("precise_QA") - eq.change_resolution(M=4, N=4, L=4, M_grid=8, N_grid=8, L_grid=8) - fig = plot_3d( - eq, - "B*n", - field=ToroidalMagneticField(1, 1), - grid=LinearGrid(M=30, N=30, NFP=1, endpoint=True), - ) +def test_plot_logo(): + """Test plotting the DESC logo.""" + fig, ax = plot_logo() return fig diff --git a/tests/test_profiles.py b/tests/test_profiles.py index b8b2513fdd..f81d6e2d42 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -371,7 +371,7 @@ def test_default_profiles(self): np.testing.assert_allclose(mp(x), 0) np.testing.assert_allclose(zp(x), 0) - @pytest.mark.unit + @pytest.mark.regression def test_solve_with_combined(self): """Make sure combined profiles work correctly for solving equilibrium. diff --git a/tests/test_stability_funs.py b/tests/test_stability_funs.py index e4633c4831..54aa43e278 100644 --- a/tests/test_stability_funs.py +++ b/tests/test_stability_funs.py @@ -42,13 +42,13 @@ def assert_all_close( np.testing.assert_allclose(y1[interval], y2[interval], rtol=rtol, atol=atol) -def get_vmec_data(stellarator, quantity): +def get_vmec_data(path, quantity): """Get data from a VMEC wout.nc file. Parameters ---------- - stellarator : str - The equilibrium's fixture. + path : str + Path to VMEC file. quantity: str Name of the quantity to return. @@ -60,7 +60,7 @@ def get_vmec_data(stellarator, quantity): Variable from VMEC output. """ - f = Dataset(str(stellarator["vmec_nc_path"])) + f = Dataset(path) rho = np.sqrt(f.variables["phi"] / np.array(f.variables["phi"])[-1]) q = np.array(f.variables[quantity]) f.close() @@ -80,13 +80,11 @@ def test_mercier_vacuum(): @pytest.mark.unit -@pytest.mark.solve -def test_compute_d_shear(DSHAPE_current, HELIOTRON_ex): +def test_compute_d_shear(): """Test that D_shear has a stabilizing effect and matches VMEC.""" - def test(stellarator, rho_range=(0, 1), rtol=1e-12, atol=0.0): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] - rho, d_shear_vmec = get_vmec_data(stellarator, "DShear") + def test(eq, vmec, rho_range=(0, 1), rtol=1e-12, atol=0.0): + rho, d_shear_vmec = get_vmec_data(vmec, "DShear") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) d_shear = grid.compress(eq.compute("D_shear", grid=grid)["D_shear"]) @@ -95,20 +93,22 @@ def test(stellarator, rho_range=(0, 1), rtol=1e-12, atol=0.0): ), "D_shear should always have a stabilizing effect." assert_all_close(d_shear, d_shear_vmec, rho, rho_range, rtol, atol) - test(DSHAPE_current, (0.3, 0.9), atol=0.01, rtol=0.1) - test(HELIOTRON_ex) + test( + desc.examples.get("DSHAPE_CURRENT"), + ".//tests//inputs//wout_DSHAPE.nc", + (0.3, 0.9), + atol=0.01, + rtol=0.1, + ) + test(desc.examples.get("HELIOTRON"), ".//tests//inputs//wout_HELIOTRON.nc") @pytest.mark.unit -@pytest.mark.solve -def test_compute_d_current(DSHAPE_current, HELIOTRON_ex): +def test_compute_d_current(): """Test calculation of D_current stability criterion against VMEC.""" - def test( - stellarator, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL - ): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] - rho, d_current_vmec = get_vmec_data(stellarator, "DCurr") + def test(eq, vmec, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL): + rho, d_current_vmec = get_vmec_data(vmec, "DCurr") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) d_current = grid.compress(eq.compute("D_current", grid=grid)["D_current"]) @@ -118,20 +118,27 @@ def test( ) assert_all_close(d_current, d_current_vmec, rho, rho_range, rtol, atol) - test(DSHAPE_current, (0.3, 0.9), rtol=1e-1, atol=1e-2) - test(HELIOTRON_ex, (0.25, 0.85), rtol=1e-1) + test( + desc.examples.get("DSHAPE_CURRENT"), + ".//tests//inputs//wout_DSHAPE.nc", + (0.3, 0.9), + rtol=1e-1, + atol=1e-2, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.25, 0.85), + rtol=1e-1, + ) @pytest.mark.unit -@pytest.mark.solve -def test_compute_d_well(DSHAPE_current, HELIOTRON_ex): +def test_compute_d_well(): """Test calculation of D_well stability criterion against VMEC.""" - def test( - stellarator, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL - ): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] - rho, d_well_vmec = get_vmec_data(stellarator, "DWell") + def test(eq, vmec, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL): + rho, d_well_vmec = get_vmec_data(vmec, "DWell") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) d_well = grid.compress(eq.compute("D_well", grid=grid)["D_well"]) @@ -140,22 +147,38 @@ def test( ) assert_all_close(d_well, d_well_vmec, rho, rho_range, rtol, atol) - test(DSHAPE_current, (0.3, 0.9), rtol=1e-1) - test(HELIOTRON_ex, (0.01, 0.45), rtol=1.75e-1) - test(HELIOTRON_ex, (0.45, 0.6), atol=7.2e-1) - test(HELIOTRON_ex, (0.6, 0.99), rtol=2e-2) + test( + desc.examples.get("DSHAPE_CURRENT"), + ".//tests//inputs//wout_DSHAPE.nc", + (0.3, 0.9), + rtol=1e-1, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.01, 0.45), + rtol=1.75e-1, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.45, 0.6), + atol=7.2e-1, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.6, 0.99), + rtol=2e-2, + ) @pytest.mark.unit -@pytest.mark.solve -def test_compute_d_geodesic(DSHAPE_current, HELIOTRON_ex): +def test_compute_d_geodesic(): """Test that D_geodesic has a destabilizing effect and matches VMEC.""" - def test( - stellarator, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL - ): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] - rho, d_geodesic_vmec = get_vmec_data(stellarator, "DGeod") + def test(eq, vmec, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL): + rho, d_geodesic_vmec = get_vmec_data(vmec, "DGeod") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) d_geodesic = grid.compress(eq.compute("D_geodesic", grid=grid)["D_geodesic"]) @@ -164,21 +187,33 @@ def test( ), "D_geodesic should always have a destabilizing effect." assert_all_close(d_geodesic, d_geodesic_vmec, rho, rho_range, rtol, atol) - test(DSHAPE_current, (0.3, 0.9), rtol=1e-1) - test(HELIOTRON_ex, (0.15, 0.825), rtol=1.2e-1) - test(HELIOTRON_ex, (0.85, 0.95), atol=1.2e-1) + test( + desc.examples.get("DSHAPE_CURRENT"), + ".//tests//inputs//wout_DSHAPE.nc", + (0.3, 0.9), + rtol=1e-1, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.15, 0.825), + rtol=1.2e-1, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.85, 0.95), + atol=1.2e-1, + ) @pytest.mark.unit @pytest.mark.solve -def test_compute_d_mercier(DSHAPE_current, HELIOTRON_ex): +def test_compute_d_mercier(): """Test calculation of D_Mercier stability criterion against VMEC.""" - def test( - stellarator, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL - ): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] - rho, d_mercier_vmec = get_vmec_data(stellarator, "DMerc") + def test(eq, vmec, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL): + rho, d_mercier_vmec = get_vmec_data(vmec, "DMerc") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) d_mercier = grid.compress(eq.compute("D_Mercier", grid=grid)["D_Mercier"]) @@ -188,18 +223,33 @@ def test( ) assert_all_close(d_mercier, d_mercier_vmec, rho, rho_range, rtol, atol) - test(DSHAPE_current, (0.3, 0.9), rtol=1e-1, atol=1e-2) - test(HELIOTRON_ex, (0.1, 0.325), rtol=1.3e-1) - test(HELIOTRON_ex, (0.325, 0.95), rtol=5e-2) + test( + desc.examples.get("DSHAPE_CURRENT"), + ".//tests//inputs//wout_DSHAPE.nc", + (0.3, 0.9), + rtol=1e-1, + atol=1e-2, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.1, 0.325), + rtol=1.3e-1, + ) + test( + desc.examples.get("HELIOTRON"), + ".//tests//inputs//wout_HELIOTRON.nc", + (0.325, 0.95), + rtol=5e-2, + ) @pytest.mark.unit @pytest.mark.solve -def test_compute_magnetic_well(DSHAPE_current, HELIOTRON_ex): +def test_compute_magnetic_well(): """Test that D_well and magnetic_well match signs under finite pressure.""" - def test(stellarator, rho=np.linspace(0, 1, 128)): - eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] + def test(eq, rho=np.linspace(0, 1, 128)): grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) d_well = grid.compress(eq.compute("D_well", grid=grid)["D_well"]) magnetic_well = grid.compress( @@ -210,8 +260,8 @@ def test(stellarator, rho=np.linspace(0, 1, 128)): <= MAX_SIGN_DIFF ) - test(DSHAPE_current) - test(HELIOTRON_ex) + test(desc.examples.get("DSHAPE_CURRENT")) + test(desc.examples.get("HELIOTRON")) @pytest.mark.unit diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index 2770f27c93..32e3307efd 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -287,6 +287,7 @@ def test_constant_offset_surface_rot_ellipse(self): rtol=2e-2, ) + @pytest.mark.unit def test_position(self): """Tests for position on surface.""" s = FourierRZToroidalSurface() diff --git a/tests/test_vmec.py b/tests/test_vmec.py index a327fc412f..d7ae22f2b4 100644 --- a/tests/test_vmec.py +++ b/tests/test_vmec.py @@ -4,6 +4,7 @@ import pytest from netCDF4 import Dataset +import desc.examples from desc.basis import DoubleFourierSeries, FourierZernikeBasis from desc.equilibrium import EquilibriaFamily, Equilibrium from desc.examples import get @@ -874,12 +875,11 @@ def test_vmec_save_2(VMEC_save): @pytest.mark.unit -@pytest.mark.solve @pytest.mark.mpl_image_compare(tolerance=1) -def test_plot_vmec_comparison(SOLOVEV): +def test_plot_vmec_comparison(): """Test that DESC and VMEC flux surface plots match.""" - eq = EquilibriaFamily.load(load_from=str(SOLOVEV["desc_h5_path"]))[-1] - fig, ax = VMECIO.plot_vmec_comparison(eq, str(SOLOVEV["vmec_nc_path"])) + eq = desc.examples.get("SOLOVEV") + fig, ax = VMECIO.plot_vmec_comparison(eq, "tests/inputs/wout_SOLOVEV.nc") return fig @@ -911,7 +911,7 @@ def test_vmec_boundary_subspace(DummyStellarator): np.testing.assert_allclose(zbs_ref, np.abs(zbs) > tol) -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.solve def test_write_vmec_input(TmpDir): """Test generated VMEC input file gives the original equilibrium when solved."""