From fd6b731a70098c6fa548c518b9889a847fdf4548 Mon Sep 17 00:00:00 2001 From: unalmis Date: Sat, 30 May 2026 01:15:44 -0700 Subject: [PATCH 1/2] cleanup --- desc/compute/_laplace.py | 481 ++++++++++++++++-------------- desc/magnetic_fields/__init__.py | 2 + desc/magnetic_fields/_laplace.py | 8 +- desc/objectives/_free_boundary.py | 253 ++++++++++------ tests/test_integrals.py | 112 ++++--- tests/test_objective_funs.py | 74 ++++- 6 files changed, 548 insertions(+), 382 deletions(-) diff --git a/desc/compute/_laplace.py b/desc/compute/_laplace.py index 2556ca9879..fcdc26521e 100644 --- a/desc/compute/_laplace.py +++ b/desc/compute/_laplace.py @@ -1,13 +1,14 @@ -"""Compute functions for multiply connected geometry Laplace solver. +"""Compute functions for multiply connected Laplace solver as described in [1]_. References ---------- - [1] Unalmis et al. New high-order accurate free surface stellarator +.. [1] Unalmis et al. New high-order accurate free surface stellarator equilibria optimization and boundary integral methods in DESC. """ from functools import partial +from typing import NamedTuple, Optional import equinox as eqx import jax @@ -26,66 +27,93 @@ get_interpolator, singular_integral, ) -from desc.utils import apply, cross, dot, errorif +from desc.utils import cross, dot, errorif from .data_index import register_compute_fun -_doc = { - "Phi_0": """jnp.ndarray : - Initial guess for iteration. - """, - "xtol": """float : - Absolute and relative error for the linear solve. Default is ``1e-8``. - """, - "maxiter": """int : - Maximum number of iterations for iterative scalar potential solves. - If non-positive then the materialized matrix will be inverted instead. - If positive, then performs that many iterations until ``maxiter`` or an - error tolerance of ``xtol`` is reached. Ten iterations should suffice for - the default GMRES solve. Default is ``10``. - """, - "solve_method": """str : - Method to use for the scalar potential solve. One of ``"auto"``, - ``"fixed_point"``, ``"gmres"``, or ``"least_squares"``. - If ``"auto"``, then uses GMRES when ``maxiter`` is positive and the - problem supports it, otherwise use the least-squares solve. Default is - ``"auto"``. - """, - "full_output": """bool - Whether to compute the maximum error ``Phi error`` and store the number of - iterations ``num iter`` used by the scalar potential solver. Default is - ``False``. - """, - "chunk_size": """int or None : - Size to split integral computation into chunks. - If no chunking should be done or the chunk size is the full input - then supply ``None``. Default is ``None``. - Recommend to verify computation with ``chunk_size`` set to a - small number due to bugs in JAX or XLA. - """, - "_D_quad": """bool - Set to ``True`` to perform double layer potential quadrature without removing - singularities. Default is ``False``. This is intended for developer use. - """, -} - - -def _select_potential_solve_method(solve_method, maxiter, problem): - if solve_method == "auto": - if (maxiter > 0) and (problem != "interior Neumann"): - return "gmres" - return "least_squares" - errorif( - solve_method not in {"fixed_point", "gmres", "least_squares"}, - msg="solve_method must be one of 'auto', 'fixed_point', 'gmres', " - f"or 'least_squares', got {solve_method!r}.", - ) - errorif( - solve_method in {"fixed_point", "gmres"} and (problem == "interior Neumann"), - msg=f"solve_method={solve_method!r} is not supported for interior Neumann " - "problems. Use solve_method='least_squares' instead.", - ) - return solve_method + +class Options(NamedTuple): + """Laplace solver options.""" + + Phi_0: Optional[jax.Array] = None + """Initial guess for iteration.""" + + atol: float = 1e-7 + """Absolute error tolerance for the iterative linear solve. Default is ``1e-7``.""" + + rtol: float = 1e-6 + """Relative error tolerance for the iterative linear solve. Default is ``1e-6``.""" + + max_steps: int = 10 + """Maximum number of steps for iterative linear solve. + + Typically converges in 2 iterations. Default max value is ``10``. + """ + + problem: str = "interior Neumann" + """Boundary value problem to solve. + + One of ``"interior Neumann"``, ``"exterior Neumann"``, or ``"interior Dirichlet"``. + (In some routines this may be determined automatically.) + """ + + solve_method: str = "auto" + """Method to use for the scalar potential solve. + + One of ``"auto"``, ``"gmres"``, or ``"direct"``. If ``"auto"``, then uses + GMRES when the problem supports it, otherwise uses the direct solve. Default + is ``"auto"``. If GMRES errors due to incompatibility with old JAX versions, + ``"fixed_point"`` can be selected instead. + """ + + full_output: bool = False + """Whether to return diagnostic output of the iterative potential solve. + + If ``True``, computes the maximum error ``Phi error`` and stores the number + of steps ``num_steps`` used by the scalar potential solver. Default is + ``False``. + """ + + chunk_size: Optional[int] = None + """Size to split integral computation into chunks. + + If no chunking should be done or the chunk size is the full input then + supply ``None``. Default is ``None``. Recommend to verify computation with + ``chunk_size`` set to a small number due to bugs in JAX or XLA. + """ + + B_coil_chunk_size: Optional[int] = None + """Size to split coil integral computation into chunks. + + If no chunking should be done or the chunk size is the full input then + supply ``None``. Default is ``None``. + """ + + D_quad: bool = False + """Developer option for double-layer potential quadrature. + + Set to ``True`` to perform double-layer potential quadrature without removing + singularities. Default is ``False``. + """ + + @staticmethod + def select_solver(options): + """Pick the solver based on the problem.""" + solve_method = options.solve_method + is_interior_neumann = options.problem == "interior Neumann" + if solve_method == "auto": + solve_method = "direct" if is_interior_neumann else "gmres" + errorif( + solve_method not in {"fixed_point", "gmres", "direct"}, + msg="solve_method must be one of 'auto', 'fixed_point', 'gmres', " + f"or 'direct', got {solve_method!r}.", + ) + errorif( + solve_method in {"fixed_point", "gmres"} and is_interior_neumann, + msg=f"solve_method={solve_method!r} is not supported for interior Neumann " + "problems. Use solve_method='direct' instead.", + ) + return options._replace(solve_method=solve_method) def _D_plus_half( @@ -142,19 +170,10 @@ def _D_plus_half( return result -def _least_squares_compute_potential( - boundary_condition, - potential_data, - source_data, - interpolator, - basis, - problem, - chunk_size=None, - _D_quad=False, - **kwargs, +@eqx.filter_jit +def _direct_solve( + boundary_condition, potential_data, source_data, interpolator, basis, options ): - assert problem in {"interior Neumann", "exterior Neumann", "interior Dirichlet"} - potential_grid = interpolator.eval_grid source_grid = interpolator.source_grid @@ -172,7 +191,12 @@ def _least_squares_compute_potential( Phi = basis.evaluate(potential_grid) potential_data["Phi(x) (periodic)"] = Phi source_data["Phi (periodic)"] = ( - Phi if (potential_grid == source_grid) else basis.evaluate(source_grid) + Phi + if ( + potential_grid.num_theta == source_grid.num_theta + and potential_grid.num_zeta == source_grid.num_zeta + ) + else basis.evaluate(source_grid) ) D = _D_plus_half( @@ -180,12 +204,12 @@ def _least_squares_compute_potential( source_data, interpolator, basis, - chunk_size, + options.chunk_size, prune_data=False, - _D_quad=_D_quad, + _D_quad=options.D_quad, ) assert D.shape == (potential_grid.num_nodes, basis.num_modes) - if problem == "exterior Neumann" or problem == "interior Dirichlet": + if options.problem in ("exterior Neumann", "interior Dirichlet"): D -= Phi if not well_posed: well_posed = None @@ -200,62 +224,10 @@ def _least_squares_compute_potential( ).value -def _iteration_operator(Phi, args): - """Equation 3.12 in [1].""" - gamma, potential_data, source_data, interpolator, chunk_size, xi = args - potential_data["Phi(x) (periodic)"] = Phi - source_data["Phi (periodic)"] = Phi - return ( - _D_plus_half( - potential_data, - source_data, - interpolator, - chunk_size=chunk_size, - prune_data=False, - ) - + (xi - 1) * Phi - - gamma - ) / xi - - -def _linear_potential_operator( - Phi, - potential_data, - source_data, - interpolator, - chunk_size, -): - """Equation solved by the iterative linear solver.""" - potential_data["Phi(x) (periodic)"] = Phi - source_data["Phi (periodic)"] = Phi - return ( - _D_plus_half( - potential_data, - source_data, - interpolator, - chunk_size=chunk_size, - prune_data=False, - ) - - Phi - ) - - @eqx.filter_jit -def _fixed_point_potential( - boundary_condition, - potential_data, - source_data, - interpolator, - Phi_0=None, - *, - xtol=1e-8, - maxiter=10, - solve_method="gmres", - full_output=False, - chunk_size=None, +def _iterative_solve( + boundary_condition, potential_data, source_data, interpolator, options ): - assert solve_method in {"gmres", "fixed_point"} - potential_grid = interpolator.eval_grid source_grid = interpolator.source_grid @@ -266,56 +238,124 @@ def _fixed_point_potential( source_grid, _kernel_dipole_plus_half, ) + Phi_0 = options.Phi_0 if Phi_0 is None: Phi_0 = jnp.ones(potential_grid.num_nodes) + assert Phi_0.size == potential_grid.num_nodes - if solve_method == "gmres": + if options.solve_method == "gmres": operator = lx.FunctionLinearOperator( partial( _linear_potential_operator, potential_data=potential_data, source_data=source_data, interpolator=interpolator, - chunk_size=chunk_size, + chunk_size=options.chunk_size, ), jax.ShapeDtypeStruct(Phi_0.shape, Phi_0.dtype), ) solution = lx.linear_solve( operator, boundary_condition, - solver=lx.GMRES(rtol=xtol, atol=xtol, max_steps=maxiter), + solver=lx.GMRES( + rtol=options.rtol, + atol=options.atol, + max_steps=options.max_steps, + ), options={"y0": Phi_0}, throw=False, ) - if full_output: + if options.full_output: err = operator.mv(solution.value) - boundary_condition return solution.value, (err, solution.stats["num_steps"]) return solution.value + # Some JAX versions fail to transpose scan, so we keep fixed point. xi = 2 / 3 args = ( boundary_condition, potential_data, source_data, interpolator, - chunk_size, + options.chunk_size, xi, ) solution = optx.fixed_point( _iteration_operator, - optx.FixedPointIteration(rtol=xtol, atol=xtol), + optx.FixedPointIteration(rtol=options.rtol, atol=options.atol), Phi_0, args, - max_steps=maxiter if maxiter > 0 else None, - adjoint=optx.ImplicitAdjoint(lx.GMRES(rtol=xtol, atol=xtol, max_steps=maxiter)), + max_steps=options.max_steps, + adjoint=optx.ImplicitAdjoint( + lx.GMRES( + rtol=options.rtol, + atol=options.atol, + max_steps=options.max_steps, + ) + ), throw=False, ) - if full_output: + if options.full_output: err = _iteration_operator(solution.value, args) - solution.value return solution.value, (err, solution.stats["num_steps"]) return solution.value +def _iteration_operator(Phi, args): + """Equation 3.12 in [1].""" + gamma, potential_data, source_data, interpolator, chunk_size, xi = args + potential_data["Phi(x) (periodic)"] = Phi + source_data["Phi (periodic)"] = _interp( + Phi, interpolator.eval_grid, interpolator.source_grid + ) + return ( + _D_plus_half( + potential_data, + source_data, + interpolator, + chunk_size=chunk_size, + prune_data=False, + ) + + (xi - 1) * Phi + - gamma + ) / xi + + +def _linear_potential_operator( + Phi, potential_data, source_data, interpolator, chunk_size +): + """Equation solved by the iterative linear solver.""" + potential_data["Phi(x) (periodic)"] = Phi + source_data["Phi (periodic)"] = _interp( + Phi, interpolator.eval_grid, interpolator.source_grid + ) + return ( + _D_plus_half( + potential_data, + source_data, + interpolator, + chunk_size=chunk_size, + prune_data=False, + ) + - Phi + ) + + +def _interp(x, input_grid, output_grid): + if ( + input_grid.num_theta == output_grid.num_theta + and input_grid.num_zeta == output_grid.num_zeta + ): + return x + return rfft_interp2d( + input_grid.meshgrid_reshape(x, "rtz")[0], + output_grid.num_theta, + output_grid.num_zeta, + dx=2 * jnp.pi / input_grid.num_theta, + dy=2 * jnp.pi / input_grid.num_zeta / input_grid.NFP, + ).ravel(order="F") + + @register_compute_fun( name="interpolator", label="", @@ -344,26 +384,14 @@ def _interpolator(params, transforms, profiles, data, **kwargs): potential_grid = kwargs.get("potential_grid", grid) data["interpolator"] = get_interpolator(potential_grid, grid, data, **kwargs) - if potential_grid == grid: - data["potential data"] = apply(data, subset=("R", "phi", "Z")) - else: - dt = 2 * jnp.pi / grid.num_theta - dz = 2 * jnp.pi / grid.num_zeta / grid.NFP - - # TODO: just interpolate Rb_mn, Zb_mn, and omegab_mn onto potential grid - # to avoid interpolation on oversampled grid - def fun(x): - return rfft_interp2d( - grid.meshgrid_reshape(x, "rtz")[0], - potential_grid.num_theta, - potential_grid.num_zeta, - dx=dt, - dy=dz, - ).ravel(order="F") - - data["potential data"] = apply(data, fun, ("R", "omega", "Z")) - zeta = potential_grid.nodes[:, 2] - data["potential data"]["phi"] = zeta + data["potential data"]["omega"] + # TODO: interpolate Rb_mn, Zb_mn, and omegab_mn directly + data["potential data"] = { + "R": _interp(data["R"], grid, potential_grid), + "omega": _interp(data["omega"], grid, potential_grid), + "Z": _interp(data["Z"], grid, potential_grid), + } + zeta = potential_grid.nodes[:, 2] + data["potential data"]["phi"] = zeta + data["potential data"]["omega"] return data @@ -403,17 +431,18 @@ def _potential_grid_position(params, transforms, profiles, data, **kwargs): resolution_requirement="tz", grid_requirement={"can_fft2": True}, parameterization="desc.magnetic_fields._laplace.SourceFreeField", - chunk_size=_doc["chunk_size"], + options=Options.__doc__, public=False, ) def _S_B0_n(params, transforms, profiles, data, **kwargs): # noqa: unused dependency + options = kwargs.get("options", Options()) data["S[B0*n]"] = singular_integral( data.get("potential data", data), data, data["interpolator"], _kernel_monopole, - chunk_size=kwargs.get("chunk_size", None), + chunk_size=options.chunk_size, ).squeeze(-1) return data @@ -434,41 +463,35 @@ def _S_B0_n(params, transforms, profiles, data, **kwargs): resolution_requirement="tz", grid_requirement={"can_fft2": True}, parameterization="desc.magnetic_fields._laplace.SourceFreeField", - problem='str : Problem to solve in {"interior Neumann", "exterior Neumann"}.', - **_doc, + options=Options.__doc__, ) def _scalar_potential_mn_Neumann(params, transforms, profiles, data, **kwargs): # noqa: unused dependency - solve_method = _select_potential_solve_method( - kwargs.get("solve_method", "auto"), - kwargs.get("maxiter", 10), - kwargs["problem"], - ) + options = Options.select_solver(kwargs.get("options", Options())) - if solve_method in {"fixed_point", "gmres"}: - solve_kwargs = apply(kwargs, subset=_doc) - solve_kwargs["solve_method"] = solve_method - data["Phi (periodic)"] = _fixed_point_potential( + if options.solve_method == "direct": + data["Phi_mn"] = _direct_solve( data["S[B0*n]"], - data, + data.get("potential data", data), data, data["interpolator"], - **solve_kwargs, + transforms["Phi"].basis, + options, ) - if kwargs.get("full_output", False): - data["Phi (periodic)"], (err, data["num iter"]) = data["Phi (periodic)"] - data["Phi error"] = jnp.abs(err).max() - - data["Phi_mn"] = transforms["Phi"].fit(data["Phi (periodic)"]) else: - data["Phi_mn"] = _least_squares_compute_potential( + data["Phi (periodic)"] = _iterative_solve( data["S[B0*n]"], data.get("potential data", data), data, data["interpolator"], - transforms["Phi"].basis, - **kwargs, + options, ) + if options.full_output: + data["Phi (periodic)"], (err, data["num_steps"]) = data["Phi (periodic)"] + data["Phi error"] = jnp.abs(err).max() + + assert data["Phi (periodic)"].size == transforms["Phi"].grid.num_nodes + data["Phi_mn"] = transforms["Phi"].fit(data["Phi (periodic)"]) return data @@ -553,8 +576,8 @@ def _pot_Phi_z_periodic(params, transforms, profiles, data, **kwargs): ) def _virtual_surface_current_periodic(params, transforms, profiles, data, **kwargs): data["K_vc (periodic)"] = -( - data["Phi_t (periodic)"][:, jnp.newaxis] * data["n_rho x grad(theta)"] - + data["Phi_z (periodic)"][:, jnp.newaxis] * data["n_rho x grad(zeta)"] + data["Phi_t (periodic)"][:, None] * data["n_rho x grad(theta)"] + + data["Phi_z (periodic)"][:, None] * data["n_rho x grad(zeta)"] ) return data @@ -603,11 +626,11 @@ def _Phi_error(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="num iter", - label="\\text{number of iterations}", + name="num_steps", + label="\\text{number of steps}", units="", units_long="", - description="Magnetic scalar potential number of iterations for inversion", + description="Magnetic scalar potential number of steps for inversion", dim=0, coordinates="", params=[], @@ -617,7 +640,7 @@ def _Phi_error(params, transforms, profiles, data, **kwargs): parameterization="desc.magnetic_fields._laplace.SourceFreeField", public=False, ) -def _Phi_num_iter(params, transforms, profiles, data, **kwargs): +def _Phi_num_steps(params, transforms, profiles, data, **kwargs): # noqa: unused dependency return data @@ -676,8 +699,8 @@ def _pot_Phi_z(params, transforms, profiles, data, **kwargs): ) def _virtual_surface_current(params, transforms, profiles, data, **kwargs): data["K_vc"] = -( - data["Phi_t"][:, jnp.newaxis] * data["n_rho x grad(theta)"] - + data["Phi_z"][:, jnp.newaxis] * data["n_rho x grad(zeta)"] + data["Phi_t"][:, None] * data["n_rho x grad(theta)"] + + data["Phi_z"][:, None] * data["n_rho x grad(zeta)"] ) return data @@ -717,19 +740,18 @@ def _K_vc_squared(params, transforms, profiles, data, **kwargs): resolution_requirement="tz", grid_requirement={"can_fft2": True}, parameterization="desc.magnetic_fields._laplace.SourceFreeField", - chunk_size=_doc["chunk_size"], + options=Options.__doc__, eval_interpolator="""_BIESTInterpolator : Interpolator from source grid to evaluation grid on boundary. If not given, default is to interpolate to source grid. """, - problem='str : Problem to solve in {"interior Neumann", "exterior Neumann"}.', on_boundary="bool : Whether RpZcoords are on boundary surface.", public=False, ) def _grad_potential(params, transforms, profiles, data, RpZ_data, **kwargs): # noqa: unused dependency - chunk_size = kwargs.get("chunk_size", None) - sign = 1 - 2 * int("exterior" in kwargs.get("problem", "")) + options = kwargs.get("options", Options()) + sign = 1 - 2 * int("exterior" in options.problem) if kwargs["on_boundary"]: RpZ_data["∇φ"] = ( @@ -740,7 +762,7 @@ def _grad_potential(params, transforms, profiles, data, RpZ_data, **kwargs): data, kwargs.get("eval_interpolator", data.get("interpolator", None)), _kernel_BS_plus_grad_S, - chunk_size=chunk_size, + chunk_size=options.chunk_size, ) ) else: @@ -756,7 +778,7 @@ def _grad_potential(params, transforms, profiles, data, RpZ_data, **kwargs): st=jnp.nan, sz=jnp.nan, kernel=_kernel_BS_plus_grad_S, - chunk_size=chunk_size, + chunk_size=options.chunk_size, ) return RpZ_data @@ -795,14 +817,15 @@ def _total_B(params, transforms, profiles, data, RpZ_data, **kwargs): parameterization="desc.magnetic_fields._laplace.SourceFreeField", B0="_MagneticField : Field object to compute with.", field_grid="Grid : Source grid used to compute magnetic field.", - chunk_size=_doc["chunk_size"], + options=Options.__doc__, ) def _B0_dot_n(params, transforms, profiles, data, **kwargs): + options = kwargs.get("options", Options()) data["B0*n"] = dot( kwargs["B0"].compute_magnetic_field( coords=data["x"], source_grid=kwargs.get("field_grid", None), - chunk_size=kwargs.get("chunk_size", None), + chunk_size=options.chunk_size, ), data["n_rho"], ) @@ -824,15 +847,16 @@ def _B0_dot_n(params, transforms, profiles, data, **kwargs): parameterization="desc.magnetic_fields._laplace.SourceFreeField", B0="_MagneticField : Field object to compute with.", field_grid="Grid : Source grid used to compute magnetic field.", - chunk_size=_doc["chunk_size"], + options=Options.__doc__, public=False, ) def _B0_field(params, transforms, profiles, data, RpZ_data, **kwargs): + options = kwargs.get("options", Options()) coords = jnp.column_stack([RpZ_data["R"], RpZ_data["phi"], RpZ_data["Z"]]) RpZ_data["B0"] = kwargs["B0"].compute_magnetic_field( coords=coords, source_grid=kwargs.get("field_grid", None), - chunk_size=kwargs.get("chunk_size", None), + chunk_size=options.chunk_size, ) return RpZ_data @@ -850,15 +874,16 @@ def _B0_field(params, transforms, profiles, data, RpZ_data, **kwargs): profiles=[], data=["x"], parameterization="desc.magnetic_fields._laplace.SourceFreeField", - B_coil_chunk_size=_doc["chunk_size"], + options=Options.__doc__, B_coil="_MagneticField : Field object to compute with.", field_grid="Grid : Source grid used to compute magnetic field.", ) def _B_coil_field(params, transforms, profiles, data, **kwargs): + options = kwargs.get("options", Options()) data["B_coil"] = kwargs["B_coil"].compute_magnetic_field( coords=data["x"], source_grid=kwargs.get("field_grid", None), - chunk_size=kwargs.get("B_coil_chunk_size", None), + chunk_size=options.B_coil_chunk_size, ) return data @@ -894,7 +919,7 @@ def _n_rho_x_B_coil(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], data=["e_zeta", "B_coil"], - chunk_size=_doc["chunk_size"], + options=Options.__doc__, parameterization="desc.magnetic_fields._laplace.FreeSurfaceOuterField", ) def _Y_coil(params, transforms, profiles, data, **kwargs): @@ -932,13 +957,13 @@ def _Phi_mn_coil(params, transforms, profiles, data, **kwargs): basis = transforms["Phi_coil"].basis # could compute these in objective build # and avoid computing if they are passed in as kwargs - _t = basis.evaluate(grid, [0, 1, 0])[:, jnp.newaxis] - _z = basis.evaluate(grid, [0, 0, 1])[:, jnp.newaxis] + _t = basis.evaluate(grid, [0, 1, 0])[:, None] + _z = basis.evaluate(grid, [0, 0, 1])[:, None] mat = lx.MatrixLinearOperator( ( - _t * data["n_rho x grad(theta)"][..., jnp.newaxis] - + _z * data["n_rho x grad(zeta)"][..., jnp.newaxis] + _t * data["n_rho x grad(theta)"][..., None] + + _z * data["n_rho x grad(zeta)"][..., None] ).reshape(grid.num_nodes * 3, basis.num_modes) ) @@ -1024,43 +1049,38 @@ def _Phi_coil(params, transforms, profiles, data, **kwargs): resolution_requirement="tz", grid_requirement={"can_fft2": True}, parameterization="desc.magnetic_fields._laplace.FreeSurfaceOuterField", - **_doc, + options=Options.__doc__, ) def _scalar_potential_mn_free_surface(params, transforms, profiles, data, **kwargs): # noqa: unused dependency - - boundary_condition = data["S[B0*n]"] - data["Phi_coil (periodic)"] - solve_method = _select_potential_solve_method( - kwargs.get("solve_method", "auto"), - kwargs.get("maxiter", 10), - "interior Dirichlet", + options = Options.select_solver( + kwargs.get("options", Options())._replace(problem="interior Dirichlet") ) - if solve_method in {"fixed_point", "gmres"}: - solve_kwargs = apply(kwargs, subset=_doc) - solve_kwargs["solve_method"] = solve_method - data["Phi (periodic)"] = _fixed_point_potential( + boundary_condition = data["S[B0*n]"] - data["Phi_coil (periodic)"] + if options.solve_method == "direct": + data["Phi_mn"] = _direct_solve( boundary_condition, - data, + data.get("potential data", data), data, data["interpolator"], - **solve_kwargs, + transforms["Phi"].basis, + options, ) - if kwargs.get("full_output", False): - data["Phi (periodic)"], (err, data["num iter"]) = data["Phi (periodic)"] - data["Phi error"] = jnp.abs(err).max() - - data["Phi_mn"] = transforms["Phi"].fit(data["Phi (periodic)"]) else: - data["Phi_mn"] = _least_squares_compute_potential( + data["Phi (periodic)"] = _iterative_solve( boundary_condition, data.get("potential data", data), data, data["interpolator"], - transforms["Phi"].basis, - problem="interior Dirichlet", - **kwargs, + options, ) + if options.full_output: + data["Phi (periodic)"], (err, data["num_steps"]) = data["Phi (periodic)"] + data["Phi error"] = jnp.abs(err).max() + + assert data["Phi (periodic)"].size == transforms["Phi"].grid.num_nodes + data["Phi_mn"] = transforms["Phi"].fit(data["Phi (periodic)"]) return data @@ -1079,11 +1099,12 @@ def _scalar_potential_mn_free_surface(params, transforms, profiles, data, **kwar resolution_requirement="tz", grid_requirement={"can_fft2": True}, parameterization="desc.magnetic_fields._laplace.FreeSurfaceOuterField", - chunk_size=_doc["chunk_size"], + options=Options.__doc__, public=False, ) def _gamma_potential(params, transforms, profiles, data, **kwargs): # noqa: unused dependency + options = kwargs.get("options", Options()) data["Phi(x) (periodic)"] = data["Phi (periodic)"] # Left hand side of equation 5.15 in [1] computed by evaluating # the right hand side. This is used for testing. @@ -1091,6 +1112,6 @@ def _gamma_potential(params, transforms, profiles, data, **kwargs): data, data, data["interpolator"], - chunk_size=kwargs.get("chunk_size", None), + chunk_size=options.chunk_size, ) return data diff --git a/desc/magnetic_fields/__init__.py b/desc/magnetic_fields/__init__.py index c018bbb11e..8f2498926c 100644 --- a/desc/magnetic_fields/__init__.py +++ b/desc/magnetic_fields/__init__.py @@ -1,5 +1,7 @@ """Classes for Magnetic Fields.""" +from desc.compute._laplace import Options + from ._core import ( MagneticFieldFromUser, OmnigenousField, diff --git a/desc/magnetic_fields/_laplace.py b/desc/magnetic_fields/_laplace.py index 53add07ce2..07ba623c3d 100644 --- a/desc/magnetic_fields/_laplace.py +++ b/desc/magnetic_fields/_laplace.py @@ -1,8 +1,8 @@ -"""High order accurate multiply connected geometry Laplace solver. +"""High order accurate multiply connected geometry Laplace solver as described in [1]_. References ---------- - [1] Unalmis et al. New high-order accurate free surface stellarator +.. [1] Unalmis et al. New high-order accurate free surface stellarator equilibria optimization and boundary integral methods in DESC. """ @@ -18,7 +18,7 @@ class SourceFreeField(FourierRZToroidalSurface): """Compute source free magnetic fields. Implements the Neumann formulation in multiply connected - geometry described in [1]. + geometry described in [1]_. Let 𝒳 be an open set with continuously differentiable closed boundary ∂𝒳. This class solves the following @@ -235,7 +235,7 @@ class FreeSurfaceOuterField(SourceFreeField): """Compute field on outer plasma for free surface. Implements the interior Dirichlet formulation in multiply connected - geometry described in [1]. + geometry described in [1]_. Parameters ---------- diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index 08a000be8b..f3c0dea027 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -6,11 +6,14 @@ from desc.backend import jnp from desc.compute import get_profiles, get_transforms +from desc.compute._laplace import Options as LaplaceOptions from desc.compute.utils import _compute as compute_fun from desc.grid import LinearGrid from desc.integrals import get_interpolator, virtual_casing_biot_savart +from desc.io import IOAble from desc.nestor import Nestor from desc.objectives.objective_funs import _Objective, collect_docs +from desc.optimizable import Optimizable, optimizable_parameter from desc.utils import ( PRINT_WIDTH, Timer, @@ -25,6 +28,27 @@ from .normalization import compute_scaling_factors +class _FreeSurfaceSheetCurrent(IOAble, Optimizable): + """Optimizable toroidal sheet-current parameter for FreeSurfaceError.""" + + _io_attrs_ = ["_I_sheet"] + + def __init__(self): + self.I_sheet = 0.0 + + @optimizable_parameter + @property + def I_sheet(self): + """float: Net toroidal sheet current determining a circulation of Phi.""" + return self._I_sheet + + @I_sheet.setter + def I_sheet(self, new): + new = jnp.asarray(new) + assert new.size == 1 + self._I_sheet = new.squeeze() + + class VacuumBoundaryError(_Objective): """Target for free boundary conditions on LCFS for vacuum equilibrium. @@ -873,17 +897,17 @@ def _print(fmt, fmax, fmin, fmean, f0max, f0min, f0mean, norm, unit): class FreeSurfaceError(_Objective): - """Target for free surface ideal MHD equilirium conditions on LCFS. + """Target for free surface ideal MHD equilirium as described in [1]_. References ---------- - [1] Unalmis et al. New high-order accurate free surface stellarator - equilibria optimization and boundary integral methods in DESC. + .. [1] Unalmis et al. New high-order accurate free surface stellarator + equilibria optimization and boundary integral methods in DESC. Notes ----- - Performance is expected to improve significantly by resolving GitHub issues - #1034 and #2171. + Performance is expected to improve significantly by resolving GitHub + issues #1034 and #2171. Parameters ---------- @@ -894,10 +918,20 @@ class FreeSurfaceError(_Objective): assumes ``field._B_coil`` is the magnetic field due to coils. If is an instance of ``SourceFreeField`` then assumes ``field._B0`` is the magnetic field due to coils. + The net toroidal sheet current ``I_sheet`` is an optimizable scalar + parameter initialized to zero. eval_grid : Grid Evaluation points on boundary to evaluate objective error. - Default is ``grid``, but it is likely best to use a grid - with much lower resolution. + Also implicitly determines the Fourier resolution of the potential. + Default is ``grid``, but it is better to use a grid with much lower resolution. + + If reverse mode differentiation is being used, it is of great benefit for + the objective residual to be a lower dimensional item. In such cases, to avoid + reducing the Fourier resolution of the potential (and hence reducing the + accuracy of the objective), it is better to instead use a loss function + that reduces the dimension of the residual before computing the derivative + relevant for optimization. This can be a mean squared error over all points + of the evaluation grid or mean absolute error over blocks of the grid, etc. grid : Grid Grid for the integral transforms. Tensor-product grid in (θ, ζ) with uniformly spaced nodes @@ -908,31 +942,8 @@ class FreeSurfaceError(_Objective): Default is default grid of coil magnetic field. q : int Order of integration on the local singular grid. - xtol : float - Absolute and relative error for the linear solve. Default is ``1e-8``. - maxiter : int - Maximum number of iterations for iterative scalar potential solves. - If non-positive then the materialized matrix will be inverted instead. - If positive, then performs that many iterations until ``maxiter`` or an - error tolerance of ``xtol`` is reached. Ten iterations should suffice for - the default GMRES solve. Default is ``10``. - solve_method : str - Method to use for the scalar potential solve. One of ``"auto"``, - ``"fixed_point"``, ``"gmres"``, or ``"least_squares"``. - Default is ``"gmres"``, initialized from one GMRES solve during ``build``. - chunk_size : int or None - Size to split integral computation into chunks. - If no chunking should be done or the chunk size is the full input - then supply ``None``. Default is ``None``. - Recommend to verify computation with ``chunk_size`` set to a - small number due to bugs in JAX or XLA. - B_coil_chunk_size : int or None - Size to split coil integral computation into chunks. - If no chunking should be done or the chunk size is the full input - then supply ``None``. Default is ``None``. - I_sheet : float - Net toroidal sheet current determining a circulation of Φ. - Default is zero. + options : LaplaceOptions + Options for the Laplace solver. """ @@ -951,13 +962,10 @@ class FreeSurfaceError(_Objective): "_B_coil", "_use_same_grid", "_q", - "_xtol", - "_maxiter", - "_solve_method", - "_chunk_size", - "_B_coil_chunk_size", + "_options", "_grad_keys", "_inner_keys", + "_I_sheet", "_reuseable_keys", ] @@ -972,12 +980,7 @@ def __init__( grid=None, coil_grid=None, q=None, - xtol=1e-8, - maxiter=10, - solve_method="gmres", - chunk_size=None, - B_coil_chunk_size=None, - I_sheet=0.0, + options=None, target=None, bounds=None, weight=1, @@ -1014,6 +1017,29 @@ def __init__( msg=f"N_Phi_coil = {getattr(field, 'N_Phi_coil', 0)} > {grid.N} = grid.N.", ) eval_grid = setdefault(eval_grid, grid) + assert eval_grid.can_fft2 + errorif( + field.M_Phi > eval_grid.M, + msg=f"M_Phi = {field.M_Phi} > {eval_grid.M} = eval_grid.M.", + ) + errorif( + field.N_Phi > eval_grid.N, + msg=f"N_Phi = {field.N_Phi} > {eval_grid.N} = eval_grid.N.", + ) + errorif( + not self._is_neumann and field.M_Phi_coil > eval_grid.M, + msg=( + f"M_Phi_coil = {getattr(field, 'M_Phi_coil', 0)} > " + f"{eval_grid.M} = eval_grid.M." + ), + ) + errorif( + not self._is_neumann and field.N_Phi_coil > eval_grid.N, + msg=( + f"N_Phi_coil = {getattr(field, 'N_Phi_coil', 0)} > " + f"{eval_grid.N} = eval_grid.N." + ), + ) self._field = field self._B_coil = field._B0 if self._is_neumann else field._B_coil @@ -1022,18 +1048,32 @@ def __init__( self._coil_grid = coil_grid self._use_same_grid = grid.equiv(eval_grid) self._q = q - self._xtol = xtol - self._maxiter = maxiter - self._solve_method = solve_method - self._chunk_size = chunk_size - self._B_coil_chunk_size = B_coil_chunk_size + self._I_sheet = _FreeSurfaceSheetCurrent() + if options is None: + options = LaplaceOptions() + options = options._replace( + problem="exterior Neumann" if self._is_neumann else "interior Dirichlet" + ) + self._options = options self._grad_keys = ["grad(theta)", "grad(zeta)", "n_rho"] - self._inner_keys = ["|B|^2", "p", "I", "|e_theta x e_zeta|"] + self._grad_keys + self._inner_keys = [ + "|B|^2", + "p", + "I", + "R", + "phi", + "omega", + "Z", + "|e_theta x e_zeta|", + ] + self._grad_keys self._reuseable_keys = [ "0", "R", + "phi", + "omega", "R_t", "R_z", + "Z", "Z_t", "Z_z", "e_theta", @@ -1044,10 +1084,9 @@ def __init__( "omega_z", "|e_theta x e_zeta|", ] - self._I_sheet = I_sheet super().__init__( - things=eq, + things=[eq, self._I_sheet], target=target, bounds=bounds, weight=weight, @@ -1079,7 +1118,9 @@ def build(self, use_jit=True, verbose=1): grad_transforms = eq_transforms else: source_transforms = get_transforms("Phi_mn", self._field, grid=self._grid) - grad_transforms = get_transforms(self._grad_keys, eq, grid=self._grid) + grad_transforms = get_transforms( + self._grad_keys + ["phi", "omega", "Z"], eq, grid=self._grid + ) data, _ = self._field.compute( ["interpolator"] if self._is_neumann else ["interpolator", "Y_coil"], @@ -1087,7 +1128,8 @@ def build(self, use_jit=True, verbose=1): q=self._q, transforms=source_transforms, B_coil=self._B_coil, - B_coil_chunk_size=self._B_coil_chunk_size, + options=self._options, + potential_grid=self._eval_grid, ) # No net poloidal current in equation 4.13 of [1]. self._field.Y = 0.0 if self._is_neumann else data["Y_coil"] @@ -1095,6 +1137,9 @@ def build(self, use_jit=True, verbose=1): initial_guess = self._compute_initial_guess( eq, source_transforms, + eval_transforms, + profiles, + self._I_sheet.params_dict, data["interpolator"], ) @@ -1125,14 +1170,15 @@ def _compute_initial_guess( self, eq, source_transforms, + eval_transforms, + profiles, + I_sheet_params, interpolator, ): - """Compute the GMRES potential used to initialize iterative solves.""" + """Compute the potential used to initialize iterative solves.""" params = eq.params_dict source_grid = self._grid - source_keys = list( - dict.fromkeys(self._reuseable_keys + self._grad_keys + ["I"]) - ) + source_keys = self._reuseable_keys + ["grad(theta)", "grad(zeta)", "I"] source_data = eq.compute( source_keys, grid=source_grid, @@ -1140,11 +1186,14 @@ def _compute_initial_guess( field_params = { "R_lmn": params["Rb_lmn"], "Z_lmn": params["Zb_lmn"], - "I": source_data["I"][source_grid.unique_rho_idx[-1]] + self._I_sheet, + "I": source_data["I"][source_grid.unique_rho_idx[-1]] + + I_sheet_params["I_sheet"][0], "Y": self._field.Y, } data = {key: source_data[key] for key in self._reuseable_keys} data["interpolator"] = interpolator + if not self._use_same_grid: + data["potential data"] = eq.compute(["R", "phi", "Z"], grid=self._eval_grid) data["B0*n"] = self._phi_sec_dot_n(field_params, source_data) if self._is_neumann: data, _ = self._field.compute( @@ -1153,33 +1202,40 @@ def _compute_initial_guess( params=field_params, transforms=source_transforms, data=data, - B_coil_chunk_size=self._B_coil_chunk_size, + options=self._options, B_coil=self._B_coil, field_grid=self._coil_grid, ) data["B0*n"] += dot(data["B_coil"], data["n_rho"]) + elif not self._use_same_grid: + potential_field_data, _ = self._field.compute( + "Phi_coil (periodic)", + grid=self._eval_grid, + params=field_params, + transforms=eval_transforms, + data={"Y_coil": self._field.Y}, + options=self._options, + B_coil=self._B_coil, + field_grid=self._coil_grid, + ) + data["Phi_coil (periodic)"] = potential_field_data["Phi_coil (periodic)"] - problem = {"problem": "exterior Neumann"} if self._is_neumann else {} - data, _ = self._field.compute( - "Phi (periodic)", - grid=source_grid, - params=field_params, - transforms=source_transforms, + data = compute_fun( + self._field, + "Phi_mn", + field_params, + eval_transforms, + profiles, data=data, - xtol=self._xtol, - maxiter=self._maxiter, - solve_method="gmres", - chunk_size=self._chunk_size, - B_coil_chunk_size=self._B_coil_chunk_size, + options=self._options._replace(solve_method="gmres"), B_coil=self._B_coil, field_grid=self._coil_grid, - **problem, ) # We differentiate through the solution, not the initial guess, # so we stop the gradient for numerical stability. return stop_gradient(data["Phi (periodic)"]) - def compute(self, params, constants=None): + def compute(self, params, I_sheet_params, constants=None): """Compute boundary error. Parameters @@ -1187,6 +1243,8 @@ def compute(self, params, constants=None): params : dict Dictionary of equilibrium degrees of freedom, e.g. ``Equilibrium.params_dict``. + I_sheet_params : dict + Dictionary containing the optimizable sheet current ``I_sheet``. constants : dict Dictionary of constant data, e.g. transforms, profiles etc. Defaults to ``self.constants``. @@ -1199,6 +1257,7 @@ def compute(self, params, constants=None): """ constants = self._get_deprecated_constants(constants) eq = self.things[0] + options = LaplaceOptions(*self._options) inner = compute_fun( eq, @@ -1211,7 +1270,8 @@ def compute(self, params, constants=None): "R_lmn": params["Rb_lmn"], "Z_lmn": params["Zb_lmn"], # This is I_plasma + I_sheet. - "I": inner["I"][self._eval_grid.unique_rho_idx[-1]] + self._I_sheet, + "I": inner["I"][self._eval_grid.unique_rho_idx[-1]] + + I_sheet_params["I_sheet"][0], "Y": self._field.Y, } outer = {key: inner[key] for key in self._reuseable_keys} @@ -1224,12 +1284,28 @@ def compute(self, params, constants=None): constants["eval_transforms"], constants["profiles"], data=outer, - B_coil_chunk_size=self._B_coil_chunk_size, + options=options, + B_coil=self._B_coil, + field_grid=self._coil_grid, + ) + elif not self._use_same_grid: + outer = compute_fun( + self._field, + "Phi_coil (periodic)", + field_params, + constants["eval_transforms"], + constants["profiles"], + data=outer, + options=options, B_coil=self._B_coil, field_grid=self._coil_grid, ) - problem = {"problem": "exterior Neumann"} if self._is_neumann else {} + potential_data = ( + None + if self._use_same_grid + else {key: inner[key] for key in ["R", "phi", "Z"]} + ) if self._use_same_grid: outer["interpolator"] = constants["interpolator"] @@ -1237,15 +1313,20 @@ def compute(self, params, constants=None): if self._is_neumann: outer["B0*n"] += dot(outer["B_coil"], inner["n_rho"]) else: + source_keys = self._grad_keys + ["phi", "omega", "Z"] grads = compute_fun( eq, - self._grad_keys, + source_keys, params, constants["grad_transforms"], constants["profiles"], ) data = {key: grads[key] for key in self._reuseable_keys} data["interpolator"] = constants["interpolator"] + if potential_data is not None: + data["potential data"] = potential_data + if not self._is_neumann: + data["Phi_coil (periodic)"] = outer["Phi_coil (periodic)"] data["B0*n"] = self._phi_sec_dot_n(field_params, grads) if self._is_neumann: data = compute_fun( @@ -1255,7 +1336,7 @@ def compute(self, params, constants=None): constants["source_transforms"], constants["profiles"], data=data, - B_coil_chunk_size=self._B_coil_chunk_size, + options=options, B_coil=self._B_coil, field_grid=self._coil_grid, ) @@ -1265,18 +1346,12 @@ def compute(self, params, constants=None): self._field, "Phi_mn", field_params, - constants["source_transforms"], + constants["eval_transforms"], constants["profiles"], data=data, - xtol=self._xtol, - maxiter=self._maxiter, - solve_method=self._solve_method, - Phi_0=constants["initial_guess"], - chunk_size=self._chunk_size, - B_coil_chunk_size=self._B_coil_chunk_size, + options=options._replace(Phi_0=constants["initial_guess"]), B_coil=self._B_coil, field_grid=self._coil_grid, - **problem, )["Phi_mn"] outer = compute_fun( @@ -1286,15 +1361,9 @@ def compute(self, params, constants=None): constants["eval_transforms"], constants["profiles"], data=outer, - xtol=self._xtol, - maxiter=self._maxiter, - solve_method=self._solve_method, - Phi_0=constants["initial_guess"], - chunk_size=self._chunk_size, - B_coil_chunk_size=self._B_coil_chunk_size, + options=options._replace(Phi_0=constants["initial_guess"]), B_coil=self._B_coil, field_grid=self._coil_grid, - **problem, ) if self._is_neumann: outer["K_vc"] -= outer["n_rho x B_coil"] diff --git a/tests/test_integrals.py b/tests/test_integrals.py index a42da36391..b1b1d0ed6b 100644 --- a/tests/test_integrals.py +++ b/tests/test_integrals.py @@ -61,11 +61,9 @@ _kernel_nr_over_r3, ) from desc.integrals.surface_integral import _get_grid_surface -from desc.magnetic_fields import ( - FreeSurfaceOuterField, - SourceFreeField, - ToroidalMagneticField, -) +from desc.magnetic_fields import FreeSurfaceOuterField +from desc.magnetic_fields import Options as LaplaceOptions +from desc.magnetic_fields import SourceFreeField, ToroidalMagneticField from desc.transform import Transform from desc.utils import dot, errorif, rpz2xyz, safediv @@ -748,12 +746,28 @@ def compute_magnetic_field(self, coords, source_grid, chunk_size): @pytest.mark.unit @pytest.mark.parametrize( - "surface, M, N, solve_method, maxiter, chunk_size, just_err", + "surface, M, N, solve_method, max_steps, chunk_size, just_err", [ pytest.param( - None, 16, 16, "least_squares", -1, 500, False, id="least-squares" + None, + 16, + 16, + "direct", + 10, + 500, + False, + id="direct", + ), + pytest.param( + None, + 16, + 16, + "fixed_point", + 40, + 500, + False, + id="fixed-point", ), - pytest.param(None, 16, 16, "fixed_point", 40, 500, False, id="fixed-point"), pytest.param(None, 16, 16, "gmres", 40, 500, False, id="gmres"), ], ) @@ -763,10 +777,11 @@ def test_interior_Dirichlet( M, N, solve_method, - maxiter, + max_steps, chunk_size, just_err, - xtol=1e-12, + atol=1e-12, + rtol=1e-12, ): """Test multiply connected interior Dirichlet Laplace solver.""" if surface is None: @@ -786,21 +801,26 @@ def test_interior_Dirichlet( B_coil=TestLaplace._Z_hat_field(), ) assert field.M != grid.M and field.N != grid.N + + keys = ["Phi error", "num_steps"] if just_err else "γ potential" data, _ = field.compute( - ["Phi error", "num iter"] if just_err else "γ potential", + keys, grid, - maxiter=int(maxiter), - solve_method=solve_method, - full_output=True, - chunk_size=chunk_size, - xtol=xtol, + options=LaplaceOptions( + max_steps=int(max_steps), + solve_method=solve_method, + full_output=True, + chunk_size=chunk_size, + atol=atol, + rtol=rtol, + ), ) - if maxiter > 0: + if "num_steps" in data: print() - print(data["num iter"]) + print(data["num_steps"]) print(data["Phi error"]) if just_err: - return data["num iter"], data["Phi error"] + return data["num_steps"], data["Phi error"] np.testing.assert_allclose(data["Y_coil"], 0, atol=1e-12) np.testing.assert_allclose(data["Phi_coil (periodic)"], data["Z"]) np.testing.assert_allclose(data["γ potential"], data["Z"], atol=1e-6) @@ -811,22 +831,22 @@ def test_convergence_run_fixed_point( surface=get("W7-X").surface, M=30, N=30, - maxiter=np.array([1, 2, 3, 4, 5]), + max_steps=np.array([1, 2, 3, 4, 5]), chunk_size=1000, name="convergence-fp_W7-X", ): """Stores errors for potential in name.pkl for plotting analysis.""" - num_iter = [] + num_steps = [] Phi_err = [] print() - for i in maxiter: + for i in max_steps: n, e = self.test_interior_Dirichlet( - surface, M, N, "gmres", i, chunk_size, just_err=True + surface, M, N, "gmres", i, chunk_size, True ) - num_iter.append(n) + num_steps.append(n) Phi_err.append(e) - print(f"Resolution num iter={n} is done with error={e}.") - data = {"num iter": np.asarray(num_iter), "Phi error": np.asarray(Phi_err)} + print(f"Resolution num_steps={n} is done with error={e}.") + data = {"num_steps": np.asarray(num_steps), "Phi error": np.asarray(Phi_err)} with open(f"{name}.pkl", "wb") as file: pickle.dump(data, file) @@ -860,11 +880,11 @@ def test_convergence_plot_fixed_point(self, name="convergence-fp_W7-X"): fig, ax = plt.subplots() gmres_label = None # fp_label = r"$\xi=2/3$" # noqa: E800 - ax.semilogy(data["num iter"], data["Phi error"], marker="o", label=gmres_label) + ax.semilogy(data["num_steps"], data["Phi error"], marker="o", label=gmres_label) ax.axhline(1e-7, color="black", label="Stop tolerance") - ax.set_xlabel(r"Number of gmres iterations in inversion for $\Phi$") + ax.set_xlabel(r"Number of gmres steps in inversion for $\Phi$") ax.set_ylabel("Absolute error") - ax.set_title(r"$\Phi$ error vs. gmres iterations " + name.split("_", 1)[1]) + ax.set_title(r"$\Phi$ error vs. gmres steps " + name.split("_", 1)[1]) ax.legend(loc="upper right", frameon=True) fig.tight_layout() plt.savefig(f"{name}.pdf") @@ -906,12 +926,14 @@ def test_interior_Neumann( data=data, RpZ_data=RpZ_data, RpZ_grid=RpZ_grid, - problem="interior Neumann", - solve_method=solve_method, on_boundary=True, - maxiter=10, - chunk_size=chunk_size, - _D_quad=_D_quad, + options=LaplaceOptions( + problem="interior Neumann", + solve_method=solve_method, + max_steps=10, + chunk_size=chunk_size, + D_quad=_D_quad, + ), ) err = np.ptp(data["Z"] - data["Phi"]) if just_err: @@ -1012,7 +1034,7 @@ def test_convergence_plot(self, name="convergence_W7-X"): plt.savefig(f"{name}.pdf") @pytest.mark.unit - def test_exterior_Neumann(self, maxiter=30, chunk_size=1000): + def test_exterior_Neumann(self, max_steps=30, chunk_size=1000): """Test Laplacian solver in exterior.""" # Fourier spectrum of G(x) becomes very wide at large R0 (e.g. 10 is large). R0 = 2 @@ -1034,15 +1056,17 @@ def test_exterior_Neumann(self, maxiter=30, chunk_size=1000): ["∇φ", "Phi", "x", "n_rho"], grid, data=data, - problem="exterior Neumann", on_boundary=True, - maxiter=maxiter, - full_output=True, - chunk_size=chunk_size, + options=LaplaceOptions( + problem="exterior Neumann", + max_steps=max_steps, + full_output=True, + chunk_size=chunk_size, + ), basis="xyz", ) assert data is RpZ_data - print("num iterations:", data["num iter"]) + print("num steps :", data["num_steps"]) print("Phi error :", data["Phi error"]) np.testing.assert_allclose( @@ -1122,10 +1146,10 @@ def test_dommaschk_vacuum(self, chunk_size=50): data=data, RpZ_data=RpZ_data, RpZ_grid=RpZ_grid, - problem="interior Neumann", on_boundary=True, - maxiter=0, - chunk_size=chunk_size, + options=LaplaceOptions( + problem="interior Neumann", max_steps=10, chunk_size=chunk_size + ), warn_fft=False, ) np.testing.assert_allclose(dot(RpZ_data["B"], RpZ_data["n_rho"]), 0, atol=5e-4) @@ -1143,7 +1167,7 @@ def test_dommaschk_vacuum(self, chunk_size=50): data=data, RpZ_grid=LinearGrid(rho=0.5, M=10, N=10, NFP=eq.NFP), on_boundary=False, - chunk_size=chunk_size, + options=LaplaceOptions(chunk_size=chunk_size), ) assert np.isfinite(RpZ_data["B"]).all() diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 16a95e42b8..e4684e9ba2 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -26,6 +26,7 @@ initialize_modular_coils, ) from desc.compute import get_transforms +from desc.compute._laplace import Options as LaplaceOptions from desc.equilibrium import Equilibrium from desc.examples import get from desc.geometry import FourierPlanarCurve, FourierRZToroidalSurface, FourierXYZCurve @@ -2188,7 +2189,7 @@ def test_objective_against_compute_ballooning(self): np.testing.assert_allclose(obj.compute(eq.params_dict), lam) @pytest.mark.unit - @pytest.mark.parametrize("solve_method", ["fixed_point", "gmres", "least_squares"]) + @pytest.mark.parametrize("solve_method", ["fixed_point", "gmres", "direct"]) def test_objective_against_compute_free_surface_error(self, solve_method): """Test FreeSurfaceError against the underlying |K_vc|^2 compute quantity.""" eq = get("W7-X") @@ -2200,7 +2201,7 @@ def test_objective_against_compute_free_surface_error(self, solve_method): field, grid=grid, eval_grid=grid, - solve_method=solve_method, + options=LaplaceOptions(solve_method=solve_method), ) obj.build(verbose=0) @@ -2230,19 +2231,49 @@ def test_objective_against_compute_free_surface_error(self, solve_method): transforms=obj._constants["eval_transforms"], data=outer_data, override_grid=False, - xtol=obj._xtol, - maxiter=obj._maxiter, - solve_method=solve_method, - Phi_0=obj._constants["initial_guess"], - chunk_size=obj._chunk_size, - B_coil_chunk_size=obj._B_coil_chunk_size, + options=obj._options._replace( + solve_method=solve_method, + Phi_0=obj._constants["initial_guess"], + ), B_coil=B, ) expected = (outer["|K_vc|^2"] - inner["|B|^2"] - 2 * mu_0 * inner["p"]) * inner[ "|e_theta x e_zeta|" ] - np.testing.assert_allclose(obj.compute(eq.params_dict), expected) + np.testing.assert_allclose( + obj.compute(eq.params_dict, obj._I_sheet.params_dict), expected + ) + + @pytest.mark.unit + def test_free_surface_error_optimizes_sheet_current(self): + """Test FreeSurfaceError exposes I_sheet as an optimizable parameter.""" + eq = get("W7-X") + grid = LinearGrid(rho=np.array([1.0]), M=2, N=2, NFP=eq.NFP, sym=False) + field = FreeSurfaceOuterField( + eq.surface, M=1, N=1, B_coil=ToroidalMagneticField(5, 1) + ) + obj = ObjectiveFunction( + FreeSurfaceError( + eq, + field, + grid=grid, + eval_grid=grid, + options=LaplaceOptions(solve_method="direct"), + ) + ) + obj.build(verbose=0) + + assert obj.things[1].optimizable_params == ["I_sheet"] + x0 = obj.x() + idx = obj.things[0].dim_x + obj.things[1].x_idx["I_sheet"][0] + grad = obj.grad(x0) + assert np.isfinite(grad[idx]) + assert not np.isclose(grad[idx], 0) + + step = 1e-5 * np.sign(grad[idx]) + x1 = x0.at[idx].add(-step) + assert obj.compute_scalar(x1) < obj.compute_scalar(x0) @pytest.mark.unit def test_generic_with_kwargs(self): @@ -4167,9 +4198,9 @@ def test_objective_no_nangrad_free_surface_error(self, flag): eq = get("W7-X") B = ToroidalMagneticField(5, 1) field = ( - FreeSurfaceOuterField(eq.surface, 3, 3, B_coil=B) + FreeSurfaceOuterField(eq.surface, 2, 2, B_coil=B) if flag - else SourceFreeField(eq.surface, 3, 3, B0=B) + else SourceFreeField(eq.surface, 2, 2, B0=B) ) obj = ObjectiveFunction( FreeSurfaceError( @@ -4177,13 +4208,32 @@ def test_objective_no_nangrad_free_surface_error(self, flag): field, eval_grid=LinearGrid(M=2, N=2, NFP=eq.NFP), grid=LinearGrid(M=3, N=3, NFP=eq.NFP), - solve_method="fixed_point", + options=LaplaceOptions(solve_method="fixed_point"), ) ) obj.build() g = obj.grad(obj.x()) assert not np.any(np.isnan(g)), "free surface error" + @pytest.mark.unit + @pytest.mark.parametrize("flag", [True, False]) + def test_free_surface_error_field_resolution_eval_grid(self, flag): + """Free surface error requires Phi resolution to fit on eval_grid.""" + eq = get("W7-X") + B = ToroidalMagneticField(5, 1) + field = ( + FreeSurfaceOuterField(eq.surface, 3, 3, B_coil=B) + if flag + else SourceFreeField(eq.surface, 3, 3, B0=B) + ) + with pytest.raises(ValueError, match="M_Phi"): + FreeSurfaceError( + eq, + field, + eval_grid=LinearGrid(M=2, N=2, NFP=eq.NFP), + grid=LinearGrid(M=3, N=3, NFP=eq.NFP), + ) + @pytest.mark.unit def test_objective_no_nanjac_boundary_error_kinetic_profiles(self): """Test BoundaryError with kinetic profiles. Related to GH Issue #1712.""" From 4bd850235d6f090698d26b99e49dd614a8851456 Mon Sep 17 00:00:00 2001 From: unalmis Date: Sat, 30 May 2026 01:52:58 -0700 Subject: [PATCH 2/2] . --- desc/objectives/_free_boundary.py | 45 ++++++++++++++++++------------- tests/test_objective_funs.py | 2 +- 2 files changed, 28 insertions(+), 19 deletions(-) diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index f3c0dea027..94b1f887f3 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -922,8 +922,8 @@ class FreeSurfaceError(_Objective): parameter initialized to zero. eval_grid : Grid Evaluation points on boundary to evaluate objective error. - Also implicitly determines the Fourier resolution of the potential. - Default is ``grid``, but it is better to use a grid with much lower resolution. + Also determines the Fourier resolution of the potential. + Default is ``LinearGrid(M=field.M_Phi,N=field.N_Phi,NFP=grid.NFP,sym=False)``. If reverse mode differentiation is being used, it is of great benefit for the objective residual to be a lower dimensional item. In such cases, to avoid @@ -1016,15 +1016,19 @@ def __init__( not self._is_neumann and field.N_Phi_coil > grid.N, msg=f"N_Phi_coil = {getattr(field, 'N_Phi_coil', 0)} > {grid.N} = grid.N.", ) - eval_grid = setdefault(eval_grid, grid) + if eval_grid is None: + eval_grid = LinearGrid( + M=field.M_Phi, N=field.N_Phi, NFP=grid.NFP, sym=False + ) assert eval_grid.can_fft2 + errorif( - field.M_Phi > eval_grid.M, - msg=f"M_Phi = {field.M_Phi} > {eval_grid.M} = eval_grid.M.", + field.M_Phi != eval_grid.M, + msg=f"M_Phi = {field.M_Phi} != {eval_grid.M} = eval_grid.M.", ) errorif( - field.N_Phi > eval_grid.N, - msg=f"N_Phi = {field.N_Phi} > {eval_grid.N} = eval_grid.N.", + field.N_Phi != eval_grid.N, + msg=f"N_Phi = {field.N_Phi} != {eval_grid.N} = eval_grid.N.", ) errorif( not self._is_neumann and field.M_Phi_coil > eval_grid.M, @@ -1051,10 +1055,12 @@ def __init__( self._I_sheet = _FreeSurfaceSheetCurrent() if options is None: options = LaplaceOptions() + else: + options = LaplaceOptions(*options) options = options._replace( problem="exterior Neumann" if self._is_neumann else "interior Dirichlet" ) - self._options = options + self._options = tuple(options) # DESC is dumb and casts NamedTuples to Tuples self._grad_keys = ["grad(theta)", "grad(zeta)", "n_rho"] self._inner_keys = [ "|B|^2", @@ -1110,6 +1116,7 @@ def build(self, use_jit=True, verbose=1): """ eq = self.things[0] + options = LaplaceOptions(*self._options) eq_transforms = get_transforms(self._inner_keys, eq, grid=self._eval_grid) eval_transforms = get_transforms("|K_vc|^2", self._field, grid=self._eval_grid) @@ -1128,7 +1135,7 @@ def build(self, use_jit=True, verbose=1): q=self._q, transforms=source_transforms, B_coil=self._B_coil, - options=self._options, + options=options, potential_grid=self._eval_grid, ) # No net poloidal current in equation 4.13 of [1]. @@ -1176,6 +1183,7 @@ def _compute_initial_guess( interpolator, ): """Compute the potential used to initialize iterative solves.""" + options = LaplaceOptions(*self._options) params = eq.params_dict source_grid = self._grid source_keys = self._reuseable_keys + ["grad(theta)", "grad(zeta)", "I"] @@ -1202,7 +1210,7 @@ def _compute_initial_guess( params=field_params, transforms=source_transforms, data=data, - options=self._options, + options=options, B_coil=self._B_coil, field_grid=self._coil_grid, ) @@ -1214,7 +1222,7 @@ def _compute_initial_guess( params=field_params, transforms=eval_transforms, data={"Y_coil": self._field.Y}, - options=self._options, + options=options, B_coil=self._B_coil, field_grid=self._coil_grid, ) @@ -1222,12 +1230,12 @@ def _compute_initial_guess( data = compute_fun( self._field, - "Phi_mn", + "Phi (periodic)", field_params, eval_transforms, profiles, data=data, - options=self._options._replace(solve_method="gmres"), + options=options._replace(solve_method="gmres"), B_coil=self._B_coil, field_grid=self._coil_grid, ) @@ -1257,7 +1265,9 @@ def compute(self, params, I_sheet_params, constants=None): """ constants = self._get_deprecated_constants(constants) eq = self.things[0] - options = LaplaceOptions(*self._options) + options = LaplaceOptions(*self._options)._replace( + Phi_0=constants["initial_guess"] + ) inner = compute_fun( eq, @@ -1313,10 +1323,9 @@ def compute(self, params, I_sheet_params, constants=None): if self._is_neumann: outer["B0*n"] += dot(outer["B_coil"], inner["n_rho"]) else: - source_keys = self._grad_keys + ["phi", "omega", "Z"] grads = compute_fun( eq, - source_keys, + self._grad_keys + ["phi", "omega", "Z"], params, constants["grad_transforms"], constants["profiles"], @@ -1349,7 +1358,7 @@ def compute(self, params, I_sheet_params, constants=None): constants["eval_transforms"], constants["profiles"], data=data, - options=options._replace(Phi_0=constants["initial_guess"]), + options=options, B_coil=self._B_coil, field_grid=self._coil_grid, )["Phi_mn"] @@ -1361,7 +1370,7 @@ def compute(self, params, I_sheet_params, constants=None): constants["eval_transforms"], constants["profiles"], data=outer, - options=options._replace(Phi_0=constants["initial_guess"]), + options=options, B_coil=self._B_coil, field_grid=self._coil_grid, ) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index e4684e9ba2..9b772b85d4 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -2231,7 +2231,7 @@ def test_objective_against_compute_free_surface_error(self, solve_method): transforms=obj._constants["eval_transforms"], data=outer_data, override_grid=False, - options=obj._options._replace( + options=LaplaceOptions(*obj._options)._replace( solve_method=solve_method, Phi_0=obj._constants["initial_guess"], ),