From 9b58f01ba57d118fee9a7389ff6a2b7d51cffe23 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 2 Dec 2024 11:15:53 -0500 Subject: [PATCH 1/7] add NFP attribute to a few fields --- desc/magnetic_fields/_core.py | 24 ++++++++++++++++++++++-- desc/magnetic_fields/_dommaschk.py | 12 +++++++++++- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 29299d7a1..6312acb71 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1948,12 +1948,22 @@ class ScalarPotentialField(_MagneticField): R,phi,Z are arrays of cylindrical coordinates. params : dict, optional default parameters to pass to potential function + NFP : int, optional + Whether the field has a discrete periodicity. This is only used when making + a ``SplineMagneticField`` from this field using its ``from_field`` method, + or when saving this field as an mgrid file using the ``save_mgrid`` method. """ - def __init__(self, potential, params=None): + def __init__(self, potential, params=None, NFP=1): self._potential = potential self._params = params + self._NFP = NFP + + @property + def NFP(self): + """int: Number of (toroidal) field periods.""" + return self._NFP def compute_magnetic_field( self, coords, params=None, basis="rpz", source_grid=None, transforms=None @@ -2042,12 +2052,22 @@ class VectorPotentialField(_MagneticField): R,phi,Z are arrays of cylindrical coordinates. params : dict, optional default parameters to pass to potential function + NFP : int, optional + Whether the field has a discrete periodicity. This is only used when making + a ``SplineMagneticField`` from this field using its ``from_field`` method, + or when saving this field as an mgrid file using the ``save_mgrid`` method. """ - def __init__(self, potential, params=None): + def __init__(self, potential, params=None, NFP=1): self._potential = potential self._params = params + self._NFP = NFP + + @property + def NFP(self): + """int: Number of (toroidal) field periods.""" + return self._NFP def _compute_A_or_B( self, diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py index 009c7fb72..a26dbff23 100644 --- a/desc/magnetic_fields/_dommaschk.py +++ b/desc/magnetic_fields/_dommaschk.py @@ -38,6 +38,10 @@ class DommaschkPotentialField(ScalarPotentialField): d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms B0: float scale strength of the magnetic field's 1/R portion + NFP : int, optional + Whether the field has a discrete periodicity. This is only used when making + a ``SplineMagneticField`` from this field using its ``from_field`` method, + or when saving this field as an mgrid file using the ``save_mgrid`` method. """ @@ -50,6 +54,7 @@ def __init__( c_arr=jnp.array([0.0]), d_arr=jnp.array([0.0]), B0=1.0, + NFP=1, ): ms = jnp.atleast_1d(jnp.asarray(ms)) ls = jnp.atleast_1d(jnp.asarray(ls)) @@ -68,6 +73,11 @@ def __init__( jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1 ), "B0 should be a scalar value!" + ms_over_NFP = ms / NFP + assert jnp.allclose( + ms_over_NFP, ms_over_NFP.astype(int) + ), "To enforce desired NFP, `ms` should be all integer multiples of NFP" + params = {} params["ms"] = ms params["ls"] = ls @@ -77,7 +87,7 @@ def __init__( params["d_arr"] = d_arr params["B0"] = B0 - super().__init__(dommaschk_potential, params) + super().__init__(dommaschk_potential, params, NFP) @classmethod def fit_magnetic_field( # noqa: C901 From 65d562f3af3ea6e051a304282e829476b8360276 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 2 Dec 2024 11:36:37 -0500 Subject: [PATCH 2/7] add NFP check to from_field SplineMagneticField method --- desc/magnetic_fields/_core.py | 6 ++++-- tests/test_magnetic_fields.py | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 6312acb71..59d40a9d4 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1885,7 +1885,7 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False): @classmethod def from_field( - cls, field, R, phi, Z, params=None, method="cubic", extrap=False, NFP=1 + cls, field, R, phi, Z, params=None, method="cubic", extrap=False, NFP=None ): """Create a splined magnetic field from another field for faster evaluation. @@ -1903,7 +1903,8 @@ def from_field( extrap : bool whether to extrapolate splines beyond specified R,phi,Z NFP : int, optional - Number of toroidal field periods. + Number of toroidal field periods. If not provided, will default to 1 or + the provided field's NFP, if it has that attribute. """ R, phi, Z = map(np.asarray, (R, phi, Z)) @@ -1911,6 +1912,7 @@ def from_field( shp = rr.shape coords = np.array([rr.flatten(), pp.flatten(), zz.flatten()]).T BR, BP, BZ = field.compute_magnetic_field(coords, params, basis="rpz").T + NFP = setdefault(field.NFP, 1, hasattr(field, "_NFP")) try: AR, AP, AZ = field.compute_magnetic_vector_potential( coords, params, basis="rpz" diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 6484615bc..c6500f90c 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -1045,11 +1045,11 @@ def test_fourier_current_potential_field_coil_cut_warnings(self): @pytest.mark.unit def test_spline_field(self, tmpdir_factory): """Test accuracy of spline magnetic field.""" - field1 = ScalarPotentialField(phi_lm, args) + field1 = ScalarPotentialField(phi_lm, args, NFP=5) R = np.linspace(0.5, 1.5, 20) Z = np.linspace(-1.5, 1.5, 20) p = np.linspace(0, 2 * np.pi / 5, 40) - field2 = SplineMagneticField.from_field(field1, R, p, Z, NFP=5) + field2 = SplineMagneticField.from_field(field1, R, p, Z) np.testing.assert_allclose( field1([1.0, 1.0, 1.0]), field2([1.0, 1.0, 1.0]), rtol=1e-2, atol=1e-2 From 612c74a170afad1fdc7ef457caff7e676992273d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 2 Dec 2024 16:07:14 -0500 Subject: [PATCH 3/7] fix attribute check --- desc/magnetic_fields/_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 59d40a9d4..b8af799d8 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1912,7 +1912,7 @@ def from_field( shp = rr.shape coords = np.array([rr.flatten(), pp.flatten(), zz.flatten()]).T BR, BP, BZ = field.compute_magnetic_field(coords, params, basis="rpz").T - NFP = setdefault(field.NFP, 1, hasattr(field, "_NFP")) + NFP = getattr(field, "_NFP", 1) try: AR, AP, AZ = field.compute_magnetic_vector_potential( coords, params, basis="rpz" From ed1a9c1276b2ea2b44a00094eabff4e21e5ec34c Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 2 Dec 2024 16:11:30 -0500 Subject: [PATCH 4/7] update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 58671a6bb..3545e2175 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ New Features * use of both this and the ``QuadraticFlux`` objective allows for REGCOIL solutions to be obtained through the optimization framework, and combined with other objectives as well. - Changes local area weighting of Bn in QuadraticFlux objective to be the square root of the local area element (Note that any existing optimizations using this objective may need different weights to achieve the same result now.) - Adds a new tutorial showing how to use``REGCOIL`` features. +- Adds an ``NFP`` attribute to ``ScalarPotentialField``, ``VectorPotentialField`` and ``DommaschkPotentialField``, to allow ``SplineMagneticField.from_field`` and ``MagneticField.save_mgrid`` to efficiently take advantage of the discrete toroidal symmetry of these fields, if present. Bug Fixes From 6b378bbb24cc47817e778154e2254b422f8d6ab0 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 2 Dec 2024 16:14:27 -0500 Subject: [PATCH 5/7] add test --- tests/test_magnetic_fields.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index c6500f90c..134408bd4 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -1450,11 +1450,13 @@ def test_dommaschk_field_errors(): b_arr = [1] c_arr = [1] d_arr = [1, 1] # length is not equal to the rest - with pytest.raises(AssertionError): + with pytest.raises(AssertionError, match="size"): DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) - d_arr = [1] + d_arr = [1] # test with incorrect NFP + with pytest.raises(AssertionError, match="desired NFP"): + DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr, NFP=2) ms = [-1] # negative mode number - with pytest.raises(AssertionError): + with pytest.raises(AssertionError, match=">= 0"): DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) From 0d4c33143d67da5841ea2e15c8217e6cb83e6802 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 2 Dec 2024 17:13:24 -0500 Subject: [PATCH 6/7] bug fix for proximal scalar optimization --- desc/optimize/_constraint_wrappers.py | 19 ++++++++++++++ tests/test_optimizer.py | 38 +++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py index dc2e3033c..7695a671b 100644 --- a/desc/optimize/_constraint_wrappers.py +++ b/desc/optimize/_constraint_wrappers.py @@ -836,6 +836,25 @@ def compute_scaled_error(self, x, constants=None): xopt, _ = self._update_equilibrium(x, store=False) return self._objective.compute_scaled_error(xopt, constants[0]) + def compute_scalar(self, x, constants=None): + """Compute the sum of squares error. + + Parameters + ---------- + x : ndarray + State vector. + constants : list + Constant parameters passed to sub-objectives. + + Returns + ------- + f : float + Objective function scalar value. + + """ + f = jnp.sum(self.compute_scaled_error(x, constants=constants) ** 2) / 2 + return f + def compute_unscaled(self, x, constants=None): """Compute the raw value of the objective function. diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 4f0cce0e0..9a4aac89b 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -324,6 +324,44 @@ def test_no_iterations(): np.testing.assert_allclose(x0, out2["x"]) +@pytest.mark.regression +@pytest.mark.optimize +def test_proximal_scalar(): + """Test that proximal scalar optimization works.""" + # test fix for GH issue #1403 + + # optimize to reduce DSHAPE volume from 100 m^3 to 90 m^3 + eq = desc.examples.get("DSHAPE") + optimizer = Optimizer("proximal-fmintr") # proximal scalar optimizer + R_modes = np.vstack( + ( + [0, 0, 0], + eq.surface.R_basis.modes[ + np.max(np.abs(eq.surface.R_basis.modes), 1) > 1, : + ], + ) + ) + Z_modes = eq.surface.Z_basis.modes[ + np.max(np.abs(eq.surface.Z_basis.modes), 1) > 1, : + ] + objective = ObjectiveFunction(Volume(eq=eq, target=90)) # scalar objective function + constraints = ( + FixBoundaryR(eq=eq, modes=R_modes), + FixBoundaryZ(eq=eq, modes=Z_modes), + FixIota(eq=eq), + FixPressure(eq=eq), + FixPsi(eq=eq), + ForceBalance(eq=eq), # force balance constraint for proximal projection + ) + [eq], _ = optimizer.optimize( + things=eq, + objective=objective, + constraints=constraints, + verbose=3, + ) + np.testing.assert_allclose(eq.compute("V")["V"], 90) + + @pytest.mark.regression @pytest.mark.slow @pytest.mark.optimize From f106527526982058214781ef3ba4af5d235b0198 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 3 Dec 2024 10:37:30 -0500 Subject: [PATCH 7/7] update Changelog --- CHANGELOG.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 58671a6bb..1b784b5fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ Changelog ========= New Features + - Adds ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file - Adds function ``solve_regularized_surface_current`` to ``desc.magnetic_fields`` module that implements the REGCOIL algorithm (Landreman, (2017)) for surface current normal field optimization * Can specify the tuple ``current_helicity=(M_coil, N_coil)`` to determine if resulting contours correspond to helical topology (both ``(M_coil, N_coil)`` not equal to 0), modular (``N_coil`` equal to 0 and ``M_coil`` nonzero) or windowpane/saddle (``M_coil`` and ``N_coil`` both zero) @@ -17,10 +18,10 @@ New Features Bug Fixes -- Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version +- Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version. - Changes ``FixLambdaGauge`` constraint to now enforce zero flux surface average for lambda, instead of enforcing lambda(rho,0,0)=0 as it was incorrectly doing before. - Fixes bug in ``softmin/softmax`` implementation. - +- Fixes bug that occured when using ``ProximalProjection`` with a scalar optimization algorithm. v0.12.3 -------