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"],
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 thedata
dictionary returned byEquilibrium.compute
, and is also the argument passed tocompute
to calculate the quantity. IE,Equilibrium.compute("J^rho")
will return a dictionary containingJ^rho
as well as all the intermediate quantities needed to compute it. General conventions are that covariant components of a vector are calledX_rho
etc, contravariant componentsX^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 formatunits_long
: SI units of the quantity, spelled outdescription
: A short description of the quantitydim
: Dimensionality of the quantity. Vectors (such as magnetic field) havedim=3
, local scalar quantities (such as vector components, pressure, volume element, etc) havedim=1
, global scalars (such as total volume, aspect ratio, etc) havedim=0
params
: list of strings ofEquilibrium
parameters needed to compute the quantity such asR_lmn
,Z_lmn
etc. These will be passed into the compute function as a dictionary in theparams
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 onR_t
, this dependency should be specified as a data dependency (see below), while the function to computeR_t
itself will depend onR_lmn
transforms
: a dictionary of whattransform
objects are needed, with keys being the quantity being transformed (R
,p
, etc) and the values are a list of derivative orders needed inrho
,theta
,zeta
. IE, if the quantity requires \(R_{\rho}{\zeta}{\zeta}\), thentransforms
should be{"R": [[1, 0, 2]]}
indicating a first derivative inrho
and a second derivative inzeta
. 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 ofProfile
quantities needed, such aspressure
oriota
. 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 havecoordinates="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 fieldB_zeta_t
, the toroidal derivative of the covariant poloidal component of the magnetic fieldB_theta_z
, and the 3-D volume Jacobian determinantsqrt(g)
. These dependencies will be passed to the compute function as a dictionary in thedata
argument. Note that this only includes direct dependencies (things that are used in this function). For example, we needsqrt(g)
, which itself depends one_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 inaxis_limit_data
. The dependencies specified in this list are marked to be computed only when there is a node at the magnetic axis.parameterization
: what sorts of DESC objects is this function for. Most functions will just be forEquilibrium
, but some methods may also be fordesc.geometry.Curve
, or specific types egdesc.geometry.FourierRZCurve
.kwargs
: If the compute function requires any additional arguments they should be specified likekwarg="thing"
where the value is the name of the keyword argument that will be passed to the compute function. 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 evalutates 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 intests/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 thenot_implemented_limits
set intests/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_funs.py
file.
This can be done with the command pytest -k test_compute_everything tests/test_compute_funs.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.pkl
file which is updated automatically when a new quantity is detected.