diff --git a/.test_durations b/.test_durations index 29003036c2..7f01447a16 100644 --- a/.test_durations +++ b/.test_durations @@ -1,420 +1,445 @@ { - "tests/test_axis_limits.py::TestAxisLimits::test_compute_axis_limit_api": 2.5391948223114014, - "tests/test_axis_limits.py::TestAxisLimits::test_b_mag_fsa": 7.966105818748474, - "tests/test_axis_limits.py::TestAxisLimits::test_rotational_transform": 11.269968628883362, - "tests/test_backend.py::test_put": 0.19738554954528809, - "tests/test_backend.py::test_sign": 0.19631361961364746, - "tests/test_basis.py::TestBasis::test_polyder": 0.1083458662033081, - "tests/test_basis.py::TestBasis::test_polyval": 0.10746932029724121, - "tests/test_basis.py::TestBasis::test_zernike_coeffs": 0.1800525188446045, - "tests/test_basis.py::TestBasis::test_polyval_exact": 0.10746932029724121, - "tests/test_basis.py::TestBasis::test_powers": 0.1132669448852539, - "tests/test_basis.py::TestBasis::test_zernike_radial": 0.8333420753479004, - "tests/test_basis.py::TestBasis::test_fourier": 0.2550234794616699, - "tests/test_basis.py::TestBasis::test_power_series": 0.11781167984008789, - "tests/test_basis.py::TestBasis::test_double_fourier": 0.2743496894836426, - "tests/test_basis.py::TestBasis::test_change_resolution": 0.15073871612548828, - "tests/test_basis.py::TestBasis::test_repr": 0.415041446685791, - "tests/test_basis.py::TestBasis::test_zernike_indexing": 0.11620235443115234, - "tests/test_basis.py::TestBasis::test_derivative_not_in_basis_zeros": 0.3892948627471924, - "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_analytic": 2.3723795413970947, - "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_Kim": 7.052055478096008, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_second_pass": 0.661003828048706, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_2_3": 1.5928778648376465, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_4_5": 4.930520176887512, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_tokamak_benchmark": 4.571721792221069, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QA": 5.316905498504639, - "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QH": 5.3142324686050415, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_normalization": 27.70256757736206, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_resolution": 62.08576798439026, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_iota": 185.59204578399658, - "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_current": 226.8012100458145, - "tests/test_bootstrap.py::test_bootstrap_objective_build": 12.902662992477417, - "tests/test_coils.py::TestCoil::test_biot_savart": 1.2084215879440308, - "tests/test_coils.py::TestCoil::test_properties": 1.0765130519866943, - "tests/test_coils.py::TestCoilSet::test_linspaced_linear": 1.440415859222412, - "tests/test_coils.py::TestCoilSet::test_linspaced_angular": 3.4112319946289062, - "tests/test_coils.py::TestCoilSet::test_from_symmetry": 5.091866493225098, - "tests/test_coils.py::TestCoilSet::test_properties": 1.0765130519866943, - "tests/test_coils.py::TestCoilSet::test_dunder_methods": 0.6133010387420654, - "tests/test_coils.py::test_repr": 0.415041446685791, - "tests/test_compute_funs.py::test_total_volume": 3.275167226791382, - "tests/test_compute_funs.py::test_enclosed_volumes": 2.591869831085205, - "tests/test_compute_funs.py::test_surface_areas": 2.0627825260162354, - "tests/test_compute_funs.py::test_surface_areas_2": 2.0627825260162354, - "tests/test_compute_funs.py::test_magnetic_field_derivatives": 9.587448120117188, - "tests/test_compute_funs.py::test_metric_derivatives": 2.72632372379303, - "tests/test_compute_funs.py::test_magnetic_pressure_gradient": 4.304385185241699, - "tests/test_compute_funs.py::test_currents": 6.562957286834717, - "tests/test_compute_funs.py::test_BdotgradB": 3.5210081338882446, - "tests/test_compute_funs.py::test_boozer_transform": 4.997244834899902, - "tests/test_compute_funs.py::test_compute_grad_p_volume_avg": 1.6617939472198486, - "tests/test_compute_funs.py::test_compare_quantities_to_vmec": 11.2726891040802, - "tests/test_compute_funs.py::test_compute_everything": 84.62090516090393, - "tests/test_compute_funs.py::test_compute_averages": 8.572983503341675, - "tests/test_compute_utils.py::TestComputeUtils::test_compress_expand_inverse_op": 0.19367456436157227, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals": 1.827927827835083, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_area_unweighted": 1.4398181438446045, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_area_weighted": 6.354798316955566, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_identity_op": 5.77168881893158, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_homomorphism": 7.06252646446228, - "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_shortcut": 5.325978398323059, - "tests/test_compute_utils.py::TestComputeUtils::test_min_max": 0.8481179475784302, - "tests/test_configuration.py::TestConstructor::test_defaults": 0.9968388080596924, - "tests/test_configuration.py::TestConstructor::test_supplied_objects": 3.4648375511169434, - "tests/test_configuration.py::TestConstructor::test_dict": 1.9488677978515625, - "tests/test_configuration.py::TestConstructor::test_asserts": 0.7387478351593018, - "tests/test_configuration.py::TestConstructor::test_supplied_coeffs": 0.8749246597290039, - "tests/test_configuration.py::TestInitialGuess::test_default_set": 2.1531221866607666, - "tests/test_configuration.py::TestInitialGuess::test_errors": 1.3456809520721436, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_other": 1.3564915657043457, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface": 1.705718994140625, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface2": 1.705718994140625, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_points": 3.1299737691879272, - "tests/test_configuration.py::TestInitialGuess::test_NFP_error": 0.9642298221588135, - "tests/test_configuration.py::TestInitialGuess::test_guess_from_file": 2.598916530609131, - "tests/test_configuration.py::TestGetSurfaces::test_get_rho_surface": 1.1722404956817627, - "tests/test_configuration.py::TestGetSurfaces::test_get_zeta_surface": 1.0646436214447021, - "tests/test_configuration.py::TestGetSurfaces::test_get_theta_surface": 0.8025240898132324, - "tests/test_configuration.py::TestGetSurfaces::test_asserts": 0.7387478351593018, - "tests/test_configuration.py::test_magnetic_axis": 3.3196523189544678, - "tests/test_configuration.py::test_is_nested": 2.3451576232910156, - "tests/test_configuration.py::test_is_nested_theta": 1.8744008541107178, - "tests/test_configuration.py::test_get_profile": 4.964652061462402, - "tests/test_configuration.py::test_kinetic_errors": 1.9296765327453613, - "tests/test_constrain_current.py::TestConstrainCurrent::test_compute_rotational_transform": 22.117167949676514, - "tests/test_curves.py::TestRZCurve::test_length": 1.2188292741775513, - "tests/test_curves.py::TestRZCurve::test_curvature": 1.0212721824645996, - "tests/test_curves.py::TestRZCurve::test_torsion": 0.9370831251144409, - "tests/test_curves.py::TestRZCurve::test_frenet": 1.167402982711792, - "tests/test_curves.py::TestRZCurve::test_coords": 0.873279333114624, - "tests/test_curves.py::TestRZCurve::test_misc": 1.248878836631775, - "tests/test_curves.py::TestRZCurve::test_asserts": 0.7387478351593018, - "tests/test_curves.py::TestXYZCurve::test_length": 1.2188292741775513, - "tests/test_curves.py::TestXYZCurve::test_curvature": 1.0212721824645996, - "tests/test_curves.py::TestXYZCurve::test_torsion": 0.9370831251144409, - "tests/test_curves.py::TestXYZCurve::test_frenet": 1.167402982711792, - "tests/test_curves.py::TestXYZCurve::test_coords": 0.873279333114624, - "tests/test_curves.py::TestXYZCurve::test_misc": 1.248878836631775, - "tests/test_curves.py::TestXYZCurve::test_asserts": 0.7387478351593018, - "tests/test_curves.py::TestPlanarCurve::test_length": 1.2188292741775513, - "tests/test_curves.py::TestPlanarCurve::test_curvature": 1.0212721824645996, - "tests/test_curves.py::TestPlanarCurve::test_torsion": 0.9370831251144409, - "tests/test_curves.py::TestPlanarCurve::test_frenet": 1.167402982711792, - "tests/test_curves.py::TestPlanarCurve::test_coords": 0.873279333114624, - "tests/test_curves.py::TestPlanarCurve::test_misc": 1.248878836631775, - "tests/test_curves.py::TestPlanarCurve::test_asserts": 0.7387478351593018, - "tests/test_derivatives.py::TestDerivative::test_finite_diff_vec": 0.7251818180084229, - "tests/test_derivatives.py::TestDerivative::test_finite_diff_scalar": 0.7568480968475342, - "tests/test_derivatives.py::TestDerivative::test_auto_diff": 0.8526492118835449, - "tests/test_derivatives.py::TestDerivative::test_compare_AD_FD": 0.7734799385070801, - "tests/test_derivatives.py::TestDerivative::test_fd_hessian": 0.6189804077148438, - "tests/test_derivatives.py::TestDerivative::test_block_jacobian": 0.8143317699432373, - "tests/test_derivatives.py::TestJVP::test_autodiff_jvp": 0.7789533138275146, - "tests/test_derivatives.py::TestJVP::test_finitediff_jvp": 0.7164802551269531, - "tests/test_derivatives.py::TestJVP::test_autodiff_jvp2": 0.7789533138275146, - "tests/test_derivatives.py::TestJVP::test_finitediff_jvp2": 0.7164802551269531, - "tests/test_derivatives.py::TestJVP::test_autodiff_jvp3": 0.7789533138275146, - "tests/test_derivatives.py::TestJVP::test_finitediff_jvp3": 0.6344552040100098, - "tests/test_equilibrium.py::test_compute_geometry": 1.9436662197113037, - "tests/test_equilibrium.py::test_compute_theta_coords": 5.6960673332214355, - "tests/test_equilibrium.py::test_compute_flux_coords": 20.213435173034668, - "tests/test_equilibrium.py::test_map_coordinates": 14.131819248199463, - "tests/test_equilibrium.py::test_to_sfl": 26.418137550354004, - "tests/test_equilibrium.py::test_continuation_resolution": 156.08914399147034, - "tests/test_equilibrium.py::test_grid_resolution_warning": 290.5435914993286, - "tests/test_equilibrium.py::test_eq_change_grid_resolution": 1.3629605770111084, - "tests/test_equilibrium.py::test_resolution": 4.879020810127258, - "tests/test_equilibrium.py::test_equilibrium_from_near_axis": 19.78379797935486, - "tests/test_equilibrium.py::test_poincare_solve_not_implemented": 1.0482544898986816, - "tests/test_equilibrium.py::test_equilibriafamily_constructor": 2.873791217803955, - "tests/test_equilibrium.py::test_change_NFP": 7.288805723190308, - "tests/test_examples.py::test_SOLOVEV_vacuum": 2.2789478302001953, - "tests/test_examples.py::test_SOLOVEV_results": 6.3239336013793945, - "tests/test_examples.py::test_DSHAPE_results": 9.376911878585815, - "tests/test_examples.py::test_DSHAPE_current_results": 9.261709928512573, - "tests/test_examples.py::test_HELIOTRON_results": 10.087924718856812, - "tests/test_examples.py::test_HELIOTRON_vac_results": 11.681270122528076, - "tests/test_examples.py::test_precise_QH_results": 29.177587032318115, - "tests/test_examples.py::test_HELIOTRON_vac2_results": 41.4894198179245, - "tests/test_examples.py::test_force_balance_grids": 155.66855549812317, - "tests/test_examples.py::test_solve_bounds": 25.855547785758972, - "tests/test_examples.py::test_1d_optimization": 45.99489974975586, - "tests/test_examples.py::test_1d_optimization_old": 45.99489974975586, - "tests/test_examples.py::test_qh_optimization1": 2177.9836943149567, - "tests/test_examples.py::test_qh_optimization2": 1525.166443824768, - "tests/test_examples.py::test_qh_optimization3": 1084.8827364444733, - "tests/test_examples.py::test_ATF_results": 828.3503400087357, - "tests/test_examples.py::test_ESTELL_results": 1949.1876405477524, - "tests/test_examples.py::test_simsopt_QH_comparison": 723.4929859638214, - "tests/test_examples.py::test_NAE_QSC_solve": 149.5543862581253, - "tests/test_examples.py::test_NAE_QIC_solve": 299.23267793655396, - "tests/test_examples.py::TestGetExample::test_missing_example": 0.3255997896194458, - "tests/test_examples.py::TestGetExample::test_example_get_eq": 1.617514967918396, - "tests/test_examples.py::TestGetExample::test_example_get_eqf": 1.617514967918396, - "tests/test_examples.py::TestGetExample::test_example_get_boundary": 4.021717548370361, - "tests/test_examples.py::TestGetExample::test_example_get_pressure": 3.003194570541382, - "tests/test_examples.py::TestGetExample::test_example_get_iota": 3.9420523643493652, - "tests/test_examples.py::TestGetExample::test_example_get_current": 3.4969444274902344, - "tests/test_geometry.py::test_rotation_matrix": 0.7729271650314331, - "tests/test_geometry.py::test_xyz2rpz": 0.7970137596130371, - "tests/test_geometry.py::test_rpz2xyz": 0.6938657760620117, - "tests/test_grid.py::TestGrid::test_custom_grid": 0.674670934677124, - "tests/test_grid.py::TestGrid::test_linear_grid": 0.6298344135284424, - "tests/test_grid.py::TestGrid::test_linear_grid_spacing_consistency": 0.6298344135284424, - "tests/test_grid.py::TestGrid::test_linear_grid_symmetric_nodes_consistency": 0.6298344135284424, - "tests/test_grid.py::TestGrid::test_linear_grid_spacing_two_nodes": 0.6298344135284424, - "tests/test_grid.py::TestGrid::test_spacing_when_duplicate_node_is_removed": 0.3456343412399292, - "tests/test_grid.py::TestGrid::test_node_spacing_non_sym": 0.5532927513122559, - "tests/test_grid.py::TestGrid::test_symmetry_spacing_basic": 0.346652626991272, - "tests/test_grid.py::TestGrid::test_node_spacing_sym": 0.5870494842529297, - "tests/test_grid.py::TestGrid::test_concentric_grid": 0.6478164196014404, - "tests/test_grid.py::TestGrid::test_quadrature_grid": 0.6671996116638184, - "tests/test_grid.py::TestGrid::test_concentric_grid_high_res": 0.6478164196014404, - "tests/test_grid.py::TestGrid::test_quad_grid_volume_integration": 1.406738519668579, - "tests/test_grid.py::TestGrid::test_repr": 0.415041446685791, - "tests/test_grid.py::TestGrid::test_change_resolution": 0.15073871612548828, - "tests/test_grid.py::TestGrid::test_enforce_symmetry_sum": 1.4409854412078857, - "tests/test_grid.py::TestGrid::test_enforce_symmetry": 1.6080760955810547, - "tests/test_grid.py::TestGrid::test_symmetry_surface_average_1": 6.71723997592926, - "tests/test_grid.py::TestGrid::test_symmetry_surface_average_2": 3.8190581798553467, - "tests/test_grid.py::TestGrid::test_symmetry_volume_integral": 0.39383578300476074, - "tests/test_grid.py::test_find_most_rational_surfaces": 0.6310467720031738, - "tests/test_grid.py::test_find_least_rational_surfaces": 2.2816789150238037, - "tests/test_input_output.py::test_vmec_input": 0.7020907402038574, - "tests/test_input_output.py::test_write_desc_input_Nones": 0.3594777584075928, - "tests/test_input_output.py::test_near_axis_input_files": 6.730593204498291, - "tests/test_input_output.py::test_vmec_input_surface_threshold": 4.727794408798218, - "tests/test_input_output.py::TestInputReader::test_no_input_file": 0.6845892667770386, - "tests/test_input_output.py::TestInputReader::test_nonexistant_input_file": 0.6784034967422485, - "tests/test_input_output.py::TestInputReader::test_min_input": 0.6427890062332153, - "tests/test_input_output.py::TestInputReader::test_np_environ": 0.6391646862030029, - "tests/test_input_output.py::TestInputReader::test_quiet_verbose": 0.6901319026947021, - "tests/test_input_output.py::TestInputReader::test_vacuum_objective_with_iota_yields_current": 0.48584890365600586, - "tests/test_input_output.py::test_writer_given_filename": 0.6769628524780273, - "tests/test_input_output.py::test_writer_given_file": 0.6777760982513428, - "tests/test_input_output.py::test_writer_close_on_delete": 0.6699178218841553, - "tests/test_input_output.py::test_writer_write_dict": 0.6447567939758301, - "tests/test_input_output.py::test_writer_write_list": 0.6857888698577881, - "tests/test_input_output.py::test_writer_write_obj": 0.6838473081588745, - "tests/test_input_output.py::test_reader_given_filename": 0.678086519241333, - "tests/test_input_output.py::test_reader_given_file": 0.6793563365936279, - "tests/test_input_output.py::test_reader_read_obj": 0.679969310760498, - "tests/test_input_output.py::test_pickle_io": 1.3660407066345215, - "tests/test_input_output.py::test_ascii_io": 6.870228290557861, - "tests/test_input_output.py::test_copy": 1.1949867010116577, - "tests/test_input_output.py::test_save_none": 1.2435132265090942, - "tests/test_input_output.py::test_load_eq_without_current": 1.809807300567627, - "tests/test_interpolate.py::TestInterp1D::test_interp1d": 2.6422271728515625, - "tests/test_interpolate.py::TestInterp1D::test_interp1d_extrap_periodic": 2.2251296043395996, - "tests/test_interpolate.py::TestInterp1D::test_interp1d_monotonic": 1.8299287557601929, - "tests/test_interpolate.py::TestInterp2D::test_interp2d": 2.0778234004974365, - "tests/test_interpolate.py::TestInterp3D::test_interp3d": 2.5942535400390625, - "tests/test_linear_objectives.py::test_LambdaGauge_sym": 1.0124235153198242, - "tests/test_linear_objectives.py::test_LambdaGauge_asym": 1.276787281036377, - "tests/test_linear_objectives.py::test_bc_on_interior_surfaces": 22.0314724445343, - "tests/test_linear_objectives.py::test_constrain_bdry_with_only_one_mode": 0.8122732639312744, - "tests/test_linear_objectives.py::test_constrain_asserts": 49.25658893585205, - "tests/test_linear_objectives.py::test_fixed_mode_solve": 23.21229350566864, - "tests/test_linear_objectives.py::test_fixed_modes_solve": 17.56595468521118, - "tests/test_linear_objectives.py::test_fixed_axis_and_theta_SFL_solve": 16.71259582042694, - "tests/test_linear_objectives.py::test_factorize_linear_constraints_asserts": 2.9949564933776855, - "tests/test_linear_objectives.py::test_build_init": 3.613085985183716, - "tests/test_linear_objectives.py::test_kinetic_constraints": 1.579378604888916, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes": 24.015464782714844, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target": 23.561565041542053, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_axis": 10.262390851974487, - "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target_axis": 15.298805952072144, - "tests/test_linear_objectives.py::test_FixBoundary_with_single_weight": 1.0754834413528442, - "tests/test_linear_objectives.py::test_FixBoundary_passed_target_no_passed_modes_error": 0.5215559005737305, - "tests/test_linear_objectives.py::test_FixAxis_passed_target_no_passed_modes_error": 0.6616849899291992, - "tests/test_linear_objectives.py::test_FixMode_passed_target_no_passed_modes_error": 0.5265893936157227, - "tests/test_linear_objectives.py::test_FixSumModes_passed_target_too_long": 0.36019039154052734, - "tests/test_linear_objectives.py::test_FixMode_False_or_None_modes": 0.35930633544921875, - "tests/test_linear_objectives.py::test_FixSumModes_False_or_None_modes": 0.3607978820800781, - "tests/test_linear_objectives.py::test_FixAxis_util_correct_objectives": 0.3655357360839844, - "tests/test_magnetic_fields.py::TestMagneticFields::test_basic_fields": 0.9740085601806641, - "tests/test_magnetic_fields.py::TestMagneticFields::test_scalar_field": 1.149132251739502, - "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field": 6.306258201599121, - "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field_axisym": 2.356407880783081, - "tests/test_magnetic_fields.py::TestMagneticFields::test_field_line_integrate": 2.4446654319763184, - "tests/test_objective_funs.py::TestObjectiveFunction::test_generic": 3.971553087234497, - "tests/test_objective_funs.py::TestObjectiveFunction::test_objective_from_user": 1.045462965965271, - "tests/test_objective_funs.py::TestObjectiveFunction::test_volume": 1.670881986618042, - "tests/test_objective_funs.py::TestObjectiveFunction::test_aspect_ratio": 3.188887119293213, - "tests/test_objective_funs.py::TestObjectiveFunction::test_elongation": 4.361514568328857, - "tests/test_objective_funs.py::TestObjectiveFunction::test_energy": 3.293406128883362, - "tests/test_objective_funs.py::TestObjectiveFunction::test_target_iota": 1.8545842170715332, - "tests/test_objective_funs.py::TestObjectiveFunction::test_toroidal_current": 3.9505090713500977, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qa_boozer": 9.19603419303894, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qh_boozer": 10.426319360733032, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_twoterm": 9.440323829650879, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_tripleproduct": 3.9321413040161133, - "tests/test_objective_funs.py::TestObjectiveFunction::test_isodynamicity": 1.9491896629333496, - "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_boozer_grids": 2.60112202167511, - "tests/test_objective_funs.py::TestObjectiveFunction::test_mercier_stability": 28.899383544921875, - "tests/test_objective_funs.py::TestObjectiveFunction::test_magnetic_well": 8.674297094345093, - "tests/test_objective_funs.py::test_derivative_modes": 65.33874678611755, - "tests/test_objective_funs.py::test_rejit": 1.0335636138916016, - "tests/test_objective_funs.py::test_generic_compute": 1.504560947418213, - "tests/test_objective_funs.py::test_getter_setter": 0.6585021018981934, - "tests/test_objective_funs.py::test_bounds_format": 0.5785126686096191, - "tests/test_objective_funs.py::test_target_profiles": 4.193389654159546, - "tests/test_objective_funs.py::test_plasma_vessel_distance": 8.535668134689331, - "tests/test_objective_funs.py::test_mean_curvature": 11.42268705368042, - "tests/test_objective_funs.py::test_principal_curvature": 6.662958383560181, - "tests/test_objective_funs.py::test_field_scale_length": 41.40760099887848, - "tests/test_objective_funs.py::test_profile_objective_print": 4.369129657745361, - "tests/test_objective_funs.py::test_plasma_vessel_distance_print": 8.535668134689331, - "tests/test_objective_funs.py::test_rebuild": 63.19959998130798, - "tests/test_objective_funs.py::test_jvp_scaled": 4.356067419052124, - "tests/test_objective_funs.py::test_objective_target_bounds": 1.8331291675567627, - "tests/test_objective_funs.py::test_jax_softmax_and_softmin": 1.0549657344818115, - "tests/test_optimizer.py::TestFmin::test_convex_full_hess_dogleg": 2.526404619216919, - "tests/test_optimizer.py::TestFmin::test_convex_full_hess_subspace": 2.8654987812042236, - "tests/test_optimizer.py::TestFmin::test_convex_full_hess_exact": 2.479990243911743, - "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_dogleg": 1.7849326133728027, - "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_subspace": 2.16491436958313, - "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_exact": 0.9446674585342407, - "tests/test_optimizer.py::TestSGD::test_sgd_convex": 1.8555312156677246, - "tests/test_optimizer.py::TestLSQTR::test_lsqtr_exact": 2.2272286415100098, - "tests/test_optimizer.py::test_no_iterations": 1.8963100910186768, - "tests/test_optimizer.py::test_overstepping": 77.34604501724243, - "tests/test_optimizer.py::test_maxiter_1_and_0_solve": 44.85545372962952, - "tests/test_optimizer.py::test_scipy_fail_message": 22.373021602630615, - "tests/test_optimizer.py::test_not_implemented_error": 0.47948431968688965, - "tests/test_optimizer.py::test_wrappers": 15.040279388427734, - "tests/test_optimizer.py::test_all_optimizers": 41.606842041015625, - "tests/test_optimizer.py::test_scipy_constrained_solve": 277.8078351020813, - "tests/test_optimizer.py::test_solve_with_x_scale": 37.2519291639328, - "tests/test_optimizer.py::test_bounded_optimization": 5.217513084411621, - "tests/test_perturbations.py::test_perturbation_orders": 953.4179646968842, - "tests/test_perturbations.py::test_perturb_with_float_without_error": 4.8092429637908936, - "tests/test_perturbations.py::test_optimal_perturb": 65.33948040008545, - "tests/test_plotting.py::test_kwarg_warning": 1.3703835010528564, - "tests/test_plotting.py::test_1d_p": 2.3615925312042236, - "tests/test_plotting.py::test_1d_fsa_consistency": 24.824483036994934, - "tests/test_plotting.py::test_1d_dpdr": 2.201986074447632, - "tests/test_plotting.py::test_1d_iota": 3.1074817180633545, - "tests/test_plotting.py::test_1d_iota_radial": 5.149364233016968, - "tests/test_plotting.py::test_1d_logpsi": 1.8761942386627197, - "tests/test_plotting.py::test_2d_logF": 6.402953386306763, - "tests/test_plotting.py::test_2d_g_tz": 3.1659910678863525, - "tests/test_plotting.py::test_2d_g_rz": 4.875414133071899, - "tests/test_plotting.py::test_2d_lambda": 2.5986695289611816, - "tests/test_plotting.py::test_3d_B": 4.344858169555664, - "tests/test_plotting.py::test_3d_J": 5.620574951171875, - "tests/test_plotting.py::test_3d_tz": 5.750120162963867, - "tests/test_plotting.py::test_3d_rz": 4.065810203552246, - "tests/test_plotting.py::test_3d_rt": 4.8052815198898315, - "tests/test_plotting.py::test_fsa_I": 5.484919428825378, - "tests/test_plotting.py::test_fsa_G": 4.416098594665527, - "tests/test_plotting.py::test_fsa_F_normalized": 5.8995115756988525, - "tests/test_plotting.py::test_section_J": 6.3657331466674805, - "tests/test_plotting.py::test_section_Z": 4.5892980098724365, - "tests/test_plotting.py::test_section_R": 2.22560715675354, - "tests/test_plotting.py::test_section_F": 5.307340860366821, - "tests/test_plotting.py::test_section_F_normalized_vac": 4.596460819244385, - "tests/test_plotting.py::test_section_logF": 4.5653135776519775, - "tests/test_plotting.py::test_plot_surfaces": 10.047749519348145, - "tests/test_plotting.py::test_plot_surfaces_no_theta": 2.538058280944824, - "tests/test_plotting.py::test_plot_boundary": 5.097471356391907, - "tests/test_plotting.py::test_plot_boundaries": 6.346785306930542, - "tests/test_plotting.py::test_plot_comparison": 12.538764119148254, - "tests/test_plotting.py::test_plot_comparison_no_theta": 3.9259798526763916, - "tests/test_plotting.py::test_plot_con_basis": 2.3857457637786865, - "tests/test_plotting.py::test_plot_cov_basis": 2.3011581897735596, - "tests/test_plotting.py::test_plot_magnetic_tension": 4.121660947799683, - "tests/test_plotting.py::test_plot_magnetic_pressure": 3.4924914836883545, - "tests/test_plotting.py::test_plot_gradpsi": 2.1703832149505615, - "tests/test_plotting.py::test_plot_normF_2d": 4.971156120300293, - "tests/test_plotting.py::test_plot_normF_section": 3.7530224323272705, - "tests/test_plotting.py::test_plot_coefficients": 2.3922486305236816, - "tests/test_plotting.py::test_plot_logo": 2.0202622413635254, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_linear": 0.9127442836761475, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_quad": 0.9276149272918701, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_jacobi": 0.9906930923461914, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb1": 0.9803032875061035, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb2": 0.9678679704666138, - "tests/test_plotting.py::TestPlotGrid::test_plot_grid_ocs": 0.8971970081329346, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_powerseries": 0.9137201309204102, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierseries": 1.0473346710205078, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_doublefourierseries": 7.4358683824539185, - "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierzernike": 24.747427105903625, - "tests/test_plotting.py::TestPlotFieldLines::test_find_idx": 0.847517728805542, - "tests/test_plotting.py::TestPlotFieldLines::test_plot_field_line": 8.131266117095947, - "tests/test_plotting.py::TestPlotFieldLines::test_plot_field_lines": 12.099355578422546, - "tests/test_plotting.py::test_plot_boozer_modes": 12.544284343719482, - "tests/test_plotting.py::test_plot_boozer_surface": 6.270490884780884, - "tests/test_plotting.py::test_plot_qs_error": 39.44042158126831, - "tests/test_plotting.py::test_plot_coils": 4.225110769271851, - "tests/test_plotting.py::test_plot_b_mag": 11.594449162483215, - "tests/test_plotting.py::test_plot_surfaces_HELIOTRON": 10.438359498977661, - "tests/test_profiles.py::TestProfiles::test_same_result": 82.2510439157486, - "tests/test_profiles.py::TestProfiles::test_close_values": 7.7879040241241455, - "tests/test_profiles.py::TestProfiles::test_repr": 0.415041446685791, - "tests/test_profiles.py::TestProfiles::test_get_set": 1.786001205444336, - "tests/test_profiles.py::TestProfiles::test_auto_sym": 0.770451545715332, - "tests/test_profiles.py::TestProfiles::test_sum_profiles": 3.679048776626587, - "tests/test_profiles.py::TestProfiles::test_product_profiles": 1.9519731998443604, - "tests/test_profiles.py::TestProfiles::test_product_profiles_derivative": 1.9519731998443604, - "tests/test_profiles.py::TestProfiles::test_scaled_profiles": 0.9374949932098389, - "tests/test_profiles.py::TestProfiles::test_profile_errors": 1.9069312810897827, - "tests/test_profiles.py::TestProfiles::test_profile_conversion": 2.3668887615203857, - "tests/test_profiles.py::TestProfiles::test_default_profiles": 1.7826738357543945, - "tests/test_profiles.py::TestProfiles::test_solve_with_combined": 28.72600507736206, - "tests/test_profiles.py::TestProfiles::test_kinetic_pressure": 2.191779136657715, - "tests/test_stability_funs.py::test_mercier_vacuum": 4.502842307090759, - "tests/test_stability_funs.py::test_compute_d_shear": 13.590898275375366, - "tests/test_stability_funs.py::test_compute_d_current": 15.9405517578125, - "tests/test_stability_funs.py::test_compute_d_well": 23.603630542755127, - "tests/test_stability_funs.py::test_compute_d_geodesic": 18.99722194671631, - "tests/test_stability_funs.py::test_compute_d_mercier": 18.758171796798706, - "tests/test_stability_funs.py::test_compute_magnetic_well": 13.469860553741455, - "tests/test_stability_funs.py::test_mercier_print": 5.0646820068359375, - "tests/test_stability_funs.py::test_magwell_print": 6.498149871826172, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_area": 3.246496081352234, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_normal": 1.4772305488586426, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_misc": 1.248878836631775, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_input_file": 3.929783344268799, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_near_axis": 1.233553409576416, - "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_curvature": 1.0212721824645996, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_area": 3.246496081352234, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_normal": 1.4772305488586426, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_misc": 1.248878836631775, - "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_curvature": 1.0212721824645996, - "tests/test_surfaces.py::test_surface_orientation": 2.6663854122161865, - "tests/test_transform.py::TestTransform::test_eq": 1.1583225727081299, - "tests/test_transform.py::TestTransform::test_transform_order_error": 0.8315334320068359, - "tests/test_transform.py::TestTransform::test_profile": 0.9312019348144531, - "tests/test_transform.py::TestTransform::test_surface": 1.6691875457763672, - "tests/test_transform.py::TestTransform::test_volume": 1.670881986618042, - "tests/test_transform.py::TestTransform::test_set_grid": 1.1346402168273926, - "tests/test_transform.py::TestTransform::test_set_basis": 1.1718149185180664, - "tests/test_transform.py::TestTransform::test_fft": 1.367253065109253, - "tests/test_transform.py::TestTransform::test_direct_fft_equal": 5.504265308380127, - "tests/test_transform.py::TestTransform::test_project": 3.4143770933151245, - "tests/test_transform.py::TestTransform::test_fft_warnings": 1.367253065109253, - "tests/test_transform.py::TestTransform::test_direct2_warnings": 1.4736188650131226, - "tests/test_transform.py::TestTransform::test_fit_direct1": 2.491840362548828, - "tests/test_transform.py::TestTransform::test_fit_direct2": 1.5647196769714355, - "tests/test_transform.py::TestTransform::test_fit_fft": 1.7684881687164307, - "tests/test_transform.py::TestTransform::test_empty_grid": 1.4022634029388428, - "tests/test_transform.py::TestTransform::test_Z_projection": 3.716538190841675, - "tests/test_utils.py::test_isalmostequal": 0.6772017478942871, - "tests/test_utils.py::test_islinspaced": 0.6764607429504395, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd": 0.8442978858947754, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd_sin_series": 0.7731931209564209, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev": 0.8971633911132812, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev_sin_sym": 0.7842483520507812, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_identities_inverse": 0.7787743806838989, - "tests/test_vmec.py::TestVMECIO::test_ptolemy_linear_transform": 1.0929808616638184, - "tests/test_vmec.py::TestVMECIO::test_fourier_to_zernike": 1.564780831336975, - "tests/test_vmec.py::TestVMECIO::test_zernike_to_fourier": 1.6053714752197266, - "tests/test_vmec.py::test_load_then_save": 45.55219221115112, - "tests/test_vmec.py::test_vmec_save_asym": 30.782756328582764, - "tests/test_vmec.py::test_vmec_save_1": 1.1307034492492676, - "tests/test_vmec.py::test_vmec_save_2": 1.2249751091003418, - "tests/test_vmec.py::test_plot_vmec_comparison": 7.043433904647827, - "tests/test_vmec.py::test_vmec_boundary_subspace": 1.366995096206665 + "tests/test_axis_limits.py::TestAxisLimits::test_axis_limit_api": 1.0464872860029573, + "tests/test_axis_limits.py::TestAxisLimits::test_limit_continuity": 27.198641567003506, + "tests/test_backend.py::test_put": 0.26291312499233754, + "tests/test_backend.py::test_sign": 0.3399597820098279, + "tests/test_backend.py::test_vmap": 0.2981577630052925, + "tests/test_basis.py::TestBasis::test_basis_resolutions_assert_integers": 0.22994619000382954, + "tests/test_basis.py::TestBasis::test_change_resolution": 0.2459372000012081, + "tests/test_basis.py::TestBasis::test_chebyshev": 0.29989392299467, + "tests/test_basis.py::TestBasis::test_derivative_not_in_basis_zeros": 0.4326746090009692, + "tests/test_basis.py::TestBasis::test_double_fourier": 0.32092916200781474, + "tests/test_basis.py::TestBasis::test_fourier": 0.30613131100835744, + "tests/test_basis.py::TestBasis::test_polyder": 0.24179753699718276, + "tests/test_basis.py::TestBasis::test_polyval": 0.2531717070014565, + "tests/test_basis.py::TestBasis::test_polyval_exact": 3.324124009006482, + "tests/test_basis.py::TestBasis::test_powers": 0.22775781300879316, + "tests/test_basis.py::TestBasis::test_power_series": 0.2452292009984376, + "tests/test_basis.py::TestBasis::test_repr": 0.2417543640040094, + "tests/test_basis.py::TestBasis::test_zernike_coeffs": 0.2811490600070101, + "tests/test_basis.py::TestBasis::test_zernike_indexing": 0.24884329400811112, + "tests/test_basis.py::TestBasis::test_zernike_radial": 0.24388470799749484, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_2_3": 0.9647520499929669, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_figures_4_5": 1.8681837440090021, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_second_pass": 0.5326963499974227, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QA": 3.8826869990080013, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_QH": 1.972220620009466, + "tests/test_bootstrap.py::TestBootstrapCompute::test_Redl_sfincs_tokamak_benchmark": 3.0387523969984613, + "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_analytic": 1.6943289549963083, + "tests/test_bootstrap.py::TestBootstrapCompute::test_trapped_fraction_Kim": 5.952490649993706, + "tests/test_bootstrap.py::test_bootstrap_objective_build": 8.188893187005306, + "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_current": 133.29894527600845, + "tests/test_bootstrap.py::TestBootstrapObjectives::test_bootstrap_consistency_iota": 168.88556157100538, + "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_normalization": 13.785245222003141, + "tests/test_bootstrap.py::TestBootstrapObjectives::test_BootstrapRedlConsistency_resolution": 94.1242997080044, + "tests/test_coils.py::TestCoilSet::test_dunder_methods": 0.6621696810034337, + "tests/test_coils.py::TestCoilSet::test_from_symmetry": 2.789724836999085, + "tests/test_coils.py::TestCoilSet::test_linspaced_angular": 2.2386740690053557, + "tests/test_coils.py::TestCoilSet::test_linspaced_linear": 1.3177181810024194, + "tests/test_coils.py::TestCoilSet::test_properties": 1.3532936060000793, + "tests/test_coils.py::TestCoil::test_biot_savart": 1.119896318006795, + "tests/test_coils.py::TestCoil::test_properties": 0.5668602200094028, + "tests/test_coils.py::test_repr": 0.6105128710041754, + "tests/test_compute_funs.py::test_BdotgradB": 4.226151214003039, + "tests/test_compute_funs.py::test_boozer_transform": 1.326571810997848, + "tests/test_compute_funs.py::test_compare_quantities_to_vmec": 9.29286608100665, + "tests/test_compute_funs.py::test_compute_averages": 4.706544874003157, + "tests/test_compute_funs.py::test_compute_grad_p_volume_avg": 0.792815654000151, + "tests/test_compute_funs.py::test_currents": 4.251725698006339, + "tests/test_compute_funs.py::test_enclosed_volumes": 2.006099311998696, + "tests/test_compute_funs.py::test_compute_everything": 18.26, + "tests/test_compute_funs.py::test_magnetic_field_derivatives": 9.015901343002042, + "tests/test_compute_funs.py::test_magnetic_pressure_gradient": 3.835518025000056, + "tests/test_compute_funs.py::test_metric_derivatives": 1.6676708760060137, + "tests/test_compute_funs.py::test_surface_areas": 1.1399255659998744, + "tests/test_compute_funs.py::test_surface_areas_2": 1.1644101690035313, + "tests/test_compute_funs.py::test_total_volume": 2.900268169003539, + "tests/test_compute_utils.py::TestComputeUtils::test_line_length": 4.150421753001865, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_area": 1.9687216179954703, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_against_shortcut": 2.2439278719903086, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_homomorphism": 1.7925866679943283, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_identity_op": 1.230343893999816, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_averages_vector_functions": 0.8377738929921179, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals": 11.814357757000835, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals_against_shortcut": 0.3332841070005088, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_integrals_transform": 0.4919979200058151, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_min_max": 0.5911533130056341, + "tests/test_compute_utils.py::TestComputeUtils::test_surface_variance": 0.6174815310005215, + "tests/test_compute_utils.py::TestComputeUtils::test_symmetry_surface_average_1": 3.6187084570046864, + "tests/test_compute_utils.py::TestComputeUtils::test_symmetry_surface_average_2": 1.750842968998768, + "tests/test_compute_utils.py::TestDataIndex::test_data_index_deps": 1.4365679520051344, + "tests/test_configuration.py::TestConstructor::test_asserts": 1.4832063389985706, + "tests/test_configuration.py::TestConstructor::test_defaults": 0.8779114760036464, + "tests/test_configuration.py::TestConstructor::test_dict": 1.3487972259972594, + "tests/test_configuration.py::TestConstructor::test_supplied_coeffs": 0.907749498001067, + "tests/test_configuration.py::TestConstructor::test_supplied_objects": 1.648549320991151, + "tests/test_configuration.py::test_get_profile": 2.733136356000614, + "tests/test_configuration.py::TestGetSurfaces::test_asserts": 0.8378489299866487, + "tests/test_configuration.py::TestGetSurfaces::test_get_rho_surface": 1.0137959559942828, + "tests/test_configuration.py::TestGetSurfaces::test_get_theta_surface": 0.8322402060002787, + "tests/test_configuration.py::TestGetSurfaces::test_get_zeta_surface": 1.0238518139958614, + "tests/test_configuration.py::TestInitialGuess::test_default_set": 0.929045880009653, + "tests/test_configuration.py::TestInitialGuess::test_errors": 1.124446292000357, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_file": 1.6884469270080444, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_other": 1.0910271920001833, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_points": 2.034145499994338, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface": 1.0073084750038106, + "tests/test_configuration.py::TestInitialGuess::test_guess_from_surface2": 1.0473684090029565, + "tests/test_configuration.py::TestInitialGuess::test_NFP_error": 0.9449550400022417, + "tests/test_configuration.py::test_is_nested": 1.855856258996937, + "tests/test_configuration.py::test_is_nested_theta": 1.7979977849900024, + "tests/test_configuration.py::test_kinetic_errors": 1.6062586749976617, + "tests/test_configuration.py::test_magnetic_axis": 315.59931112699996, + "tests/test_constrain_current.py::TestConstrainCurrent::test_current_to_iota_and_back": 320.77903574299125, + "tests/test_constrain_current.py::TestConstrainCurrent::test_iota_to_current_and_back": 71.87296323300689, + "tests/test_curves.py::TestPlanarCurve::test_asserts": 0.9115305340019404, + "tests/test_curves.py::TestPlanarCurve::test_coords": 0.9283086049981648, + "tests/test_curves.py::TestPlanarCurve::test_curvature": 0.9611980159970699, + "tests/test_curves.py::TestPlanarCurve::test_frenet": 0.9492236939986469, + "tests/test_curves.py::TestPlanarCurve::test_length": 1.0377524949944927, + "tests/test_curves.py::TestPlanarCurve::test_misc": 0.9228372060024412, + "tests/test_curves.py::TestPlanarCurve::test_torsion": 0.9288547650066903, + "tests/test_curves.py::TestRZCurve::test_asserts": 0.8825806309978361, + "tests/test_curves.py::TestRZCurve::test_coords": 0.9266123410052387, + "tests/test_curves.py::TestRZCurve::test_curvature": 1.0539713780017337, + "tests/test_curves.py::TestRZCurve::test_frenet": 0.9682074919983279, + "tests/test_curves.py::TestRZCurve::test_length": 1.3079071410029428, + "tests/test_curves.py::TestRZCurve::test_misc": 1.056210798007669, + "tests/test_curves.py::TestRZCurve::test_to_FourierXYZCurve": 2.460215673992934, + "tests/test_curves.py::TestRZCurve::test_torsion": 1.068913770999643, + "tests/test_curves.py::TestXYZCurve::test_asserts": 0.9161841050008661, + "tests/test_curves.py::TestXYZCurve::test_coords": 0.9731651219990454, + "tests/test_curves.py::TestXYZCurve::test_curvature": 0.9922603250015527, + "tests/test_curves.py::TestXYZCurve::test_frenet": 1.0665025709968177, + "tests/test_curves.py::TestXYZCurve::test_length": 1.104770219004422, + "tests/test_curves.py::TestXYZCurve::test_misc": 1.0032778279928607, + "tests/test_curves.py::TestXYZCurve::test_torsion": 1.0174006340021151, + "tests/test_derivatives.py::TestDerivative::test_auto_diff": 0.9765123139950447, + "tests/test_derivatives.py::TestDerivative::test_block_jacobian": 1.034016514000541, + "tests/test_derivatives.py::TestDerivative::test_compare_AD_FD": 0.8883194609952625, + "tests/test_derivatives.py::TestDerivative::test_fd_hessian": 0.8966648439964047, + "tests/test_derivatives.py::TestDerivative::test_finite_diff_scalar": 0.9362148300060653, + "tests/test_derivatives.py::TestDerivative::test_finite_diff_vec": 0.9248911340037012, + "tests/test_derivatives.py::TestDerivative::test_jac_looped": 0.9687139650050085, + "tests/test_derivatives.py::TestJVP::test_autodiff_jvp": 1.0231309980081278, + "tests/test_derivatives.py::TestJVP::test_autodiff_jvp2": 1.073198546997446, + "tests/test_derivatives.py::TestJVP::test_autodiff_jvp3": 1.0986065899924142, + "tests/test_derivatives.py::TestJVP::test_finitediff_jvp": 0.9253449819952948, + "tests/test_derivatives.py::TestJVP::test_finitediff_jvp2": 0.900767724000616, + "tests/test_derivatives.py::TestJVP::test_finitediff_jvp3": 0.9104889360023662, + "tests/test_derivatives.py::TestJVP::test_vjp": 1.0234685109971906, + "tests/test_equilibrium.py::test_change_NFP": 2.898086309003702, + "tests/test_equilibrium.py::test_compute_flux_coords": 23.83571945699805, + "tests/test_equilibrium.py::test_compute_geometry": 1.4316535710022436, + "tests/test_equilibrium.py::test_compute_theta_coords": 2.5141915600033826, + "tests/test_equilibrium.py::test_continuation_resolution": 46.93961339700763, + "tests/test_equilibrium.py::test_eq_change_grid_resolution": 1.5142954219918465, + "tests/test_equilibrium.py::test_eq_change_symmetry": 2.128670626996609, + "tests/test_equilibrium.py::test_equilibriafamily_constructor": 2.471822694998991, + "tests/test_equilibrium.py::test_equilibrium_from_near_axis": 12.983962494996376, + "tests/test_equilibrium.py::test_grid_resolution_warning": 47.99903397500748, + "tests/test_equilibrium.py::test_map_coordinates2": 14.646045742010756, + "tests/test_equilibrium.py::test_map_coordinates": 23.662408988995594, + "tests/test_equilibrium.py::test_poincare_solve_not_implemented": 1.2679396940002334, + "tests/test_equilibrium.py::test_resolution": 4.065810970991151, + "tests/test_equilibrium.py::test_to_sfl": 12.789840172998083, + "tests/test_examples.py::test_1d_optimization": 35.35925663400121, + "tests/test_examples.py::test_1d_optimization_old": 89.90233920099854, + "tests/test_examples.py::test_ATF_results": 523.5010129150032, + "tests/test_examples.py::test_DSHAPE_current_results": 47.86380493100296, + "tests/test_examples.py::test_DSHAPE_results": 4.265998300994397, + "tests/test_examples.py::test_ESTELL_results": 1526.5112719009994, + "tests/test_examples.py::test_force_balance_grids": 33.44704783799534, + "tests/test_examples.py::TestGetExample::test_example_get_boundary": 3.0332657730032224, + "tests/test_examples.py::TestGetExample::test_example_get_current": 2.5184027979921666, + "tests/test_examples.py::TestGetExample::test_example_get_eq": 1.5714667919964995, + "tests/test_examples.py::TestGetExample::test_example_get_eqf": 1.7227903209932265, + "tests/test_examples.py::TestGetExample::test_example_get_iota": 3.125554414007638, + "tests/test_examples.py::TestGetExample::test_example_get_pressure": 2.2749817880030605, + "tests/test_examples.py::TestGetExample::test_missing_example": 1.0887831300060498, + "tests/test_examples.py::test_HELIOTRON_results": 344.15801551400364, + "tests/test_examples.py::test_HELIOTRON_vac2_results": 396.8973683439908, + "tests/test_examples.py::test_HELIOTRON_vac_results": 5.1467031560023315, + "tests/test_examples.py::test_NAE_QIC_solve": 196.8964859569969, + "tests/test_examples.py::test_NAE_QSC_solve": 128.90866155899857, + "tests/test_examples.py::test_precise_QH_results": 14.028632663997996, + "tests/test_examples.py::test_qh_optimization1": 788.0094411460013, + "tests/test_examples.py::test_qh_optimization2": 2276.018544077997, + "tests/test_examples.py::test_qh_optimization3": 906.6888439869945, + "tests/test_examples.py::test_simsopt_QH_comparison": 476.23041609398933, + "tests/test_examples.py::test_SOLOVEV_results": 19.63613601900579, + "tests/test_examples.py::test_SOLOVEV_vacuum": 24.96038271800353, + "tests/test_examples.py::test_solve_bounds": 21.704924048994144, + "tests/test_geometry.py::test_rotation_matrix": 1.1176765419950243, + "tests/test_geometry.py::test_rpz2xyz": 1.1326016200036975, + "tests/test_geometry.py::test_xyz2rpz": 1.1871705459998338, + "tests/test_grid.py::test_find_least_rational_surfaces": 1.8505934839995462, + "tests/test_grid.py::test_find_most_rational_surfaces": 1.097906677008723, + "tests/test_grid.py::TestGrid::test_change_resolution": 1.094492779011489, + "tests/test_grid.py::TestGrid::test_compress_expand_inverse_op": 1.1097727920059697, + "tests/test_grid.py::TestGrid::test_concentric_grid": 1.076786325997091, + "tests/test_grid.py::TestGrid::test_concentric_grid_high_res": 1.124056073007523, + "tests/test_grid.py::TestGrid::test_custom_grid": 1.1073047609970672, + "tests/test_grid.py::TestGrid::test_enforce_symmetry": 1.1002924089989392, + "tests/test_grid.py::TestGrid::test_linear_grid": 1.0651826030007214, + "tests/test_grid.py::TestGrid::test_linear_grid_spacing_consistency": 1.093027342998539, + "tests/test_grid.py::TestGrid::test_linear_grid_spacing_two_nodes": 1.0554385290015489, + "tests/test_grid.py::TestGrid::test_linear_grid_symmetric_nodes_consistency": 1.1000499370056787, + "tests/test_grid.py::TestGrid::test_node_spacing_non_sym": 1.2667479969968554, + "tests/test_grid.py::TestGrid::test_node_spacing_sym": 1.2707626549963607, + "tests/test_grid.py::TestGrid::test_quad_grid_volume_integration": 1.427309056991362, + "tests/test_grid.py::TestGrid::test_quadrature_grid": 1.0930434850015445, + "tests/test_grid.py::TestGrid::test_repr": 1.0830842589930398, + "tests/test_grid.py::TestGrid::test_spacing_when_duplicate_node_is_removed": 1.1066608410037588, + "tests/test_grid.py::TestGrid::test_symmetry_spacing_basic": 1.1336049260062282, + "tests/test_grid.py::TestGrid::test_symmetry_volume_integral": 2.514414705998206, + "tests/test_input_output.py::test_ascii_io": 5.198659340006998, + "tests/test_input_output.py::test_copy": 1.3791897580085788, + "tests/test_input_output.py::TestInputReader::test_min_input": 1.0578299679982592, + "tests/test_input_output.py::TestInputReader::test_no_input_file": 1.0799451610073447, + "tests/test_input_output.py::TestInputReader::test_nonexistant_input_file": 1.0830367329981527, + "tests/test_input_output.py::TestInputReader::test_np_environ": 1.0720524069984094, + "tests/test_input_output.py::TestInputReader::test_quiet_verbose": 1.0945736310022767, + "tests/test_input_output.py::TestInputReader::test_vacuum_objective_with_iota_yields_current": 1.0986467959955917, + "tests/test_input_output.py::test_load_eq_without_current": 1.6469790199989802, + "tests/test_input_output.py::test_near_axis_input_files": 4.51630549300171, + "tests/test_input_output.py::test_pickle_io": 1.4404222100056359, + "tests/test_input_output.py::test_reader_given_file": 1.093506653996883, + "tests/test_input_output.py::test_reader_given_filename": 1.1337760599926696, + "tests/test_input_output.py::test_reader_read_obj": 1.1413002549961675, + "tests/test_input_output.py::test_save_none": 1.428643749008188, + "tests/test_input_output.py::test_vmec_input": 1.203770918000373, + "tests/test_input_output.py::test_vmec_input_surface_threshold": 3.6363580570032354, + "tests/test_input_output.py::test_write_desc_input_Nones": 1.0831025239967857, + "tests/test_input_output.py::test_writer_close_on_delete": 1.1210618789991713, + "tests/test_input_output.py::test_writer_given_file": 1.1198129140029778, + "tests/test_input_output.py::test_writer_given_filename": 1.130466985006933, + "tests/test_input_output.py::test_writer_write_dict": 1.0748073310023756, + "tests/test_input_output.py::test_writer_write_list": 1.0967376879998483, + "tests/test_input_output.py::test_writer_write_obj": 1.1515619350102497, + "tests/test_interpolate.py::TestInterp1D::test_interp1d": 2.785034660992096, + "tests/test_interpolate.py::TestInterp1D::test_interp1d_extrap_periodic": 2.0204221369931474, + "tests/test_interpolate.py::TestInterp1D::test_interp1d_monotonic": 2.1549058759919717, + "tests/test_interpolate.py::TestInterp2D::test_interp2d": 3.111412221995124, + "tests/test_interpolate.py::TestInterp3D::test_interp3d": 4.214047164990916, + "tests/test_linear_objectives.py::test_bc_on_interior_surfaces": 9.72851990700292, + "tests/test_linear_objectives.py::test_build_init": 2.3891700799940736, + "tests/test_linear_objectives.py::test_constrain_asserts": 4.155162373994244, + "tests/test_linear_objectives.py::test_constrain_bdry_with_only_one_mode": 1.1757937200018205, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes": 19.727867216002778, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target": 15.954220396997698, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_and_passed_target_axis": 6.309178292998695, + "tests/test_linear_objectives.py::test_correct_indexing_passed_modes_axis": 9.230813657006365, + "tests/test_linear_objectives.py::test_factorize_linear_constraints_asserts": 2.1265679609932704, + "tests/test_linear_objectives.py::test_FixAxis_passed_target_no_passed_modes_error": 1.2264831729989965, + "tests/test_linear_objectives.py::test_FixAxis_util_correct_objectives": 1.1656612229926395, + "tests/test_linear_objectives.py::test_FixBoundary_passed_target_no_passed_modes_error": 1.2405282569961855, + "tests/test_linear_objectives.py::test_FixBoundary_with_single_weight": 1.6325562220008578, + "tests/test_linear_objectives.py::test_fixed_axis_and_theta_SFL_solve": 13.493089123003301, + "tests/test_linear_objectives.py::test_fixed_mode_solve": 18.450932453008136, + "tests/test_linear_objectives.py::test_fixed_modes_solve": 12.158118179999292, + "tests/test_linear_objectives.py::test_FixMode_False_or_None_modes": 1.1487968879955588, + "tests/test_linear_objectives.py::test_FixMode_passed_target_no_passed_modes_error": 1.2425332969869487, + "tests/test_linear_objectives.py::test_FixNAE_util_correct_objectives": 1.6786472919993685, + "tests/test_linear_objectives.py::test_FixSumModes_False_or_None_modes": 1.159057088996633, + "tests/test_linear_objectives.py::test_FixSumModes_passed_target_too_long": 1.1688318220112706, + "tests/test_linear_objectives.py::test_kinetic_constraints": 1.6900724829974934, + "tests/test_linear_objectives.py::test_LambdaGauge_asym": 1.5686675739998464, + "tests/test_linear_objectives.py::test_LambdaGauge_sym": 1.527167931009899, + "tests/test_magnetic_fields.py::TestMagneticFields::test_basic_fields": 1.4081961929914542, + "tests/test_magnetic_fields.py::TestMagneticFields::test_field_line_integrate": 1.6778842540079495, + "tests/test_magnetic_fields.py::TestMagneticFields::test_scalar_field": 1.65062886600208, + "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field": 5.274317914998392, + "tests/test_magnetic_fields.py::TestMagneticFields::test_spline_field_axisym": 2.402778241004853, + "tests/test_objective_funs.py::test_bounds_format": 1.3528379410054185, + "tests/test_objective_funs.py::test_derivative_modes": 69.70801763399504, + "tests/test_objective_funs.py::test_field_scale_length": 18.89917086299829, + "tests/test_objective_funs.py::test_generic_compute": 1.7866184119993704, + "tests/test_objective_funs.py::test_getter_setter": 1.439185113995336, + "tests/test_objective_funs.py::test_jax_softmax_and_softmin": 1.7536369260051288, + "tests/test_objective_funs.py::test_jvp_scaled": 2.9055520399997476, + "tests/test_objective_funs.py::test_mean_curvature": 5.022188588998688, + "tests/test_objective_funs.py::TestObjectiveFunction::test_aspect_ratio": 1.8728012360079447, + "tests/test_objective_funs.py::TestObjectiveFunction::test_elongation": 2.9014134329845547, + "tests/test_objective_funs.py::TestObjectiveFunction::test_energy": 2.1287094459985383, + "tests/test_objective_funs.py::TestObjectiveFunction::test_generic": 2.347110872004123, + "tests/test_objective_funs.py::TestObjectiveFunction::test_isodynamicity": 2.1565887089964235, + "tests/test_objective_funs.py::TestObjectiveFunction::test_magnetic_well": 2.7287971429977915, + "tests/test_objective_funs.py::TestObjectiveFunction::test_mercier_stability": 3.942152444004023, + "tests/test_objective_funs.py::TestObjectiveFunction::test_objective_from_user": 1.4539140529959695, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qa_boozer": 5.621530480006186, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qh_boozer": 8.190650623000693, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_boozer_grids": 2.3479965919905226, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_tripleproduct": 4.811457171999791, + "tests/test_objective_funs.py::TestObjectiveFunction::test_qs_twoterm": 2.9408952170051634, + "tests/test_objective_funs.py::TestObjectiveFunction::test_target_iota": 1.8612304310008767, + "tests/test_objective_funs.py::TestObjectiveFunction::test_toroidal_current": 2.2384538919941406, + "tests/test_objective_funs.py::TestObjectiveFunction::test_volume": 1.8912741800013464, + "tests/test_objective_funs.py::test_objective_target_bounds": 2.1945660939964, + "tests/test_objective_funs.py::test_plasma_vessel_distance": 6.458658501993341, + "tests/test_objective_funs.py::test_plasma_vessel_distance_print": 2.50811629799864, + "tests/test_objective_funs.py::test_principal_curvature": 4.480132026990759, + "tests/test_objective_funs.py::test_profile_objective_print": 3.439420919996337, + "tests/test_objective_funs.py::test_rebuild": 32.801968941996165, + "tests/test_objective_funs.py::test_rejit": 1.6128012940098415, + "tests/test_objective_funs.py::test_target_profiles": 2.6087243350048084, + "tests/test_objective_funs.py::test_vjp": 16.296108381000522, + "tests/test_optimizer.py::test_all_optimizers": 32.81177038900205, + "tests/test_optimizer.py::test_auglag": 20.11058654500812, + "tests/test_optimizer.py::test_bounded_optimization": 3.22657639899262, + "tests/test_optimizer.py::test_constrained_AL_lsq": 33.22712204700656, + "tests/test_optimizer.py::test_constrained_AL_scalar": 140.65198725400114, + "tests/test_optimizer.py::TestFmin::test_convex_full_hess_dogleg": 2.6184097320001456, + "tests/test_optimizer.py::TestFmin::test_convex_full_hess_exact": 1.5315024910014472, + "tests/test_optimizer.py::TestFmin::test_convex_full_hess_subspace": 1.8314242110063788, + "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_dogleg": 2.8757154310005717, + "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_exact": 1.6096570660010912, + "tests/test_optimizer.py::TestFmin::test_rosenbrock_bfgs_subspace": 1.8989509789971635, + "tests/test_optimizer.py::TestLSQTR::test_lsqtr_exact": 3.123739970003953, + "tests/test_optimizer.py::test_maxiter_1_and_0_solve": 25.147945896002057, + "tests/test_optimizer.py::test_no_iterations": 2.776398852991406, + "tests/test_optimizer.py::test_not_implemented_error": 1.4681641200077138, + "tests/test_optimizer.py::test_overstepping": 186.77352349199646, + "tests/test_optimizer.py::test_scipy_constrained_solve": 211.6421550160012, + "tests/test_optimizer.py::test_scipy_fail_message": 15.634397530004208, + "tests/test_optimizer.py::TestSGD::test_sgd_convex": 2.09538175000489, + "tests/test_optimizer.py::test_solve_with_x_scale": 19.66171781401499, + "tests/test_optimizer.py::test_wrappers": 9.776084370991157, + "tests/test_perturbations.py::test_optimal_perturb": 28.115295303003222, + "tests/test_perturbations.py::test_perturbation_orders": 188.6908808980079, + "tests/test_perturbations.py::test_perturb_axis": 17.513416502006294, + "tests/test_perturbations.py::test_perturb_with_float_without_error": 3.92209741500119, + "tests/test_plotting.py::test_1d_dpdr": 2.2400733989998116, + "tests/test_plotting.py::test_1d_fsa_consistency": 22.040991275003762, + "tests/test_plotting.py::test_1d_iota": 3.370761713995307, + "tests/test_plotting.py::test_1d_iota_radial": 4.615922213000886, + "tests/test_plotting.py::test_1d_logpsi": 2.1383116389915813, + "tests/test_plotting.py::test_1d_p": 2.214763081006822, + "tests/test_plotting.py::test_2d_g_rz": 2.596926157006237, + "tests/test_plotting.py::test_2d_g_tz": 2.5065802530007204, + "tests/test_plotting.py::test_2d_lambda": 2.6026459350032383, + "tests/test_plotting.py::test_2d_logF": 6.397114165003586, + "tests/test_plotting.py::test_3d_B": 3.2318124579978758, + "tests/test_plotting.py::test_3d_J": 3.591661351005314, + "tests/test_plotting.py::test_3d_rt": 2.5911154890054604, + "tests/test_plotting.py::test_3d_rz": 2.692471773007128, + "tests/test_plotting.py::test_3d_tz": 3.5040650200025993, + "tests/test_plotting.py::test_fsa_F_normalized": 5.3021594420060865, + "tests/test_plotting.py::test_fsa_G": 2.885779743999592, + "tests/test_plotting.py::test_fsa_I": 5.236025021004025, + "tests/test_plotting.py::test_kwarg_future_warning": 4.8477915990006295, + "tests/test_plotting.py::test_kwarg_warning": 2.1860916759906104, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_doublefourierseries": 4.149428076998447, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierseries": 1.9759187230083626, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_fourierzernike": 4.060661367991997, + "tests/test_plotting.py::TestPlotBasis::test_plot_basis_powerseries": 1.8642975820039283, + "tests/test_plotting.py::test_plot_b_mag": 7.586959128995659, + "tests/test_plotting.py::test_plot_boozer_modes": 17.22079940600088, + "tests/test_plotting.py::test_plot_boozer_surface": 5.794906193994393, + "tests/test_plotting.py::test_plot_boundaries": 5.867537543999788, + "tests/test_plotting.py::test_plot_boundary": 4.464614589000121, + "tests/test_plotting.py::test_plot_coefficients": 2.2505050100007793, + "tests/test_plotting.py::test_plot_coils": 4.563949970004614, + "tests/test_plotting.py::test_plot_comparison": 4.837778633998823, + "tests/test_plotting.py::test_plot_comparison_no_theta": 2.7345728670043172, + "tests/test_plotting.py::test_plot_con_basis": 2.414726741008053, + "tests/test_plotting.py::test_plot_cov_basis": 2.2958074069974828, + "tests/test_plotting.py::TestPlotFieldLines::test_find_idx": 1.9428987050123396, + "tests/test_plotting.py::TestPlotFieldLines::test_plot_field_line": 4.67791478500294, + "tests/test_plotting.py::TestPlotFieldLines::test_plot_field_lines": 4.8310439880006015, + "tests/test_plotting.py::test_plot_fsa_axis_limit": 13.354158946000098, + "tests/test_plotting.py::test_plot_gradpsi": 2.3928257479929016, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb1": 1.9617451439989964, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_cheb2": 1.9928445369951078, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_jacobi": 1.9893377359985607, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_linear": 1.969201037005405, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_ocs": 1.9722542530071223, + "tests/test_plotting.py::TestPlotGrid::test_plot_grid_quad": 1.9618056860053912, + "tests/test_plotting.py::test_plot_logo": 2.8646091110131238, + "tests/test_plotting.py::test_plot_magnetic_pressure": 3.0180758710048394, + "tests/test_plotting.py::test_plot_magnetic_tension": 3.2902509289924637, + "tests/test_plotting.py::test_plot_normF_2d": 5.089891604002332, + "tests/test_plotting.py::test_plot_normF_section": 4.5190368670009775, + "tests/test_plotting.py::test_plot_qs_error": 47.9670878859979, + "tests/test_plotting.py::test_plot_surfaces": 6.834042420996411, + "tests/test_plotting.py::test_plot_surfaces_HELIOTRON": 14.414476540994656, + "tests/test_plotting.py::test_plot_surfaces_no_theta": 2.4581562590028625, + "tests/test_plotting.py::test_section_F": 3.33160220800346, + "tests/test_plotting.py::test_section_F_normalized_vac": 4.956865172003745, + "tests/test_plotting.py::test_section_J": 6.029406550995191, + "tests/test_plotting.py::test_section_logF": 3.166178615996614, + "tests/test_plotting.py::test_section_R": 2.4493891010060906, + "tests/test_plotting.py::test_section_Z": 2.3918421199996374, + "tests/test_profiles.py::TestProfiles::test_auto_sym": 1.9677782630096772, + "tests/test_profiles.py::TestProfiles::test_close_values": 6.00612420100515, + "tests/test_profiles.py::TestProfiles::test_default_profiles": 2.5178508079989115, + "tests/test_profiles.py::TestProfiles::test_get_set": 2.279912530008005, + "tests/test_profiles.py::TestProfiles::test_kinetic_pressure": 2.463389439995808, + "tests/test_profiles.py::TestProfiles::test_product_profiles": 2.213361353002256, + "tests/test_profiles.py::TestProfiles::test_product_profiles_derivative": 2.0030428019963438, + "tests/test_profiles.py::TestProfiles::test_profile_conversion": 2.923352846999478, + "tests/test_profiles.py::TestProfiles::test_profile_errors": 2.4316903650033055, + "tests/test_profiles.py::TestProfiles::test_repr": 2.4147915160065168, + "tests/test_profiles.py::TestProfiles::test_same_result": 80.89257165599702, + "tests/test_profiles.py::TestProfiles::test_scaled_profiles": 2.0555991330038523, + "tests/test_profiles.py::TestProfiles::test_solve_with_combined": 20.07468461299868, + "tests/test_profiles.py::TestProfiles::test_SplineProfile_methods": 3.1647336239984725, + "tests/test_profiles.py::TestProfiles::test_sum_profiles": 2.9485370149923256, + "tests/test_stability_funs.py::test_compute_d_current": 11.241232370994112, + "tests/test_stability_funs.py::test_compute_d_geodesic": 16.7818530020013, + "tests/test_stability_funs.py::test_compute_d_mercier": 18.654368460003752, + "tests/test_stability_funs.py::test_compute_d_shear": 9.922470428995439, + "tests/test_stability_funs.py::test_compute_d_well": 13.664260720994207, + "tests/test_stability_funs.py::test_compute_magnetic_well": 11.858983706995787, + "tests/test_stability_funs.py::test_magwell_print": 6.690227218001382, + "tests/test_stability_funs.py::test_mercier_print": 5.474188908003271, + "tests/test_stability_funs.py::test_mercier_vacuum": 3.184365107001213, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_area": 4.135052845995233, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_curvature": 2.32408962000045, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_input_file": 4.449585768990801, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_from_near_axis": 2.4458258809972904, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_misc": 2.970418883007369, + "tests/test_surfaces.py::TestFourierRZToroidalSurface::test_normal": 2.714608484006021, + "tests/test_surfaces.py::test_surface_orientation": 3.401730249992397, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_area": 3.797193332000461, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_curvature": 2.297793673002161, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_misc": 2.702328484003374, + "tests/test_surfaces.py::TestZernikeRZToroidalSection::test_normal": 2.8036721339958603, + "tests/test_transform.py::test_transform_pytree": 3.133346861002792, + "tests/test_transform.py::TestTransform::test_direct2_warnings": 2.1174599409932853, + "tests/test_transform.py::TestTransform::test_direct_fft_equal": 6.308722748006403, + "tests/test_transform.py::TestTransform::test_empty_grid": 2.5340789450056036, + "tests/test_transform.py::TestTransform::test_eq": 2.3977073869973538, + "tests/test_transform.py::TestTransform::test_fft": 2.7988818130033906, + "tests/test_transform.py::TestTransform::test_fft_warnings": 3.2662716299964814, + "tests/test_transform.py::TestTransform::test_fit_direct1": 3.2259791489923373, + "tests/test_transform.py::TestTransform::test_fit_direct2": 2.6846433930040803, + "tests/test_transform.py::TestTransform::test_fit_fft": 3.023515369000961, + "tests/test_transform.py::TestTransform::test_profile": 2.18462181299401, + "tests/test_transform.py::TestTransform::test_project": 4.3892835219958215, + "tests/test_transform.py::TestTransform::test_set_basis": 2.4614632760058157, + "tests/test_transform.py::TestTransform::test_set_grid": 2.356843399000354, + "tests/test_transform.py::TestTransform::test_surface": 2.839725963996898, + "tests/test_transform.py::TestTransform::test_transform_order_error": 2.152247540994722, + "tests/test_transform.py::TestTransform::test_volume_chebyshev": 2.6900079129991354, + "tests/test_transform.py::TestTransform::test_volume_zernike": 2.239503749005962, + "tests/test_transform.py::TestTransform::test_Z_projection": 3.2358590119984, + "tests/test_utils.py::test_isalmostequal": 2.100865389998944, + "tests/test_utils.py::test_islinspaced": 2.143538215997978, + "tests/test_vmec.py::test_axis_surf_after_load": 11.875969582004473, + "tests/test_vmec.py::test_load_then_save": 36.62609588299529, + "tests/test_vmec.py::test_plot_vmec_comparison": 5.689872365001065, + "tests/test_vmec.py::test_vmec_boundary_subspace": 2.52556586600258, + "tests/test_vmec.py::TestVMECIO::test_fourier_to_zernike": 2.604328910994809, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identities_inverse": 2.2736523150015273, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd": 2.192678698003874, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_fwd_sin_series": 2.193070601002546, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev": 2.144266970004537, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_identity_rev_sin_sym": 2.1817471730028046, + "tests/test_vmec.py::TestVMECIO::test_ptolemy_linear_transform": 2.8358562020002864, + "tests/test_vmec.py::TestVMECIO::test_zernike_to_fourier": 2.1473985249976977, + "tests/test_vmec.py::test_vmec_load_profiles": 17.117205705013475, + "tests/test_vmec.py::test_vmec_save_1": 23.67683808500442, + "tests/test_vmec.py::test_vmec_save_2": 2.403565439002705, + "tests/test_vmec.py::test_vmec_save_asym": 24.341486097000598 } diff --git a/desc/backend.py b/desc/backend.py index d5cffad24b..7f5dff1431 100644 --- a/desc/backend.py +++ b/desc/backend.py @@ -201,7 +201,7 @@ def fori_loop(lower, upper, body_fun, init_val): val = body_fun(i, val) return val - def cond(pred, true_fun, false_fun, operand): + def cond(pred, true_fun, false_fun, *operand): """Conditionally apply true_fun or false_fun. This version is for the numpy backend, for jax backend see jax.lax.cond @@ -227,9 +227,9 @@ def cond(pred, true_fun, false_fun, operand): """ if pred: - return true_fun(operand) + return true_fun(*operand) else: - return false_fun(operand) + return false_fun(*operand) def switch(index, branches, operand): """Apply exactly one of branches given by index. diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index f2e214dcfa..610dccbf80 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -5,11 +5,11 @@ Parameters ---------- params : dict of ndarray - Parameters from the equilibrium, such as R_lmn, Z_lmn, i_l, p_l, etc + Parameters from the equilibrium, such as R_lmn, Z_lmn, i_l, p_l, etc. transforms : dict of Transform - Transforms for R, Z, lambda, etc + Transforms for R, Z, lambda, etc. profiles : dict of Profile - Profile objects for pressure, iota, current, etc + Profile objects for pressure, iota, current, etc. data : dict of ndarray Data computed so far, generally output from other compute functions kwargs : dict @@ -59,8 +59,8 @@ # import the compute module. def _build_data_index(): - for p in data_index.keys(): - for key in data_index[p].keys(): + for p in data_index: + for key in data_index[p]: full = { "data": get_data_deps(key, p, has_axis=False), "transforms": get_derivs(key, p, has_axis=False), diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 39b8cd2ce3..07b5fea055 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -26,10 +26,10 @@ transforms={}, profiles=[], coordinates="rtz", - data=["B"], + data=["B", "|B|"], ) def _b(params, transforms, profiles, data, **kwargs): - data["b"] = (data["B"].T / jnp.linalg.norm(data["B"], axis=-1)).T + data["b"] = (data["B"].T / data["|B|"]).T return data @@ -47,6 +47,8 @@ def _b(params, transforms, profiles, data, **kwargs): data=["e_theta/sqrt(g)", "e_zeta"], ) def _e_sup_rho(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞^ρ | ρ=0 }. data["e^rho"] = cross(data["e_theta/sqrt(g)"], data["e_zeta"]) return data @@ -196,8 +198,14 @@ def _e_sup_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_sup_theta_times_sqrt_g(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞^θ √g | ρ=0 }. data["e^theta*sqrt(g)"] = cross(data["e_zeta"], data["e_rho"]) return data @@ -299,6 +307,8 @@ def _e_sup_theta_z(params, transforms, profiles, data, **kwargs): data=["e_rho", "e_theta/sqrt(g)"], ) def _e_sup_zeta(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞^ζ | ρ=0 }. data["e^zeta"] = cross(data["e_rho"], data["e_theta/sqrt(g)"]) return data @@ -453,6 +463,8 @@ def _e_sub_phi(params, transforms, profiles, data, **kwargs): data=["R", "R_r", "Z_r", "omega_r"], ) def _e_sub_rho(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞ᵨ | ρ=0 }. data["e_rho"] = jnp.array([data["R_r"], data["R"] * data["omega_r"], data["Z_r"]]).T return data @@ -1386,6 +1398,8 @@ def _e_sub_theta(params, transforms, profiles, data, **kwargs): axis_limit_data=["e_theta_r", "sqrt(g)_r"], ) def _e_sub_theta_over_sqrt_g(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞_θ / √g | ρ=0 }. data["e_theta/sqrt(g)"] = transforms["grid"].replace_at_axis( (data["e_theta"].T / data["sqrt(g)"]).T, lambda: (data["e_theta_r"].T / data["sqrt(g)_r"]).T, @@ -1426,6 +1440,8 @@ def _e_sub_theta_pest(params, transforms, profiles, data, **kwargs): data=["R", "R_r", "R_rt", "R_t", "Z_rt", "omega_r", "omega_rt", "omega_t"], ) def _e_sub_theta_r(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { ∂ᵨ 𝐞_θ | ρ=0 } data["e_theta_r"] = jnp.array( [ -data["R"] * data["omega_t"] * data["omega_r"] + data["R_rt"], @@ -3428,16 +3444,22 @@ def _gradpsi(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta", "|e_theta x e_zeta|"], + axis_limit_data=["e_theta_r", "|e_theta x e_zeta|_r"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.core.Surface", ], ) def _n_rho(params, transforms, profiles, data, **kwargs): - # equal to e^rho / |e^rho| but works correctly for surfaces as well that don't have - # contravariant basis defined - data["n_rho"] = ( - cross(data["e_theta"], data["e_zeta"]) / data["|e_theta x e_zeta|"][:, None] + # Equal to 𝐞^ρ / ‖𝐞^ρ‖ but works correctly for surfaces as well that don't + # have contravariant basis defined. + data["n_rho"] = transforms["grid"].replace_at_axis( + (cross(data["e_theta"], data["e_zeta"]).T / data["|e_theta x e_zeta|"]).T, + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞^ρ / ‖𝐞^ρ‖ | ρ=0 }. + lambda: ( + cross(data["e_theta_r"], data["e_zeta"]).T / data["|e_theta x e_zeta|_r"] + ).T, ) return data @@ -3460,9 +3482,11 @@ def _n_rho(params, transforms, profiles, data, **kwargs): ], ) def _n_theta(params, transforms, profiles, data, **kwargs): + # Equal to 𝐞^θ / ‖𝐞^θ‖ but works correctly for surfaces as well that don't + # have contravariant basis defined. data["n_theta"] = ( - cross(data["e_zeta"], data["e_rho"]) / data["|e_zeta x e_rho|"][:, None] - ) + cross(data["e_zeta"], data["e_rho"]).T / data["|e_zeta x e_rho|"] + ).T return data @@ -3478,13 +3502,21 @@ def _n_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta", "|e_rho x e_theta|"], + axis_limit_data=["e_theta_r", "|e_rho x e_theta|_r"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.core.Surface", ], ) def _n_zeta(params, transforms, profiles, data, **kwargs): - data["n_zeta"] = ( - cross(data["e_rho"], data["e_theta"]) / data["|e_rho x e_theta|"][:, None] + # Equal to 𝐞^ζ / ‖𝐞^ζ‖ but works correctly for surfaces as well that don't + # have contravariant basis defined. + data["n_zeta"] = transforms["grid"].replace_at_axis( + (cross(data["e_rho"], data["e_theta"]).T / data["|e_rho x e_theta|"]).T, + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { 𝐞^ζ / ‖𝐞^ζ‖ | ρ=0 }. + lambda: ( + cross(data["e_rho"], data["e_theta_r"]).T / data["|e_rho x e_theta|_r"] + ).T, ) return data diff --git a/desc/compute/_bootstrap.py b/desc/compute/_bootstrap.py index fffeecc46e..a755d6b0da 100644 --- a/desc/compute/_bootstrap.py +++ b/desc/compute/_bootstrap.py @@ -1,16 +1,25 @@ -"""Compute functions for bootstrap current.""" +"""Compute functions for bootstrap current. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" from scipy.constants import elementary_charge from scipy.special import roots_legendre from ..backend import fori_loop, jnp from .data_index import register_compute_fun -from .utils import compress, expand, surface_averages_map +from .utils import surface_averages_map @register_compute_fun( name="trapped fraction", - label="1 - \\frac{3}{4} \\langle B^2 \\rangle \\int_0^{1/Bmax} " + label="1 - \\frac{3}{4} \\langle |B|^2 \\rangle \\int_0^{1/Bmax} " "\\frac{\\lambda\\; d\\lambda}{\\langle \\sqrt{1 - \\lambda B} \\rangle}", units="~", units_long="None", @@ -20,22 +29,18 @@ transforms={"grid": []}, profiles=[], coordinates="r", - data=["sqrt(g)", "V_r(r)", "|B|", "", "max_tz |B|"], + data=["sqrt(g)", "V_r(r)", "|B|", "<|B|^2>", "max_tz |B|"], + axis_limit_data=["sqrt(g)_r", "V_rr(r)"], n_gauss="n_gauss", ) def _trapped_fraction(params, transforms, profiles, data, **kwargs): - r""" - Evaluate the effective trapped particle fraction. + """Evaluate the effective trapped particle fraction. Compute the effective fraction of trapped particles, which enters - several formulae for neoclassical transport. The trapped fraction - ``f_t`` has a standard definition in neoclassical theory: - - .. math:: - f_t = 1 - \frac{3}{4} \langle B^2 \rangle \int_0^{1/Bmax} - \frac{\lambda\; d\lambda}{\langle \sqrt{1 - \lambda B} \rangle} - - where :math:`\langle \ldots \rangle` is a flux surface average. + several formulae for neoclassical transport. + The trapped fraction fₜ has a standard definition in neoclassical theory: + fₜ = 1 − 3/4 〈|B|²〉 ∫₀¹/ᴮᵐᵃˣ λ / 〈√(1 − λ B)〉 dλ + where 〈 ⋯ 〉 is a flux surface average. """ # Get nodes and weights for Gauss-Legendre integration: n_gauss = kwargs.get("n_gauss", 20) @@ -45,18 +50,20 @@ def _trapped_fraction(params, transforms, profiles, data, **kwargs): lambda_weights = jnp.asarray(base_weights * 0.5) grid = transforms["grid"] - Bmax = data["max_tz |B|"] - modB_over_Bmax = data["|B|"] / Bmax - sqrt_g = data["sqrt(g)"] - Bmax_squared = compress(grid, Bmax * Bmax) - V_r = compress(grid, data["V_r(r)"]) + modB_over_Bmax = data["|B|"] / data["max_tz |B|"] + Bmax_squared = grid.compress(data["max_tz |B|"]) ** 2 + # to resolve indeterminate form of limit at magnetic axis + sqrt_g = grid.replace_at_axis(data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True) + V_r = grid.compress( + grid.replace_at_axis(data["V_r(r)"], lambda: data["V_rr(r)"], copy=True) + ) compute_surface_averages = surface_averages_map(grid, expand_out=False) # Sum over the lambda grid points, using fori_loop for efficiency. def body_fun(jlambda, lambda_integral): flux_surf_avg_term = compute_surface_averages( jnp.sqrt(1 - lambd[jlambda] * modB_over_Bmax), - sqrt_g, + sqrt_g=sqrt_g, denominator=V_r, ) return lambda_integral + lambda_weights[jlambda] * lambd[jlambda] / ( @@ -64,21 +71,15 @@ def body_fun(jlambda, lambda_integral): ) lambda_integral = fori_loop(0, n_gauss, body_fun, jnp.zeros(grid.num_rho)) - - trapped_fraction = 1 - 0.75 * compress(grid, data[""]) * lambda_integral - data["trapped fraction"] = expand(grid, trapped_fraction) + data["trapped fraction"] = 1 - 0.75 * data["<|B|^2>"] * grid.expand(lambda_integral) return data -def j_dot_B_Redl( - geom_data, - profile_data, - helicity_N=None, -): - r"""Compute the bootstrap current. +def j_dot_B_Redl(geom_data, profile_data, helicity_N=None): + """Compute the bootstrap current 〈𝐉 ⋅ 𝐁〉. - (specifically :math:`\langle\vec{J}\cdot\vec{B}\rangle`) using the formulae in - Redl et al, Physics of Plasmas 28, 022502 (2021). This formula for + Compute 〈𝐉 ⋅ 𝐁〉 using the formulae in + Redl et al., Physics of Plasmas 28, 022502 (2021). This formula for the bootstrap current is valid in axisymmetry, quasi-axisymmetry, and quasi-helical symmetry, but not in other stellarators. @@ -91,7 +92,7 @@ def j_dot_B_Redl( - iota: 1D array with the rotational transform. - epsilon: 1D array with the effective inverse aspect ratio to use in the Redl formula. - - psi_edge: float, the boundary toroidal flux, divided by (2 pi). + - psi_edge: float, the boundary toroidal flux, divided by 2π. - f_t: 1D array with the effective trapped particle fraction The argument ``profile_data`` is a Dictionary that should contain the @@ -153,7 +154,7 @@ def j_dot_B_Redl( geometry_factor = abs(R / (iota - helicity_N)) nu_e = ( geometry_factor - * (6.921e-18) + * 6.921e-18 * ne * Zeff * ln_Lambda_e @@ -161,7 +162,7 @@ def j_dot_B_Redl( ) nu_i = ( geometry_factor - * (4.90e-18) + * 4.90e-18 * ni * (Zeff**4) * ln_Lambda_ii @@ -339,53 +340,47 @@ def j_dot_B_Redl( helicity="helicity", ) def _compute_J_dot_B_Redl(params, transforms, profiles, data, **kwargs): - r"""Compute the bootstrap current. + """Compute the bootstrap current 〈𝐉 ⋅ 𝐁〉. - (specifically :math:`\langle\vec{J}\cdot\vec{B}\rangle`) using the formulae in - Redl et al, Physics of Plasmas 28, 022502 (2021). This formula for + Compute 〈𝐉 ⋅ 𝐁〉 using the formulae in + Redl et al., Physics of Plasmas 28, 022502 (2021). This formula for the bootstrap current is valid in axisymmetry, quasi-axisymmetry, and quasi-helical symmetry, but not in other stellarators. """ grid = transforms["grid"] - # Note that the geom_data dictionary provided to j_dot_B_Redl() # contains info only as a function of rho, not theta or zeta, # i.e. on the compressed grid. In contrast, "data" contains # quantities on a 3D grid even for quantities that are flux # functions. - geom_data = {} - geom_data["f_t"] = compress(grid, data["trapped fraction"]) - geom_data["epsilon"] = compress(grid, data["effective r/R0"]) - geom_data["G"] = compress(grid, data["G"]) - geom_data["I"] = compress(grid, data["I"]) - geom_data["iota"] = compress(grid, data["iota"]) - geom_data["<1/|B|>"] = compress(grid, data["<1/|B|>"]) + geom_data = { + "f_t": grid.compress(data["trapped fraction"]), + "epsilon": grid.compress(data["effective r/R0"]), + "G": grid.compress(data["G"]), + "I": grid.compress(data["I"]), + "iota": grid.compress(data["iota"]), + "<1/|B|>": grid.compress(data["<1/|B|>"]), + "psi_edge": params["Psi"] / (2 * jnp.pi), + } geom_data["R"] = (geom_data["G"] + geom_data["iota"] * geom_data["I"]) * geom_data[ "<1/|B|>" ] - geom_data["psi_edge"] = params["Psi"] / (2 * jnp.pi) - - profile_data = {} - profile_data["rho"] = compress(grid, data["rho"]) - profile_data["ne"] = compress(grid, data["ne"]) - profile_data["ne_r"] = compress(grid, data["ne_r"]) - profile_data["Te"] = compress(grid, data["Te"]) - profile_data["Te_r"] = compress(grid, data["Te_r"]) - profile_data["Ti"] = compress(grid, data["Ti"]) - profile_data["Ti_r"] = compress(grid, data["Ti_r"]) + profile_data = { + "rho": grid.compress(data["rho"]), + "ne": grid.compress(data["ne"]), + "ne_r": grid.compress(data["ne_r"]), + "Te": grid.compress(data["Te"]), + "Te_r": grid.compress(data["Te_r"]), + "Ti": grid.compress(data["Ti"]), + "Ti_r": grid.compress(data["Ti_r"]), + } if profiles["atomic_number"] is None: - Zeff = jnp.ones(grid.num_rho) + profile_data["Zeff"] = jnp.ones(grid.num_rho) else: - Zeff = compress(grid, data["Zeff"]) - profile_data["Zeff"] = Zeff + profile_data["Zeff"] = grid.compress(data["Zeff"]) helicity = kwargs.get("helicity", (1, 0)) helicity_N = helicity[1] - - j_dot_B_data = j_dot_B_Redl( - geom_data, - profile_data, - helicity_N, - ) - data[" Redl"] = expand(grid, j_dot_B_data[""]) + j_dot_B_data = j_dot_B_Redl(geom_data, profile_data, helicity_N) + data[" Redl"] = grid.expand(j_dot_B_data[""]) return data diff --git a/desc/compute/_equil.py b/desc/compute/_equil.py index 24bb3a94ea..3a6daddb2d 100644 --- a/desc/compute/_equil.py +++ b/desc/compute/_equil.py @@ -1,4 +1,13 @@ -"""Compute functions for equilibrium objectives, ie Force and MHD energy.""" +"""Compute functions for equilibrium objectives, i.e. Force and MHD energy. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" from scipy.constants import mu_0 @@ -20,9 +29,39 @@ profiles=[], coordinates="rtz", data=["sqrt(g)", "B_zeta_t", "B_theta_z"], + axis_limit_data=["sqrt(g)_r", "B_zeta_rt", "B_theta_rz"], + parameterization="desc.equilibrium.equilibrium.Equilibrium", ) def _J_sup_rho(params, transforms, profiles, data, **kwargs): - data["J^rho"] = (data["B_zeta_t"] - data["B_theta_z"]) / (mu_0 * data["sqrt(g)"]) + # At the magnetic axis, + # ∂_θ (𝐁 ⋅ 𝐞_ζ) - ∂_ζ (𝐁 ⋅ 𝐞_θ) = 𝐁 ⋅ (∂_θ 𝐞_ζ - ∂_ζ 𝐞_θ) = 0 + # because the partial derivatives commute. So 𝐉^ρ is of the indeterminate + # form 0/0 and we may compute the limit as follows. + data["J^rho"] = ( + transforms["grid"].replace_at_axis( + (data["B_zeta_t"] - data["B_theta_z"]) / data["sqrt(g)"], + lambda: (data["B_zeta_rt"] - data["B_theta_rz"]) / data["sqrt(g)_r"], + ) + ) / mu_0 + return data + + +@register_compute_fun( + name="J^theta*sqrt(g)", + label="J^{\\theta} \\sqrt{g}", + units="A", + units_long="Amperes", + description="Contravariant poloidal component of plasma current density," + " weighted by 3-D volume Jacobian", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B_rho_z", "B_zeta_r"], +) +def _J_sup_theta_sqrt_g(params, transforms, profiles, data, **kwargs): + data["J^theta*sqrt(g)"] = (data["B_rho_z"] - data["B_zeta_r"]) / mu_0 return data @@ -37,10 +76,10 @@ def _J_sup_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "B_rho_z", "B_zeta_r"], + data=["sqrt(g)", "J^theta*sqrt(g)"], ) def _J_sup_theta(params, transforms, profiles, data, **kwargs): - data["J^theta"] = (data["B_rho_z"] - data["B_zeta_r"]) / (mu_0 * data["sqrt(g)"]) + data["J^theta"] = data["J^theta*sqrt(g)"] / data["sqrt(g)"] return data @@ -56,9 +95,19 @@ def _J_sup_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["sqrt(g)", "B_theta_r", "B_rho_t"], + axis_limit_data=["sqrt(g)_r", "B_theta_rr", "B_rho_rt"], ) def _J_sup_zeta(params, transforms, profiles, data, **kwargs): - data["J^zeta"] = (data["B_theta_r"] - data["B_rho_t"]) / (mu_0 * data["sqrt(g)"]) + # At the magnetic axis, + # ∂ᵨ (𝐁 ⋅ 𝐞_θ) - ∂_θ (𝐁 ⋅ 𝐞ᵨ) = 𝐁 ⋅ (∂ᵨ 𝐞_θ - ∂_θ 𝐞ᵨ) = 0 + # because the partial derivatives commute. So 𝐉^ζ is of the indeterminate + # form 0/0 and we may compute the limit as follows. + data["J^zeta"] = ( + transforms["grid"].replace_at_axis( + (data["B_theta_r"] - data["B_rho_t"]) / data["sqrt(g)"], + lambda: (data["B_theta_rr"] - data["B_rho_rt"]) / data["sqrt(g)_r"], + ) + ) / mu_0 return data @@ -73,23 +122,107 @@ def _J_sup_zeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["J^rho", "J^theta", "J^zeta", "e_rho", "e_theta", "e_zeta"], + data=[ + "J^rho", + "J^zeta", + "J^theta*sqrt(g)", + "e_rho", + "e_zeta", + "e_theta/sqrt(g)", + ], ) def _J(params, transforms, profiles, data, **kwargs): data["J"] = ( data["J^rho"] * data["e_rho"].T - + data["J^theta"] * data["e_theta"].T + + data["J^theta*sqrt(g)"] * data["e_theta/sqrt(g)"].T + data["J^zeta"] * data["e_zeta"].T ).T return data +@register_compute_fun( + name="J*sqrt(g)", + label="\\mathbf{J} \\sqrt{g}", + units="A m", + units_long="Ampere meters", + description="Plasma current density weighted by 3-D volume Jacobian", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "B_rho_z", + "B_theta_r", + "B_zeta_t", + "B_rho_t", + "B_theta_z", + "B_zeta_r", + "e_rho", + "e_theta", + "e_zeta", + ], +) +def _J_sqrt_g(params, transforms, profiles, data, **kwargs): + data["J*sqrt(g)"] = ( + (data["B_zeta_t"] - data["B_theta_z"]) * data["e_rho"].T + + (data["B_rho_z"] - data["B_zeta_r"]) * data["e_theta"].T + + (data["B_theta_r"] - data["B_rho_t"]) * data["e_zeta"].T + ).T / mu_0 + return data + + +@register_compute_fun( + name="(J*sqrt(g))_r", + label="\\partial_{\\rho} (\\mathbf{J} \\sqrt{g})", + units="A m", + units_long="Ampere meters", + description="Plasma current density weighted by 3-D volume Jacobian," + " radial derivative", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "B_rho_z", + "B_rho_rz", + "B_theta_r", + "B_theta_rr", + "B_zeta_t", + "B_zeta_rt", + "B_rho_t", + "B_rho_rt", + "B_theta_z", + "B_theta_rz", + "B_zeta_r", + "B_zeta_rr", + "e_rho", + "e_theta", + "e_zeta", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + ], +) +def _J_sqrt_g_r(params, transforms, profiles, data, **kwargs): + data["(J*sqrt(g))_r"] = ( + (data["B_zeta_rt"] - data["B_theta_rz"]) * data["e_rho"].T + + (data["B_zeta_t"] - data["B_theta_z"]) * data["e_rho_r"].T + + (data["B_rho_rz"] - data["B_zeta_rr"]) * data["e_theta"].T + + (data["B_rho_z"] - data["B_zeta_r"]) * data["e_theta_r"].T + + (data["B_theta_rr"] - data["B_rho_rt"]) * data["e_zeta"].T + + (data["B_theta_r"] - data["B_rho_t"]) * data["e_zeta_r"].T + ).T / mu_0 + return data + + @register_compute_fun( name="J_R", label="J_{R}", units="A \\cdot m^{-2}", units_long="Amperes / square meter", - description="Radial componenet of plasma current density in lab frame", + description="Radial component of plasma current density in lab frame", dim=1, params=[], transforms={}, @@ -107,7 +240,7 @@ def _J_R(params, transforms, profiles, data, **kwargs): label="J_{\\phi}", units="A \\cdot m^{-2}", units_long="Amperes / square meter", - description="Toroidal componenet of plasma current density in lab frame", + description="Toroidal component of plasma current density in lab frame", dim=1, params=[], transforms={}, @@ -125,7 +258,7 @@ def _J_phi(params, transforms, profiles, data, **kwargs): label="J_{Z}", units="A \\cdot m^{-2}", units_long="Amperes / square meter", - description="Vertical componenet of plasma current density in lab frame", + description="Vertical component of plasma current density in lab frame", dim=1, params=[], transforms={}, @@ -143,7 +276,7 @@ def _J_Z(params, transforms, profiles, data, **kwargs): label="|\\mathbf{J}|", units="A \\cdot m^{-2}", units_long="Amperes / square meter", - description="Magnitue of plasma current density", + description="Magnitude of plasma current density", dim=1, params=[], transforms={}, @@ -234,19 +367,26 @@ def _J_dot_B(params, transforms, profiles, data, **kwargs): label="\\langle \\mathbf{J} \\cdot \\mathbf{B} \\rangle", units="N / m^{3}", units_long="Newtons / cubic meter", - description="Flux surface average of current density dotted into magnetic field", + description="Flux surface average of current density dotted into magnetic field " + + "(note units are not Amperes)", dim=1, params=[], transforms={}, profiles=[], coordinates="r", - data=["J*B", "sqrt(g)"], + data=["J*sqrt(g)", "B", "V_r(r)"], + axis_limit_data=["(J*sqrt(g))_r", "V_rr(r)"], ) def _J_dot_B_fsa(params, transforms, profiles, data, **kwargs): + J = transforms["grid"].replace_at_axis( + data["J*sqrt(g)"], lambda: data["(J*sqrt(g))_r"], copy=True + ) data[""] = surface_averages( transforms["grid"], - data["J*B"], - sqrt_g=data["sqrt(g)"], + dot(J, data["B"]), # sqrt(g) factor pushed into J + denominator=transforms["grid"].replace_at_axis( + data["V_r(r)"], lambda: data["V_rr(r)"], copy=True + ), ) return data @@ -280,12 +420,10 @@ def _J_parallel(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["p_r", "sqrt(g)", "B^theta", "B^zeta", "J^theta", "J^zeta"], + data=["p_r", "(curl(B)xB)_rho"], ) def _F_rho(params, transforms, profiles, data, **kwargs): - data["F_rho"] = -data["p_r"] + data["sqrt(g)"] * ( - data["B^zeta"] * data["J^theta"] - data["B^theta"] * data["J^zeta"] - ) + data["F_rho"] = data["(curl(B)xB)_rho"] / mu_0 - data["p_r"] return data @@ -300,10 +438,10 @@ def _F_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "B^zeta", "J^rho"], + data=["F_helical", "B^zeta"], ) def _F_theta(params, transforms, profiles, data, **kwargs): - data["F_theta"] = -data["sqrt(g)"] * data["B^zeta"] * data["J^rho"] + data["F_theta"] = -data["B^zeta"] * data["F_helical"] return data @@ -318,28 +456,28 @@ def _F_theta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "B^theta", "J^rho"], + data=["B^theta", "F_helical"], ) def _F_zeta(params, transforms, profiles, data, **kwargs): - data["F_zeta"] = data["sqrt(g)"] * data["B^theta"] * data["J^rho"] + data["F_zeta"] = data["B^theta"] * data["F_helical"] return data @register_compute_fun( name="F_helical", label="F_{helical}", - units="N \\cdot m^{-2}", - units_long="Newtons / square meter", + units="A", + units_long="Amperes", description="Covariant helical component of force balance error", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "J^rho"], + data=["B_zeta_t", "B_theta_z"], ) def _F_helical(params, transforms, profiles, data, **kwargs): - data["F_helical"] = data["sqrt(g)"] * data["J^rho"] + data["F_helical"] = (data["B_zeta_t"] - data["B_theta_z"]) / mu_0 return data @@ -354,12 +492,13 @@ def _F_helical(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["F_rho", "F_theta", "F_zeta", "e^rho", "e^theta", "e^zeta"], + data=["F_rho", "F_zeta", "e^rho", "e^zeta", "B^zeta", "J^rho", "e^theta*sqrt(g)"], ) def _F(params, transforms, profiles, data, **kwargs): + # F_theta e^theta refactored as below to resolve indeterminacy at axis. data["F"] = ( data["F_rho"] * data["e^rho"].T - + data["F_theta"] * data["e^theta"].T + - data["B^zeta"] * data["J^rho"] * data["e^theta*sqrt(g)"].T + data["F_zeta"] * data["e^zeta"].T ).T return data @@ -416,7 +555,7 @@ def _Fmag_vol(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["B^theta", "B^zeta", "e^theta", "e^zeta"], ) -def _e_helical(params, transforms, profiles, data, **kwargs): +def _e_sup_helical(params, transforms, profiles, data, **kwargs): data["e^helical"] = ( data["B^zeta"] * data["e^theta"].T - data["B^theta"] * data["e^zeta"].T ).T @@ -436,7 +575,7 @@ def _e_helical(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["e^helical"], ) -def _helical_mag(params, transforms, profiles, data, **kwargs): +def _e_sup_helical_mag(params, transforms, profiles, data, **kwargs): data["|e^helical|"] = jnp.linalg.norm(data["e^helical"], axis=-1) return data diff --git a/desc/compute/_field.py b/desc/compute/_field.py index 760b00a26f..5bd2e74d27 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -1,15 +1,24 @@ -"""Compute functions for magnetic field quantities.""" +"""Compute functions for magnetic field quantities. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" from scipy.constants import mu_0 -from desc.backend import jnp, put +from desc.backend import jnp from .data_index import register_compute_fun from .utils import ( cross, dot, surface_averages, - surface_integrals, + surface_integrals_map, surface_max, surface_min, ) @@ -30,12 +39,10 @@ axis_limit_data=["psi_rr", "sqrt(g)_r"], ) def _B0(params, transforms, profiles, data, **kwargs): - data["B0"] = data["psi_r"] / data["sqrt(g)"] - if transforms["grid"].axis.size: - limit = data["psi_rr"] / data["sqrt(g)_r"] - data["B0"] = put( - data["B0"], transforms["grid"].axis, limit[transforms["grid"].axis] - ) + data["B0"] = transforms["grid"].replace_at_axis( + data["psi_r"] / data["sqrt(g)"], + lambda: data["psi_rr"] / data["sqrt(g)_r"], + ) return data @@ -68,10 +75,12 @@ def _B_sup_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "iota", "lambda_z"], + data=["B0", "iota", "lambda_z", "omega_z"], ) def _B_sup_theta(params, transforms, profiles, data, **kwargs): - data["B^theta"] = data["B0"] * (data["iota"] - data["lambda_z"]) + data["B^theta"] = data["B0"] * ( + data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] + ) return data @@ -86,10 +95,12 @@ def _B_sup_theta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "lambda_t"], + data=["B0", "iota", "lambda_t", "omega_t"], ) def _B_sup_zeta(params, transforms, profiles, data, **kwargs): - data["B^zeta"] = data["B0"] * (1 + data["lambda_t"]) + data["B^zeta"] = data["B0"] * ( + -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + ) return data @@ -179,11 +190,16 @@ def _B_Z(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "psi_rr", "sqrt(g)", "sqrt(g)_r"], + axis_limit_data=["psi_rrr", "sqrt(g)_rr"], ) def _B0_r(params, transforms, profiles, data, **kwargs): - data["B0_r"] = ( - data["psi_rr"] / data["sqrt(g)"] - - data["psi_r"] * data["sqrt(g)_r"] / data["sqrt(g)"] ** 2 + data["B0_r"] = transforms["grid"].replace_at_axis( + (data["psi_rr"] * data["sqrt(g)"] - data["psi_r"] * data["sqrt(g)_r"]) + / data["sqrt(g)"] ** 2, + lambda: ( + data["psi_rrr"] * data["sqrt(g)_r"] - data["psi_rr"] * data["sqrt(g)_rr"] + ) + / (2 * data["sqrt(g)_r"] ** 2), ) return data @@ -193,18 +209,34 @@ def _B0_r(params, transforms, profiles, data, **kwargs): label="\\partial_{\\rho} B^{\\theta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant poloidal component of magnetic field, derivative wrt " - + "radial coordinate", + description=( + "Contravariant poloidal component of magnetic field, derivative wrt radial" + " coordinate" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_r", "iota", "iota_r", "lambda_z", "lambda_rz"], + data=[ + "B0", + "B0_r", + "iota", + "iota_r", + "lambda_rz", + "lambda_z", + "omega_rz", + "omega_z", + ], ) def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): - data["B^theta_r"] = data["B0_r"] * (data["iota"] - data["lambda_z"]) + ( - data["B0"] * (data["iota_r"] - data["lambda_rz"]) + data["B^theta_r"] = data["B0"] * ( + data["iota"] * data["omega_rz"] + + data["iota_r"] * data["omega_z"] + + data["iota_r"] + - data["lambda_rz"] + ) + data["B0_r"] * ( + data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -214,19 +246,32 @@ def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): label="\\partial_{\\rho} B^{\\zeta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant toroidal component of magnetic field, derivative wrt " - + "radial coordinate", + description=( + "Contravariant toroidal component of magnetic field, derivative wrt radial" + " coordinate" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_r", "lambda_t", "lambda_rt"], + data=[ + "B0", + "B0_r", + "iota", + "iota_r", + "lambda_rt", + "lambda_t", + "omega_rt", + "omega_t", + ], ) def _B_sup_zeta_r(params, transforms, profiles, data, **kwargs): - data["B^zeta_r"] = ( - data["B0_r"] * (1 + data["lambda_t"]) + data["B0"] * data["lambda_rt"] - ) + data["B^zeta_r"] = data["B0"] * ( + -data["iota"] * data["omega_rt"] + - data["iota_r"] * data["omega_t"] + + data["lambda_rt"] + ) + data["B0_r"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) return data @@ -273,10 +318,14 @@ def _B_r(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["psi_r", "sqrt(g)_t", "sqrt(g)"], + data=["psi_r", "sqrt(g)", "sqrt(g)_t"], + axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rt"], ) def _B0_t(params, transforms, profiles, data, **kwargs): - data["B0_t"] = -data["psi_r"] * data["sqrt(g)_t"] / data["sqrt(g)"] ** 2 + data["B0_t"] = transforms["grid"].replace_at_axis( + -data["psi_r"] * data["sqrt(g)_t"] / data["sqrt(g)"] ** 2, + lambda: -data["psi_rr"] * data["sqrt(g)_rt"] / data["sqrt(g)_r"] ** 2, + ) return data @@ -285,19 +334,22 @@ def _B0_t(params, transforms, profiles, data, **kwargs): label="\\partial_{\\theta} B^{\\theta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant poloidal component of magnetic field, derivative wrt " - + "poloidal angle", + description=( + "Contravariant poloidal component of magnetic field, derivative wrt poloidal" + " coordinate" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_t", "iota", "lambda_z", "lambda_tz"], + data=["B0", "B0_t", "iota", "lambda_tz", "lambda_z", "omega_tz", "omega_z"], ) def _B_sup_theta_t(params, transforms, profiles, data, **kwargs): - data["B^theta_t"] = ( - data["B0_t"] * (data["iota"] - data["lambda_z"]) - - data["B0"] * data["lambda_tz"] + data["B^theta_t"] = data["B0"] * ( + data["iota"] * data["omega_tz"] - data["lambda_tz"] + ) + data["B0_t"] * ( + data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -307,19 +359,21 @@ def _B_sup_theta_t(params, transforms, profiles, data, **kwargs): label="\\partial_{\\theta} B^{\\zeta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant toroidal component of magnetic field, derivative wrt " - + "poloidal angle", + description=( + "Contravariant toroidal component of magnetic field, derivative wrt poloidal" + " coordinate" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_t", "lambda_t", "lambda_tt"], + data=["B0", "B0_t", "iota", "lambda_t", "lambda_tt", "omega_t", "omega_tt"], ) def _B_sup_zeta_t(params, transforms, profiles, data, **kwargs): - data["B^zeta_t"] = ( - data["B0_t"] * (1 + data["lambda_t"]) + data["B0"] * data["lambda_tt"] - ) + data["B^zeta_t"] = data["B0"] * ( + -data["iota"] * data["omega_tt"] + data["lambda_tt"] + ) + data["B0_t"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) return data @@ -367,9 +421,13 @@ def _B_t(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "sqrt(g)", "sqrt(g)_z"], + axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rz"], ) def _B0_z(params, transforms, profiles, data, **kwargs): - data["B0_z"] = -data["psi_r"] * data["sqrt(g)_z"] / data["sqrt(g)"] ** 2 + data["B0_z"] = transforms["grid"].replace_at_axis( + -data["psi_r"] * data["sqrt(g)_z"] / data["sqrt(g)"] ** 2, + lambda: -data["psi_rr"] * data["sqrt(g)_rz"] / data["sqrt(g)_r"] ** 2, + ) return data @@ -378,19 +436,22 @@ def _B0_z(params, transforms, profiles, data, **kwargs): label="\\partial_{\\zeta} B^{\\theta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant poloidal component of magnetic field, derivative wrt " - + "toroidal angle", + description=( + "Contravariant poloidal component of magnetic field, derivative wrt toroidal" + " coordinate" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_z", "iota", "lambda_z", "lambda_zz"], + data=["B0", "B0_z", "iota", "lambda_z", "lambda_zz", "omega_z", "omega_zz"], ) def _B_sup_theta_z(params, transforms, profiles, data, **kwargs): - data["B^theta_z"] = ( - data["B0_z"] * (data["iota"] - data["lambda_z"]) - - data["B0"] * data["lambda_zz"] + data["B^theta_z"] = data["B0"] * ( + data["iota"] * data["omega_zz"] - data["lambda_zz"] + ) + data["B0_z"] * ( + data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -400,19 +461,21 @@ def _B_sup_theta_z(params, transforms, profiles, data, **kwargs): label="\\partial_{\\zeta} B^{\\zeta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant toroidal component of magnetic field, derivative wrt " - + "toroidal angle", + description=( + "Contravariant toroidal component of magnetic field, derivative wrt toroidal" + " coordinate" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_z", "lambda_t", "lambda_tz"], + data=["B0", "B0_z", "iota", "lambda_t", "lambda_tz", "omega_t", "omega_tz"], ) def _B_sup_zeta_z(params, transforms, profiles, data, **kwargs): - data["B^zeta_z"] = ( - data["B0_z"] * (1 + data["lambda_t"]) + data["B0"] * data["lambda_tz"] - ) + data["B^zeta_z"] = data["B0"] * ( + -data["iota"] * data["omega_tz"] + data["lambda_tz"] + ) + data["B0_z"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) return data @@ -460,13 +523,23 @@ def _B_z(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "psi_rr", "psi_rrr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_rr"], + axis_limit_data=["sqrt(g)_rrr"], ) def _B0_rr(params, transforms, profiles, data, **kwargs): - data["B0_rr"] = ( - data["psi_rrr"] / data["sqrt(g)"] - - 2 * data["psi_rr"] * data["sqrt(g)_r"] / data["sqrt(g)"] ** 2 - - data["psi_r"] * data["sqrt(g)_rr"] / data["sqrt(g)"] ** 2 - + 2 * data["psi_r"] * data["sqrt(g)_r"] ** 2 / data["sqrt(g)"] ** 3 + data["B0_rr"] = transforms["grid"].replace_at_axis( + ( + data["psi_rrr"] * data["sqrt(g)"] ** 2 + - 2 * data["psi_rr"] * data["sqrt(g)_r"] * data["sqrt(g)"] + - data["psi_r"] * data["sqrt(g)_rr"] * data["sqrt(g)"] + + 2 * data["psi_r"] * data["sqrt(g)_r"] ** 2 + ) + / data["sqrt(g)"] ** 3, + lambda: ( + 3 * data["sqrt(g)_rr"] ** 2 * data["psi_rr"] + - 2 * data["sqrt(g)_rrr"] * data["sqrt(g)_r"] * data["psi_rr"] + - 3 * data["psi_rrr"] * data["sqrt(g)_r"] * data["sqrt(g)_rr"] + ) + / (6 * data["sqrt(g)_r"] ** 3), ) return data @@ -475,9 +548,11 @@ def _B0_rr(params, transforms, profiles, data, **kwargs): name="B^theta_rr", label="\\partial_{\\rho\\rho} B^{\\theta}", units="T \\cdot m^{-1}", - units_long="Tesla / meters", - description="Contravariant poloidal component of magnetic field, second derivative " - + "wrt radial coordinate", + units_long="Tesla / meter", + description=( + "Contravariant poloidal component of magnetic field, second derivative wrt" + " radial and radial coordinates" + ), dim=1, params=[], transforms={}, @@ -490,16 +565,34 @@ def _B0_rr(params, transforms, profiles, data, **kwargs): "iota", "iota_r", "iota_rr", - "lambda_z", - "lambda_rz", "lambda_rrz", + "lambda_rz", + "lambda_z", + "omega_rrz", + "omega_rz", + "omega_z", ], ) def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): data["B^theta_rr"] = ( - data["B0_rr"] * (data["iota"] - data["lambda_z"]) - + 2 * data["B0_r"] * (data["iota_r"] - data["lambda_rz"]) - + data["B0"] * (data["iota_rr"] - data["lambda_rrz"]) + data["B0"] + * ( + data["iota"] * data["omega_rrz"] + + 2 * data["iota_r"] * data["omega_rz"] + + data["iota_rr"] * data["omega_z"] + + data["iota_rr"] + - data["lambda_rrz"] + ) + + 2 + * data["B0_r"] + * ( + data["iota"] * data["omega_rz"] + + data["iota_r"] * data["omega_z"] + + data["iota_r"] + - data["lambda_rz"] + ) + + data["B0_rr"] + * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -508,21 +601,48 @@ def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): name="B^zeta_rr", label="\\partial_{\\rho\\rho} B^{\\zeta}", units="T \\cdot m^{-1}", - units_long="Tesla / meters", - description="Contravariant toroidal component of magnetic field, second derivative " - + "wrt radial coordinate", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, second derivative wrt" + " radial and radial coordinates" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_r", "B0_rr", "lambda_t", "lambda_rt", "lambda_rrt"], + data=[ + "B0", + "B0_r", + "B0_rr", + "iota", + "iota_r", + "iota_rr", + "lambda_rrt", + "lambda_rt", + "lambda_t", + "omega_rrt", + "omega_rt", + "omega_t", + ], ) def _B_sup_zeta_rr(params, transforms, profiles, data, **kwargs): data["B^zeta_rr"] = ( - data["B0_rr"] * (1 + data["lambda_t"]) - + 2 * data["B0_r"] * data["lambda_rt"] - + data["B0"] * data["lambda_rrt"] + -data["B0"] + * ( + data["iota"] * data["omega_rrt"] + + 2 * data["iota_r"] * data["omega_rt"] + + data["iota_rr"] * data["omega_t"] + - data["lambda_rrt"] + ) + - 2 + * data["B0_r"] + * ( + data["iota"] * data["omega_rt"] + + data["iota_r"] * data["omega_t"] + - data["lambda_rt"] + ) + + data["B0_rr"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -577,12 +697,16 @@ def _B_rr(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "sqrt(g)", "sqrt(g)_t", "sqrt(g)_tt"], + axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rt", "sqrt(g)_rtt"], ) def _B0_tt(params, transforms, profiles, data, **kwargs): - data["B0_tt"] = -( + data["B0_tt"] = transforms["grid"].replace_at_axis( data["psi_r"] - / data["sqrt(g)"] ** 2 - * (data["sqrt(g)_tt"] - 2 * data["sqrt(g)_t"] ** 2 / data["sqrt(g)"]) + * (2 * data["sqrt(g)_t"] ** 2 - data["sqrt(g)"] * data["sqrt(g)_tt"]) + / data["sqrt(g)"] ** 3, + lambda: data["psi_rr"] + * (2 * data["sqrt(g)_rt"] ** 2 - data["sqrt(g)_r"] * data["sqrt(g)_rtt"]) + / data["sqrt(g)_r"] ** 3, ) return data @@ -592,20 +716,34 @@ def _B0_tt(params, transforms, profiles, data, **kwargs): label="\\partial_{\\theta\\theta} B^{\\theta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant poloidal component of magnetic field, second " - + "derivative wrt poloidal angle", + description=( + "Contravariant poloidal component of magnetic field, second derivative wrt" + " poloidal and poloidal coordinates" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_t", "B0_tt", "iota", "lambda_z", "lambda_tz", "lambda_ttz"], + data=[ + "B0", + "B0_t", + "B0_tt", + "iota", + "lambda_ttz", + "lambda_tz", + "lambda_z", + "omega_ttz", + "omega_tz", + "omega_z", + ], ) def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): data["B^theta_tt"] = ( - data["B0_tt"] * (data["iota"] - data["lambda_z"]) - - 2 * data["B0_t"] * data["lambda_tz"] - - data["B0"] * data["lambda_ttz"] + data["B0"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) + + 2 * data["B0_t"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + + data["B0_tt"] + * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -615,20 +753,33 @@ def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): label="\\partial_{\\theta\\theta} B^{\\zeta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant toroidal component of magnetic field, second " - + "derivative wrt poloidal angle", + description=( + "Contravariant toroidal component of magnetic field, second derivative wrt" + " poloidal and poloidal coordinates" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_t", "B0_tt", "lambda_t", "lambda_tt", "lambda_ttt"], + data=[ + "B0", + "B0_t", + "B0_tt", + "iota", + "lambda_t", + "lambda_tt", + "lambda_ttt", + "omega_t", + "omega_tt", + "omega_ttt", + ], ) def _B_sup_zeta_tt(params, transforms, profiles, data, **kwargs): data["B^zeta_tt"] = ( - data["B0_tt"] * (1 + data["lambda_t"]) - + 2 * data["B0_t"] * data["lambda_tt"] - + data["B0"] * data["lambda_ttt"] + -data["B0"] * (data["iota"] * data["omega_ttt"] - data["lambda_ttt"]) + - 2 * data["B0_t"] * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + + data["B0_tt"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -683,12 +834,16 @@ def _B_tt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "sqrt(g)", "sqrt(g)_z", "sqrt(g)_zz"], + axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rz", "sqrt(g)_rzz"], ) def _B0_zz(params, transforms, profiles, data, **kwargs): - data["B0_zz"] = -( + data["B0_zz"] = transforms["grid"].replace_at_axis( data["psi_r"] - / data["sqrt(g)"] ** 2 - * (data["sqrt(g)_zz"] - 2 * data["sqrt(g)_z"] ** 2 / data["sqrt(g)"]) + * (2 * data["sqrt(g)_z"] ** 2 - data["sqrt(g)"] * data["sqrt(g)_zz"]) + / data["sqrt(g)"] ** 3, + lambda: data["psi_rr"] + * (2 * data["sqrt(g)_rz"] ** 2 - data["sqrt(g)_r"] * data["sqrt(g)_rzz"]) + / data["sqrt(g)_r"] ** 3, ) return data @@ -698,20 +853,34 @@ def _B0_zz(params, transforms, profiles, data, **kwargs): label="\\partial_{\\zeta\\zeta} B^{\\theta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant poloidal component of magnetic field, second " - + "derivative wrt toroidal angle", + description=( + "Contravariant poloidal component of magnetic field, second derivative wrt" + " toroidal and toroidal coordinates" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_z", "B0_zz", "iota", "lambda_z", "lambda_zz", "lambda_zzz"], + data=[ + "B0", + "B0_z", + "B0_zz", + "iota", + "lambda_z", + "lambda_zz", + "lambda_zzz", + "omega_z", + "omega_zz", + "omega_zzz", + ], ) def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): data["B^theta_zz"] = ( - data["B0_zz"] * (data["iota"] - data["lambda_z"]) - - 2 * data["B0_z"] * data["lambda_zz"] - - data["B0"] * data["lambda_zzz"] + data["B0"] * (data["iota"] * data["omega_zzz"] - data["lambda_zzz"]) + + 2 * data["B0_z"] * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) + + data["B0_zz"] + * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -721,20 +890,33 @@ def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): label="\\partial_{\\zeta\\zeta} B^{\\zeta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant toroidal component of magnetic field, second " - + "derivative wrt toroidal angle", + description=( + "Contravariant toroidal component of magnetic field, second derivative wrt" + " toroidal and toroidal coordinates" + ), dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_z", "B0_zz", "lambda_t", "lambda_tz", "lambda_tzz"], + data=[ + "B0", + "B0_z", + "B0_zz", + "iota", + "lambda_t", + "lambda_tz", + "lambda_tzz", + "omega_t", + "omega_tz", + "omega_tzz", + ], ) def _B_sup_zeta_zz(params, transforms, profiles, data, **kwargs): data["B^zeta_zz"] = ( - data["B0_zz"] * (1 + data["lambda_t"]) - + 2 * data["B0_z"] * data["lambda_tz"] - + data["B0"] * data["lambda_tzz"] + -data["B0"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) + - 2 * data["B0_z"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + + data["B0_zz"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -789,16 +971,25 @@ def _B_zz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "psi_rr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_t", "sqrt(g)_rt"], + axis_limit_data=["psi_rrr", "sqrt(g)_rr", "sqrt(g)_rrt"], ) def _B0_rt(params, transforms, profiles, data, **kwargs): - data["B0_rt"] = ( - -data["psi_rr"] * data["sqrt(g)_t"] / data["sqrt(g)"] ** 2 - - data["psi_r"] * data["sqrt(g)_rt"] / data["sqrt(g)"] ** 2 - + 2 - * data["psi_r"] - * data["sqrt(g)_r"] - * data["sqrt(g)_t"] - / data["sqrt(g)"] ** 3 + data["B0_rt"] = transforms["grid"].replace_at_axis( + ( + -data["sqrt(g)"] + * (data["psi_rr"] * data["sqrt(g)_t"] + data["psi_r"] * data["sqrt(g)_rt"]) + + 2 * data["psi_r"] * data["sqrt(g)_r"] * data["sqrt(g)_t"] + ) + / data["sqrt(g)"] ** 3, + lambda: ( + -data["sqrt(g)_r"] + * ( + data["psi_rrr"] * data["sqrt(g)_rt"] + + data["psi_rr"] * data["sqrt(g)_rrt"] + ) + + 2 * data["psi_rr"] * data["sqrt(g)_rr"] * data["sqrt(g)_rt"] + ) + / (2 * data["sqrt(g)_r"] ** 3), ) return data @@ -807,9 +998,11 @@ def _B0_rt(params, transforms, profiles, data, **kwargs): name="B^theta_rt", label="\\partial_{\\rho\\theta} B^{\\theta}", units="T \\cdot m^{-1}", - units_long="Tesla / meters", - description="Contravariant poloidal component of magnetic field, second " - + "derivative wrt radial coordinate and poloidal angle", + units_long="Tesla / meter", + description=( + "Contravariant poloidal component of magnetic field, second derivative wrt" + " radial and poloidal coordinates" + ), dim=1, params=[], transforms={}, @@ -818,22 +1011,38 @@ def _B0_rt(params, transforms, profiles, data, **kwargs): data=[ "B0", "B0_r", - "B0_t", "B0_rt", + "B0_t", "iota", "iota_r", - "lambda_z", + "lambda_rtz", "lambda_rz", "lambda_tz", - "lambda_rtz", + "lambda_z", + "omega_rtz", + "omega_rz", + "omega_tz", + "omega_z", ], ) def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): data["B^theta_rt"] = ( - data["B0_rt"] * (data["iota"] - data["lambda_z"]) - - data["B0_r"] * data["lambda_tz"] - + data["B0_t"] * (data["iota_r"] - data["lambda_rz"]) - - data["B0"] * data["lambda_rtz"] + data["B0"] + * ( + data["iota"] * data["omega_rtz"] + + data["iota_r"] * data["omega_tz"] + - data["lambda_rtz"] + ) + + data["B0_r"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + + data["B0_t"] + * ( + data["iota"] * data["omega_rz"] + + data["iota_r"] * data["omega_z"] + + data["iota_r"] + - data["lambda_rz"] + ) + + data["B0_rt"] + * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -842,9 +1051,11 @@ def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): name="B^zeta_rt", label="\\partial_{\\rho\\theta} B^{\\zeta}", units="T \\cdot m^{-1}", - units_long="Tesla / meters", - description="Contravariant toroidal component of magnetic field, second " - + "derivative wrt radial coordinate and poloidal angle", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, second derivative wrt" + " radial and poloidal coordinates" + ), dim=1, params=[], transforms={}, @@ -853,20 +1064,36 @@ def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): data=[ "B0", "B0_r", - "B0_t", "B0_rt", - "lambda_t", - "lambda_tt", + "B0_t", + "iota", + "iota_r", "lambda_rt", "lambda_rtt", + "lambda_t", + "lambda_tt", + "omega_rt", + "omega_rtt", + "omega_t", + "omega_tt", ], ) def _B_sup_zeta_rt(params, transforms, profiles, data, **kwargs): data["B^zeta_rt"] = ( - data["B0_rt"] * (1 + data["lambda_t"]) - + data["B0_r"] * data["lambda_tt"] - + data["B0_t"] * data["lambda_rt"] - + data["B0"] * data["lambda_rtt"] + -data["B0"] + * ( + data["iota"] * data["omega_rtt"] + + data["iota_r"] * data["omega_tt"] + - data["lambda_rtt"] + ) + - data["B0_r"] * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + - data["B0_t"] + * ( + data["iota"] * data["omega_rt"] + + data["iota_r"] * data["omega_t"] + - data["lambda_rt"] + ) + + data["B0_rt"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -928,15 +1155,22 @@ def _B_rt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "sqrt(g)", "sqrt(g)_t", "sqrt(g)_z", "sqrt(g)_tz"], + axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rt", "sqrt(g)_rz", "sqrt(g)_rtz"], ) def _B0_tz(params, transforms, profiles, data, **kwargs): - data["B0_tz"] = -( + data["B0_tz"] = transforms["grid"].replace_at_axis( data["psi_r"] - / data["sqrt(g)"] ** 2 * ( - data["sqrt(g)_tz"] - - 2 * data["sqrt(g)_t"] * data["sqrt(g)_z"] / data["sqrt(g)"] + 2 * data["sqrt(g)_t"] * data["sqrt(g)_z"] + - data["sqrt(g)_tz"] * data["sqrt(g)"] ) + / data["sqrt(g)"] ** 3, + lambda: data["psi_rr"] + * ( + 2 * data["sqrt(g)_rt"] * data["sqrt(g)_rz"] + - data["sqrt(g)_rtz"] * data["sqrt(g)_r"] + ) + / data["sqrt(g)_r"] ** 3, ) return data @@ -946,8 +1180,10 @@ def _B0_tz(params, transforms, profiles, data, **kwargs): label="\\partial_{\\theta\\zeta} B^{\\theta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant poloidal component of magnetic field, second " - + "derivative wrt poloidal and toroidal angles", + description=( + "Contravariant poloidal component of magnetic field, second derivative wrt" + " poloidal and toroidal coordinates" + ), dim=1, params=[], transforms={}, @@ -956,21 +1192,26 @@ def _B0_tz(params, transforms, profiles, data, **kwargs): data=[ "B0", "B0_t", - "B0_z", "B0_tz", + "B0_z", "iota", - "lambda_z", - "lambda_zz", "lambda_tz", "lambda_tzz", + "lambda_z", + "lambda_zz", + "omega_tz", + "omega_tzz", + "omega_z", + "omega_zz", ], ) def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): data["B^theta_tz"] = ( - data["B0_tz"] * (data["iota"] - data["lambda_z"]) - - data["B0_t"] * data["lambda_zz"] - - data["B0_z"] * data["lambda_tz"] - - data["B0"] * data["lambda_tzz"] + data["B0"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) + + data["B0_t"] * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) + + data["B0_z"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + + data["B0_tz"] + * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -980,8 +1221,10 @@ def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): label="\\partial_{\\theta\\zeta} B^{\\zeta}", units="T \\cdot m^{-1}", units_long="Tesla / meter", - description="Contravariant toroidal component of magnetic field, second " - + "derivative wrt poloidal and toroidal angles", + description=( + "Contravariant toroidal component of magnetic field, second derivative wrt" + " poloidal and toroidal coordinates" + ), dim=1, params=[], transforms={}, @@ -990,20 +1233,25 @@ def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): data=[ "B0", "B0_t", - "B0_z", "B0_tz", + "B0_z", + "iota", "lambda_t", "lambda_tt", - "lambda_tz", "lambda_ttz", + "lambda_tz", + "omega_t", + "omega_tt", + "omega_ttz", + "omega_tz", ], ) def _B_sup_zeta_tz(params, transforms, profiles, data, **kwargs): data["B^zeta_tz"] = ( - data["B0_tz"] * (1 + data["lambda_t"]) - + data["B0_t"] * data["lambda_tz"] - + data["B0_z"] * data["lambda_tt"] - + data["B0"] * data["lambda_ttz"] + -data["B0"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) + - data["B0_t"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + - data["B0_z"] * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + + data["B0_tz"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -1064,16 +1312,25 @@ def _B_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["psi_r", "psi_rr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_z", "sqrt(g)_rz"], + axis_limit_data=["psi_rrr", "sqrt(g)_rr", "sqrt(g)_rrz"], ) def _B0_rz(params, transforms, profiles, data, **kwargs): - data["B0_rz"] = ( - -data["psi_rr"] * data["sqrt(g)_z"] / data["sqrt(g)"] ** 2 - - data["psi_r"] * data["sqrt(g)_rz"] / data["sqrt(g)"] ** 2 - + 2 - * data["psi_r"] - * data["sqrt(g)_r"] - * data["sqrt(g)_z"] - / data["sqrt(g)"] ** 3 + data["B0_rz"] = transforms["grid"].replace_at_axis( + ( + -data["sqrt(g)"] + * (data["psi_rr"] * data["sqrt(g)_z"] + data["psi_r"] * data["sqrt(g)_rz"]) + + 2 * data["psi_r"] * data["sqrt(g)_r"] * data["sqrt(g)_z"] + ) + / data["sqrt(g)"] ** 3, + lambda: ( + -data["sqrt(g)_r"] + * ( + data["psi_rrr"] * data["sqrt(g)_rz"] + + data["psi_rr"] * data["sqrt(g)_rrz"] + ) + + 2 * data["psi_rr"] * data["sqrt(g)_rr"] * data["sqrt(g)_rz"] + ) + / (2 * data["sqrt(g)_r"] ** 3), ) return data @@ -1082,9 +1339,11 @@ def _B0_rz(params, transforms, profiles, data, **kwargs): name="B^theta_rz", label="\\partial_{\\rho\\zeta} B^{\\theta}", units="T \\cdot m^{-1}", - units_long="Tesla / meters", - description="Contravariant poloidal component of magnetic field, second " - + "derivative wrt radial coordinate and toroidal angle", + units_long="Tesla / meter", + description=( + "Contravariant poloidal component of magnetic field, second derivative wrt" + " radial and toroidal coordinates" + ), dim=1, params=[], transforms={}, @@ -1093,22 +1352,38 @@ def _B0_rz(params, transforms, profiles, data, **kwargs): data=[ "B0", "B0_r", - "B0_z", "B0_rz", + "B0_z", "iota", "iota_r", - "lambda_z", "lambda_rz", - "lambda_zz", "lambda_rzz", + "lambda_z", + "lambda_zz", + "omega_rz", + "omega_rzz", + "omega_z", + "omega_zz", ], ) def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): data["B^theta_rz"] = ( - data["B0_rz"] * (data["iota"] - data["lambda_z"]) - - data["B0_r"] * data["lambda_zz"] - + data["B0_z"] * (data["iota_r"] - data["lambda_rz"]) - - data["B0"] * data["lambda_rzz"] + data["B0"] + * ( + data["iota"] * data["omega_rzz"] + + data["iota_r"] * data["omega_zz"] + - data["lambda_rzz"] + ) + + data["B0_r"] * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) + + data["B0_z"] + * ( + data["iota"] * data["omega_rz"] + + data["iota_r"] * data["omega_z"] + + data["iota_r"] + - data["lambda_rz"] + ) + + data["B0_rz"] + * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -1117,9 +1392,11 @@ def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): name="B^zeta_rz", label="\\partial_{\\rho\\zeta} B^{\\zeta}", units="T \\cdot m^{-1}", - units_long="Tesla / meters", - description="Contravariant toroidal component of magnetic field, second " - + "derivative wrt radial coordinate and toroidal angle", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, second derivative wrt" + " radial and toroidal coordinates" + ), dim=1, params=[], transforms={}, @@ -1128,20 +1405,36 @@ def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): data=[ "B0", "B0_r", - "B0_z", "B0_rz", - "lambda_t", + "B0_z", + "iota", + "iota_r", "lambda_rt", - "lambda_tz", "lambda_rtz", + "lambda_t", + "lambda_tz", + "omega_rt", + "omega_rtz", + "omega_t", + "omega_tz", ], ) def _B_sup_zeta_rz(params, transforms, profiles, data, **kwargs): data["B^zeta_rz"] = ( - data["B0_rz"] * (1 + data["lambda_t"]) - + data["B0_r"] * data["lambda_tz"] - + data["B0_z"] * data["lambda_rt"] - + data["B0"] * data["lambda_rtz"] + -data["B0"] + * ( + data["iota"] * data["omega_rtz"] + + data["iota_r"] * data["omega_tz"] + - data["lambda_rtz"] + ) + - data["B0_r"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + - data["B0_z"] + * ( + data["iota"] * data["omega_rt"] + + data["iota_r"] * data["omega_t"] + - data["lambda_rt"] + ) + + data["B0_rz"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -2261,19 +2554,16 @@ def _B_mag_rz(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=[ - "|B|_r", - "|B|_t", - "|B|_z", - "e^rho", - "e^theta", - "e^zeta", - ], + data=["|B|_r", "|B|_t", "|B|_z", "e^rho", "e^theta*sqrt(g)", "e^zeta", "sqrt(g)"], + axis_limit_data=["|B|_rt", "sqrt(g)_r"], ) def _grad_B(params, transforms, profiles, data, **kwargs): data["grad(|B|)"] = ( data["|B|_r"] * data["e^rho"].T - + data["|B|_t"] * data["e^theta"].T + + transforms["grid"].replace_at_axis( + data["|B|_t"] / data["sqrt(g)"], lambda: data["|B|_rt"] / data["sqrt(g)_r"] + ) + * data["e^theta*sqrt(g)"].T + data["|B|_z"] * data["e^zeta"].T ).T return data @@ -2331,21 +2621,23 @@ def _B_rms(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=["sqrt(g)", "|B|", "V_r(r)"], + data=["sqrt(g)", "|B|"], + axis_limit_data=["sqrt(g)_r"], ) def _B_fsa(params, transforms, profiles, data, **kwargs): data["<|B|>"] = surface_averages( transforms["grid"], data["|B|"], - data["sqrt(g)"], - denominator=data["V_r(r)"], + sqrt_g=transforms["grid"].replace_at_axis( + data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True + ), ) return data @register_compute_fun( - name="", - label="\\langle B^2 \\rangle", + name="<|B|^2>", + label="\\langle |B|^2 \\rangle", units="T^2", units_long="Tesla squared", description="Flux surface average magnetic field squared", @@ -2354,21 +2646,23 @@ def _B_fsa(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=["sqrt(g)", "|B|^2", "V_r(r)"], + data=["sqrt(g)", "|B|^2"], + axis_limit_data=["sqrt(g)_r"], ) def _B2_fsa(params, transforms, profiles, data, **kwargs): - data[""] = surface_averages( + data["<|B|^2>"] = surface_averages( transforms["grid"], data["|B|^2"], - data["sqrt(g)"], - denominator=data["V_r(r)"], + sqrt_g=transforms["grid"].replace_at_axis( + data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True + ), ) return data @register_compute_fun( name="<1/|B|>", - label="\\langle 1/B \\rangle", + label="\\langle 1/|B| \\rangle", units="T^{-1}", units_long="1 / Tesla", description="Flux surface averaged inverse field strength", @@ -2377,21 +2671,23 @@ def _B2_fsa(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=["sqrt(g)", "|B|", "V_r(r)"], + data=["sqrt(g)", "|B|"], + axis_limit_data=["sqrt(g)_r"], ) def _1_over_B_fsa(params, transforms, profiles, data, **kwargs): data["<1/|B|>"] = surface_averages( transforms["grid"], 1 / data["|B|"], - data["sqrt(g)"], - denominator=data["V_r(r)"], + sqrt_g=transforms["grid"].replace_at_axis( + data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True + ), ) return data @register_compute_fun( - name="_r", - label="\\partial_{\\rho} \\langle B^2 \\rangle", + name="<|B|^2>_r", + label="\\partial_{\\rho} \\langle |B|^2 \\rangle", units="T^2", units_long="Tesla squared", description="Flux surface average magnetic field squared, radial derivative", @@ -2400,32 +2696,29 @@ def _1_over_B_fsa(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=[ - "sqrt(g)", - "sqrt(g)_r", - "B", - "B_r", - "|B|^2", - "", - "V_r(r)", - "V_rr(r)", - ], + data=["sqrt(g)", "sqrt(g)_r", "B", "B_r", "|B|^2", "V_r(r)", "V_rr(r)"], + axis_limit_data=["sqrt(g)_rr", "V_rrr(r)"], ) def _B2_fsa_r(params, transforms, profiles, data, **kwargs): - data["_r"] = ( - surface_integrals( - transforms["grid"], - data["sqrt(g)_r"] * data["|B|^2"] - + data["sqrt(g)"] * 2 * dot(data["B"], data["B_r"]), + integrate = surface_integrals_map(transforms["grid"]) + B2_r = 2 * dot(data["B"], data["B_r"]) + num = integrate(data["sqrt(g)"] * data["|B|^2"]) + num_r = integrate(data["sqrt(g)_r"] * data["|B|^2"] + data["sqrt(g)"] * B2_r) + data["<|B|^2>_r"] = transforms["grid"].replace_at_axis( + (num_r * data["V_r(r)"] - num * data["V_rr(r)"]) / data["V_r(r)"] ** 2, + lambda: ( + integrate(data["sqrt(g)_rr"] * data["|B|^2"] + 2 * data["sqrt(g)_r"] * B2_r) + * data["V_rr(r)"] + - num_r * data["V_rrr(r)"] ) - - data["V_rr(r)"] * data[""] - ) / data["V_r(r)"] + / (2 * data["V_rr(r)"] ** 2), + ) return data @register_compute_fun( name="grad(|B|^2)_rho", - label="(\\nabla B^{2})_{\\rho}", + label="(\\nabla |B|^{2})_{\\rho}", units="T^{2}", units_long="Tesla squared", description="Covariant radial component of magnetic pressure gradient", @@ -2457,7 +2750,7 @@ def _gradB2_rho(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="grad(|B|^2)_theta", - label="(\\nabla B^{2})_{\\theta}", + label="(\\nabla |B|^{2})_{\\theta}", units="T^{2}", units_long="Tesla squared", description="Covariant poloidal component of magnetic pressure gradient", @@ -2489,7 +2782,7 @@ def _gradB2_theta(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="grad(|B|^2)_zeta", - label="(\\nabla B^{2})_{\\zeta}", + label="(\\nabla |B|^{2})_{\\zeta}", units="T^{2}", units_long="Tesla squared", description="Covariant toroidal component of magnetic pressure gradient", @@ -2521,7 +2814,7 @@ def _gradB2_zeta(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="grad(|B|^2)", - label="\\nabla B^{2}", + label="\\nabla |B|^{2}", units="T^{2} \\cdot m^{-1}", units_long="Tesla squared / meters", description="Magnetic pressure gradient", @@ -2530,27 +2823,16 @@ def _gradB2_zeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=[ - "grad(|B|^2)_rho", - "grad(|B|^2)_theta", - "grad(|B|^2)_zeta", - "e^rho", - "e^theta", - "e^zeta", - ], + data=["|B|", "grad(|B|)"], ) def _gradB2(params, transforms, profiles, data, **kwargs): - data["grad(|B|^2)"] = ( - data["grad(|B|^2)_rho"] * data["e^rho"].T - + data["grad(|B|^2)_theta"] * data["e^theta"].T - + data["grad(|B|^2)_zeta"] * data["e^zeta"].T - ).T + data["grad(|B|^2)"] = 2 * (data["|B|"] * data["grad(|B|)"].T).T return data @register_compute_fun( name="|grad(|B|^2)|/2mu0", - label="|\\nabla B^{2}/(2\\mu_0)|", + label="|\\nabla |B|^{2}/(2\\mu_0)|", units="N \\cdot m^{-3}", units_long="Newton / cubic meter", description="Magnitude of magnetic pressure gradient", @@ -2570,7 +2852,7 @@ def _gradB2mag(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="<|grad(|B|^2)|/2mu0>_vol", - label="\\langle |\\nabla B^{2}/(2\\mu_0)| \\rangle_{vol}", + label="\\langle |\\nabla |B|^{2}/(2\\mu_0)| \\rangle_{vol}", units="N \\cdot m^{-3}", units_long="Newtons per cubic meter", description="Volume average of magnitude of magnetic pressure gradient", @@ -2602,14 +2884,12 @@ def _gradB2mag_vol(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "B^theta", "B^zeta", "J^theta", "J^zeta"], + data=["B^theta", "B^zeta", "B_rho_z", "B_zeta_r", "B_theta_r", "B_rho_t"], ) def _curl_B_x_B_rho(params, transforms, profiles, data, **kwargs): - data["(curl(B)xB)_rho"] = ( - mu_0 - * data["sqrt(g)"] - * (data["B^zeta"] * data["J^theta"] - data["B^theta"] * data["J^zeta"]) - ) + data["(curl(B)xB)_rho"] = data["B^zeta"] * ( + data["B_rho_z"] - data["B_zeta_r"] + ) - data["B^theta"] * (data["B_theta_r"] - data["B_rho_t"]) return data @@ -2624,10 +2904,10 @@ def _curl_B_x_B_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "B^zeta", "J^rho"], + data=["B^zeta", "B_zeta_t", "B_theta_z"], ) def _curl_B_x_B_theta(params, transforms, profiles, data, **kwargs): - data["(curl(B)xB)_theta"] = -mu_0 * data["sqrt(g)"] * data["B^zeta"] * data["J^rho"] + data["(curl(B)xB)_theta"] = -data["B^zeta"] * (data["B_zeta_t"] - data["B_theta_z"]) return data @@ -2642,10 +2922,10 @@ def _curl_B_x_B_theta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["sqrt(g)", "B^theta", "J^rho"], + data=["B^theta", "B_zeta_t", "B_theta_z"], ) def _curl_B_x_B_zeta(params, transforms, profiles, data, **kwargs): - data["(curl(B)xB)_zeta"] = mu_0 * data["sqrt(g)"] * data["B^theta"] * data["J^rho"] + data["(curl(B)xB)_zeta"] = data["B^theta"] * (data["B_zeta_t"] - data["B_theta_z"]) return data @@ -2662,17 +2942,19 @@ def _curl_B_x_B_zeta(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=[ "(curl(B)xB)_rho", - "(curl(B)xB)_theta", + "B^zeta", + "J^rho", "(curl(B)xB)_zeta", "e^rho", - "e^theta", + "e^theta*sqrt(g)", "e^zeta", ], ) def _curl_B_x_B(params, transforms, profiles, data, **kwargs): + # (curl(B)xB)_theta e^theta refactored to resolve indeterminacy at axis. data["curl(B)xB"] = ( data["(curl(B)xB)_rho"] * data["e^rho"].T - + data["(curl(B)xB)_theta"] * data["e^theta"].T + - mu_0 * data["B^zeta"] * data["J^rho"] * data["e^theta*sqrt(g)"].T + data["(curl(B)xB)_zeta"] * data["e^zeta"].T ).T return data @@ -2809,7 +3091,38 @@ def _B_dot_gradB(params, transforms, profiles, data, **kwargs): return data -# TODO: (B*grad(|B|))_r +@register_compute_fun( + name="(B*grad(|B|))_r", + label="\\partial_{\\theta} (\\mathbf{B} \\cdot \\nabla B)", + units="T^2 \\cdot m^{-1}", + units_long="Tesla squared / meters", + description="", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "B^theta", + "B^zeta", + "|B|_t", + "|B|_z", + "B^theta_r", + "B^zeta_r", + "|B|_rt", + "|B|_rz", + ], +) +def _B_dot_gradB_r(params, transforms, profiles, data, **kwargs): + data["(B*grad(|B|))_r"] = ( + data["B^theta_r"] * data["|B|_t"] + + data["B^theta"] * data["|B|_rt"] + + data["B^zeta_r"] * data["|B|_z"] + + data["B^zeta"] * data["|B|_rz"] + ) + return data + + @register_compute_fun( name="(B*grad(|B|))_t", label="\\partial_{\\theta} (\\mathbf{B} \\cdot \\nabla B)", @@ -2924,18 +3237,14 @@ def _min_tz_modB(params, transforms, profiles, data, **kwargs): data=["max_tz |B|", "min_tz |B|"], ) def _effective_r_over_R0(params, transforms, profiles, data, **kwargs): - r""" - Compute an effective local inverse aspect ratio. + """Compute an effective local inverse aspect ratio. This effective local inverse aspect ratio epsilon is defined by - - .. math:: - \frac{Bmax}{Bmin} = \frac{1 + \epsilon}{1 - \epsilon} + Bmax / Bmin = (1 + ε) / (1 − ε). This definition is motivated by the fact that this formula would be true in the case of circular cross-section surfaces in - axisymmetry with :math:`B \propto 1/R` and :math:`R = (1 + - \epsilon \cos\theta) R_0`. + axisymmetry with B ∝ 1/R and R = (1 + ε cos θ) R₀. """ w = data["max_tz |B|"] / data["min_tz |B|"] data["effective r/R0"] = (w - 1) / (w + 1) @@ -3010,13 +3319,21 @@ def _kappa_g(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B_r", "B_t", "B_z", "e^rho", "e^theta", "e^zeta"], + data=["B_r", "B_t", "B_z", "e^rho", "e^theta*sqrt(g)", "e^zeta", "sqrt(g)"], + axis_limit_data=["B_rt", "sqrt(g)_r"], ) def _grad_B_vec(params, transforms, profiles, data, **kwargs): + B_t_over_sqrt_g = transforms["grid"].replace_at_axis( + (data["B_t"].T / data["sqrt(g)"]).T, + lambda: (data["B_rt"].T / data["sqrt(g)_r"]).T, + ) data["grad(B)"] = ( - (data["B_r"][:, None, :] * data["e^rho"][:, :, None]) - + (data["B_t"][:, None, :] * data["e^theta"][:, :, None]) - + (data["B_z"][:, None, :] * data["e^zeta"][:, :, None]) + (data["B_r"][:, jnp.newaxis, :] * data["e^rho"][:, :, jnp.newaxis]) + + ( + B_t_over_sqrt_g[:, jnp.newaxis, :] + * data["e^theta*sqrt(g)"][:, :, jnp.newaxis] + ) + + (data["B_z"][:, jnp.newaxis, :] * data["e^zeta"][:, :, jnp.newaxis]) ) return data diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 157e2bb0f5..622697afe1 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -1,3 +1,14 @@ +"""Compute functions for quantities with obvious geometric meaning. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" + from desc.backend import jnp from .data_index import register_compute_fun @@ -105,12 +116,31 @@ def _V_r_of_r(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=["sqrt(g)_r", "sqrt(g)"], + data=["sqrt(g)_r"], ) def _V_rr_of_r(params, transforms, profiles, data, **kwargs): - data["V_rr(r)"] = surface_integrals( - transforms["grid"], data["sqrt(g)_r"] * jnp.sign(data["sqrt(g)"]) - ) + # The sign of sqrt(g) is enforced to be non-negative. + data["V_rr(r)"] = surface_integrals(transforms["grid"], data["sqrt(g)_r"]) + return data + + +@register_compute_fun( + name="V_rrr(r)", + label="\\partial_{\\rho\\rho\\rho} V(\\rho)", + units="m^{3}", + units_long="cubic meters", + description="Volume enclosed by flux surfaces, third derivative wrt radial " + + "coordinate", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["sqrt(g)_rr"], +) +def _V_rrr_of_r(params, transforms, profiles, data, **kwargs): + # The sign of sqrt(g) is enforced to be non-negative. + data["V_rrr(r)"] = surface_integrals(transforms["grid"], data["sqrt(g)_rr"]) return data @@ -207,6 +237,45 @@ def _S_of_r(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="S_r(r)", + label="\\partial_{\\rho} S(\\rho)", + units="m^{2}", + units_long="square meters", + description="Surface area of flux surfaces, derivative wrt radial coordinate", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["|e_theta x e_zeta|_r"], +) +def _S_r_of_r(params, transforms, profiles, data, **kwargs): + data["S_r(r)"] = surface_integrals(transforms["grid"], data["|e_theta x e_zeta|_r"]) + return data + + +@register_compute_fun( + name="S_rr(r)", + label="\\partial_{\\rho\\rho} S(\\rho)", + units="m^{2}", + units_long="square meters", + description="Surface area of flux surfaces, second derivative wrt radial" + " coordinate", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["|e_theta x e_zeta|_rr"], +) +def _S_rr_of_r(params, transforms, profiles, data, **kwargs): + data["S_rr(r)"] = surface_integrals( + transforms["grid"], data["|e_theta x e_zeta|_rr"] + ) + return data + + @register_compute_fun( name="R0", label="R_{0}", @@ -276,7 +345,7 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs): ) def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs): max_rho = transforms["grid"].nodes[transforms["grid"].unique_rho_idx[-1], 0] - P = ( # perimeter + P = ( # perimeter at rho=1 line_integrals( transforms["grid"], jnp.sqrt(data["g_tt"]), @@ -416,6 +485,9 @@ def _curvature_k1_rho(params, transforms, profiles, data, **kwargs): c = L * N - M**2 r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) + # In the axis limit, the matrix of the first fundamental form is singular. + # The diagonal of the shape operator becomes unbounded, + # so the eigenvalues do not exist. data["curvature_k1_rho"] = jnp.maximum(r1, r2) data["curvature_k2_rho"] = jnp.minimum(r1, r2) return data @@ -452,6 +524,9 @@ def _curvature_k2_rho(params, transforms, profiles, data, **kwargs): c = L * N - M**2 r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) + # In the axis limit, the matrix of the first fundamental form is singular. + # The diagonal of the shape operator becomes unbounded, + # so the eigenvalues do not exist. data["curvature_k1_rho"] = jnp.maximum(r1, r2) data["curvature_k2_rho"] = jnp.minimum(r1, r2) return data @@ -774,6 +849,9 @@ def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs): c = L * N - M**2 r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) + # In the axis limit, the matrix of the first fundamental form is singular. + # The diagonal of the shape operator becomes unbounded, + # so the eigenvalues do not exist. data["curvature_k1_zeta"] = jnp.maximum(r1, r2) data["curvature_k2_zeta"] = jnp.minimum(r1, r2) return data @@ -810,6 +888,9 @@ def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs): c = L * N - M**2 r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) + # In the axis limit, the matrix of the first fundamental form is singular. + # The diagonal of the shape operator becomes unbounded, + # so the eigenvalues do not exist. data["curvature_k1_zeta"] = jnp.maximum(r1, r2) data["curvature_k2_zeta"] = jnp.minimum(r1, r2) return data diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index 49b65552f9..9088144c09 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -1,3 +1,14 @@ +"""Compute functions related to the metric tensor of the coordinate system. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" + from desc.backend import jnp from .data_index import register_compute_fun @@ -33,11 +44,11 @@ def _sqrtg(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_theta_PEST", "e_zeta"], + data=["e_rho", "e_theta_PEST", "e_phi"], ) def _sqrtg_pest(params, transforms, profiles, data, **kwargs): data["sqrt(g)_PEST"] = dot( - data["e_rho"], cross(data["e_theta_PEST"], data["e_zeta"]) + data["e_rho"], cross(data["e_theta_PEST"], data["e_phi"]) ) return data @@ -61,7 +72,80 @@ def _sqrtg_pest(params, transforms, profiles, data, **kwargs): ) def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs): data["|e_theta x e_zeta|"] = jnp.linalg.norm( - cross(data["e_theta"], data["e_zeta"]), axis=1 + cross(data["e_theta"], data["e_zeta"]), axis=-1 + ) + return data + + +@register_compute_fun( + name="|e_theta x e_zeta|_r", + label="\\partial_{\\rho} |e_{\\theta} \\times e_{\\zeta}|", + units="m^{2}", + units_long="square meters", + description="2D Jacobian determinant for constant rho surface" + + " derivative wrt radial coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], +) +def _e_theta_x_e_zeta_r(params, transforms, profiles, data, **kwargs): + a = cross(data["e_theta"], data["e_zeta"]) + a_r = cross(data["e_theta_r"], data["e_zeta"]) + cross( + data["e_theta"], data["e_zeta_r"] + ) + # The limit of a sequence and the norm function can be interchanged + # because norms are continuous functions. Likewise with dot product. + # Then lim ‖𝐞^ρ‖ = ‖ lim 𝐞^ρ ‖ ≠ 0 + # lim (𝐞^ρ ⋅ a_r / ‖𝐞^ρ‖) = lim 𝐞^ρ ⋅ lim a_r / lim ‖𝐞^ρ‖ + # The vectors converge to be parallel. + data["|e_theta x e_zeta|_r"] = transforms["grid"].replace_at_axis( + dot(a, a_r) / jnp.linalg.norm(a, axis=-1), lambda: jnp.linalg.norm(a_r, axis=-1) + ) + return data + + +@register_compute_fun( + name="|e_theta x e_zeta|_rr", + label="\\partial_{\\rho \\rho} |e_{\\theta} \\times e_{\\zeta}|", + units="m^{2}", + units_long="square meters", + description="2D Jacobian determinant for constant rho surface" + + " second derivative wrt radial coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r", "e_theta_rr", "e_zeta_rr"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], +) +def _e_theta_x_e_zeta_rr(params, transforms, profiles, data, **kwargs): + a = cross(data["e_theta"], data["e_zeta"]) + a_r = cross(data["e_theta_r"], data["e_zeta"]) + cross( + data["e_theta"], data["e_zeta_r"] + ) + a_rr = ( + cross(data["e_theta_rr"], data["e_zeta"]) + + 2 * cross(data["e_theta_r"], data["e_zeta_r"]) + + cross(data["e_theta"], data["e_zeta_rr"]) + ) + norm_a = jnp.linalg.norm(a, axis=-1) + norm_a_r = jnp.linalg.norm(a_r, axis=-1) + # The limit eventually reduces to a form where the technique used to compute + # lim |e_theta x e_zeta|_r can be applied. + data["|e_theta x e_zeta|_rr"] = transforms["grid"].replace_at_axis( + (norm_a_r**2 + dot(a, a_rr) - (dot(a, a_r) / norm_a) ** 2) / norm_a, + lambda: dot(a_r, a_rr) / norm_a_r, ) return data @@ -77,16 +161,14 @@ def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_zeta"], + data=["e^theta*sqrt(g)"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.core.Surface", ], ) def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs): - data["|e_zeta x e_rho|"] = jnp.linalg.norm( - cross(data["e_zeta"], data["e_rho"]), axis=1 - ) + data["|e_zeta x e_rho|"] = jnp.linalg.norm(data["e^theta*sqrt(g)"], axis=-1) return data @@ -109,7 +191,80 @@ def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs): ) def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs): data["|e_rho x e_theta|"] = jnp.linalg.norm( - cross(data["e_rho"], data["e_theta"]), axis=1 + cross(data["e_rho"], data["e_theta"]), axis=-1 + ) + return data + + +@register_compute_fun( + name="|e_rho x e_theta|_r", + label="\\partial_{\\rho} |e_{\\rho} \\times e_{\\theta}|", + units="m^{2}", + units_long="square meters", + description="2D Jacobian determinant for constant zeta surface" + " derivative wrt radial coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_theta", "e_rho_r", "e_theta_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], +) +def _e_rho_x_e_theta_r(params, transforms, profiles, data, **kwargs): + a = cross(data["e_rho"], data["e_theta"]) + a_r = cross(data["e_rho_r"], data["e_theta"]) + cross( + data["e_rho"], data["e_theta_r"] + ) + # The limit of a sequence and the norm function can be interchanged + # because norms are continuous functions. Likewise with dot product. + # Then lim ‖𝐞^ζ‖ = ‖ lim 𝐞^ζ ‖ ≠ 0 + # lim (𝐞^ζ ⋅ a_r / ‖𝐞^ζ‖) = lim 𝐞^ζ ⋅ lim a_r / lim ‖𝐞^ζ‖ + # The vectors converge to be parallel. + data["|e_rho x e_theta|_r"] = transforms["grid"].replace_at_axis( + dot(a, a_r) / jnp.linalg.norm(a, axis=-1), lambda: jnp.linalg.norm(a_r, axis=-1) + ) + return data + + +@register_compute_fun( + name="|e_rho x e_theta|_rr", + label="\\partial_{\\rho \\rho} |e_{\\rho} \\times e_{\\theta}|", + units="m^{2}", + units_long="square meters", + description="2D Jacobian determinant for constant zeta surface" + + " second derivative wrt radial coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_theta", "e_rho_r", "e_theta_r", "e_rho_rr", "e_theta_rr"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], +) +def _e_rho_x_e_theta_rr(params, transforms, profiles, data, **kwargs): + a = cross(data["e_rho"], data["e_theta"]) + a_r = cross(data["e_rho_r"], data["e_theta"]) + cross( + data["e_rho"], data["e_theta_r"] + ) + a_rr = ( + cross(data["e_rho_rr"], data["e_theta"]) + + 2 * cross(data["e_rho_r"], data["e_theta_r"]) + + cross(data["e_rho"], data["e_theta_rr"]) + ) + norm_a = jnp.linalg.norm(a, axis=-1) + norm_a_r = jnp.linalg.norm(a_r, axis=-1) + # The limit eventually reduces to a form where the technique used to compute + # lim |e_rho x e_theta|_r can be applied. + data["|e_rho x e_theta|_rr"] = transforms["grid"].replace_at_axis( + (norm_a_r**2 + dot(a, a_rr) - (dot(a, a_r) / norm_a) ** 2) / norm_a, + lambda: dot(a_r, a_rr) / norm_a_r, ) return data @@ -219,6 +374,122 @@ def _sqrtg_rr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="sqrt(g)_rrr", + label="\\partial_{\\rho\\rho\\rho} \\sqrt{g}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of flux coordinate system, third derivative wrt " + + "radial coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_rho", + "e_theta", + "e_zeta", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + "e_rho_rr", + "e_theta_rr", + "e_zeta_rr", + "e_rho_rrr", + "e_theta_rrr", + "e_zeta_rrr", + ], +) +def _sqrtg_rrr(params, transforms, profiles, data, **kwargs): + data["sqrt(g)_rrr"] = ( + dot(data["e_rho_rrr"], cross(data["e_theta"], data["e_zeta"])) + + dot(data["e_rho"], cross(data["e_theta_rrr"], data["e_zeta"])) + + dot(data["e_rho"], cross(data["e_theta"], data["e_zeta_rrr"])) + + 3 * dot(data["e_rho_rr"], cross(data["e_theta_r"], data["e_zeta"])) + + 3 * dot(data["e_rho_rr"], cross(data["e_theta"], data["e_zeta_r"])) + + 3 * dot(data["e_rho_r"], cross(data["e_theta_rr"], data["e_zeta"])) + + 3 * dot(data["e_rho"], cross(data["e_theta_rr"], data["e_zeta_r"])) + + 3 * dot(data["e_rho_r"], cross(data["e_theta"], data["e_zeta_rr"])) + + 3 * dot(data["e_rho"], cross(data["e_theta_r"], data["e_zeta_rr"])) + + 6 * dot(data["e_rho_r"], cross(data["e_theta_r"], data["e_zeta_r"])) + ) + return data + + +@register_compute_fun( + name="sqrt(g)_rrt", + label="\\partial_{\\rho\\rho\\theta} \\sqrt{g}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of flux coordinate system, third derivative wrt " + + "radial coordinate twice and poloidal angle once", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_rho", + "e_theta", + "e_zeta", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + "e_rho_t", + "e_theta_t", + "e_zeta_t", + "e_rho_rt", + "e_theta_rt", + "e_zeta_rt", + "e_rho_rr", + "e_theta_rr", + "e_zeta_rr", + "e_rho_rrt", + "e_theta_rrt", + "e_zeta_rrt", + ], +) +def _sqrtg_rrt(params, transforms, profiles, data, **kwargs): + data["sqrt(g)_rrt"] = ( + dot(data["e_rho_rrt"], cross(data["e_theta"], data["e_zeta"])) + + dot( + data["e_rho_rr"], + cross(data["e_theta_t"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_t"]), + ) + + 2 + * dot( + data["e_rho_rt"], + cross(data["e_theta_r"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_r"]), + ) + + 2 + * dot( + data["e_rho_r"], + cross(data["e_theta_rt"], data["e_zeta"]) + + cross(data["e_theta_r"], data["e_zeta_t"]) + + cross(data["e_theta_t"], data["e_zeta_r"]) + + cross(data["e_theta"], data["e_zeta_rt"]), + ) + + dot( + data["e_rho_t"], + cross(data["e_theta_rr"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_rr"]), + ) + + dot( + data["e_rho"], + cross(data["e_theta_rrt"], data["e_zeta"]) + + 2 * cross(data["e_theta_rt"], data["e_zeta_r"]) + + cross(data["e_theta_rr"], data["e_zeta_t"]) + + 2 * cross(data["e_theta_r"], data["e_zeta_rt"]) + + cross(data["e_theta_t"], data["e_zeta_rr"]) + + cross(data["e_theta"], data["e_zeta_rrt"]), + ) + ) + return data + + @register_compute_fun( name="sqrt(g)_tt", label="\\partial_{\\theta\\theta} \\sqrt{g}", @@ -255,6 +526,62 @@ def _sqrtg_tt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="sqrt(g)_rtt", + label="\\partial_{\\rho\\theta\\theta} \\sqrt{g}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of flux coordinate system, third derivative wrt" + + " radial coordinate once and poloidal angle twice.", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_rho", + "e_theta", + "e_zeta", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + "e_rho_t", + "e_theta_t", + "e_zeta_t", + "e_rho_rt", + "e_theta_rt", + "e_zeta_rt", + "e_rho_tt", + "e_theta_tt", + "e_zeta_tt", + "e_rho_rtt", + "e_theta_rtt", + "e_zeta_rtt", + ], +) +def _sqrtg_rtt(params, transforms, profiles, data, **kwargs): + data["sqrt(g)_rtt"] = ( + dot(data["e_rho_rtt"], cross(data["e_theta"], data["e_zeta"])) + + dot(data["e_rho_r"], cross(data["e_theta_tt"], data["e_zeta"])) + + dot(data["e_rho_r"], cross(data["e_theta"], data["e_zeta_tt"])) + + 2 * dot(data["e_rho_rt"], cross(data["e_theta_t"], data["e_zeta"])) + + 2 * dot(data["e_rho_rt"], cross(data["e_theta"], data["e_zeta_t"])) + + 2 * dot(data["e_rho_r"], cross(data["e_theta_t"], data["e_zeta_t"])) + + dot(data["e_rho_tt"], cross(data["e_theta_r"], data["e_zeta"])) + + dot(data["e_rho"], cross(data["e_theta_rtt"], data["e_zeta"])) + + dot(data["e_rho"], cross(data["e_theta_r"], data["e_zeta_tt"])) + + 2 * dot(data["e_rho_t"], cross(data["e_theta_rt"], data["e_zeta"])) + + 2 * dot(data["e_rho"], cross(data["e_theta_rt"], data["e_zeta_t"])) + + dot(data["e_rho_tt"], cross(data["e_theta"], data["e_zeta_r"])) + + dot(data["e_rho"], cross(data["e_theta_tt"], data["e_zeta_r"])) + + dot(data["e_rho"], cross(data["e_theta"], data["e_zeta_rtt"])) + + 2 * dot(data["e_rho_t"], cross(data["e_theta_t"], data["e_zeta_r"])) + + 2 * dot(data["e_rho_t"], cross(data["e_theta"], data["e_zeta_rt"])) + + 2 * dot(data["e_rho"], cross(data["e_theta_t"], data["e_zeta_rt"])) + ) + return data + + @register_compute_fun( name="sqrt(g)_zz", label="\\partial_{\\zeta\\zeta} \\sqrt{g}", @@ -291,6 +618,62 @@ def _sqrtg_zz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="sqrt(g)_rzz", + label="\\partial_{\\rho\\zeta\\zeta} \\sqrt{g}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of flux coordinate system, third derivative wrt " + + "radial coordinate once and toroidal angle twice", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_rho", + "e_theta", + "e_zeta", + "e_rho_z", + "e_theta_z", + "e_zeta_z", + "e_rho_zz", + "e_theta_zz", + "e_zeta_zz", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + "e_rho_rz", + "e_theta_rz", + "e_zeta_rz", + "e_rho_rzz", + "e_theta_rzz", + "e_zeta_rzz", + ], +) +def _sqrtg_rzz(params, transforms, profiles, data, **kwargs): + data["sqrt(g)_rzz"] = ( + dot(data["e_rho_rzz"], cross(data["e_theta"], data["e_zeta"])) + + dot(data["e_rho_r"], cross(data["e_theta_zz"], data["e_zeta"])) + + dot(data["e_rho_r"], cross(data["e_theta"], data["e_zeta_zz"])) + + 2 * dot(data["e_rho_rz"], cross(data["e_theta_z"], data["e_zeta"])) + + 2 * dot(data["e_rho_rz"], cross(data["e_theta"], data["e_zeta_z"])) + + 2 * dot(data["e_rho_r"], cross(data["e_theta_z"], data["e_zeta_z"])) + + dot(data["e_rho_zz"], cross(data["e_theta_r"], data["e_zeta"])) + + dot(data["e_rho"], cross(data["e_theta_rzz"], data["e_zeta"])) + + dot(data["e_rho"], cross(data["e_theta_r"], data["e_zeta_zz"])) + + 2 * dot(data["e_rho_z"], cross(data["e_theta_rz"], data["e_zeta"])) + + 2 * dot(data["e_rho_z"], cross(data["e_theta_r"], data["e_zeta_z"])) + + 2 * dot(data["e_rho"], cross(data["e_theta_rz"], data["e_zeta_z"])) + + dot(data["e_rho_zz"], cross(data["e_theta"], data["e_zeta_r"])) + + dot(data["e_rho"], cross(data["e_theta_zz"], data["e_zeta_r"])) + + dot(data["e_rho"], cross(data["e_theta"], data["e_zeta_rzz"])) + + 2 * dot(data["e_rho_z"], cross(data["e_theta"], data["e_zeta_rz"])) + + 2 * dot(data["e_rho"], cross(data["e_theta_z"], data["e_zeta_rz"])) + ) + return data + + @register_compute_fun( name="sqrt(g)_rt", label="\\partial_{\\rho\\theta} \\sqrt{g}", @@ -323,7 +706,6 @@ def _sqrtg_rt(params, transforms, profiles, data, **kwargs): dot(data["e_rho_rt"], cross(data["e_theta"], data["e_zeta"])) + dot(data["e_rho_r"], cross(data["e_theta_t"], data["e_zeta"])) + dot(data["e_rho_r"], cross(data["e_theta"], data["e_zeta_t"])) - + dot(data["e_rho_t"], cross(data["e_theta_r"], data["e_zeta"])) + dot(data["e_rho"], cross(data["e_theta_rt"], data["e_zeta"])) + dot(data["e_rho"], cross(data["e_theta_r"], data["e_zeta_t"])) + dot(data["e_rho_t"], cross(data["e_theta"], data["e_zeta_r"])) @@ -367,7 +749,6 @@ def _sqrtg_tz(params, transforms, profiles, data, **kwargs): + dot(data["e_rho_z"], cross(data["e_theta"], data["e_zeta_t"])) + dot(data["e_rho_t"], cross(data["e_theta_z"], data["e_zeta"])) + dot(data["e_rho"], cross(data["e_theta_tz"], data["e_zeta"])) - + dot(data["e_rho"], cross(data["e_theta_z"], data["e_zeta_t"])) + dot(data["e_rho_t"], cross(data["e_theta"], data["e_zeta_z"])) + dot(data["e_rho"], cross(data["e_theta_t"], data["e_zeta_z"])) + dot(data["e_rho"], cross(data["e_theta"], data["e_zeta_tz"])) @@ -375,6 +756,96 @@ def _sqrtg_tz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="sqrt(g)_rtz", + label="\\partial_{\\rho\\theta\\zeta} \\sqrt{g}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of flux coordinate system, third derivative wrt " + + "radial, poloidal, and toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_rho", + "e_theta", + "e_zeta", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + "e_rho_t", + "e_theta_t", + "e_zeta_t", + "e_rho_rt", + "e_theta_rt", + "e_zeta_rt", + "e_rho_z", + "e_theta_z", + "e_zeta_z", + "e_rho_rz", + "e_theta_rz", + "e_zeta_rz", + "e_rho_tz", + "e_theta_tz", + "e_zeta_tz", + "e_rho_rtz", + "e_theta_rtz", + "e_zeta_rtz", + ], +) +def _sqrtg_rtz(params, transforms, profiles, data, **kwargs): + data["sqrt(g)_rtz"] = ( + dot(data["e_rho_rtz"], cross(data["e_theta"], data["e_zeta"])) + + dot( + data["e_rho_rz"], + cross(data["e_theta_t"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_t"]), + ) + + dot( + data["e_rho_rt"], + cross(data["e_theta_z"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_z"]), + ) + + dot( + data["e_rho_r"], + cross(data["e_theta_tz"], data["e_zeta"]) + + cross(data["e_theta_t"], data["e_zeta_z"]) + + cross(data["e_theta"], data["e_zeta_tz"]), + ) + + dot( + data["e_rho_tz"], + cross(data["e_theta_r"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_r"]), + ) + + dot( + data["e_rho_z"], + cross(data["e_theta_rt"], data["e_zeta"]) + + cross(data["e_theta_r"], data["e_zeta_t"]) + + cross(data["e_theta"], data["e_zeta_rt"]), + ) + + dot( + data["e_rho_t"], + cross(data["e_theta_rz"], data["e_zeta"]) + + cross(data["e_theta_z"], data["e_zeta_r"]) + + cross(data["e_theta"], data["e_zeta_rz"]), + ) + + dot( + data["e_rho"], + cross(data["e_theta_rtz"], data["e_zeta"]) + + cross(data["e_theta_tz"], data["e_zeta_r"]) + + cross(data["e_theta_rz"], data["e_zeta_t"]) + + cross(data["e_theta_z"], data["e_zeta_rt"]) + + cross(data["e_theta_rt"], data["e_zeta_z"]) + + cross(data["e_theta_t"], data["e_zeta_rz"]) + + cross(data["e_theta_r"], data["e_zeta_tz"]) + + cross(data["e_theta"], data["e_zeta_rtz"]), + ) + ) + return data + + @register_compute_fun( name="sqrt(g)_rz", label="\\partial_{\\rho\\zeta} \\sqrt{g}", @@ -410,13 +881,85 @@ def _sqrtg_rz(params, transforms, profiles, data, **kwargs): + dot(data["e_rho_z"], cross(data["e_theta_r"], data["e_zeta"])) + dot(data["e_rho"], cross(data["e_theta_rz"], data["e_zeta"])) + dot(data["e_rho"], cross(data["e_theta_r"], data["e_zeta_z"])) - + dot(data["e_rho_z"], cross(data["e_theta"], data["e_zeta_r"])) + dot(data["e_rho"], cross(data["e_theta_z"], data["e_zeta_r"])) + dot(data["e_rho"], cross(data["e_theta"], data["e_zeta_rz"])) ) return data +@register_compute_fun( + name="sqrt(g)_rrz", + label="\\partial_{\\rho\\rho\\zeta} \\sqrt{g}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of flux coordinate system, third derivative wrt " + + "radial coordinate twice and toroidal angle once", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_rho", + "e_theta", + "e_zeta", + "e_rho_r", + "e_theta_r", + "e_zeta_r", + "e_rho_z", + "e_theta_z", + "e_zeta_z", + "e_rho_rr", + "e_rho_rz", + "e_theta_rr", + "e_theta_rz", + "e_zeta_rz", + "e_zeta_rr", + "e_rho_rrz", + "e_theta_rrz", + "e_zeta_rrz", + ], +) +def _sqrtg_rrz(params, transforms, profiles, data, **kwargs): + data["sqrt(g)_rrz"] = ( + dot(data["e_rho_rrz"], cross(data["e_theta"], data["e_zeta"])) + + dot( + data["e_rho_rr"], + cross(data["e_theta_z"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_z"]), + ) + + 2 + * dot( + data["e_rho_rz"], + cross(data["e_theta_r"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_r"]), + ) + + 2 + * dot( + data["e_rho_r"], + cross(data["e_theta_rz"], data["e_zeta"]) + + cross(data["e_theta_r"], data["e_zeta_z"]) + + cross(data["e_theta_z"], data["e_zeta_r"]) + + cross(data["e_theta"], data["e_zeta_rz"]), + ) + + dot( + data["e_rho_z"], + cross(data["e_theta_rr"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_rr"]), + ) + + dot( + data["e_rho"], + cross(data["e_theta_rrz"], data["e_zeta"]) + + cross(data["e_theta_rr"], data["e_zeta_z"]) + + 2 * cross(data["e_theta_r"], data["e_zeta_rz"]) + + 2 * cross(data["e_theta_rz"], data["e_zeta_r"]) + + cross(data["e_theta_z"], data["e_zeta_rr"]) + + cross(data["e_theta"], data["e_zeta_rrz"]), + ) + ) + return data + + @register_compute_fun( name="g_rr", label="g_{\\rho\\rho}", @@ -619,6 +1162,27 @@ def _g_sub_tt_rr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="g_tt_rrr", + label="\\partial_{\\rho\\rho\\rho} g_{\\theta\\theta}", + units="m^{2}", + units_long="square meters", + description="Poloidal/Poloidal element of covariant metric tensor, third " + + "derivative wrt rho", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "e_theta_r", "e_theta_rr", "e_theta_rrr"], +) +def _g_sub_tt_rrr(params, transforms, profiles, data, **kwargs): + data["g_tt_rrr"] = 6 * dot(data["e_theta_rr"], data["e_theta_r"]) + 2 * dot( + data["e_theta"], data["e_theta_rrr"] + ) + return data + + @register_compute_fun( name="g_tz_rr", label="\\partial_{\\rho\\rho} g_{\\theta\\zeta}", @@ -642,6 +1206,39 @@ def _g_sub_tz_rr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="g_tz_rrr", + label="\\partial_{\\rho\\rho\\rho} g_{\\theta\\zeta}", + units="m^{2}", + units_long="square meters", + description="Poloidal/Toroidal element of covariant metric tensor, third " + + "derivative wrt rho", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "e_theta", + "e_zeta", + "e_theta_r", + "e_zeta_r", + "e_theta_rr", + "e_zeta_rr", + "e_theta_rrr", + "e_zeta_rrr", + ], +) +def _g_sub_tz_rrr(params, transforms, profiles, data, **kwargs): + data["g_tz_rrr"] = ( + dot(data["e_theta_rrr"], data["e_zeta"]) + + 3 * dot(data["e_theta_rr"], data["e_zeta_r"]) + + 3 * dot(data["e_theta_r"], data["e_zeta_rr"]) + + dot(data["e_theta"], data["e_zeta_rrr"]) + ) + return data + + @register_compute_fun( name="g^rr", label="g^{\\rho\\rho}", @@ -1128,6 +1725,42 @@ def _gradrho(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="|grad(psi)|", + label="|\\nabla\\psi|", + units="Wb / m", + units_long="Webers per meter", + description="Toroidal flux gradient (normalized by 2pi) magnitude", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|grad(psi)|^2"], +) +def _gradpsi_mag(params, transforms, profiles, data, **kwargs): + data["|grad(psi)|"] = jnp.sqrt(data["|grad(psi)|^2"]) + return data + + +@register_compute_fun( + name="|grad(psi)|^2", + label="|\\nabla\\psi|^{2}", + units="(Wb / m)^{2}", + units_long="Webers squared per square meter", + description="Toroidal flux gradient (normalized by 2pi) magnitude squared", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["grad(psi)"], +) +def _gradpsi_mag2(params, transforms, profiles, data, **kwargs): + data["|grad(psi)|^2"] = dot(data["grad(psi)"], data["grad(psi)"]) + return data + + @register_compute_fun( name="|grad(theta)|", label="|\\nabla \\theta|", diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index 17fc73b189..8eaea6c343 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -1,9 +1,20 @@ +"""Compute functions for profiles. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" + from scipy.constants import elementary_charge, mu_0 -from desc.backend import jnp, put +from desc.backend import jnp from .data_index import register_compute_fun -from .utils import compress, cumtrapz, dot, expand, surface_averages +from .utils import cumtrapz, dot, surface_averages, surface_integrals @register_compute_fun( @@ -97,60 +108,6 @@ def _psi_rrr(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="grad(psi)", - label="\\nabla\\psi", - units="Wb / m", - units_long="Webers per meter", - description="Toroidal flux gradient (normalized by 2pi)", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["psi_r", "e^rho"], -) -def _gradpsi(params, transforms, profiles, data, **kwargs): - data["grad(psi)"] = (data["psi_r"] * data["e^rho"].T).T - return data - - -@register_compute_fun( - name="|grad(psi)|^2", - label="|\\nabla\\psi|^{2}", - units="(Wb / m)^{2}", - units_long="Webers squared per square meter", - description="Toroidal flux gradient (normalized by 2pi) magnitude squared", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["grad(psi)"], -) -def _gradpsi_mag2(params, transforms, profiles, data, **kwargs): - data["|grad(psi)|^2"] = dot(data["grad(psi)"], data["grad(psi)"]) - return data - - -@register_compute_fun( - name="|grad(psi)|", - label="|\\nabla\\psi|", - units="Wb / m", - units_long="Webers per meter", - description="Toroidal flux gradient (normalized by 2pi) magnitude", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["|grad(psi)|^2"], -) -def _gradpsi_mag(params, transforms, profiles, data, **kwargs): - data["|grad(psi)|"] = jnp.sqrt(data["|grad(psi)|^2"]) - return data - - @register_compute_fun( name="chi_r", label="\\partial_{\\rho} \\chi", @@ -180,15 +137,12 @@ def _chi_r(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=[ - "chi_r", - ], + data=["chi_r", "rho"], ) def _chi(params, transforms, profiles, data, **kwargs): - chi_r = compress(transforms["grid"], data["chi_r"], surface_label="rho") - rho = transforms["grid"].nodes[transforms["grid"].unique_rho_idx, 0] - chi = cumtrapz(chi_r, rho, initial=0) - data["chi"] = expand(transforms["grid"], chi, surface_label="rho") + chi_r = transforms["grid"].compress(data["chi_r"]) + chi = cumtrapz(chi_r, transforms["grid"].compress(data["rho"]), initial=0) + data["chi"] = transforms["grid"].expand(chi) return data @@ -234,6 +188,27 @@ def _Te_r(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Te_rr", + label="\\partial_{\\rho \\rho} T_e", + units="eV", + units_long="electron-Volts", + description="Electron temperature, second radial derivative", + dim=1, + params=["Te_l"], + transforms={}, + profiles=["electron_temperature"], + coordinates="r", + data=["0"], +) +def _Te_rr(params, transforms, profiles, data, **kwargs): + if profiles["electron_temperature"] is not None: + data["Te_rr"] = profiles["electron_temperature"].compute(params["Te_l"], dr=2) + else: + data["Te_rr"] = jnp.nan * data["0"] + return data + + @register_compute_fun( name="ne", label="n_e", @@ -276,6 +251,27 @@ def _ne_r(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="ne_rr", + label="\\partial_{\\rho \\rho} n_e", + units="m^{-3}", + units_long="1 / cubic meters", + description="Electron density, second radial derivative", + dim=1, + params=["ne_l"], + transforms={}, + profiles=["electron_density"], + coordinates="r", + data=["0"], +) +def _ne_rr(params, transforms, profiles, data, **kwargs): + if profiles["electron_density"] is not None: + data["ne_rr"] = profiles["electron_density"].compute(params["ne_l"], dr=2) + else: + data["ne_rr"] = jnp.nan * data["0"] + return data + + @register_compute_fun( name="Ti", label="T_i", @@ -318,6 +314,27 @@ def _Ti_r(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Ti_rr", + label="\\partial_{\\rho \\rho} T_i", + units="eV", + units_long="electron-Volts", + description="Ion temperature, second radial derivative", + dim=1, + params=["Ti_l"], + transforms={}, + profiles=["ion_temperature"], + coordinates="r", + data=["0"], +) +def _Ti_rr(params, transforms, profiles, data, **kwargs): + if profiles["ion_temperature"] is not None: + data["Ti_rr"] = profiles["ion_temperature"].compute(params["Ti_l"], dr=2) + else: + data["Ti_rr"] = jnp.nan * data["0"] + return data + + @register_compute_fun( name="Zeff", label="Z_{eff}", @@ -492,43 +509,21 @@ def _gradp_mag_vol(params, transforms, profiles, data, **kwargs): units_long="None", description="Rotational transform (normalized by 2pi)", dim=1, - params=["i_l", "c_l"], + params=["i_l"], transforms={"grid": []}, profiles=["iota", "current"], coordinates="r", - data=["psi_r", "iota_zero_current_num", "iota_zero_current_den"], - axis_limit_data=["psi_rr"], + data=["iota_den", "iota_num"], + axis_limit_data=["iota_den_r", "iota_num_r"], ) def _iota(params, transforms, profiles, data, **kwargs): - # The rotational transform is computed from the toroidal current profile using - # equation 11 in S.P. Hishman & J.T. Hogan (1986) - # doi:10.1016/0021-9991(86)90197-X. Their "zero current algorithm" is supplemented - # with an additional term to account for finite net toroidal currents. Note that - # the flux surface average integrals in their formula should not be weighted by a - # coordinate Jacobian factor, meaning the sqrt(g) terms in the denominators of - # these averages will not be canceled out. if profiles["iota"] is not None: data["iota"] = profiles["iota"].compute(params["i_l"], dr=0) elif profiles["current"] is not None: - # current_term = 2*pi * I / params["Psi"]_r = mu_0 / 2*pi * current / psi_r - current_term = ( - mu_0 - / (2 * jnp.pi) - * profiles["current"].compute(params["c_l"], dr=0) - / data["psi_r"] - ) - if transforms["grid"].axis.size: - limit = ( - mu_0 - / (2 * jnp.pi) - * profiles["current"].compute(params["c_l"], dr=2) - / data["psi_rr"] - ) - current_term = put( - current_term, transforms["grid"].axis, limit[transforms["grid"].axis] - ) - data["iota"] = (current_term + data["iota_zero_current_num"]) / ( - data["iota_zero_current_den"] + # See the document attached to GitHub pull request #556 for the math. + data["iota"] = transforms["grid"].replace_at_axis( + data["iota_num"] / data["iota_den"], + lambda: data["iota_num_r"] / data["iota_den_r"], ) return data @@ -540,45 +535,30 @@ def _iota(params, transforms, profiles, data, **kwargs): units_long="None", description="Rotational transform (normalized by 2pi), first radial derivative", dim=1, - params=["i_l", "c_l"], + params=["i_l"], transforms={"grid": []}, profiles=["iota", "current"], coordinates="r", - data=[ - "iota", - "psi_r", - "psi_rr", - "iota_zero_current_num_r", - "iota_zero_current_den", - "iota_zero_current_den_r", - ], + data=["iota_den", "iota_den_r", "iota_num", "iota_num_r"], + axis_limit_data=["iota_den_rr", "iota_num_rr"], ) def _iota_r(params, transforms, profiles, data, **kwargs): - # The rotational transform is computed from the toroidal current profile using - # equation 11 in S.P. Hishman & J.T. Hogan (1986) - # doi:10.1016/0021-9991(86)90197-X. Their "zero current algorithm" is supplemented - # with an additional term to account for finite net toroidal currents. Note that - # the flux surface average integrals in their formula should not be weighted by a - # coordinate Jacobian factor, meaning the sqrt(g) terms in the denominators of - # these averages will not be canceled out. if profiles["iota"] is not None: data["iota_r"] = profiles["iota"].compute(params["i_l"], dr=1) elif profiles["current"] is not None: - current_term = ( - mu_0 - / (2 * jnp.pi) - * profiles["current"].compute(params["c_l"], dr=0) - / data["psi_r"] + # See the document attached to GitHub pull request #556 for the math. + data["iota_r"] = transforms["grid"].replace_at_axis( + ( + data["iota_num_r"] * data["iota_den"] + - data["iota_num"] * data["iota_den_r"] + ) + / data["iota_den"] ** 2, + lambda: ( + data["iota_num_rr"] * data["iota_den_r"] + - data["iota_num_r"] * data["iota_den_rr"] + ) + / (2 * data["iota_den_r"] ** 2), ) - current_term_r = ( - mu_0 / (2 * jnp.pi) * profiles["current"].compute(params["c_l"], dr=1) - - current_term * data["psi_rr"] - ) / data["psi_r"] - data["iota_r"] = ( - current_term_r - + data["iota_zero_current_num_r"] - - data["iota"] * data["iota_zero_current_den_r"] - ) / data["iota_zero_current_den"] return data @@ -589,277 +569,721 @@ def _iota_r(params, transforms, profiles, data, **kwargs): units_long="None", description="Rotational transform (normalized by 2pi), second radial derivative", dim=1, - params=["i_l", "c_l"], + params=["i_l"], transforms={"grid": []}, profiles=["iota", "current"], coordinates="r", data=[ - "iota", - "iota_r", - "psi_r", - "psi_rr", - "psi_rrr", - "iota_zero_current_num_rr", - "iota_zero_current_den", - "iota_zero_current_den_r", - "iota_zero_current_den_rr", + "iota_den", + "iota_den_r", + "iota_den_rr", + "iota_num", + "iota_num_r", + "iota_num_rr", ], + axis_limit_data=["iota_den_rrr", "iota_num_rrr"], ) def _iota_rr(params, transforms, profiles, data, **kwargs): - # The rotational transform is computed from the toroidal current profile using - # equation 11 in S.P. Hishman & J.T. Hogan (1986) - # doi:10.1016/0021-9991(86)90197-X. Their "zero current algorithm" is supplemented - # with an additional term to account for finite net toroidal currents. Note that - # the flux surface average integrals in their formula should not be weighted by a - # coordinate Jacobian factor, meaning the sqrt(g) terms in the denominators of - # these averages will not be canceled out. if profiles["iota"] is not None: data["iota_rr"] = profiles["iota"].compute(params["i_l"], dr=2) elif profiles["current"] is not None: - current_term = ( - mu_0 - / (2 * jnp.pi) - * profiles["current"].compute(params["c_l"], dr=0) - / data["psi_r"] + # See the document attached to GitHub pull request #556 for the math. + data["iota_rr"] = transforms["grid"].replace_at_axis( + ( + data["iota_num_rr"] * data["iota_den"] ** 2 + - 2 * data["iota_num_r"] * data["iota_den"] * data["iota_den_r"] + + 2 * data["iota_num"] * data["iota_den_r"] ** 2 + - data["iota_num"] * data["iota_den"] * data["iota_den_rr"] + ) + / data["iota_den"] ** 3, + lambda: ( + 2 * data["iota_num_rrr"] * data["iota_den_r"] ** 2 + - 3 * data["iota_num_rr"] * data["iota_den_r"] * data["iota_den_rr"] + + 3 * data["iota_num_r"] * data["iota_den_rr"] ** 2 + - 2 * data["iota_num_r"] * data["iota_den_r"] * data["iota_den_rrr"] + ) + / (6 * data["iota_den_r"] ** 3), ) - current_term_r = ( - mu_0 / (2 * jnp.pi) * profiles["current"].compute(params["c_l"], dr=1) - - current_term * data["psi_rr"] - ) / data["psi_r"] - current_term_rr = ( - mu_0 / (2 * jnp.pi) * profiles["current"].compute(params["c_l"], dr=2) - - 2 * current_term_r * data["psi_rr"] - - current_term * data["psi_rrr"] - ) / data["psi_r"] - data["iota_rr"] = ( - current_term_rr - + data["iota_zero_current_num_rr"] - - 2 * data["iota_r"] * data["iota_zero_current_den_r"] - - data["iota"] * data["iota_zero_current_den_rr"] - ) / data["iota_zero_current_den"] return data @register_compute_fun( - name="iota_zero_current_num", - label="\\iota_{0~\\mathrm{numerator}}", + name="iota_num", + label="\\iota_{\\mathrm{numerator}}", units="m^{-1}", units_long="inverse meters", - description="Zero toroidal current rotational transform numerator", + description="Numerator of rotational transform formula", dim=1, - params=[], + params=["c_l"], transforms={"grid": []}, - profiles=[], + profiles=["current"], coordinates="r", - data=["lambda_t", "lambda_z", "g_tt", "g_tz", "sqrt(g)"], - axis_limit_data=[ + data=["0", "lambda_z", "g_tt", "lambda_t", "g_tz", "sqrt(g)", "psi_r"], +) +def _iota_num(params, transforms, profiles, data, **kwargs): + """Numerator of rotational transform formula. + + Computes 𝛼 + 𝛽 as defined in the document attached to the description + of GitHub pull request #556. 𝛼 supplements the rotational transform with an + additional term to account for the enclosed net toroidal current. + """ + if profiles["current"] is None: + data["iota_num"] = jnp.nan * data["0"] + return data + + # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current + alpha = ( + 2 + * jnp.pi + * mu_0 + * profiles["current"].compute(params["c_l"], dr=0) + / data["psi_r"] + ) + beta = surface_integrals( + transforms["grid"], + (data["lambda_z"] * data["g_tt"] - (1 + data["lambda_t"]) * data["g_tz"]) + / data["sqrt(g)"], + ) + data["iota_num"] = transforms["grid"].replace_at_axis(alpha + beta, 0) + return data + + +@register_compute_fun( + name="iota_num_r", + label="\\partial_{\\rho} \\iota_{\\mathrm{numerator}}", + units="m^{-1}", + units_long="inverse meters", + description="Numerator of rotational transform formula, first radial derivative", + dim=1, + params=["c_l"], + transforms={"grid": []}, + profiles=["current"], + coordinates="r", + data=[ + "0", + "lambda_t", "lambda_rt", - "g_tt_rr", + "lambda_z", + "lambda_rz", + "g_tt", + "g_tt_r", + "g_tz", "g_tz_r", - "g_tz_rr", + "sqrt(g)", "sqrt(g)_r", - "sqrt(g)_rr", + "psi_r", + "psi_rr", ], + axis_limit_data=["g_tt_rr", "g_tz_rr", "sqrt(g)_rr", "psi_rrr"], ) -def _iota_zero_current_num(params, transforms, profiles, data, **kwargs): - num = ( +def _iota_num_r(params, transforms, profiles, data, **kwargs): + """Numerator of rotational transform formula, first radial derivative. + + Computes d(𝛼+𝛽)/d𝜌 as defined in the document attached to the description + of GitHub pull request #556. 𝛼 supplements the rotational transform with an + additional term to account for the enclosed net toroidal current. + """ + if profiles["current"] is None: + data["iota_num_r"] = jnp.nan * data["0"] + return data + + current_r = profiles["current"].compute(params["c_l"], dr=1) + # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current + alpha_r = ( + jnp.pi + * mu_0 + * transforms["grid"].replace_at_axis( + 2 + * ( + current_r * data["psi_r"] + - profiles["current"].compute(params["c_l"], dr=0) * data["psi_rr"] + ) + / data["psi_r"] ** 2, + lambda: ( + profiles["current"].compute(params["c_l"], dr=2) * data["psi_rr"] + - current_r * data["psi_rrr"] + ) + / data["psi_rr"] ** 2, + ) + ) + beta = ( data["lambda_z"] * data["g_tt"] - (1 + data["lambda_t"]) * data["g_tz"] ) / data["sqrt(g)"] - data["iota_zero_current_num"] = surface_averages(transforms["grid"], num) - - if transforms["grid"].axis.size: - limit = ( - (data["g_tz_r"] * data["sqrt(g)_rr"] * (1 + data["lambda_t"])) - / data["sqrt(g)_r"] ** 2 - ) + ( - data["lambda_z"] * data["g_tt_rr"] - - 2 * data["g_tz_r"] * data["lambda_rt"] - - (1 + data["lambda_t"]) * data["g_tz_rr"] - ) / data[ - "sqrt(g)_r" - ] - limit = surface_averages(transforms["grid"], limit) - data["iota_zero_current_num"] = put( - data["iota_zero_current_num"], - transforms["grid"].axis, - limit[transforms["grid"].axis], + beta_r = transforms["grid"].replace_at_axis( + ( + data["lambda_rz"] * data["g_tt"] + + data["lambda_z"] * data["g_tt_r"] + - data["lambda_rt"] * data["g_tz"] + - (1 + data["lambda_t"]) * data["g_tz_r"] + - beta * data["sqrt(g)_r"] ) - + / data["sqrt(g)"], + lambda: ( + (1 + data["lambda_t"]) + * data["g_tz_r"] + * data["sqrt(g)_rr"] + / (2 * data["sqrt(g)_r"] ** 2) + + ( + data["lambda_z"] * data["g_tt_rr"] + - 2 * data["lambda_rt"] * data["g_tz_r"] + - (1 + data["lambda_t"]) * data["g_tz_rr"] + ) + / (2 * data["sqrt(g)_r"]) + ), + ) + beta_r = surface_integrals(transforms["grid"], beta_r) + data["iota_num_r"] = alpha_r + beta_r return data @register_compute_fun( - name="iota_zero_current_num_r", - label="\\partial_{\\rho} \\iota_{0~\\mathrm{numerator}}", + name="iota_num_rr", + label="\\partial_{\\rho\\rho} \\iota_{\\mathrm{numerator}}", units="m^{-1}", units_long="inverse meters", - description="Zero toroidal current rotational transform numerator," - " first radial derivative", + description="Numerator of rotational transform formula, second radial derivative", dim=1, - params=[], + params=["c_l"], transforms={"grid": []}, - profiles=[], + profiles=["current"], coordinates="r", data=[ + "0", "lambda_t", "lambda_rt", + "lambda_rrt", "lambda_z", "lambda_rz", + "lambda_rrz", "g_tt", "g_tt_r", + "g_tt_rr", "g_tz", "g_tz_r", + "g_tz_rr", "sqrt(g)", "sqrt(g)_r", + "sqrt(g)_rr", + "psi_r", + "psi_rr", + "psi_rrr", ], + axis_limit_data=["sqrt(g)_rrr", "g_tt_rrr", "g_tz_rrr"], ) -def _iota_zero_current_num_r(params, transforms, profiles, data, **kwargs): - num = ( +def _iota_num_rr(params, transforms, profiles, data, **kwargs): + """Numerator of rotational transform formula, second radial derivative. + + Computes d2(𝛼+𝛽)/d𝜌2 as defined in the document attached to the description + of GitHub pull request #556. 𝛼 supplements the rotational transform with an + additional term to account for the enclosed net toroidal current. + """ + if profiles["current"] is None: + data["iota_num_rr"] = jnp.nan * data["0"] + return data + + current_r = profiles["current"].compute(params["c_l"], dr=1) + current_rr = profiles["current"].compute(params["c_l"], dr=2) + # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current + alpha_rr = ( + jnp.pi + * mu_0 + * transforms["grid"].replace_at_axis( + 2 * current_rr / data["psi_r"] + - 4 * current_r * data["psi_rr"] / data["psi_r"] ** 2 + + 2 + * profiles["current"].compute(params["c_l"], dr=0) + * (2 * data["psi_rr"] ** 2 - data["psi_rrr"] * data["psi_r"]) + / data["psi_r"] ** 3, + lambda: 2 + * profiles["current"].compute(params["c_l"], dr=3) + / (3 * data["psi_rr"]) + - current_rr * data["psi_rrr"] / data["psi_rr"] ** 2 + + current_r * data["psi_rrr"] ** 2 / data["psi_rr"] ** 3, + ) + ) + beta = ( data["lambda_z"] * data["g_tt"] - (1 + data["lambda_t"]) * data["g_tz"] ) / data["sqrt(g)"] - num_r = ( + beta_r = ( data["lambda_rz"] * data["g_tt"] + data["lambda_z"] * data["g_tt_r"] - data["lambda_rt"] * data["g_tz"] - (1 + data["lambda_t"]) * data["g_tz_r"] - - num * data["sqrt(g)_r"] + - beta * data["sqrt(g)_r"] ) / data["sqrt(g)"] - data["iota_zero_current_num_r"] = surface_averages(transforms["grid"], num_r) - # TODO: limit at axis + beta_rr = transforms["grid"].replace_at_axis( + ( + data["lambda_rrz"] * data["g_tt"] + + 2 * data["lambda_rz"] * data["g_tt_r"] + + data["lambda_z"] * data["g_tt_rr"] + - data["lambda_rrt"] * data["g_tz"] + - 2 * data["lambda_rt"] * data["g_tz_r"] + - (1 + data["lambda_t"]) * data["g_tz_rr"] + - 2 * beta_r * data["sqrt(g)_r"] + - beta * data["sqrt(g)_rr"] + ) + / data["sqrt(g)"], + lambda: ( + 2 + * data["sqrt(g)_r"] ** 2 + * ( + 3 * data["g_tt_rr"] * data["lambda_rz"] + + data["g_tt_rrr"] * data["lambda_z"] + - 3 * data["g_tz_rr"] * data["lambda_rt"] + - 3 * data["g_tz_r"] * data["lambda_rrt"] + - data["g_tz_rrr"] * (1 + data["lambda_t"]) + ) + + data["sqrt(g)_r"] + * ( + 3 + * data["sqrt(g)_rr"] + * ( + 2 * data["g_tz_r"] * data["lambda_rt"] + - data["g_tt_rr"] * data["lambda_t"] + + data["g_tz_rr"] * (1 + data["lambda_t"]) + ) + + 2 * data["sqrt(g)_rrr"] * data["g_tz_r"] * (1 + data["lambda_t"]) + ) + - 3 * data["sqrt(g)_rr"] ** 2 * data["g_tz_r"] * (1 + data["lambda_t"]) + ) + / (6 * data["sqrt(g)_r"] ** 3), + ) + beta_rr = surface_integrals(transforms["grid"], beta_rr) + data["iota_num_rr"] = alpha_rr + beta_rr return data @register_compute_fun( - name="iota_zero_current_num_rr", - label="\\partial_{\\rho\\rho} \\iota_{0~\\mathrm{numerator}}", + name="iota_num_rrr", + label="\\partial_{\\rho\\rho\\rho} \\iota_{\\mathrm{numerator}}", units="m^{-1}", units_long="inverse meters", - description="Zero toroidal current rotational transform numerator," - " second radial derivative", + description="Numerator of rotational transform formula, third radial derivative", dim=1, - params=[], + params=["c_l"], transforms={"grid": []}, - profiles=[], + profiles=["current"], coordinates="r", data=[ + "0", "lambda_t", "lambda_rt", "lambda_rrt", + "lambda_rrrt", "lambda_z", "lambda_rz", "lambda_rrz", + "lambda_rrrz", "g_tt", "g_tt_r", "g_tt_rr", + "g_tt_rrr", "g_tz", "g_tz_r", "g_tz_rr", + "g_tz_rrr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_rr", + "sqrt(g)_rrr", + "psi_r", + "psi_rr", + "psi_rrr", ], ) -def _iota_zero_current_num_rr(params, transforms, profiles, data, **kwargs): - num = ( +def _iota_num_rrr(params, transforms, profiles, data, **kwargs): + """Numerator of rotational transform formula, third radial derivative. + + Computes d3(𝛼+𝛽)/d𝜌3 as defined in the document attached to the description + of GitHub pull request #556. 𝛼 supplements the rotational transform with an + additional term to account for the enclosed net toroidal current. + """ + if profiles["current"] is None: + data["iota_num_rrr"] = jnp.nan * data["0"] + return data + + current_r = profiles["current"].compute(params["c_l"], dr=1) + current_rr = profiles["current"].compute(params["c_l"], dr=2) + current_rrr = profiles["current"].compute(params["c_l"], dr=3) + # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current + alpha_rrr = ( + jnp.pi + * mu_0 + * transforms["grid"].replace_at_axis( + 2 * current_rrr / data["psi_r"] + - 6 * current_rr * data["psi_rr"] / data["psi_r"] ** 2 + + 6 + * current_r + * ( + 2 * data["psi_r"] * data["psi_rr"] ** 2 + - data["psi_rrr"] * data["psi_r"] ** 2 + ) + / data["psi_r"] ** 4 + + 12 + * profiles["current"].compute(params["c_l"], dr=0) + * (data["psi_rrr"] * data["psi_rr"] * data["psi_r"] - data["psi_rr"] ** 3) + / data["psi_r"] ** 4, + lambda: profiles["current"].compute(params["c_l"], dr=4) + / (2 * data["psi_rr"]) + - current_rrr * data["psi_rrr"] / data["psi_rr"] ** 2 + + 3 * current_rr * data["psi_rrr"] ** 2 / (2 * data["psi_rr"] ** 3) + - 3 * current_r * data["psi_rrr"] ** 3 / (2 * data["psi_rr"] ** 4), + ) + ) + beta = ( data["lambda_z"] * data["g_tt"] - (1 + data["lambda_t"]) * data["g_tz"] ) / data["sqrt(g)"] - num_r = ( + beta_r = ( data["lambda_rz"] * data["g_tt"] + data["lambda_z"] * data["g_tt_r"] - data["lambda_rt"] * data["g_tz"] - (1 + data["lambda_t"]) * data["g_tz_r"] - - num * data["sqrt(g)_r"] + - beta * data["sqrt(g)_r"] ) / data["sqrt(g)"] - num_rr = ( + beta_rr = ( data["lambda_rrz"] * data["g_tt"] + 2 * data["lambda_rz"] * data["g_tt_r"] + data["lambda_z"] * data["g_tt_rr"] - data["lambda_rrt"] * data["g_tz"] - 2 * data["lambda_rt"] * data["g_tz_r"] - (1 + data["lambda_t"]) * data["g_tz_rr"] - - 2 * num_r * data["sqrt(g)_r"] - - num * data["sqrt(g)_rr"] + - 2 * beta_r * data["sqrt(g)_r"] + - beta * data["sqrt(g)_rr"] ) / data["sqrt(g)"] - data["iota_zero_current_num_rr"] = surface_averages(transforms["grid"], num_rr) - # TODO: limit at axis + beta_rrr = transforms["grid"].replace_at_axis( + ( + data["lambda_rrrz"] * data["g_tt"] + + 3 * data["lambda_rrz"] * data["g_tt_r"] + + 3 * data["lambda_rz"] * data["g_tt_rr"] + + data["lambda_z"] * data["g_tt_rrr"] + - data["lambda_rrrt"] * data["g_tz"] + - 3 * data["lambda_rrt"] * data["g_tz_r"] + - 3 * data["lambda_rt"] * data["g_tz_rr"] + - (1 + data["lambda_t"]) * data["g_tz_rrr"] + - 3 * beta_rr * data["sqrt(g)_r"] + - 3 * beta_r * data["sqrt(g)_rr"] + - beta * data["sqrt(g)_rrr"] + ) + / data["sqrt(g)"], + # Todo: axis limit of beta_rrr + # Computed with four applications of l’Hôpital’s rule. + # Requires sqrt(g)_rrrr and fourth derivatives of basis vectors. + jnp.nan, + ) + beta_rrr = surface_integrals(transforms["grid"], beta_rrr) + # force limit to nan until completed because integration replaces nan with 0 + data["iota_num_rrr"] = alpha_rrr + transforms["grid"].replace_at_axis( + beta_rrr, jnp.nan + ) return data @register_compute_fun( - name="iota_zero_current_den", - label="\\iota_{0~\\mathrm{denominator}}", + name="iota_den", + label="\\iota_{\\mathrm{denominator}}", units="m^{-1}", units_long="inverse meters", - description="Zero toroidal current rotational transform denominator", + description="Denominator of rotational transform formula", dim=1, params=[], transforms={"grid": []}, - profiles=[], + profiles=["current"], coordinates="r", - data=["g_tt", "sqrt(g)"], - axis_limit_data=["g_tt_rr", "sqrt(g)_r"], + data=["0", "g_tt", "g_tz", "sqrt(g)", "omega_t", "omega_z"], ) -def _iota_zero_current_den(params, transforms, profiles, data, **kwargs): - den = data["g_tt"] / data["sqrt(g)"] - data["iota_zero_current_den"] = surface_averages(transforms["grid"], den) +def _iota_den(params, transforms, profiles, data, **kwargs): + """Denominator of rotational transform formula. + + Computes 𝛾 as defined in the document attached to the description + of GitHub pull request #556. + """ + if profiles["current"] is None: + data["iota_den"] = jnp.nan * data["0"] + return data + + gamma = ( + (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"] + ) / data["sqrt(g)"] + # Assumes toroidal stream function behaves such that the magnetic axis limit + # of gamma is zero (as it would if omega = 0 identically). + gamma = transforms["grid"].replace_at_axis( + surface_integrals(transforms["grid"], gamma), 0 + ) + data["iota_den"] = gamma + return data - if transforms["grid"].axis.size: - limit = surface_averages( - transforms["grid"], data["g_tt_rr"] / data["sqrt(g)_r"] - ) - data["iota_zero_current_den"] = put( - data["iota_zero_current_den"], - transforms["grid"].axis, - limit[transforms["grid"].axis], - ) +@register_compute_fun( + name="iota_den_r", + label="\\partial_{\\rho} \\iota_{\\mathrm{denominator}}", + units="m^{-1}", + units_long="inverse meters", + description="Denominator of rotational transform formula, first radial derivative", + dim=1, + params=[], + transforms={"grid": []}, + profiles=["current"], + coordinates="r", + data=[ + "0", + "g_tt", + "g_tt_r", + "g_tz", + "g_tz_r", + "sqrt(g)", + "sqrt(g)_r", + "omega_t", + "omega_rt", + "omega_z", + "omega_rz", + ], + axis_limit_data=["sqrt(g)_rr", "g_tt_rr", "g_tz_rr"], +) +def _iota_den_r(params, transforms, profiles, data, **kwargs): + """Denominator of rotational transform formula, first radial derivative. + + Computes d𝛾/d𝜌 as defined in the document attached to the description + of GitHub pull request #556. + """ + if profiles["current"] is None: + data["iota_den_r"] = jnp.nan * data["0"] + return data + + gamma = ( + (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"] + ) / data["sqrt(g)"] + gamma_r = transforms["grid"].replace_at_axis( + ( + data["omega_rz"] * data["g_tt"] + + (1 + data["omega_z"]) * data["g_tt_r"] + - data["omega_rt"] * data["g_tz"] + - data["omega_t"] * data["g_tz_r"] + - gamma * data["sqrt(g)_r"] + ) + / data["sqrt(g)"], + lambda: ( + data["omega_t"] + * data["g_tz_r"] + * data["sqrt(g)_rr"] + / (2 * data["sqrt(g)_r"] ** 2) + + ( + (1 + data["omega_z"]) * data["g_tt_rr"] + - 2 * data["omega_rt"] * data["g_tz_r"] + - data["omega_t"] * data["g_tz_rr"] + ) + / (2 * data["sqrt(g)_r"]) + ), + ) + gamma_r = surface_integrals(transforms["grid"], gamma_r) + data["iota_den_r"] = gamma_r return data @register_compute_fun( - name="iota_zero_current_den_r", - label="\\partial_{\\rho} \\iota_{0~\\mathrm{denominator}}", + name="iota_den_rr", + label="\\partial_{\\rho\\rho} \\iota_{\\mathrm{denominator}}", units="m^{-1}", units_long="inverse meters", - description="Zero toroidal current rotational transform denominator," - " first radial derivative", + description="Denominator of rotational transform formula, second radial derivative", dim=1, params=[], transforms={"grid": []}, - profiles=[], + profiles=["current"], coordinates="r", - data=["g_tt", "g_tt_r", "sqrt(g)", "sqrt(g)_r"], + data=[ + "0", + "g_tt", + "g_tt_r", + "g_tt_rr", + "g_tz", + "g_tz_r", + "g_tz_rr", + "sqrt(g)", + "sqrt(g)_r", + "sqrt(g)_rr", + "omega_t", + "omega_rt", + "omega_rrt", + "omega_z", + "omega_rz", + "omega_rrz", + ], + axis_limit_data=["sqrt(g)_rrr", "g_tt_rrr", "g_tz_rrr"], ) -def _iota_zero_current_den_r(params, transforms, profiles, data, **kwargs): - den = data["g_tt"] / data["sqrt(g)"] - den_r = (data["g_tt_r"] - den * data["sqrt(g)_r"]) / data["sqrt(g)"] - data["iota_zero_current_den_r"] = surface_averages(transforms["grid"], den_r) - # TODO: limit at axis +def _iota_den_rr(params, transforms, profiles, data, **kwargs): + """Denominator of rotational transform formula, second radial derivative. + + Computes d2𝛾/d𝜌2 as defined in the document attached to the description + of GitHub pull request #556. + """ + if profiles["current"] is None: + data["iota_den_rr"] = jnp.nan * data["0"] + return data + + gamma = ( + (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"] + ) / data["sqrt(g)"] + gamma_r = ( + data["omega_rz"] * data["g_tt"] + + (1 + data["omega_z"]) * data["g_tt_r"] + - data["omega_rt"] * data["g_tz"] + - data["omega_t"] * data["g_tz_r"] + - gamma * data["sqrt(g)_r"] + ) / data["sqrt(g)"] + gamma_rr = transforms["grid"].replace_at_axis( + ( + data["omega_rrz"] * data["g_tt"] + + 2 * data["omega_rz"] * data["g_tt_r"] + + (1 + data["omega_z"]) * data["g_tt_rr"] + - data["omega_rrt"] * data["g_tz"] + - 2 * data["omega_rt"] * data["g_tz_r"] + - data["omega_t"] * data["g_tz_rr"] + - 2 * gamma_r * data["sqrt(g)_r"] + - gamma * data["sqrt(g)_rr"] + ) + / data["sqrt(g)"], + lambda: ( + 2 + * data["sqrt(g)_r"] ** 2 + * ( + 3 * data["g_tt_rr"] * data["omega_rz"] + + data["g_tt_rrr"] * (1 + data["omega_z"]) + - 3 * data["g_tz_rr"] * data["omega_rt"] + - 3 * data["g_tz_r"] * data["omega_rrt"] + - data["g_tz_rrr"] * data["omega_t"] + ) + + data["sqrt(g)_r"] + * ( + 3 + * data["sqrt(g)_rr"] + * ( + 2 * data["g_tz_r"] * data["omega_rt"] + - data["g_tt_rr"] * (1 + data["omega_z"]) + + data["g_tz_rr"] * data["omega_t"] + ) + + 2 * data["sqrt(g)_rrr"] * data["g_tz_r"] * data["omega_t"] + ) + - 3 * data["sqrt(g)_rr"] ** 2 * data["g_tz_r"] * data["omega_t"] + ) + / (6 * data["sqrt(g)_r"] ** 3), + ) + gamma_rr = surface_integrals(transforms["grid"], gamma_rr) + data["iota_den_rr"] = gamma_rr return data @register_compute_fun( - name="iota_zero_current_den_rr", - label="\\partial_{\\rho\\rho} \\iota_{0~\\mathrm{denominator}}", + name="iota_den_rrr", + label="\\partial_{\\rho\\rho\\rho} \\iota_{\\mathrm{denominator}}", units="m^{-1}", units_long="inverse meters", - description="Zero toroidal current rotational transform denominator," - " second radial derivative", + description="Denominator of rotational transform formula, third radial derivative", dim=1, params=[], transforms={"grid": []}, - profiles=[], + profiles=["current"], coordinates="r", - data=["g_tt", "g_tt_r", "g_tt_rr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_rr"], + data=[ + "0", + "g_tt", + "g_tt_r", + "g_tt_rr", + "g_tt_rrr", + "g_tz", + "g_tz_r", + "g_tz_rr", + "g_tz_rrr", + "sqrt(g)", + "sqrt(g)_r", + "sqrt(g)_rr", + "sqrt(g)_rrr", + "omega_t", + "omega_rt", + "omega_rrt", + "omega_rrrt", + "omega_z", + "omega_rz", + "omega_rrz", + "omega_rrrz", + ], ) -def _iota_zero_current_den_rr(params, transforms, profiles, data, **kwargs): - den = data["g_tt"] / data["sqrt(g)"] - den_r = (data["g_tt_r"] - den * data["sqrt(g)_r"]) / data["sqrt(g)"] - den_rr = ( - data["g_tt_rr"] - 2 * den_r * data["sqrt(g)_r"] - den * data["sqrt(g)_rr"] +def _iota_den_rrr(params, transforms, profiles, data, **kwargs): + """Denominator of rotational transform formula, third radial derivative. + + Computes d3𝛾/d𝜌3 as defined in the document attached to the description + of GitHub pull request #556. + """ + if profiles["current"] is None: + data["iota_den_rrr"] = jnp.nan * data["0"] + return data + + gamma = ( + (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"] + ) / data["sqrt(g)"] + gamma_r = ( + data["omega_rz"] * data["g_tt"] + + (1 + data["omega_z"]) * data["g_tt_r"] + - data["omega_rt"] * data["g_tz"] + - data["omega_t"] * data["g_tz_r"] + - gamma * data["sqrt(g)_r"] ) / data["sqrt(g)"] - data["iota_zero_current_den_rr"] = surface_averages(transforms["grid"], den_rr) - # TODO: limit at axis + gamma_rr = ( + data["omega_rrz"] * data["g_tt"] + + 2 * data["omega_rz"] * data["g_tt_r"] + + (1 + data["omega_z"]) * data["g_tt_rr"] + - data["omega_rrt"] * data["g_tz"] + - 2 * data["omega_rt"] * data["g_tz_r"] + - data["omega_t"] * data["g_tz_rr"] + - 2 * gamma_r * data["sqrt(g)_r"] + - gamma * data["sqrt(g)_rr"] + ) / data["sqrt(g)"] + gamma_rrr = transforms["grid"].replace_at_axis( + ( + data["omega_rrrz"] * data["g_tt"] + + 3 * data["omega_rrz"] * data["g_tt_r"] + + 3 * data["omega_rz"] * data["g_tt_rr"] + + (1 + data["omega_z"]) * data["g_tt_rrr"] + - data["omega_rrrt"] * data["g_tz"] + - 3 * data["omega_rrt"] * data["g_tz_r"] + - 3 * data["omega_rt"] * data["g_tz_rr"] + - data["omega_t"] * data["g_tz_rrr"] + - 3 * gamma_rr * data["sqrt(g)_r"] + - 3 * gamma_r * data["sqrt(g)_rr"] + - gamma * data["sqrt(g)_rrr"] + ) + / data["sqrt(g)"], + # Todo: axis limit + # Computed with four applications of l’Hôpital’s rule. + # Requires sqrt(g)_rrrr and fourth derivatives of basis vectors. + jnp.nan, + ) + gamma_rrr = surface_integrals(transforms["grid"], gamma_rrr) + # force limit to nan until completed because integration replaces nan with 0 + data["iota_den_rrr"] = transforms["grid"].replace_at_axis(gamma_rrr, jnp.nan) + return data + + +@register_compute_fun( + name="iota_psi", + label="\\partial_{\\psi} \\iota", + units="Wb^{-1}", + units_long="Inverse Webers", + description="Rotational transform, radial derivative wrt toroidal flux", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["iota_r", "psi_r"], + axis_limit_data=["iota_rr", "psi_rr"], +) +def _iota_psi(params, transforms, profiles, data, **kwargs): + # Existence of limit at magnetic axis requires ∂ᵨ iota = 0 at axis. + # Assume iota may be expanded as an even power series of ρ so that this + # condition is satisfied. + data["iota_psi"] = transforms["grid"].replace_at_axis( + data["iota_r"] / data["psi_r"], lambda: data["iota_rr"] / data["psi_rr"] + ) return data diff --git a/desc/compute/_qs.py b/desc/compute/_qs.py index d8c32432c1..a082945bb5 100644 --- a/desc/compute/_qs.py +++ b/desc/compute/_qs.py @@ -1,4 +1,13 @@ -"""Compute functions for quasisymmetry objectives.""" +"""Compute functions for quasisymmetry objectives. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" import numpy as np @@ -70,9 +79,9 @@ def _w_mn(params, transforms, profiles, data, **kwargs): # indices of matching modes in w and B bases # need to use np instead of jnp here as jnp.where doesn't work under jit # even if the args are static - ib, iw = np.where((Bm[:, None] == -wm) * (Bn[:, None] == wn) * (wm != 0)) + ib, iw = np.where((Bm[:, None] == -wm) & (Bn[:, None] == wn) & (wm != 0)) jb, jw = np.where( - (Bm[:, None] == wm) * (Bn[:, None] == -wn) * (wm == 0) * (wn != 0) + (Bm[:, None] == wm) & (Bn[:, None] == -wn) & (wm == 0) & (wn != 0) ) w_mn = put(w_mn, iw, sign(wn[iw]) * data["B_theta_mn"][ib] / jnp.abs(wm[iw])) w_mn = put(w_mn, jw, sign(wm[jw]) * data["B_zeta_mn"][jb] / jnp.abs(NFP * wn[jw])) @@ -268,9 +277,7 @@ def _B_mn(params, transforms, profiles, data, **kwargs): norm = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1)) data["|B|_mn"] = ( norm # 1 if m=n=0, 2 if m=0 or n=0, 4 if m!=0 and n!=0 - * jnp.matmul( - transforms["B"].basis.evaluate(nodes).T, data["sqrt(g)_B"] * data["|B|"] - ) + * (transforms["B"].basis.evaluate(nodes).T @ (data["sqrt(g)_B"] * data["|B|"])) / transforms["B"].grid.num_nodes ) return data @@ -306,23 +313,12 @@ def _B_modes(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=[ - "iota", - "psi_r", - "sqrt(g)", - "B_theta", - "B_zeta", - "|B|_t", - "|B|_z", - "G", - "I", - "B*grad(|B|)", - ], + data=["iota", "B0", "B_theta", "B_zeta", "|B|_t", "|B|_z", "G", "I", "B*grad(|B|)"], helicity="helicity", ) def _f_C(params, transforms, profiles, data, **kwargs): M, N = kwargs.get("helicity", (1, 0)) - data["f_C"] = (M * data["iota"] - N) * (data["psi_r"] / data["sqrt(g)"]) * ( + data["f_C"] = (M * data["iota"] - N) * data["B0"] * ( data["B_zeta"] * data["|B|_t"] - data["B_theta"] * data["|B|_z"] ) - (M * data["G"] + N * data["I"]) * data["B*grad(|B|)"] return data @@ -340,17 +336,10 @@ def _f_C(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=[ - "psi_r", - "sqrt(g)", - "|B|_t", - "|B|_z", - "(B*grad(|B|))_t", - "(B*grad(|B|))_z", - ], + data=["B0", "|B|_t", "|B|_z", "(B*grad(|B|))_t", "(B*grad(|B|))_z"], ) def _f_T(params, transforms, profiles, data, **kwargs): - data["f_T"] = (data["psi_r"] / data["sqrt(g)"]) * ( + data["f_T"] = data["B0"] * ( data["|B|_t"] * data["(B*grad(|B|))_z"] - data["|B|_z"] * data["(B*grad(|B|))_t"] ) @@ -359,7 +348,7 @@ def _f_T(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="isodynamicity", - label="1/B^2 (\\mathbf{b} \\times \\nabla B) \\cdot \\nabla \\psi", + label="1/|B|^2 (\\mathbf{b} \\times \\nabla B) \\cdot \\nabla \\psi", units="~", units_long="None", description="Measure of cross field drift at each point, " diff --git a/desc/compute/_stability.py b/desc/compute/_stability.py index d90df7c449..ea93734248 100644 --- a/desc/compute/_stability.py +++ b/desc/compute/_stability.py @@ -1,11 +1,20 @@ -"""Compute functions for Mercier stability objectives.""" +"""Compute functions for Mercier stability objectives. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" from scipy.constants import mu_0 from desc.backend import jnp from .data_index import register_compute_fun -from .utils import dot, surface_integrals +from .utils import dot, surface_integrals_map @register_compute_fun( @@ -19,12 +28,12 @@ transforms={}, profiles=[], coordinates="r", - data=["iota_r", "psi_r"], + data=["iota_psi"], ) def _D_shear(params, transforms, profiles, data, **kwargs): - # Implements equations 4.16 in M. Landreman & R. Jorge (2020) + # Implements equation 4.16 in M. Landreman & R. Jorge (2020) # doi:10.1017/S002237782000121X. - data["D_shear"] = (data["iota_r"] / (4 * jnp.pi * data["psi_r"])) ** 2 + data["D_shear"] = data["iota_psi"] ** 2 / (16 * jnp.pi**2) return data @@ -39,20 +48,34 @@ def _D_shear(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=["psi_r", "iota_r", "B", "J", "G", "I_r", "|grad(psi)|", "|e_theta x e_zeta|"], + data=[ + "psi_r", + "iota_psi", + "B", + "J", + "G", + "I_r", + "|grad(psi)|", + "|e_theta x e_zeta|", + ], ) def _D_current(params, transforms, profiles, data, **kwargs): - # Implements equations 4.17 in M. Landreman & R. Jorge (2020) + # Implements equation 4.17 in M. Landreman & R. Jorge (2020) # doi:10.1017/S002237782000121X. - Xi = mu_0 * data["J"] - jnp.atleast_2d(data["I_r"] / data["psi_r"]).T * data["B"] + Xi = mu_0 * data["J"] - (data["I_r"] / data["psi_r"] * data["B"].T).T + integrate = surface_integrals_map(transforms["grid"]) data["D_current"] = ( -jnp.sign(data["G"]) / (2 * jnp.pi) ** 4 - * data["iota_r"] - / data["psi_r"] - * surface_integrals( - transforms["grid"], - data["|e_theta x e_zeta|"] / data["|grad(psi)|"] ** 3 * dot(Xi, data["B"]), + * data["iota_psi"] + * transforms["grid"].replace_at_axis( + integrate( + data["|e_theta x e_zeta|"] + / data["|grad(psi)|"] ** 3 + * dot(Xi, data["B"]) + ), + # Todo: implement equivalent of equation 4.3 in desc coordinates + jnp.nan, ) ) return data @@ -82,28 +105,29 @@ def _D_current(params, transforms, profiles, data, **kwargs): ], ) def _D_well(params, transforms, profiles, data, **kwargs): - # Implements equations 4.18 in M. Landreman & R. Jorge (2020) + # Implements equation 4.18 in M. Landreman & R. Jorge (2020) # doi:10.1017/S002237782000121X. + integrate = surface_integrals_map(transforms["grid"]) dp_dpsi = mu_0 * data["p_r"] / data["psi_r"] d2V_dpsi2 = ( - data["V_rr(r)"] - data["V_r(r)"] * data["psi_rr"] / data["psi_r"] - ) / data["psi_r"] ** 2 + data["V_rr(r)"] * data["psi_r"] - data["V_r(r)"] * data["psi_rr"] + ) / data["psi_r"] ** 3 data["D_well"] = ( dp_dpsi * ( jnp.sign(data["psi"]) * d2V_dpsi2 - dp_dpsi - * surface_integrals( - transforms["grid"], - data["|e_theta x e_zeta|"] / (data["|B|^2"] * data["|grad(psi)|"]), + * integrate( + data["|e_theta x e_zeta|"] / (data["|B|^2"] * data["|grad(psi)|"]) ) ) - * surface_integrals( - transforms["grid"], - data["|e_theta x e_zeta|"] * data["|B|^2"] / data["|grad(psi)|"] ** 3, + * integrate( + data["|e_theta x e_zeta|"] * data["|B|^2"] / data["|grad(psi)|"] ** 3 ) / (2 * jnp.pi) ** 6 ) + # Axis limit does not exist as ∂ᵨ ψ and ‖∇ ψ‖ terms dominate so that D_well + # is of the order ρ⁻² near axis. return data @@ -121,26 +145,33 @@ def _D_well(params, transforms, profiles, data, **kwargs): data=["|grad(psi)|", "J*B", "|B|^2", "|e_theta x e_zeta|"], ) def _D_geodesic(params, transforms, profiles, data, **kwargs): - # Implements equations 4.19 in M. Landreman & R. Jorge (2020) + # Implements equation 4.19 in M. Landreman & R. Jorge (2020) # doi:10.1017/S002237782000121X. - data["D_geodesic"] = ( - surface_integrals( - transforms["grid"], - data["|e_theta x e_zeta|"] * mu_0 * data["J*B"] / data["|grad(psi)|"] ** 3, - ) - ** 2 - - surface_integrals( - transforms["grid"], - data["|e_theta x e_zeta|"] * data["|B|^2"] / data["|grad(psi)|"] ** 3, - ) - * surface_integrals( - transforms["grid"], - data["|e_theta x e_zeta|"] - * mu_0**2 - * data["J*B"] ** 2 - / (data["|B|^2"] * data["|grad(psi)|"] ** 3), + integrate = surface_integrals_map(transforms["grid"]) + data["D_geodesic"] = transforms["grid"].replace_at_axis( + ( + integrate( + data["|e_theta x e_zeta|"] + * mu_0 + * data["J*B"] + / data["|grad(psi)|"] ** 3 + ) + ** 2 + - integrate( + data["|e_theta x e_zeta|"] * data["|B|^2"] / data["|grad(psi)|"] ** 3 + ) + * integrate( + data["|e_theta x e_zeta|"] + * mu_0**2 + * data["J*B"] ** 2 + / (data["|B|^2"] * data["|grad(psi)|"] ** 3), + ) ) - ) / (2 * jnp.pi) ** 6 + / (2 * jnp.pi) ** 6, + jnp.nan, # enforce manually because our integration replaces nan with 0 + ) + # Axis limit does not exist as ‖∇ ψ‖ terms dominate so that D_geodesic + # is of the order ρ⁻² near axis. return data @@ -159,11 +190,12 @@ def _D_geodesic(params, transforms, profiles, data, **kwargs): data=["D_shear", "D_current", "D_well", "D_geodesic"], ) def _D_Mercier(params, transforms, profiles, data, **kwargs): - # Implements equations 4.20 in M. Landreman & R. Jorge (2020) + # Implements equation 4.20 in M. Landreman & R. Jorge (2020) # doi:10.1017/S002237782000121X. data["D_Mercier"] = ( data["D_shear"] + data["D_current"] + data["D_well"] + data["D_geodesic"] ) + # The axis limit does not exist as D_Mercier is of the order ρ⁻² near axis. return data @@ -179,24 +211,20 @@ def _D_Mercier(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="r", - data=[ - "V(r)", - "V_r(r)", - "p_r", - "", - "_r", - ], + data=["V(r)", "V_r(r)", "p_r", "<|B|^2>", "<|B|^2>_r"], ) def _magnetic_well(params, transforms, profiles, data, **kwargs): # Implements equation 3.2 in M. Landreman & R. Jorge (2020) # doi:10.1017/S002237782000121X. - # pressure = thermal + magnetic = 2 mu_0 p + B^2 + # pressure = thermal + magnetic = 2 mu_0 p + |B|^2 # The surface average operation is an additive homomorphism. # Thermal pressure is constant over a rho surface. # surface average(pressure) = thermal + surface average(magnetic) - data["magnetic well"] = ( + # The sign of sqrt(g) is enforced to be non-negative. + data["magnetic well"] = transforms["grid"].replace_at_axis( data["V(r)"] - * (2 * mu_0 * data["p_r"] + data["_r"]) - / (data["V_r(r)"] * data[""]) + * (2 * mu_0 * data["p_r"] + data["<|B|^2>_r"]) + / (data["V_r(r)"] * data["<|B|^2>"]), + 0, # coefficient of limit is V_r / V_rr, rest is finite ) return data diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index c6b3981e72..7c45d7df1c 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -141,6 +141,30 @@ def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar return data +@register_compute_fun( + name="e_rho_rr", + label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector," + " second derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_rho_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_rr"] = coords + return data + + @register_compute_fun( name="e_rho_t", label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", @@ -210,6 +234,30 @@ def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kw return data +@register_compute_fun( + name="e_theta_rr", + label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector," + " second derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_theta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_theta_rr"] = coords + return data + + @register_compute_fun( name="e_theta_t", label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", @@ -293,6 +341,30 @@ def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa return data +@register_compute_fun( + name="e_zeta_rr", + label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector," + " second derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_zeta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_rr"] = coords + return data + + @register_compute_fun( name="e_zeta_t", label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", @@ -498,6 +570,37 @@ def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar return data +@register_compute_fun( + name="e_rho_rr", + label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector," + " second derivative wrt radial coordinate", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[3, 0, 0]], + "Z": [[3, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_rho_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=3) + Z = transforms["Z"].transform(params["Z_lmn"], dr=3) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["e_rho_rr"] = coords + return data + + @register_compute_fun( name="e_rho_t", label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", @@ -581,6 +684,37 @@ def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw return data +@register_compute_fun( + name="e_theta_rr", + label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector," + " second derivative wrt radial coordinate", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[2, 1, 0]], + "Z": [[2, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_theta_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=2, dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=2, dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["e_theta_rr"] = coords + return data + + @register_compute_fun( name="e_theta_t", label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", @@ -657,6 +791,30 @@ def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwa return data +@register_compute_fun( + name="e_zeta_rr", + label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector," + " second derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_zeta_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_rr"] = coords + return data + + @register_compute_fun( name="e_zeta_t", label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index 9670cd535a..59a9025e1d 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -46,7 +46,7 @@ def register_compute_fun( a flux function, etc. data : list of str Names of other items in the data index needed to compute qty. - parameterization: str + parameterization: str or list of str Name of desc types the method is valid for. eg 'desc.geometry.FourierXYZCurve' or `desc.equilibrium.Equilibrium`. axis_limit_data : list of str @@ -84,6 +84,10 @@ def _decorator(func): flag = False for base_class, superclasses in _class_inheritance.items(): if p in superclasses or p == base_class: + if name in data_index[base_class]: + raise ValueError( + f"Already registered function with parameterization {p} and name {name}." + ) data_index[base_class][name] = d.copy() flag = True if not flag: diff --git a/desc/compute/utils.py b/desc/compute/utils.py index 4957a0d147..4ce3dbdae6 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -66,7 +66,7 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** names : str or array-like of str Name(s) of the quantity(s) to compute. params : dict of ndarray - Parameters from the equilibrium, such as R_lmn, Z_lmn, i_l, p_l, etc + Parameters from the equilibrium, such as R_lmn, Z_lmn, i_l, p_l, etc. Defaults to attributes of self. transforms : dict of Transform Transforms for R, Z, lambda, etc. Default is to build from grid @@ -504,120 +504,6 @@ def cross(a, b, axis=-1): return jnp.cross(a, b, axis=axis) -def _get_grid_surface(grid, surface_label): - """Return grid quantities associated with the given surface label. - - Parameters - ---------- - grid : Grid - Collocation grid containing the nodes to evaluate at. - surface_label : str - The surface label of rho, theta, or zeta. - - Returns - ------- - unique_size : ndarray - The number of the unique values of the surface_label. - inverse_idx : ndarray - Indexing array to go from unique values to full grid. - spacing : ndarray - The relevant columns of grid.spacing. - has_endpoint_dupe : bool - Whether this surface label's nodes have a duplicate at the endpoint - of a periodic domain. (e.g. a node at 0 and 2π). - - """ - assert surface_label in {"rho", "theta", "zeta"} - if surface_label == "rho": - unique_size = grid.num_rho - inverse_idx = grid.inverse_rho_idx - spacing = grid.spacing[:, 1:] - has_endpoint_dupe = False - elif surface_label == "theta": - unique_size = grid.num_theta - inverse_idx = grid.inverse_theta_idx - spacing = grid.spacing[:, [0, 2]] - has_endpoint_dupe = (grid.nodes[grid.unique_theta_idx[0], 1] == 0) & ( - grid.nodes[grid.unique_theta_idx[-1], 1] == 2 * np.pi - ) - else: - unique_size = grid.num_zeta - inverse_idx = grid.inverse_zeta_idx - spacing = grid.spacing[:, :2] - has_endpoint_dupe = (grid.nodes[grid.unique_zeta_idx[0], 2] == 0) & ( - grid.nodes[grid.unique_zeta_idx[-1], 2] == 2 * np.pi / grid.NFP - ) - return unique_size, inverse_idx, spacing, has_endpoint_dupe - - -def compress(grid, x, surface_label="rho"): - """Compress x by returning only the elements at unique surface_label indices. - - Parameters - ---------- - grid : Grid - Collocation grid containing the nodes to evaluate at. - x : ndarray - The array to compress. - Should usually represent a surface function (a function constant over a surface) - in an array that matches the grid's pattern. - surface_label : str - The surface label of rho, theta, or zeta. - - Returns - ------- - compress_x : ndarray - x[grid.unique_surface_label_indices] - This array will be sorted such that the - first element corresponds to the value associated with the smallest surface - last element corresponds to the value associated with the largest surface - - """ - assert surface_label in {"rho", "theta", "zeta"} - assert len(x) == grid.num_nodes - if surface_label == "rho": - return x[grid.unique_rho_idx] - if surface_label == "theta": - return x[grid.unique_theta_idx] - if surface_label == "zeta": - return x[grid.unique_zeta_idx] - - -def expand(grid, x, surface_label="rho"): - """Expand x by duplicating elements to match the grid's pattern. - - Parameters - ---------- - grid : Grid - Collocation grid containing the nodes to evaluate at. - x : ndarray - Stores the values of a surface function (a function constant over a surface) - for all unique surfaces of the specified label on the grid. - - len(x) should be grid.num_surface_label - - x should be sorted such that the - first element corresponds to the value associated with the smallest surface - last element corresponds to the value associated with the largest surface - surface_label : str - The surface label of rho, theta, or zeta. - - Returns - ------- - expand_x : ndarray - X expanded to match the grid's pattern. - - """ - assert surface_label in {"rho", "theta", "zeta"} - if surface_label == "rho": - assert len(x) == grid.num_rho - return x[grid.inverse_rho_idx] - if surface_label == "theta": - assert len(x) == grid.num_theta - return x[grid.inverse_theta_idx] - if surface_label == "zeta": - assert len(x) == grid.num_zeta - return x[grid.inverse_zeta_idx] - - def cumtrapz(y, x=None, dx=1.0, axis=-1, initial=None): """Cumulatively integrate y(x) using the composite trapezoidal rule. @@ -694,6 +580,52 @@ def tupleset(t, i, value): return res +def _get_grid_surface(grid, surface_label): + """Return grid quantities associated with the given surface label. + + Parameters + ---------- + grid : Grid + Collocation grid containing the nodes to evaluate at. + surface_label : str + The surface label of rho, theta, or zeta. + + Returns + ------- + unique_size : int + The number of the unique values of the surface_label. + inverse_idx : ndarray + Indexing array to go from unique values to full grid. + spacing : ndarray + The relevant columns of grid.spacing. + has_endpoint_dupe : bool + Whether this surface label's nodes have a duplicate at the endpoint + of a periodic domain. (e.g. a node at 0 and 2π). + + """ + assert surface_label in {"rho", "theta", "zeta"} + if surface_label == "rho": + unique_size = grid.num_rho + inverse_idx = grid.inverse_rho_idx + spacing = grid.spacing[:, 1:] + has_endpoint_dupe = False + elif surface_label == "theta": + unique_size = grid.num_theta + inverse_idx = grid.inverse_theta_idx + spacing = grid.spacing[:, [0, 2]] + has_endpoint_dupe = (grid.nodes[grid.unique_theta_idx[0], 1] == 0) & ( + grid.nodes[grid.unique_theta_idx[-1], 1] == 2 * np.pi + ) + else: + unique_size = grid.num_zeta + inverse_idx = grid.inverse_zeta_idx + spacing = grid.spacing[:, :2] + has_endpoint_dupe = (grid.nodes[grid.unique_zeta_idx[0], 2] == 0) & ( + grid.nodes[grid.unique_zeta_idx[-1], 2] == 2 * np.pi / grid.NFP + ) + return unique_size, inverse_idx, spacing, has_endpoint_dupe + + def line_integrals( grid, q=jnp.array([1.0]), @@ -792,7 +724,7 @@ def surface_integrals(grid, q=jnp.array([1.0]), surface_label="rho", expand_out= Notes ----- - It is assumed that the integration surface has area 4π^2 when the + It is assumed that the integration surface has area 4π² when the surface label is rho and area 2π when the surface label is theta or zeta. You may want to multiply the input by the surface area Jacobian. @@ -867,8 +799,8 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True): grid, surface_label ) - # Todo: Define masks as a sparse matrix once sparse matrices are - # are no longer experimental in jax. + # Todo: Define masks as a sparse matrix once sparse matrices are no longer + # experimental in jax. # The ith row of masks is True only at the indices which correspond to the # ith surface. The integral over the ith surface is the dot product of the # ith row vector and the vector of integrands of all surfaces. @@ -883,7 +815,7 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True): # surface will have the correct total area of π+π = 2π. # An edge case exists if the duplicate surface has nodes with # different values for the surface label, which only occurs when - # has_endpoint_dupe is true. If has_endpoint_dupe is true, this grid + # has_endpoint_dupe is true. If ``has_endpoint_dupe`` is true, this grid # has a duplicate surface at surface_label=0 and # surface_label=max surface value. Although the modulo of these values # are equal, their numeric values are not, so the integration @@ -893,9 +825,8 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True): # previous paragraph. masks = cond( has_endpoint_dupe, - lambda masks: put(masks, jnp.asarray([0, -1]), masks[0] | masks[-1]), - lambda masks: masks, - masks, + lambda: put(masks, jnp.array([0, -1]), masks[0] | masks[-1]), + lambda: masks, ) spacing = jnp.prod(spacing, axis=1) @@ -904,7 +835,7 @@ def _surface_integrals(q=jnp.array([1.0])): Notes ----- - It is assumed that the integration surface has area 4π^2 when the + It is assumed that the integration surface has area 4π² when the surface label is rho and area 2π when the surface label is theta or zeta. You may want to multiply the input by the surface area Jacobian. @@ -932,6 +863,7 @@ def _surface_integrals(q=jnp.array([1.0])): Surface integral of the input over each surface in the grid. """ + axis_to_move = (jnp.ndim(q) == 3) * 2 integrands = (spacing * jnp.nan_to_num(q).T).T # `integrands` may have shape (g.size, f.size, v.size), where # g is the grid function depending on the integration variables @@ -963,11 +895,10 @@ def _surface_integrals(q=jnp.array([1.0])): # shape (v.size, g.size, f.size). As we expect f.size >> v.size, the # integration is in theory faster since numpy optimizes large matrix # products. However, timing results showed no difference. - axis_to_move = (integrands.ndim == 3) * 2 integrals = jnp.moveaxis( masks @ jnp.moveaxis(integrands, axis_to_move, 0), 0, axis_to_move ) - return expand(grid, integrals, surface_label) if expand_out else integrals + return grid.expand(integrals, surface_label) if expand_out else integrals return _surface_integrals @@ -1011,9 +942,10 @@ def surface_averages( surface_label : str The surface label of rho, theta, or zeta to compute the average over. denominator : ndarray - This can optionally be supplied to avoid redundant computations. - Volume enclosed by the surfaces, derivative wrt the surface label. - This array should succeed broadcasting with arrays of size + By default, the denominator is computed as the surface integral of + ``sqrt_g``. This parameter can optionally be supplied to avoid + redundant computations or to use a different denominator to compute + the average. This array should broadcast with arrays of size ``grid.num_nodes`` (``grid.num_surface_label``) if ``expand_out`` is true (false). expand_out : bool @@ -1051,7 +983,7 @@ def surface_averages_map(grid, surface_label="rho", expand_out=True): ------- function : callable Method to compute any surface average of the input ``q`` and optionally - the volume Jacobian ``sqrt_g`` over each surface in the grid with code: + the volume Jacobian ``sqrt_g`` over each surface in the grid with code: ``function(q, sqrt_g)``. """ @@ -1085,9 +1017,10 @@ def _surface_averages(q, sqrt_g=jnp.array([1.0]), denominator=None): sqrt_g : ndarray Coordinate system Jacobian determinant; see ``data_index["sqrt(g)"]``. denominator : ndarray - This can optionally be supplied to avoid redundant computations. - Volume enclosed by the surfaces, derivative wrt the surface label. - This array should succeed broadcasting with arrays of size + By default, the denominator is computed as the surface integral of + ``sqrt_g``. This parameter can optionally be supplied to avoid + redundant computations or to use a different denominator to compute + the average. This array should broadcast with arrays of size ``grid.num_nodes`` (``grid.num_surface_label``) if ``expand_out`` is true (false). @@ -1110,11 +1043,11 @@ def _surface_averages(q, sqrt_g=jnp.array([1.0]), denominator=None): ) averages = (numerator.T / denominator).T if expand_out: - averages = expand(grid, averages, surface_label) + averages = grid.expand(averages, surface_label) else: if expand_out: # implies denominator given with size grid.num_nodes - numerator = expand(grid, numerator, surface_label) + numerator = grid.expand(numerator, surface_label) averages = (numerator.T / denominator).T return averages @@ -1130,15 +1063,15 @@ def surface_integrals_transform(grid, surface_label="rho"): five variables, the returned method computes an integral transform, reducing ``q`` to a set of functions of at most three variables. - Define the domain D = u_1 × u_2 × u_3 and the codomain C = u_4 × u_5 × u_6. - For every surface of constant u_1 in the domain, the returned method - evaluates the transform T_{u_1}: (u_2 × u_3) × C → C, where T_{u_1} projects - away the parameters u_2 and u_3 via an integration of the given kernel - function K_{u_1} over the corresponding surface of constant u_1. + Define the domain D = u₁ × u₂ × u₃ and the codomain C = u₄ × u₅ × u₆. + For every surface of constant u₁ in the domain, the returned method + evaluates the transform Tᵤ₁ : u₂ × u₃ × C → C, where Tᵤ₁ projects + away the parameters u₂ and u₃ via an integration of the given kernel + function Kᵤ₁ over the corresponding surface of constant u₁. Notes ----- - It is assumed that the integration surface has area 4π^2 when the + It is assumed that the integration surface has area 4π² when the surface label is rho and area 2π when the surface label is theta or zeta. You may want to multiply the input ``q`` by the surface area Jacobian. @@ -1150,7 +1083,7 @@ def surface_integrals_transform(grid, surface_label="rho"): surface_label : str The surface label of rho, theta, or zeta to compute the integration over. These correspond to the domain parameters discussed in this method's - description. In particular, ``surface_label`` names u_1. + description. In particular, ``surface_label`` names u₁. Returns ------- @@ -1189,8 +1122,8 @@ def surface_integrals_transform(grid, surface_label="rho"): Output ------ Each element along the first dimension of the returned array, stores - T_{u_1} for a particular surface of constant u_1 in the given grid. - The order is sorted in increasing order of the values which specify u_1. + Tᵤ₁ for a particular surface of constant u₁ in the given grid. + The order is sorted in increasing order of the values which specify u₁. If ``q`` is one-dimensional, the returned array has shape (grid.num_surface_label, ). @@ -1220,31 +1153,19 @@ def surface_variance( surface_label="rho", expand_out=True, ): - r"""Compute the weighted sample variance of ``q`` on each surface of the grid. - - Computes the following quantity on each surface of the grid. - - .. math:: - \frac{n_e}{n_e - b} - \frac{ \sum_{i=1}^{n} (q_i - \bar{q})^2 w_i }{ \sum_{i=1}^{n} w_i } - - where - :math:`w_i` is the weight assigned to :math:`q_i` given by the product - of ``weights[i]`` and the differential surface area element (not already - weighted by the area Jacobian) at the node where ``q[i]`` is evaluated, - :math:`\bar{q}` is the weighted mean of :math:`q`, - :math:`b` is 0 if the biased sample variance is to be returned and 1 otherwise, - :math:`n` is the number of samples on a surface, and - :math:`n_e` is the effective number of samples on a surface defined as - - .. math:: - (\sum_{i=1}^{n} w_i)^2 / (\sum_{i=1}^{n} w_i^2) + """Compute the weighted sample variance of ``q`` on each surface of the grid. - As the weights :math:`w_i` approach each other, :math:`n_e` approaches - :math:`n`, and the output converges to + Computes nₑ / (nₑ − b) * (∑ᵢ₌₁ⁿ (qᵢ − q̅)² wᵢ) / (∑ᵢ₌₁ⁿ wᵢ). + wᵢ is the weight assigned to qᵢ given by the product of ``weights[i]`` and + the differential surface area element (not already weighted by the area + Jacobian) at the node where qᵢ is evaluated, + q̅ is the weighted mean of q, + b is 0 if the biased sample variance is to be returned and 1 otherwise, + n is the number of samples on a surface, and + nₑ ≝ (∑ᵢ₌₁ⁿ wᵢ)² / ∑ᵢ₌₁ⁿ wᵢ² is the effective number of samples. - .. math:: - \frac{1}{n - b} \sum_{i=1}^{n} (q_i - \bar{q})^2 + As the weights wᵢ approach each other, nₑ approaches n, and the output + converges to ∑ᵢ₌₁ⁿ (qᵢ − q̅)² / (n − b). Notes ----- @@ -1263,7 +1184,7 @@ def surface_variance( https://en.wikipedia.org/wiki/Inverse-variance_weighting. The unbiased sample variance for this case is obtained by replacing the effective number of samples in the formula this function implements, - :math:`n_e`, with the actual number of samples :math:`n`. + nₑ, with the actual number of samples n. The third case is when the weights denote the integer frequency of each sample. See @@ -1371,8 +1292,6 @@ def body(i, mins): return mins mins = fori_loop(0, inverse_idx.size, body, mins) - # The above implementation was benchmarked to be more efficient, after jit - # compilation, than the alternative given in the two lines below. - # masks = inverse_idx == jnp.arange(unique_size)[:, jnp.newaxis] # noqa: E501,E800 - # mins = jnp.amin(x[jnp.newaxis, :], axis=1, initial=jnp.inf, where=masks) # noqa: E501,E800 - return expand(grid, mins, surface_label) + # The above implementation was benchmarked to be more efficient than + # alternatives without explicit loops in GitHub pull request #501. + return grid.expand(mins, surface_label) diff --git a/desc/equilibrium/coords.py b/desc/equilibrium/coords.py index 5491403390..e4363a7225 100644 --- a/desc/equilibrium/coords.py +++ b/desc/equilibrium/coords.py @@ -72,7 +72,7 @@ def map_coordinates( # noqa: C901 basis_derivs = [f"{X}_{d}" for X in inbasis for d in ("r", "t", "z")] for key in basis_derivs: assert ( - key in data_index["desc.equilibrium.equilibrium.Equilibrium"].keys() + key in data_index["desc.equilibrium.equilibrium.Equilibrium"] ), f"don't have recipe to compute partial derivative {key}" rhomin = kwargs.pop("rhomin", tol / 10) diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index cd537083c8..51d5a4ca68 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -13,14 +13,7 @@ from desc.basis import FourierZernikeBasis, fourier, zernike_radial from desc.compute import compute as compute_fun from desc.compute import data_index -from desc.compute.utils import ( - compress, - expand, - get_data_deps, - get_params, - get_profiles, - get_transforms, -) +from desc.compute.utils import get_data_deps, get_params, get_profiles, get_transforms from desc.geometry import ( FourierRZCurve, FourierRZToroidalSurface, @@ -381,6 +374,10 @@ def _set_up(self): for attribute in self._io_attrs_: if not hasattr(self, attribute): setattr(self, attribute, None) + if self.current is not None and hasattr(self.current, "_get_transform"): + # Need to rebuild derivative matrices to get higher order derivatives + # on equilibrium's saved before GitHub pull request #586. + self.current._transform = self.current._get_transform(self.current.grid) def __repr__(self): """String form of the object.""" @@ -657,7 +654,7 @@ def get_profile(self, name, grid=None, **kwargs): grid = QuadratureGrid(self.L_grid, self.M_grid, self.N_grid, self.NFP) data = self.compute(name, grid=grid, **kwargs) x = data[name] - x = compress(grid, x, surface_label="rho") + x = grid.compress(x, surface_label="rho") return SplineProfile( x, grid.nodes[grid.unique_rho_idx, 0], grid=grid, name=name ) @@ -808,9 +805,9 @@ def compute( data=None, **kwargs, ) - # need to make this data broadcastable with the data on the original grid + # need to make this data broadcast with the data on the original grid data1d = { - key: expand(grid, compress(grid1d, val)) + key: grid.expand(grid1d.compress(val)) for key, val in data1d.items() if key in dep1d } @@ -859,7 +856,7 @@ def map_coordinates( # noqa: C901 same as the compute function data key guess : None or ndarray, shape(k,3) Initial guess for the computational coordinates ['rho', 'theta', 'zeta'] - coresponding to coords in inbasis. If None, heuristics are used based on + corresponding to coords in inbasis. If None, heuristics are used based on in basis and a nearest neighbor search on a coarse grid. period : tuple of float Assumed periodicity for each quantity in inbasis. @@ -1467,25 +1464,24 @@ def from_near_axis( if ntheta is None: ntheta = 2 * M + 1 - inputs = {} - inputs["Psi"] = np.pi * r**2 * na_eq.Bbar - inputs["NFP"] = na_eq.nfp - inputs["L"] = L - inputs["M"] = M - inputs["N"] = N - inputs["sym"] = not na_eq.lasym - inputs["spectral_indexing "] = spectral_indexing - inputs["pressure"] = np.array( - [[0, -na_eq.p2 * r**2], [2, na_eq.p2 * r**2]] - ) - inputs["iota"] = None - inputs["current"] = np.array([[2, 2 * np.pi / mu_0 * na_eq.I2 * r**2]]) - inputs["axis"] = FourierRZCurve( - R_n=np.concatenate((np.flipud(na_eq.rs[1:]), na_eq.rc)), - Z_n=np.concatenate((np.flipud(na_eq.zs[1:]), na_eq.zc)), - NFP=na_eq.nfp, - ) - inputs["surface"] = None + inputs = { + "Psi": np.pi * r**2 * na_eq.Bbar, + "NFP": na_eq.nfp, + "L": L, + "M": M, + "N": N, + "sym": not na_eq.lasym, + "spectral_indexing ": spectral_indexing, + "pressure": np.array([[0, -na_eq.p2 * r**2], [2, na_eq.p2 * r**2]]), + "iota": None, + "current": np.array([[2, 2 * np.pi / mu_0 * na_eq.I2 * r**2]]), + "axis": FourierRZCurve( + R_n=np.concatenate((np.flipud(na_eq.rs[1:]), na_eq.rc)), + Z_n=np.concatenate((np.flipud(na_eq.zs[1:]), na_eq.zc)), + NFP=na_eq.nfp, + ), + "surface": None, + } except AttributeError as e: raise ValueError("Input must be a pyQSC or pyQIC solution.") from e @@ -1526,7 +1522,6 @@ def from_near_axis( Z_1D = np.zeros((grid.num_nodes,)) L_1D = np.zeros((grid.num_nodes,)) for rho_i in rho: - idx = idx = np.where(grid.nodes[:, 0] == rho_i)[0] R_2D, Z_2D, phi0_2D = na_eq.Frenet_to_cylindrical(r * rho_i, ntheta) phi_cyl_ax = np.linspace( 0, 2 * np.pi / na_eq.nfp, na_eq.nphi, endpoint=False @@ -1534,6 +1529,7 @@ def from_near_axis( nu_B_ax = na_eq.nu_spline(phi_cyl_ax) phi_B = phi_cyl_ax + nu_B_ax nu_B = phi_B - phi0_2D + idx = np.nonzero(grid.nodes[:, 0] == rho_i)[0] R_1D[idx] = R_2D.flatten(order="F") Z_1D[idx] = Z_2D.flatten(order="F") L_1D[idx] = nu_B.flatten(order="F") * na_eq.iota diff --git a/desc/objectives/_bootstrap.py b/desc/objectives/_bootstrap.py index e02ab1ab6c..9052f2eea0 100644 --- a/desc/objectives/_bootstrap.py +++ b/desc/objectives/_bootstrap.py @@ -7,7 +7,6 @@ from desc.backend import jnp from desc.compute import compute as compute_fun from desc.compute import get_params, get_profiles, get_transforms -from desc.compute.utils import compress from desc.grid import LinearGrid from desc.utils import Timer @@ -219,11 +218,8 @@ def compute(self, *args, **kwargs): profiles=constants["profiles"], helicity=constants["helicity"], ) - - return compress( - constants["transforms"]["grid"], - data[""] - data[" Redl"], - surface_label="rho", + return constants["transforms"]["grid"].compress( + data[""] - data[" Redl"] ) def _scale(self, *args, **kwargs): @@ -231,10 +227,8 @@ def _scale(self, *args, **kwargs): constants = kwargs.get("constants", None) if constants is None: constants = self.constants - w = compress( - constants["transforms"]["grid"], - constants["transforms"]["grid"].spacing[:, 0], - surface_label="rho", + w = constants["transforms"]["grid"].compress( + constants["transforms"]["grid"].spacing[:, 0] ) return super()._scale(*args, **kwargs) * jnp.sqrt(w) diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index 86c3e14a8c..c9c1487744 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -595,7 +595,7 @@ def compute(self, *args, **kwargs): class Energy(_Objective): """MHD energy. - W = integral( B^2 / (2*mu0) + p / (gamma - 1) ) dV (J) + W = integral( ||B||^2 / (2*mu0) + p / (gamma - 1) ) dV (J) Parameters ---------- diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 503a965a60..e9bfbf6650 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -5,7 +5,7 @@ from desc.backend import jnp from desc.compute import compute as compute_fun from desc.compute import data_index -from desc.compute.utils import compress, get_params, get_profiles, get_transforms +from desc.compute.utils import get_params, get_profiles, get_transforms from desc.grid import LinearGrid, QuadratureGrid from desc.profiles import Profile from desc.utils import Timer @@ -62,7 +62,7 @@ class ObjectiveFromUser(_Objective): -------- .. code-block:: python - from desc.compute.utils import surface_averages, compress + from desc.compute.utils import surface_averages def myfun(grid, data): # This will compute the flux surface average of the function # R*B_T from the Grad-Shafranov equation @@ -70,7 +70,7 @@ def myfun(grid, data): f_fsa = surface_averages(grid, f, sqrt_g=data['sqrt_g']) # this has the FSA values on the full grid, but we just want # the unique values: - return compress(grid, f_fsa) + return grid.compress(f_fsa) myobj = ObjectiveFromUser(myfun) @@ -126,14 +126,14 @@ def build(self, eq=None, use_jit=True, verbose=1): else: grid = self._grid - def getvars(fun): + def get_vars(fun): pattern = r"data\[(.*?)\]" src = inspect.getsource(fun) variables = re.findall(pattern, src) - variables = [s.replace("'", "").replace('"', "") for s in variables] + variables = list({s.strip().strip("'").strip('"') for s in variables}) return variables - self._data_keys = getvars(self._fun) + self._data_keys = get_vars(self._fun) dummy_data = {} p = "desc.equilibrium.equilibrium.Equilibrium" for key in self._data_keys: @@ -479,19 +479,15 @@ def compute(self, *args, **kwargs): transforms=constants["transforms"], profiles=constants["profiles"], ) - return compress( - constants["transforms"]["grid"], data["current"], surface_label="rho" - ) + return constants["transforms"]["grid"].compress(data["current"]) def _scale(self, *args, **kwargs): """Compute and apply the target/bounds, weighting, and normalization.""" constants = kwargs.get("constants", None) if constants is None: constants = self.constants - w = compress( - constants["transforms"]["grid"], - constants["transforms"]["grid"].spacing[:, 0], - surface_label="rho", + w = constants["transforms"]["grid"].compress( + constants["transforms"]["grid"].spacing[:, 0] ) return super()._scale(*args, **kwargs) * jnp.sqrt(w) @@ -669,19 +665,15 @@ def compute(self, *args, **kwargs): transforms=constants["transforms"], profiles=constants["profiles"], ) - return compress( - constants["transforms"]["grid"], data["iota"], surface_label="rho" - ) + return constants["transforms"]["grid"].compress(data["iota"]) def _scale(self, *args, **kwargs): """Compute and apply the target/bounds, weighting, and normalization.""" constants = kwargs.get("constants", None) if constants is None: constants = self.constants - w = compress( - constants["transforms"]["grid"], - constants["transforms"]["grid"].spacing[:, 0], - surface_label="rho", + w = constants["transforms"]["grid"].compress( + constants["transforms"]["grid"].spacing[:, 0] ) return super()._scale(*args, **kwargs) * jnp.sqrt(w) diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index c81f26194b..b65c931ab2 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -5,7 +5,6 @@ from desc.backend import jnp from desc.compute import compute as compute_fun from desc.compute import get_params, get_profiles, get_transforms -from desc.compute.utils import compress from desc.grid import LinearGrid from desc.utils import Timer @@ -178,19 +177,15 @@ def compute(self, *args, **kwargs): transforms=constants["transforms"], profiles=constants["profiles"], ) - return compress( - constants["transforms"]["grid"], data["D_Mercier"], surface_label="rho" - ) + return constants["transforms"]["grid"].compress(data["D_Mercier"]) def _scale(self, *args, **kwargs): """Compute and apply the target/bounds, weighting, and normalization.""" constants = kwargs.get("constants", None) if constants is None: constants = self.constants - w = compress( - constants["transforms"]["grid"], - constants["transforms"]["grid"].spacing[:, 0], - surface_label="rho", + w = constants["transforms"]["grid"].compress( + constants["transforms"]["grid"].spacing[:, 0] ) return super()._scale(*args, **kwargs) * jnp.sqrt(w) @@ -380,19 +375,15 @@ def compute(self, *args, **kwargs): transforms=constants["transforms"], profiles=constants["profiles"], ) - return compress( - constants["transforms"]["grid"], data["magnetic well"], surface_label="rho" - ) + return constants["transforms"]["grid"].compress(data["magnetic well"]) def _scale(self, *args, **kwargs): """Compute and apply the target/bounds, weighting, and normalization.""" constants = kwargs.get("constants", None) if constants is None: constants = self.constants - w = compress( - constants["transforms"]["grid"], - constants["transforms"]["grid"].spacing[:, 0], - surface_label="rho", + w = constants["transforms"]["grid"].compress( + constants["transforms"]["grid"].spacing[:, 0] ) return super()._scale(*args, **kwargs) * jnp.sqrt(w) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 1b3a0b31d4..82419f1ebd 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -40,7 +40,7 @@ def update_target(self, eq): class BoundaryRSelfConsistency(_Objective): - """Ensure that the boundary and interior surfaces are self consistent. + """Ensure that the boundary and interior surfaces are self-consistent. Note: this constraint is automatically applied when needed, and does not need to be included by the user. diff --git a/desc/optimize/aug_lagrangian.py b/desc/optimize/aug_lagrangian.py index d5ed684765..c4654c9e2e 100644 --- a/desc/optimize/aug_lagrangian.py +++ b/desc/optimize/aug_lagrangian.py @@ -1,4 +1,4 @@ -"""Augmented Langrangian for scalar valued objectives.""" +"""Augmented Lagrangian for scalar valued objectives.""" from scipy.optimize import BFGS, NonlinearConstraint, OptimizeResult from termcolor import colored @@ -32,7 +32,7 @@ def fmin_auglag( # noqa: C901 - FIXME: simplify this maxiter=None, options={}, ): - """Minimize a function with constraints using an augmented Langrangian method. + """Minimize a function with constraints using an augmented Lagrangian method. Parameters ---------- diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py index 90f7836954..ed3ad42da6 100644 --- a/desc/optimize/aug_lagrangian_ls.py +++ b/desc/optimize/aug_lagrangian_ls.py @@ -1,4 +1,4 @@ -"""Augmented Langrangian for vector valued objectives.""" +"""Augmented Lagrangian for vector valued objectives.""" from scipy.optimize import NonlinearConstraint, OptimizeResult @@ -30,7 +30,7 @@ def lsq_auglag( # noqa: C901 - FIXME: simplify this maxiter=None, options={}, ): - """Minimize a function with constraints using an augmented Langrangian method. + """Minimize a function with constraints using an augmented Lagrangian method. The objective function is assumed to be vector valued, and is minimized in the least squares sense. diff --git a/desc/optimize/fmin_scalar.py b/desc/optimize/fmin_scalar.py index 4498f1b681..be0d7704ee 100644 --- a/desc/optimize/fmin_scalar.py +++ b/desc/optimize/fmin_scalar.py @@ -90,7 +90,7 @@ def fmintr( # noqa: C901 - FIXME: simplify this If None, the termination by this condition is disabled. gtol : float or None, optional Absolute tolerance for termination by the norm of the gradient. - Optimizer teriminates when ``max(abs(g)) < gtol``. + Optimizer terminates when ``max(abs(g)) < gtol``. If None, the termination by this condition is disabled. verbose : {0, 1, 2}, optional * 0 (default) : work silently. diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py index 2ac5b15856..73d76ecacf 100644 --- a/desc/optimize/least_squares.py +++ b/desc/optimize/least_squares.py @@ -80,7 +80,7 @@ def lsqtr( # noqa: C901 - FIXME: simplify this If None, the termination by this condition is disabled. gtol : float or None, optional Absolute tolerance for termination by the norm of the gradient. - Optimizer teriminates when ``max(abs(g)) < gtol``. + Optimizer terminates when ``max(abs(g)) < gtol``. If None, the termination by this condition is disabled. verbose : {0, 1, 2}, optional * 0 (default) : work silently. diff --git a/desc/optimize/stochastic.py b/desc/optimize/stochastic.py index ac918920f0..3848c0c52f 100644 --- a/desc/optimize/stochastic.py +++ b/desc/optimize/stochastic.py @@ -55,7 +55,7 @@ def sgd( If None, the termination by this condition is disabled. gtol : float or None, optional Absolute tolerance for termination by the norm of the gradient. - Optimizer teriminates when ``max(abs(g)) < gtol``. + Optimizer terminates when ``max(abs(g)) < gtol``. If None, the termination by this condition is disabled. verbose : {0, 1, 2}, optional * 0 (default) : work silently. diff --git a/desc/plotting.py b/desc/plotting.py index 342c71275d..1b21f0df43 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -16,7 +16,7 @@ from desc.backend import sign from desc.basis import fourier, zernike_radial_poly from desc.compute import data_index, get_transforms -from desc.compute.utils import compress, surface_averages +from desc.compute.utils import surface_averages_map from desc.grid import Grid, LinearGrid from desc.utils import flatten_list, parse_argname_change from desc.vmec_utils import ptolemy_linear_transform @@ -241,11 +241,11 @@ def _get_plot_axes(grid): """ plot_axes = [0, 1, 2] - if np.unique(grid.nodes[:, 0]).size == 1: + if grid.num_rho == 1: plot_axes.remove(0) - if np.unique(grid.nodes[:, 1]).size == 1: + if grid.num_theta == 1: plot_axes.remove(1) - if np.unique(grid.nodes[:, 2]).size == 1: + if grid.num_zeta == 1: plot_axes.remove(2) return tuple(plot_axes) @@ -280,11 +280,7 @@ def _compute(eq, name, grid, component=None, reshape=True): "Z", ], f"component must be one of [None, 'R', 'phi', 'Z'], got {component}" - components = { - "R": 0, - "phi": 1, - "Z": 2, - } + components = {"R": 0, "phi": 1, "Z": 2} label = data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["label"] @@ -531,9 +527,8 @@ def plot_1d(eq, name, grid=None, log=False, ax=None, return_data=False, **kwargs ax.set_xlabel(xlabel, fontsize=xlabel_fontsize) ax.set_ylabel(label, fontsize=ylabel_fontsize) _set_tight_layout(fig) - plot_data = {} - plot_data[xlabel.strip("$").strip("\\")] = grid.nodes[:, plot_axes[0]] - plot_data[name] = data + plot_data = {xlabel.strip("$").strip("\\"): grid.nodes[:, plot_axes[0]], name: data} + if return_data: return fig, ax, plot_data @@ -686,10 +681,12 @@ def plot_2d( ) ) _set_tight_layout(fig) - plot_data = {} - plot_data[xlabel.strip("$").strip("\\")] = xx - plot_data[ylabel.strip("$").strip("\\")] = yy - plot_data[name] = data + plot_data = { + xlabel.strip("$").strip("\\"): xx, + ylabel.strip("$").strip("\\"): yy, + name: data, + } + if norm_F: plot_data["normalization"] = np.nanmean(np.abs(norm_data)) else: @@ -875,11 +872,7 @@ def plot_3d( ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) - plot_data = {} - plot_data["X"] = X - plot_data["Y"] = Y - plot_data["Z"] = Z - plot_data[name] = data + plot_data = {"X": X, "Y": Y, "Z": Z, name: data} if return_data: return fig, ax, plot_data @@ -974,16 +967,8 @@ def plot_fsa( """ if np.isscalar(rho) and (int(rho) == rho): - if ( - data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["coordinates"] - == "r" - ): - # OK to plot origin for most quantities denoted as functions of rho - rho = np.flipud(np.linspace(1, 0, rho + 1, endpoint=True)) - else: - rho = np.linspace(1 / rho, 1, rho) - else: - rho = np.atleast_1d(rho) + rho = np.linspace(0, 1, rho + 1) + rho = np.atleast_1d(rho) if M is None: M = eq.M_grid if N is None: @@ -994,38 +979,81 @@ def plot_fsa( fig, ax = _format_ax(ax, figsize=kwargs.pop("figsize", (4, 4))) grid = LinearGrid(M=M, N=N, NFP=eq.NFP, rho=rho) + + p = "desc.equilibrium.equilibrium.Equilibrium" + if "<" + name + ">" in data_index[p]: + # If we identify the quantity to plot as something in data_index, then + # we may be able to compute more involved magnetic axis limits. + deps = data_index[p]["<" + name + ">"]["dependencies"]["data"] + if with_sqrt_g == ("sqrt(g)" in deps or "V_r(r)" in deps): + # When we denote a quantity as ```` in data_index, we have + # marked it a surface average of ``name``. This does not specify + # the type of surface average however (i.e. with/without the sqrt(g) + # factor). The above conditional guard should ensure that the + # surface average we have the recipe to compute in data_index is the + # desired surface average. + name = "<" + name + ">" values, label = _compute( eq, name, grid, kwargs.pop("component", None), reshape=False ) label = label.split("~") if ( - data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["coordinates"] - == "r" + data_index[p][name]["coordinates"] == "r" + or data_index[p][name]["coordinates"] == "" ): # If the quantity is a surface function, averaging it again has no # effect, regardless of whether sqrt(g) is used. # So we avoid surface averaging it and forgo the <> around the label. label = r"$ " + label[0][1:] + r" ~" + "~".join(label[1:]) plot_data_ylabel_key = f"{name}" - values = compress(grid, values) - elif with_sqrt_g: - # flux surface average - label = r"$\langle " + label[0][1:] + r" \rangle~" + "~".join(label[1:]) - plot_data_ylabel_key = f"<{name}>_fsa" - sqrt_g, _ = _compute(eq, "sqrt(g)", grid, reshape=False) - values = surface_averages(grid, q=values, sqrt_g=sqrt_g, expand_out=False) + if data_index[p][name]["coordinates"] == "r": + values = grid.compress(values) else: - # theta average - label = ( - r"$\langle " + label[0][1:] + r" \rangle_{\theta}~" + "~".join(label[1:]) - ) + compute_surface_averages = surface_averages_map(grid, expand_out=False) + if with_sqrt_g: # flux surface average + sqrt_g = _compute(eq, "sqrt(g)", grid, reshape=False)[0] + # Attempt to compute the magnetic axis limit. + # Compute derivative depending on various naming schemes. + # e.g. B -> B_r, V(r) -> V_r(r), S_r(r) -> S_rr(r) + schemes = ( + name + "_r", + name[:-3] + "_r" + name[-3:], + name[:-3] + "r" + name[-3:], + ) + values_r = next( + ( + _compute(eq, x, grid, reshape=False)[0] + for x in schemes + if x in data_index[p] + ), + np.nan, + ) + if (np.isfinite(values) & np.isfinite(values_r))[grid.axis].all(): + # Otherwise cannot compute axis limit in this agnostic manner. + sqrt_g = grid.replace_at_axis( + sqrt_g, _compute(eq, "sqrt(g)_r", grid, reshape=False)[0], copy=True + ) + averages = compute_surface_averages(values, sqrt_g=sqrt_g) + label = r"$\langle " + label[0][1:] + r" \rangle~" + "~".join(label[1:]) + else: # theta average + averages = compute_surface_averages(values) + label = ( + r"$\langle " + + label[0][1:] + + r" \rangle_{\theta}~" + + "~".join(label[1:]) + ) + # True if values has nan on a given surface. + is_nan = compute_surface_averages(np.isnan(values)).astype(bool) + # The integration replaced nan with 0. + # Put them back to avoid misleading plot (e.g. cusp near the magnetic axis). + values = np.where(is_nan, np.nan, averages) plot_data_ylabel_key = f"<{name}>_fsa" - values = surface_averages(grid, q=values, expand_out=False) if norm_F: # normalize force by B pressure gradient norm_name = kwargs.pop("norm_name", "<|grad(|B|^2)|/2mu0>_vol") - norm_data, _ = _compute(eq, norm_name, grid, reshape=False) + norm_data = _compute(eq, norm_name, grid, reshape=False)[0] values = values / np.nanmean(np.abs(norm_data)) # normalize if log: values = np.abs(values) # ensure data is positive for log plot @@ -1048,22 +1076,14 @@ def plot_fsa( ax.set_ylabel( "%s / %s" % ( - "$" - + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["label"] - + "$", - "$" - + data_index["desc.equilibrium.equilibrium.Equilibrium"][norm_name][ - "label" - ] - + "$", + "$" + data_index[p][name]["label"] + "$", + "$" + data_index[p][norm_name]["label"] + "$", ), fontsize=ylabel_fontsize, ) _set_tight_layout(fig) - plot_data = {} - plot_data["rho"] = rho - plot_data[plot_data_ylabel_key] = values + plot_data = {"rho": rho, plot_data_ylabel_key: values} if norm_F: plot_data["normalization"] = np.nanmean(np.abs(norm_data)) else: @@ -1274,10 +1294,7 @@ def plot_section( ) _set_tight_layout(fig) - plot_data = {} - plot_data["R"] = R - plot_data["Z"] = Z - plot_data[name] = data + plot_data = {"R": R, "Z": Z, name: data} if norm_F: plot_data["normalization"] = np.nanmean(np.abs(norm_data)) else: @@ -1390,7 +1407,7 @@ def plot_surfaces(eq, rho=8, theta=8, phi=None, ax=None, return_data=False, **kw plot_theta = bool(theta) nfp = eq.NFP if isinstance(rho, numbers.Integral): - rho = np.linspace(0, 1, rho + 1) # offset to ignore axis + rho = np.linspace(0, 1, rho + 1) rho = np.atleast_1d(rho) if isinstance(theta, numbers.Integral): theta = np.linspace(0, 2 * np.pi, theta, endpoint=False) @@ -1464,13 +1481,7 @@ def plot_surfaces(eq, rho=8, theta=8, phi=None, ax=None, return_data=False, **kw figh = 5 * rows if figsize is None: figsize = (figw, figh) - fig, ax = _format_ax( - ax, - rows=rows, - cols=cols, - figsize=figsize, - equal=True, - ) + fig, ax = _format_ax(ax, rows=rows, cols=cols, figsize=figsize, equal=True) ax = np.atleast_1d(ax).flatten() for i in range(nphi): @@ -1483,11 +1494,7 @@ def plot_surfaces(eq, rho=8, theta=8, phi=None, ax=None, return_data=False, **kw lw=theta_lw, ) ax[i].plot( - Rr[:, :, i], - Zr[:, :, i], - color=rho_color, - linestyle=rho_ls, - lw=rho_lw, + Rr[:, :, i], Zr[:, :, i], color=rho_color, linestyle=rho_ls, lw=rho_lw ) ax[i].plot( Rr[:, -1, i], @@ -1524,7 +1531,7 @@ def plot_surfaces(eq, rho=8, theta=8, phi=None, ax=None, return_data=False, **kw return fig, ax -def plot_boundary(eq, phi=None, plot_axis=False, ax=None, return_data=False, **kwargs): +def plot_boundary(eq, phi=None, plot_axis=True, ax=None, return_data=False, **kwargs): """Plot stellarator boundary at multiple toroidal coordinates. Parameters @@ -1536,7 +1543,7 @@ def plot_boundary(eq, phi=None, plot_axis=False, ax=None, return_data=False, **k If an integer, plot that many contours linearly spaced in [0,2pi). Default is 1 contour for axisymmetric equilibria or 4 for non-axisymmetry. plot_axis : bool - Whether or not to plot the magnetic axis locations. Default is False. + Whether to plot the magnetic axis locations. Default is True. ax : matplotlib AxesSubplot, optional Axis to plot on. return_data : bool @@ -1649,13 +1656,7 @@ def plot_boundary(eq, phi=None, plot_axis=False, ax=None, return_data=False, **k ), ) if rho[0] == 0: - ax.scatter( - R[0, 0, i], - Z[0, 0, i], - color=colors[i], - marker=marker, - s=size, - ) + ax.scatter(R[0, 0, i], Z[0, 0, i], color=colors[i], marker=marker, s=size) ax.set_xlabel(_AXIS_LABELS_RPZ[0], fontsize=xlabel_fontsize) ax.set_ylabel(_AXIS_LABELS_RPZ[2], fontsize=ylabel_fontsize) @@ -1766,11 +1767,7 @@ def plot_boundaries(eqs, labels=None, phi=None, ax=None, return_data=False, **kw plot_data["Z"] = [] for i in range(neq): - grid_kwargs = { - "NFP": eqs[i].NFP, - "theta": 100, - "zeta": phi, - } + grid_kwargs = {"NFP": eqs[i].NFP, "theta": 100, "zeta": phi} grid = _get_grid(**grid_kwargs) nr, nt, nz = grid.num_rho, grid.num_theta, grid.num_zeta grid = Grid( @@ -1793,11 +1790,7 @@ def plot_boundaries(eqs, labels=None, phi=None, ax=None, return_data=False, **kw for j in range(nz - 1): (line,) = ax.plot( - R[:, -1, j], - Z[:, -1, j], - color=colors[i], - linestyle=ls[i], - lw=lw[i], + R[:, -1, j], Z[:, -1, j], color=colors[i], linestyle=ls[i], lw=lw[i] ) if j == 0: line.set_label(labels[i]) @@ -1938,13 +1931,7 @@ def plot_comparison( figh = 5 * rows if figsize is None: figsize = (figw, figh) - fig, ax = _format_ax( - ax, - rows=rows, - cols=cols, - figsize=figsize, - equal=True, - ) + fig, ax = _format_ax(ax, rows=rows, cols=cols, figsize=figsize, equal=True) ax = np.atleast_1d(ax).flatten() plot_data = {} @@ -2058,9 +2045,7 @@ def plot_coils(coils, grid=None, ax=None, return_data=False, **kwargs): color = [color] fig, ax = _format_ax(ax, True, figsize=figsize) if grid is None: - grid_kwargs = { - "zeta": np.linspace(0, 2 * np.pi, 50), - } + grid_kwargs = {"zeta": np.linspace(0, 2 * np.pi, 50)} grid = _get_grid(**grid_kwargs) def flatten_coils(coilset): @@ -2371,20 +2356,21 @@ def plot_boozer_surface( with warnings.catch_warnings(): warnings.simplefilter("ignore") data = eq.compute("|B|_mn", grid=grid_compute, transforms=transforms_compute) - iota = compress(grid_compute, data["iota"]) + iota = grid_compute.compress(data["iota"]) data = transforms_plot["B"].transform(data["|B|_mn"]) data = data.reshape((grid_plot.num_theta, grid_plot.num_zeta), order="F") fig, ax = _format_ax(ax, figsize=kwargs.pop("figsize", None)) divider = make_axes_locatable(ax) - contourf_kwargs = {} - contourf_kwargs["norm"] = matplotlib.colors.Normalize() - contourf_kwargs["levels"] = kwargs.pop( - "levels", np.linspace(np.nanmin(data), np.nanmax(data), ncontours) - ) - contourf_kwargs["cmap"] = kwargs.pop("cmap", "jet") - contourf_kwargs["extend"] = "both" + contourf_kwargs = { + "norm": matplotlib.colors.Normalize(), + "levels": kwargs.pop( + "levels", np.linspace(np.nanmin(data), np.nanmax(data), ncontours) + ), + "cmap": kwargs.pop("cmap", "jet"), + "extend": "both", + } assert ( len(kwargs) == 0 @@ -2429,10 +2415,7 @@ def plot_boozer_surface( ax.set_title(r"$|\mathbf{B}|~(T)$", fontsize=title_fontsize) _set_tight_layout(fig) - plot_data = {} - plot_data["zeta_Boozer"] = zz - plot_data["theta_Boozer"] = tt - plot_data["|B|"] = data + plot_data = {"zeta_Boozer": zz, "theta_Boozer": tt, "|B|": data} if return_data: return fig, ax, plot_data @@ -2547,7 +2530,6 @@ def plot_qs_error( # noqa: 16 fxn too complex R0 = data["R0"] B0 = np.mean(data["|B|"] * data["sqrt(g)"]) / np.mean(data["sqrt(g)"]) - data = None f_B = np.array([]) f_C = np.array([]) f_T = np.array([]) @@ -2730,7 +2712,7 @@ def plot_grid(grid, return_data=False, **kwargs): ), f"plot_grid got unexpected keyword argument: {kwargs.keys()}" # node locations - nodes = grid.nodes[np.where(grid.nodes[:, 2] == 0)] + nodes = grid.nodes[grid.nodes[:, 2] == 0] ax.scatter(nodes[:, 1], nodes[:, 0], s=4) ax.set_ylim(0, 1) ax.set_xticks( @@ -2777,9 +2759,7 @@ def plot_grid(grid, return_data=False, **kwargs): ) _set_tight_layout(fig) - plot_data = {} - plot_data["rho"] = nodes[:, 0] - plot_data["theta"] = nodes[:, 1] + plot_data = {"rho": nodes[:, 0], "theta": nodes[:, 1]} if return_data: return fig, ax, plot_data @@ -2834,16 +2814,14 @@ def plot_basis(basis, return_data=False, **kwargs): title_fontsize = kwargs.pop("title_fontsize", None) if basis.__class__.__name__ == "PowerSeries": + # fixme: lmax unused lmax = abs(basis.modes[:, 0]).max() grid = LinearGrid(rho=100, endpoint=True) r = grid.nodes[:, 0] fig, ax = plt.subplots(figsize=kwargs.get("figsize", (6, 4))) f = basis.evaluate(grid.nodes) - plot_data = {} - plot_data["l"] = basis.modes[:, 0] - plot_data["amplitude"] = [] - plot_data["rho"] = r + plot_data = {"l": basis.modes[:, 0], "amplitude": [], "rho": r} for fi, l in zip(f.T, basis.modes[:, 0]): ax.plot(r, fi, label="$l={:d}$".format(int(l))) @@ -2864,6 +2842,7 @@ def plot_basis(basis, return_data=False, **kwargs): return fig, ax elif basis.__class__.__name__ == "FourierSeries": + # fixme nmax unused nmax = abs(basis.modes[:, 2]).max() grid = LinearGrid(zeta=100, NFP=basis.NFP, endpoint=True) z = grid.nodes[:, 2] @@ -2974,16 +2953,12 @@ def plot_basis(basis, return_data=False, **kwargs): mmax = abs(basis.modes[:, 1]).max().astype(int) grid = LinearGrid(rho=100, theta=100, endpoint=True) - r = np.unique(grid.nodes[:, 0]) - v = np.unique(grid.nodes[:, 1]) + r = grid.nodes[grid.unique_rho_idx, 0] + v = grid.nodes[grid.unique_theta_idx, 1] fig = plt.figure(figsize=kwargs.get("figsize", (3 * mmax, 3 * lmax / 2))) - plot_data = {} - - plot_data["amplitude"] = [] - plot_data["rho"] = r - plot_data["theta"] = v + plot_data = {"amplitude": [], "rho": r, "theta": v} ax = {i: {} for i in range(lmax + 1)} ratios = np.ones(2 * (mmax + 1) + 1) @@ -2992,7 +2967,7 @@ def plot_basis(basis, return_data=False, **kwargs): lmax + 2, 2 * (mmax + 1) + 1, width_ratios=ratios ) - modes = basis.modes[np.where(basis.modes[:, 2] == 0)] + modes = basis.modes[basis.modes[:, 2] == 0] plot_data["l"] = basis.modes[:, 0] plot_data["m"] = basis.modes[:, 1] Zs = basis.evaluate(grid.nodes, modes=modes) @@ -3338,6 +3313,7 @@ def plot_field_lines_sfl( ) """ + # TODO: can this be removed now? if rho == 0: raise NotImplementedError( "Currently does not support field line tracing of the magnetic axis, " diff --git a/desc/profiles.py b/desc/profiles.py index c6db8b407a..26169d22b9 100644 --- a/desc/profiles.py +++ b/desc/profiles.py @@ -575,9 +575,10 @@ class PowerSeriesProfile(Profile): Parameters ---------- params: array-like - Coefficients of the series. If modes is not supplied, assumed to be in ascending - order with no missing values. If modes is given, coefficients can be in any - order or indexing. + Coefficients of the series. Assumed to be zero if not specified. + If modes is not supplied, assumed to be in ascending order with no + missing values. If modes is given, coefficients can be in any order or + indexing. modes : array-like Mode numbers for the associated coefficients. eg a[modes[i]] = params[i] grid : Grid @@ -590,9 +591,11 @@ class PowerSeriesProfile(Profile): _io_attrs_ = Profile._io_attrs_ + ["_basis", "_transform"] - def __init__(self, params=[0], modes=None, grid=None, sym="auto", name=""): + def __init__(self, params=None, modes=None, grid=None, sym="auto", name=""): super().__init__(grid, name) + if params is None: + params = [0] params = np.atleast_1d(params) if sym == "auto": # sym = "even" if all odd modes are zero, else sym = False @@ -629,7 +632,7 @@ def _get_transform(self, grid): transform = Transform( grid, self.basis, - derivs=np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]]), + derivs=np.array([[i, 0, 0] for i in range(5)]), build=True, ) return transform @@ -799,11 +802,11 @@ class SplineProfile(Profile): _io_attrs_ = Profile._io_attrs_ + ["_knots", "_method"] - def __init__( - self, values=[0, 0, 0], knots=None, grid=None, method="cubic2", name="" - ): + def __init__(self, values=None, knots=None, grid=None, method="cubic2", name=""): super().__init__(grid, name) + if values is None: + values = [0, 0, 0] values = np.atleast_1d(values) if knots is None: knots = np.linspace(0, 1, values.size) @@ -914,9 +917,11 @@ class MTanhProfile(Profile): """ - def __init__(self, params=[0, 0, 1, 1, 0], grid=None, name=""): + def __init__(self, params=None, grid=None, name=""): super().__init__(grid, name) + if params is None: + params = [0, 0, 1, 1, 0] self._params = params def __repr__(self): @@ -1165,9 +1170,11 @@ class FourierZernikeProfile(Profile): _io_attrs_ = Profile._io_attrs_ + ["_basis", "_transform"] - def __init__(self, params=[0], modes=None, grid=None, sym="auto", NFP=1, name=""): + def __init__(self, params=None, modes=None, grid=None, sym="auto", NFP=1, name=""): super().__init__(grid, name) + if params is None: + params = [0] params = np.atleast_1d(params) if modes is None: @@ -1211,7 +1218,7 @@ def _get_transform(self, grid): transform = Transform( grid, self.basis, - derivs=np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]]), + derivs=np.array([[i, 0, 0] for i in range(5)]), build=True, ) return transform diff --git a/desc/transform.py b/desc/transform.py index dd2ab14eb3..945c8d1fb4 100644 --- a/desc/transform.py +++ b/desc/transform.py @@ -106,9 +106,9 @@ def _get_derivatives(self, derivs): """ if isinstance(derivs, int) and derivs >= 0: derivatives = combination_permutation(3, derivs, False) - elif np.atleast_1d(derivs).ndim == 1 and len(derivs) == 3: + elif np.ndim(derivs) == 1 and len(derivs) == 3: derivatives = np.asarray(derivs).reshape((1, 3)) - elif np.atleast_2d(derivs).ndim == 2 and np.atleast_2d(derivs).shape[1] == 3: + elif np.ndim(derivs) == 2 and np.atleast_2d(derivs).shape[1] == 3: derivatives = np.atleast_2d(derivs) else: raise NotImplementedError( diff --git a/desc/utils.py b/desc/utils.py index 2a2b4d6855..0512a2f8e0 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -496,11 +496,7 @@ def setdefault(val, default, cond=None): If cond is None, then it checks if val is not None, returning val or default accordingly. """ - if cond is None: - cond = val is not None - if cond: - return val - return default + return val if cond or (cond is None and val is not None) else default def isnonnegint(x): diff --git a/desc/vmec.py b/desc/vmec.py index d3714cfb7d..79e8aebd77 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -11,7 +11,7 @@ from desc.basis import DoubleFourierSeries from desc.compat import ensure_positive_jacobian -from desc.compute.utils import compress, surface_averages +from desc.compute.utils import surface_averages from desc.equilibrium import Equilibrium from desc.geometry import FourierRZToroidalSurface from desc.grid import Grid, LinearGrid @@ -431,9 +431,8 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify if eq.iota is not None: iotaf[:] = -eq.iota(r_full) # negative sign for negative Jacobian else: - # value closest to axis will be nan grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=r_full, NFP=NFP) - iotaf[:] = -compress(grid, eq.compute("iota", grid=grid)["iota"]) + iotaf[:] = -grid.compress(eq.compute("iota", grid=grid)["iota"]) q_factor = file.createVariable("q_factor", np.float64, ("radius",)) q_factor.long_name = "inverse rotational transform on full mesh" @@ -448,7 +447,7 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify iotas[1:] = -eq.iota(r_half) # negative sign for negative Jacobian else: grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=r_half, NFP=NFP) - iotas[1:] = -compress(grid, eq.compute("iota", grid=grid)["iota"]) + iotas[1:] = -grid.compress(eq.compute("iota", grid=grid)["iota"]) phi = file.createVariable("phi", np.float64, ("radius",)) phi.long_name = "toroidal flux" @@ -545,33 +544,42 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify # grid for computing radial profile data grid = LinearGrid(M=eq.M_grid, N=eq.M_grid, NFP=eq.NFP, sym=eq.sym, rho=r_full) data = eq.compute( - ["", "I", "G", "", "sqrt(g)", "J^theta", "J^zeta", "D_Mercier"], + [ + "<|B|^2>", + "I", + "G", + "", + "sqrt(g)", + "J^theta*sqrt(g)", + "J^zeta", + "D_Mercier", + ], grid=grid, ) bdotb = file.createVariable("bdotb", np.float64, ("radius",)) bdotb.long_name = "flux surface average of magnetic field squared" bdotb.units = "T^2" - bdotb[:] = compress(grid, data[""]) + bdotb[:] = grid.compress(data["<|B|^2>"]) bdotb[0] = 0 # currents buco = file.createVariable("buco", np.float64, ("radius",)) buco.long_name = "Boozer toroidal current I" buco.units = "T*m" - buco[:] = compress(grid, data["I"]) + buco[:] = grid.compress(data["I"]) buco[0] = 0 bvco = file.createVariable("bvco", np.float64, ("radius",)) bvco.long_name = "Boozer poloidal current G" bvco.units = "T*m" - bvco[:] = compress(grid, data["G"]) + bvco[:] = grid.compress(data["G"]) bvco[0] = 0 jdotb = file.createVariable("jdotb", np.float64, ("radius",)) jdotb.long_name = "flux surface average of J*B" jdotb.units = "N/m^3" - jdotb[:] = compress(grid, data[""]) + jdotb[:] = grid.compress(data[""]) jdotb[0] = 0 jcuru = file.createVariable("jcuru", np.float64, ("radius",)) @@ -579,7 +587,7 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify jcuru.units = "A/m^3" jcuru[:] = surface_averages( grid, - data["sqrt(g)"] * data["J^theta"] / (2 * data["rho"]), + data["J^theta*sqrt(g)"] / (2 * data["rho"]), sqrt_g=data["sqrt(g)"], expand_out=False, ) @@ -600,38 +608,38 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify DShear = file.createVariable("DShear", np.float64, ("radius",)) DShear.long_name = "Mercier stability criterion magnetic shear term" DShear.units = "1/Wb^2" - DShear[:] = compress(grid, data["D_shear"]) + DShear[:] = grid.compress(data["D_shear"]) DShear[0] = 0 DCurr = file.createVariable("DCurr", np.float64, ("radius",)) DCurr.long_name = "Mercier stability criterion toroidal current term" DCurr.units = "1/Wb^2" - DCurr[:] = compress(grid, data["D_current"]) + DCurr[:] = grid.compress(data["D_current"]) DCurr[0] = 0 DWell = file.createVariable("DWell", np.float64, ("radius",)) DWell.long_name = "Mercier stability criterion magnetic well term" DWell.units = "1/Wb^2" - DWell[:] = compress(grid, data["D_well"]) + DWell[:] = grid.compress(data["D_well"]) DWell[0] = 0 DGeod = file.createVariable("DGeod", np.float64, ("radius",)) DGeod.long_name = "Mercier stability criterion geodesic curvature term" DGeod.units = "1/Wb^2" - DGeod[:] = compress(grid, data["D_geodesic"]) + DGeod[:] = grid.compress(data["D_geodesic"]) DGeod[0] = 0 DMerc = file.createVariable("DMerc", np.float64, ("radius",)) DMerc.long_name = "Mercier stability criterion" DMerc.units = "1/Wb^2" - DMerc[:] = compress(grid, data["D_Mercier"]) + DMerc[:] = grid.compress(data["D_Mercier"]) DMerc[0] = 0 timer.stop("parameters") if verbose > 1: timer.disp("parameters") - # indepentent variables (exact conversion) + # independent variables (exact conversion) # R axis idx = np.where(eq.R_basis.modes[:, 1] == 0)[0] @@ -785,12 +793,14 @@ def fullfit(x): # half grid quantities half_grid = LinearGrid(M=M_nyq, N=N_nyq, NFP=NFP, rho=r_half) data_half_grid = eq.compute( - ["J", "|B|", "B_rho", "B_theta", "B_zeta"], grid=half_grid + ["J", "|B|", "B_rho", "B_theta", "B_zeta", "sqrt(g)"], grid=half_grid ) # full grid quantities full_grid = LinearGrid(M=M_nyq, N=N_nyq, NFP=NFP, rho=r_full) - data_full_grid = eq.compute(["J", "B_rho", "B_theta", "B_zeta"], grid=full_grid) + data_full_grid = eq.compute( + ["J", "B_rho", "B_theta", "B_zeta", "J^theta*sqrt(g)"], grid=full_grid + ) # Jacobian timer.start("Jacobian") @@ -993,7 +1003,7 @@ def fullfit(x): bsubsmns[0, :] = ( # linear extrapolation for coefficient at the magnetic axis s[1, :] - (s[2, :] - s[1, :]) / (s_full[2] - s_full[1]) * s_full[1] ) - # TODO: evaluate current at rho=0 nodes instead of extrapolation + # Todo: evaluate current at rho=0 nodes instead of extrapolation if not eq.sym: bsubsmnc[:, :] = c bsubsmnc[0, :] = ( @@ -1089,10 +1099,9 @@ def fullfit(x): if verbose > 1: timer.disp("B_zeta") - # J^theta * sqrt(g) # noqa: E800 - timer.start("J^theta") + timer.start("J^theta*sqrt(g)") if verbose > 0: - print("Saving J^theta") + print("Saving J^theta*sqrt(g)") currumnc = file.createVariable( "currumnc", np.float64, ("radius", "mn_mode_nyq") ) @@ -1111,11 +1120,7 @@ def fullfit(x): m = full_basis.modes[:, 1] n = full_basis.modes[:, 2] data = ( - ( - data_full_grid["J^theta"] - * data_full_grid["sqrt(g)"] - / (2 * data_full_grid["rho"]) - ) + (data_full_grid["J^theta*sqrt(g)"] / (2 * data_full_grid["rho"])) .reshape( (full_grid.num_theta, full_grid.num_rho, full_grid.num_zeta), order="F" ) @@ -1133,20 +1138,19 @@ def fullfit(x): currumnc[0, :] = ( # linear extrapolation for coefficient at the magnetic axis s[1, :] - (c[2, :] - c[1, :]) / (s_full[2] - s_full[1]) * s_full[1] ) - # TODO: evaluate current at rho=0 nodes instead of extrapolation + # Todo: evaluate current at rho=0 nodes instead of extrapolation if not eq.sym: currumns[:, :] = s currumns[0, :] = ( s[1, :] - (s[2, :] - s[1, :]) / (s_full[2] - s_full[1]) * s_full[1] ) - timer.stop("J^theta") + timer.stop("J^theta*sqrt(g)") if verbose > 1: - timer.disp("J^theta") + timer.disp("J^theta*sqrt(g)") - # J^zeta * sqrt(g) # noqa: E800 - timer.start("J^zeta") + timer.start("J^zeta*sqrt(g)") if verbose > 0: - print("Saving J^zeta") + print("Saving J^zeta*sqrt(g)") currvmnc = file.createVariable( "currvmnc", np.float64, ("radius", "mn_mode_nyq") ) @@ -1187,22 +1191,22 @@ def fullfit(x): currvmnc[0, :] = -( # linear extrapolation for coefficient at the magnetic axis s[1, :] - (c[2, :] - c[1, :]) / (s_full[2] - s_full[1]) * s_full[1] ) - # TODO: evaluate current at rho=0 nodes instead of extrapolation + # Todo: evaluate current at rho=0 nodes instead of extrapolation if not eq.sym: currvmns[:, :] = -s currumns[0, :] = -( s[1, :] - (s[2, :] - s[1, :]) / (s_full[2] - s_full[1]) * s_full[1] ) - timer.stop("J^zeta") + timer.stop("J^zeta*sqrt(g)") if verbose > 1: - timer.disp("J^zeta") + timer.disp("J^zeta*sqrt(g)") # TODO: these output quantities need to be added bdotgradv = file.createVariable("bdotgradv", np.float64, ("radius",)) bdotgradv[:] = np.zeros((file.dimensions["radius"].size,)) bdotgradv.long_name = "Not Implemented: This output is hard-coded to 0!" bdotgradv.units = "None" - # beta_vol = something like p/(B^2-p) ? It's not _vol(r) + # beta_vol = something like p/(|B|^2 - p) ? It's not _vol(r) beta_vol = file.createVariable("beta_vol", np.float64, ("radius",)) beta_vol[:] = np.zeros((file.dimensions["radius"].size,)) beta_vol.long_name = "Not Implemented: This output is hard-coded to 0!" diff --git a/docs/adding_compute_funs.rst b/docs/adding_compute_funs.rst index a2d2cf18c4..651d4828c5 100644 --- a/docs/adding_compute_funs.rst +++ b/docs/adding_compute_funs.rst @@ -4,8 +4,8 @@ Adding new physics quantities All calculation of physics quantities takes place in ``desc.compute`` -As an example, we'll walk through the calculation of the radial component of the MHD -force :math:`F_\rho` +As an example, we'll walk through the calculation of the contravariant radial +component of the plasma current density :math:`J^\rho` The full code is below: :: @@ -13,76 +13,93 @@ The full code is below: from desc.data_index import register_compute_fun @register_compute_fun( - name="F_rho", - label="F_{\\rho}", - units="N \\cdot m^{-2}", - units_long="Newtons / square meter", - description="Covariant radial component of force balance error", + name="J^rho", + label="J^{\\rho}", + units="A \\cdot m^{-3}", + units_long="Amperes / cubic meter", + description="Contravariant radial component of plasma current density", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["p_r", "sqrt(g)", "B^theta", "B^zeta", "J^theta", "J^zeta"], - parameterization="desc.equilibrium.Equilibrium" + data=["sqrt(g)", "B_zeta_t", "B_theta_z"], + axis_limit_data=["sqrt(g)_r", "B_zeta_rt", "B_theta_rz"], + parameterization="desc.equilibrium.equilibrium.Equilibrium", ) - def _F_rho(params, transforms, profiles, data, **kwargs): - data["F_rho"] = -data["p_r"] + data["sqrt(g)"] * ( - data["B^zeta"] * data["J^theta"] - data["B^theta"] * data["J^zeta"] - ) + def _J_sup_rho(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, + # ∂_θ (𝐁 ⋅ 𝐞_ζ) - ∂_ζ (𝐁 ⋅ 𝐞_θ) = 𝐁 ⋅ (∂_θ 𝐞_ζ - ∂_ζ 𝐞_θ) = 0 + # because the partial derivatives commute. So 𝐉^ρ is of the indeterminate + # form 0/0 and we may compute the limit as follows. + data["J^rho"] = ( + transforms["grid"].replace_at_axis( + (data["B_zeta_t"] - data["B_theta_z"]) / data["sqrt(g)"], + lambda: (data["B_zeta_rt"] - data["B_theta_rz"]) / data["sqrt(g)_r"], + ) + ) / mu_0 return data The decorator ``register_compute_fun`` tells DESC that the quantity exists and contains -metadata about the quanity. The necessary fields are detailed below: +metadata about the quantity. The necessary fields are detailed below: -* ``name``: A short, meaningfull name that is used elsewhere in the code to refer to the - quanity. This name will appear in the ``data`` dictionary returned by ``Equilibrium.compute``, - and is also the argument passed to ``compute`` to calculate the quanity. IE, - ``Equilibrium.compute("F_rho")`` will return a dictionary containing ``F_rho`` as well +* ``name``: A short, meaningful name that is used elsewhere in the code to refer to the + quantity. This name will appear in the ``data`` dictionary returned by ``Equilibrium.compute``, + and is also the argument passed to ``compute`` to calculate the quantity. IE, + ``Equilibrium.compute("J^rho")`` will return a dictionary containing ``J^rho`` as well as all the intermediate quantities needed to compute it. General conventions are that covariant components of a vector are called ``X_rho`` etc, contravariant components ``X^rho`` etc, and derivatives by a single character subscript, ``X_r`` etc for :math:`\partial_{\rho} X` * ``label``: a LaTeX style label used for plotting. -* ``units``: SI units of the quanity in LaTeX format -* ``units_long``: SI units of the quanity, spelled out +* ``units``: SI units of the quantity in LaTeX format +* ``units_long``: SI units of the quantity, spelled out * ``description``: A short description of the quantity * ``dim``: Dimensionality of the quantity. Vectors (such as magnetic field) have ``dim=3``, local scalar quantities (such as vector components, pressure, volume element, etc) have ``dim=1``, global scalars (such as total volume, aspect ratio, etc) have ``dim=0`` -* ``params``: list of strings of ``Equilibrium`` parameters needed to compute the quanity +* ``params``: list of strings of ``Equilibrium`` parameters needed to compute the quantity such as ``R_lmn``, ``Z_lmn`` etc. These will be passed into the compute function as a dictionary in the ``params`` argument. Note that this only includes direct dependencies (things that are used in this function). For most quantities, this will be an empty list. For example, if the function relies on ``R_t``, this dependency should be specified as a - data dependecy (see below), while the function to compute ``R_t`` itself will depend on + data dependency (see below), while the function to compute ``R_t`` itself will depend on ``R_lmn`` * ``transforms``: a dictionary of what ``transform`` objects are needed, with keys being the quantity being transformed (``R``, ``p``, etc) and the values are a list of derivative - orders needed in ``rho``, ``theta``, ``zeta``. IE, if the quanity requires + orders needed in ``rho``, ``theta``, ``zeta``. IE, if the quantity requires :math:`R_{\rho}{\zeta}{\zeta}`, then ``transforms`` should be ``{"R": [[1, 0, 2]]}`` indicating a first derivative in ``rho`` and a second derivative in ``zeta``. Note that this only includes direct dependencies (things that are used in this function). For most - quantites this will be an empty dictionary. + quantities this will be an empty dictionary. * ``profiles``: List of string of ``Profile`` quantities needed, such as ``pressure`` or ``iota``. Note that this only includes direct dependencies (things that are used in - this function). For most quantites this will be an empty list. -* ``coordinates``: String denoting which coordinate the quanity depends on. Most will be - ``"rtz"`` indicating it is a funciton of :math:`\rho, \theta, \zeta`. Profiles and flux surface + this function). For most quantities this will be an empty list. +* ``coordinates``: String denoting which coordinate the quantity depends on. Most will be + ``"rtz"`` indicating it is a function of :math:`\rho, \theta, \zeta`. Profiles and flux surface quantities would have ``coordinates="r"`` indicating it only depends on :math:`\rho` -* ``data``: What other physics quantites are needed to compute this quanity. In our - example, we need the radial derivative of pressure ``p_r``, the Jacobian determinant - ``sqrt(g)``, and contravariant components of current and magnetic field. These dependencies - will be passed to the compute function as a dictionary in the ``data`` argument. Note - that this only includes direct dependencies (things that are used in this function). - For example, we need ``sqrt(g)``, which itself depends on ``e_rho``, etc. But we don'take - need to specify ``e_rho`` here, that dependency is determined automatically at runtime. +* ``data``: What other physics quantities are needed to compute this quantity. In our + example, we need the poloidal (theta) derivative of the covariant toroidal (zeta) component + of the magnetic field ``B_zeta_t``, the toroidal derivative of the covariant poloidal + component of the magnetic field ``B_theta_z``, and the 3-D volume Jacobian determinant + ``sqrt(g)``. These dependencies will be passed to the compute function as a dictionary + in the ``data`` argument. Note that this only includes direct dependencies (things that + are used in this function). For example, we need ``sqrt(g)``, which itself depends on + ``e_rho``, etc. But we don't need to specify those sub-dependencies here. +* ``axis_limit_data``: Some quantities require additional work to compute at the + magnetic axis. A Python lambda function is used to lazily compute the magnetic + axis limits of these quantities. These lambda functions are evaluated only when + the computational grid has a node on the magnetic axis to avoid potentially + expensive computations. In our example, we need the radial derivatives of each + the quantities mentioned above to evaluate the magnetic axis limit. These dependencies + are specified in ``axis_limit_data``. The dependencies specified in this list are + marked to be computed only when there is a node at the magnetic axis. * ``parameterization``: what sorts of DESC objects is this function for. Most functions will just be for ``Equilibrium``, but some methods may also be for ``desc.geometry.Curve``, or specific types eg ``desc.geometry.FourierRZCurve``. * ``kwargs``: If the compute function requires any additional arguments they should be specified like ``kwarg="thing"`` where the value is the name of the keyword argument - that will be passed to the compute function. Most quantites do not take kwargs. + that will be passed to the compute function. Most quantities do not take kwargs. The function itself should have a signature of the form diff --git a/docs/write_optimizers.py b/docs/write_optimizers.py index 9b868b1565..23f51be9ba 100644 --- a/docs/write_optimizers.py +++ b/docs/write_optimizers.py @@ -21,15 +21,15 @@ writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") writer.writeheader() - keys = optimizers.keys() - for key in keys: - d = {} - d["Name"] = "``" + key + "``" - d["Scalar"] = optimizers[key]["scalar"] - d["Equality Constraints"] = optimizers[key]["equality_constraints"] - d["Inequality Constraints"] = optimizers[key]["inequality_constraints"] - d["Stochastic"] = optimizers[key]["stochastic"] - d["Hessian"] = optimizers[key]["hessian"] - d["GPU"] = optimizers[key]["GPU"] - d["Description"] = optimizers[key]["description"] + for key, val in optimizers.items(): + d = { + "Name": "``" + key + "``", + "Scalar": val["scalar"], + "Equality Constraints": val["equality_constraints"], + "Inequality Constraints": val["inequality_constraints"], + "Stochastic": val["stochastic"], + "Hessian": val["hessian"], + "GPU": val["GPU"], + "Description": val["description"], + } writer.writerow(d) diff --git a/docs/write_variables.py b/docs/write_variables.py index 2887f8167f..3c22784509 100644 --- a/docs/write_variables.py +++ b/docs/write_variables.py @@ -26,12 +26,13 @@ def write_csv(parameterization): datidx = data_index[parameterization] keys = datidx.keys() for key in keys: - d = {} - d["Name"] = "``" + key + "``" - d["Label"] = ":math:`" + datidx[key]["label"].replace("$", "") + "`" - d["Units"] = datidx[key]["units_long"] - d["Description"] = datidx[key]["description"] - d["Module"] = "``" + datidx[key]["fun"].__module__ + "``" + d = { + "Name": "``" + key + "``", + "Label": ":math:`" + datidx[key]["label"].replace("$", "") + "`", + "Units": datidx[key]["units_long"], + "Description": datidx[key]["description"], + "Module": "``" + datidx[key]["fun"].__module__ + "``", + } # stuff like |x| is interpreted as a substitution by rst, need to escape d["Description"] = _escape(d["Description"]) diff --git a/tests/baseline/test_1d_dpdr.png b/tests/baseline/test_1d_dpdr.png index 9a55a46373..d949bc1d93 100644 Binary files a/tests/baseline/test_1d_dpdr.png and b/tests/baseline/test_1d_dpdr.png differ diff --git a/tests/baseline/test_1d_iota.png b/tests/baseline/test_1d_iota.png index b0a9af7faa..83cf2ec3cf 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 455e4826cd..6137343a14 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 fd5f1cfa7c..6fbb627e9d 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 a8b54c0243..cfc2704eb1 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 9ab6156378..01bb27a699 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 9643807e72..ab2dceea24 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 index ec7949dd2d..c81798568e 100644 Binary files a/tests/baseline/test_2d_lambda.png and b/tests/baseline/test_2d_lambda.png differ diff --git a/tests/baseline/test_2d_logF.png b/tests/baseline/test_2d_logF.png index 07f3f85288..94089aaf86 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_3d_B.png b/tests/baseline/test_3d_B.png index f46d30ad43..99d6ee866b 100644 Binary files a/tests/baseline/test_3d_B.png and b/tests/baseline/test_3d_B.png differ diff --git a/tests/baseline/test_3d_J.png b/tests/baseline/test_3d_J.png index 71bfca4a73..e38133335b 100644 Binary files a/tests/baseline/test_3d_J.png and b/tests/baseline/test_3d_J.png differ diff --git a/tests/baseline/test_3d_rt.png b/tests/baseline/test_3d_rt.png index 14672aa6a6..d9e2ed191c 100644 Binary files a/tests/baseline/test_3d_rt.png and b/tests/baseline/test_3d_rt.png differ diff --git a/tests/baseline/test_3d_rz.png b/tests/baseline/test_3d_rz.png index 602347f21a..8ba4e6eb5d 100644 Binary files a/tests/baseline/test_3d_rz.png and b/tests/baseline/test_3d_rz.png differ diff --git a/tests/baseline/test_3d_tz.png b/tests/baseline/test_3d_tz.png index cc117f4717..dea5cd2b02 100644 Binary files a/tests/baseline/test_3d_tz.png and b/tests/baseline/test_3d_tz.png differ diff --git a/tests/baseline/test_Redl_figures_2_3.png b/tests/baseline/test_Redl_figures_2_3.png index 2f2d172422..525d569f43 100644 Binary files a/tests/baseline/test_Redl_figures_2_3.png and b/tests/baseline/test_Redl_figures_2_3.png differ diff --git a/tests/baseline/test_Redl_figures_4_5.png b/tests/baseline/test_Redl_figures_4_5.png index 7dd08b393c..52fbbbc549 100644 Binary files a/tests/baseline/test_Redl_figures_4_5.png and b/tests/baseline/test_Redl_figures_4_5.png differ diff --git a/tests/baseline/test_Redl_sfincs_QA.png b/tests/baseline/test_Redl_sfincs_QA.png index 978a2567d4..4b2559cc7d 100644 Binary files a/tests/baseline/test_Redl_sfincs_QA.png and b/tests/baseline/test_Redl_sfincs_QA.png differ diff --git a/tests/baseline/test_Redl_sfincs_QH.png b/tests/baseline/test_Redl_sfincs_QH.png index ce062a9ce5..15798d91d9 100644 Binary files a/tests/baseline/test_Redl_sfincs_QH.png and b/tests/baseline/test_Redl_sfincs_QH.png differ diff --git a/tests/baseline/test_Redl_sfincs_tokamak_benchmark.png b/tests/baseline/test_Redl_sfincs_tokamak_benchmark.png index 90a735aad9..3184aea467 100644 Binary files a/tests/baseline/test_Redl_sfincs_tokamak_benchmark.png and b/tests/baseline/test_Redl_sfincs_tokamak_benchmark.png differ diff --git a/tests/baseline/test_fsa_F_normalized.png b/tests/baseline/test_fsa_F_normalized.png index 8e1e40f7c1..d25f683891 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 index 197c5616cb..4f5362354f 100644 Binary files a/tests/baseline/test_fsa_G.png and b/tests/baseline/test_fsa_G.png differ diff --git a/tests/baseline/test_fsa_I.png b/tests/baseline/test_fsa_I.png index 9b814b0928..c317fa5f84 100644 Binary files a/tests/baseline/test_fsa_I.png and b/tests/baseline/test_fsa_I.png differ diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png index 6686c58672..fc2feb35b9 100644 Binary files a/tests/baseline/test_plot_b_mag.png and b/tests/baseline/test_plot_b_mag.png differ diff --git a/tests/baseline/test_plot_basis_doublefourierseries.png b/tests/baseline/test_plot_basis_doublefourierseries.png index 5f50a8fdbf..b1658d8dc6 100644 Binary files a/tests/baseline/test_plot_basis_doublefourierseries.png and b/tests/baseline/test_plot_basis_doublefourierseries.png differ diff --git a/tests/baseline/test_plot_basis_fourierseries.png b/tests/baseline/test_plot_basis_fourierseries.png index c214eb9f99..fa5dfddb44 100644 Binary files a/tests/baseline/test_plot_basis_fourierseries.png and b/tests/baseline/test_plot_basis_fourierseries.png differ diff --git a/tests/baseline/test_plot_basis_fourierzernike.png b/tests/baseline/test_plot_basis_fourierzernike.png index 6680fc5745..5b522db51b 100644 Binary files a/tests/baseline/test_plot_basis_fourierzernike.png and b/tests/baseline/test_plot_basis_fourierzernike.png differ diff --git a/tests/baseline/test_plot_basis_powerseries.png b/tests/baseline/test_plot_basis_powerseries.png index 158a54d924..605b227683 100644 Binary files a/tests/baseline/test_plot_basis_powerseries.png and b/tests/baseline/test_plot_basis_powerseries.png differ diff --git a/tests/baseline/test_plot_boozer_modes.png b/tests/baseline/test_plot_boozer_modes.png index c9c78ee11d..6eff45d707 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_surface.png b/tests/baseline/test_plot_boozer_surface.png index 9d682674c8..34bdaa475f 100644 Binary files a/tests/baseline/test_plot_boozer_surface.png and b/tests/baseline/test_plot_boozer_surface.png differ diff --git a/tests/baseline/test_plot_boundaries.png b/tests/baseline/test_plot_boundaries.png index a8281325b9..6bbdb926bc 100644 Binary files a/tests/baseline/test_plot_boundaries.png and b/tests/baseline/test_plot_boundaries.png differ diff --git a/tests/baseline/test_plot_boundary.png b/tests/baseline/test_plot_boundary.png index 5bc17b7779..b19adb53ee 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 161f5a3e1d..f0bfc75d93 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_coils.png b/tests/baseline/test_plot_coils.png index 5acd6e15b1..6d1f12594a 100644 Binary files a/tests/baseline/test_plot_coils.png and b/tests/baseline/test_plot_coils.png differ diff --git a/tests/baseline/test_plot_comparison.png b/tests/baseline/test_plot_comparison.png index 6e7fcd2cc2..bf6cb0fdf1 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 0980cb29e9..be62236776 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_con_basis.png b/tests/baseline/test_plot_con_basis.png index db092a9f15..5e92a57c90 100644 Binary files a/tests/baseline/test_plot_con_basis.png and b/tests/baseline/test_plot_con_basis.png differ diff --git a/tests/baseline/test_plot_cov_basis.png b/tests/baseline/test_plot_cov_basis.png index 8ed63125e6..e8d9c94d5c 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 index 648f766f06..5d51729e83 100644 Binary files a/tests/baseline/test_plot_field_line.png and b/tests/baseline/test_plot_field_line.png differ diff --git a/tests/baseline/test_plot_field_lines.png b/tests/baseline/test_plot_field_lines.png index 528cdbc9ec..130c24bd6d 100644 Binary files a/tests/baseline/test_plot_field_lines.png and b/tests/baseline/test_plot_field_lines.png differ diff --git a/tests/baseline/test_plot_gradpsi.png b/tests/baseline/test_plot_gradpsi.png index 7bcb4777a5..bfce808654 100644 Binary files a/tests/baseline/test_plot_gradpsi.png and b/tests/baseline/test_plot_gradpsi.png differ diff --git a/tests/baseline/test_plot_grid_cheb1.png b/tests/baseline/test_plot_grid_cheb1.png index 6acb499bbc..76da79e7cc 100644 Binary files a/tests/baseline/test_plot_grid_cheb1.png and b/tests/baseline/test_plot_grid_cheb1.png differ diff --git a/tests/baseline/test_plot_grid_cheb2.png b/tests/baseline/test_plot_grid_cheb2.png index 8c2df30d35..ce56f802bd 100644 Binary files a/tests/baseline/test_plot_grid_cheb2.png and b/tests/baseline/test_plot_grid_cheb2.png differ diff --git a/tests/baseline/test_plot_grid_jacobi.png b/tests/baseline/test_plot_grid_jacobi.png index c0a14442f9..9a07db6de2 100644 Binary files a/tests/baseline/test_plot_grid_jacobi.png and b/tests/baseline/test_plot_grid_jacobi.png differ diff --git a/tests/baseline/test_plot_grid_linear.png b/tests/baseline/test_plot_grid_linear.png index 5418c649ac..3fc34b7056 100644 Binary files a/tests/baseline/test_plot_grid_linear.png and b/tests/baseline/test_plot_grid_linear.png differ diff --git a/tests/baseline/test_plot_grid_ocs.png b/tests/baseline/test_plot_grid_ocs.png index 8c7bbe9a34..c78f0a9bb7 100644 Binary files a/tests/baseline/test_plot_grid_ocs.png and b/tests/baseline/test_plot_grid_ocs.png differ diff --git a/tests/baseline/test_plot_grid_quad.png b/tests/baseline/test_plot_grid_quad.png index 1ff2801630..a6a1d5e40f 100644 Binary files a/tests/baseline/test_plot_grid_quad.png and b/tests/baseline/test_plot_grid_quad.png differ diff --git a/tests/baseline/test_plot_logo.png b/tests/baseline/test_plot_logo.png index add27e5b78..2d81c70997 100644 Binary files a/tests/baseline/test_plot_logo.png and b/tests/baseline/test_plot_logo.png differ diff --git a/tests/baseline/test_plot_magnetic_pressure.png b/tests/baseline/test_plot_magnetic_pressure.png index 676b29cffd..776c0e1c58 100644 Binary files a/tests/baseline/test_plot_magnetic_pressure.png and b/tests/baseline/test_plot_magnetic_pressure.png differ diff --git a/tests/baseline/test_plot_magnetic_tension.png b/tests/baseline/test_plot_magnetic_tension.png index c163d54e24..48b4e32cc0 100644 Binary files a/tests/baseline/test_plot_magnetic_tension.png and b/tests/baseline/test_plot_magnetic_tension.png differ diff --git a/tests/baseline/test_plot_normF_2d.png b/tests/baseline/test_plot_normF_2d.png index 0a5b276dba..e69db5e89f 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 c27b57f3d2..6f973f6321 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_surfaces.png b/tests/baseline/test_plot_surfaces.png index 0e264beecc..6f9c9c12d5 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 7e7ffb592c..e483c1654a 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 a9821a4400..7b2ff0dc93 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_plot_vac_normF_2d.png b/tests/baseline/test_plot_vac_normF_2d.png deleted file mode 100644 index 3379594184..0000000000 Binary files a/tests/baseline/test_plot_vac_normF_2d.png and /dev/null differ diff --git a/tests/baseline/test_plot_vac_normF_section.png b/tests/baseline/test_plot_vac_normF_section.png deleted file mode 100644 index 121a45617f..0000000000 Binary files a/tests/baseline/test_plot_vac_normF_section.png and /dev/null differ diff --git a/tests/baseline/test_plot_vmec_comparison.png b/tests/baseline/test_plot_vmec_comparison.png index 75323c5bf8..72d07abf30 100644 Binary files a/tests/baseline/test_plot_vmec_comparison.png and b/tests/baseline/test_plot_vmec_comparison.png differ diff --git a/tests/baseline/test_section_F.png b/tests/baseline/test_section_F.png index 37a703eb4a..c0fd3a0a20 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 index 2831750fca..0fb596e2d1 100644 Binary files a/tests/baseline/test_section_F_normalized_vac.png and b/tests/baseline/test_section_F_normalized_vac.png differ diff --git a/tests/baseline/test_section_J.png b/tests/baseline/test_section_J.png index 67f679123b..106a3203d1 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 index aa55b1a19f..7157a55a7c 100644 Binary files a/tests/baseline/test_section_R.png and b/tests/baseline/test_section_R.png differ diff --git a/tests/baseline/test_section_Z.png b/tests/baseline/test_section_Z.png index 222077bcac..4a9b8f6d07 100644 Binary files a/tests/baseline/test_section_Z.png and b/tests/baseline/test_section_Z.png differ diff --git a/tests/baseline/test_section_logF.png b/tests/baseline/test_section_logF.png index 2d260256ff..0f19dd943d 100644 Binary files a/tests/baseline/test_section_logF.png and b/tests/baseline/test_section_logF.png differ diff --git a/tests/baseline/test_trapped_fraction_Kim.png b/tests/baseline/test_trapped_fraction_Kim.png index 698ab28aa2..faa395ef87 100644 Binary files a/tests/baseline/test_trapped_fraction_Kim.png and b/tests/baseline/test_trapped_fraction_Kim.png differ diff --git a/tests/conftest.py b/tests/conftest.py index 4e9b8acfc0..9f8fc3464f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -329,9 +329,7 @@ def DummyStellarator(tmpdir_factory): eq = Equilibrium(**inputs) eq.save(output_path) - DummyStellarator_out = { - "output_path": output_path, - } + DummyStellarator_out = {"output_path": output_path} return DummyStellarator_out diff --git a/tests/inputs/HELIOTRON_vacuum b/tests/inputs/HELIOTRON_vacuum index e9f67ca1e2..f401abbcf1 100644 --- a/tests/inputs/HELIOTRON_vacuum +++ b/tests/inputs/HELIOTRON_vacuum @@ -26,6 +26,10 @@ objective = vacuum spectral_indexing = ansi node_pattern = jacobi +# Note that the vacuum objective will optimize for current density = 0. +# The rotational transform profile coefficients below will only be used +# as an initial guess, as the rotational transform will not be constrained +# throughout the optimization. See GitHub issue #278. # pressure and rotational transform/current profiles l: 0 p = 0 i = 0 l: 2 p = 0 i = 0 diff --git a/tests/inputs/master_compute_data.pkl b/tests/inputs/master_compute_data.pkl new file mode 100644 index 0000000000..421bbbbe72 Binary files /dev/null and b/tests/inputs/master_compute_data.pkl differ diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index 0afbc21cf5..a953e5bfc1 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -1,29 +1,23 @@ """Tests for compute functions evaluated at limits.""" -import inspect -import re import numpy as np import pytest -import desc.compute from desc.compute import data_index -from desc.compute.data_index import _class_inheritance -from desc.compute.utils import surface_integrals_map +from desc.compute.utils import dot, surface_integrals_map from desc.equilibrium import Equilibrium +from desc.examples import get from desc.grid import LinearGrid # Unless mentioned in the source code of the compute function, the assumptions # made to compute the magnetic axis limit can be reduced to assuming that these # functions tend toward zero as the magnetic axis is approached and that -# d^2𝜓/(d𝜌)^2 and 𝜕√𝑔/𝜕𝜌 are both finite nonzero at the magnetic axis. -# Also d^n𝜓/(d𝜌)^n for n > 3 is assumed zero everywhere. +# d²ψ/(dρ)² and 𝜕√𝑔/𝜕𝜌 are both finite nonzero at the magnetic axis. +# Also, dⁿψ/(dρ)ⁿ for n > 3 is assumed zero everywhere. zero_limits = {"rho", "psi", "psi_r", "e_theta", "sqrt(g)", "B_t"} - not_finite_limits = { "D_Mercier", - "D_current", "D_geodesic", - "D_shear", # may not exist for all configurations "D_well", "J^theta", "curvature_H_rho", @@ -56,15 +50,17 @@ "|grad(theta)|", " Redl", # may not exist for all configurations } - -# reliant limits will be added to this set automatically not_implemented_limits = { + # reliant limits will be added to this set automatically "iota_num_rrr", "iota_den_rrr", + "D_current", } -def grow_seeds(seeds, search_space): +def grow_seeds( + seeds, search_space, parameterization="desc.equilibrium.equilibrium.Equilibrium" +): """Traverse the dependency DAG for keys in search space dependent on seeds. Parameters @@ -73,6 +69,9 @@ def grow_seeds(seeds, search_space): Keys to find paths toward. search_space : iterable Additional keys to consider returning. + parameterization: str or list of str + Name of desc types the method is valid for. eg 'desc.geometry.FourierXYZCurve' + or `desc.equilibrium.Equilibrium`. Returns ------- @@ -82,9 +81,7 @@ def grow_seeds(seeds, search_space): """ out = seeds.copy() for key in search_space: - deps = data_index["desc.equilibrium.equilibrium.Equilibrium"][key][ - "full_with_axis_dependencies" - ]["data"] + deps = data_index[parameterization][key]["full_with_axis_dependencies"]["data"] if not seeds.isdisjoint(deps): out.add(key) return out @@ -94,6 +91,7 @@ def grow_seeds(seeds, search_space): not_implemented_limits, data_index["desc.equilibrium.equilibrium.Equilibrium"].keys() - not_finite_limits, ) +not_implemented_limits.discard("D_Mercier") def _skip_this(eq, name): @@ -110,7 +108,7 @@ def _skip_this(eq, name): def assert_is_continuous( eq, - names, + names=data_index["desc.equilibrium.equilibrium.Equilibrium"].keys(), delta=5e-5, rtol=1e-4, atol=1e-6, @@ -124,7 +122,7 @@ def assert_is_continuous( ---------- eq : Equilibrium The equilibrium object used for the computation. - names : list, str + names : list of str A list of names of the quantities to test for continuity. delta: float, optional Max distance from magnetic axis. @@ -151,49 +149,56 @@ def assert_is_continuous( """ if kwargs is None: kwargs = {} - if isinstance(names, str): - names = [names] # TODO: remove when boozer transform works with multiple surfaces - names = [x for x in names if not ("Boozer" in x or "_mn" in x or x == "B modes")] + names = [ + name + for name in names + if not ( + "Boozer" in name + or "_mn" in name + or name == "B modes" + or _skip_this(eq, name) + ) + ] - num_points = 15 + num_points = 12 rho = np.linspace(start=0, stop=delta, num=num_points) grid = LinearGrid(rho=rho, M=5, N=5, NFP=eq.NFP, sym=eq.sym) - assert grid.axis.size + axis = grid.nodes[:, 0] == 0 + assert axis.any() and not axis.all() integrate = surface_integrals_map(grid, expand_out=False) - data = eq.compute(names, grid=grid) + data = eq.compute(names=names, grid=grid) - data_index_eq = data_index["desc.equilibrium.equilibrium.Equilibrium"] + p = "desc.equilibrium.equilibrium.Equilibrium" for name in names: - if _skip_this(eq, name) or data_index_eq[name]["coordinates"] == "": - # can't check continuity of global scaler quantity + if name in not_finite_limits: + assert (np.isfinite(data[name]).T != axis).all(), name + continue + else: + assert np.isfinite(data[name]).all(), name + if data_index[p][name]["coordinates"] == "": + # can't check continuity of global scalar continue # make single variable function of rho - if data_index_eq[name]["coordinates"] == "r": + if data_index[p][name]["coordinates"] == "r": # already single variable function of rho profile = grid.compress(data[name]) else: # integrate out theta and zeta dependence - profile = np.where( - # True if integrand has nan on a given surface. - integrate(np.isnan(data[name])).astype(bool), - # The integration below replaces nan with 0; put them back. - np.nan, - # Norms and integrals are continuous functions, so their composition - # cannot disrupt existing continuity. Note that the absolute value - # before the integration ensures that a discontinuous integrand does - # not become continuous once integrated. - integrate(np.abs(data[name])), - ) + # Norms and integrals are continuous functions, so their composition + # cannot disrupt existing continuity. Note that the absolute value + # before the integration ensures that a discontinuous integrand does + # not become continuous once integrated. + profile = integrate(np.abs(data[name])) fit = kwargs.get(name, {}).get("desired_at_axis", desired_at_axis) if fit is None: - if np.ndim(data_index_eq[name]["dim"]): + if np.ndim(data_index[p][name]["dim"]): # can't polyfit tensor arrays like grad(B) - fit = (profile[0] + profile[1]) / 2 + fit = profile[1] else: # fit the data to a polynomial to extrapolate to axis poly = np.polynomial.polynomial.polyfit( - rho[1:], profile[1:], deg=min(5, num_points // 3) + rho[1:], profile[1:], deg=min(4, num_points // 3) ) # constant term is the same as evaluating polynomial at rho=0 fit = poly[0] @@ -207,88 +212,9 @@ def assert_is_continuous( ) -def get_matches(fun, pattern): - """Return all matches of ``pattern`` in source code of function ``fun``.""" - src = inspect.getsource(fun) - # attempt to remove any decorator functions - # (currently works without this filter, but better to be defensive) - src = src.partition("def ")[2] - # attempt to remove comments - src = "\n".join(line.partition("#")[0] for line in src.splitlines()) - matches = pattern.findall(src) - matches = {s.strip().strip('"') for s in matches} - return matches - - -def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): - """Get parameterization of thing computed by function ``fun``.""" - pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")') - decorator = inspect.getsource(fun).partition("def ")[0] - matches = pattern.findall(decorator) - # if list was found, split strings in list, else string was found so just get that - matches = [match[0].split(",") if match[0] else [match[1]] for match in matches] - # flatten the list - matches = {s.strip().strip('"') for sublist in matches for s in sublist} - matches.discard("") - return matches if matches else {default} - - class TestAxisLimits: """Tests for compute functions evaluated at limits.""" - @pytest.mark.unit - def test_data_index_deps(self): - """Ensure developers do not add extra (or forget needed) dependencies.""" - queried_deps = {} - - pattern_names = re.compile(r"(?"], - expand( - grid, - np.array([13.0**2 + 0.5 * 2.6**2, 9.0**2 + 0.5 * 3.7**2]), + f_t_data["<|B|^2>"], + grid.expand( + np.array([13.0**2 + 0.5 * 2.6**2, 9.0**2 + 0.5 * 3.7**2]) ), ) np.testing.assert_allclose( f_t_data["<1/|B|>"], - expand( - grid, - np.array( - [ - 1 / np.sqrt(13.0**2 - 2.6**2), - 1 / np.sqrt(9.0**2 - 3.7**2), - ] - ), - ), + grid.expand(1 / np.sqrt([13.0**2 - 2.6**2, 9.0**2 - 3.7**2])), ) np.testing.assert_allclose( f_t_data["min_tz |B|"], - expand(grid, np.array([13.0 - 2.6, 9.0 - 3.7])), + grid.expand(np.array([13.0 - 2.6, 9.0 - 3.7])), rtol=1e-4, ) np.testing.assert_allclose( f_t_data["max_tz |B|"], - expand(grid, np.array([13.0 + 2.6, 9.0 + 3.7])), + grid.expand(np.array([13.0 + 2.6, 9.0 + 3.7])), rtol=1e-4, ) np.testing.assert_allclose( f_t_data["effective r/R0"], - expand(grid, np.array([2.6 / 13.0, 3.7 / 9.0])), + grid.expand(np.array([2.6 / 13.0, 3.7 / 9.0])), rtol=1e-3, ) @@ -152,60 +135,55 @@ def test_trapped_fraction_Kim(self): fig = plt.figure() def test(N, grid_type): - grid = grid_type( - L=L, - M=M, - N=N, - NFP=NFP, - ) + grid = grid_type(L=L, M=M, N=N, NFP=NFP) rho = grid.nodes[:, 0] theta = grid.nodes[:, 1] epsilon_3D = rho * epsilon_max - # Pick out unique values: - epsilon = np.array(sorted(list(set(epsilon_3D)))) + epsilon = np.unique(epsilon_3D) # Eq (A6) # noqa: E800 modB = B0 / (1 + epsilon_3D * np.cos(theta)) # For Jacobian, use eq (A7) for the theta dependence, # times an arbitrary overall scale factor sqrt_g = 6.7 * (1 + epsilon_3D * np.cos(theta)) - - f_t_data = trapped_fraction(grid, modB, sqrt_g) + # Above "Jacobian" is nonzero at magnetic axis, so set + # sqrt(g)_r as sqrt(g) to nullify automatic computation of + # limit which assumes sqrt(g) is true Jacobian and zero at the + # magnetic axis. + f_t_data = trapped_fraction(grid, modB, sqrt_g, sqrt_g_r=sqrt_g) # Eq (C18) in Kim et al: f_t_Kim = 1.46 * np.sqrt(epsilon) - 0.46 * epsilon np.testing.assert_allclose( - f_t_data["min_tz |B|"], expand(grid, B0 / (1 + epsilon)) + f_t_data["min_tz |B|"], grid.expand(B0 / (1 + epsilon)) ) # Looser tolerance for Bmax since there is no grid point there: Bmax = B0 / (1 - epsilon) np.testing.assert_allclose( - f_t_data["max_tz |B|"], expand(grid, Bmax), rtol=0.001 + f_t_data["max_tz |B|"], grid.expand(Bmax), rtol=0.001 ) np.testing.assert_allclose( - f_t_data["effective r/R0"], expand(grid, epsilon), rtol=1e-4 + f_t_data["effective r/R0"], grid.expand(epsilon), rtol=1e-4 ) # Eq (A8): - fsa_B2 = B0 * B0 / np.sqrt(1 - epsilon**2) + fsa_B2 = B0**2 / np.sqrt(1 - epsilon**2) np.testing.assert_allclose( - f_t_data[""], - expand(grid, fsa_B2), - rtol=1e-6, + f_t_data["<|B|^2>"], grid.expand(fsa_B2), rtol=1e-6 ) np.testing.assert_allclose( - f_t_data["<1/|B|>"], expand(grid, (2 + epsilon**2) / (2 * B0)) + f_t_data["<1/|B|>"], grid.expand((2 + epsilon**2) / (2 * B0)) ) # Note the loose tolerance for this next test since we do not expect precise # agreement. np.testing.assert_allclose( - f_t_data["trapped fraction"], expand(grid, f_t_Kim), rtol=0.1, atol=0.07 + f_t_data["trapped fraction"], grid.expand(f_t_Kim), rtol=0.1, atol=0.07 ) # Now compute f_t numerically by a different algorithm: modB = modB.reshape((grid.num_zeta, grid.num_rho, grid.num_theta)) sqrt_g = sqrt_g.reshape((grid.num_zeta, grid.num_rho, grid.num_theta)) - fourpisq = 4 * np.pi * np.pi + fourpisq = 4 * np.pi**2 d_V_d_rho = np.mean(sqrt_g, axis=(0, 2)) / fourpisq f_t = np.zeros(grid.num_rho) for jr in range(grid.num_rho): @@ -221,7 +199,7 @@ def integrand(lambd): f_t[jr] = 1 - 0.75 * fsa_B2[jr] * integral[0] np.testing.assert_allclose( - compress(grid, f_t_data["trapped fraction"])[1:], + grid.compress(f_t_data["trapped fraction"])[1:], f_t[1:], rtol=0.001, atol=0.001, @@ -229,7 +207,7 @@ def integrand(lambd): plt.plot(epsilon, f_t_Kim, "b", label="Kim") plt.plot( - epsilon, compress(grid, f_t_data["trapped fraction"]), "r", label="desc" + epsilon, grid.compress(f_t_data["trapped fraction"]), "r", label="desc" ) plt.plot(epsilon, f_t, ":g", label="Alternative algorithm") @@ -287,7 +265,7 @@ def test_Redl_second_pass(self): # Sauter eq (18b)-(18c): nu_e = abs( R - * (6.921e-18) + * 6.921e-18 * ne_rho * Zeff_rho * ln_Lambda_e @@ -295,7 +273,7 @@ def test_Redl_second_pass(self): ) nu_i = abs( R - * (4.90e-18) + * 4.90e-18 * ni_rho * (Zeff_rho**4) * ln_Lambda_ii @@ -521,7 +499,7 @@ def test_Redl_figures_2_3(self): # Sauter eq (18b), but without the iota factor: nu_e_without_iota = ( R - * (6.921e-18) + * 6.921e-18 * ne_rho * Zeff_rho * ln_Lambda_e @@ -658,7 +636,7 @@ def test_Redl_figures_4_5(self): # Sauter eq (18b), but without the q = 1/iota factor: nu_e_without_iota = ( R - * (6.921e-18) + * 6.921e-18 * ne_rho * Zeff_rho * ln_Lambda_e @@ -913,12 +891,8 @@ def test_Redl_sfincs_tokamak_benchmark(self): ) grid = LinearGrid(rho=rho, M=eq.M, N=eq.N, NFP=eq.NFP) - data = eq.compute( - " Redl", - grid=grid, - helicity=helicity, - ) - J_dot_B_Redl = compress(grid, data[" Redl"]) + data = eq.compute(" Redl", grid=grid, helicity=helicity) + J_dot_B_Redl = grid.compress(data[" Redl"]) # The relative error is a bit larger at the boundary, where the # absolute magnitude is quite small, so drop those points. @@ -1004,12 +978,8 @@ def test_Redl_sfincs_QA(self): ) grid = LinearGrid(rho=rho, M=eq.M, N=eq.N, NFP=eq.NFP) - data = eq.compute( - " Redl", - grid=grid, - helicity=helicity, - ) - J_dot_B_Redl = compress(grid, data[" Redl"]) + data = eq.compute(" Redl", grid=grid, helicity=helicity) + J_dot_B_Redl = grid.compress(data[" Redl"]) np.testing.assert_allclose(J_dot_B_Redl[1:-1], J_dot_B_sfincs[1:-1], rtol=0.1) @@ -1098,12 +1068,8 @@ def test_Redl_sfincs_QH(self): ) grid = LinearGrid(rho=rho, M=eq.M, N=eq.N, NFP=eq.NFP) - data = eq.compute( - " Redl", - grid=grid, - helicity=helicity, - ) - J_dot_B_Redl = compress(grid, data[" Redl"]) + data = eq.compute(" Redl", grid=grid, helicity=helicity) + J_dot_B_Redl = grid.compress(data[" Redl"]) np.testing.assert_allclose(J_dot_B_Redl[1:-1], J_dot_B_sfincs[1:-1], rtol=0.1) @@ -1167,10 +1133,7 @@ def test_BootstrapRedlConsistency_normalization(self): ) # The equilibrium need not be in force balance, so no need to solve(). grid = QuadratureGrid( - L=LMN_resolution, - M=LMN_resolution, - N=LMN_resolution, - NFP=eq.NFP, + L=LMN_resolution, M=LMN_resolution, N=LMN_resolution, NFP=eq.NFP ) obj = ObjectiveFunction( BootstrapRedlConsistency(eq=eq, grid=grid, helicity=helicity) @@ -1255,19 +1218,9 @@ def test_BootstrapRedlConsistency_resolution(self, DSHAPE_current): 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, - ) + grid = grid_type(L=L, M=M, N=N, NFP=eq.NFP, **kwargs) obj = ObjectiveFunction( - BootstrapRedlConsistency( - eq=eq, - grid=grid, - helicity=helicity, - ), + BootstrapRedlConsistency(eq=eq, grid=grid, helicity=helicity) ) obj.build() scalar_objective = obj.compute_scalar(obj.x(eq)) @@ -1380,11 +1333,7 @@ def test_bootstrap_consistency_iota(self, TmpDir): NFP=eq.NFP, ) objective = ObjectiveFunction( - BootstrapRedlConsistency( - eq=eq, - grid=grid, - helicity=helicity, - ) + BootstrapRedlConsistency(eq=eq, grid=grid, helicity=helicity) ) eq, _ = eq.optimize( verbose=3, @@ -1400,13 +1349,9 @@ def test_bootstrap_consistency_iota(self, TmpDir): scalar_objective = objective.compute_scalar(objective.x(eq)) assert scalar_objective < 3e-5 - data = eq.compute( - ["", " Redl"], - grid=grid, - helicity=helicity, - ) - J_dot_B_MHD = compress(grid, data[""]) - J_dot_B_Redl = compress(grid, data[" Redl"]) + data = eq.compute(["", " Redl"], grid=grid, helicity=helicity) + J_dot_B_MHD = grid.compress(data[""]) + J_dot_B_Redl = grid.compress(data[" Redl"]) assert np.max(J_dot_B_MHD) < 4e5 assert np.max(J_dot_B_MHD) > 0 @@ -1494,18 +1439,9 @@ def test_bootstrap_consistency_current(self, TmpDir): ) # grid for bootstrap consistency objective: - grid = QuadratureGrid( - L=current_L * 2, - M=eq.M * 2, - N=eq.N * 2, - NFP=eq.NFP, - ) + grid = QuadratureGrid(L=current_L * 2, M=eq.M * 2, N=eq.N * 2, NFP=eq.NFP) objective = ObjectiveFunction( - BootstrapRedlConsistency( - eq=eq, - grid=grid, - helicity=helicity, - ) + BootstrapRedlConsistency(eq=eq, grid=grid, helicity=helicity) ) eq, _ = eq.optimize( verbose=3, @@ -1524,13 +1460,9 @@ def test_bootstrap_consistency_current(self, TmpDir): scalar_objective = objective.compute_scalar(objective.x(eq)) assert scalar_objective < 3e-5 - data = eq.compute( - ["", " Redl"], - grid=grid, - helicity=helicity, - ) - J_dot_B_MHD = compress(grid, data[""]) - J_dot_B_Redl = compress(grid, data[" Redl"]) + data = eq.compute(["", " Redl"], grid=grid, helicity=helicity) + J_dot_B_MHD = grid.compress(data[""]) + J_dot_B_Redl = grid.compress(data[" Redl"]) assert np.max(J_dot_B_MHD) < 4e5 assert np.max(J_dot_B_MHD) > 0 diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 274087759d..4db2dd3c3c 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1,14 +1,15 @@ """Tests for compute functions.""" +import pickle + import numpy as np import pytest from scipy.io import netcdf_file from scipy.signal import convolve2d -import desc.examples from desc.compute import data_index, rpz2xyz_vec -from desc.compute.utils import compress from desc.equilibrium import EquilibriaFamily, Equilibrium +from desc.examples import get from desc.geometry import ( FourierPlanarCurve, FourierRZCurve, @@ -51,7 +52,7 @@ def test_total_volume(DummyStellarator): ) grid = LinearGrid(M=12, N=12, NFP=eq.NFP, sym=eq.sym) # rho = 1 - lcfs_volume = eq.compute("V(r)", grid=grid)["V(r)"].mean() + lcfs_volume = eq.compute("V(r)", grid=grid)["V(r)"] total_volume = eq.compute("V")["V"] # default quadrature grid np.testing.assert_allclose(lcfs_volume, total_volume) @@ -59,25 +60,29 @@ def test_total_volume(DummyStellarator): @pytest.mark.unit def test_enclosed_volumes(): """Test that the volume enclosed by flux surfaces matches analytic formulas.""" + R0 = 10 surf = FourierRZToroidalSurface( - R_lmn=[10, 1, 0.2], + R_lmn=[R0, 1, 0.2], Z_lmn=[-2, -0.2], modes_R=[[0, 0], [1, 0], [0, 1]], modes_Z=[[-1, 0], [0, -1]], ) + # 𝐞(ρ, θ, ζ) = R(ρ, θ, ζ) 𝐫 + Z(ρ, θ, ζ) 𝐳 + # V(ρ) = ∯ dθ dζ (∂_θ 𝐞 × ∂_ζ 𝐞) ⋅ (0, 0, Z) + # = ∯ dθ dζ (R₀ + ρ cos θ + 0.2 cos ζ) (2 ρ² sin²θ − 0.2 ρ sin θ sin ζ) + np.testing.assert_allclose(4 * R0 * np.pi**2, surf.compute(["V"])["V"]) eq = Equilibrium(surface=surf) # elliptical cross-section with torsion - rho = np.linspace(1 / 128, 1, 128) + rho = np.linspace(0, 1, 64) grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - data = eq.compute(["R0", "V(r)", "V_r(r)", "V_rr(r)"], grid=grid) - np.testing.assert_allclose( - 4 * data["R0"] * (np.pi * rho) ** 2, compress(grid, data["V(r)"]) - ) + data = eq.compute(["R0", "V(r)", "V_r(r)", "V_rr(r)", "V_rrr(r)"], grid=grid) np.testing.assert_allclose( - 8 * data["R0"] * np.pi**2 * rho, compress(grid, data["V_r(r)"]) + 4 * data["R0"] * (np.pi * rho) ** 2, grid.compress(data["V(r)"]) ) np.testing.assert_allclose( - 8 * data["R0"] * np.pi**2, compress(grid, data["V_rr(r)"]) + 8 * data["R0"] * np.pi**2 * rho, grid.compress(data["V_r(r)"]) ) + np.testing.assert_allclose(8 * data["R0"] * np.pi**2, data["V_rr(r)"]) + np.testing.assert_allclose(0, data["V_rrr(r)"], atol=2e-14) @pytest.mark.unit @@ -90,21 +95,25 @@ def test_enclosed_areas(): modes_Z=[[-1, 0], [0, -1]], ) eq = Equilibrium(surface=surf) # elliptical cross-section with torsion - rho = np.linspace(1 / 128, 1, 128) + rho = np.linspace(0, 1, 64) grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) data = eq.compute(["A(r)"], grid=grid) - np.testing.assert_allclose(2 * np.pi * rho**2, compress(grid, data["A(r)"])) + # area = π a b = 2 π ρ² + np.testing.assert_allclose(2 * np.pi * rho**2, grid.compress(data["A(r)"])) @pytest.mark.unit def test_surface_areas(): """Test that the flux surface areas match known analytic formulas.""" eq = Equilibrium() # torus - rho = np.linspace(1 / 128, 1, 128) + rho = np.linspace(0, 1, 64) grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - data = eq.compute(["S(r)", "R0"], grid=grid) - S = 4 * data["R0"] * np.pi**2 * rho - np.testing.assert_allclose(S, compress(grid, data["S(r)"])) + data = eq.compute(["R0", "S(r)", "S_r(r)", "S_rr(r)"], grid=grid) + np.testing.assert_allclose( + 4 * data["R0"] * np.pi**2 * rho, grid.compress(data["S(r)"]) + ) + np.testing.assert_allclose(4 * data["R0"] * np.pi**2, data["S_r(r)"]) + np.testing.assert_allclose(0, data["S_rr(r)"], atol=3e-12) @pytest.mark.unit @@ -153,7 +162,7 @@ def test_elongation(): eq1 = Equilibrium() # elongation = 1 eq2 = Equilibrium(surface=surf2) # elongation = 2 eq3 = Equilibrium(surface=surf3) # elongation = 3 - rho = np.linspace(1 / 128, 1, 128) + rho = np.linspace(0, 1, 128) grid = LinearGrid(M=eq3.M_grid, N=eq3.N_grid, NFP=eq3.NFP, sym=eq3.sym, rho=rho) data1 = eq1.compute(["a_major/a_minor"], grid=grid) data2 = eq2.compute(["a_major/a_minor"], grid=grid) @@ -1083,31 +1092,26 @@ def test_BdotgradB(DummyStellarator): load_from=str(DummyStellarator["output_path"]), file_format="hdf5" ) - # partial derivative wrt theta - num_theta = 120 - grid = LinearGrid(NFP=eq.NFP, theta=num_theta) - dtheta = grid.nodes[1, 1] - data = eq.compute(["B*grad(|B|)", "(B*grad(|B|))_t"], grid=grid) - Btilde_t = np.convolve(data["B*grad(|B|)"], FD_COEF_1_4, "same") / dtheta - np.testing.assert_allclose( - data["(B*grad(|B|))_t"][2:-2], - Btilde_t[2:-2], - rtol=2e-2, - atol=2e-2 * np.mean(np.abs(data["(B*grad(|B|))_t"])), - ) + def test_partial_derivative(name): + cases = { + "r": {"label": "rho", "column_id": 0}, + "t": {"label": "theta", "column_id": 1}, + "z": {"label": "zeta", "column_id": 2}, + }[name[-1]] + grid = LinearGrid(NFP=eq.NFP, **{cases["label"]: 120}) + dx = grid.nodes[1, cases["column_id"]] + data = eq.compute(["B*grad(|B|)", name], grid=grid) + Btilde_x = np.convolve(data["B*grad(|B|)"], FD_COEF_1_4, "same") / dx + np.testing.assert_allclose( + actual=data[name][2:-2], + desired=Btilde_x[2:-2], + rtol=2e-2, + atol=2e-2 * np.mean(np.abs(data[name])), + ) - # partial derivative wrt zeta - num_zeta = 120 - grid = LinearGrid(NFP=eq.NFP, zeta=num_zeta) - dzeta = grid.nodes[1, 2] - data = eq.compute(["B*grad(|B|)", "(B*grad(|B|))_z"], grid=grid) - Btilde_z = np.convolve(data["B*grad(|B|)"], FD_COEF_1_4, "same") / dzeta - np.testing.assert_allclose( - data["(B*grad(|B|))_z"][2:-2], - Btilde_z[2:-2], - rtol=2e-2, - atol=2e-2 * np.mean(np.abs(data["(B*grad(|B|))_z"])), - ) + test_partial_derivative("(B*grad(|B|))_r") + test_partial_derivative("(B*grad(|B|))_t") + test_partial_derivative("(B*grad(|B|))_z") # TODO: add test with stellarator example @@ -1181,60 +1185,97 @@ def test_compare_quantities_to_vmec(): 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 = compress(grid, data[""]) - - # Drop first point since desc gives NaN: - np.testing.assert_allclose(J_dot_B_desc[1:], J_dot_B_vmec[1:], rtol=0.005) + 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(): - """Make sure we can compute everything without errors.""" - eq = Equilibrium(1, 1, 1) - grid = LinearGrid(1, 1, 1) - for key in data_index["desc.equilibrium.equilibrium.Equilibrium"].keys(): - data = eq.compute(key, grid=grid) - assert key in data - - -@pytest.mark.unit -def test_curve_compute_everything(): - """Make sure we can compute every curve thing without errors.""" - curves = { - "desc.geometry.curve.FourierXYZCurve": FourierXYZCurve(), - "desc.geometry.curve.FourierRZCurve": FourierRZCurve(), - "desc.geometry.curve.FourierPlanarCurve": FourierPlanarCurve(), + """Test that the computations on this branch agree with those on master. + + Also make sure we can compute everything without errors. + """ + elliptic_cross_section_with_torsion = { + "R_lmn": [10, 1, 0.2], + "Z_lmn": [-2, -0.2], + "modes_R": [[0, 0], [1, 0], [0, 1]], + "modes_Z": [[-1, 0], [0, -1]], } - - for p, thing in curves.items(): - for key in data_index[p].keys(): - data = thing.compute(key) - assert key in data - - -@pytest.mark.unit -def test_surface_compute_everything(): - """Make sure we can compute every surface thing without errors.""" - surfaces = { - "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface(), - "desc.geometry.surface.ZernikeRZToroidalSection": ZernikeRZToroidalSection(), + things = { + # equilibria + "desc.equilibrium.equilibrium.Equilibrium": get("W7-X"), + # curves + "desc.geometry.curve.FourierXYZCurve": FourierXYZCurve( + X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] + ), + "desc.geometry.curve.FourierRZCurve": FourierRZCurve( + R_n=[10, 1, 0.2], Z_n=[-2, -0.2], modes_R=[0, 1, 2], modes_Z=[-1, -2], NFP=2 + ), + "desc.geometry.curve.FourierPlanarCurve": FourierPlanarCurve( + center=[10, 1, 3], normal=[1, 2, 3], r_n=[1, 2, 3], modes=[0, 1, 2] + ), + # surfaces + "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface( + **elliptic_cross_section_with_torsion + ), + "desc.geometry.surface.ZernikeRZToroidalSection": ZernikeRZToroidalSection( + **elliptic_cross_section_with_torsion + ), } + # use this low resolution grid for equilibria to reduce file size + grid = LinearGrid( + # include magnetic axis + rho=np.linspace(0, 1, 10), + M=5, + N=5, + NFP=things["desc.equilibrium.equilibrium.Equilibrium"].NFP, + sym=things["desc.equilibrium.equilibrium.Equilibrium"].sym, + ) + grid = {"desc.equilibrium.equilibrium.Equilibrium": {"grid": grid}} + + with open("tests/inputs/master_compute_data.pkl", "rb") as file: + master_data = pickle.load(file) + this_branch_data = {} + update_master_data = False + error = False + + for p in things: + this_branch_data[p] = things[p].compute( + list(data_index[p].keys()), **grid.get(p, {}) + ) + # make sure we can compute everything + assert this_branch_data[p].keys() == data_index[p].keys(), p + # compare against master branch + for name in this_branch_data[p]: + if p in master_data and name in master_data[p]: + try: + np.testing.assert_allclose( + actual=this_branch_data[p][name], + desired=master_data[p][name], + atol=1e-12, + err_msg=f"Parameterization: {p}. Name: {name}.", + ) + except AssertionError as e: + error = True + print(e) + else: + update_master_data = True - for p, thing in surfaces.items(): - for key in data_index[p].keys(): - data = thing.compute(key) - assert key in data + if not error and update_master_data: + with open("tests/inputs/master_compute_data.pkl", "wb") as file: + pickle.dump(this_branch_data, file) + assert not error @pytest.mark.unit def test_compute_averages(): """Test that computing averages uses the correct grid.""" - eq = desc.examples.get("HELIOTRON") - Vr = eq.get_profile("V_r(r)") + eq = get("HELIOTRON") + V_r = eq.get_profile("V_r(r)") rho = np.linspace(0.01, 1, 20) grid = LinearGrid(rho=rho, NFP=eq.NFP) out = eq.compute("V_r(r)", grid=grid) - np.testing.assert_allclose(Vr(rho), out["V_r(r)"], rtol=1e-4) + np.testing.assert_allclose(V_r(rho), out["V_r(r)"], rtol=1e-4) eq = Equilibrium(1, 1, 1) grid = LinearGrid(rho=[0.3], theta=[np.pi / 3], zeta=[0]) diff --git a/tests/test_compute_utils.py b/tests/test_compute_utils.py index 517d3dc21e..a31f8d391c 100644 --- a/tests/test_compute_utils.py +++ b/tests/test_compute_utils.py @@ -1,12 +1,17 @@ -"""Tests for surface averaging etc.""" +"""Tests compute utilities.""" + +import inspect +import re import numpy as np import pytest +import desc.compute +from desc.basis import FourierZernikeBasis +from desc.compute import data_index +from desc.compute.data_index import _class_inheritance from desc.compute.utils import ( _get_grid_surface, - compress, - expand, line_integrals, surface_averages, surface_integrals, @@ -17,104 +22,148 @@ ) from desc.examples import get from desc.grid import ConcentricGrid, LinearGrid, QuadratureGrid +from desc.transform import Transform + + +class TestDataIndex: + """Tests for things related to data_index.""" + + @staticmethod + def get_matches(fun, pattern): + """Return all matches of ``pattern`` in source code of function ``fun``.""" + src = inspect.getsource(fun) + # attempt to remove any decorator functions + # (currently works without this filter, but better to be defensive) + src = src.partition("def ")[2] + # attempt to remove comments + src = "\n".join(line.partition("#")[0] for line in src.splitlines()) + matches = pattern.findall(src) + matches = {s.strip().strip('"') for s in matches} + return matches + + @staticmethod + def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): + """Get parameterization of thing computed by function ``fun``.""" + pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")') + decorator = inspect.getsource(fun).partition("def ")[0] + matches = pattern.findall(decorator) + # if list was found, split strings in list, else string was found so get that + matches = [match[0].split(",") if match[0] else [match[1]] for match in matches] + # flatten the list + matches = {s.strip().strip('"') for sublist in matches for s in sublist} + matches.discard("") + return matches if matches else {default} - -def benchmark_surface_integrals(grid, q=np.array([1.0]), surface_label="rho"): - """Compute a surface integral for each surface in the grid. - - Notes - ----- - It is assumed that the integration surface has area 4π^2 when the - surface label is rho and area 2π when the surface label is theta or - zeta. You may want to multiply q by the surface area Jacobian. - - Parameters - ---------- - grid : Grid - Collocation grid containing the nodes to evaluate at. - q : ndarray - Quantity to integrate. - The first dimension of the array should have size ``grid.num_nodes``. - - When ``q`` is 1-dimensional, the intention is to integrate, - over the domain parameterized by rho, theta, and zeta, - a scalar function over the previously mentioned domain. - - When ``q`` is 2-dimensional, the intention is to integrate, - over the domain parameterized by rho, theta, and zeta, - a vector-valued function over the previously mentioned domain. - - When ``q`` is 3-dimensional, the intention is to integrate, - over the domain parameterized by rho, theta, and zeta, - a matrix-valued function over the previously mentioned domain. - surface_label : str - The surface label of rho, theta, or zeta to compute the integration over. - - Returns - ------- - integrals : ndarray - Surface integral of the input over each surface in the grid. - - """ - _, _, spacing, has_endpoint_dupe = _get_grid_surface(grid, surface_label) - weights = (spacing.prod(axis=1) * np.nan_to_num(q).T).T - - surfaces = {} - nodes = grid.nodes[:, {"rho": 0, "theta": 1, "zeta": 2}[surface_label]] - # collect node indices for each surface_label surface - for grid_row_idx, surface_label_value in enumerate(nodes): - surfaces.setdefault(surface_label_value, []).append(grid_row_idx) - # integration over non-contiguous elements - integrals = [] - for _, surface_idx in sorted(surfaces.items()): - integrals.append(weights[surface_idx].sum(axis=0)) - if has_endpoint_dupe: - integrals[0] += integrals[-1] - integrals[-1] = integrals[0] - return np.asarray(integrals) + @pytest.mark.unit + def test_data_index_deps(self): + """Ensure developers do not add extra (or forget needed) dependencies.""" + queried_deps = {} + + pattern_names = re.compile(r"(?" 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 + ) + + 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.solve @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) diff --git a/tests/test_stability_funs.py b/tests/test_stability_funs.py index de61bedb71..9326575423 100644 --- a/tests/test_stability_funs.py +++ b/tests/test_stability_funs.py @@ -6,7 +6,6 @@ import desc.examples import desc.io -from desc.compute.utils import compress from desc.equilibrium import Equilibrium from desc.grid import LinearGrid from desc.objectives import MagneticWell, MercierStability @@ -17,7 +16,7 @@ MAX_SIGN_DIFF = 5 -def all_close( +def assert_all_close( y1, y2, rho, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL ): """Test that the values of y1 and y2, over a given range are close enough. @@ -39,7 +38,7 @@ def all_close( """ minimum, maximum = rho_range - interval = np.where((minimum < rho) & (rho < maximum))[0] + interval = (minimum < rho) & (rho < maximum) np.testing.assert_allclose(y1[interval], y2[interval], rtol=rtol, atol=atol) @@ -62,8 +61,8 @@ def get_vmec_data(stellarator, quantity): """ f = Dataset(str(stellarator["vmec_nc_path"])) - rho = np.sqrt(f.variables["phi"] / np.asarray(f.variables["phi"])[-1]) - q = np.asarray(f.variables[quantity]) + rho = np.sqrt(f.variables["phi"] / np.array(f.variables["phi"])[-1]) + q = np.array(f.variables[quantity]) f.close() return rho, q @@ -72,11 +71,12 @@ def get_vmec_data(stellarator, quantity): def test_mercier_vacuum(): """Test that the Mercier stability criteria are 0 without pressure.""" eq = Equilibrium() - np.testing.assert_allclose(eq.compute("D_shear")["D_shear"], 0) - np.testing.assert_allclose(eq.compute("D_current")["D_current"], 0) - np.testing.assert_allclose(eq.compute("D_well")["D_well"], 0) - np.testing.assert_allclose(eq.compute("D_geodesic")["D_geodesic"], 0) - np.testing.assert_allclose(eq.compute("D_Mercier")["D_Mercier"], 0) + data = eq.compute(["D_shear", "D_current", "D_well", "D_geodesic", "D_Mercier"]) + np.testing.assert_allclose(data["D_shear"], 0) + np.testing.assert_allclose(data["D_current"], 0) + np.testing.assert_allclose(data["D_well"], 0) + np.testing.assert_allclose(data["D_geodesic"], 0) + np.testing.assert_allclose(data["D_Mercier"], 0) @pytest.mark.unit @@ -84,16 +84,16 @@ def test_mercier_vacuum(): def test_compute_d_shear(DSHAPE_current, HELIOTRON_ex): """Test that D_shear has a stabilizing effect and matches VMEC.""" - def test(stellarator, rho_range=(0, 1), rtol=1e-12, atol=0): + 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") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - d_shear = compress(grid, eq.compute("D_shear", grid=grid)["D_shear"]) + d_shear = grid.compress(eq.compute("D_shear", grid=grid)["D_shear"]) assert np.all( - d_shear[np.isfinite(d_shear)] >= 0 + d_shear[bool(grid.axis.size) :] >= 0 ), "D_shear should always have a stabilizing effect." - all_close(d_shear, d_shear_vmec, rho, rho_range, rtol, atol) + 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) @@ -110,13 +110,13 @@ def test( eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] rho, d_current_vmec = get_vmec_data(stellarator, "DCurr") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - d_current = compress(grid, eq.compute("D_current", grid=grid)["D_current"]) + d_current = grid.compress(eq.compute("D_current", grid=grid)["D_current"]) assert ( - len(np.where(np.sign(d_current) != np.sign(d_current_vmec))[0]) + np.nonzero(np.sign(d_current) != np.sign(d_current_vmec))[0].size <= MAX_SIGN_DIFF ) - all_close(d_current, d_current_vmec, rho, rho_range, rtol, atol) + 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) @@ -133,12 +133,12 @@ def test( eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] rho, d_well_vmec = get_vmec_data(stellarator, "DWell") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - d_well = compress(grid, eq.compute("D_well", grid=grid)["D_well"]) + d_well = grid.compress(eq.compute("D_well", grid=grid)["D_well"]) assert ( - len(np.where(np.sign(d_well) != np.sign(d_well_vmec))[0]) <= MAX_SIGN_DIFF + np.nonzero(np.sign(d_well) != np.sign(d_well_vmec))[0].size <= MAX_SIGN_DIFF ) - all_close(d_well, d_well_vmec, rho, rho_range, rtol, atol) + 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) @@ -157,12 +157,12 @@ def test( eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] rho, d_geodesic_vmec = get_vmec_data(stellarator, "DGeod") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - d_geodesic = compress(grid, eq.compute("D_geodesic", grid=grid)["D_geodesic"]) + d_geodesic = grid.compress(eq.compute("D_geodesic", grid=grid)["D_geodesic"]) assert np.all( - d_geodesic[np.isfinite(d_geodesic)] <= 0 + d_geodesic[bool(grid.axis.size) :] <= 0 ), "D_geodesic should always have a destabilizing effect." - all_close(d_geodesic, d_geodesic_vmec, rho, rho_range, rtol, atol) + 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) @@ -180,13 +180,13 @@ def test( eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] rho, d_mercier_vmec = get_vmec_data(stellarator, "DMerc") grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - d_mercier = compress(grid, eq.compute("D_Mercier", grid=grid)["D_Mercier"]) + d_mercier = grid.compress(eq.compute("D_Mercier", grid=grid)["D_Mercier"]) assert ( - len(np.where(np.sign(d_mercier) != np.sign(d_mercier_vmec))[0]) + np.nonzero(np.sign(d_mercier) != np.sign(d_mercier_vmec))[0].size <= MAX_SIGN_DIFF ) - all_close(d_mercier, d_mercier_vmec, rho, rho_range, rtol, atol) + 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) @@ -201,13 +201,13 @@ def test_compute_magnetic_well(DSHAPE_current, HELIOTRON_ex): def test(stellarator, rho=np.linspace(0, 1, 128)): eq = desc.io.load(load_from=str(stellarator["desc_h5_path"]))[-1] grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho) - d_well = compress(grid, eq.compute("D_well", grid=grid)["D_well"]) - magnetic_well = compress( - grid, eq.compute("magnetic well", grid=grid)["magnetic well"] + d_well = grid.compress(eq.compute("D_well", grid=grid)["D_well"]) + magnetic_well = grid.compress( + eq.compute("magnetic well", grid=grid)["magnetic well"] ) - assert ( - len(np.where(np.sign(d_well) != np.sign(magnetic_well))[0]) <= MAX_SIGN_DIFF + np.nonzero(np.sign(d_well) != np.sign(magnetic_well))[0].size + <= MAX_SIGN_DIFF ) test(DSHAPE_current) @@ -268,7 +268,7 @@ def test_magwell_print(capsys): obj = MagneticWell(eq=eq, grid=grid) obj.build() - magwell = compress(grid, eq.compute("magnetic well", grid=grid)["magnetic well"]) + magwell = grid.compress(eq.compute("magnetic well", grid=grid)["magnetic well"]) f = obj.compute(*obj.xs(eq)) np.testing.assert_allclose(f, magwell) diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index 010e4d72d3..3da0b7e105 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -41,19 +41,8 @@ def test_misc(self): np.testing.assert_allclose(Z, 0) c.set_coeffs(0, 0, 5, None) c.set_coeffs(-1, 0, None, 2) - np.testing.assert_allclose( - c.R_lmn, - [ - 5, - 1, - ], - ) - np.testing.assert_allclose( - c.Z_lmn, - [ - 2, - ], - ) + np.testing.assert_allclose(c.R_lmn, [5, 1]) + np.testing.assert_allclose(c.Z_lmn, [2]) s = c.copy() assert s.eq(c) @@ -167,19 +156,8 @@ def test_misc(self): np.testing.assert_allclose(Z, 0) c.set_coeffs(0, 0, 5, None) c.set_coeffs(1, -1, None, 2) - np.testing.assert_allclose( - c.R_lmn, - [ - 5, - 1, - ], - ) - np.testing.assert_allclose( - c.Z_lmn, - [ - 2, - ], - ) + np.testing.assert_allclose(c.R_lmn, [5, 1]) + np.testing.assert_allclose(c.Z_lmn, [2]) with pytest.raises(ValueError): c.set_coeffs(0, 0, None, 2) s = c.copy() diff --git a/tests/test_vmec.py b/tests/test_vmec.py index abca87455c..c318a3bd6e 100644 --- a/tests/test_vmec.py +++ b/tests/test_vmec.py @@ -5,7 +5,6 @@ from netCDF4 import Dataset from desc.basis import DoubleFourierSeries, FourierZernikeBasis -from desc.compute.utils import compress from desc.equilibrium import EquilibriaFamily, Equilibrium from desc.grid import LinearGrid from desc.vmec import VMECIO @@ -276,10 +275,10 @@ def test_vmec_load_profiles(TmpDir): data_iota = eq_iota.compute(["iota", "current"], grid=grid) data_current = eq_current.compute(["iota", "current"], grid=grid) - iota_iota = compress(grid, data_iota["iota"]) - iota_current = compress(grid, data_current["iota"]) - current_iota = compress(grid, data_iota["current"]) - current_current = compress(grid, data_current["current"]) + iota_iota = grid.compress(data_iota["iota"]) + iota_current = grid.compress(data_current["iota"]) + current_iota = grid.compress(data_iota["current"]) + current_current = grid.compress(data_current["current"]) np.testing.assert_allclose(iota_iota, iota_current, rtol=2e-2) np.testing.assert_allclose(current_current, current_iota, rtol=2e-2) @@ -371,7 +370,7 @@ def test_vmec_save_asym(TmpDir): VMECIO.save(eq, output_path) -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.slow def test_vmec_save_1(VMEC_save): """Tests that saving in NetCDF format agrees with VMEC.""" @@ -556,7 +555,7 @@ def test_vmec_save_1(VMEC_save): ) -@pytest.mark.unit +@pytest.mark.regression @pytest.mark.slow def test_vmec_save_2(VMEC_save): """Tests that saving in NetCDF format agrees with VMEC."""