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

ko/coils quad flux #726

Merged
merged 85 commits into from
Apr 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
85 commits
Select commit Hold shift + click to select a range
72f8dc5
initial adjustment of QuadraticFlux
kianorr Oct 25, 2023
b8e414d
updated constants
kianorr Oct 27, 2023
16e6356
changed test slightly
kianorr Oct 27, 2023
99a29a9
fixed bugs
kianorr Nov 4, 2023
749f101
formatting
kianorr Nov 8, 2023
c09d9b6
changed src_grid to eval_grid in weights
kianorr Nov 20, 2023
5413e37
updated tests
kianorr Nov 26, 2023
2dfca99
added precise_QA and ext_field to things
kianorr Dec 5, 2023
9acfde1
update test with more resolution
kianorr Dec 27, 2023
388f9b0
update source res in test
kianorr Dec 31, 2023
54ff47c
changed resolution
kianorr Dec 31, 2023
06a58b2
Merge branch 'rc/freeb' into ko/coils_quad_flux
kianorr Jan 3, 2024
2b00a3a
update init and add back singular_integral
kianorr Jan 3, 2024
de4ab18
updated documentation
kianorr Jan 3, 2024
bbc6b02
added eq_fixed option and logic
kianorr Jan 4, 2024
f231773
obtain eq differently in build
kianorr Jan 8, 2024
3ed48eb
add docs
kianorr Jan 8, 2024
324d9dd
add eq_params and field_params
kianorr Jan 8, 2024
cc6d3e0
added more resolution to grids and sym_Phi=sin
kianorr Jan 8, 2024
d8826d3
updated docstring with new params
kianorr Jan 9, 2024
cd5ee7f
update test with comments
kianorr Jan 10, 2024
f8a12aa
added field_fixed logic
kianorr Jan 11, 2024
747f7d0
fixed comparison tests and added more logic
kianorr Jan 11, 2024
ebc5a93
added more to field_fixed test
kianorr Jan 12, 2024
cbd6101
deleted eval_data under fixed_field
kianorr Jan 12, 2024
c48c683
updated field_fixed=True test (all tests passing)
kianorr Jan 14, 2024
af1b73d
refactored objective function
kianorr Jan 15, 2024
b3a754d
updated comments
kianorr Jan 15, 2024
7379b12
added constraints so that eq_fixed=field_fixed=False passes
kianorr Jan 16, 2024
05c9a0a
updated third test with nonaxiym field and eq
kianorr Jan 19, 2024
b11b825
refactored some of compute and changed docs
kianorr Jan 19, 2024
8fc54fc
added code to use params in MagneticFieldclasses
kianorr Jan 19, 2024
10425cc
added field_fixed=False with different field
kianorr Jan 19, 2024
de375cf
merge base and fix conflicts
kianorr Jan 23, 2024
dbc74c9
attempt to fix mistake with mean_iota function
kianorr Jan 23, 2024
9f951f4
use diffent QA eq for nonaxisymmetric case
kianorr Jan 24, 2024
077a2d1
delete int()s
kianorr Jan 24, 2024
faf4c82
Merge branch 'rc/freeb' into ko/coils_quad_flux
kianorr Jan 25, 2024
d83336d
update comments and docstring
kianorr Jan 26, 2024
0b30edb
add back int() for src grid
kianorr Jan 26, 2024
6be1706
fix nonaxisymmetric field modes
kianorr Jan 29, 2024
1d4e1c8
Merge branch 'rc/freeb' into ko/coils_quad_flux
kianorr Jan 30, 2024
355e2b9
separate tests into functions
kianorr Jan 31, 2024
80bdbae
delete xyz2rpz to fix bugs
kianorr Jan 31, 2024
4771fb5
add vacuum flag and test
kianorr Feb 1, 2024
9ee9116
make test class to reuse code
kianorr Feb 1, 2024
e03a088
add docs for vacuum flag
kianorr Feb 1, 2024
76e8880
solve equilibrium
kianorr Feb 1, 2024
cc6601a
update vacuum test
kianorr Feb 3, 2024
f2a9035
Merge branch 'rc/freeb' into ko/coils_quad_flux
kianorr Feb 3, 2024
91a235a
Merge branch 'rc/freeb' into ko/coils_quad_flux
kianorr Feb 9, 2024
0d3e2e6
add phi to data keys and adjust singular_integral
kianorr Feb 9, 2024
c50bc8c
Merge branch 'master' into ko/coils_quad_flux
f0uriest Feb 24, 2024
a73fceb
Merge branch 'master' into ko/coils_quad_flux
f0uriest Feb 25, 2024
b8c2070
address some reviews
kianorr Feb 27, 2024
5d7ca1d
get rid of eq_fixed and field_fixed options
kianorr Feb 28, 2024
1fc8aa5
increase default source resolution
kianorr Feb 28, 2024
f19edf6
delete loop option and update docstring
kianorr Feb 28, 2024
b5962bc
Merge branch 'master' into ko/coils_quad_flux
kianorr Feb 28, 2024
4f5d94d
delete unneeded eq in conftest.py
kianorr Feb 28, 2024
2d61252
add verbose arg for code coverage
kianorr Feb 28, 2024
55b6eca
refactor build()
kianorr Mar 1, 2024
dbd2beb
Merge branch 'master' into ko/coils_quad_flux
dpanici Mar 6, 2024
e1f178a
address some reviews
kianorr Mar 6, 2024
039272c
Merge branch 'ko/coils_quad_flux' of https://github.com/PlasmaControl…
kianorr Mar 6, 2024
e41b22c
swap eq and ext_field
kianorr Mar 6, 2024
d80f411
Merge branch 'master' into ko/coils_quad_flux
kianorr Mar 7, 2024
bd0f84d
Merge branch 'master' into ko/coils_quad_flux
kianorr Mar 13, 2024
7e0622f
Merge branch 'master' into ko/coils_quad_flux
f0uriest Apr 2, 2024
e8ff7a4
Rename args, use helper func to compute Bplasma
f0uriest Apr 2, 2024
265cde5
Use precomputed equilibria for quadratic flux tests
f0uriest Apr 2, 2024
26eed1d
Merge branch 'master' into ko/coils_quad_flux
dpanici Apr 2, 2024
372fbe3
Merge branch 'master' into ko/coils_quad_flux
f0uriest Apr 4, 2024
aa1a5d3
Add QuadraticFlux to nangrad and scalar resolution tests
f0uriest Apr 4, 2024
bd8eb25
Refactor quadratic flux tests
f0uriest Apr 4, 2024
70c4b0c
Merge branch 'master' into ko/coils_quad_flux
f0uriest Apr 4, 2024
c980284
update docstring to match order of args in init:
dpanici Apr 4, 2024
e843d94
add defualt grid info
dpanici Apr 4, 2024
7a1ba53
Move QuadraticFlux to coil objectives
f0uriest Apr 4, 2024
638272f
Minor cleanup
f0uriest Apr 4, 2024
608f2e9
Merge branch 'master' into ko/coils_quad_flux
f0uriest Apr 5, 2024
6318ec9
Merge branch 'master' into ko/coils_quad_flux
f0uriest Apr 5, 2024
0d84c4e
formatting changes to reduce number of lines of code
daniel-dudt Apr 5, 2024
1f27e86
Add check for fixed eq to optimizer
f0uriest Apr 5, 2024
851e658
remove extra space in error message
daniel-dudt Apr 5, 2024
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
2 changes: 1 addition & 1 deletion desc/objectives/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Classes defining objectives for equilibrium and optimization."""

from ._bootstrap import BootstrapRedlConsistency
from ._coils import CoilCurvature, CoilLength, CoilTorsion
from ._coils import CoilCurvature, CoilLength, CoilTorsion, QuadraticFlux
from ._equilibrium import (
CurrentDensity,
Energy,
Expand Down
212 changes: 207 additions & 5 deletions desc/objectives/_coils.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@
tree_structure,
tree_unflatten,
)
from desc.compute import get_transforms
from desc.compute import compute as compute_fun
from desc.compute import get_profiles, get_transforms
from desc.grid import LinearGrid, _Grid
from desc.singularities import compute_B_plasma
from desc.utils import Timer, errorif

from .normalization import compute_scaling_factors
Expand Down Expand Up @@ -46,7 +48,7 @@ class _CoilObjective(_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. Operates over all coils, not each individial coil.
normalization. Operates over all coils, not each individual coil.
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
Expand Down Expand Up @@ -270,7 +272,7 @@ class CoilLength(_CoilObjective):
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. Operates over all coils, not each individial coil.
normalization. Operates over all coils, not each individual coil.
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
Expand Down Expand Up @@ -400,7 +402,7 @@ class CoilCurvature(_CoilObjective):
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. Operates over all coils, not each individial coil.
normalization. Operates over all coils, not each individual coil.
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
Expand Down Expand Up @@ -515,7 +517,7 @@ class CoilTorsion(_CoilObjective):
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. Operates over all coils, not each individial coil.
normalization. Operates over all coils, not each individual coil.
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
Expand Down Expand Up @@ -597,3 +599,203 @@ def compute(self, params, constants=None):
data = tree_flatten(data, is_leaf=lambda x: isinstance(x, dict))[0]
out = jnp.concatenate([dat["torsion"] for dat in data])
return out


class QuadraticFlux(_Objective):
"""Target B*n = 0 on LCFS.

