Merge branch 'master' into ku/biot

CHANGELOG.md

@@ -18,6 +18,7 @@ New Features
 - Changes hessian computation to use chunked ``jacfwd`` and ``jacrev``, allowing ``jac_chunk_size`` to now reduce hessian memory usage as well.
 - Adds an option to ``VMECIO.save`` to specify the grid resolution in real space.
 - Adds a new objective ``desc.objectives.CoilIntegratedCurvature`` for targeting convex coils.
+- Adds the example "reactor_QA", which is similar to "precise_QA" but with self-consistent bootstrap current at finite beta.
 - Allows non-proximal optimizers to  handle optimizing more than one ``Equilibrium`` object simultaneously.
 - Adds batching feature to singular integrals.
 - ``desc.objectives.CoilSetMinDistance`` and ``desc.objectives.PlasmaCoilSetMinDistance`` now include the option to use a softmin which can give smoother gradients. They also both now have a ``dist_chunk_size`` option to break up the distance calculation into smaller pieces to save memory
@@ -31,11 +32,12 @@ for compatibility with other codes which expect such files from the Booz_Xform c
 Bug Fixes
 
 - Small bug fix to use the correct normalization length ``a`` in the BallooningStability objective.
-- Fixed I/O bug when saving/loading ``_Profile`` classes that do not have a ``_params`` attribute.
+- Fixes I/O bug when saving/loading ``_Profile`` classes that do not have a ``_params`` attribute.
 - Minor bugs described in [#1323](https://github.com/PlasmaControl/DESC/pull/1323).
 - Corrects basis vectors computations made on surface objects [#1175](https://github.com/PlasmaControl/DESC/pull/1175).
+- Allows keyword arguments to be passed to ``GenericObjective`` and ``ObjectiveFromUser``.
 - Fixes bug where ``save_in_makegrid_format`` function did not correctly account for ``CoilSet`` objects with NFP>1 or sym=True attributes, and so would not save all the coils.
-- Fix issue with interpolator for singular integrals [#1522](https://github.com/PlasmaControl/DESC/issues/1522) and additional checks [1519](https://github.com/PlasmaControl/DESC/issues/1519).
+- Fixes issue with interpolator for singular integrals [#1522](https://github.com/PlasmaControl/DESC/issues/1522) and additional checks [1519](https://github.com/PlasmaControl/DESC/issues/1519).
 - Fixes the coil currents in ``desc.coils.initialize_modular_coils`` to now give the correct expected linking current.
 - ``desc.objectives.PlasmaVesselDistance`` now correctly accounts for multiple field periods on both the equilibrium and the vessel surface. Previously it only considered distances within a single field period.
 - Sets ``os.environ["JAX_PLATFORMS"] = "cpu"`` instead of ``os.environ["JAX_PLATFORM_NAME"] = "cpu"`` when doing ``set_device("cpu")``.

desc/examples/__init__.py

@@ -4,6 +4,7 @@
 
 import desc.io
 from desc.backend import execute_on_cpu
+from desc.equilibrium import EquilibriaFamily
 
 
 def listall():
@@ -54,6 +55,8 @@ def get(name, data=None):
     assert os.path.exists(path)
 
     eqf = desc.io.load(path)
+    if not isinstance(eqf, EquilibriaFamily):
+        eqf = EquilibriaFamily(eqf)
 
     if data is None:
         return eqf[-1]

desc/examples/precise_QA.py

@@ -1,4 +1,8 @@
-"""Example script for recreating the precise QA configuration of Landreman and Paul."""
+"""Example script for recreating the "precise QA" configuration of Landreman and Paul.
+
+Note that this resembles their optimization process in SIMSOPT, but the final optimized
+equilibrium is slightly different from their VMEC solution.
+"""
 
 from desc import set_device
 

desc/examples/precise_QH.py

@@ -1,4 +1,8 @@
-"""Example script for recreating the precise QH configuration of Landreman and Paul."""
+"""Example script for recreating the "precise QH" configuration of Landreman and Paul.
+
+Note that this resembles their optimization process in SIMSOPT, but the final optimized
+equilibrium is slightly different from their VMEC solution.
+"""
 
 from desc import set_device
 

desc/examples/reactor_QA.py

@@ -0,0 +1,92 @@
+"""Example script for creating a QA configuration at the scale of the ARIES-CS reactor.
+
+It has the boundary of the "precise QA" configuration, but is solved at finite-beta
+and optimized for a self-consistent bootstratp current using the temperature and density
+profiles of Landreman, Buller, and Drevlak.
+"""
+
+from desc import set_device
+
+set_device("gpu")
+
+import numpy as np
+
+from desc.compat import rescale
+from desc.examples import get
+from desc.grid import LinearGrid
+from desc.objectives import (
+    BootstrapRedlConsistency,
+    FixAtomicNumber,
+    FixBoundaryR,
+    FixBoundaryZ,
+    FixCurrent,
+    FixElectronDensity,
+    FixElectronTemperature,
+    FixIonTemperature,
+    FixPsi,
+    ForceBalance,
+    ObjectiveFunction,
+)
+from desc.profiles import PowerSeriesProfile
+
+# same boundary as the precise_QA equilibrium, but at finite beta
+eq = get("precise_QA")
+
+# kinetic profiles
+a = 1.7  # minor radius (ARIES-CS scale)
+B = 5.86  # volume average B (ARIES-CS scale)
+n0 = 2.38e20  # density on axis (<beta>=2.5%)
+T0 = 9.45e3  # density on axis (<beta>=2.5%)
+eq.pressure = None
+eq.current = PowerSeriesProfile(np.zeros(eq.L + 1), np.arange(eq.L + 1), sym=False)
+eq.atomic_number = PowerSeriesProfile([1], [0])
+eq.electron_density = PowerSeriesProfile(
+    n0 * np.array([1, -1]), np.array([0, 10]), sym=True
+)
+eq.electron_temperature = PowerSeriesProfile(
+    T0 * np.array([1, -1]), np.array([0, 2]), sym=True
+)
+eq.ion_temperature = PowerSeriesProfile(
+    T0 * np.array([1, -1]), np.array([0, 2]), sym=True
+)
+
+# solve at finite beta
+eq, _ = eq.solve(verbose=3, copy=True)
+
+# scale to ARIES-CS reactor size
+eq = rescale(eq, L=("a", a), B=("<B>", B), scale_pressure=False, copy=True, verbose=1)
+
+# optimize for self-consistent bootstrap current
+grid = LinearGrid(
+    rho=np.linspace(1 / eq.L_grid, 1, eq.L_grid) - 1 / (2 * eq.L_grid),
+    M=eq.M_grid,
+    N=eq.N_grid,
+    NFP=eq.NFP,
+    sym=True,
+)
+objective = ObjectiveFunction(
+    BootstrapRedlConsistency(eq=eq, grid=grid, helicity=(1, 0))
+)
+constraints = (
+    FixAtomicNumber(eq=eq),
+    FixBoundaryR(eq=eq),
+    FixBoundaryZ(eq=eq),
+    FixCurrent(eq=eq, indices=np.array([0, 1])),
+    FixElectronDensity(eq=eq),
+    FixElectronTemperature(eq=eq),
+    FixIonTemperature(eq=eq),
+    FixPsi(eq=eq),
+    ForceBalance(eq=eq),
+)
+eq, _ = eq.optimize(
+    objective=objective,
+    constraints=constraints,
+    optimizer="proximal-lsq-exact",
+    maxiter=50,
+    gtol=1e-16,
+    verbose=3,
+    copy=True,
+)
+
+# save equilibrium
+eq.save("reactor_QA_output.h5")

desc/examples/regenerate_all_equilibria.py

@@ -68,3 +68,4 @@
 if not args.noprecise:
     spr.run(["python3", "-u", "precise_QA.py"])
     spr.run(["python3", "-u", "precise_QH.py"])
+    spr.run(["python3", "-u", "reactor_QA.py"])

desc/objectives/_generic.py

@@ -11,7 +11,7 @@
 from desc.compute.utils import _parse_parameterization, get_profiles, get_transforms
 from desc.grid import QuadratureGrid
 from desc.optimizable import OptimizableCollection
-from desc.utils import errorif, parse_argname_change
+from desc.utils import errorif, parse_argname_change, setdefault
 
 from .linear_objectives import _FixedObjective
 from .objective_funs import _Objective, collect_docs
@@ -29,6 +29,8 @@ class GenericObjective(_Objective):
     grid : Grid, optional
         Collocation grid containing the nodes to evaluate at. Defaults to
         ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium.
+    compute_kwargs : dict
+        Optional keyword arguments passed to core compute function, eg ``helicity``.
 
     """
 
@@ -37,6 +39,7 @@ class GenericObjective(_Objective):
     )
 
     _print_value_fmt = "Generic objective value: "
+    _static_attrs = ["_compute_kwargs"]
 
     def __init__(
         self,
@@ -52,6 +55,7 @@ def __init__(
         grid=None,
         name="generic",
         jac_chunk_size=None,
+        compute_kwargs=None,
         **kwargs,
     ):
         errorif(
@@ -64,6 +68,7 @@ def __init__(
             target = 0
         self.f = f
         self._grid = grid
+        self._compute_kwargs = setdefault(compute_kwargs, {})
         super().__init__(
             things=thing,
             target=target,
@@ -110,7 +115,9 @@ def build(self, use_jit=True, verbose=1):
         else:
             self._dim_f = grid.num_nodes * np.prod(data_index[self._p][self.f]["dim"])
         profiles = get_profiles(self.f, obj=thing, grid=grid)
-        transforms = get_transforms(self.f, obj=thing, grid=grid)
+        transforms = get_transforms(
+            self.f, obj=thing, grid=grid, **self._compute_kwargs
+        )
         self._constants = {
             "transforms": transforms,
             "profiles": profiles,
@@ -142,6 +149,7 @@ def compute(self, params, constants=None):
             params=params,
             transforms=constants["transforms"],
             profiles=constants["profiles"],
+            **self._compute_kwargs,
         )
         f = data[self.f]
         if self._coordinates == "r":
@@ -279,6 +287,8 @@ class ObjectiveFromUser(_Objective):
     grid : Grid, optional
         Collocation grid containing the nodes to evaluate at. Defaults to
         ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium.
+    compute_kwargs : dict
+        Optional keyword arguments passed to core compute function, eg ``helicity``.
 
     """
 
@@ -308,6 +318,7 @@ def myfun(grid, data):
 
     _units = "(Unknown)"
     _print_value_fmt = "Custom objective value: "
+    _static_attrs = ["_compute_kwargs"]
 
     def __init__(
         self,
@@ -323,6 +334,7 @@ def __init__(
         grid=None,
         name="custom",
         jac_chunk_size=None,
+        compute_kwargs=None,
         **kwargs,
     ):
         errorif(
@@ -335,6 +347,7 @@ def __init__(
             target = 0
         self._fun = fun
         self._grid = grid
+        self._compute_kwargs = setdefault(compute_kwargs, {})
         super().__init__(
             things=thing,
             target=target,
@@ -395,7 +408,9 @@ def get_vars(fun):
         self._dim_f = jax.eval_shape(self._fun_wrapped, dummy_data).size
         self._scalar = self._dim_f == 1
         profiles = get_profiles(self._data_keys, obj=thing, grid=grid)
-        transforms = get_transforms(self._data_keys, obj=thing, grid=grid)
+        transforms = get_transforms(
+            self._data_keys, obj=thing, grid=grid, **self._compute_kwargs
+        )
         self._constants = {
             "transforms": transforms,
             "profiles": profiles,
@@ -429,6 +444,7 @@ def compute(self, params, constants=None):
             params=params,
             transforms=constants["transforms"],
             profiles=constants["profiles"],
+            **self._compute_kwargs,
         )
         f = self._fun_wrapped(data)
         return f

setup.cfg

@@ -11,7 +11,6 @@ versionfile_build = desc/_version.py
 tag_prefix = v
 parentdir_prefix = desc-
 
-
 [coverage:run]
 # Settings to control coverage.py or pytest with coverage flags: "--cov" (requires pytest-cov)
 # we only care about the coverage of the source itself
@@ -24,6 +23,7 @@ omit =
      desc/_version.py
      desc/examples/precise_QH.py
      desc/examples/precise_QA.py
+     desc/examples/reactor_QA.py
      desc/examples/regenerate_all_equilibria.py
 
 [coverage:report]
@@ -36,7 +36,6 @@ exclude_lines =
 
 ignore_errors = True
 
-
 [tool:pytest]
 markers=
 	unit: marks small unit tests
@@ -55,7 +54,6 @@ filterwarnings=
     # ignore DeprecationWarnings from dependency packages
     ignore::DeprecationWarning:(?!desc)
 
-
 [flake8]
 # Primarily ignoring whitespace, indentation, and commenting etiquette that black does not catch
 # These will be fixed in a code-cleaning branch in the future
@@ -80,8 +78,9 @@ per-file-ignores =
     # need imports in weird order for selecting device before benchmarks
     tests/benchmarks/benchmark*.py: E402
     # stop complaining about setting gpu before import other desc stuff
-    desc/examples/precise_QH.py: E402
     desc/examples/precise_QA.py: E402
+    desc/examples/precise_QH.py: E402
+    desc/examples/reactor_QA.py: E402
 max-line-length = 88
 exclude =
     docs/*

tests/test_objective_funs.py

@@ -1754,6 +1754,47 @@ def test_objective_compute(self):
             obj.compute(eq.params_dict), grid.compress(data["Gamma_c"])
         )
 
+    @pytest.mark.unit
+    def test_generic_with_kwargs(self):
+        """Test GenericObjective with keyword arguments. Related to issue #1224."""
+        eq = desc.examples.get("reactor_QA")
+
+        def fusion_gain(grid, data):
+            p_out = data["P_fusion"]
+            p_in = data["P_ISS04"]
+            Q = p_out / p_in
+            return Q
+
+        obj_p0_out = FusionPower(eq, fuel="DT")
+        obj_p0_out.build()
+        p0_out = obj_p0_out.compute(*obj_p0_out.xs(eq))
+
+        obj_p0_in = HeatingPowerISS04(eq, H_ISS04=1.2, gamma=0)
+        obj_p0_in.build()
+        p0_in = obj_p0_in.compute(*obj_p0_in.xs(eq))
+
+        q0 = p0_out / p0_in
+
+        obj_p1_out = GenericObjective("P_fusion", eq, compute_kwargs={"fuel": "DT"})
+        obj_p1_out.build()
+        p1_out = obj_p1_out.compute(*obj_p1_out.xs(eq))
+
+        obj_p1_in = GenericObjective(
+            "P_ISS04", eq, compute_kwargs={"H_ISS04": 1.2, "gamma": 0}
+        )
+        obj_p1_in.build()
+        p1_in = obj_p1_in.compute(*obj_p1_in.xs(eq))
+
+        obj_q1 = ObjectiveFromUser(
+            fusion_gain, eq, compute_kwargs={"fuel": "DT", "H_ISS04": 1.2, "gamma": 0}
+        )
+        obj_q1.build()
+        q1 = obj_q1.compute(*obj_q1.xs(eq))
+
+        np.testing.assert_allclose(p0_out, p1_out)
+        np.testing.assert_allclose(p0_in, p1_in)
+        np.testing.assert_allclose(q0, q1)
+
 
 @pytest.mark.regression
 def test_derivative_modes():