Adding new physics quantities

All calculation of physics quantities takes place in desc.compute

As an example, we’ll walk through the calculation of the contravariant radial component of the plasma current density \(J^\rho\)

The full code is below:

from desc.data_index import register_compute_fun

@register_compute_fun(
    name="J^rho",
    label="J^{\\rho}",
    units="A \\cdot m^{-3}",
    units_long="Amperes / cubic meter",
    description="Contravariant radial component of plasma current density",
    dim=1,
    params=[],
    transforms={},
    profiles=[],
    coordinates="rtz",
    data=["sqrt(g)", "B_zeta_t", "B_theta_z"],
    axis_limit_data=["sqrt(g)_r", "B_zeta_rt", "B_theta_rz"],
    resolution_requirement="",
    parameterization="desc.equilibrium.equilibrium.Equilibrium",
)
def _J_sup_rho(params, transforms, profiles, data, **kwargs):
    # At the magnetic axis,
    # ∂_θ (𝐁 ⋅ 𝐞_ζ) - ∂_ζ (𝐁 ⋅ 𝐞_θ) = 𝐁 ⋅ (∂_θ 𝐞_ζ - ∂_ζ 𝐞_θ) = 0
    # because the partial derivatives commute. So 𝐉^ρ is of the indeterminate
    # form 0/0 and we may compute the limit as follows.
    data["J^rho"] = (
        transforms["grid"].replace_at_axis(
            (data["B_zeta_t"] - data["B_theta_z"]) / data["sqrt(g)"],
            lambda: (data["B_zeta_rt"] - data["B_theta_rz"]) / data["sqrt(g)_r"],
        )
    ) / mu_0
    return data

The decorator register_compute_fun tells DESC that the quantity exists and contains metadata about the quantity. The necessary fields are detailed below:

  • name: A short, meaningful name that is used elsewhere in the code to refer to the quantity. This name will appear in the data dictionary returned by Equilibrium.compute, and is also the argument passed to compute to calculate the quantity. IE, Equilibrium.compute("J^rho") will return a dictionary containing J^rho as well as all the intermediate quantities needed to compute it. General conventions are that covariant components of a vector are called X_rho etc, contravariant components X^rho etc, and derivatives by a single character subscript, X_r etc for \(\partial_{\rho} X\)

  • label: a LaTeX style label used for plotting.

  • units: SI units of the quantity in LaTeX format

  • units_long: SI units of the quantity, spelled out

  • description: A short description of the quantity

  • dim: Dimensionality of the quantity. Vectors (such as magnetic field) have dim=3, local scalar quantities (such as vector components, pressure, volume element, etc) have dim=1, global scalars (such as total volume, aspect ratio, etc) have dim=0

  • params: list of strings of Equilibrium parameters needed to compute the quantity such as R_lmn, Z_lmn etc. These will be passed into the compute function as a dictionary in the params argument. Note that this only includes direct dependencies (things that are used in this function). For most quantities, this will be an empty list. For example, if the function relies on R_t, this dependency should be specified as a data dependency (see below), while the function to compute R_t itself will depend on R_lmn

  • transforms: a dictionary of what transform objects are needed, with keys being the quantity being transformed (R, p, etc) and the values are a list of derivative orders needed in rho, theta, zeta. IE, if the quantity requires \(R_{\rho}{\zeta}{\zeta}\), then transforms should be {"R": [[1, 0, 2]]} indicating a first derivative in rho and a second derivative in zeta. Note that this only includes direct dependencies (things that are used in this function). For most quantities this will be an empty dictionary.

  • profiles: List of string of Profile quantities needed, such as pressure or iota. Note that this only includes direct dependencies (things that are used in this function). For most quantities this will be an empty list.

  • coordinates: String denoting which coordinate the quantity depends on. Most will be "rtz" indicating it is a function of \(\rho, \theta, \zeta\). Profiles and flux surface quantities would have coordinates="r" indicating it only depends on \(\rho\)

  • data: What other physics quantities are needed to compute this quantity. In our example, we need the poloidal (theta) derivative of the covariant toroidal (zeta) component of the magnetic field B_zeta_t, the toroidal derivative of the covariant poloidal component of the magnetic field B_theta_z, and the 3-D volume Jacobian determinant sqrt(g). These dependencies will be passed to the compute function as a dictionary in the data argument. Note that this only includes direct dependencies (things that are used in this function). For example, we need sqrt(g), which itself depends on e_rho, etc. But we don’t need to specify those sub-dependencies here.

  • axis_limit_data: Some quantities require additional work to compute at the magnetic axis. A Python lambda function is used to lazily compute the magnetic axis limits of these quantities. These lambda functions are evaluated only when the computational grid has a node on the magnetic axis to avoid potentially expensive computations. In our example, we need the radial derivatives of each the quantities mentioned above to evaluate the magnetic axis limit. These dependencies are specified in axis_limit_data. The dependencies specified in this list are marked to be computed only when there is a node at the magnetic axis.

  • resolution_requirement: Resolution requirements in coordinates. I.e. “r” expects radial resolution in the grid. Likewise, “rtz” is shorthand for “rho, theta, zeta” and indicates the computation expects a grid with radial, poloidal, and toroidal resolution. If the computation simply performs pointwise operations, instead of a reduction (such as integration) over a coordinate, then an empty string may be used to indicate no requirements.

  • parameterization: what sorts of DESC objects is this function for. Most functions will just be for Equilibrium, but some methods may also be for desc.geometry.core.Curve, or specific types eg desc.geometry.curve.FourierRZCurve. If a quantity is computed differently for a subclass versus a superclass, then one may define a compute function for the superclass (e.g. for desc.geometry.Curve) which will be used for that class and any of its subclasses, and then if a specific subclass requires a different method, one may define a second compute function for the same quantity, with a parameterization for that subclass (e.g. desc.geometry.curve.SplineXYZCurve). See the compute definitions for the length quantity in compute/_curve.py for an example of this, which is similar to the inheritance structure of Python classes.

  • kwargs: If the compute function requires any additional arguments they should be specified like kwarg="description" where kwarg is replaced by the actual keyword argument, and "description" is a string describing what it is. Most quantities do not take kwargs.