Uses virtual casing to find plasma component of B and penalizes
(B_coil + B_plasma)*n. The equilibrium is kept fixed while the
field is unfixed.

Parameters
----------
eq : Equilibrium
Equilibrium upon whose surface the normal field error will be minimized.
The equilibrium is kept fixed during the optimization with this objective.
field : MagneticField
External field produced by coils or other source, which will be optimized to
minimize the normal field error on the provided equilibrium's surface.
target : float, ndarray, optional
Target value(s) of the objective. Only used if bounds is None.
len(target) must be equal to Objective.dim_f.
Default target is zero.
bounds : tuple, optional
Lower and upper bounds on the objective. Overrides target.
len(bounds[0]) and len(bounds[1]) must be equal to Objective.dim_f
weight : float, ndarray, optional
Weighting to apply to the Objective, relative to other Objectives.
len(weight) must be equal to Objective.dim_f
normalize : bool
Whether to compute the error in physical units or non-dimensionalize.
normalize_target : bool
Whether target and bounds 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.
source_grid : Grid, optional
Collocation grid containing the nodes for plasma source terms.
Default grid is detailed in the docs for ``compute_B_plasma``
eval_grid : Grid, optional
Collocation grid containing the nodes on the plasma surface at which the
magnetic field is being calculated and where to evaluate Bn errors.
Default grid is: LinearGrid(rho=np.array([1.0]), M=eq.M_grid, N=eq.N_grid,
NFP=int(eq.NFP), sym=False)
field_grid : Grid, optional
Grid used to discretize field (e.g. grid for the magnetic field source from
coils). Default grid is determined by the specific MagneticField object, see
the docs of that object's ``compute_magnetic_field`` method for more detail.
vacuum : bool
If true, B_plasma (the contribution to the normal field on the boundary from the
plasma currents) is set to zero.
name : str
Name of the objective function.

