Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

Add puddle objective to avoid local min in current potential #1507

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions desc/objectives/__init__.py
Original file line number Diff line number Diff line change
@@ -9,6 +9,7 @@
CoilSetLinkingNumber,
CoilSetMinDistance,
CoilTorsion,
CurrentPotentialPuddle,
LinkingCurrentConsistency,
PlasmaCoilSetMinDistance,
QuadraticFlux,
186 changes: 186 additions & 0 deletions desc/objectives/_coils.py
Original file line number Diff line number Diff line change
@@ -2279,3 +2279,189 @@

K_mag = safenorm(surface_data["K"], axis=-1)
return K_mag * jnp.sqrt(surface_data["|e_theta x e_zeta|"])


class CurrentPotentialPuddle(_Objective):
"""Target the current potential derivatives to avoid saddle contours.

compute::

|Φ_t| + |Φ_z|

where Φ is the winding surface current potential, and the
subscripts represent derivatives with respect to the
poloidal angle theta and the toroidal angle zeta.

This is intended to be used with a surface current::

K = n x ∇ Φ
Φ(θ,ζ) = Φₛᵥ(θ,ζ) + Gζ/2π + Iθ/2π

i.e. a FourierCurrentPotentialField in order to
prevent local extrema in the current potential, which
makes the potential more amenable to creating modular
or helical coils without the need for saddle coils.

Parameters
----------
surface_current_field : FourierCurrentPotentialField
Surface current which is producing the magnetic field, the parameters
of this will be optimized to minimize the objective.
target : {float, ndarray}, optional
Target value(s) of the objective. Only used if bounds is None.
Must be broadcastable to Objective.dim_f.
bounds : tuple of {float, ndarray}, optional
Lower and upper bounds on the objective. Overrides target.
Both bounds must be broadcastable to to Objective.dim_f
weight : {float, ndarray}, optional
Weighting to apply to the Objective, relative to other Objectives.
Must be broadcastable to to Objective.dim_f
When used with QuadraticFlux objective, this acts as the regularization
parameter, with 0 corresponding to no regularization. The larger this
parameter is, the less complex the surface current will be, but the
worse the normal field.
normalize : bool, optional
Whether to compute the error in physical units or non-dimensionalize.
Note: has no effect on this objective.
normalize_target : bool
Whether target should be normalized before comparing to computed values.
if `normalize` is `True` and the target is in physical units, this should also
be set to True.
Note: has no effect on this objective.
loss_function : {None, 'mean', 'min', 'max'}, optional
Loss function to apply to the objective values once computed. This loss function
is called on the raw compute value, before any shifting, scaling, or
normalization. Note: has no effect for this objective
deriv_mode : {"auto", "fwd", "rev"}
Specify how to compute jacobian matrix, either forward mode or reverse mode AD.
"auto" selects forward or reverse mode based on the size of the input and output
of the objective. Has no effect on self.grad or self.hess which always use
reverse mode and forward over reverse mode respectively.
source_grid : Grid, optional
Collocation grid containing the nodes to evaluate current source at on
the winding surface. If used in conjunction with the QuadraticFlux objective,
with the same ``source_grid``, this replicates the REGCOIL algorithm described
in [1]_.
name : str, optional
Name of the objective function.
"""

_coordinates = ""
_units = ""
_print_value_fmt = "Current Potential Gradient: {:10.3e} "

def __init__(
self,
surface_current_field,
target=None,
bounds=None,
weight=1,
normalize=True,
normalize_target=True,
loss_function=None,
deriv_mode="auto",
source_grid=None,
name="current-potential-puddle",
):
if target is None and bounds is None:
target = 0
assert hasattr(

Check warning on line 2368 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2366-L2368

Added lines #L2366 - L2368 were not covered by tests
surface_current_field, "Phi_mn"
), "surface_current_field must be a FourierCurrentPotentialField"
self._surface_current_field = surface_current_field
self._source_grid = source_grid

Check warning on line 2372 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2371-L2372

Added lines #L2371 - L2372 were not covered by tests

super().__init__(

Check warning on line 2374 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2374

Added line #L2374 was not covered by tests
things=[surface_current_field],
target=target,
bounds=bounds,
weight=weight,
normalize=normalize,
normalize_target=normalize_target,
loss_function=loss_function,
deriv_mode=deriv_mode,
name=name,
)

def build(self, use_jit=True, verbose=1):
"""Build constant arrays.

Parameters
----------
use_jit : bool, optional
Whether to just-in-time compile the objective and derivatives.
verbose : int, optional
Level of output.

"""
surface_current_field = self.things[0]

Check warning on line 2397 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2397

Added line #L2397 was not covered by tests

if self._source_grid is None:
source_grid = LinearGrid(

Check warning on line 2400 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2399-L2400

Added lines #L2399 - L2400 were not covered by tests
M=surface_current_field._M_Phi * 3 + 1,
N=surface_current_field._N_Phi * 3 + 1,
NFP=surface_current_field.NFP,
)
else:
source_grid = self._source_grid

Check warning on line 2406 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2406

Added line #L2406 was not covered by tests

if not np.allclose(source_grid.nodes[:, 0], 1):
warnings.warn("Source grid includes off-surface pts, should be rho=1")

Check warning on line 2409 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2408-L2409

Added lines #L2408 - L2409 were not covered by tests

# source_grid.num_nodes for the regularization cost
self._dim_f = source_grid.num_nodes
self._surface_data_keys = ["Phi_t", "Phi_z"]

Check warning on line 2413 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2412-L2413

Added lines #L2412 - L2413 were not covered by tests

timer = Timer()
if verbose > 0:
print("Precomputing transforms")
timer.start("Precomputing transforms")

Check warning on line 2418 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2415-L2418

Added lines #L2415 - L2418 were not covered by tests

surface_transforms = get_transforms(

Check warning on line 2420 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2420

Added line #L2420 was not covered by tests
self._surface_data_keys,
obj=surface_current_field,
grid=source_grid,
has_axis=source_grid.axis.size,
)

self._constants = {

Check warning on line 2427 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2427

Added line #L2427 was not covered by tests
"surface_transforms": surface_transforms,
}

timer.stop("Precomputing transforms")
if verbose > 1:
timer.disp("Precomputing transforms")

Check warning on line 2433 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2431-L2433

Added lines #L2431 - L2433 were not covered by tests

super().build(use_jit=use_jit, verbose=verbose)

Check warning on line 2435 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2435

Added line #L2435 was not covered by tests

def compute(self, surface_params=None, constants=None):
"""Compute surface current regularization.

Parameters
----------
surface_params : dict
Dictionary of surface degrees of freedom,
eg FourierCurrentPotential.params_dict
constants : dict
Dictionary of constant data, eg transforms, profiles etc. Defaults to
self.constants

Returns
-------
f : ndarray
The surface current density magnitude on the source surface.

"""
if constants is None:
constants = self.constants

Check warning on line 2456 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2455-L2456

Added lines #L2455 - L2456 were not covered by tests

surface_data = compute_fun(

Check warning on line 2458 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2458

Added line #L2458 was not covered by tests
self._surface_current_field,
self._surface_data_keys,
params=surface_params,
transforms=constants["surface_transforms"],
profiles={},
basis="xyz",
)

return jnp.abs(surface_data["Phi_t"]) + jnp.abs(surface_data["Phi_z"])

Check warning on line 2467 in desc/objectives/_coils.py

Codecov / codecov/patch

desc/objectives/_coils.py#L2467

Added line #L2467 was not covered by tests