Skip to content

Commit

Permalink
remove unneeded changes to build methods
Browse files Browse the repository at this point in the history
  • Loading branch information
@dpanici
dpanici committed Jul 2, 2024
1 parent 3e352c3 commit 8edeecf
Show file tree
Hide file tree
Showing 3 changed files with 0 additions and 99 deletions.
37 changes: 0 additions & 37 deletions desc/objectives/_coils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1225,31 +1225,13 @@ def build(self, use_jit=True, verbose=1):
Bplasma = compute_B_plasma(
eq, eval_grid, self._source_grid, normal_only=True
)
field = self._field
# FIXME: need to add case for CurrentPotentialField as well
# FIXME: This will fail if we have a SumMagneticField where
# one of the constituents is a (Fourier)CurrentPotentialField
if hasattr(field, "Phi_mn"):
# make the transform for the CurrentPotentialField
if self._field_grid is None:
self._field_grid = LinearGrid(
M=30 + 2 * max(field.M, field.M_Phi),
N=30 + 2 * max(field.N, field.N_Phi),
NFP=field.NFP,
)
field_transforms = get_transforms(
["K", "x"], obj=field, grid=self._field_grid
)
else:
field_transforms = None

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

timer.stop("Precomputing transforms")
Expand Down Expand Up @@ -1294,7 +1276,6 @@ def compute(self, field_params, constants=None):
source_grid=constants["field_grid"],
basis="rpz",
params=field_params,
transforms=constants["field_transforms"],
)
B_ext = jnp.sum(B_ext * eval_data["n_rho"], axis=-1)
f = (B_ext + B_plasma) * eval_data["|e_theta x e_zeta|"]
Expand Down Expand Up @@ -1450,30 +1431,13 @@ def build(self, use_jit=True, verbose=1):

plasma_coords = jnp.array([data["R"], data["phi"], data["Z"]]).T

field = self._field
# FIXME: need to add case for CurrentPotentialField as well
if hasattr(field, "Phi_mn"):
# make the transform for the CurrentPotentialField
if self._field_grid is None:
self._field_grid = LinearGrid(
M=30 + 2 * max(field.M, field.M_Phi),
N=30 + 2 * max(field.N, field.N_Phi),
NFP=field.NFP,
)
field_transforms = get_transforms(
["K", "x"], obj=field, grid=self._field_grid
)
else:
field_transforms = None

self._constants = {
"plasma_coords": plasma_coords,
"equil_data": data,
"quad_weights": 1.0,
"field": self._field,
"field_grid": self._field_grid,
"eval_grid": eval_grid,
"field_transforms": field_transforms,
}

timer.stop("Precomputing transforms")
Expand Down Expand Up @@ -1511,7 +1475,6 @@ def compute(self, field_params=None, constants=None):
basis="rpz",
source_grid=constants["field_grid"],
params=field_params,
transforms=constants["field_transforms"],
)
grid = constants["eval_grid"]

Expand Down
6 changes: 0 additions & 6 deletions desc/objectives/_free_boundary.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -240,9 +240,6 @@ def compute(self, eq_params, field_params=None, constants=None):
x = jnp.array([data["R"], data["phi"], data["Z"]]).T
# can always pass in field params. If they're None, it just uses the
# defaults for the given field.
# FIXME: need to add case for CurrentPotentialField as well,
# if external field is a CurrentPotentialField then transforms
# must be built in build then passed in here
Bext = constants["field"].compute_magnetic_field(
x, source_grid=self._field_grid, basis="rpz", params=field_params
)
Expand Down Expand Up @@ -689,9 +686,6 @@ def compute(self, eq_params, field_params=None, constants=None):
x = jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T
# can always pass in field params. If they're None, it just uses the
# defaults for the given field.
# FIXME: need to add case for CurrentPotentialField as well,
# if external field is a CurrentPotentialField then transforms
# must be built in build then passed in here
Bext = constants["field"].compute_magnetic_field(
x, source_grid=self._field_grid, basis="rpz", params=field_params
)
Expand Down
56 changes: 0 additions & 56 deletions tests/test_optimizer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,6 @@
from desc.equilibrium import Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.grid import LinearGrid
from desc.io import load
from desc.magnetic_fields import FourierCurrentPotentialField
from desc.objectives import (
AspectRatio,
Energy,
Expand All @@ -38,9 +36,7 @@
MeanCurvature,
ObjectiveFunction,
PlasmaVesselDistance,
QuadraticFlux,
QuasisymmetryTripleProduct,
ToroidalFlux,
Volume,
get_fixed_boundary_constraints,
)
Expand Down Expand Up @@ -1322,55 +1318,3 @@ def test_LinearConstraint_jacobian():
np.testing.assert_allclose(vjp_unscaled, vjp1, rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(vjp_unscaled, vjp2, rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(vjp_unscaled, vjp3, rtol=1e-12, atol=1e-12)


@pytest.mark.unit
def test_quad_flux_with_surface_current_field():
"""Test that QuadraticFlux does not throw an error when field has transforms."""
# this happens because in QuadraticFlux.compute, field.compute_magnetic_field
# is called. If the field needs transforms to evaluate, then these transforms
# will be created on the fly if they are not provided, resulting in an error
# This tests the fix where the transforms are precomputed and passed in
# for the FourierCurrentPotentialField class specifically.
eq = load("./tests/inputs/vacuum_circular_tokamak.h5")
field = FourierCurrentPotentialField.from_surface(
eq.surface, Phi_mn=[1, 0], modes_Phi=[[0, 0], [1, 1]], M_Phi=1, N_Phi=1
)
obj = ObjectiveFunction(
QuadraticFlux(
eq=eq,
field=field,
vacuum=True,
eval_grid=LinearGrid(M=2, N=2, sym=True),
field_grid=LinearGrid(M=2, N=2),
),
)
constraints = FixParameters(field, {"I": True, "G": True})
opt = Optimizer("lsq-exact")
# this should run without an error
(field_modular_opt,), result = opt.optimize(
field, objective=obj, constraints=constraints, maxiter=1, copy=True
)


@pytest.mark.unit
def test_tor_flux_with_surface_current_field():
"""Test that ToroidalFlux does not throw an error when field has transforms."""
eq = load("./tests/inputs/vacuum_circular_tokamak.h5")
field = FourierCurrentPotentialField.from_surface(
eq.surface, Phi_mn=[1, 0], modes_Phi=[[0, 0], [1, 1]], M_Phi=1, N_Phi=1
)
obj = ObjectiveFunction(
ToroidalFlux(
eq=eq,
field=field,
eval_grid=LinearGrid(L=2, M=2, sym=True),
field_grid=LinearGrid(M=2, N=2),
),
)
constraints = FixParameters(field, {"I": True, "G": True})
opt = Optimizer("fmintr")
# this should run without an error
(field_modular_opt,), result = opt.optimize(
field, objective=obj, constraints=constraints, maxiter=1, copy=True
)

0 comments on commit 8edeecf

Please sign in to comment.