"""

_scalar = False
_linear = False
_print_value_fmt = "Boundary normal field error: {:10.3e} "
_units = "(T m^2)"
_coordinates = "rtz"

def __init__(
self,
eq,
field,
target=None,
bounds=None,
weight=1,
normalize=True,
normalize_target=True,
source_grid=None,
eval_grid=None,
field_grid=None,
vacuum=False,
name="Quadratic flux",
):
if target is None and bounds is None:
target = 0
self._source_grid = source_grid
self._eval_grid = eval_grid
self._eq = eq
self._field = field
self._field_grid = field_grid
self._vacuum = vacuum
things = [field]
super().__init__(
things=things,
target=target,
bounds=bounds,
weight=weight,
normalize=normalize,
normalize_target=normalize_target,
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.

"""
eq = self._eq

if self._eval_grid is None:
eval_grid = LinearGrid(
rho=np.array([1.0]),
M=eq.M_grid,
N=eq.N_grid,
NFP=int(eq.NFP),
sym=False,
)
self._eval_grid = eval_grid
else:
eval_grid = self._eval_grid

self._data_keys = ["R", "Z", "n_rho", "phi", "|e_theta x e_zeta|"]

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

self._dim_f = eval_grid.num_nodes

w = eval_grid.weights
w *= jnp.sqrt(eval_grid.num_nodes)

