Data structure

To take full advantage of pyAQSC, it is helpful to understand its two main data structures, ChiPhiFunc and ChiPhiEpsFunc.

ChiPhiFunc

ChiPhiFunc is an array-like class. It is the primary data structure pyAQSC performs computations with. It is traced in JAX for auto-differentiation.

ChiPhiFunc manages functions of helical and/or toroidal angle $(\chi, \phi)$. As discussed in Ref.2, all such functions in this expansions can be represented as even/odd Fourier series of $\chi$ with $\phi$-dependent coefficients: $$F_n(\chi, \phi) = \sum_{m=0|1}^n e^{im} F_{n,m}(\phi) + e^{-im} F_{n,-m}(\phi)$$ The $\phi$ dependence of $F_{n,\pm m}(\phi)$ are stored on $n$ uniformly spaced grid points over one field period, $\phi = 0, \frac{2\pi}{n_{fp}}\frac{1}{n}, ..., \frac{2\pi}{n_{fp}}\frac{(n-1)}{n}$, where $n_{fp}$ is the number of field period.

Numerical operations between ChiPhiFunc's with different grid number or nfp are not supported.

Internally, the coefficients are stored as the 2D numpy array, ChiPhiFunc.content, whose axis=0 corresponds to $m=-n, -n+2, ..., n$, and axis=1 corresponds to $\phi$ grid points. To evaluate a ChiPhiFunc, use ChiPhiFunc.eval(chi, phi).

ChiPhiFunc supports a number of numerical operations, including: - Basic arithmetic operations +, -, *, /, **. - Derivatives and anti-derivatives ChiPhiFunc.dphi(), ChiPhiFunc.dchi(), ChiPhiFunc.antid_chi(), ChiPhiFunc.integrate_phi_fft() - Indexing operation ChiPhiFunc[m] that extracts the $m$-th harmonic coefficient $F_{n,m}(\phi)$ as a $\chi$-independent ChiPhiFunc. - Low-pass filter in $\phi$, ChiPhiFunc.filter() - Lambda functions output, ChiPhiFunc.get_lambda() - Simple ODE solves

For more information, please refer to ChiPhiFunc API.

ChiPhiEpsFunc

ChiPhiEpsFunc is a list-like class containing scalars, ChiPhiFunc's, and ChiPhiFuncSpecial's. It is the primary data structure pyAQSC stores results in. To evaluate a ChiPhiEpsFunc, use ChiPhiEpsFunc.eval(psi, chi, phi, sq_eps_series). The $\epsilon$, $\chi$, $\phi$ derivative of a ChiPhiEpsFunc can be calculated with ChiPhiEpsFunc.deps(), ChiPhiEpsFunc.dchi() and ChiPhiEpsFunc.dphi().

ChiPhiEpsFunc manages the power-Fourier series: $$F(\psi, \chi, \phi)\\ =\sum_{n=0}^{n_{max}}\epsilon^nF_n(\chi,\phi) (\epsilon\equiv\sqrt{\psi})\\ =\sum_{n=0}^{n_{max}}\epsilon^n\sum_{m=0|1}^n e^{im} F_{n,m}(\phi) + e^{-im} F_{n,-m}(\phi).$$

ChiPhiEpsFunc supports __getitem__ and append, but not __setitem__. ChiPhiEpsFunc[n] Extracts the $n$-th order power coefficient, which may be a scalar or a ChiPhiFunc.

ChiPhiEpsFunc differs from list, in that ChiPhiEpsFunc.__getitem__ return ChiPhiFunc(nfp=0) when the index is negative or out-of-bound. This is required to work with the expressions generated by Maxima scripts.

ChiPhiEpsFunc supports a number of numerical operations, including: - Basic arithmetic operations +, -, * with scalars, ChiPhiFuncs and ChiPhiEpsFuncs. - Derivatives and anti-derivatives ChiPhiEpsFunc.dphi(), ChiPhiEpsFunc.dchi(), ChiPhiEpsFunc.deps()

Naming convention

Quantities are named by name_dependence. - Quantities with coef in its name are power series (ChiPhiEpsFunc). E.g., B_theta_coef_cp, iota_coef - Quantities with c and/or p at the end has $\chi$ and/or $\phi$ dependence. E.g., B_theta_coef_cp, B_denom_coef_c, kap_p. They are either ChiPhiFunc's or ChiPhiEpsFunc's containing ChiPhiFunc's. - The only exception is $dl/d\phi$, dl_p. This quantity is a float constant.

Note: zero and error handling

ChiPhiFunc with ChiPhiFunc.nfp<=0 is reserved for zero and error handling. This is implemented because JAX cannot determine traced variables' value and has no runtime error support.

nfp<0 means a run-time error. The list of all implemented errors is:

-2: Invalid/mismatched nfp
-3: Invalid mode number
-4: Mismatched even/oddness/direct division by chi-dependent ChiPhiFunc
-5: Operation between ChiPhiFunc and another item that's not a constant or ChiPhiFunc
-6: Operation between 2 ChiPhiFuncs with mismatching lengths.
-7: Invalid content shape
-8: /zero (special ChiPhiFunc)
-9: Incorrect argument for pow
-10: Incorrect argument for dchi
-11: Incorrect argument for dphi
-12: DEPRECIATED
-13: Invalid object to diff()
-14: Inconsistent type in ChiPhiEpsFunc
-15: Insufficient argument in Equilibrium
-16: Filter mode not recognized
-17: Incorrect looped equation even/odd-ness

nfp=0 means a ChiPhiFunc is known to be zero. Because JAX cannot distinguish the value of variables, expressions such as

int(0)*<even ChiPhiFunc>+<odd ChiPhiFunc>

will throw an error even though it should produce an odd ChiPhiFunc. This happens because JAX cannot tell if a traced integer is zero or not, or overload * for static inputs. As a result, 0*<even ChiPhiFunc> yields an even ChiPhiFunc with all zero coefficients.