diff --git a/CHANGELOG.md b/CHANGELOG.md index d36274333..a45a6d410 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,13 @@ Changelog ========= +New Feature + +- Adds a new profile class ``PowerProfile`` for raising profiles to a power. + +v0.13.0 +------- + 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 @@ -14,6 +21,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. - Adds ``SurfaceQuadraticFlux`` objective which minimizes the quadratic magnetic flux through a ``FourierRZToroidalSurface`` object, allowing for optimizing for Quadratic flux minimizing (QFM) surfaces. - Allows ``ToroidalFlux`` objective to accept ``FourierRZToroidalSurface`` so it can be used to specify the toroidal flux through a QFM surface. - Adds ``eq_fixed`` flag to ``ToroidalFlux`` to allow for the equilibrium/QFM surface to vary during optimization, useful for single-stage optimizations. @@ -21,13 +29,12 @@ New Features - Adds ``rotate_zeta`` function to ``desc.compat`` to rotate an ``Equilibrium`` around Z axis. - Adds an option ``scaled_termination`` (defaults to True) to all of the desc optimizers to measure the norms for ``xtol`` and ``gtol`` in the scaled norm provided by ``x_scale`` (which defaults to using an adaptive scaling based on the Jacobian or Hessian). This should make things more robust when optimizing parameters with widely different magnitudes. The old behavior can be recovered by passing ``options={"scaled_termination": False}``. - 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 ------- diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 29299d7a1..b8af799d8 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 = getattr(field, "_NFP", 1) try: AR, AP, AZ = field.compute_magnetic_vector_potential( coords, params, basis="rpz" @@ -1948,12 +1950,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 +2054,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 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/desc/profiles.py b/desc/profiles.py index 064871e58..ae1684abb 100644 --- a/desc/profiles.py +++ b/desc/profiles.py @@ -244,6 +244,17 @@ def __sub__(self, x): """Subtract another profile from this one.""" return self.__add__(-x) + def __pow__(self, x): + """Raise this profile to a power.""" + if np.isscalar(x): + return PowerProfile(x, self) + else: + raise NotImplementedError() + + def __rpow__(self, x): + """Raise this profile to a power.""" + return self.__pow__(x) + class ScaledProfile(_Profile): """Profile times a constant value. @@ -252,10 +263,10 @@ class ScaledProfile(_Profile): Parameters ---------- - profile : Profile - Base profile to scale. scale : float Scale factor. + profile : Profile + Base profile to scale. """ @@ -335,6 +346,128 @@ def __repr__(self): return s +class PowerProfile(_Profile): + """Profile raised to a power. + + f_1(x) = f(x)**a + + Parameters + ---------- + power : float + Exponent of the new profile. + profile : Profile + Base profile to raise to a power. + + """ + + _io_attrs_ = _Profile._io_attrs_ + ["_profile", "_power"] + + def __init__(self, power, profile, **kwargs): + assert isinstance( + profile, _Profile + ), "profile in a PowerProfile must be a Profile or subclass, got {}.".format( + str(profile) + ) + assert np.isscalar(power), "power must be a scalar." + + self._profile = profile.copy() + self._power = power + + self._check_params() + + kwargs.setdefault("name", profile.name) + super().__init__(**kwargs) + + def _check_params(self, params=None): + """Check params and throw warnings or errors if necessary.""" + params = self.params if params is None else params + power, params = self._parse_params(params) + warnif( + power < 0, + UserWarning, + "This profile may be undefined at some points because power < 0.", + ) + + @property + def params(self): + """ndarray: Parameters for computation [power, profile.params].""" + return jnp.concatenate([jnp.atleast_1d(self._power), self._profile.params]) + + @params.setter + def params(self, x): + self._check_params(x) + self._power, self._profile.params = self._parse_params(x) + + def _parse_params(self, x): + if x is None: + power = self._power + params = self._profile.params + elif isinstance(x, (tuple, list)) and len(x) == 2: + params = x[1] + power = x[0] + elif np.isscalar(x): + power = x + params = self._profile.params + elif len(x) == len(self._profile.params): + power = self._power + params = x + elif len(x) == len(self.params): + power = x[0] + params = x[1:] + else: + raise ValueError("Got wrong number of parameters for PowerProfile") + return power, params + + def compute(self, grid, params=None, dr=0, dt=0, dz=0): + """Compute values of profile at specified nodes. + + Parameters + ---------- + grid : Grid + locations to compute values at. + params : array-like + Parameters to use. If not given, uses the + values given by the self.params attribute. + dr, dt, dz : int + derivative order in rho, theta, zeta. + + Returns + ------- + values : ndarray + values of the profile or its derivative at the points specified. + + """ + if dt > 0 or dz > 0: + raise NotImplementedError( + "Poloidal and toroidal derivatives of PowerProfile have not been " + + "implemented yet." + ) + power, params = self._parse_params(params) + f0 = self._profile.compute(grid, params, 0, dt, dz) + if dr >= 1: + df1 = self._profile.compute(grid, params, 1, dt, dz) # df/dr + fn1 = self.compute(grid, (power - 1, params), 0, dt, dz) # f^(n-1) + if dr >= 2: + df2 = self._profile.compute(grid, params, 2, dt, dz) # d^2f/dr^2 + fn2 = self.compute(grid, (power - 2, params), 0, dt, dz) # f^(n-2) + if dr == 0: + f = f0**power + elif dr == 1: + f = power * fn1 * df1 + elif dr == 2: + f = power * ((power - 1) * fn2 * df1**2 + fn1 * df2) + else: + raise NotImplementedError("dr > 2 not implemented for PowerProfile!") + return f + + def __repr__(self): + """Get the string form of the object.""" + s = super().__repr__() + s = s[:-1] + s += ", power={})".format(self._power) + return s + + class SumProfile(_Profile): """Sum of two or more Profiles. @@ -724,17 +857,24 @@ def __init__(self, params=None, name=""): params = [0, 1, 1] self._params = np.atleast_1d(params) + self._check_params() + + def _check_params(self, params=None): + """Check params and throw warnings or errors if necessary.""" + params = self.params if params is None else params errorif( - self._params.size != 3, ValueError, "params must be an array of size 3." + params.size != 3, + ValueError, + f"params must be an array of size 3, got {len(params)}.", ) warnif( - self._params[1] < 1, + params[1] < 1, UserWarning, "Derivatives of this profile will be infinite at rho=0 " + "because params[1] < 1.", ) warnif( - self._params[2] < 1, + params[2] < 1, UserWarning, "Derivatives of this profile will be infinite at rho=1 " + "because params[2] < 1.", @@ -748,10 +888,8 @@ def params(self): @params.setter def params(self, new): new = jnp.atleast_1d(jnp.asarray(new)) - if new.size == 3: - self._params = jnp.asarray(new) - else: - raise ValueError(f"params should be an array of size 3, got {len(new)}.") + self._check_params(new) + self._params = new def compute(self, grid, params=None, dr=0, dt=0, dz=0): """Compute values of profile at specified nodes. diff --git a/docs/api.rst b/docs/api.rst index 1755b7a64..fa32edb86 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -292,14 +292,16 @@ Profiles :recursive: :template: class.rst - desc.profiles.PowerSeriesProfile - desc.profiles.TwoPowerProfile - desc.profiles.SplineProfile + desc.profiles.FourierZernikeProfile desc.profiles.HermiteSplineProfile desc.profiles.MTanhProfile + desc.profiles.PowerProfile + desc.profiles.PowerSeriesProfile + desc.profiles.ProductProfile desc.profiles.ScaledProfile + desc.profiles.SplineProfile desc.profiles.SumProfile - desc.profiles.ProductProfile + desc.profiles.TwoPowerProfile Transform ********* diff --git a/docs/api_equilibrium.rst b/docs/api_equilibrium.rst index 5600382c6..30a04d730 100644 --- a/docs/api_equilibrium.rst +++ b/docs/api_equilibrium.rst @@ -56,9 +56,12 @@ profiles together by addition, multiplication, or scaling. desc.profiles.PowerSeriesProfile desc.profiles.SplineProfile desc.profiles.MTanhProfile + desc.profiles.TwoPowerProfile + desc.profiles.HermiteSplineProfile desc.profiles.ScaledProfile - desc.profiles.SumProfile desc.profiles.ProductProfile + desc.profiles.SumProfile + desc.profiles.PowerProfile Utilities diff --git a/docs/api_plotting.rst b/docs/api_plotting.rst index f104735ad..4c15c760d 100644 --- a/docs/api_plotting.rst +++ b/docs/api_plotting.rst @@ -27,6 +27,7 @@ Plotting Flux Surfaces desc.plotting.plot_comparison desc.plotting.plot_boundary desc.plotting.plot_boundaries + desc.plotting.poincare_plot Contour Plots of 2-D data diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 6484615bc..134408bd4 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 @@ -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) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index edf3bf1a4..77544020e 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 diff --git a/tests/test_profiles.py b/tests/test_profiles.py index bf75489a6..ef14c166f 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -189,6 +189,7 @@ def test_repr(self): assert "SumProfile" in str(pp + zp) assert "ProductProfile" in str(pp * zp) assert "ScaledProfile" in str(2 * zp) + assert "PowerProfile" in str(zp**2) @pytest.mark.unit def test_get_set(self): @@ -347,6 +348,47 @@ def test_scaled_profiles(self): np.testing.assert_allclose(pp.params, [1, -2, 1]) np.testing.assert_allclose(f(x), 8 * (pp(x)), atol=1e-3) + @pytest.mark.unit + def test_powered_profiles(self): + """Test raising profiles to a power.""" + pp = PowerSeriesProfile( + modes=np.array([0, 1, 2, 4]), params=np.array([1, 0, -2, 1]), sym="auto" + ) + + f = pp**3 + x = np.linspace(0, 1, 50) + np.testing.assert_allclose(f(x), (pp(x)) ** 3, atol=1e-3) + + params = f.params + assert params[0] == 3 + assert all(params[1:] == pp.params) + + f.params = 2 + np.testing.assert_allclose(f(x), (pp(x)) ** 2, atol=1e-3) + + f.params = 0.5 + np.testing.assert_allclose(f(x), np.sqrt(pp(x)), atol=1e-3) + + @pytest.mark.unit + def test_powered_profiles_derivative(self): + """Test that powered profiles computes the derivative correctly.""" + x = np.linspace(0, 1, 50) + p1 = PowerSeriesProfile( + modes=np.array([0, 1, 2, 4]), params=np.array([1, 3, -2, 4]), sym="auto" + ) + p2 = p1 * p1 + p3 = p1 * p2 + + f3 = p1**3 + np.testing.assert_allclose(f3(x, dr=0), p3(x, dr=0)) + np.testing.assert_allclose(f3(x, dr=1), p3(x, dr=1)) + np.testing.assert_allclose(f3(x, dr=2), p3(x, dr=2)) + + f2 = f3 ** (2 / 3) + np.testing.assert_allclose(f2(x, dr=0), p2(x, dr=0)) + np.testing.assert_allclose(f2(x, dr=1), p2(x, dr=1)) + np.testing.assert_allclose(f2(x, dr=2), p2(x, dr=2)) + @pytest.mark.unit def test_profile_errors(self): """Test error checking when creating and working with profiles.""" @@ -383,6 +425,10 @@ def test_profile_errors(self): tp.compute(grid, dr=3) with pytest.raises(NotImplementedError): mp.compute(grid, dr=3) + with pytest.raises(UserWarning): + tp.params = [1, 0.3, 0.7] + with pytest.raises(UserWarning): + a = sp**-1 @pytest.mark.unit def test_default_profiles(self):