eval_profiles = get_profiles(self._data_keys, obj=eq, grid=eval_grid)
eval_transforms = get_transforms(self._data_keys, obj=eq, grid=eval_grid)
eval_data = compute_fun(
"desc.equilibrium.equilibrium.Equilibrium",
self._data_keys,
params=eq.params_dict,
transforms=eval_transforms,
profiles=eval_profiles,
)

# pre-compute B_plasma because we are assuming eq is fixed
if self._vacuum:
Bplasma = jnp.zeros(eval_grid.num_nodes)

else:
Bplasma = compute_B_plasma(
eq, eval_grid, self._source_grid, normal_only=True
)

self._constants = {
"field": self._field,
"field_grid": self._field_grid,
"quad_weights": w,
"eval_data": eval_data,
"B_plasma": Bplasma,
}

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

if self._normalize:
scales = compute_scaling_factors(eq)
self._normalization = scales["B"] * scales["R0"] * scales["a"]

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

def compute(self, field_params, constants=None):
"""Compute boundary force error.

Parameters
----------
field_params : dict
Dictionary of the external field's degrees of freedom.
constants : dict
Dictionary of constant data, eg transforms, profiles etc. Defaults to
self.constants

Returns
-------
f : ndarray
Bnorm from B_ext and B_plasma

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

# B_plasma from equilibrium precomputed
eval_data = constants["eval_data"]
B_plasma = constants["B_plasma"]

x = jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T

# B_ext is not pre-computed because field is not fixed
B_ext = constants["field"].compute_magnetic_field(
x, source_grid=constants["field_grid"], basis="rpz", params=field_params
)
B_ext = jnp.sum(B_ext * eval_data["n_rho"], axis=-1)
f = (B_ext + B_plasma) * eval_data["|e_theta x e_zeta|"]
return f
20 changes: 9 additions & 11 deletions desc/objectives/getters.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

import numpy as np

from desc.utils import is_any_instance

from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance
from .linear_objectives import (
AxisRSelfConsistency,
Expand Down Expand Up @@ -220,10 +222,6 @@

def maybe_add_self_consistency(thing, constraints):
"""Add self consistency constraints if needed."""

def _is_any_instance(things, cls):
return any([isinstance(t, cls) for t in things])

# Equilibrium
if (
hasattr(thing, "Ra_n")
Expand All @@ -232,22 +230,22 @@
and hasattr(thing, "Zb_lmn")
and hasattr(thing, "L_lmn")
):
if not _is_any_instance(constraints, BoundaryRSelfConsistency):
if not is_any_instance(constraints, BoundaryRSelfConsistency):
constraints += (BoundaryRSelfConsistency(eq=thing),)
if not _is_any_instance(constraints, BoundaryZSelfConsistency):
if not is_any_instance(constraints, BoundaryZSelfConsistency):
constraints += (BoundaryZSelfConsistency(eq=thing),)
if not _is_any_instance(constraints, FixLambdaGauge):
if not is_any_instance(constraints, FixLambdaGauge):
constraints += (FixLambdaGauge(eq=thing),)
if not _is_any_instance(constraints, AxisRSelfConsistency):
if not is_any_instance(constraints, AxisRSelfConsistency):
constraints += (AxisRSelfConsistency(eq=thing),)
if not _is_any_instance(constraints, AxisZSelfConsistency):
if not is_any_instance(constraints, AxisZSelfConsistency):
constraints += (AxisZSelfConsistency(eq=thing),)

# Curve
elif hasattr(thing, "shift") and hasattr(thing, "rotmat"):
if not _is_any_instance(constraints, FixCurveShift):
if not is_any_instance(constraints, FixCurveShift):

Check warning on line 246 in desc/objectives/getters.py

View check run for this annotation

Codecov / codecov/patch

desc/objectives/getters.py#L246

Added line #L246 was not covered by tests
constraints += (FixCurveShift(curve=thing),)
if not _is_any_instance(constraints, FixCurveRotation):
if not is_any_instance(constraints, FixCurveRotation):

Check warning on line 248 in desc/objectives/getters.py

View check run for this annotation

Codecov / codecov/patch

desc/objectives/getters.py#L248

Added line #L248 was not covered by tests
constraints += (FixCurveRotation(curve=thing),)

return constraints
24 changes: 23 additions & 1 deletion desc/optimize/optimizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,15 @@
maybe_add_self_consistency,
)
from desc.objectives.utils import combine_args
from desc.utils import Timer, flatten_list, get_instance, unique_list, warnif
from desc.utils import (
Timer,
errorif,
flatten_list,
get_instance,
is_any_instance,
unique_list,
warnif,
)

from ._constraint_wrappers import LinearConstraintProjection, ProximalProjection

Expand Down Expand Up @@ -147,6 +155,11 @@ def optimize( # noqa: C901 - FIXME: simplify this
"""
if not isinstance(constraints, (tuple, list)):
constraints = (constraints,)
errorif(
not isinstance(objective, ObjectiveFunction),
TypeError,
"objective should be of type ObjectiveFunction.",
)