The function itself should have a signature of the form

_foo(params, transforms, profiles, data, **kwargs)

Our convention is to start the function name with an underscore and have the it be something like the name attribute, but name of the function doesn’t actually matter as long as it is registered. params, transforms, profiles, and data are dictionaries containing the needed dependencies specified by the decorator. The function itself should do any calculation needed using these dependencies and then add the output to the data dictionary and return it. The key in the data dictionary should match the name of the quantity.

Once a new quantity is added to the desc.compute module, there are two final steps involving the testing suite which must be checked. The first step is implementing the correct axis limit, or marking it as not finite or not implemented. We can check whether the axis limit currently evaluates as finite by computing the quantity on a grid with nodes at the axis.

from desc.examples import get
from desc.grid import LinearGrid
import numpy as np

eq = get("HELIOTRON")
grid = LinearGrid(rho=np.array([0.0]), M=4, N=8, axis=True)
new_quantity = eq.compute(name="new_quantity_name", grid=grid)["new_quantity_name"]
print(np.isfinite(new_quantity).all())

if False is printed, then the limit of the quantity does not evaluate as finite which can be due to 3 reasons:

  • The limit is actually not finite, in which case please add the new quantity to the not_finite_limits set in tests/test_axis_limits.py.

  • The new quantity has an indeterminate expression at the magnetic axis, in which case you should try to implement the correct limit as done in the example for J^rho above. If you wish to skip implementing the limit at the magnetic axis, please add the new quantity to the not_implemented_limits set in tests/test_axis_limits.py.

  • The new quantity includes a dependency whose limit at the magnetic axis has not been implemented. The tests automatically detect this, so no further action is needed from developers in this case.

The second step is to run the test_compute_everything test located in the tests/test_compute_everything.py file. This can be done with the command pytest tests/test_compute_everything.py. This test is a regression test to ensure that compute quantities in each new update of DESC do not differ significantly from previous versions of DESC. Since the new quantity did not exist in previous versions of DESC, one must run this test and commit the outputted tests/inputs/master_compute_data_rpz.pkl file which is updated automatically when a new quantity is detected.

Compute function may take additional **kwargs arguments to provide more information to the function. One example of this kind of compute function is P_ISS04 which has a keyword argument H_ISS04.

@register_compute_fun(
  name="P_ISS04",
  label="P_{ISS04}",
  units="W",
  units_long="Watts",
  description="Heating power required by the ISS04 energy confinement time scaling",
  dim=0,
  params=[],
  transforms={"grid": []},
  profiles=[],
  coordinates="",
  data=["a", "iota", "rho", "R0", "W_p", "<ne>_vol", "<|B|>_axis"],
  method="str: Interpolation method. Default 'cubic'.",
  H_ISS04="float: ISS04 confinement enhancement factor. Default 1.",
)
def _P_ISS04(params, transforms, profiles, data, **kwargs):
    rho = transforms["grid"].compress(data["rho"], surface_label="rho")
    iota = transforms["grid"].compress(data["iota"], surface_label="rho")
    fx = {}
    if "iota_r" in data:
        fx["fx"] = transforms["grid"].compress(
            data["iota_r"]
        )  # noqa: unused dependency
    iota_23 = interp1d(2 / 3, rho, iota, method=kwargs.get("method", "cubic"), **fx)
    data["P_ISS04"] = 1e6 * (  # MW -> W
        jnp.abs(data["W_p"] / 1e6)  # J -> MJ
        / (
            0.134
            * data["a"] ** 2.28  # m
            * data["R0"] ** 0.64  # m
            * (data["<ne>_vol"] / 1e19) ** 0.54  # 1/m^3 -> 1e19/m^3
            * data["<|B|>_axis"] ** 0.84  # T
            * iota_23**0.41
            * kwargs.get("H_ISS04", 1)
        )
    ) ** (1 / 0.39)
    return data

This function can be called by following notation,

from desc.compute.utils import _compute as compute_fun

# Compute P_ISS04
# specify gamma and H_ISS04 values as keyword arguments
data = compute_fun(
          "desc.equilibrium.equilibrium.Equilibrium",
          "P_ISS04",
          params=params,
          transforms=transforms,
          profiles=profiles,
          gamma=gamma,
          H_ISS04=H_ISS04,
      )
P_ISS04 = data["P_ISS04"]

Note: Here we used _compute instead of compute to be able to call this function inside a jitted objective function. However, for normal use both functions should work. **kwargs can also be passed to eq.compute.

data = eq.compute(names="P_ISS04", gamma=gamma, H_ISS04=H_ISS04)
P_ISS04 = data["P_ISS04"]