Dev guide

desc/compute/_core.py

@@ -1843,7 +1843,7 @@ def _omega_rr(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="omega_rrr",
-    label="\\partial_{\rho \\rho \\rho} \\omega",
+    label="\\partial_{\\rho \\rho \\rho} \\omega",
     units="rad",
     units_long="radians",
     description="Toroidal stream function, third radial derivative",
@@ -1861,7 +1861,7 @@ def _omega_rrr(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="omega_rrrr",
-    label="\\partial_{\rho \\rho \\rho \\rho} \\omega",
+    label="\\partial_{\\rho \\rho \\rho \\rho} \\omega",
     units="rad",
     units_long="radians",
     description="Toroidal stream function, fourth radial derivative",
@@ -1879,7 +1879,7 @@ def _omega_rrrr(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="omega_rrrt",
-    label="\\partial_{\rho \\rho \\rho \\theta} \\omega",
+    label="\\partial_{\\rho \\rho \\rho \\theta} \\omega",
     units="rad",
     units_long="radians",
     description="Toroidal stream function, fourth derivative wrt radial coordinate"
@@ -1898,7 +1898,7 @@ def _omega_rrrt(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="omega_rrrz",
-    label="\\partial_{\rho \\rho \\rho \\zeta} \\omega",
+    label="\\partial_{\\rho \\rho \\rho \\zeta} \\omega",
     units="rad",
     units_long="radians",
     description="Toroidal stream function, fourth derivative wrt radial coordinate"

desc/compute/_equil.py

@@ -465,7 +465,7 @@ def _F_zeta(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="F_helical",
-    label="F_{helical}",
+    label="F_{\\mathrm{helical}}",
     units="A",
     units_long="Amperes",
     description="Covariant helical component of force balance error",

desc/compute/_field.py

@@ -3225,7 +3225,7 @@ def _min_tz_modB(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="effective r/R0",
-    label="(r / R_0)_{effective}",
+    label="(r / R_0)_{\\mathrm{effective}}",
     units="~",
     units_long="None",
     description="Effective local inverse aspect ratio, based on max and min |B|",

desc/compute/_geometry.py

@@ -332,7 +332,7 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="a_major/a_minor",
-    label="a_{major} / a_{minor}",
+    label="a_{\\mathrm{major}} / a_{\\mathrm{minor}}",
     units="~",
     units_long="None",
     description="Maximum elongation",

desc/compute/_metric.py

@@ -55,7 +55,7 @@ def _sqrtg_pest(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_theta x e_zeta|",
-    label="|e_{\\theta} \\times e_{\\zeta}|",
+    label="|\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant rho surface",
@@ -79,7 +79,7 @@ def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_theta x e_zeta|_r",
-    label="\\partial_{\\rho} |e_{\\theta} \\times e_{\\zeta}|",
+    label="\\partial_{\\rho} |\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant rho surface"
@@ -113,7 +113,7 @@ def _e_theta_x_e_zeta_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_theta x e_zeta|_rr",
-    label="\\partial_{\\rho \\rho} |e_{\\theta} \\times e_{\\zeta}|",
+    label="\\partial_{\\rho\\rho} |\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant rho surface"
@@ -152,7 +152,7 @@ def _e_theta_x_e_zeta_rr(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_zeta x e_rho|",
-    label="|e_{\\zeta} \\times e_{\\rho}|",
+    label="|\\mathbf{e}_{\\zeta} \\times \\mathbf{e}_{\\rho}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant theta surface",
@@ -174,7 +174,7 @@ def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_rho x e_theta|",
-    label="|e_{\\rho} \\times e_{\\theta}|",
+    label="|\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\theta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant zeta surface",
@@ -198,7 +198,7 @@ def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_rho x e_theta|_r",
-    label="\\partial_{\\rho} |e_{\\rho} \\times e_{\\theta}|",
+    label="\\partial_{\\rho} |\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\theta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant zeta surface"
@@ -232,7 +232,7 @@ def _e_rho_x_e_theta_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_rho x e_theta|_rr",
-    label="\\partial_{\\rho \\rho} |e_{\\rho} \\times e_{\\theta}|",
+    label="\\partial_{\\rho \\rho} |\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\theta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant zeta surface"
@@ -1349,7 +1349,7 @@ def _g_sup_tz(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rr_r",
-    label="g^{\\rho}{\\rho}_{\\rho}",
+    label="\\partial_{\\rho} g^{\\rho}{\\rho}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Radial element of contravariant metric tensor, "
@@ -1368,7 +1368,7 @@ def _g_sup_rr_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rt_r",
-    label="g^{\\rho}{\\theta}_{\\rho}",
+    label="\\partial_{\\rho} g^{\\rho}{\\theta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Poloidal element of contravariant metric tensor, "
@@ -1389,7 +1389,7 @@ def _g_sup_rt_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rz_r",
-    label="g^{\\rho}{\\zeta}_{\\rho}",
+    label="\\partial_{\\rho} g^{\\rho}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Toroidal element of contravariant metric tensor, "
@@ -1410,7 +1410,7 @@ def _g_sup_rz_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^tt_r",
-    label="g^{\\theta}{\\theta}_{\\rho}",
+    label="\\partial_{\\rho} g^{\\theta}{\\theta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Poloidal/Poloidal element of contravariant metric tensor, "
@@ -1429,7 +1429,7 @@ def _g_sup_tt_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^tz_r",
-    label="g^{\\theta}{\\zeta}_{\\rho}",
+    label="\\partial_{\\rho} g^{\\theta}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Poloidal/Toroidal element of contravariant metric tensor, "
@@ -1450,7 +1450,7 @@ def _g_sup_tz_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^zz_r",
-    label="g^{\\zeta}{\\zeta}_{\\rho}",
+    label="\\partial_{\\rho} g^{\\zeta}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Toroidal/Toroidal element of contravariant metric tensor, "
@@ -1469,7 +1469,7 @@ def _g_sup_zz_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rr_t",
-    label="g^{\\rho}{\\rho}_{\\theta}",
+    label="\\partial_{\\theta} g^{\\rho}{\\rho}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Radial element of contravariant metric tensor, "
@@ -1488,7 +1488,7 @@ def _g_sup_rr_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rt_t",
-    label="g^{\\rho}{\\theta}_{\\theta}",
+    label="\\partial_{\\theta} g^{\\rho}{\\theta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Poloidal element of contravariant metric tensor, "
@@ -1509,7 +1509,7 @@ def _g_sup_rt_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rz_t",
-    label="g^{\\rho}{\\zeta}_{\\theta}",
+    label="\\partial_{\\theta} g^{\\rho}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Toroidal element of contravariant metric tensor, "
@@ -1530,7 +1530,7 @@ def _g_sup_rz_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^tt_t",
-    label="g^{\\theta}{\\theta}_{\\theta}",
+    label="\\partial_{\\theta} g^{\\theta}{\\theta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Poloidal/Poloidal element of contravariant metric tensor, "
@@ -1549,7 +1549,7 @@ def _g_sup_tt_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^tz_t",
-    label="g^{\\theta}{\\zeta}_{\\theta}",
+    label="\\partial_{\\theta} g^{\\theta}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Poloidal/Toroidal element of contravariant metric tensor, "
@@ -1570,7 +1570,7 @@ def _g_sup_tz_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^zz_t",
-    label="g^{\\zeta}{\\zeta}_{\\theta}",
+    label="\\partial_{\\theta} g^{\\zeta}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Toroidal/Toroidal element of contravariant metric tensor, "
@@ -1589,7 +1589,7 @@ def _g_sup_zz_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rr_z",
-    label="g^{\\rho}{\\rho}_{\\zeta}",
+    label="\\partial_{\\zeta} g^{\\rho}{\\rho}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Radial element of contravariant metric tensor, "
@@ -1608,7 +1608,7 @@ def _g_sup_rr_z(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rt_z",
-    label="g^{\\rho}{\\theta}_{\\zeta}",
+    label="\\partial_{\\zeta} g^{\\rho}{\\theta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Poloidal element of contravariant metric tensor, "
@@ -1629,7 +1629,7 @@ def _g_sup_rt_z(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^rz_z",
-    label="g^{\\rho}{\\zeta}_{\\zeta}",
+    label="\\partial_{\\zeta} g^{\\rho}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Radial/Toroidal element of contravariant metric tensor, "
@@ -1650,7 +1650,7 @@ def _g_sup_rz_z(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^tt_z",
-    label="g^{\\theta}{\\theta}_{\\zeta}",
+    label="\\partial_{\\zeta} g^{\\theta}{\\theta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Poloidal/Poloidal element of contravariant metric tensor, "
@@ -1669,7 +1669,7 @@ def _g_sup_tt_z(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^tz_z",
-    label="g^{\\theta}{\\zeta}_{\\zeta}",
+    label="\\partial_{\\zeta} g^{\\theta}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Poloidal/Toroidal element of contravariant metric tensor, "
@@ -1690,7 +1690,7 @@ def _g_sup_tz_z(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="g^zz_z",
-    label="g^{\\zeta}{\\zeta}_{\\zeta}",
+    label="\\partial_{\\zeta} g^{\\zeta}{\\zeta}",
     units="m^-2",
     units_long="inverse square meters",
     description="Toroidal/Toroidal element of contravariant metric tensor, "

docs/adding_objectives.rst

@@ -88,6 +88,7 @@ A full example objective with comments describing key points is given below:
             name="QS triple product",
         ):
             # we don't have to do much here, mostly just call ``super().__init__()``
+            # to inherit common initialization logic from ``desc.objectives._Objective``
             if target is None and bounds is None:
                 target = 0 # default target value
             self.grid = grid
@@ -118,9 +119,13 @@ A full example objective with comments describing key points is given below:
                 self.grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)
 
             # dim_f = size of the output vector returned by self.compute
-            # usually the same as self.grid.num_nodes, unless you're doing some downsampling
-            # or averaging etc.
-            self._dim_f = self.grid.num_nodes
+            # self.compute refers to the objective's own compute method
+            # Typically an objective returns the output of a quantity computed in
+            # ``desc.compute``, with some additional scale factor.
+            # In these cases dim_f should match the size of the quantity calculated in
+            # ``desc.compute`` (for example self.grid.num_nodes).
+            # If the objective does post-processing on the quantity, like downsampling or
+            # averaging, then dim_f should be changed accordingly.
             # What data from desc.compute is needed? Here we want the QS triple product.
             self._data_keys = ["f_T"]
             # what arguments should be passed to self.compute

docs/dev_guide/backend.rst

@@ -0,0 +1,23 @@
+Backend
+-------
+
+
+DESC uses JAX for faster compile times, automatic differentiation, and other scientific computing tools.
+The purpose of ``backend.py`` is to determine whether DESC may take advantage of JAX and GPUs or default to standard ``numpy`` and CPUs.
+
+JAX provides a ``numpy`` style API for array operations.
+In many cases, to take advantage of JAX, one only needs to replace calls to ``numpy`` with calls to ``jax.numpy``.
+A convenient way to do this is with the import statement ``import jax.numpy as jnp``.
+
+Of course if such an import statement is used in DESC, and DESC is run on a machine where JAX is not installed, then a runtime error is thrown.
+We would prefer if DESC still works on machines where JAX is not installed.
+With that goal, in functions which can benefit from JAX, we use the following import statement: ``from desc.backend import jnp``.
+``desc.backend.jnp`` is an alias to ``jax.numpy`` if JAX is installed and ``numpy`` otherwise.
+
+While ``jax.numpy`` attempts to serve as a drop in replacement for ``numpy``, it imposes some constraints on how the code is written.
+For example, ``jax.numpy`` arrays are immutable.
+This means in-place updates to elements in arrays is not possible.
+To update elements in ``jax.numpy`` arrays, memory needs to be allocated to create a new array with the updated element.
+Similarly, JAX's JIT compilation requires code flow structures such as loops and conditionals to be written in a specific way.
+
+The utility functions in ``desc.backend`` provide a simple interface to perform these operations.