# get unique things
things, indices = unique_list(flatten_list(things, flatten_tuple=True))
Expand All @@ -161,10 +174,19 @@ def optimize( # noqa: C901 - FIXME: simplify this

# need local import to avoid circular dependencies
from desc.equilibrium import Equilibrium
from desc.objectives import QuadraticFlux

# eq may be None
eq = get_instance(things, Equilibrium)
if eq is not None:
# check if stage 2 objectives are here:
errorif(
is_any_instance(constraints, QuadraticFlux)
or is_any_instance(objective.objectives, QuadraticFlux),
ValueError,
"QuadraticFlux objective assumes Equilibrium is fixed but Equilibrium "
+ "is in things to optimize.",
)
# save these for later
eq_params_init = eq.params_dict.copy()

Expand Down
19 changes: 3 additions & 16 deletions desc/singularities.py
Original file line number Diff line number Diff line change
Expand Up @@ -834,16 +834,7 @@ def compute_B_plasma(eq, eval_grid, source_grid=None, normal_only=False):
sym=False,
)

data_keys = [
"K_vc",
"B",
"R",
"phi",
"Z",
"e^rho",
"n_rho",
"|e_theta x e_zeta|",
]
data_keys = ["K_vc", "B", "R", "phi", "Z", "e^rho", "n_rho", "|e_theta x e_zeta|"]
eval_data = eq.compute(data_keys, grid=eval_grid)
source_data = eq.compute(data_keys, grid=source_grid)

Expand All @@ -860,13 +851,9 @@ def compute_B_plasma(eq, eval_grid, source_grid=None, normal_only=False):
interpolator = DFTInterpolator(eval_grid, source_grid, s, q)
if hasattr(eq.surface, "Phi_mn"):
source_data["K_vc"] += eq.surface.compute("K", grid=source_grid)["K"]
Bplasma = virtual_casing_biot_savart(
eval_data,
source_data,
interpolator,
)
Bplasma = virtual_casing_biot_savart(eval_data, source_data, interpolator)
# need extra factor of B/2 bc we're evaluating on plasma surface
Bplasma = Bplasma + eval_data["B"] / 2
if normal_only:
Bplasma = jnp.sum(Bplasma * source_data["n_rho"], axis=1)
Bplasma = jnp.sum(Bplasma * eval_data["n_rho"], axis=1)
return Bplasma
Loading
Loading