diff --git a/docs/physics/base_fluid_models/RMHD/weak_form.md b/docs/physics/base_fluid_models/RMHD/weak_form.md
new file mode 100644
index 0000000000..236f079243
--- /dev/null
+++ b/docs/physics/base_fluid_models/RMHD/weak_form.md
@@ -0,0 +1,2055 @@
+---
+title: "Model 600 weak form"
+nav_order: 3
+parent: "Reduced MHD models"
+grand_parent: "Base Fluid Models"
+layout: default
+render_with_liquid: false
+---
+
+# Model 600: equations in weak form
+
+This page documents the equations of the model-600 element routine
+`models/model600/mod_elt_matrix_fft.f90` in weak form. It covers the equations
+for $\psi$, $u$, $j$, $\omega$, $\rho$, $v_\parallel$, $\rho_{imp}$, $\rho_n$,
+$T_i$, $T_e$ and $T$; see [Status](#status) for what is not covered yet.
+
+The equations are written, as Python, in
+`util/equation_codegen/src/jorek_equations/model600.py`, and the equation
+checker compares them with the Fortran term by term. All math on this page is
+generated from that file (see [Regenerating this page](#regenerating-this-page)).
+Where the Python and the Fortran differ, the difference is known and recorded:
+it is listed under [Known differences](#known-differences) and explained in
+[`JOREK_FINDINGS.md`](https://github.com/iterorganization/JOREK/blob/develop/util/equation_codegen/JOREK_FINDINGS.md).
+
+**The integral is implied.** Each equation is multiplied by a test function
+$v$ (a finite-element basis function, not a velocity) and integrated over the
+plasma volume, $\int \ldots\,\mathrm{d}V$ with
+$\mathrm{d}V = w_V\,\mathrm{d}s\,\mathrm{d}t_{\mathrm{el}}\,\mathrm{d}\phi$ and
+$w_V = R\,\mathcal{J}$, where $(s, t_{\mathrm{el}})$ are the element
+coordinates and $\mathcal{J}$ their Jacobian. (The code calls the second
+element coordinate `t`; on this page $t$ is always time.) The integral, the
+common volume weight $w_V$ and the coordinate differentials are not written,
+so $\frac{\partial}{\partial t}\big(v\,\rho\big) = v\,S_p + \ldots$ stands for
+$\frac{\partial}{\partial t}\int v\,\rho\,\mathrm{d}V = \int \big(v\,S_p + \ldots\big)\,\mathrm{d}V$.
+Terms with derivatives of $v$ come from integration by parts; see
+[Conventions](#conventions) for the assumptions behind this.
+
+$\nabla_{\mathrm{pol}}$ is the poloidal gradient and $[a,b]$ the poloidal
+bracket. The term groups written $\mathcal{C}$, $\mathcal{D}$, $\mathcal{T}$,
+$Q$, ... are defined under [Shared operators](#shared-operators), and the
+other symbols under [Notation](#notation). Refer also to the
+[RMHD model](rmhd_model.md), [notation](../../notation.md) and
+[normalization](../../normalization.md) pages.
+
+## Variables
+
+The model is configured at compile time by the logical parameters in
+`models/model600/mod_model_settings.f90` (set with `util/config.sh`). When a
+variable is switched off, the element routine replaces it by a constant:
+$\rho = 1$, $v_\parallel = 0$, $\rho_n = 0$, $\rho_{imp} = 0$, and its
+equation is not assembled. The equations below are written with every variable
+present.
+
+| Python | Symbol | Meaning | Present when |
+|---|---|---|---|
+| `psi` | $\psi$ | poloidal magnetic flux | always |
+| `u` | $u$ | stream function of the $E\times B$ velocity | always |
+| `j` | $j$ | toroidal current, $j = -R\,\mathbf{j}\cdot\mathbf{e}_\phi = \Delta^*\psi$; it includes a factor $R$ (see the [RMHD model](rmhd_model.md)) | always |
+| `omega` | $\omega$ | vorticity, $\omega = \nabla^2_{\mathrm{pol}} u$ (the [vorticity definition](#eq-w) below) | always |
+| `rho` | $\rho$ | total mass density, impurities included, in units of the main-ion mass | `with_rho = .true.` |
+| `rhoimp` | $\rho_{imp}$ | impurity mass density, in the same units | `with_impurities = .true.` |
+| `rhon` | $\rho_n$ | neutral density | `with_neutrals = .true.` |
+| `vpar` | $v_\parallel$ | parallel velocity per unit field: the parallel flow is $v_\parallel\mathbf{B}$ | `with_vpar = .true.` |
+| `T` | $T$ | total temperature $T_i + T_e$ (single-temperature model, where $T_i = T_e = T/2$) | `with_TiTe = .false.` |
+| `Ti` | $T_i$ | ion temperature (two-temperature model) | `with_TiTe = .true.` |
+| `Te` | $T_e$ | electron temperature (two-temperature model) | `with_TiTe = .true.` |
+
+Some quantities take a different temperature in the two models
+(`with_TiTe` selects the two-temperature model):
+
+| Python | Symbol | Two-temperature model | Single-temperature model |
+|---|---|---|---|
+| `T_or_Te` | $\check{T}_e$ | $T_e$ | $T$ |
+| `Te_gen` | $\hat{T}_e$ | $T_e$ | $T/2$ |
+| `ion_temperature` | $\tilde{T}_i$ | $T_i$ | $T/2$ |
+
+In the single-temperature model this distinction is deliberate: the rate and
+charge-state functions ($S_{ion}$, $S_{rec}$, $\alpha_e$, $\eta$, ...) take
+$\check{T}_e = T$, while the electron pressure of the induction equation and
+the ion pressure of the diamagnetic terms take $T/2$. Do not "correct" one to
+match the other.
+
+Derived quantities that appear in several equations:
+
+| Symbol | Definition | Meaning |
+|---|---|---|
+| $\rho_{main}$ | $\rho - \rho_{imp}$ | main-ion mass density |
+| $n_i$ | $\rho + \alpha_i\,\rho_{imp}$ | ion density, $\alpha_i = m_i/m_{imp} - 1$ |
+| $n_e$ | $\rho + \alpha_e(\check{T}_e)\,\rho_{imp}$ | electron density, $\alpha_e = (m_i/m_{imp})\,Z_{imp} - 1$ |
+| $n_e^{\mathrm{c}}$ | $\rho^{\mathrm{c}} + \alpha_e(\check{T}_e)\,\rho_{imp}^{\mathrm{c}}$ | electron density from the corrected densities (not $\mathrm{corr}(n_e)$) |
+| $x^{\mathrm{c}}$ | $\mathrm{corr}(x)$ | density $x$ corrected for negative values (`corr_neg_dens`); $x^{\mathrm{c}} = x$ where $x$ is well above zero |
+| $p_e(\Theta)$ | $\rho\,\Theta + \rho_{imp}\,A_e(\Theta)$ | electron pressure, $A_e(\Theta) = \alpha_e(\Theta)\,\Theta$; $\Theta = T_e$ in the electron energy equation and $\Theta = \hat{T}_e$ in the induction equation |
+| $p_i$ | $n_i\,\tilde{T}_i$ | ion pressure; written $p_i^{dia}$ in the diamagnetic terms |
+| $p$ | two-temperature: $\rho\,(T_i + T_e) + \rho_{imp}\,\big(\alpha_i\,T_i + A_e(T_e)\big)$; single-temperature: $\rho\,T + \rho_{imp}\,A_{imp}(T)$ | total pressure, $A_{imp}(T) = \alpha_{imp}(T)\,T$ with $\alpha_{imp} = \frac{1}{2}(m_i/m_{imp})(Z_{imp}+1) - 1$ |
+| $W_{ion}$ | $E_{ion}(\check{T}_e)\,\rho_{imp} + E_{ion}^{bg}\,(\rho - \rho_{imp})$ | ionization potential energy of the impurities and of the main ions |
+| $B^2$ | $(F_0^2 + \lvert\nabla_{\mathrm{pol}}\psi\rvert^2)/R^2$ | square of the magnetic field |
+
+## Equations
+
+
+
+
+
+### Induction equation (`var_psi`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(\frac{1}{R^{2}}\,v\,\psi\Big) \;=\; & \frac{1}{R}\,v\,[\psi,u]^{st} \\
+&- \frac{1}{R^{2}}\,v\,F_0\,\partial_{\phi} u \\
+&+ \frac{1}{R^{2}}\,v\,\eta(\check{T}_e, \rho, \rho_{imp})\,\left(j - j_{src} - j_b\right) \\
+&+ \frac{1}{R}\,\eta_{num}(\check{T}_e)\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} j\right) \\
+&- \frac{1}{R^{2}}\,v\,\eta(\check{T}_e, \rho, \rho_{imp})\,j_{RE}^{ind} \\
+&- \frac{2}{\rho^{\mathrm{c}}\,B^2(\psi)\,R^{3}}\,v\,\tau_{IC}\,F_0^{2}\,[\psi,p_e]^{st} \\
+&+ \frac{2}{\rho^{\mathrm{c}}\,B^2(\psi)\,R^{4}}\,v\,\tau_{IC}\,F_0^{3}\,\partial_{\phi} p_e
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-psi)
+
+
+
+### Perpendicular momentum equation (`var_u`)
+
+$$
+\begin{aligned}
+R^{2}\,\left(-\rho^{\mathrm{c}}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \partial_t u\right) - f_{cons}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right)\,\partial_t \rho\right) \;=\; &- \frac{1}{2}\,R\,\left(\left(\partial_{R} u\right)^{2} + \left(\partial_{Z} u\right)^{2}\right)\,[v,\hat{\rho}] \\
+&- R^{3}\,\rho\,\omega\,[v,u]^{st} \\
+&+ \frac{1}{R}\,v\,[\psi,j]^{st} \\
+&- \frac{1}{R^{2}}\,v\,F_0\,\partial_{\phi} j \\
+&+ R\,[v,p]^{st} \\
+&- \mu_\perp(\check{T}_e)\,R^{2}\,f_\mu^{old}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \omega\right) \\
+&- 2\,\mu_\perp(\check{T}_e)\,R\,f_\mu^{new}\,\omega\,\partial_{R} v \\
+&- \mu_\perp(\check{T}_e)\,f_\mu^{new}\,\left(\partial_{R} v\,\partial_{R}\partial_{\phi}\partial_{\phi} u + \partial_{Z} v\,\partial_{Z}\partial_{\phi}\partial_{\phi} u\right) \\
+&- \frac{1}{R}\,\mu_{num}(\check{T}_e)\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \omega \\
+&- 2\,v\,\tau_{IC}\,R^{3}\,[p_i,\omega]^{st} \\
+&- 2\,\tau_{IC}\,R^{2}\,\partial_{Z} p_i\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right) \\
+&- 2\,v\,\tau_{IC}\,R^{3}\,\left(\partial_{R}\partial_{Z} u\,\left(\partial_{R}\partial_{R} p_i - \partial_{Z}\partial_{Z} p_i\right) - \partial_{R}\partial_{Z} p_i\,\left(\partial_{R}\partial_{R} u - \partial_{Z}\partial_{Z} u\right)\right) \\
+&+ \mu_\perp'(\check{T}_e)\,W_{dia}\,\left(\nabla_{\mathrm{pol}} \begin{cases} T_i & \text{if } \texttt{with TiTe} \\ \frac{T}{2} & \text{otherwise} \end{cases}\cdot\nabla_{\mathrm{pol}} v\right) \\
+&+ \mu_\perp(\check{T}_e)\,W_{dia}\,\nabla^2_{\mathrm{pol}} v \\
+&+ f_{cons}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right)\,\left(-R\,[\hat{\rho},u] + F_0\,\left(\rho\,\partial_{\phi} v_\parallel + v_\parallel\,\partial_{\phi} \rho\right) + R\,\rho\,[v_\parallel,\psi] + R\,v_\parallel\,[\rho,\psi]\right) \\
+&- v\,\left(P_\parallel^{RE} + P_\perp^{RE}\right) \\
+&+ R\,\left(-\Pi_R^{RE}\,\partial_{Z} v + \Pi_Z^{RE}\,\partial_{R} v\right) \\
+&+ R^{2}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right)\,\left(\left(1 - \delta_{n}\right)\,S_{neut} + \left(1 - f_{cons}\right)\,\left(S_p + S_{pellet} + S_{bg,drift} + S_{imp,drift}\right)\right) \\
+&- \frac{1}{4}\,c_{TG}^{u}\,R^{4}\,\rho\,[\omega,u]\,[v,u]\,\Delta t \\
+&- \frac{1}{4}\,c_{TG}^{u}\,R^{2}\,\omega\,f_{cons}\,[\hat{\rho},u]\,[v,u]\,\Delta t
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-u)
+
+
+
+### Current definition (`var_zj`)
+
+$$
+\begin{aligned}
+0 \;=\; &\frac{1}{R^{2}}\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \psi + v\,j\right)
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-zj)
+
+
+
+### Vorticity definition (`var_w`)
+
+$$
+\begin{aligned}
+0 \;=\; &\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u + v\,\omega
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-w)
+
+
+
+### Density equation (`var_rho`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(v\,\rho\Big) \;=\; & v\,\left(S_p + S_{pellet} + S_{bg,drift} + S_{imp,drift} + S_\rho^{aux}\right) \\
+&+ v\,n_e\,\rho_n\,S_{ion}(\check{T}_e) \\
+&- v\,n_e\,\rho_{main}\,S_{rec}(\check{T}_e) \\
+&+ v\,\mathcal{C}_u(\rho) \\
+&+ v\,\mathcal{C}_\parallel(\rho) \\
+&+ \left(D_\parallel^{\mathrm{tot}} - D_\perp\right)\,\mathcal{D}_\parallel(v, \rho_{main}) \\
+&- D_\perp\,\mathcal{D}_{tot}(v, \rho_{main}) \\
+&+ \left(D_{\parallel,imp}^{\mathrm{tot}} - D_{\perp,imp}\right)\,\mathcal{D}_\parallel(v, \rho_{imp}) \\
+&- D_{\perp,imp}\,\mathcal{D}_{tot}(v, \rho_{imp}) \\
+&+ c_{TG}^{\rho}\,\mathcal{T}_n(v, \rho) \\
+&- D_{\perp,num}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \rho \\
+&+ 4\,v\,\tau_{IC}\,\partial_{Z} p_i^{dia} \\
+&- \frac{1}{\sqrt{\left(\partial_{R} \psi\right)^{2} + \left(\partial_{Z} \psi\right)^{2}}}\,V_{pinch}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \psi\right)\,\rho
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-rho)
+
+
+
+### Parallel velocity equation (`var_vpar`)
+
+$$
+\begin{aligned}
+v\,\left(\rho^{\mathrm{c}}\,B^2(\psi)\,\partial_t v_\parallel + \frac{1}{2}\,\rho^{\mathrm{c}}\,v_\parallel\,\partial_t B^2_{pol}(\psi) + f_{cons}\,v_\parallel\,B^2(\psi)\,\partial_t \rho\right) \;=\; &- v\,\left(\mathbf{B}\cdot\nabla p\right) \\
+&+ \frac{1}{2}\,v_\parallel^{2}\,B^2\,\left(\rho\,\left(\mathbf{B}\cdot\nabla v\right) + v\,\left(\mathbf{B}\cdot\nabla \rho\right)\right) \\
+&+ f_{cons}\,v\,v_\parallel\,B^2\,\left(\frac{[\hat{\rho},u]}{R} - \mathbf{B}\cdot\nabla \left(\rho\,v_\parallel\right)\right) \\
+&- \mu_{\parallel,num}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} v_\parallel \\
+&- \frac{1}{R^{2}\,B^2}\,\mu_{\parallel\parallel}\,F_0^{2}\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)\,\left(\mathbf{B}\cdot\nabla v\right) \\
+&- \mu_\parallel^{\mathrm{eff}}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}}(v_\parallel - V_{rot})\right) \\
+&+ v\,M_\parallel^{aux} \\
+&- v\,S_n\,v_\parallel\,B^2\,\left(1 - f_{cons}\right) \\
+&- v\,S_\rho^{aux}\,v_\parallel\,B^2\,\left(1 - f_{cons}\right) \\
+&+ \left(1 - \delta_{n}\right)\,v\,n_e\,v_\parallel\,B^2\,\left(\left(\rho - \rho_{imp}\right)\,S_{rec}(\check{T}_e) - \rho_n\,S_{ion}(\check{T}_e)\right) \\
+&- \frac{1}{4}\,c_{TG}^{v_\parallel}\,\rho\,v_\parallel^{2}\,B^2\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)\,\left(\mathbf{B}\cdot\nabla v\right)\,\Delta t \\
+&- \frac{1}{4}\,c_{TG}^{v_\parallel}\,v\,v_\parallel^{2}\,B^2\,\left(1 - f_{cons}\right)\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)\,\left(\mathbf{B}\cdot\nabla \rho\right)\,\Delta t \\
+&- \frac{1}{4}\,c_{TG}^{v_\parallel}\,v_\parallel^{3}\,B^2\,f_{cons}\,\left(\mathbf{B}\cdot\nabla \rho\right)\,\left(\mathbf{B}\cdot\nabla v\right)\,\Delta t \\
+&+ \frac{1}{\sqrt{\left(\partial_{R} \psi\right)^{2} + \left(\partial_{Z} \psi\right)^{2}}}\,V_{pinch}\,\left(\nabla_{\mathrm{pol}} \psi\cdot\nabla_{\mathrm{pol}} v_\parallel\right)\,\rho\,v
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-vpar)
+
+
+
+### Impurity density equation (`var_rhoimp`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(v\,\rho_{imp}\Big) \;=\; & v\,S_{imp,drift} \\
+&+ v\,\mathcal{C}_u(\rho_{imp}) \\
+&+ v\,\mathcal{C}_\parallel(\rho_{imp}) \\
+&+ \left(D_{\parallel,imp}^{\mathrm{tot}} - D_{\perp,imp}\right)\,\mathcal{D}_\parallel(v, \rho_{imp}) \\
+&- D_{\perp,imp}\,\mathcal{D}_{tot}(v, \rho_{imp}) \\
+&+ c_{TG}^{\rho_{imp}}\,\mathcal{T}_n(v, \rho_{imp}) \\
+&- D_{\perp,num}^{n}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \rho_{imp}
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-rhoimp)
+
+
+
+### Neutral density equation (`var_rhon`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(v\,\rho_n\Big) \;=\; &- D_{n,R}\,\partial_{R} v\,\partial_{R} \rho_n \\
+&- D_{n,Z}\,\partial_{Z} v\,\partial_{Z} \rho_n \\
+&- \frac{1}{R^{2}}\,D_{n,\phi}\,\partial_{\phi} v\,\partial_{\phi} \rho_n \\
+&+ \delta_{n}\,v\,\left(\mathcal{C}_u(\rho_n) + \mathcal{C}_\parallel(\rho_n)\right) \\
+&- v\,n_e\,\rho_n\,S_{ion}(\check{T}_e) \\
+&+ v\,n_e\,\rho_{main}\,S_{rec}(\check{T}_e) \\
+&+ v\,S_n^{drift} \\
+&- D_{\perp,num}^{n}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \rho_n
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-rhon)
+
+
+
+### Ion energy equation (`var_Ti`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(v\,\left(\rho^{\mathrm{c}} + \alpha_i\,\rho_{imp}^{\mathrm{c}}\right)\,T_i\Big) \;=\; & v\,\mathcal{C}_u^{p}(p_i) \\
+&+ v\,\mathcal{C}_\parallel^{p}(p_i) \\
+&+ \left(\kappa_{\parallel,i}(T_i) - \kappa_{\perp,i}(\rho)\right)\,\mathcal{D}_\parallel(v, T_i) \\
+&- \kappa_{\perp,i}(\rho)\,\mathcal{D}_{tot}(v, T_i) \\
+&+ v\,\left(H_i + H_i^{aux}\right) \\
+&+ v\,Q_{ie}(T_i, T_e, \rho, \rho_{imp}) \\
+&- v\,T_i\,\left(\rho^{\mathrm{c}}\right)^{2}\,S_{rec}(T_e) \\
+&+ \frac{1}{2}\,v\,\left(\Gamma - 1\right)\,\left(v_\parallel^{2}\,B^2(\psi) + R^{2}\,\left(\left(\partial_{R} u\right)^{2} + \left(\partial_{Z} u\right)^{2}\right)\right)\,S_{kin}(T_e) \\
+&+ Q_{\mu\parallel}(v) \\
+&+ Q_{\mu\perp}(v, \mu_{heat}(T_e)) \\
+&+ Q_{kin}(v) \\
+&+ \mathcal{T}_p(v, p_i; c_{TG}^{T_i}) \\
+&- \kappa_{\perp,num}^{i}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} T_i \\
+&+ \mathcal{H}(v, \epsilon_i^{\mathrm{floor}}(T_i), T_i^{\mathrm{floor}}(T_i), T_{\mathrm{min}})
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-Ti)
+
+
+
+### Electron energy equation (`var_Te`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(v\,\left(\rho^{\mathrm{c}}\,T_e + \rho_{imp}^{\mathrm{c}}\,A_e(T_e) + \left(\Gamma - 1\right)\,W_{ion}\right)\Big) \;=\; & v\,\mathcal{C}_u^{p}(p_e) \\
+&+ v\,\mathcal{C}_\parallel^{p}(p_e) \\
+&+ \left(\kappa_{\parallel,e}(T_e) - \kappa_{\perp,e}(\rho)\right)\,\mathcal{D}_\parallel(v, T_e) \\
+&- \kappa_{\perp,e}(\rho)\,\mathcal{D}_{tot}(v, T_e) \\
+&+ v\,\left(H_e + H_e^{aux} + P_{teleport}\right) \\
+&- v\,\xi_{ion}\,\left(\rho + \alpha_e(T_e)\,\rho_{imp}\right)\,\rho_n\,S_{ion}(T_e) \\
+&+ v\,Q_{ei}(T_i, T_e, \rho, \rho_{imp}) \\
+&+ v\,\left(\Gamma - 1\right)\,\eta_{ohm}(T_e, \rho, \rho_{imp})\,\left(\frac{1}{R}\left(j - j_{RE}\right)\right)^{2} \\
+&- v\,n_e^{\mathrm{c}}\,\rho_n^{\mathrm{c}}\,L_{rays}(T_e) \\
+&- v\,n_e^{\mathrm{c}}\,\left(\rho^{\mathrm{c}} - \rho_{imp}^{\mathrm{c}}\right)\,L_{cont}(T_e) \\
+&- v\,n_e^{\mathrm{c}}\,f_{rad}^{bg}(T_e) \\
+&- v\,n_e^{\mathrm{c}}\,\rho_{imp}^{\mathrm{c}}\,L_{rad}(T_e) \\
+&+ \mathcal{Q}_{ion}(v, W_{ion}, T_e) \\
+&+ \mathcal{T}_p(v, p_e; c_{TG}^{T_e}) \\
+&- \kappa_{\perp,num}^{e}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} T_e \\
+&+ \mathcal{H}(v, \epsilon_e^{\mathrm{floor}}(T_e), T_e^{\mathrm{floor}}(T_e), T_{\mathrm{min}})
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-Te)
+
+
+
+### Single-temperature energy equation (`var_T`)
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial t}\Big(v\,\left(\rho^{\mathrm{c}}\,T + \rho_{imp}^{\mathrm{c}}\,A_{imp}(T) + \left(\Gamma - 1\right)\,W_{ion}\right)\Big) \;=\; & v\,\mathcal{C}_u^{p}(p) \\
+&+ v\,\mathcal{C}_\parallel^{p}(p) \\
+&+ \left(\kappa_\parallel(T) - \kappa_\perp(\rho)\right)\,\mathcal{D}_\parallel(v, T) \\
+&- \kappa_\perp(\rho)\,\mathcal{D}_{tot}(v, T) \\
+&+ v\,\left(H + H^{aux} + P_{teleport}\right) \\
+&- v\,\xi_{ion}\,\left(\rho + \alpha_e(T)\,\rho_{imp}\right)\,\rho_n\,S_{ion}(T) \\
+&- \frac{1}{2}\,\left(\Gamma - 1\right)\,v\,T\,\left(\rho^{\mathrm{c}}\right)^{2}\,S_{rec}(T) \\
+&+ v\,\left(\Gamma - 1\right)\,\eta_{ohm}(T, \rho, \rho_{imp})\,\left(\frac{1}{R}\left(j - j_{RE}\right)\right)^{2} \\
+&- v\,n_e^{\mathrm{c}}\,\rho_n^{\mathrm{c}}\,L_{rays}(T) \\
+&- v\,n_e^{\mathrm{c}}\,\left(\rho^{\mathrm{c}} - \rho_{imp}^{\mathrm{c}}\right)\,L_{cont}(T) \\
+&- v\,n_e^{\mathrm{c}}\,f_{rad}^{bg}(T) \\
+&- v\,n_e^{\mathrm{c}}\,\rho_{imp}^{\mathrm{c}}\,L_{rad}(T) \\
+&+ \mathcal{Q}_{ion}(v, W_{ion}, T) \\
+&+ \frac{1}{2}\,v\,\left(\Gamma - 1\right)\,\left(v_\parallel^{2}\,B^2(\psi) + R^{2}\,\left(\left(\partial_{R} u\right)^{2} + \left(\partial_{Z} u\right)^{2}\right)\right)\,S_{kin}(T) \\
+&+ Q_{\mu\parallel}(v) \\
+&+ Q_{\mu\perp}(v, \mu_{heat}(T)) \\
+&+ Q_{kin}(v) \\
+&+ \mathcal{T}_p(v, p; c_{TG}^{T}) \\
+&- \kappa_{\perp,num}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} T \\
+&+ \mathcal{H}(v, \epsilon^{\mathrm{floor}}(T), T^{\mathrm{floor}}(T), T_{\mathrm{min}})
+\end{aligned}
+$$
+
+[Definitions and term groups](#details-T)
+
+
+
+
+## Status and options
+
+### Status
+
+| Item | Equations | Status |
+|---|---|---|
+| Differences with the Fortran | several | **Recorded**, see [Known differences](#known-differences). |
+| Neoclassical terms | $u$, $v_\parallel$ | **Not implemented.** The Fortran neoclassical terms (profiles `amu_neo_prof`, `aki_neo_prof`) are not transcribed, and the checker does not audit the NEO branch. The `include_neo` argument of `momentum_equation_2` has no effect. |
+| Inward pinch | $\rho$, $v_\parallel$ | **Under review.** Whether the $v_\parallel$ pinch term belongs there is open, and the Fortran Jacobian of both pinch terms disagrees with the checker (findings 4 and 9). |
+| Momentum of the neutral sources | $v_\parallel$ | **Suspected defect, to confirm.** Probably needs the factor $(1-f_{cons})$ that the other particle-source momentum terms carry. |
+| Freezing $u$ in the time derivative | $u$ | **To confirm.** The element routine holds $u$ at its current value in the time derivative of the conservative-form term. |
+| Impurity pressure in the $\tau_{IC}$ tangents | $\psi$, $u$, $\rho$ | **Accepted.** The residual is complete; the Fortran Jacobian omits the impurity part (findings 1 and 2). Impurities and $\tau_{IC} \neq 0$ are not used together. |
+
+#### Known differences
+
+The findings still recorded in the checker's reference
+(`util/equation_codegen/reference/model600_discrepancies.json`), and the
+blocks of the element routine they affect. A finding with residual blocks
+changes the converged solution; one with only Jacobian blocks affects only
+the Newton convergence. The table is generated from the reference.
+
+
+
+| Finding | Residual blocks | Jacobian blocks |
+|---:|---|---|
+| [1](https://github.com/iterorganization/JOREK/blob/develop/util/equation_codegen/JOREK_FINDINGS.md) | none | `amat(var_rho,var_rhoimp)`, `amat(var_rho,var_t)`, `amat(var_rho,var_ti)`, `amat(var_u,var_rhoimp)`, `amat(var_u,var_t)`, `amat(var_u,var_ti)` |
+| [2](https://github.com/iterorganization/JOREK/blob/develop/util/equation_codegen/JOREK_FINDINGS.md) | none | `amat(var_psi,var_rhoimp)`, `amat(var_psi,var_t)`, `amat(var_psi,var_te)`, `amat_n(var_psi,var_rhoimp)`, `amat_n(var_psi,var_t)`, `amat_n(var_psi,var_te)` |
+| [4](https://github.com/iterorganization/JOREK/blob/develop/util/equation_codegen/JOREK_FINDINGS.md) | none | `amat(var_rho,var_psi)`, `amat(var_vpar,var_psi)` |
+| [9](https://github.com/iterorganization/JOREK/blob/develop/util/equation_codegen/JOREK_FINDINGS.md) | none | `amat(var_vpar,var_rho)`, `amat(var_vpar,var_vpar)` |
+
+
+
+### Options
+
+The compile-time switches of `mod_model_settings.f90` (see
+[Variables](#variables)) are genuine branches: they remove a variable and its
+equation, and in the element routine they gate whole blocks of terms. All other
+features are switched by the coefficients below, as they are named in the
+element routine (the input parameters that set them are documented with the
+model); a disabled feature has a zero coefficient, and its terms are still
+assembled.
+
+| Feature | Controlled by | Mechanism |
+|---|---|---|
+| Two-temperature model | `with_TiTe` | branch: equations for $T_i$, $T_e$ instead of $T$ |
+| Density, parallel velocity, neutrals, impurities | `with_rho`, `with_vpar`, `with_neutrals`, `with_impurities` | branch: the variable is replaced by a constant and its equation is not assembled |
+| Diamagnetic terms | $\tau_{IC}$ (`tauIC`) | coefficient |
+| Conservative form of the momentum equations | $f_{cons}$ (`fact_conservative_u`) | coefficient, $0$ or $1$ |
+| Convection of the neutrals with the plasma flow, and the momentum they carry | $\delta_n$ (`delta_n_convection`) | the neutral equation convects $\rho_n$ with $\delta_n$ (on for $\delta_n = 1$); the momentum of the neutral sources carries $(1-\delta_n)$ (on for $\delta_n = 0$) |
+| Inward pinch | $V_{pinch}$ (`V_prof_pinch`) | profile, zero when off |
+| Taylor-Galerkin stabilization | $c_{TG}$ (`tgnum_*`) | coefficient per equation |
+| Numerical hyper-diffusion | $D_{\perp,num}$, $\kappa_{\perp,num}$, $\eta_{num}$, $\mu_{num}$, $\mu_{\parallel,num}$ | coefficients |
+| Shock capturing | $D_{\parallel,sc}\,\tau_{sc}$, $D_{\parallel,imp,sc}\,\tau_{sc}$, $\mu_{\parallel,sc}\,\tau_{sc}$ | products of a coefficient and the shock indicator $\tau_{sc}$ |
+| Runaway-electron coupling | $j_{RE}$, $P^{RE}$, $\Pi^{RE}$ (`aux_jre`, `aux_P_*_re`, `aux_divPI*_perp`) | coupling values, zero without runaways |
+| Kinetic neutral and impurity coupling | $S_\rho^{aux}$, $M_\parallel^{aux}$, $H^{aux}$ (`aux_rho0`, `aux_mom_par0`, `aux_E0*`) | coupling values, zero without the kinetic model |
+| Neoclassical terms | `amu_neo_prof`, `aki_neo_prof` in the Fortran | not transcribed (see Status) |
+
+## Conventions
+
+The weak form assumes:
+
+- $v$ is a finite-element basis function in space (of the element
+ coordinates and of $\phi$); it does not depend on time.
+- The fields are periodic in $\phi$.
+- The grid, hence $R$ and $\mathcal{J}$, does not change in time.
+- Integration by parts moves derivatives onto $v$; the boundary terms it
+ produces are not part of these volume terms. Boundary conditions are imposed
+ separately (`models/model600/mod_boundary_conditions.f90`).
+
+**Time derivatives.** The equations at the top of the page write every time
+derivative out. In the momentum and parallel velocity equations it is not the
+derivative of a single quantity: the parallel velocity equation, for example,
+is the projection of the momentum equation on $\mathbf{B}$ and has
+$\rho^{\mathrm{c}}\big(B^2\,\partial_t v_\parallel + \frac{1}{2}v_\parallel\,\partial_t B^2_{pol}\big)$.
+In `model600.py` such a term is written as a mass functional $A$ with
+frozen factors (`freeze(x)`, shown as $\overline{x}$), and $\partial_t$ acts only on the factors without an overline:
+$\partial_t(\overline{a}\,b)$ means $a\,\partial_t b$. The symbolic model never
+uses the value of $A$ itself, only this derivative, and the checker derives the
+Jacobian from it with the same rule. Where the density correction is inactive
+($\rho^{\mathrm{c}} = \rho$), the momentum term
+$-\rho\,\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}}\partial_t u - f_{cons}\,\partial_t\rho\,\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u$
+is the time derivative of $\rho\,\nabla_{\mathrm{pol}} u$ for $f_{cons} = 1$
+(conservative form), and $\rho\,\partial_t\nabla_{\mathrm{pol}} u$ for
+$f_{cons} = 0$. Overlines appear only in these two mass functionals.
+
+How the other terms are linearized for the Newton iteration follows from the
+equations, apart from the state-dependent functions: the table
+[Linearization of the state-dependent functions](#linearization-of-the-state-dependent-functions)
+says which derivatives the Jacobian uses for each of them.
+
+In `model600.py` an evolution equation is written as a pair $A$, $B$ meaning
+$\frac{\partial}{\partial t}\int A = \int B$, and a constraint as $C$ meaning
+$\int C = 0$. There $A$, $B$ and $C$ include the volume weight
+$w_V = R\,\mathcal{J}$ (the Python name is `dV`), where $\mathcal{J}$ is the
+Jacobian of the element coordinates $(s,t_{\mathrm{el}})$; the equations at the top of the
+page show them divided by $w_V$. The time discretization and the Jacobian (the
+`amat` blocks) are not written here: the checker derives them by
+differentiating $A$ and $B$ with respect to each variable.
+
+- $[a,b] = \partial_R a\,\partial_Z b - \partial_Z a\,\partial_R b$ is the
+ poloidal bracket.
+- $[a,b]^{st} = \big(\partial_s a\,\partial_{t_{\mathrm{el}}} b - \partial_{t_{\mathrm{el}}} a\,\partial_s b\big)/\mathcal{J}$
+ is the same bracket computed in the element coordinates:
+ $[a,b]^{st} = [a,b]^{RZ} = [a,b]$. The element routine uses both spellings
+ (in `model600.py`, `poiss_bracket_st(a, b) / xjac` and `poiss_bracket(a, b)`);
+ they are equal, but the checker keeps them apart because they expand into
+ different monomials. Where the element-coordinate bracket appears without
+ its $1/\mathcal{J}$, it is written $\mathcal{J}\,[a,b]^{st}$.
+- $\nabla_{\mathrm{pol}} = (\partial_R, \partial_Z)$ is the **poloidal** gradient: it has
+ no toroidal component, and
+ $\nabla_{\mathrm{pol}} a\cdot\nabla_{\mathrm{pol}} b = \partial_R a\,\partial_R b + \partial_Z a\,\partial_Z b$.
+ Toroidal derivatives are
+ always written explicitly as $\partial_\phi$. Likewise
+ $\nabla^2_{\mathrm{pol}} f = \partial_R\partial_R f + \partial_Z\partial_Z f + \partial_R f / R$
+ is the poloidal part of the Laplacian.
+- $\mathbf{B}\cdot\nabla$ is the full three-dimensional derivative along the field,
+ including its toroidal part (see [Parallel gradient](#parallel-gradient)).
+- $\overline{x}$ in a mass functional $A$ is a factor that $\partial_t$ does
+ not act on (see above).
+- The two diffusion operators carry opposite signs, as in `model600.py`:
+ $\mathcal{D}_\parallel(v,n) = -(\mathbf{B}\cdot\nabla v)(\mathbf{B}\cdot\nabla n)/B^2$
+ is the integrated-by-parts form of $+\nabla\cdot(\mathbf{b}\,\mathbf{b}\cdot\nabla n)$,
+ while $\mathcal{D}_{tot}(v,f) = \nabla v\cdot\nabla f + \ldots$ is that of $-\nabla^2 f$.
+ A parallel diffusion therefore appears as $+(D_\parallel - D_\perp)\,\mathcal{D}_\parallel(v, n)$
+ and a perpendicular one as $-D_\perp\,\mathcal{D}_{tot}(v, n)$.
+- $\Gamma$ is the ratio of specific heats. The element routine spells it both
+ `GAMMA` and `gamma`; Fortran does not distinguish them.
+- `with_TiTe` selects the two-temperature model (fields $T_i$, $T_e$) instead
+ of the single-temperature one (field $T$); it is both a switch of the
+ element routine and an argument of the equation functions of `model600.py`. `st_form` only selects which
+ spelling of a bracket the element routine uses; it does not change the
+ physics. The checker needs it because the two spellings expand into
+ different monomials.
+
+## Notation
+
+Names that are not listed are rendered from their Python spelling:
+`D_par_local` becomes $D_{\mathrm{par},\mathrm{local}}$.
+
+
+
+### Work values
+
+| Python | Symbol |
+|---|---|
+| `T_or_Te` | $\check{T}_e$ |
+| `Te_gen` | $\hat{T}_e$ |
+| `ion_temperature` | $\tilde{T}_i$ |
+| `R` | $R$ |
+| `xjac` | $\mathcal{J}$ |
+| `dV` | $w_V$ |
+| `F0` | $F_0$ |
+| `tauIC` | $\tau_{IC}$ |
+| `tstep` | $\Delta t$ |
+| `GAMMA` | $\Gamma$ |
+| `gamma` | $\Gamma$ |
+| `ne` | $n_e$ |
+| `rho_main` | $\rho_{main}$ |
+| `rho_corr` | $\rho^{\mathrm{c}}$ |
+| `rho_hat` | $\hat{\rho}$ |
+| `bb2` | $B^2$ |
+| `Pe` | $p_e$ |
+| `pressure` | $p$ |
+| `pi` | $p_i$ |
+| `W_dia` | $W_{dia}$ |
+| `ion_density` | $n_i$ |
+| `ion_pressure` | $p_i$ |
+| `electron_pressure` | $p_e$ |
+| `electron_density` | $n_e^{\mathrm{c}}$ |
+| `ionization_energy` | $W_{ion}$ |
+| `visco_par_eff` | $\mu_\parallel^{\mathrm{eff}}$ |
+| `source_dens_tot` | $S_n$ |
+| `D_par_tot` | $D_\parallel^{\mathrm{tot}}$ |
+| `D_par_imp_tot` | $D_{\parallel,imp}^{\mathrm{tot}}$ |
+| `d_par_tot` | $\hat{D}_\parallel$ |
+| `d_par_imp_tot` | $\hat{D}_{\parallel,imp}$ |
+| `alpha_e_value` | $\alpha_e$ |
+| `sion_rate` | $S_{ion}$ |
+| `srec_rate` | $S_{rec}$ |
+| `grad_v` | $\nabla_{\mathrm{pol}} v$ |
+| `grad_u` | $\nabla_{\mathrm{pol}} u$ |
+| `grad_omega` | $\nabla_{\mathrm{pol}} \omega$ |
+| `grad_vpar` | $\nabla_{\mathrm{pol}} v_\parallel$ |
+| `lap_v` | $\nabla^2_{\mathrm{pol}} v$ |
+| `lap_omega` | $\nabla^2_{\mathrm{pol}} \omega$ |
+| `neutral_sources` | $S_{neut}$ |
+| `current_source` | $j_{src}$ |
+| `Jb` | $j_b$ |
+| `visco_par` | $\mu_\parallel$ |
+| `visco_par_num` | $\mu_{\parallel,num}$ |
+| `visco_par_par` | $\mu_{\parallel\parallel}$ |
+| `visco_par_sc_num` | $\mu_{\parallel,sc}$ |
+| `D_prof` | $D_\perp$ |
+| `D_prof_imp` | $D_{\perp,imp}$ |
+| `D_par_local` | $D_\parallel$ |
+| `D_par_local_imp` | $D_{\parallel,imp}$ |
+| `D_par_sc_num` | $D_{\parallel,sc}$ |
+| `D_par_imp_sc_num` | $D_{\parallel,imp,sc}$ |
+| `tau_sc` | $\tau_{sc}$ |
+| `V_prof_pinch` | $V_{pinch}$ |
+| `fact_conservative_u` | $f_{cons}$ |
+| `delta_n_convection` | $\delta_{n}$ |
+| `dV_dpsi_source` | $\frac{\mathrm{d}V_{rot}}{\mathrm{d}\psi}$ |
+| `particle_source` | $S_p$ |
+| `source_pellet` | $S_{pellet}$ |
+| `source_bg_drift` | $S_{bg,drift}$ |
+| `source_imp_drift` | $S_{imp,drift}$ |
+| `aux_rho0` | $S_\rho^{aux}$ |
+| `aux_mom_par0` | $M_\parallel^{aux}$ |
+| `heat_source_total` | $H$ |
+| `heat_source_i` | $H_i$ |
+| `heat_source_e` | $H_e$ |
+| `aux_E0` | $H^{aux}$ |
+| `aux_E0_Ti` | $H_i^{aux}$ |
+| `aux_E0_Te` | $H_e^{aux}$ |
+| `power_dens_teleport_ju` | $P_{teleport}$ |
+| `ksi_ion_norm` | $\xi_{ion}$ |
+| `implicit_heat_source` | $c_{\mathrm{floor}}$ |
+| `T_min_neg` | $T_{\mathrm{min}}$ |
+| `Tie_min_neg` | $T_{\mathrm{min}}$ |
+| `aux_jre` | $j_{RE}$ |
+| `aux_jre_ind` | $j_{RE}^{ind}$ |
+| `aux_P_par_re` | $P_\parallel^{RE}$ |
+| `aux_P_perp_re` | $P_\perp^{RE}$ |
+| `aux_divPIR_perp` | $\Pi_R^{RE}$ |
+| `aux_divPIZ_perp` | $\Pi_Z^{RE}$ |
+| `visco_fact_old` | $f_\mu^{old}$ |
+| `visco_fact_new` | $f_\mu^{new}$ |
+| `visco_par_heating` | $\mu_{\parallel,heat}$ |
+| `D_perp_num_psin` | $D_{\perp,num}$ |
+| `Dn0x` | $D_{n,R}$ |
+| `Dn0y` | $D_{n,Z}$ |
+| `Dn0p` | $D_{n,\phi}$ |
+| `source_neutral_drift` | $S_n^{drift}$ |
+| `Dn_perp_num` | $D_{\perp,num}^{n}$ |
+| `ZK_perp_num_psin` | $\kappa_{\perp,num}$ |
+| `ZK_i_perp_num_psin` | $\kappa_{\perp,num}^{i}$ |
+| `ZK_e_perp_num_psin` | $\kappa_{\perp,num}^{e}$ |
+| `tgnum_u` | $c_{TG}^{u}$ |
+| `tgnum_rho` | $c_{TG}^{\rho}$ |
+| `tgnum_vpar` | $c_{TG}^{v_\parallel}$ |
+| `tgnum_rhoimp` | $c_{TG}^{\rho_{imp}}$ |
+| `tgnum_T` | $c_{TG}^{T}$ |
+| `tgnum_Ti` | $c_{TG}^{T_i}$ |
+| `tgnum_Te` | $c_{TG}^{T_e}$ |
+| `psi_gradient` | $\nabla_{\mathrm{pol}} \psi$ |
+| `rotation_shear` | $\nabla_{\mathrm{pol}}(v_\parallel - V_{rot})$ |
+
+### State-dependent functions
+
+A function of the state is shown with its arguments, e.g. $S_{ion}(T)$;
+$x^{\mathrm{c}}$ is a density corrected for negative values.
+
+| Python | Symbol |
+|---|---|
+| `eta` | $\eta$ |
+| `eta_num_T` | $\eta_{num}$ |
+| `visco` | $\mu_\perp$ |
+| `visco_num` | $\mu_{num}$ |
+| `alpha_i_state` | $\alpha_i$ |
+| `dvisco_state` | $\mu_\perp'$ |
+| `Sion_rate` | $S_{ion}$ |
+| `Srec_rate` | $S_{rec}$ |
+| `alpha_e_state` | $\alpha_e$ |
+| `alpha_e_bis_state` | $\alpha_e'$ |
+| `alpha_e_temperature` | $A_e$ |
+| `alpha_imp_bis_state` | $\alpha_{imp}'$ |
+| `alpha_imp_temperature` | $A_{imp}$ |
+| `corr_neg_dens` | ${0}^{\mathrm{c}}$ |
+| `corr_neg_dens_imp` | ${0}^{\mathrm{c}}$ |
+| `corr_neg_dens_n` | ${0}^{\mathrm{c}}$ |
+| `W_dia_single` | $W_{dia}$ |
+| `W_dia_two` | $W_{dia}$ |
+| `ZKi_par` | $\kappa_{\parallel,i}$ |
+| `ZKi_perp` | $\kappa_{\perp,i}$ |
+| `visco_heating` | $\mu_{heat}$ |
+| `Ti_e_exchange` | $Q_{ie}$ |
+| `Ti_floor` | $T_i^{\mathrm{floor}}$ |
+| `Ti_floor_exp` | $\epsilon_i^{\mathrm{floor}}$ |
+| `E_ion_bg_state` | $E_{ion}^{bg}$ |
+| `E_ion_state` | $E_{ion}$ |
+| `ZKe_par` | $\kappa_{\parallel,e}$ |
+| `ZKe_perp` | $\kappa_{\perp,e}$ |
+| `eta_ohm_e` | $\eta_{ohm}$ |
+| `LradDrays` | $L_{rays}$ |
+| `LradDcont` | $L_{cont}$ |
+| `frad_bg_state` | $f_{rad}^{bg}$ |
+| `Lrad_state` | $L_{rad}$ |
+| `Te_i_exchange` | $Q_{ei}$ |
+| `Te_floor` | $T_e^{\mathrm{floor}}$ |
+| `Te_floor_exp` | $\epsilon_e^{\mathrm{floor}}$ |
+| `ZK_par` | $\kappa_\parallel$ |
+| `ZK_perp` | $\kappa_\perp$ |
+| `T_floor` | $T^{\mathrm{floor}}$ |
+| `T_floor_exp` | $\epsilon^{\mathrm{floor}}$ |
+
+
+
+### Linearization of the state-dependent functions
+
+Generated from the `external_function` declarations of `model600.py`. The
+Jacobian differentiates a function through the derivatives the element
+routine supplies for it; "current branch" means the derivative of the branch
+(or clipping regime) selected at the current state, so a clipped constant has
+zero derivative.
+
+
+
+| DSL | Arguments | Fortran value | Derivatives supplied | Jacobian |
+|---|---|---|---|---|
+| `eta` | `T`, `rho`, `rhoimp` | `eta_T` | `deta_dT` for `T`, `deta_dr0` for `rho`, `deta_drimp0` for `rhoimp` | differentiated (current branch) |
+| `eta_num_T` | `T` | `eta_num_T` | `deta_num_dT` for `T` | differentiated (current branch) |
+| `visco` | `T` | `visco_T` | `dvisco_dT` for `T` | differentiated (current branch) |
+| `visco_num` | `T` | `visco_num_T` | none | not differentiated |
+| `alpha_i_state` | none | `alpha_i` | none | not differentiated |
+| `dvisco_state` | `T` | `dvisco_dT` | `d2visco_dT2` for `T` | differentiated (current branch) |
+| `Sion_rate` | `T` | `Sion_T` | `dSion_dT` for `T` | differentiated (current branch) |
+| `Srec_rate` | `T` | `Srec_T` | `dSrec_dT` for `T` | differentiated (current branch) |
+| `alpha_e_state` | `T` | `alpha_e` | `dalpha_e_dT` for `T` | differentiated (current branch) |
+| `alpha_e_bis_state` | `T` | `alpha_e_bis` | `alpha_e_tri` for `T` | differentiated (current branch) |
+| `alpha_e_temperature` | `T` | `alpha_e_T` | `alpha_e_bis` for `T` | differentiated (current branch) |
+| `alpha_imp_bis_state` | `T` | `alpha_imp_bis` | `alpha_imp_tri` for `T` | differentiated (current branch) |
+| `alpha_imp_temperature` | `T` | `alpha_imp_T` | `alpha_imp_bis` for `T` | differentiated (current branch) |
+| `corr_neg_dens` | `rho` | `corr_neg_dens` | `dr0_corr_dn` for `rho` | differentiated (current branch) |
+| `W_dia_single` | `rho`, `T` | `W_dia` | `W_dia_rho` for `rho`, `W_dia_T` for `T` | differentiated (current branch) |
+| `W_dia_two` | `rho`, `Ti` | `W_dia` | `W_dia_rho` for `rho`, `W_dia_Ti` for `Ti` | differentiated (current branch) |
+| `ZKi_par` | `Ti` | `ZKi_par_T` | `dZKi_par_dT` for `Ti` | differentiated (current branch) |
+| `ZKi_perp` | `rho` | `ZKi_prof` | `dZKi_prof_drho` for `rho` | differentiated (current branch) |
+| `visco_heating` | `Te` | `visco_T_heating` | `dvisco_dT_heating` for `Te` | differentiated (current branch) |
+| `Ti_e_exchange` | `Ti`, `Te`, `rho`, `rhoimp` | `dTi_e` | `ddTi_e_dTi` for `Ti`, `ddTi_e_dTe` for `Te`, `ddTi_e_drho` for `rho`, `ddTi_e_drhoimp` for `rhoimp` | differentiated (current branch) |
+| `Ti_floor` | `Ti` | `Ti0_floor` | `dTi_floor` for `Ti` | differentiated (current branch) |
+| `Ti_floor_exp` | `Ti` | `Ti_floor_exp` | `dTi_floor_exp` for `Ti` | differentiated (current branch) |
+| `corr_neg_dens_imp` | `rhoimp` | `corr_neg_dens_imp` | `drimp0_corr_dn` for `rhoimp` | differentiated (current branch) |
+| `E_ion_bg_state` | none | `E_ion_bg` | none | not differentiated |
+| `ZKe_par` | `Te` | `ZKe_par_T` | `dZKe_par_dT` for `Te` | differentiated (current branch) |
+| `ZKe_perp` | `rho` | `ZKe_prof` | `dZKe_prof_drho` for `rho` | differentiated (current branch) |
+| `eta_ohm_e` | `Te`, `rho`, `rhoimp` | `eta_T_ohm` | `deta_dT_ohm` for `Te`, `deta_dr0_ohm` for `rho`, `deta_drimp0_ohm` for `rhoimp` | differentiated (current branch) |
+| `dE_ion_dT_state` | `Te` | `dE_ion_dT` | none | not differentiated |
+| `E_ion_state` | `Te` | `E_ion` | `dE_ion_dT` for `Te` | differentiated (current branch) |
+| `LradDrays` | `Te` | `LradDrays_T` | `dLradDrays_dT` for `Te` | differentiated (current branch) |
+| `LradDcont` | `Te` | `LradDcont_corr` | `dLradDcont_dT_corr` for `Te` | differentiated (current branch) |
+| `frad_bg_state` | `Te` | `frad_bg` | `dfrad_bg_dT` for `Te` | differentiated (current branch) |
+| `Lrad_state` | `Te` | `Lrad` | `dLrad_dT` for `Te` | differentiated (current branch) |
+| `Te_i_exchange` | `Ti`, `Te`, `rho`, `rhoimp` | `dTe_i` | `ddTe_i_dTi` for `Ti`, `ddTe_i_dTe` for `Te`, `ddTe_i_drho` for `rho`, `ddTe_i_drhoimp` for `rhoimp` | differentiated (current branch) |
+| `Te_floor` | `Te` | `Te0_floor` | `dTe_floor` for `Te` | differentiated (current branch) |
+| `Te_floor_exp` | `Te` | `Te_floor_exp` | `dTe_floor_exp` for `Te` | differentiated (current branch) |
+| `corr_neg_dens_n` | `rhon` | `corr_neg_dens_n` | `drn0_corr_dn` for `rhon` | differentiated (current branch) |
+| `ZK_par` | `T` | `ZK_par_T` | `dZK_par_dT` for `T` | differentiated (current branch) |
+| `ZK_perp` | `rho` | `ZK_prof` | `dZK_prof_drho` for `rho` | differentiated (current branch) |
+| `T_floor` | `T` | `T0_floor` | `dT_floor` for `T` | differentiated (current branch) |
+| `T_floor_exp` | `T` | `T_floor_exp` | `dT_floor_exp` for `T` | differentiated (current branch) |
+
+
+
+## Shared operators
+
+Term groups that appear in several equations, grouped by topic. The symbol on
+the left of each definition is how the group is written in the equations
+above; $f$, $n$ and $p$ stand for any scalar, density and pressure. Each is a
+`_helper` function of `model600.py`.
+
+
+
+### Magnetic field
+
+#### Magnetic field strength
+
+The square of the total field, with $\mathbf{B} = F_0\,\nabla\phi + \nabla\psi\times\nabla\phi$.
+
+$$
+\begin{aligned}
+B^2(\psi) &= \frac{1}{R^{2}}\left(F_0^{2} + \left(\partial_{R} \psi\right)^{2} + \left(\partial_{Z} \psi\right)^{2}\right)
+\end{aligned}
+$$
+
+Source: `_B2`.
+
+#### Poloidal magnetic field strength
+
+The square of the poloidal field, $\lvert\nabla_{\mathrm{pol}}\psi\times\nabla\phi\rvert^2$. Since $F_0$ is constant, $\partial_t B^2 = \partial_t B^2_{pol}$.
+
+$$
+\begin{aligned}
+B^2_{pol}(\psi) &= \frac{1}{R^{2}}\left(\left(\partial_{R} \psi\right)^{2} + \left(\partial_{Z} \psi\right)^{2}\right)
+\end{aligned}
+$$
+
+Source: `_B2_pol`.
+
+#### Parallel gradient
+
+The derivative along the field. With `st_form` the bracket is written in element coordinates, $[f,\psi]^{st}$, which is the same quantity.
+
+$$
+\begin{aligned}
+\mathbf{B}\cdot\nabla f &= \frac{1}{R}\left(\frac{1}{R}\,F_0\,\partial_{\phi} f + [f,\psi]\right)
+\end{aligned}
+$$
+
+Source: `_B_dot_grad`.
+
+#### Weighted density
+
+The density weighted by $R^2$, as it enters the momentum equation.
+
+$$
+\begin{aligned}
+\hat{\rho} &= R^{2}\,\rho
+\end{aligned}
+$$
+
+Source: `_rho_hat`.
+
+### Convection
+
+#### Convection by the ExB flow
+
+Advection and compression of a density by the $E\times B$ flow of stream function $u$.
+
+$$
+\begin{aligned}
+\mathcal{C}_u(n) &= R\,[n,u] + 2\,n\,\partial_{Z} u
+\end{aligned}
+$$
+
+Source: `_u_convection`.
+
+#### Parallel convection
+
+Advection and compression of a density by the parallel flow $v_\parallel$.
+
+$$
+\begin{aligned}
+\mathcal{C}_\parallel(n) &= -v_\parallel\,\left(\mathbf{B}\cdot\nabla n\right) - n\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)
+\end{aligned}
+$$
+
+Source: `_parallel_convection`.
+
+#### Pressure convection by the ExB flow
+
+The same for a pressure: the compression carries the factor $\Gamma$.
+
+$$
+\begin{aligned}
+\mathcal{C}_u^{p}(p) &= R\,[p,u] + 2\,\Gamma\,p\,\partial_{Z} u
+\end{aligned}
+$$
+
+Source: `_press_u_convection`.
+
+#### Parallel pressure convection
+
+The same for a pressure along the field.
+
+$$
+\begin{aligned}
+\mathcal{C}_\parallel^{p}(p) &= -v_\parallel\,\left(\mathbf{B}\cdot\nabla p\right) - \Gamma\,p\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)
+\end{aligned}
+$$
+
+Source: `_press_parallel_convection`.
+
+### Diffusion
+
+#### Parallel diffusion
+
+Diffusion along the magnetic field, after integration by parts.
+
+$$
+\begin{aligned}
+\mathcal{D}_\parallel(v, n) &= -\frac{1}{B^2(\psi)}\,\left(\mathbf{B}\cdot\nabla v\right)\,\left(\mathbf{B}\cdot\nabla n\right)
+\end{aligned}
+$$
+
+Source: `_par_diff_intg_by_parts`.
+
+#### Total diffusion
+
+Isotropic diffusion, including the toroidal derivative, after integration by parts.
+
+$$
+\begin{aligned}
+\mathcal{D}_{tot}(v, f) &= \nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} f + \frac{1}{R^{2}}\,\partial_{\phi} v\,\partial_{\phi} f
+\end{aligned}
+$$
+
+Source: `_diffusion_tot_intg_by_parts`.
+
+### Numerical stabilization
+
+#### Taylor-Galerkin stabilization of a density
+
+Taylor-Galerkin stabilization of the convection by the $E\times B$ and parallel flows.
+
+$$
+\begin{aligned}
+\mathcal{T}_n(v, n) &= -\frac{1}{4}\,R^{2}\,[n,u]\,[v,u]\,\Delta t - \frac{1}{4}\,v_\parallel^{2}\,\left(\mathbf{B}\cdot\nabla n\right)\,\left(\mathbf{B}\cdot\nabla v\right)\,\Delta t
+\end{aligned}
+$$
+
+Source: `_dens_tgnum_intg_by_parts`.
+
+#### Taylor-Galerkin stabilization of an energy equation
+
+The same for a pressure, with the coefficient $c$ of the equation.
+
+$$
+\begin{aligned}
+\mathcal{T}_p(v, p; c) &= -\frac{1}{4}\,c\,R^{2}\,[p,u]\,[v,u]\,\Delta t - \frac{1}{4}\,c\,v_\parallel^{2}\,\left(\mathbf{B}\cdot\nabla p\right)\,\left(\mathbf{B}\cdot\nabla v\right)\,\Delta t
+\end{aligned}
+$$
+
+Source: `_energy_tgnum`.
+
+#### Heating at the temperature floor
+
+Implicit heat source that keeps a temperature away from its floor $T_{\mathrm{min}}$. $\epsilon$ and $T^{\mathrm{floor}}$ are the floor functions of the element routine (`Ti_floor_exp`, `Ti0_floor` and their electron and single-temperature counterparts).
+
+$$
+\begin{aligned}
+\mathcal{H}(v, \epsilon, T^{\mathrm{floor}}, T_{\mathrm{min}}) &= c_{\mathrm{floor}}\,\left(\Gamma - 1\right)\,v\,\left(\frac{1}{2}\,T_{\mathrm{min}}\,\left(1 + \epsilon\right) - T^{\mathrm{floor}}\right)
+\end{aligned}
+$$
+
+Source: `_heating_floor`.
+
+### Diamagnetic terms
+
+#### Diamagnetic ion pressure
+
+The ion pressure of the diamagnetic terms, as built in `construct_pressure`. The single-temperature model evolves the total temperature, so the ion temperature is $T/2$ there. $\alpha_i$ does not depend on the temperature.
+
+$$
+\begin{aligned}
+\tilde{T}_i &= \begin{cases} T_i & \text{if } \texttt{with TiTe} \\ \frac{T}{2} & \text{otherwise} \end{cases} \\
+p_i^{dia} &= \left(\rho + \rho_{imp}\,\alpha_i\right)\,\tilde{T}_i
+\end{aligned}
+$$
+
+Source: `_diamagnetic_pressure`.
+
+#### Diamagnetic viscosity
+
+The element routine's `W_dia`. It is built from the ion pressure alone, so it has no explicit electron-temperature dependence.
+
+$$
+\begin{aligned}
+W_{dia} &= \begin{cases} W_{dia}(\rho, T_i) & \text{if } \texttt{with TiTe} \\ W_{dia}(\rho, T) & \text{otherwise} \end{cases}
+\end{aligned}
+$$
+
+Source: `_diamagnetic_viscosity`.
+
+### Sources and heating
+
+#### Particle sources releasing kinetic energy
+
+Particle sources whose kinetic energy is released into the ions: ionization of neutrals and the external sources.
+
+$$
+\begin{aligned}
+S_{kin}(T) &= \left(\rho + \alpha_e(T)\,\rho_{imp}\right)\,\rho_n\,S_{ion}(T) + S_p + S_{pellet} + S_{bg,drift} + S_{imp,drift}
+\end{aligned}
+$$
+
+Source: `_released_kinetic_energy`.
+
+#### Parallel viscous heating
+
+Heating by the parallel viscosity.
+
+$$
+\begin{aligned}
+Q_{\mu\parallel}(v) &= \left(\Gamma - 1\right)\,\mu_{\parallel,heat}\,\left(v\,\left(\nabla_{\mathrm{pol}} v_\parallel\cdot\nabla_{\mathrm{pol}} v_\parallel\right) + v_\parallel\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} v_\parallel\right)\right)
+\end{aligned}
+$$
+
+Source: `_parallel_viscous_heating`.
+
+#### Perpendicular viscous heating
+
+Heating by the perpendicular viscosity of the $u$ flow, with the heating viscosity $\mu_{heat}$.
+
+$$
+\begin{aligned}
+Q_{\mu\perp}(v, \mu_{heat}) &= -\left(\Gamma - 1\right)\,v\,\mu_{heat}\,R^{2}\,f_\mu^{old}\,\left(\nabla_{\mathrm{pol}} u\cdot\nabla_{\mathrm{pol}} \omega\right) - 2\,\left(\Gamma - 1\right)\,v\,\mu_{heat}\,R\,f_\mu^{new}\,\omega\,\partial_{R} u - \left(\Gamma - 1\right)\,v\,\mu_{heat}\,f_\mu^{new}\,\left(\partial_{R} u\,\partial_{R}\partial_{\phi}\partial_{\phi} u + \partial_{Z} u\,\partial_{Z}\partial_{\phi}\partial_{\phi} u\right)
+\end{aligned}
+$$
+
+Source: `_u_viscous_heating`.
+
+#### Coupling to the kinetic neutral and impurity model
+
+Energy and momentum handed over by the kinetic neutral and impurity model.
+
+$$
+\begin{aligned}
+Q_{kin}(v) &= \frac{1}{2}\,\left(\Gamma - 1\right)\,v\,S_\rho^{aux}\,v_\parallel^{2}\,B^2(\psi) - \left(\Gamma - 1\right)\,v\,M_\parallel^{aux}\,v_\parallel
+\end{aligned}
+$$
+
+Source: `_kinetic_coupling`.
+
+#### Transport of the ionization potential energy
+
+The ionization potential energy $W$ is convected like a density, and diffuses with the impurity and main-ion densities that carry it.
+
+$$
+\begin{aligned}
+\hat{D}_\parallel &= D_\parallel + D_{\parallel,sc}\,\tau_{sc} - D_\perp \\
+\hat{D}_{\parallel,imp} &= D_{\parallel,imp} + D_{\parallel,imp,sc}\,\tau_{sc} - D_{\perp,imp} \\
+B^2 &= B^2(\psi) \\
+\mathcal{Q}_{ion}(v, W, T) &= \left(\Gamma - 1\right)\Big(v\,R\,[W,u]^{st} \\
+&\qquad + 2\,v\,W\,\partial_{Z} u \\
+&\qquad - \frac{1}{R^{2}}\,v\,F_0\,v_\parallel\,\partial_{\phi} W \\
+&\qquad - \frac{1}{R}\,v\,v_\parallel\,[W,\psi]^{st} \\
+&\qquad - v\,W\,\left(\mathbf{B}\cdot\nabla v_\parallel\right) \\
+&\qquad - \frac{1}{B^2}\,E_{ion}(T)\,\hat{D}_{\parallel,imp}\,\left(\mathbf{B}\cdot\nabla v\right)\,\left(\mathbf{B}\cdot\nabla \rho_{imp}\right) \\
+&\qquad - E_{ion}(T)\,D_{\perp,imp}\,\mathcal{D}_{tot}(v, \rho_{imp}) \\
+&\qquad - \frac{1}{B^2}\,E_{ion}^{bg}\,\hat{D}_\parallel\,\left(\mathbf{B}\cdot\nabla v\right)\,\left(\mathbf{B}\cdot\nabla \left(\rho - \rho_{imp}\right)\right) \\
+&\qquad - E_{ion}^{bg}\,D_\perp\,\mathcal{D}_{tot}(v, \rho - \rho_{imp})\Big)
+\end{aligned}
+$$
+
+Source: `_ionization_energy_transport`.
+
+
+
+
+## Equation details
+
+The same equations, group by group, as they are written in `model600.py`: the
+local definitions, the mass functional $A$ and the right-hand side $B$ (or the
+constraint $C$), which include the volume weight $w_V$. Each term group is a
+commented block of the corresponding Python function. Inside a mass
+functional, $\overline{x}$ marks a factor that $\partial_t$ does not act on
+(see [Conventions](#conventions)).
+
+
+
+
+
+### Induction equation (`var_psi`)
+
+Source: `induction_equation_1` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\check{T}_e &= \begin{cases} T_e & \text{if } \texttt{with TiTe} \\ T & \text{otherwise} \end{cases} \\
+\hat{T}_e &= \begin{cases} T_e & \text{if } \texttt{with TiTe} \\ \frac{T}{2} & \text{otherwise} \end{cases} \\
+p_e &= \rho\,\hat{T}_e + \rho_{imp}\,A_e(\hat{T}_e) \\
+w_V &= R\,\mathcal{J}
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = \frac{1}{R^{2}}\,v\,\psi\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### -B.grad u
+
+$$
+\begin{aligned}
+&+ \frac{1}{R}\,v\,[\psi,u]^{st} \\
+&- \frac{1}{R^{2}}\,v\,F_0\,\partial_{\phi} u
+\end{aligned}
+$$
+
+##### Eta*j
+
+$$
+\begin{aligned}
+&+ \frac{1}{R^{2}}\,v\,\eta(\check{T}_e, \rho, \rho_{imp})\,\left(j - j_{src} - j_b\right)
+\end{aligned}
+$$
+
+##### Hyper-resistivity
+
+$$
+\begin{aligned}
+&+ \frac{1}{R}\,\eta_{num}(\check{T}_e)\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} j\right)
+\end{aligned}
+$$
+
+##### Runaway current coupling, -eta*j_RE
+
+$$
+\begin{aligned}
+&- \frac{1}{R^{2}}\,v\,\eta(\check{T}_e, \rho, \rho_{imp})\,j_{RE}^{ind}
+\end{aligned}
+$$
+
+##### B.grad Pe diamagnetic term (no impurity contribution)
+
+$$
+\begin{aligned}
+&- \frac{2}{\rho^{\mathrm{c}}\,B^2(\psi)\,R^{3}}\,v\,\tau_{IC}\,F_0^{2}\,[\psi,p_e]^{st} \\
+&+ \frac{2}{\rho^{\mathrm{c}}\,B^2(\psi)\,R^{4}}\,v\,\tau_{IC}\,F_0^{3}\,\partial_{\phi} p_e
+\end{aligned}
+$$
+
+
+
+### Perpendicular momentum equation (`var_u`)
+
+Source: `momentum_equation_2` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\check{T}_e &= \begin{cases} T_e & \text{if } \texttt{with TiTe} \\ T & \text{otherwise} \end{cases} \\
+\alpha_e &= \alpha_e(\check{T}_e) \\
+S_{ion} &= S_{ion}(\check{T}_e) \\
+S_{rec} &= S_{rec}(\check{T}_e) \\
+w_V &= R\,\mathcal{J} \\
+&\text{if } \texttt{with TiTe}\text{:} \\
+\quad p &= \rho\,\left(T_i + T_e\right) + \rho_{imp}\,\left(\alpha_i\,T_i + A_e(T_e)\right) \\
+&\text{otherwise:} \\
+\quad p &= \rho\,T + \rho_{imp}\,A_{imp}(T) \\
+p_i &= p_i^{dia} \\
+n_e &= \rho + \alpha_e\,\rho_{imp} \\
+S_{neut} &= n_e\,\rho_n\,S_{ion} - n_e\,\left(\rho - \rho_{imp}\right)\,S_{rec}
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = R^{2}\,w_V\,\left(-\overline{\rho}^{\mathrm{c}}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right) - f_{cons}\,\rho\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \overline{u}\right)\right)
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Perpendicular inertia (poloidal Jacobian of the kinetic energy)
+
+$$
+\begin{aligned}
+&- \frac{1}{2}\,R\,\left(\left(\partial_{R} u\right)^{2} + \left(\partial_{Z} u\right)^{2}\right)\,[v,\hat{\rho}]
+\end{aligned}
+$$
+
+##### Vorticity advection
+
+$$
+\begin{aligned}
+&- R^{3}\,\rho\,\omega\,[v,u]^{st}
+\end{aligned}
+$$
+
+##### JxB force term. which here is B.gra j
+
+$$
+\begin{aligned}
+&+ \frac{1}{R}\,v\,[\psi,j]^{st} \\
+&- \frac{1}{R^{2}}\,v\,F_0\,\partial_{\phi} j
+\end{aligned}
+$$
+
+##### Pressure advection
+
+$$
+\begin{aligned}
+&+ R\,[v,p]^{st}
+\end{aligned}
+$$
+
+##### Perpendicular and toroidal viscosity
+
+$$
+\begin{aligned}
+&- \mu_\perp(\check{T}_e)\,R^{2}\,f_\mu^{old}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \omega\right) \\
+&- 2\,\mu_\perp(\check{T}_e)\,R\,f_\mu^{new}\,\omega\,\partial_{R} v \\
+&- \mu_\perp(\check{T}_e)\,f_\mu^{new}\,\left(\partial_{R} v\,\partial_{R}\partial_{\phi}\partial_{\phi} u + \partial_{Z} v\,\partial_{Z}\partial_{\phi}\partial_{\phi} u\right) \\
+&- \frac{1}{R}\,\mu_{num}(\check{T}_e)\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \omega
+\end{aligned}
+$$
+
+##### Diamagnetic pressure advection and diamagnetic viscosity
+
+$$
+\begin{aligned}
+&- 2\,v\,\tau_{IC}\,R^{3}\,[p_i,\omega]^{st} \\
+&- 2\,\tau_{IC}\,R^{2}\,\partial_{Z} p_i\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right) \\
+&- 2\,v\,\tau_{IC}\,R^{3}\,\left(\partial_{R}\partial_{Z} u\,\left(\partial_{R}\partial_{R} p_i - \partial_{Z}\partial_{Z} p_i\right) - \partial_{R}\partial_{Z} p_i\,\left(\partial_{R}\partial_{R} u - \partial_{Z}\partial_{Z} u\right)\right) \\
+&+ \mu_\perp'(\check{T}_e)\,W_{dia}\,\left(\nabla_{\mathrm{pol}} \begin{cases} T_i & \text{if } \texttt{with TiTe} \\ \frac{T}{2} & \text{otherwise} \end{cases}\cdot\nabla_{\mathrm{pol}} v\right) \\
+&+ \mu_\perp(\check{T}_e)\,W_{dia}\,\nabla^2_{\mathrm{pol}} v
+\end{aligned}
+$$
+
+##### Conservative form of the momentum equation
+
+$$
+\begin{aligned}
+&+ f_{cons}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right)\,\left(-R\,[\hat{\rho},u] + F_0\,\left(\rho\,\partial_{\phi} v_\parallel + v_\parallel\,\partial_{\phi} \rho\right) + R\,\rho\,[v_\parallel,\psi] + R\,v_\parallel\,[\rho,\psi]\right)
+\end{aligned}
+$$
+
+##### Runaway/auxiliary pressure contributions
+
+$$
+\begin{aligned}
+&- v\,\left(P_\parallel^{RE} + P_\perp^{RE}\right) \\
+&+ R\,\left(-\Pi_R^{RE}\,\partial_{Z} v + \Pi_Z^{RE}\,\partial_{R} v\right)
+\end{aligned}
+$$
+
+##### Momentum carried by the neutral and particle sources
+
+$$
+\begin{aligned}
+&+ R^{2}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u\right)\,\left(\left(1 - \delta_{n}\right)\,S_{neut} + \left(1 - f_{cons}\right)\,\left(S_p + S_{pellet} + S_{bg,drift} + S_{imp,drift}\right)\right)
+\end{aligned}
+$$
+
+##### Taylor-Galerkin stabilization
+
+$$
+\begin{aligned}
+&- \frac{1}{4}\,c_{TG}^{u}\,R^{4}\,\rho\,[\omega,u]\,[v,u]\,\Delta t \\
+&- \frac{1}{4}\,c_{TG}^{u}\,R^{2}\,\omega\,f_{cons}\,[\hat{\rho},u]\,[v,u]\,\Delta t
+\end{aligned}
+$$
+
+**Note:** Neo-classical terms still missing!
+
+
+
+### Current definition (`var_zj`)
+
+Source: `current_constraint_equation_zj` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Constraint
+
+$$
+C = \frac{1}{R}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \psi + v\,j\right)\,\mathcal{J}
+$$
+
+
+
+### Vorticity definition (`var_w`)
+
+Source: `vorticity_constraint_equation_w` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Constraint
+
+$$
+C = \left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} u + v\,\omega\right)\,R\,\mathcal{J}
+$$
+
+
+
+### Density equation (`var_rho`)
+
+Source: `density_equation_rho` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\check{T}_e &= \begin{cases} T_e & \text{if } \texttt{with TiTe} \\ T & \text{otherwise} \end{cases} \\
+\rho_{main} &= \rho - \rho_{imp} \\
+n_e &= \rho + \alpha_e(\check{T}_e)\,\rho_{imp} \\
+D_\parallel^{\mathrm{tot}} &= D_\parallel + D_{\parallel,sc}\,\tau_{sc} \\
+D_{\parallel,imp}^{\mathrm{tot}} &= D_{\parallel,imp} + D_{\parallel,imp,sc}\,\tau_{sc} \\
+w_V &= R\,\mathcal{J}
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\rho\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Particle sources
+
+$$
+\begin{aligned}
+&+ v\,\left(S_p + S_{pellet} + S_{bg,drift} + S_{imp,drift} + S_\rho^{aux}\right)
+\end{aligned}
+$$
+
+##### Sources/sinks due to neutrals
+
+$$
+\begin{aligned}
+&+ v\,n_e\,\rho_n\,S_{ion}(\check{T}_e) \\
+&- v\,n_e\,\rho_{main}\,S_{rec}(\check{T}_e)
+\end{aligned}
+$$
+
+##### Convection/compression by u variable (ExB)
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_u(\rho)
+\end{aligned}
+$$
+
+##### Parallel convection/compression by vpar
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_\parallel(\rho)
+\end{aligned}
+$$
+
+##### Parallel and perpendicular diffusion of main ions
+
+$$
+\begin{aligned}
+&+ \left(D_\parallel^{\mathrm{tot}} - D_\perp\right)\,\mathcal{D}_\parallel(v, \rho_{main}) \\
+&- D_\perp\,\mathcal{D}_{tot}(v, \rho_{main})
+\end{aligned}
+$$
+
+##### Parallel and perpendicular diffusion of impurities (total mass)
+
+$$
+\begin{aligned}
+&+ \left(D_{\parallel,imp}^{\mathrm{tot}} - D_{\perp,imp}\right)\,\mathcal{D}_\parallel(v, \rho_{imp}) \\
+&- D_{\perp,imp}\,\mathcal{D}_{tot}(v, \rho_{imp})
+\end{aligned}
+$$
+
+##### Numerical stabilization
+
+$$
+\begin{aligned}
+&+ c_{TG}^{\rho}\,\mathcal{T}_n(v, \rho) \\
+&- D_{\perp,num}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \rho
+\end{aligned}
+$$
+
+##### Diamagnetic drift
+
+$$
+\begin{aligned}
+&+ 4\,v\,\tau_{IC}\,\partial_{Z} p_i^{dia}
+\end{aligned}
+$$
+
+##### Pinch term, not sure about this one here
+
+$$
+\begin{aligned}
+&- \frac{1}{\sqrt{\left(\partial_{R} \psi\right)^{2} + \left(\partial_{Z} \psi\right)^{2}}}\,V_{pinch}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}} \psi\right)\,\rho
+\end{aligned}
+$$
+
+
+
+### Parallel velocity equation (`var_vpar`)
+
+The equation is the projection of the momentum equation on $\mathbf{B}$. With the parallel flow $v_\parallel\mathbf{B}$ and $\mathbf{B}\cdot\partial_t\mathbf{B} = \frac{1}{2}\partial_t B^2 = \frac{1}{2}\partial_t B^2_{pol}$, $\mathbf{B}\cdot\rho\,\partial_t(v_\parallel\mathbf{B}) = \rho\,\big(B^2\,\partial_t v_\parallel + \frac{1}{2}v_\parallel\,\partial_t B^2_{pol}\big)$: the time derivative acts separately on $v_\parallel$ and on $B^2_{pol}$, and is not the derivative of $\rho\,v_\parallel B^2$. In the conservative form ($f_{cons} = 1$) it adds $v_\parallel B^2\,\partial_t\rho$. The element routine uses the corrected density $\rho^{\mathrm{c}}$ in the first two terms.
+
+Source: `parallel_velocity_equation_vpar` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\check{T}_e &= \begin{cases} T_e & \text{if } \texttt{with TiTe} \\ T & \text{otherwise} \end{cases} \\
+B^2 &= B^2(\psi) \\
+w_V &= R\,\mathcal{J} \\
+n_e &= \rho + \alpha_e(\check{T}_e)\,\rho_{imp} \\
+p &= \begin{cases} \rho\,\left(T_i + T_e\right) & \text{if } \texttt{with TiTe} \\ \rho\,T & \text{otherwise} \end{cases} \\
+&\text{if } \texttt{with TiTe}\text{:} \\
+\quad p &\mathrel{+}= \rho_{imp}\,\left(\alpha_i\,T_i + A_e(T_e)\right) \\
+&\text{otherwise:} \\
+\quad p &\mathrel{+}= \rho_{imp}\,A_{imp}(T) \\
+\mu_\parallel^{\mathrm{eff}} &= \mu_\parallel + \mu_{\parallel,sc}\,\tau_{sc} \\
+S_n &= S_p + S_{pellet} + S_{bg,drift} + S_{imp,drift} \\
+\nabla_{\mathrm{pol}}(v_\parallel - V_{rot}) &= \nabla_{\mathrm{pol}} v_\parallel - \frac{\mathrm{d}V_{rot}}{\mathrm{d}\psi}\,\nabla_{\mathrm{pol}} \psi
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\left(\overline{\rho}^{\mathrm{c}}\,v_\parallel\,B^2(\overline{\psi}) + \frac{1}{2}\,\overline{\rho}^{\mathrm{c}}\,\overline{v_\parallel}\,B^2_{pol}(\psi) + f_{cons}\,\rho\,\overline{v_\parallel}\,B^2(\overline{\psi})\right)\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Parallel pressure gradient
+
+$$
+\begin{aligned}
+&- v\,\left(\mathbf{B}\cdot\nabla p\right)
+\end{aligned}
+$$
+
+##### Parallel advection of the kinetic energy, 0.5*v_par**2*B**2
+
+$$
+\begin{aligned}
+&+ \frac{1}{2}\,v_\parallel^{2}\,B^2\,\left(\rho\,\left(\mathbf{B}\cdot\nabla v\right) + v\,\left(\mathbf{B}\cdot\nabla \rho\right)\right)
+\end{aligned}
+$$
+
+##### Term to obtain conservative form of momentum equation
+
+$$
+\begin{aligned}
+&+ f_{cons}\,v\,v_\parallel\,B^2\,\left(\frac{[\hat{\rho},u]}{R} - \mathbf{B}\cdot\nabla \left(\rho\,v_\parallel\right)\right)
+\end{aligned}
+$$
+
+##### Numerical and physical parallel viscosities
+
+$$
+\begin{aligned}
+&- \mu_{\parallel,num}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} v_\parallel \\
+&- \frac{1}{R^{2}\,B^2}\,\mu_{\parallel\parallel}\,F_0^{2}\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)\,\left(\mathbf{B}\cdot\nabla v\right) \\
+&- \mu_\parallel^{\mathrm{eff}}\,\left(\nabla_{\mathrm{pol}} v\cdot\nabla_{\mathrm{pol}}(v_\parallel - V_{rot})\right)
+\end{aligned}
+$$
+
+##### External momentum sources
+
+$$
+\begin{aligned}
+&+ v\,M_\parallel^{aux}
+\end{aligned}
+$$
+
+##### Momentum carried by the particle sources; not active in conservative form
+
+$$
+\begin{aligned}
+&- v\,S_n\,v_\parallel\,B^2\,\left(1 - f_{cons}\right) \\
+&- v\,S_\rho^{aux}\,v_\parallel\,B^2\,\left(1 - f_{cons}\right)
+\end{aligned}
+$$
+
+##### This one should probably be multiplied by (1 - fact_conservative_u
+
+$$
+\begin{aligned}
+&+ \left(1 - \delta_{n}\right)\,v\,n_e\,v_\parallel\,B^2\,\left(\left(\rho - \rho_{imp}\right)\,S_{rec}(\check{T}_e) - \rho_n\,S_{ion}(\check{T}_e)\right)
+\end{aligned}
+$$
+
+##### Taylor-Galerkin stabilization
+
+$$
+\begin{aligned}
+&- \frac{1}{4}\,c_{TG}^{v_\parallel}\,\rho\,v_\parallel^{2}\,B^2\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)\,\left(\mathbf{B}\cdot\nabla v\right)\,\Delta t \\
+&- \frac{1}{4}\,c_{TG}^{v_\parallel}\,v\,v_\parallel^{2}\,B^2\,\left(1 - f_{cons}\right)\,\left(\mathbf{B}\cdot\nabla v_\parallel\right)\,\left(\mathbf{B}\cdot\nabla \rho\right)\,\Delta t \\
+&- \frac{1}{4}\,c_{TG}^{v_\parallel}\,v_\parallel^{3}\,B^2\,f_{cons}\,\left(\mathbf{B}\cdot\nabla \rho\right)\,\left(\mathbf{B}\cdot\nabla v\right)\,\Delta t
+\end{aligned}
+$$
+
+##### Not sure this pinch term should be here
+
+$$
+\begin{aligned}
+&+ \frac{1}{\sqrt{\left(\partial_{R} \psi\right)^{2} + \left(\partial_{Z} \psi\right)^{2}}}\,V_{pinch}\,\left(\nabla_{\mathrm{pol}} \psi\cdot\nabla_{\mathrm{pol}} v_\parallel\right)\,\rho\,v
+\end{aligned}
+$$
+
+
+
+### Impurity density equation (`var_rhoimp`)
+
+`with_TiTe` has no effect here; it is accepted so that every density equation has the same interface.
+
+Source: `impurity_density_equation_rhoimp` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+D_{\parallel,imp}^{\mathrm{tot}} &= D_{\parallel,imp} + D_{\parallel,imp,sc}\,\tau_{sc} \\
+w_V &= R\,\mathcal{J}
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\rho_{imp}\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Particle sources
+
+$$
+\begin{aligned}
+&+ v\,S_{imp,drift}
+\end{aligned}
+$$
+
+##### Convection/compression by u variable (ExB)
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_u(\rho_{imp})
+\end{aligned}
+$$
+
+##### Parallel convection/compression by vpar
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_\parallel(\rho_{imp})
+\end{aligned}
+$$
+
+##### Parallel and perpendicular diffusion
+
+$$
+\begin{aligned}
+&+ \left(D_{\parallel,imp}^{\mathrm{tot}} - D_{\perp,imp}\right)\,\mathcal{D}_\parallel(v, \rho_{imp}) \\
+&- D_{\perp,imp}\,\mathcal{D}_{tot}(v, \rho_{imp})
+\end{aligned}
+$$
+
+##### Numerical stabilization
+
+$$
+\begin{aligned}
+&+ c_{TG}^{\rho_{imp}}\,\mathcal{T}_n(v, \rho_{imp}) \\
+&- D_{\perp,num}^{n}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \rho_{imp}
+\end{aligned}
+$$
+
+
+
+### Neutral density equation (`var_rhon`)
+
+Fluid neutrals are diffused with an anisotropic diffusivity, convected with the plasma flow when $\delta_n = 1$, ionized and recombined with the same rates as in the density equation, and fed by a prescribed source.
+
+Source: `neutral_density_equation_rhon` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\check{T}_e &= \begin{cases} T_e & \text{if } \texttt{with TiTe} \\ T & \text{otherwise} \end{cases} \\
+n_e &= \rho + \alpha_e(\check{T}_e)\,\rho_{imp} \\
+\rho_{main} &= \rho - \rho_{imp} \\
+w_V &= R\,\mathcal{J}
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\rho_n\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Anisotropic diffusion
+
+$$
+\begin{aligned}
+&- D_{n,R}\,\partial_{R} v\,\partial_{R} \rho_n \\
+&- D_{n,Z}\,\partial_{Z} v\,\partial_{Z} \rho_n \\
+&- \frac{1}{R^{2}}\,D_{n,\phi}\,\partial_{\phi} v\,\partial_{\phi} \rho_n
+\end{aligned}
+$$
+
+##### Convection/compression with the plasma flow
+
+$$
+\begin{aligned}
+&+ \delta_{n}\,v\,\left(\mathcal{C}_u(\rho_n) + \mathcal{C}_\parallel(\rho_n)\right)
+\end{aligned}
+$$
+
+##### Ionization and recombination with the main/impurity ions
+
+$$
+\begin{aligned}
+&- v\,n_e\,\rho_n\,S_{ion}(\check{T}_e) \\
+&+ v\,n_e\,\rho_{main}\,S_{rec}(\check{T}_e)
+\end{aligned}
+$$
+
+##### Neutral source
+
+$$
+\begin{aligned}
+&+ v\,S_n^{drift}
+\end{aligned}
+$$
+
+##### Numerical stabilization
+
+$$
+\begin{aligned}
+&- D_{\perp,num}^{n}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} \rho_n
+\end{aligned}
+$$
+
+
+
+### Ion energy equation (`var_Ti`)
+
+Source: `ion_energy_equation_Ti` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\mathcal{J}\,[\cdot,\cdot] &= \begin{cases} \mathcal{J}\,[\cdot,\cdot]^{st} & \text{if } \texttt{st form} \\ (a, b) \mapsto [a,b]\,\mathcal{J} & \text{otherwise} \end{cases} \\
+w_V &= R\,\mathcal{J} \\
+n_i &= \rho + \alpha_i\,\rho_{imp} \\
+p_i &= n_i\,T_i
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\left(\rho^{\mathrm{c}} + \alpha_i\,\rho_{imp}^{\mathrm{c}}\right)\,T_i\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Convection/compression by u (ExB)
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_u^{p}(p_i)
+\end{aligned}
+$$
+
+##### Parallel convection/compression by vpar
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_\parallel^{p}(p_i)
+\end{aligned}
+$$
+
+##### Parallel heat conduction
+
+$$
+\begin{aligned}
+&+ \left(\kappa_{\parallel,i}(T_i) - \kappa_{\perp,i}(\rho)\right)\,\mathcal{D}_\parallel(v, T_i)
+\end{aligned}
+$$
+
+##### Perpendicular (total) heat conduction
+
+$$
+\begin{aligned}
+&- \kappa_{\perp,i}(\rho)\,\mathcal{D}_{tot}(v, T_i)
+\end{aligned}
+$$
+
+##### Heat sources/sinks
+
+$$
+\begin{aligned}
+&+ v\,\left(H_i + H_i^{aux}\right)
+\end{aligned}
+$$
+
+##### Energy exchange with electrons
+
+$$
+\begin{aligned}
+&+ v\,Q_{ie}(T_i, T_e, \rho, \rho_{imp})
+\end{aligned}
+$$
+
+##### Neutral recombination sink
+
+$$
+\begin{aligned}
+&- v\,T_i\,\left(\rho^{\mathrm{c}}\right)^{2}\,S_{rec}(T_e)
+\end{aligned}
+$$
+
+##### Kinetic energy released by particle sources (friction heating)
+
+$$
+\begin{aligned}
+&+ \frac{1}{2}\,v\,\left(\Gamma - 1\right)\,\left(v_\parallel^{2}\,B^2(\psi) + R^{2}\,\left(\left(\partial_{R} u\right)^{2} + \left(\partial_{Z} u\right)^{2}\right)\right)\,S_{kin}(T_e)
+\end{aligned}
+$$
+
+##### Parallel viscous heating
+
+$$
+\begin{aligned}
+&+ Q_{\mu\parallel}(v)
+\end{aligned}
+$$
+
+##### Perpendicular (u) viscous heating
+
+$$
+\begin{aligned}
+&+ Q_{\mu\perp}(v, \mu_{heat}(T_e))
+\end{aligned}
+$$
+
+##### Energy and momentum from the kinetic neutral/impurity model
+
+$$
+\begin{aligned}
+&+ Q_{kin}(v)
+\end{aligned}
+$$
+
+##### Numerical stabilization
+
+$$
+\begin{aligned}
+&+ \mathcal{T}_p(v, p_i; c_{TG}^{T_i}) \\
+&- \kappa_{\perp,num}^{i}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} T_i
+\end{aligned}
+$$
+
+##### Implicit heat source to avoid negative temperatures
+
+$$
+\begin{aligned}
+&+ \mathcal{H}(v, \epsilon_i^{\mathrm{floor}}(T_i), T_i^{\mathrm{floor}}(T_i), T_{\mathrm{min}})
+\end{aligned}
+$$
+
+
+
+### Electron energy equation (`var_Te`)
+
+Source: `electron_energy_equation_Te` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\mathcal{J}\,[\cdot,\cdot] &= \begin{cases} \mathcal{J}\,[\cdot,\cdot]^{st} & \text{if } \texttt{st form} \\ (a, b) \mapsto [a,b]\,\mathcal{J} & \text{otherwise} \end{cases} \\
+w_V &= R\,\mathcal{J} \\
+p_e &= \rho\,T_e + \rho_{imp}\,A_e(T_e) \\
+n_e^{\mathrm{c}} &= \rho^{\mathrm{c}} + \alpha_e(T_e)\,\rho_{imp}^{\mathrm{c}} \\
+W_{ion} &= E_{ion}(T_e)\,\rho_{imp} + E_{ion}^{bg}\,\left(\rho - \rho_{imp}\right)
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\left(\rho^{\mathrm{c}}\,T_e + \rho_{imp}^{\mathrm{c}}\,A_e(T_e) + \left(\Gamma - 1\right)\,W_{ion}\right)\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Convection/compression by u (ExB)
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_u^{p}(p_e)
+\end{aligned}
+$$
+
+##### Parallel convection/compression by vpar
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_\parallel^{p}(p_e)
+\end{aligned}
+$$
+
+##### Parallel heat conduction
+
+$$
+\begin{aligned}
+&+ \left(\kappa_{\parallel,e}(T_e) - \kappa_{\perp,e}(\rho)\right)\,\mathcal{D}_\parallel(v, T_e)
+\end{aligned}
+$$
+
+##### Perpendicular (total) heat conduction
+
+$$
+\begin{aligned}
+&- \kappa_{\perp,e}(\rho)\,\mathcal{D}_{tot}(v, T_e)
+\end{aligned}
+$$
+
+##### Heat sources/sinks
+
+$$
+\begin{aligned}
+&+ v\,\left(H_e + H_e^{aux} + P_{teleport}\right)
+\end{aligned}
+$$
+
+##### Ionization sink due to neutrals
+
+$$
+\begin{aligned}
+&- v\,\xi_{ion}\,\left(\rho + \alpha_e(T_e)\,\rho_{imp}\right)\,\rho_n\,S_{ion}(T_e)
+\end{aligned}
+$$
+
+##### Energy exchange with ions
+
+$$
+\begin{aligned}
+&+ v\,Q_{ei}(T_i, T_e, \rho, \rho_{imp})
+\end{aligned}
+$$
+
+##### Ohmic heating
+
+$$
+\begin{aligned}
+&+ v\,\left(\Gamma - 1\right)\,\eta_{ohm}(T_e, \rho, \rho_{imp})\,\left(\frac{1}{R}\left(j - j_{RE}\right)\right)^{2}
+\end{aligned}
+$$
+
+##### Neutral (line) radiation
+
+$$
+\begin{aligned}
+&- v\,n_e^{\mathrm{c}}\,\rho_n^{\mathrm{c}}\,L_{rays}(T_e)
+\end{aligned}
+$$
+
+##### Background radiation
+
+$$
+\begin{aligned}
+&- v\,n_e^{\mathrm{c}}\,\left(\rho^{\mathrm{c}} - \rho_{imp}^{\mathrm{c}}\right)\,L_{cont}(T_e) \\
+&- v\,n_e^{\mathrm{c}}\,f_{rad}^{bg}(T_e)
+\end{aligned}
+$$
+
+##### Impurity radiation
+
+$$
+\begin{aligned}
+&- v\,n_e^{\mathrm{c}}\,\rho_{imp}^{\mathrm{c}}\,L_{rad}(T_e)
+\end{aligned}
+$$
+
+##### Ionization potential transport
+
+$$
+\begin{aligned}
+&+ \mathcal{Q}_{ion}(v, W_{ion}, T_e)
+\end{aligned}
+$$
+
+##### Numerical stabilization
+
+$$
+\begin{aligned}
+&+ \mathcal{T}_p(v, p_e; c_{TG}^{T_e}) \\
+&- \kappa_{\perp,num}^{e}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} T_e
+\end{aligned}
+$$
+
+##### Implicit heat source to avoid negative temperatures
+
+$$
+\begin{aligned}
+&+ \mathcal{H}(v, \epsilon_e^{\mathrm{floor}}(T_e), T_e^{\mathrm{floor}}(T_e), T_{\mathrm{min}})
+\end{aligned}
+$$
+
+
+
+### Single-temperature energy equation (`var_T`)
+
+Source: `total_energy_equation_T` in `util/equation_codegen/src/jorek_equations/model600.py`.
+
+#### Definitions
+
+$$
+\begin{aligned}
+\mathcal{J}\,[\cdot,\cdot] &= \begin{cases} \mathcal{J}\,[\cdot,\cdot]^{st} & \text{if } \texttt{st form} \\ (a, b) \mapsto [a,b]\,\mathcal{J} & \text{otherwise} \end{cases} \\
+w_V &= R\,\mathcal{J} \\
+p &= \rho\,T + \rho_{imp}\,A_{imp}(T) \\
+n_e^{\mathrm{c}} &= \rho^{\mathrm{c}} + \alpha_e(T)\,\rho_{imp}^{\mathrm{c}} \\
+W_{ion} &= E_{ion}(T)\,\rho_{imp} + E_{ion}^{bg}\,\left(\rho - \rho_{imp}\right)
+\end{aligned}
+$$
+
+#### Time derivative
+
+$$
+A = v\,\left(\rho^{\mathrm{c}}\,T + \rho_{imp}^{\mathrm{c}}\,A_{imp}(T) + \left(\Gamma - 1\right)\,W_{ion}\right)\,w_V
+$$
+
+#### Right-hand side
+
+$$
+B = \Big(\textstyle\sum_k b_k\Big)\,w_V
+$$
+
+##### Convection/compression by u (ExB)
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_u^{p}(p)
+\end{aligned}
+$$
+
+##### Parallel convection/compression by vpar
+
+$$
+\begin{aligned}
+&+ v\,\mathcal{C}_\parallel^{p}(p)
+\end{aligned}
+$$
+
+##### Parallel heat conduction
+
+$$
+\begin{aligned}
+&+ \left(\kappa_\parallel(T) - \kappa_\perp(\rho)\right)\,\mathcal{D}_\parallel(v, T)
+\end{aligned}
+$$
+
+##### Perpendicular (total) heat conduction
+
+$$
+\begin{aligned}
+&- \kappa_\perp(\rho)\,\mathcal{D}_{tot}(v, T)
+\end{aligned}
+$$
+
+##### Heat sources/sinks
+
+$$
+\begin{aligned}
+&+ v\,\left(H + H^{aux} + P_{teleport}\right)
+\end{aligned}
+$$
+
+##### Ionization sink due to neutrals
+
+$$
+\begin{aligned}
+&- v\,\xi_{ion}\,\left(\rho + \alpha_e(T)\,\rho_{imp}\right)\,\rho_n\,S_{ion}(T)
+\end{aligned}
+$$
+
+##### Recombination sink due to neutrals
+
+$$
+\begin{aligned}
+&- \frac{1}{2}\,\left(\Gamma - 1\right)\,v\,T\,\left(\rho^{\mathrm{c}}\right)^{2}\,S_{rec}(T)
+\end{aligned}
+$$
+
+##### Ohmic heating
+
+$$
+\begin{aligned}
+&+ v\,\left(\Gamma - 1\right)\,\eta_{ohm}(T, \rho, \rho_{imp})\,\left(\frac{1}{R}\left(j - j_{RE}\right)\right)^{2}
+\end{aligned}
+$$
+
+##### Neutral (line) radiation
+
+$$
+\begin{aligned}
+&- v\,n_e^{\mathrm{c}}\,\rho_n^{\mathrm{c}}\,L_{rays}(T)
+\end{aligned}
+$$
+
+##### Background radiation
+
+$$
+\begin{aligned}
+&- v\,n_e^{\mathrm{c}}\,\left(\rho^{\mathrm{c}} - \rho_{imp}^{\mathrm{c}}\right)\,L_{cont}(T) \\
+&- v\,n_e^{\mathrm{c}}\,f_{rad}^{bg}(T)
+\end{aligned}
+$$
+
+##### Impurity radiation
+
+$$
+\begin{aligned}
+&- v\,n_e^{\mathrm{c}}\,\rho_{imp}^{\mathrm{c}}\,L_{rad}(T)
+\end{aligned}
+$$
+
+##### Ionization potential transport
+
+$$
+\begin{aligned}
+&+ \mathcal{Q}_{ion}(v, W_{ion}, T)
+\end{aligned}
+$$
+
+##### Kinetic energy released by particle sources (friction heating)
+
+$$
+\begin{aligned}
+&+ \frac{1}{2}\,v\,\left(\Gamma - 1\right)\,\left(v_\parallel^{2}\,B^2(\psi) + R^{2}\,\left(\left(\partial_{R} u\right)^{2} + \left(\partial_{Z} u\right)^{2}\right)\right)\,S_{kin}(T)
+\end{aligned}
+$$
+
+##### Parallel viscous heating
+
+$$
+\begin{aligned}
+&+ Q_{\mu\parallel}(v)
+\end{aligned}
+$$
+
+##### Perpendicular (u) viscous heating
+
+$$
+\begin{aligned}
+&+ Q_{\mu\perp}(v, \mu_{heat}(T))
+\end{aligned}
+$$
+
+##### Energy and momentum from the kinetic neutral/impurity model
+
+$$
+\begin{aligned}
+&+ Q_{kin}(v)
+\end{aligned}
+$$
+
+##### Numerical stabilization
+
+$$
+\begin{aligned}
+&+ \mathcal{T}_p(v, p; c_{TG}^{T}) \\
+&- \kappa_{\perp,num}\,\nabla^2_{\mathrm{pol}} v\,\nabla^2_{\mathrm{pol}} T
+\end{aligned}
+$$
+
+##### Implicit heat source to avoid negative temperatures
+
+$$
+\begin{aligned}
+&+ \mathcal{H}(v, \epsilon^{\mathrm{floor}}(T), T^{\mathrm{floor}}(T), T_{\mathrm{min}})
+\end{aligned}
+$$
+
+
+
+
+## Regenerating this page
+
+The equations are edited in `util/equation_codegen/src/jorek_equations/model600.py`
+only; do not edit the generated regions of this page (between the
+`BEGIN GENERATED` and `END GENERATED` comments), which are overwritten. The
+prose outside them is hand-written. How each Python name is written in LaTeX,
+and the titles and descriptions of the operators, are in
+`util/equation_codegen/src/jorek_equations/model600_notation.py`.
+
+The tools need Python 3 with `sympy`; on the ITER cluster, run
+`module load sympy/1.14.0-gfbf-2025b` first. From `util/equation_codegen`:
+
+```bash
+python3 examples/model600_docs.py render # rewrite the generated regions
+python3 examples/model600_docs.py check # change nothing; exit 1 if the page is stale
+```
+
+The unit tests and `run_test.sh` run `check`, so a change to `model600.py`
+that is not rendered into this page fails the build.
diff --git a/models/model600/mod_elt_matrix_fft.f90 b/models/model600/mod_elt_matrix_fft.f90
index c6b65b15e9..7a13f0212c 100644
--- a/models/model600/mod_elt_matrix_fft.f90
+++ b/models/model600/mod_elt_matrix_fft.f90
@@ -57,7 +57,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
real*8 :: R_axis, Z_axis, psi_axis, psi_bnd, R_xpoint(2), Z_xpoint(2), dj_dpsi, dj_dz, source_pellet, source_volume
real*8 :: Bgrad_rho_star, Bgrad_rho, Bgrad_vpar, Bgrad_T_star, Bgrad_Ti, Bgrad_Te, Bgrad_T, BB2
real*8 :: Bgrad_rho_star_psi, Bgrad_rho_psi, Bgrad_rho_rho, Bgrad_vpar_vpar, Bgrad_vpar_psi, Bgrad_T_star_psi, Bgrad_Ti_psi, Bgrad_T_psi, Bgrad_Ti_Ti, Bgrad_Te_psi, Bgrad_T_T, Bgrad_Te_Te, BB2_psi
-real*8 :: Bgrad_rho_k_star, Bgrad_T_k_star, Bgrad_Ti_Ti_n, Bgrad_Te_Te_n, Bgrad_T_T_n, Bgrad_rho_rho_n
+real*8 :: Bgrad_rho_k_star, Bgrad_T_k_star, Bgrad_Ti_Ti_n, Bgrad_Te_Te_n, Bgrad_T_T_n, Bgrad_rho_rho_n, Bgrad_vpar_vpar_n
real*8 :: Bgrad_rhoimp, Bgrad_rhoimp_psi, Bgrad_rhoimp_rhoimp, Bgrad_rhoimp_rhoimp_n
real*8 :: ZK_par_T, dZK_par_dT, ZKi_par_T, dZKi_par_dT, ZKe_par_T, dZKe_par_dT
real*8 :: D_prof, D_par_local, ZK_prof, ZKi_prof, ZKe_prof, psi_norm, theta, zeta, delta_u_x, delta_u_y, delta_ps_x, delta_ps_y
@@ -1712,13 +1712,13 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- visco_par_num * (v_xx + v_x/Bigr + v_yy)*(vpar0_xx + vpar0_x/Bigr + vpar0_yy) * BigR * xjac * tstep * factor(var_vpar,4) &
- + zeta * v * delta_g(mp,var_vpar,ms,mt) * r0_corr * F0**2 / BigR * xjac * factor(var_vpar,5) &
+ + zeta * v * delta_g(mp,var_vpar,ms,mt) * r0_corr * BB2 * BigR * xjac * factor(var_vpar,5) &
+ zeta * v * r0_corr * vpar0 * (ps0_x * delta_ps_x + ps0_y * delta_ps_y) / BigR * xjac * factor(var_vpar,5) &
! New terms coming from -(\partial_t \rho + \nabla \cdot (\rho \mathbf{v})) \mathbf{v} in RHS of momentum equation
! (see wiki: https://www.jorek.eu/wiki/doku.php?id=model500_501_555#equations):
+ fact_conservative_u * ( &
- + zeta * v * delta_rho_g * vpar0 * F0**2 / BigR * xjac &
+ + zeta * v * delta_rho_g * vpar0 * BB2 * BigR * xjac &
+ v * (r0_x_hat * u0_y - r0_y_hat * u0_x) * vpar0 * BB2 * xjac * tstep &
- v * F0 / BigR * (r0 * vpar0_p + r0_p * vpar0) * vpar0 * BB2 * xjac * tstep &
- v * r0 * (vpar0_x * ps0_y - vpar0_y * ps0_x) * vpar0 * BB2 * xjac * tstep &
@@ -2010,7 +2010,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- v * (r0 + rimp0*alpha_imp_bis) * F0 / BigR * Vpar0 * T0_p * xjac * tstep * factor(var_T,4 ) &
- v * T0 * F0 / BigR * Vpar0 * (r0_p + alpha_imp * rimp0_p) * xjac * tstep * factor(var_T,4 ) &
- - v * (r0 + rn0*alpha_imp_bis) * Vpar0 * (T0_s * ps0_t - T0_t * ps0_s) * tstep * factor(var_T,4 ) &
+ - v * (r0 + rimp0*alpha_imp_bis) * Vpar0 * (T0_s * ps0_t - T0_t * ps0_s) * tstep * factor(var_T,4 ) &
- v * T0 * Vpar0 * (r0_s * ps0_t - r0_t * ps0_s) * tstep * factor(var_T,4 ) &
- v * T0 * Vpar0 * alpha_imp * (rimp0_s * ps0_t - rimp0_t * ps0_s) * tstep * factor(var_T,4 ) &
@@ -2200,7 +2200,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
rhs_ij_k(var_rhoimp) = &
! The new diffusion scheme for the impurities
- - (D_par_local_imp-D_prof_imp) * BigR / BB2 * Bgrad_rho_k_star * (Bgrad_rhoimp) * xjac * tstep * factor(var_rhoimp,1)&
+ - ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_k_star * (Bgrad_rhoimp) * xjac * tstep * factor(var_rhoimp,1)&
- D_prof_imp * BigR * ( v_p * rimp0_p /BigR**2) * xjac * tstep * factor(var_rhoimp,2)&
- tgnum_rhoimp * 0.25d0 / BigR * vpar0**2 &
@@ -2344,6 +2344,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
Bgrad_rhoimp_rhoimp_n = ( F0 / BigR * rhoimp_p ) / BigR
Bgrad_vpar_psi = ( vpar0_x * psi_y - vpar0_y * psi_x ) / BigR
Bgrad_vpar_vpar = ( vpar_x * ps0_y - vpar_y * ps0_x ) / BigR
+ Bgrad_vpar_vpar_n = ( F0 / BigR * vpar_p ) / BigR
BB2_psi = 2.d0 * (psi_x * ps0_x + psi_y * ps0_y ) /BigR**2
@@ -2645,10 +2646,10 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ dvisco_dT * T * 2.d0 * v_x * w0 * BigR**2.d0 * visco_fact_new * xjac * theta * tstep &
+ dvisco_dT * T * (v_x*u0_xpp + v_y*u0_ypp) * BigR * visco_fact_new * xjac * theta * tstep &
- ! --- Contributions of the diamagnetic viscosity
- - dvisco_dT * bigR * W_dia_Ti * (v_x*Ti0_x + v_y*Ti0_y) * xjac * theta * tstep &
- - visco_T * bigR * W_dia_Ti * (v_xx + v_x/bigR + v_yy) * xjac * theta * tstep &
- - dvisco_dT * bigR * W_dia * (v_x*T_x + v_y*T_y ) * xjac * theta * tstep &
+ ! --- Contributions of the diamagnetic viscosity (/2 comes from dTi/dT)
+ - dvisco_dT * bigR * W_dia_Ti * (v_x*Ti0_x + v_y*Ti0_y) / 2.d0 * xjac * theta * tstep &
+ - visco_T * bigR * W_dia_Ti * (v_xx + v_x/bigR + v_yy) / 2.d0 * xjac * theta * tstep &
+ - dvisco_dT * bigR * W_dia * (v_x*T_x + v_y*T_y ) / 2.d0 * xjac * theta * tstep &
- d2visco_dT2*T * bigR * W_dia * (v_x*Ti0_x + v_y*Ti0_y) * xjac * theta * tstep &
- dvisco_dT*T * bigR * W_dia * (v_xx + v_x/bigR + v_yy) * xjac * theta * tstep &
@@ -2815,11 +2816,15 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
if ( with_TiTe ) then
amat(var_rho,var_Ti) = - v * 2.d0 * tauIC*2. * (Ti_y * r0 + Ti*r0_y) * BigR * xjac * theta * tstep
amat(var_rho,var_Te) = - v * BigR * (r0+alpha_e*rimp0) * rn0 * dSion_dT * Te * xjac * theta * tstep &
- + v * BigR * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * Te * xjac * theta * tstep
+ + v * BigR * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * Te * xjac * theta * tstep &
+ - v * BigR * dalpha_e_dT * rimp0 * rn0 * Sion_T * Te * xjac * theta * tstep &
+ + v * BigR * dalpha_e_dT * rimp0 * (r0-rimp0) * Srec_T * Te * xjac * theta * tstep
else ! (with_TiTe)
amat(var_rho,var_T) = - v * 2.d0 * tauIC * (T_y * r0 + T *r0_y) * BigR * xjac * theta * tstep &
- v * BigR * (r0+alpha_e*rimp0) * rn0 * dSion_dT * T * xjac * theta * tstep &
- + v * BigR * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * T * xjac * theta * tstep
+ + v * BigR * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * T * xjac * theta * tstep &
+ - v * BigR * dalpha_e_dT * rimp0 * rn0 * Sion_T * T * xjac * theta * tstep &
+ + v * BigR * dalpha_e_dT *rimp0 * (r0-rimp0) * Srec_T * T * xjac * theta * tstep
end if ! (with_TiTe)
if ( with_vpar ) then
@@ -2915,27 +2920,54 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2 &
* (-(ps0_s * vpar0_t - ps0_t * vpar0_s)/xjac) / BigR &
* (-(psi_s * v_t - psi_t * v_s) /xjac) * xjac * theta * tstep*tstep &
-
+
+ tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2 &
* (-(psi_s * vpar0_t - psi_t * vpar0_s)/xjac) / BigR &
* (-(ps0_s * v_t - ps0_t * v_s) /xjac) * xjac * theta * tstep*tstep &
-
+ ! d/dpsi of the F0/BigR*vpar0_p piece of Bgrad_vpar times the unvaried v-bracket
+ + tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2 &
+ * (F0 / BigR * vpar0_p) / BigR &
+ * (-(psi_s * v_t - psi_t * v_s) /xjac) * xjac * theta * tstep*tstep &
+
+ tgnum_vpar * 0.25d0 * v * Vpar0**2 * BB2 * (1.d0 - fact_conservative_u) &
* (-(ps0_s * vpar0_t - ps0_t * vpar0_s)/xjac) / BigR &
* (-(psi_s * r0_t - psi_t * r0_s) /xjac) * xjac * theta * tstep*tstep &
-
+ ! d/dpsi of the Bgrad_rho bracket times the unvaried F0/BigR*vpar0_p piece of Bgrad_vpar
+ + tgnum_vpar * 0.25d0 * v * Vpar0**2 * BB2 * (1.d0 - fact_conservative_u) &
+ * (F0 / BigR * vpar0_p) / BigR &
+ * (-(psi_s * r0_t - psi_t * r0_s) /xjac) * xjac * theta * tstep*tstep &
+
+ tgnum_vpar * 0.25d0 * v * Vpar0**2 * BB2 * (1.d0 - fact_conservative_u) &
* (-(psi_s * vpar0_t - psi_t * vpar0_s)/xjac) / BigR &
- * (-(ps0_s * r0_t - ps0_t * r0_s) /xjac) * xjac * theta * tstep*tstep &
+ * (-(ps0_s * r0_t - ps0_t * r0_s) /xjac) * xjac * theta * tstep*tstep &
+ ! d/dpsi of the unvaried Bgrad_vpar bracket times the F0/BigR*r0_p piece of Bgrad_rho
+ + tgnum_vpar * 0.25d0 * v * Vpar0**2 * BB2 * (1.d0 - fact_conservative_u) &
+ * (-(psi_s * vpar0_t - psi_t * vpar0_s)/xjac) / BigR &
+ * (F0 / BigR * r0_p) * xjac * theta * tstep*tstep &
!=============================== New TG_num terms==================================
+ tgnum_vpar * 0.25d0 * vpar0 * Vpar0**2 * BB2 * fact_conservative_u &
* (-(ps0_s * r0_t - ps0_t * r0_s)/xjac) / BigR &
* (-(psi_s * v_t - psi_t * v_s) /xjac) * xjac * theta * tstep*tstep &
+ ! d/dpsi of the unvaried v-bracket times the F0/BigR*r0_p piece of Bgrad_rho
+ + tgnum_vpar * 0.25d0 * vpar0 * Vpar0**2 * BB2 * fact_conservative_u &
+ * (F0 / BigR * r0_p) / BigR &
+ * (-(psi_s * v_t - psi_t * v_s) /xjac) * xjac * theta * tstep*tstep &
+ tgnum_vpar * 0.25d0 * vpar0 * Vpar0**2 * BB2 * fact_conservative_u &
* (-(psi_s * r0_t - psi_t * r0_s)/xjac) / BigR &
- * (-(ps0_s * v_t - ps0_t * v_s) /xjac) * xjac * theta * tstep*tstep
+ * (-(ps0_s * v_t - ps0_t * v_s) /xjac) * xjac * theta * tstep*tstep &
+
+ ! BB2 -> BB2_psi copies of the three tgnum_vpar residual terms (see finding 7)
+ + tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2_psi &
+ * (-(ps0_s * vpar0_t - ps0_t * vpar0_s)/xjac + F0 / BigR * vpar0_p) / BigR &
+ * (-(ps0_s * v_t - ps0_t * v_s) /xjac) * xjac * theta * tstep*tstep &
+ + tgnum_vpar * 0.25d0 * v * Vpar0**2 * BB2_psi * (1.d0 - fact_conservative_u) &
+ * (-(ps0_s * vpar0_t - ps0_t * vpar0_s)/xjac + F0 / BigR * vpar0_p) / BigR &
+ * (-(ps0_s * r0_t - ps0_t * r0_s) /xjac + F0 / BigR * r0_p) * xjac * theta * tstep*tstep &
+ + tgnum_vpar * 0.25d0 * vpar0 * Vpar0**2 * BB2_psi * fact_conservative_u &
+ * (-(ps0_s * r0_t - ps0_t * r0_s)/xjac + F0 / BigR * r0_p) / BigR &
+ * (-(ps0_s * v_t - ps0_t * v_s) /xjac) * xjac * theta * tstep*tstep
!===============================End of new TG_num terms============================
if (normalized_velocity_profile) then
@@ -2946,7 +2978,23 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
amat_k(var_vpar,var_psi) = - 0.5d0 * r0 * vpar0**2 * BB2_psi * F0 / BigR * v_p * xjac * theta * tstep &
- visco_par_par * F0**2 / (BigR * BB2**2) * BB2_psi * Bgrad_vpar * Bgrad_rho_k_star * xjac * theta * tstep &
- + visco_par_par * F0**2 / (BigR * BB2) * Bgrad_vpar_psi * Bgrad_rho_k_star * xjac * theta * tstep
+ + visco_par_par * F0**2 / (BigR * BB2) * Bgrad_vpar_psi * Bgrad_rho_k_star * xjac * theta * tstep &
+ ! BB2 -> BB2_psi copies of the two rhs_ij_k tgnum_vpar terms (see finding 7)
+ + tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2_psi &
+ * (-(ps0_s * vpar0_t - ps0_t * vpar0_s)/xjac + F0 / BigR * vpar0_p) / BigR &
+ * (F0 / BigR * v_p) * xjac * theta * tstep*tstep &
+ + tgnum_vpar * 0.25d0 * vpar0 * Vpar0**2 * BB2_psi * fact_conservative_u &
+ * (-(ps0_s * r0_t - ps0_t * r0_s)/xjac + F0 / BigR * r0_p) / BigR &
+ * (F0 / BigR * v_p) * xjac * theta * tstep*tstep &
+ ! Bracket differentiation of the two rhs_ij_k tgnum_vpar terms (see finding 7):
+ ! the k-channel v-factor F0/BigR*v_p has no ps0 dependence, so only
+ ! d(Bgrad_vpar)/dpsi (resp. d(Bgrad_rho)/dpsi) times it is needed.
+ + tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2 &
+ * (-(psi_s * vpar0_t - psi_t * vpar0_s)/xjac) / BigR &
+ * (F0 / BigR * v_p) * xjac * theta * tstep*tstep &
+ + tgnum_vpar * 0.25d0 * vpar0 * Vpar0**2 * BB2 * fact_conservative_u &
+ * (-(psi_s * r0_t - psi_t * r0_s)/xjac) / BigR &
+ * (F0 / BigR * v_p) * xjac * theta * tstep*tstep
amat(var_vpar,var_u) = 0.d0
!---------------------------------------- NEO
@@ -2977,7 +3025,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
! New terms coming from -(\partial_t \rho + \nabla \cdot (\rho \mathbf{v})) \mathbf{v} in RHS of momentum equation
! (see wiki: https://www.jorek.eu/wiki/doku.php?id=model500_501_555#equations):
+ fact_conservative_u * ( &
- + v * rho * vpar0 * F0**2 / BigR * xjac * (1.d0 + zeta) &
+ + v * rho * vpar0 * BB2 * BigR * xjac * (1.d0 + zeta) &
- v * (rho_x_hat * u0_y - rho_y_hat * u0_x) * vpar0 * BB2 * theta * xjac * tstep &
+ v * F0 / BigR * rho * vpar0_p * vpar0 * BB2 * theta * xjac * tstep &
+ v * rho * (vpar0_x * ps0_y - vpar0_y * ps0_x) * vpar0 * BB2 * theta * xjac * tstep &
@@ -3083,7 +3131,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
end if ! (with_TiTe) ***********************************************************
- amat(var_vpar,var_vpar) = v * Vpar * r0_corr * F0**2 / BigR * xjac * (1.d0 + zeta) &
+ amat(var_vpar,var_vpar) = v * Vpar * r0_corr * BB2 * BigR * xjac * (1.d0 + zeta) &
! New terms coming from -(\partial_t \rho + \nabla \cdot (\rho \mathbf{v})) \mathbf{v} in RHS of momentum equation
! (see wiki: https://www.jorek.eu/wiki/doku.php?id=model500_501_555#equations):
@@ -3179,12 +3227,16 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* (-(ps0_s * v_t - ps0_t * v_s) /xjac ) * xjac * theta * tstep*tstep &
+ tgnum_vpar * 0.25d0 * v * Vpar0**2 * BB2 * (1.d0 - fact_conservative_u) &
* ( + F0 / BigR * vpar_p) / BigR &
- * (-(ps0_s * r0_t - ps0_t * r0_s) /xjac + F0 / BigR * r0_p) * xjac * theta * tstep*tstep
-
+ * (-(ps0_s * r0_t - ps0_t * r0_s) /xjac + F0 / BigR * r0_p) * xjac * theta * tstep*tstep &
+
+ + visco_par_par * F0**2 / (BigR * BB2) * Bgrad_vpar_vpar_n * Bgrad_rho_star * xjac * theta * tstep
+
amat_kn(var_vpar,var_vpar) = &
+ tgnum_vpar * 0.25d0 * r0 * Vpar0**2 * BB2 &
* ( + F0 / BigR * vpar_p) / BigR &
- * ( + F0 / BigR * v_p) * xjac * theta * tstep*tstep
+ * ( + F0 / BigR * v_p) * xjac * theta * tstep*tstep &
+
+ + visco_par_par * F0**2 / (BigR * BB2) * Bgrad_vpar_vpar_n * Bgrad_rho_k_star * xjac * theta * tstep
if ( NEO ) then
amat(var_vpar,var_rho) = amat(var_vpar,var_rho) - v * amu_neo_prof(ms,mt)*BB2/(Btheta2+epsil) &
@@ -3274,6 +3326,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- v * ((GAMMA - 1.) / BigR) * vpar0**2 * (psi_x * ps0_x + psi_y * ps0_y)&
* (particle_source(ms,mt) + source_pellet + source_bg_drift + source_imp_drift) * xjac * theta * tstep &
!==============================End of friction terms=================
+ - v * ((GAMMA - 1.) / BigR) * vpar0**2 * (psi_x * ps0_x + psi_y * ps0_y) * aux_rho0 * xjac * theta * tstep &
+ tgnum_Ti* 0.25d0 / BigR * vpar0**2 &
* Ti0 * ((r0_x+alpha_i*rimp0_x) * psi_y - (r0_y+alpha_i*rimp0_y) * psi_x) &
@@ -3315,13 +3368,13 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ (GAMMA-1.) * v * 2.d0 * BigR * w0 * u_x * visco_T_heating * visco_fact_new * BigR * xjac * theta * tstep &
+ (GAMMA-1.) * v * (u_x * u0_xpp + u_y * u0_ypp) * visco_T_heating * visco_fact_new * BigR * xjac * theta * tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * Ti0* ((r0_x+alpha_i*rimp0_x) * u_y - (r0_y+alpha_i*rimp0_y) * u_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * Ti0* ((r0_x+alpha_i*rimp0_x) * u_y - (r0_y+alpha_i*rimp0_y) * u_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * (r0+alpha_i*rimp0) * (Ti0_x * u_y - Ti0_y * u_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * (r0+alpha_i*rimp0) * (Ti0_x * u_y - Ti0_y * u_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * Ti0* ((r0_x+alpha_i*rimp0_x)*u0_y - (r0_y+alpha_i*rimp0_y)*u0_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * Ti0* ((r0_x+alpha_i*rimp0_x)*u0_y - (r0_y+alpha_i*rimp0_y)*u0_x) &
* ( v_x * u_y - v_y * u_x) * xjac * theta*tstep*tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * (r0+alpha_i*rimp0) * (Ti0_x * u0_y - Ti0_y * u0_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * (r0+alpha_i*rimp0) * (Ti0_x * u0_y - Ti0_y * u0_x) &
* ( v_x * u_y - v_y * u_x) * xjac * theta*tstep*tstep
amat_nn(var_Ti,var_u) = (GAMMA-1.) * v * visco_T_heating * (u0_x * u_xpp + u0_y * u_ypp) * visco_fact_new * BigR * xjac * theta * tstep
@@ -3353,9 +3406,9 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ 2.d0 * v * BigR * rho * r0_corr * Srec_T * Ti0 * xjac * theta * tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * Ti0* (rho_x * u0_y - rho_y * u0_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * Ti0* (rho_x * u0_y - rho_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * rho * (Ti0_x * u0_y - Ti0_y * u0_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * rho * (Ti0_x * u0_y - Ti0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac* theta*tstep*tstep &
+ tgnum_Ti* 0.25d0 / BigR * vpar0**2 &
* Ti0 * (rho_x * ps0_y - rho_y * ps0_x ) &
@@ -3379,7 +3432,10 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+ tgnum_Ti* 0.25d0 / BigR * vpar0**2 &
* rho * (Ti0_x * ps0_y - Ti0_y * ps0_x + F0 / BigR * Ti0_p) &
- * ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep
+ * ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+
+ - dZKi_prof_drho * rho * BigR / BB2 * Bgrad_T_k_star * Bgrad_Ti * xjac * theta * tstep &
+ + dZKi_prof_drho * rho * BigR * (v_p * Ti0_p / BigR**2) * xjac * theta * tstep
amat_kn(var_Ti,var_rho) = + tgnum_Ti* 0.25d0 / BigR * vpar0**2 &
* Ti0 * (+ F0 / BigR * rho_p) &
@@ -3418,9 +3474,9 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ v * bigR * Ti * r0_corr * r0_corr * Srec_T * xjac * theta * tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * Ti* ((r0_x+alpha_i*rimp0_x)*u0_y - (r0_y+alpha_i*rimp0_y)*u0_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * Ti* ((r0_x+alpha_i*rimp0_x)*u0_y - (r0_y+alpha_i*rimp0_y)*u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_Ti* 0.25d0 * BigR**2 * (r0+alpha_i*rimp0) * (Ti_x * u0_y - Ti_y * u0_x) &
+ + tgnum_Ti* 0.25d0 * BigR**3 * (r0+alpha_i*rimp0) * (Ti_x * u0_y - Ti_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
+ tgnum_Ti* 0.25d0 / BigR * vpar0**2 &
* Ti * ((r0_x+alpha_i*rimp0_x)*ps0_y - (r0_y+alpha_i*rimp0_y)*ps0_x + F0 / BigR * (r0_p+alpha_i*rimp0_p)) &
@@ -3512,8 +3568,9 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- v * BigR * ((GAMMA - 1.)/2.) * vpar0**2 * BB2 &
* (rimp0*dalpha_e_dT*rn0*Sion_T) * Te * xjac * theta * tstep &
- v * BigR * ((GAMMA - 1.)/2.) * vv2 &
- * (rimp0*dalpha_e_dT*rn0*Sion_T) * Te * xjac * theta * tstep
+ * (rimp0*dalpha_e_dT*rn0*Sion_T) * Te * xjac * theta * tstep &
!==============================End of friction terms=================
+ + v * BigR * Ti0 * r0_corr * r0_corr * dSrec_dT * Te * xjac * theta * tstep
if (with_neutrals) then
!===================== Additional terms from friction terms============
@@ -3534,10 +3591,10 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ v * alpha_i * rhoimp * GAMMA * Ti0 * (vpar0_s * ps0_t - vpar0_t * ps0_s) * theta * tstep &
+ v * alpha_i * rhoimp * GAMMA * Ti0 * F0 / BigR * vpar0_p * xjac * theta * tstep &
!=========================New TG_num terms====================================
- + tgnum_Ti * 0.25d0 * BigR**2 * Ti0 * alpha_i * (rhoimp_x * u0_y - rhoimp_y * u0_x) &
+ + tgnum_Ti * 0.25d0 * BigR**3 * Ti0 * alpha_i * (rhoimp_x * u0_y - rhoimp_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Ti * 0.25d0 * BigR**2 * alpha_i * rhoimp * (Ti0_x * u0_y - Ti0_y * u0_x) &
+ + tgnum_Ti * 0.25d0 * BigR**3 * alpha_i * rhoimp * (Ti0_x * u0_y - Ti0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac* theta*tstep*tstep &
+ tgnum_Ti * 0.25d0 / BigR * vpar0**2 &
@@ -3652,13 +3709,13 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- (GAMMA-1.) * v * E_ion_bg * (r0-rimp0) * 2.d0 * BigR * u_y * xjac * theta * tstep &
!================= End ionization potential energy ===========================
- + tgnum_Te * 0.25d0 * BigR**2 * Te0* ((r0_x+alpha_e*rimp0_x) * u_y - (r0_y+alpha_e*rimp0_y) * u_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * Te0* ((r0_x+alpha_e*rimp0_x) * u_y - (r0_y+alpha_e*rimp0_y) * u_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * (r0+alpha_e_bis*rimp0) * (Te0_x * u_y - Te0_y * u_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * (r0+alpha_e_bis*rimp0) * (Te0_x * u_y - Te0_y * u_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * Te0* ((r0_x+alpha_e*rimp0_x) * u0_y - (r0_y+alpha_e*rimp0_y) * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * Te0* ((r0_x+alpha_e*rimp0_x) * u0_y - (r0_y+alpha_e*rimp0_y) * u0_x) &
* ( v_x * u_y - v_y * u_x) * xjac * theta*tstep*tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * (r0+alpha_e_bis*rimp0) * (Te0_x * u0_y - Te0_y * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * (r0+alpha_e_bis*rimp0) * (Te0_x * u0_y - Te0_y * u0_x) &
* ( v_x * u_y - v_y * u_x) * xjac * theta*tstep*tstep
amat(var_Te,var_zj) = - v * (gamma-1.d0) * eta_T_ohm * 2.d0 * zj * (zj0 - aux_jre) /(BigR**2.d0) * BigR * xjac * theta * tstep !> aux_jre from kinetic REs
@@ -3701,9 +3758,9 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ (GAMMA - 1.) * E_ion_bg * D_prof * BigR * (v_x*rho_x + v_y*rho_y ) * xjac * theta * tstep &
!================= End ionization potential energy ===========================
- + tgnum_Te * 0.25d0 * BigR**2 * Te0* (rho_x * u0_y - rho_y * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * Te0* (rho_x * u0_y - rho_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * rho * (Te0_x * u0_y - Te0_y * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * rho * (Te0_x * u0_y - Te0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac* theta*tstep*tstep &
+ tgnum_Te * 0.25d0 / BigR * vpar0**2 &
* Te0 * (rho_x * ps0_y - rho_y * ps0_x ) &
@@ -3739,7 +3796,10 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+ tgnum_Te * 0.25d0 / BigR * vpar0**2 &
* rho * (Te0_x * ps0_y - Te0_y * ps0_x + F0 / BigR * Te0_p) &
- * ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep
+ * ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+
+ - dZKe_prof_drho * rho * BigR / BB2 * Bgrad_T_k_star * Bgrad_Te * xjac * theta * tstep &
+ + dZKe_prof_drho * rho * BigR * (v_p * Te0_p / BigR**2) * xjac * theta * tstep
amat_kn(var_Te,var_rho) = &
!=============== The ionization potential energy term=========================
@@ -3763,9 +3823,16 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ (GAMMA-1.) * v * rimp0 * dE_ion_dT * Vpar0 * (Te_s*ps0_t - Te_t*ps0_s) * theta * tstep &
! New diffusive ionization energy flux term
- + (GAMMA - 1.) * dE_ion_dT * Te * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_star * (Bgrad_rhoimp) * xjac * tstep &
- + (GAMMA - 1.) * dE_ion_dT * Te * D_prof_imp * BigR * (v_x*(rimp0_x) + v_y*(rimp0_y) ) * xjac * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * Te * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_star * (Bgrad_rhoimp) * xjac * theta * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * Te * D_prof_imp * BigR * (v_x*(rimp0_x) + v_y*(rimp0_y) ) * xjac * theta * tstep &
!================= End ionization potential energy ===========================
+ ! E_ion -> dE_ion_dT*Te copies of the compression and parallel-flow terms
+ - (GAMMA-1.) * v * dE_ion_dT * Te * BigR**2 * (rimp0_s * u0_t - rimp0_t * u0_s) * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * Te * F0 / BigR * Vpar0 * rimp0_p * xjac * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * Te * Vpar0 * (rimp0_s * ps0_t - rimp0_t * ps0_s) * theta * tstep &
+ - (GAMMA-1.) * v * dE_ion_dT * Te * rimp0 * 2.d0 * BigR * u0_y * xjac * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * Te * rimp0 * (vpar0_s * ps0_t - vpar0_t * ps0_s) * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * Te * rimp0 * F0 / BigR * vpar0_p * xjac * theta * tstep &
- v * (r0 + rimp0 * alpha_e_bis) * BigR**2 * (Te_s * u0_t - Te_t * u0_s) * theta * tstep &
- v * (rimp0 * alpha_e_tri) * Te * BigR**2 * (Te0_s* u0_t - Te0_t * u0_s) * theta * tstep &
- v * Te * BigR**2 * ( r0_s * u0_t - r0_t * u0_s) * theta * tstep &
@@ -3809,15 +3876,15 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ v * BigR * Te * dalpha_e_dT * rimp0_corr**2 * Lrad * xjac * theta * tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * Te* ((r0_x+alpha_e_bis*rimp0_x)*u0_y &
+ + tgnum_Te * 0.25d0 * BigR**3 * Te* ((r0_x+alpha_e_bis*rimp0_x)*u0_y &
- (r0_y+alpha_e_bis*rimp0_y)*u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * (r0+alpha_e_bis*rimp0) * (Te_x * u0_y - Te_y * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * (r0+alpha_e_bis*rimp0) * (Te_x * u0_y - Te_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * (alpha_e_tri*rimp0)*Te* (Te0_x* u0_y - Te0_y* u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * (alpha_e_tri*rimp0)*Te* (Te0_x* u0_y - Te0_y* u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
+ tgnum_Te * 0.25d0 / BigR * vpar0**2 &
- * Te * ((r0_x+alpha_e_bis*rimp0_x)*ps0_y - (r0_y+alpha_e_bis*rimp0_y)*ps0_x + F0 / BigR * (r0_p+alpha_e*rimp0_p)) &
+ * Te * ((r0_x+alpha_e_bis*rimp0_x)*ps0_y - (r0_y+alpha_e_bis*rimp0_y)*ps0_x + F0 / BigR * (r0_p+alpha_e_bis*rimp0_p)) &
* ( v_x * ps0_y - v_y * ps0_x ) * xjac * theta * tstep * tstep &
+ tgnum_Te * 0.25d0 / BigR * vpar0**2 &
* (r0+alpha_e_bis*rimp0) * (Te_x * ps0_y - Te_y * ps0_x ) &
@@ -3831,8 +3898,8 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
!=============== The ionization potential energy term=========================
! New diffusive ionization energy flux term
- + (GAMMA - 1.) * dE_ion_dT * Te * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_k_star * (Bgrad_rhoimp) * xjac * tstep &
- + (GAMMA - 1.) * dE_ion_dT * Te * D_prof_imp * BigR * ( + v_p*(rimp0_p) /BigR**2 ) * xjac * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * Te * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_k_star * (Bgrad_rhoimp) * xjac * theta * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * Te * D_prof_imp * BigR * ( + v_p*(rimp0_p) /BigR**2 ) * xjac * theta * tstep &
!================= End ionization potential energy ===========================
@@ -3904,7 +3971,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
end if ! (with_vpar)
if (with_neutrals) then
amat(var_Te,var_rhon) = + v * BigR * (r0 + rimp0 * alpha_e) * rhon * ksi_ion_norm * Sion_T * xjac * theta * tstep &
- + v * BigR * rhon * (r0 + rimp0 * alpha_e) * LradDrays_T * xjac * theta * tstep
+ + v * BigR * rhon * (r0_corr + rimp0_corr * alpha_e) * LradDrays_T * xjac * theta * tstep
endif
if (with_impurities) then
amat(var_Te,var_rhoimp) = v * rhoimp * alpha_e * Te0 * BigR * xjac * (1.d0 + zeta) &
@@ -3930,10 +3997,10 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- (GAMMA - 1.) * E_ion_bg * D_prof * BigR * (v_x*rhoimp_x + v_y*rhoimp_y ) * xjac * theta * tstep &
!================= End ionization potential energy ===========================
!=========================New TG_num terms====================================
- + tgnum_Te * 0.25d0 * BigR**2 * Te0 * alpha_e * (rhoimp_x * u0_y - rhoimp_y * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * Te0 * alpha_e * (rhoimp_x * u0_y - rhoimp_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_Te * 0.25d0 * BigR**2 * alpha_e_bis * rhoimp * (Te0_x * u0_y - Te0_y * u0_x) &
+ + tgnum_Te * 0.25d0 * BigR**3 * alpha_e_bis * rhoimp * (Te0_x * u0_y - Te0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac* theta*tstep*tstep &
+ tgnum_Te * 0.25d0 / BigR * vpar0**2 &
@@ -3959,7 +4026,13 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
! Energy exchange term
- v * BigR * ddTe_i_drhoimp * rhoimp * xjac * theta * tstep &
-
+
+ ! Ionization sink and line/continuum radiation, alpha_e dependence on rhoimp
+ + v * BigR * ksi_ion_norm * alpha_e * rn0 * Sion_T * rhoimp * xjac * theta * tstep &
+ + v * BigR * alpha_e * rn0_corr * LradDrays_T * rhoimp * xjac * theta * tstep &
+ + v * BigR * alpha_e * (r0_corr-rimp0_corr) * LradDcont_corr * rhoimp * xjac * theta * tstep &
+ - v * BigR * (r0_corr+alpha_e*rimp0_corr) * LradDcont_corr * rhoimp * xjac * theta * tstep &
+
+ v * BigR * rhoimp * (r0_corr + 2.*alpha_e*rimp0_corr) * Lrad * xjac * theta * tstep &
+ v * BigR * rhoimp * alpha_e * frad_bg * xjac * theta * tstep
@@ -4039,6 +4112,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- v * ((GAMMA - 1.) / BigR) * vpar0**2 * (psi_x * ps0_x + psi_y * ps0_y)&
* ((r0+alpha_e*rimp0)*rn0*Sion_T + particle_source(ms,mt) + source_pellet + source_bg_drift + source_imp_drift) * xjac * theta * tstep &
!==============================End of friction terms=================
+ - v * ((GAMMA - 1.) / BigR) * vpar0**2 * (psi_x * ps0_x + psi_y * ps0_y) * aux_rho0 * xjac * theta * tstep &
+ tgnum_T * 0.25d0 / BigR * vpar0**2 &
* T0 * ((r0_x+alpha_imp*rimp0_x) * psi_y - (r0_y+alpha_imp*rimp0_y) * psi_x) &
@@ -4068,7 +4142,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* T0 * ((r0_x+alpha_imp*rimp0_x) * psi_y - (r0_y+alpha_imp*rimp0_y) * psi_x) &
* ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+ tgnum_T * 0.25d0 / BigR * vpar0**2 &
- * (r0+alpha_imp*rimp0) * (T0_x * psi_y - T0_y * psi_x) &
+ * (r0+alpha_imp_bis*rimp0) * (T0_x * psi_y - T0_y * psi_x) &
* ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep
@@ -4090,13 +4164,13 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ (GAMMA-1.) * v * 2.d0 * BigR * w0 * u_x * visco_T_heating * visco_fact_new * BigR * xjac * theta * tstep &
+ (GAMMA-1.) * v * (u_x * u0_xpp + u_y * u0_ypp) * visco_T_heating * visco_fact_new * BigR * xjac * theta * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * T0* ((r0_x+alpha_imp*rimp0_x) * u_y - (r0_y+alpha_imp*rimp0_y) * u_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * T0* ((r0_x+alpha_imp*rimp0_x) * u_y - (r0_y+alpha_imp*rimp0_y) * u_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * (r0+alpha_imp_bis*rimp0) * (T0_x * u_y - T0_y * u_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * (r0+alpha_imp_bis*rimp0) * (T0_x * u_y - T0_y * u_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * T0* ((r0_x+alpha_imp*rimp0_x)*u0_y - (r0_y+alpha_imp*rimp0_y)*u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * T0* ((r0_x+alpha_imp*rimp0_x)*u0_y - (r0_y+alpha_imp*rimp0_y)*u0_x) &
* ( v_x * u_y - v_y * u_x) * xjac * theta * tstep * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * (r0+alpha_imp_bis*rimp0) * (T0_x * u0_y - T0_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * (r0+alpha_imp_bis*rimp0) * (T0_x * u0_y - T0_y * u0_x) &
* ( v_x * u_y - v_y * u_x) * xjac * theta * tstep * tstep
amat_nn(var_T,var_u)= (GAMMA-1.) * v * visco_T_heating * (u0_x * u_xpp + u0_y * u_ypp) * visco_fact_new * BigR * xjac * theta * tstep
@@ -4148,9 +4222,9 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
- v * BigR * ((GAMMA - 1.)/2.) * vv2 * (rho*rn0*Sion_T) * xjac * theta * tstep &
!==============================End of friction terms=================
- + tgnum_T * 0.25d0 * BigR**2 * T0* (rho_x * u0_y - rho_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * T0* (rho_x * u0_y - rho_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * rho * (T0_x * u0_y - T0_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * rho * (T0_x * u0_y - T0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
+ tgnum_T * 0.25d0 / BigR * vpar0**2 &
* T0 * (rho_x * ps0_y - rho_y * ps0_x ) &
@@ -4186,8 +4260,11 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+ tgnum_T * 0.25d0 / BigR * vpar0**2 &
* rho * (T0_x * ps0_y - T0_y * ps0_x + F0 / BigR * T0_p) &
- * ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep
-
+ * ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep &
+
+ - dZK_prof_drho * rho * BigR / BB2 * Bgrad_T_k_star * Bgrad_T * xjac * theta * tstep &
+ + dZK_prof_drho * rho * BigR * (v_p * T0_p / BigR**2) * xjac * theta * tstep
+
amat_kn(var_T,var_rho) = &
!=============== The ionization potential energy term=========================
! New diffusive ionization energy flux term
@@ -4209,10 +4286,17 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ (GAMMA - 1.) * v * rimp0 * dE_ion_dT * Vpar0 * (T_s*ps0_t - T_t*ps0_s) * theta * tstep &
! New diffusive ionization energy flux term
- + (GAMMA - 1.) * dE_ion_dT * T * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_star * (Bgrad_rhoimp) * xjac * tstep &
- + (GAMMA - 1.) * dE_ion_dT * T * D_prof_imp * BigR * (v_x*(rimp0_x) + v_y*(rimp0_y) ) * xjac * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * T * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_star * (Bgrad_rhoimp) * xjac * theta * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * T * D_prof_imp * BigR * (v_x*(rimp0_x) + v_y*(rimp0_y) ) * xjac * theta * tstep &
!================= End ionization potential energy ===========================
+ ! E_ion -> dE_ion_dT*T copies of the compression and parallel-flow terms
+ - (GAMMA-1.) * v * dE_ion_dT * T * BigR**2 * (rimp0_s * u0_t - rimp0_t * u0_s) * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * T * F0 / BigR * Vpar0 * rimp0_p * xjac * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * T * Vpar0 * (rimp0_s * ps0_t - rimp0_t * ps0_s) * theta * tstep &
+ - (GAMMA-1.) * v * dE_ion_dT * T * rimp0 * 2.d0 * BigR * u0_y * xjac * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * T * rimp0 * (vpar0_s * ps0_t - vpar0_t * ps0_s) * theta * tstep &
+ + (GAMMA-1.) * v * dE_ion_dT * T * rimp0 * F0 / BigR * vpar0_p * xjac * theta * tstep &
- v * (r0 + rimp0 * alpha_imp_bis) * BigR**2 * ( T_s * u0_t - T_t * u0_s) * theta * tstep &
- v * (rimp0 * alpha_imp_tri) * T * BigR**2 * ( T0_s * u0_t - T0_t * u0_s) * theta * tstep &
- v * T * BigR**2 * ( r0_s * u0_t - r0_t * u0_s) * theta * tstep &
@@ -4257,17 +4341,22 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* ((r0+alpha_e*rimp0)*rn0*dSion_dT) * T * xjac * theta * tstep &
- v * BigR * ((GAMMA - 1.)/2.) * vv2 &
* ((r0+alpha_e*rimp0)*rn0*dSion_dT) * T * xjac * theta * tstep &
+ ! alpha_e temperature dependence of the friction source
+ - v * BigR * ((GAMMA - 1.)/2.) * vpar0**2 * BB2 &
+ * (dalpha_e_dT*rimp0*rn0*Sion_T) * T * xjac * theta * tstep &
+ - v * BigR * ((GAMMA - 1.)/2.) * vv2 &
+ * (dalpha_e_dT*rimp0*rn0*Sion_T) * T * xjac * theta * tstep &
!==============================End of friction terms=================
+ (GAMMA-1.) * v * BigR**2.d0 * ( u0_x * w0_x + u0_y * w0_y) * dvisco_dT_heating * T * visco_fact_old * BigR * xjac * theta * tstep &
+ (GAMMA-1.) * v * 2.d0 * BigR * w0 * u0_x * dvisco_dT_heating * T * visco_fact_new * BigR * xjac * theta * tstep &
+ (GAMMA-1.) * v * (u0_x * u0_xpp + u0_y * u0_ypp) * dvisco_dT_heating * T * visco_fact_new * BigR * xjac * theta * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * T* ((r0_x+alpha_imp_bis*rimp0_x) * u0_y &
+ + tgnum_T * 0.25d0 * BigR**3 * T* ((r0_x+alpha_imp_bis*rimp0_x) * u0_y &
- (r0_y+alpha_imp_bis*rimp0_y) * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * (r0+alpha_imp_bis*rimp0) * (T_x * u0_y - T_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * (r0+alpha_imp_bis*rimp0) * (T_x * u0_y - T_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
- + tgnum_T * 0.25d0 * BigR**2 * (alpha_imp_tri*rimp0)*T* (T0_x * u0_y - T0_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * (alpha_imp_tri*rimp0)*T* (T0_x * u0_y - T0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta * tstep * tstep &
+ tgnum_T * 0.25d0 / BigR * vpar0**2 &
* T * ((r0_x+alpha_imp_bis*rimp0_x) * ps0_y - (r0_y+alpha_imp_bis*rimp0_y) * ps0_x + F0 / BigR * (r0_p+alpha_imp_bis*rimp0_p)) &
@@ -4283,8 +4372,8 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ dZK_par_dT * T * BigR / BB2 * Bgrad_T_k_star * Bgrad_T * xjac * theta * tstep &
!=============== The ionization potential energy term=========================
! New diffusive ionization energy flux term
- + (GAMMA - 1.) * dE_ion_dT * T * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_k_star * (Bgrad_rhoimp) * xjac * tstep &
- + (GAMMA - 1.) * dE_ion_dT * T * D_prof_imp * BigR * ( + v_p*(rimp0_p) /BigR**2 ) * xjac * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * T * ((D_par_local_imp+D_par_imp_sc_num*tau_sc)-D_prof_imp) * BigR / BB2 * Bgrad_rho_k_star * (Bgrad_rhoimp) * xjac * theta * tstep &
+ + (GAMMA - 1.) * dE_ion_dT * T * D_prof_imp * BigR * ( + v_p*(rimp0_p) /BigR**2 ) * xjac * theta * tstep &
!================= End ionization potential energy ===========================
@@ -4369,7 +4458,7 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
* (r0+alpha_imp_bis*rimp0) * (T0_x * ps0_y - T0_y * ps0_x + F0 / BigR * T0_p) &
* ( + F0 / BigR * v_p) * xjac * theta * tstep * tstep
- amat_n(var_T,var_vpar) = + v * r0 * GAMMA * T0 * F0 / BigR * vpar_p * xjac * theta * tstep &
+ amat_n(var_T,var_vpar) = + v * (r0 + rimp0*alpha_imp) * GAMMA * T0 * F0 / BigR * vpar_p * xjac * theta * tstep &
!=============== The ionization potential energy term=========================
+ (GAMMA - 1.) * v * E_ion * rimp0 * F0 / BigR * vpar_p * xjac * theta * tstep &
+ (GAMMA - 1.) * v * E_ion_bg * (r0-rimp0) * F0 / BigR * vpar_p * xjac * theta * tstep
@@ -4410,10 +4499,10 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
!================= End ionization potential energy ===========================
!=========================New TG_num terms====================================
- + tgnum_T * 0.25d0 * BigR**2 * T0 * alpha_imp * (rhoimp_x * u0_y - rhoimp_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * T0 * alpha_imp * (rhoimp_x * u0_y - rhoimp_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac * theta*tstep*tstep &
- + tgnum_T * 0.25d0 * BigR**2 * alpha_imp_bis * rhoimp * (T0_x * u0_y - T0_y * u0_x) &
+ + tgnum_T * 0.25d0 * BigR**3 * alpha_imp_bis * rhoimp * (T0_x * u0_y - T0_y * u0_x) &
* ( v_x * u0_y - v_y * u0_x) * xjac* theta*tstep*tstep &
+ tgnum_T * 0.25d0 / BigR * vpar0**2 &
@@ -4437,6 +4526,20 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ v * alpha_imp * rhoimp * GAMMA * T0 * (vpar0_s * ps0_t - vpar0_t * ps0_s) * theta * tstep &
+ v * alpha_imp * rhoimp * GAMMA * T0 * F0 / BigR * vpar0_p * xjac * theta * tstep &
+ !===================== Additional terms from friction terms============
+ ! alpha_e rhoimp dependence of the friction source
+ - v * BigR * ((GAMMA - 1.)/2.) * vpar0**2 * BB2 &
+ * (alpha_e*rn0*Sion_T) * rhoimp * xjac * theta * tstep &
+ - v * BigR * ((GAMMA - 1.)/2.) * vv2 &
+ * (alpha_e*rn0*Sion_T) * rhoimp * xjac * theta * tstep &
+ !==============================End of friction terms=================
+
+ ! Ionization sink and line/continuum radiation, alpha_e dependence on rhoim
+ + v * BigR * ksi_ion_norm * alpha_e * rn0 * Sion_T * rhoimp * xjac * theta * tstep &
+ + v * BigR * alpha_e * rn0_corr * LradDrays_T * rhoimp * xjac * theta * tstep &
+ + v * BigR * alpha_e * (r0_corr-rimp0_corr) * LradDcont_corr * rhoimp * xjac * theta * tstep &
+ - v * BigR * (r0_corr+alpha_e*rimp0_corr) * LradDcont_corr * rhoimp * xjac * theta * tstep &
+
+ v * BigR * rhoimp * (r0_corr + 2*alpha_e*rimp0_corr) * Lrad * xjac * theta * tstep &
+ v * BigR * rhoimp * alpha_e * frad_bg * xjac * theta * tstep
@@ -4496,8 +4599,8 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
+ v * Vpar0 * (rn0_s * psi_t - rn0_t * psi_s) * theta * tstep )
amat(var_rhon,var_u) = delta_n_convection*( &
- + v * BigR**2 * ( rn0_s * u_t - rn0_t * u_s) * theta * tstep &
- + v * 2.d0 * BigR * rn0 * u_y * xjac * theta * tstep )
+ - v * BigR**2 * ( rn0_s * u_t - rn0_t * u_s) * theta * tstep &
+ - v * 2.d0 * BigR * rn0 * u_y * xjac * theta * tstep )
amat(var_rhon,var_rho) = + BigR * v * rn0 * Sion_T * rho * xjac * theta * tstep &
- BigR * v * (2.d0*r0 +(alpha_e-1.)*rimp0) * rho * Srec_T * xjac * theta * tstep
@@ -4507,10 +4610,14 @@ subroutine element_matrix_fft(element, nodes, xpoint2, xcase2, R_axis, Z_axis, p
if (with_TiTe) then
amat(var_rhon,var_Te) = + BigR * v * (r0+alpha_e*rimp0) * rn0 * dSion_dT * Te * xjac * theta * tstep &
- - BigR * v * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * Te * xjac * theta * tstep
+ + BigR * v * dalpha_e_dT * rimp0 * rn0 * Sion_T * Te * xjac * theta * tstep &
+ - BigR * v * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * Te * xjac * theta * tstep &
+ - BigR * v * dalpha_e_dT * rimp0 * (r0-rimp0) * Srec_T * Te * xjac * theta * tstep
else
- amat(var_rhon,var_T) = + BigR * v * r0 * rn0 * dSion_dT * T * xjac * theta * tstep &
- - BigR * v * r0 * r0 * dSrec_dT * T * xjac * theta * tstep
+ amat(var_rhon,var_T) = + BigR * v * (r0+alpha_e*rimp0) * rn0 * dSion_dT * T * xjac * theta * tstep &
+ + BigR * v * dalpha_e_dT * rimp0 * rn0 * Sion_T * T * xjac * theta * tstep &
+ - BigR * v * (r0+alpha_e*rimp0) * (r0-rimp0) * dSrec_dT * T * xjac * theta * tstep &
+ - BigR * v * dalpha_e_dT * rimp0 * (r0-rimp0) * Srec_T * T * xjac * theta * tstep
endif ! with_TiTe
if (with_vpar) then
diff --git a/util/equation_codegen/.gitignore b/util/equation_codegen/.gitignore
new file mode 100644
index 0000000000..09731983db
--- /dev/null
+++ b/util/equation_codegen/.gitignore
@@ -0,0 +1,6 @@
+.venv/
+.tmp/
+__pycache__/
+*.egg-info/
+build/
+dist/
diff --git a/util/equation_codegen/JOREK_FINDINGS.md b/util/equation_codegen/JOREK_FINDINGS.md
new file mode 100644
index 0000000000..5f2085de6b
--- /dev/null
+++ b/util/equation_codegen/JOREK_FINDINGS.md
@@ -0,0 +1,230 @@
+# Model-600 Jacobian: missing and inconsistent terms
+
+This file records the differences between the volume terms assembled by
+`models/model600/mod_elt_matrix_fft.f90` and the automatically linearized
+weak forms in `src/jorek_equations/model600.py`.
+
+Every finding below was obtained by generating the two term reports and
+comparing them block by block:
+
+```bash
+python3 examples/export_model600_terms.py
+python3 examples/diff_model600_reports.py
+```
+
+`diff_model600_reports.py` compares each `rhs_ij`/`amat` assignment as a
+multiset of monomials and prints the algebraic residual, so a block that only
+reorders its terms is reported as identical. The residual is evaluated after
+rewriting every `_s`/`_t` derivative through the chain rule, because the
+element routine spells the same poloidal bracket sometimes as
+`a_s*b_t - a_t*b_s` and sometimes as `xjac*(a_x*b_y - a_y*b_x)` — even between
+a residual and its own tangent. A block whose two sides differ only by that
+rewrite is reported as `SAME`.
+
+The original audit reported seventeen findings. Findings 3, 5–8 and 10–17
+have since been fixed in the Fortran (one commit per finding, `fix
+linearization: finding N, ...`), and their sections have been removed from
+this file. The neutral density row `rhon` was added afterwards and brought
+findings 18 and 19, which have also since been fixed. What is left disagrees
+in 24 blocks, for two reasons:
+
+- **Open: the inward pinch (findings 4 and 9).** These are real
+ inconsistencies in the Jacobian, but the pinch implementation as a whole
+ needs revisiting, so they are left to its developer rather than patched
+ here. They are described below.
+- **Accepted: tauIC with impurities (findings 1 and 2).** These terms only
+ exist when the diamagnetic terms (`tauIC /= 0`) and impurities are both
+ active, and that combination is not supported. They are kept in the
+ reference as known differences, not as bugs to fix. See the section below.
+
+Status of the exported rows:
+
+| Row | Residual | Jacobian |
+|---|---|---|
+| `psi` | reproduced | finding 2 (tauIC × impurities only) |
+| `u` | reproduced | finding 1 (tauIC × impurities only) |
+| `zj` | reproduced | reproduced |
+| `w` | reproduced | reproduced |
+| `rho` | reproduced | findings 1 (tauIC × impurities only) and 4 |
+| `vpar` | reproduced | findings 4 and 9 |
+| `rhoimp` | reproduced | reproduced |
+| `rhon` | reproduced | reproduced |
+| `Ti` | reproduced | reproduced |
+| `Te` | reproduced | reproduced |
+| `T` | reproduced | reproduced |
+
+Every residual is now reproduced exactly, so none of the remaining differences
+changes the converged solution. They only affect the Newton tangent.
+
+The NEO branch is excluded from the default reports (`--include-neo` enables
+it) and has not been audited here.
+
+---
+
+
+## Finding 4 — the density inward pinch is not differentiated with respect to psi
+
+**Where:** the `var_rho` block, `amat(var_rho,var_psi)`.
+
+**What happens:** the density residual contains the weak form of
+`-div(r0 * V_pinch)` with `V_pinch = -V_prof_pinch * grad(psi)/|grad(psi)|`:
+
+```fortran
+- V_prof_pinch / sqrt(psi_grad2) * (v_x * ps0_x + v_y * ps0_y) * r0 * BigR * xjac * tstep * factor(var_rho,14)
+```
+
+with `psi_grad2 = ps0_x**2 + ps0_y**2`. The pinch direction is therefore a
+function of the evolved poloidal flux, but `amat(var_rho,var_psi)` contains no
+corresponding tangent. Every other flux-dependent quantity of the same
+equation *is* differentiated there — `BB2` through `BB2_psi`, the parallel
+gradients through `Bgrad_rho_star_psi`, `Bgrad_rho_psi` and
+`Bgrad_rhoimp_psi` — which is what makes the omission stand out.
+
+**Suggested fix:** the missing contribution is the transverse projection of
+the trial flux gradient,
+
+```fortran
++ V_prof_pinch / sqrt(psi_grad2) * r0 * BigR * xjac * theta * tstep &
+ * ( (v_x * psi_x + v_y * psi_y) &
+ - (ps0_x*v_x + ps0_y*v_y) * (ps0_x*psi_x + ps0_y*psi_y) / psi_grad2 )
+```
+
+(the first bracket from varying `grad(psi)` in the dot product, the second from
+varying `1/|grad(psi)|`).
+
+The parallel-velocity equation advects `v_par` with the same pinch velocity
+and has the same gap: `amat(var_vpar,var_psi)` carries no pinch tangent
+either.
+
+**Scale:** 6 monomials in `amat(var_rho,var_psi)` and 6 in
+`amat(var_vpar,var_psi)`, in each temperature branch.
+
+**Note:** `V_prof_pinch`, `D_prof`, `D_prof_imp`, `D_par_local`,
+`D_par_local_imp`, `tau_sc` and `D_perp_num_psin` are profile work values that
+also depend on the state through `psi_norm` (and, for `tau_sc`, through the
+pressure). The element routine freezes all of them, and the generator follows
+that convention, so they produce no report differences. The pinch term is
+different: its flux dependence is explicit in the assembled expression rather
+than hidden inside a work value.
+
+---
+
+## Finding 9 — the parallel-velocity inward-pinch tangent has the wrong sign
+
+**Where:** `amat(var_vpar,var_rho)` and `amat(var_vpar,var_vpar)`.
+
+**What happens:** the residual contribution is
+
+```fortran
+rhs_ij(var_vpar) = rhs_ij(var_vpar) &
+ + V_prof_pinch / sqrt(psi_grad2) * (ps0_x*vpar0_x + ps0_y*vpar0_y) &
+ * r0 * v * BigR * xjac * tstep * factor(var_vpar,13)
+```
+
+and its two tangents repeat that sign:
+
+```fortran
+amat(var_vpar,var_rho) = ... + V_prof_pinch / sqrt(psi_grad2) * (ps0_x*vpar0_x + ps0_y*vpar0_y) * rho * v * BigR * xjac * theta * tstep
+amat(var_vpar,var_vpar) = ... + V_prof_pinch / sqrt(psi_grad2) * (ps0_x*vpar_x + ps0_y*vpar_y ) * r0 * v * BigR * xjac * theta * tstep
+```
+
+The JOREK convention is `AMAT = -theta*tstep*dB`, so both should be negative.
+The density equation gets this right for the same pinch term: its residual
+contribution is negative and `amat(var_rho,var_rho)` is positive.
+
+**Suggested fix:** flip the sign of both `amat(var_vpar,...)` pinch terms.
+
+**Note:** the residual sign itself also deserves a look. With
+`V_pinch = -V_prof_pinch * grad(psi)/|grad(psi)|` as documented in the source
+comment, the advection `rho * V_pinch . grad(v_par)` is
+`-V_prof_pinch/|grad(psi)| * rho * grad(psi).grad(v_par)`, i.e. negative,
+while the assembled term is positive. The generator follows the assembled
+residual, so this note is a reading of the source comment rather than a
+report difference.
+
+**Scale:** 2 monomials in each of `amat(var_vpar,var_rho)` and
+`amat(var_vpar,var_vpar)`, in each temperature branch.
+
+---
+
+## Accepted differences: tauIC with impurities (findings 1 and 2)
+
+The diamagnetic terms scale with `tauIC`. Their tangents are written for the
+main-ion and electron pressures `r0*Ti0` and `r0*Te0`, while the residuals use
+the full pressures `(r0 + rimp0*alpha_i)*Ti0` and `(r0 + rimp0*alpha_e)*Te0`.
+The generator therefore produces impurity-pressure terms that the element
+routine does not have:
+
+- **Finding 1**: the impurity ion pressure in the diamagnetic terms of the
+ momentum and density equations: `amat(var_u,var_Ti|T|rhoimp)` and
+ `amat(var_rho,var_Ti|T|rhoimp)`.
+- **Finding 2**: the impurity electron pressure in the diamagnetic coupling
+ of the induction equation: `amat(var_psi,var_Te|T|rhoimp)`,
+ `amat_n(var_psi,var_Te|T)`, and an `amat_n(var_psi,var_rhoimp)` that the
+ element routine does not assemble at all.
+
+Every one of these monomials carries both `tauIC` and an impurity factor
+(`rhoimp`, `rimp0` or its derivatives). Model 600 does not support running
+with impurities and `tauIC /= 0` at the same time, so the missing terms never
+contribute in a supported configuration. They are recorded in the reference as
+accepted differences, not bugs. The combination should be rejected at input
+time, but nothing in model 600 does this yet. If it ever becomes supported,
+these tangents must be completed first; the full derivation is in the history
+of this file (commit `f4f8a4e98`).
+
+---
+
+## Complete inventory
+
+Every report line that is blank on one side belongs to one of the four
+findings above. There is nothing unexplained left.
+
+| Finding | Report lines | Blocks (each in both temperature branches) |
+|---:|---:|---|
+| 1 — impurity ion pressure, tauIC terms (accepted) | 96 | `amat(var_u,var_rhoimp)` 44, `amat(var_u,var_t)` 22, `amat(var_u,var_ti)` 22, `amat(var_rho,var_rhoimp)` 4, `amat(var_rho,var_t)` 2, `amat(var_rho,var_ti)` 2 |
+| 2 — impurity electron pressure, tauIC terms (accepted) | 30 | `amat(var_psi,var_rhoimp)` 10, `amat(var_psi,var_t)` 8, `amat(var_psi,var_te)` 8, `amat_n(var_psi,var_rhoimp)` 2, `amat_n(var_psi,var_t)` 1, `amat_n(var_psi,var_te)` 1 |
+| 4 — pinch not differentiated with respect to psi (open) | 24 | `amat(var_rho,var_psi)` 12, `amat(var_vpar,var_psi)` 12 |
+| 9 — sign of the vpar pinch tangent (open) | 16 | `amat(var_vpar,var_rho)` 8, `amat(var_vpar,var_vpar)` 8 |
+| **total** | **166** | 24 blocks: 8 source-only and 158 generated-only lines |
+
+The 8 source-only lines are the wrong-sign pinch terms of finding 9, two per
+block in both temperature branches. Each has a generated partner with the
+opposite sign. All other lines are terms the generator produces and the
+element routine does not.
+
+---
+
+## Regression benchmark
+
+The inventory above is frozen in
+[`reference/model600_discrepancies.json`](reference/model600_discrepancies.json)
+and checked by [`final_test.py`](final_test.py) (or `run_test.sh`, which also
+sets up `sympy`). Each block of the reference names the finding it belongs
+to, so a failure points straight back at this document:
+
+```bash
+./run_test.sh # or: ./final_test.py
+```
+
+Fixing finding 4 or 9 in the Fortran will make the benchmark fail with `-`
+lines for the monomials that stopped disagreeing. That is the intended signal:
+re-run with `--update`, then remove the finding from this file.
+
+## Reproducing
+
+```bash
+cd util/equation_codegen
+module load sympy/1.14.0-gfbf-2025b # or any Python with sympy
+python3 examples/export_model600_terms.py
+python3 examples/diff_model600_reports.py
+```
+
+The **residual** printed under a `DIFF` block is the authoritative difference.
+It is computed in the element basis, where the two spellings of a poloidal
+bracket coincide. It is displayed with `f_s -> f_x`, `f_t -> f_y`,
+`xjac -> 1` so that it reads as a physical expression.
+
+The tool also verifies its own alignment. The blank cells of every block must
+account for that block's multiset difference, and any surplus is reported as
+`MISALIGNED`. In the aligned Markdown reports, a blank cell on the generated
+side is therefore a genuinely missing term.
diff --git a/util/equation_codegen/README.md b/util/equation_codegen/README.md
new file mode 100644
index 0000000000..0ef69c8b5e
--- /dev/null
+++ b/util/equation_codegen/README.md
@@ -0,0 +1,225 @@
+# JOREK equation code generator
+
+This package checks the hand-written linearization in a JOREK element routine
+against one derived symbolically from the same weak form. It does four things:
+
+1. **Read** the assembled `rhs_ij`/`amat` terms out of a model's
+ `mod_elt_matrix_fft.f90`.
+2. **Generate** the residual and the Jacobian blocks from the weak form, by
+ directional differentiation of the equations written in the DSL.
+3. **Compare** the two, monomial by monomial.
+4. **Report** the comparison, so that a mismatch is readable rather than a
+ wall of algebra.
+
+The symbolic layer is described in [`SPECIFICATION.md`](SPECIFICATION.md).
+Equations are written with bare fields (`psi`, `u`, `rho`, ...); the
+linearization stage decides when those become background values and when the
+differentiated one becomes a trial function.
+
+## Layout
+
+| Path | Role |
+|---|---|
+| `src/jorek_equations/symbols.py`, `operators.py`, `external.py` | fields, roles, spatial operators, externally supplied quantities |
+| `src/jorek_equations/equations.py`, `linearization.py`, `channels.py` | weak equations, the Gateaux derivative, the FFT channel split |
+| `src/jorek_equations/fortran.py` | printing a symbolic term with JOREK's names |
+| `src/jorek_equations/fortran_source.py` | step 1 and step 3: read the element routine, normalize its work variables, compare |
+| `src/jorek_equations/model199.py`, `model600.py` | the equations themselves |
+| `src/jorek_equations/model600_markdown.py` | step 4: the aligned Markdown reports |
+| `src/jorek_equations/latex_render.py`, `model600_docs.py`, `model600_notation.py` | the model-600 documentation, generated from `model600.py` |
+| `final_test.py`, `reference/` | the regression benchmark |
+
+## Setup and tests
+
+The only dependency is `sympy`. No package install is needed, because the
+scripts add `src/` to `sys.path` themselves. On the ITER cluster:
+
+```bash
+module load sympy/1.14.0-gfbf-2025b
+python3 -m unittest discover -s tests -v
+```
+
+Any other Python with `sympy` works too, for example a virtual environment
+with `pip install sympy`. `run_test.sh` tries these in order: the system
+Python, then `module load`, then a local `.venv`.
+
+## Model 199
+
+```bash
+python3 examples/check_model199.py
+```
+
+The comparison is symbolic rather than textual. A mismatch raises an error
+that includes the source expression, the generated expression, and their
+algebraic difference.
+
+## Model 600 documentation
+
+[`docs/physics/base_fluid_models/RMHD/weak_form.md`](../../docs/physics/base_fluid_models/RMHD/weak_form.md)
+documents the model-600 equations: every equation written out in full, the
+shared operators, and each equation's term groups. Its math is generated from
+`model600.py`, which stays the only place the equations are written; the
+generator reads the Python source (it does not evaluate it), so each
+`_helper` appears as a named operator and each commented block of a `B = (...)`
+sum becomes a term group. How each name is written in LaTeX, and the titles and
+descriptions of the operators, are in `model600_notation.py`. Only the regions
+between `BEGIN GENERATED` and `END GENERATED` are rewritten; the prose around
+them is hand-written.
+
+```bash
+python3 examples/model600_docs.py render # rewrite the generated regions
+python3 examples/model600_docs.py check # exit 1 if the page is stale
+```
+
+The unit tests and `run_test.sh` run `check`, so a change to `model600.py`
+that is not rendered into the page fails the build. A new equation or helper
+must also be added to `model600_notation.py`; `render` says so if it is missing.
+
+## Model 600 reports
+
+`export_model600_terms.py` writes the two line-aligned Markdown reports for
+the `psi`, `u`, `zj`, `w`, `rho`, `vpar`, `rhoimp`, `rhon`, `Ti`, `Te` and `T` rows:
+
+```bash
+python3 examples/export_model600_terms.py
+meld reports/model600_fortran_terms.md reports/model600_generated_terms.md
+```
+
+Repeat `--equation` to export fewer rows while developing one of them:
+
+```bash
+python3 examples/export_model600_terms.py --equation psi
+```
+
+`Ti` and `Te` exist only in the two-temperature branch of the element routine
+and `T` only in the other one, so each is exported into its own section.
+
+To compare the two reports block by block instead of line by line:
+
+```bash
+python3 examples/diff_model600_reports.py
+python3 examples/diff_model600_reports.py --quiet # verdicts only
+```
+
+Each assignment block is compared as a multiset of monomials, and a block that
+differs is printed together with the algebraic residual `source-generated`.
+Reordering alone therefore does not show up as a difference. The residual is
+taken after rewriting every `_s`/`_t` derivative through the chain rule, so a
+poloidal bracket written `a_s*b_t - a_t*b_s` in one report and
+`xjac*(a_x*b_y - a_y*b_x)` in the other cancels; such a block is reported as
+`SAME` rather than `DIFF`.
+
+## Temperature conventions
+
+The element routine builds every pressure in `construct_pressure`, which is
+written once for both temperature models and always uses the species
+temperatures `Ti0`/`Te0` and the per-species impurity coefficients `alpha_i`,
+`alpha_e` and `alpha_e_bis`. The one-temperature model evolves the *total*
+temperature and enters that routine with
+
+```
+Ti0 = Te0 = T0/2
+```
+
+so a single-species temperature is half of the evolved field, and the stored
+one-temperature closure values are the means
+
+```
+alpha_imp = (alpha_i + alpha_e)/2
+alpha_imp_bis = (alpha_i + alpha_e_bis)/2
+alpha_imp_tri = alpha_e_tri/4
+```
+
+The comparison therefore expands both reports in the two-species basis. The
+DSL follows the same convention: `_diamagnetic_pressure` and the induction
+equation's electron pressure use `T/2` in the one-temperature branch. This
+replaced two earlier "legacy tangent" overrides that reproduced JOREK's
+factor of one half by hand; that half is physical, not a legacy quirk.
+
+## The `final_test` benchmark
+
+The discrepancies between the element routine and the generated linearization
+are frozen in
+[`reference/model600_discrepancies.json`](reference/model600_discrepancies.json).
+For every assignment that disagrees, it records the monomials the element
+routine has and the generator does not, the monomials the generator has and
+the element routine does not, and which findings of
+[`JOREK_FINDINGS.md`](JOREK_FINDINGS.md) that block belongs to. Those two
+lists are the complete statement of the disagreement: everything else in the
+reports is identical on both sides.
+
+```bash
+./final_test.py # recompute and compare; exit 0 on a match
+./final_test.py --equation vpar # only these rows, for a quick check
+./final_test.py --update # rewrite the reference after an intended change
+```
+
+The benchmark exports the reports into a temporary directory, so it never
+touches `reports/`. It takes a few minutes, which is why it is a standalone
+script rather than part of `python -m unittest discover`; the unit tests cover
+its comparison logic and the integrity of the reference file.
+
+A failure prints one line per changed monomial: `-` for a discrepancy the
+reference expects that the run no longer produces, `+` for a new one. Both
+directions matter. A `-` line means either that somebody fixed the Fortran, in
+which case update the reference and strike the finding from
+`JOREK_FINDINGS.md`, or that the generated equation drifted and stopped
+reproducing a term it used to reproduce. A `+` line means a new disagreement:
+a change in the element routine, or a regression in the DSL. The script also
+warns when `mod_elt_matrix_fft.f90` no longer has the SHA-256 recorded in the
+reference, which usually explains the rest of the output.
+
+## Known JOREK differences
+
+With the conventions above, all eleven rows (`psi`, `u`, `zj`, `w`, `rho`,
+`vpar`, `rhoimp`, `rhon`, `Ti`, `Te` and `T`) are reproduced by the
+generator. The first audit reported seventeen findings for the first ten
+rows. Most have since been fixed in the Fortran, and every residual
+(`rhs_ij`) now agrees exactly. The assignment blocks that still differ are all
+Jacobian tangents, documented in [`JOREK_FINDINGS.md`](JOREK_FINDINGS.md):
+
+- **Open, left to the pinch developer**:
+ - Finding 4: the density and parallel-velocity inward-pinch terms are not
+ differentiated with respect to `psi` (`amat(var_rho,var_psi)`,
+ `amat(var_vpar,var_psi)`).
+ - Finding 9: the parallel-velocity pinch tangent repeats the sign of its
+ own residual instead of flipping it (`amat(var_vpar,var_rho)`,
+ `amat(var_vpar,var_vpar)`).
+- **Accepted, not bugs**:
+ - Findings 1 and 2: the impurity parts of the ion and electron pressures
+ are not differentiated in the `tauIC` diamagnetic tangents of the `u`,
+ `rho` and `psi` equations.
+ - Model 600 does not support impurities together with `tauIC /= 0`, so
+ these terms never contribute in a supported run.
+
+A source monomial and its generated counterpart are aligned on the same report
+line even when only their numeric coefficient or their power of `BigR` differs, so a term the element
+routine scales differently shows up side by side in Meld rather than as a
+source line with an empty generated cell followed by an unmatched generated
+line at the end of the block.
+
+The element routine is not consistent about how it spells a poloidal bracket:
+`rhs_ij(var_vpar)` writes the parallel kinetic-energy flux as
+`a_s*b_t - a_t*b_s` while `amat(var_vpar,var_psi)` writes the very same group
+as `xjac*(a_x*b_y - a_y*b_x)`, and the two expand into different monomials.
+The parallel-velocity equation is therefore generated in both spellings and
+each assignment is aligned against whichever one its own source block uses.
+A blank cell on the generated side of the report is then a genuinely missing
+term rather than a change of coordinates.
+
+Matching runs in three passes over a whole block — exact monomials first, then
+the ones that differ only by a coefficient, then the ones that also differ by
+a power of `BigR`. A single interleaved pass lets an early source monomial
+with no exact partner consume, through one of the relaxed keys, a generated
+monomial that a later source monomial matches exactly; the pair then drifts
+apart even though the block agrees. A fourth pass pairs terms that differ by the implicitness factor `theta` as
+well, since that is a scalar scaling like any other. `diff_model600_reports.py`
+checks the result: the blank cells of every block must account for its
+multiset difference, and any surplus is reported as `MISALIGNED`.
+
+Assignments and terms follow their order in
+`models/model600/mod_elt_matrix_fft.f90`. A source term that is not available
+from the equation generator has an empty line at the same position in the
+generated report. Generated-only terms are appended to their corresponding
+assignment block with an empty source line. NEO AMAT terms remain explicitly
+source-only until their symbolic linearization is enabled.
diff --git a/util/equation_codegen/SPECIFICATION.md b/util/equation_codegen/SPECIFICATION.md
new file mode 100644
index 0000000000..dbeb7b1a99
--- /dev/null
+++ b/util/equation_codegen/SPECIFICATION.md
@@ -0,0 +1,386 @@
+# JOREK equation-codegen specification
+
+## 1. Purpose and scope
+
+This document is the initial correctness reference for the Python equation
+code generator. The first target is an exact, auditable reproduction of the
+volume terms assembled by `models/model199/mod_elt_matrix_fft.f90`.
+
+Version 1 accepts equations that are already in weak form. It does not derive
+the weak form from a strong equation and does not perform integration by
+parts. Strong-form input, automatic integration by parts, and explicit surface
+operators are deferred to a later version.
+
+There are two distinct goals:
+
+1. **Legacy reproduction mode:** reproduce model 199, including quantities
+ that the existing Jacobian intentionally or historically treats as frozen.
+2. **Declared-physics mode:** later allow each external or derived quantity to
+ state its true dependencies explicitly and differentiate them when desired.
+
+The generator must never change from one policy to the other implicitly.
+
+## 2. Model-199 fields
+
+The ordered state vector is
+
+$$
+\mathbf q=(\psi,u,j,\omega,\rho,T).
+$$
+
+The ordering is part of the generated-code interface:
+
+| Index | DSL name | Mathematical meaning | Current-state Fortran name | Trial-function Fortran name |
+|---:|---|---|---|---|
+| 1 | `psi` | poloidal magnetic flux | `ps0` | `psi` |
+| 2 | `u` | velocity stream function | `u0` | `u` |
+| 3 | `j` | toroidal plasma current | `zj0` | `zj` |
+| 4 | `omega` | toroidal vorticity | `w0` | `w` |
+| 5 | `rho` | particle density | `r0` | `rho` |
+| 6 | `T` | total temperature | `T0` | `T` |
+
+The DSL uses `j` and `omega` as the canonical physics names. The Fortran
+backend maps them to the legacy names `zj` and `w`.
+
+## 3. State, increment, and trial-function naming
+
+For a field `q`, the DSL distinguishes the following roles:
+
+| DSL concept | Mathematical notation | Example for `psi` | Purpose |
+|---|---|---|---|
+| field declaration | $q$ | `psi = field("psi")` | identifies a state variable |
+| current/background value | $q_0\equiv q^n$ | `psi.current` | evaluates the residual |
+| symbolic increment | $\delta q$ | `delta(psi)` | directional linearization |
+| previous increment | $\delta q^{n-1}$ | `previous_delta(psi)` | multistep time term |
+| finite-element trial function | $\varphi_j$ | `trial(psi)` | matrix column after differentiation |
+
+Linearization is defined by the Gateaux derivative
+
+$$
+D F(\mathbf q_0)[\delta\mathbf q]
+=
+\left.\frac{d}{d\epsilon}
+F(\mathbf q_0+\epsilon\,\delta\mathbf q)
+\right|_{\epsilon=0}.
+$$
+
+When a Jacobian block for field $q$ is emitted, `delta(q)` and all its
+derivatives are replaced by the corresponding finite-element trial function
+and its derivatives. Background values keep the legacy Fortran suffix `0`.
+
+Equation definitions use the bare field declaration (`psi`, `u`, `rho`, ...),
+which denotes an abstract state field. The `.current` form is an explicit
+low-level state value and is useful in tests or when constructing already
+linearized expressions. The linearization stage resolves every non-varied
+abstract field to its current/background value and replaces only the selected
+field by the requested increment or trial role.
+
+Spatial variation and differentiation commute:
+
+$$
+\delta(\partial_a q)=\partial_a(\delta q),
+\qquad a\in\{R,Z,\phi\}.
+$$
+
+## 4. Coordinates and geometry
+
+### 4.1 Physical coordinates
+
+The physical coordinates are cylindrical
+
+$$
+(R,Z,\phi).
+$$
+
+The DSL coordinate names are `R`, `Z`, and `phi`, with derivative operators
+`dR`, `dZ`, and `dphi`. The legacy model-199 implementation uses `x`, `y`, and
+`p` suffixes for these derivatives:
+
+| DSL | Mathematical notation | Legacy suffix |
+|---|---|---|
+| `dR(f)` | $\partial_R f$ | `f_x` |
+| `dZ(f)` | $\partial_Z f$ | `f_y` |
+| `dphi(f)` | $\partial_\phi f$ | `f_p` |
+
+The toroidal direction is periodic. Toroidal endpoint surface terms therefore
+cancel.
+
+### 4.2 Element coordinates
+
+The two-dimensional element coordinates are $(s,t)$, with mapping
+
+$$
+R=R(s,t),\qquad Z=Z(s,t),
+$$
+
+and poloidal Jacobian
+
+$$
+J=R_s Z_t-R_t Z_s.
+$$
+
+The physical derivatives are
+
+$$
+\partial_R f=\frac{Z_t f_s-Z_s f_t}{J},\qquad
+\partial_Z f=\frac{-R_t f_s+R_s f_t}{J}.
+$$
+
+The DSL's canonical equations are written using physical derivatives. The
+Fortran backend may rewrite antisymmetric products in $(s,t)$ coordinates,
+as the existing routine does, but that is a code-generation transformation and
+must not alter the symbolic equation.
+
+The weak integrals use $dR\,dZ\,d\phi$. Cylindrical factors of $R$ remain
+explicit in the equation. At a quadrature point, $dR\,dZ$ contributes the
+mapping factor `xjac`, and Gaussian integration contributes `wst`.
+
+Geometry is frozen during physics linearization: variations of `R`, `Z`, the
+mapping, `xjac`, metric factors, basis functions, and quadrature weights are
+zero.
+
+## 5. Test- and trial-function conventions
+
+The default discretization is Galerkin: test and trial functions come from the
+same finite-element space but have independent row and column indices.
+
+```python
+v = test_function("v")
+```
+
+denotes a symbolic test function. It is not a physical field and is never
+linearized. Its derivatives `dR(v)`, `dZ(v)`, and `dphi(v)` are valid symbolic
+objects.
+
+For a residual component $\mathcal R_i$, the local row uses test basis
+function $v_i$. A variation is expanded as
+
+$$
+\delta q=\sum_j \varphi_j\,\delta q_j,
+$$
+
+and the matrix entry is
+
+$$
+M_{ij}^{(q)}=D\mathcal R_i[\varphi_j].
+$$
+
+This corresponds to the outer `(i,j)` basis loops for `v` and the inner
+`(k,l)` basis loops for the trial quantities in the current Fortran routine.
+
+Toroidal derivatives determine the intermediate FFT channel:
+
+| Derivative placement | RHS channel | Matrix channel |
+|---|---|---|
+| neither test nor trial | `RHS_p` | `ELM_p` |
+| trial only | not applicable | `ELM_n` |
+| test only | `RHS_k` | `ELM_k` |
+| both test and trial | not applicable | `ELM_kn` |
+
+The names `p`, `n`, `k`, and `kn` are legacy output names. Internally, the DSL
+must retain semantic flags such as `test_phi_order` and `trial_phi_order`.
+Unsupported derivative orders must produce an error instead of being assigned
+heuristically.
+
+## 6. Equation, RHS, and AMAT sign conventions
+
+### 6.1 Evolution equations
+
+Evolution equations 1, 2, 5, and 6 are represented as weak functionals
+
+$$
+\frac{\partial A_i(\mathbf q)}{\partial t}=B_i(\mathbf q).
+$$
+
+The state update is
+
+$$
+\mathbf q^{n+1}=\mathbf q^n+\delta\mathbf q^n.
+$$
+
+With JOREK's `theta` and `zeta` convention, the generated local system is
+
+$$
+\left[(1+\zeta)D A_i
+-\theta\,\Delta t\,D B_i\right][\delta\mathbf q^n]
+=
+\Delta t\,B_i(\mathbf q^n)
++\zeta\,D A_i[\delta\mathbf q^{n-1}].
+$$
+
+Therefore:
+
+```text
+RHS = + tstep * B(current)
+ + zeta * D(A)(previous_delta)
+
+AMAT = + (1 + zeta) * D(A)(trial)
+ - theta * tstep * D(B)(trial)
+```
+
+The existing routine rescales the stored previous increment by
+`tstep/tstep_prev` before using it. This rescaling belongs to time-history
+preparation, not symbolic differentiation.
+
+For exact legacy reproduction, the momentum time functional treats the
+current-state density coefficient `r0_hat` as frozen while differentiating the
+time term with respect to `u`. This special policy must be declared explicitly;
+it must not arise from generic differentiation.
+
+### 6.2 Algebraic constraint equations
+
+Equations 3 and 4 are algebraic definitions rather than evolution equations.
+For a weak constraint
+
+$$
+C_i(\mathbf q)=0,
+$$
+
+the Newton correction convention is
+
+$$
+D C_i(\mathbf q^n)[\delta\mathbf q]=-C_i(\mathbf q^n).
+$$
+
+Therefore:
+
+```text
+RHS = -C(current)
+AMAT = +D(C)(trial)
+```
+
+No `theta`, `zeta`, or `tstep` factor is applied to these constraint blocks.
+
+## 7. Boundary-term policy
+
+Version 1 starts from the already-integrated weak form. Its input expression is
+the volume integrand that remains after any integration by parts. It does not
+reconstruct a discarded surface term.
+
+For legacy model-199 reproduction:
+
+1. Toroidal surface terms cancel because $\phi$ is periodic.
+2. Poloidal surface terms produced by integrations by parts are absent from
+ `mod_elt_matrix_fft.f90` and are treated as discarded in the volume
+ generator.
+3. The model applies essential boundary conditions separately after volume
+ assembly. Fixed-boundary test/variation traces are consequently constrained.
+4. Free-boundary or open-boundary surface physics is outside the version-1
+ generator scope. The absence of a surface contribution must not be taken as
+ a derivation that it is zero for those cases.
+
+The discarded poloidal terms belong to these classes:
+
+| Weak volume structure | Discarded surface-term class |
+|---|---|
+| `grad(v) . grad(psi)` and `grad(v) . grad(u)` | normal flux from the elliptic current/vorticity definitions |
+| `grad(v) . grad(j)` and `grad(v) . grad(omega)` | normal numerical resistive/viscous flux |
+| `grad(v) . grad(rho)` and `grad(v) . grad(T)` | normal perpendicular particle/heat flux |
+| parallel-gradient test factor times parallel field gradient | normal anisotropic particle/heat flux |
+| derivatives moved onto `v` in bracket or pressure terms | advective or pressure boundary flux generated by that integration by parts |
+| the fourth-order numerical-viscosity weak product | boundary traces generated by the repeated integration by parts |
+
+Because version 1 receives only the final weak expression, the exact signed
+surface formula is not inferred. Each future strong-form transformation must
+return both a volume term and an explicit boundary term. Discarding that term
+will require a named boundary policy and a stated boundary condition.
+
+## 8. Linearization dependency policy
+
+Every nontrivial derived quantity or external function has one of three
+policies:
+
+- `active`: apply the chain rule using its declared derivatives;
+- `frozen`: evaluate its value at the current state and set its variation to
+ zero;
+- `piecewise_active`: select a branch at the current state, then use derivatives
+ explicitly supplied for that branch.
+
+If an active dependency has no derivative interface, generation must fail with
+a diagnostic. It must not silently freeze the dependency.
+
+### 8.1 Quantities differentiated in legacy reproduction mode
+
+| Quantity | Active dependencies | Required derivative behavior |
+|---|---|---|
+| all six state fields and their spatial derivatives | their own field | standard linear variation |
+| $P=\rho T$ | `rho`, `T` | product rule |
+| $\hat\rho=R^2\rho$ and its gradients | `rho` | geometry remains frozen |
+| $\lvert\mathbf v_E\rvert^2=R^2\lvert\nabla u\rvert^2$ (`vv2`) | `u` | differentiate both gradient factors |
+| $B^2=(F_0^2+\lvert\nabla\psi\rvert^2)/R^2$ | `psi` | includes the derivative called `BB2_psi` |
+| parallel-gradient expressions | `psi` and the transported field | product and quotient rules, including the derivative of `B^2` |
+| `eta_T` | `T` | external interface supplies `deta_dT` |
+| `eta_T_ohm` | `T` | external interface supplies `deta_dT_ohm` |
+| `visco_T` | `T` | external interface supplies `dvisco_dT` |
+
+Temperature-dependent coefficients use `piecewise_active`: the current branch
+or clipping regime is selected first, and the interface supplies the derivative
+for that branch. A clipped constant branch consequently has zero derivative.
+
+### 8.2 Quantities frozen in legacy reproduction mode
+
+| Quantity | Notes |
+|---|---|
+| coordinates, mapping, metric, `xjac`, basis functions, and quadrature weights | mesh/shape differentiation is out of scope |
+| `F0`, `gamma`, `theta`, `zeta`, `tstep`, and numerical coefficients | run parameters |
+| `D_par`, `ZK_par`, `eta_num`, and `visco_num` | scalar coefficients |
+| `psi_norm = get_psi_n(psi, Z)` | evaluated from the current state but not differentiated in model 199 |
+| `D_prof = get_dperp(psi_norm)` | frozen even though its evaluated value depends on `psi_norm` |
+| `ZK_prof = get_zkperp(psi_norm)` | frozen even though its evaluated value depends on `psi_norm` |
+| threshold-based replacements `D_prof_neg` and `ZK_prof_neg` | branch and value frozen during a Jacobian evaluation |
+| current, particle, and heat sources | frozen, including any state dependence inside their evaluation routines |
+| equilibrium and normalization data such as axis/boundary fluxes | external data |
+
+Model 199 evaluates the background density and temperature with `abs(...)`,
+but its matrix expressions use `rho` and `T` without the derivative of the
+absolute-value operation. To reproduce the implementation, background
+sanitization is treated as an external evaluation step while field increments
+use the identity tangent. This behavior must be marked as a legacy policy.
+
+### 8.3 Future explicit external-function interface
+
+An external function declaration must record:
+
+```text
+name
+arguments and declared state dependencies
+value symbol or value callback
+available first derivatives
+linearization policy
+piecewise/clipping policy, when applicable
+Fortran names for the value and derivatives
+```
+
+For example, the mathematical function `eta(T)` supplies both its value and
+`deta_dT`. Profile functions can later be changed from `frozen` to `active`
+only by an explicit specification change and the addition of the corresponding
+derivative interface.
+
+## 9. Correctness requirements for later implementation
+
+The eventual Python implementation must:
+
+1. preserve field, test, trial, and coordinate-derivative roles in the symbolic
+ expression tree;
+2. reproduce the sign rules in Section 6 before applying algebraic
+ simplification;
+3. make every external dependency policy inspectable;
+4. keep named physics terms available in generated comments and audit output;
+5. classify toroidal derivatives structurally;
+6. reject missing derivatives, unsupported toroidal orders, and undeclared
+ boundary-term deletion;
+7. compare generated model-199 terms against the existing RHS and AMAT
+ expressions using symbolic and numerical tests.
+
+## 10. Scope decisions still requiring physics review
+
+The following legacy behaviors are recorded for reproduction but should be
+reviewed before declaring them the preferred mathematical model:
+
+- freezing `D_prof`, `ZK_prof`, `psi_norm`, and state-dependent sources;
+- freezing the density coefficient in the momentum time derivative;
+- using an identity tangent after `abs(...)` background sanitization;
+- omitting poloidal surface contributions outside fixed-boundary use.
+
+Changes to these items define a different Jacobian or boundary-value problem
+and must be made as explicit specification revisions.
diff --git a/util/equation_codegen/examples/check_model199.py b/util/equation_codegen/examples/check_model199.py
new file mode 100644
index 0000000000..fe2ba433de
--- /dev/null
+++ b/util/equation_codegen/examples/check_model199.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+"""Compare all implemented model-199 equations with the Fortran source."""
+
+from pathlib import Path
+import argparse
+import sys
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+REPOSITORY_ROOT = PROJECT_ROOT.parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+
+from jorek_equations.fortran_source import ( # noqa: E402
+ EQUATION_1_ASSIGNMENTS,
+ EQUATION_2_ASSIGNMENTS,
+ EQUATION_3_ASSIGNMENTS,
+ EQUATION_4_ASSIGNMENTS,
+ EQUATION_5_ASSIGNMENTS,
+ EQUATION_6_ASSIGNMENTS,
+ FortranComparisonError,
+ compare_assignment_maps,
+ extract_fortran_assignments,
+ generated_assignment_text,
+ generated_model199_transport_assignments,
+ generated_partial_model199_assignments,
+)
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--keep-BigR-x", dest="set_BigR_x_one", action="store_false",
+ help="retain BigR_x geometry factors during comparison",
+ )
+ arguments = parser.parse_args()
+ source_file = REPOSITORY_ROOT / "models/model199/mod_elt_matrix_fft.f90"
+ assignment_names = (
+ EQUATION_1_ASSIGNMENTS + EQUATION_2_ASSIGNMENTS + EQUATION_3_ASSIGNMENTS
+ + EQUATION_4_ASSIGNMENTS + EQUATION_5_ASSIGNMENTS + EQUATION_6_ASSIGNMENTS
+ )
+ source = extract_fortran_assignments(source_file, assignment_names)
+ generated = generated_assignment_text()
+ generated.update(generated_partial_model199_assignments())
+ generated.update(generated_model199_transport_assignments())
+ try:
+ compare_assignment_maps(
+ source, generated, set_BigR_x_one=arguments.set_BigR_x_one
+ )
+ except FortranComparisonError as error:
+ print(error)
+ raise SystemExit(1)
+ print("No differences found for model-199 equations 1–6")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/util/equation_codegen/examples/diff_model600_reports.py b/util/equation_codegen/examples/diff_model600_reports.py
new file mode 100755
index 0000000000..ed7ed369d0
--- /dev/null
+++ b/util/equation_codegen/examples/diff_model600_reports.py
@@ -0,0 +1,259 @@
+#!/usr/bin/env python3
+"""Compare the two model-600 term reports block by block.
+
+Meld shows every line that differs, including the many lines that differ only
+because source and generated monomials are listed in a different order. This
+script instead compares each assignment block as a multiset of monomials and,
+where the multisets differ, prints the algebraic residual ``source-generated``.
+A block whose residual is zero contains the same physics in a different
+arrangement; a block with a non-zero residual is a real disagreement.
+"""
+
+import collections
+from collections import Counter, OrderedDict
+from pathlib import Path
+import argparse
+import re
+import sys
+
+import sympy as sp
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def read_blocks(path):
+ """Return {(section, assignment, repeat): [monomial, ...]}."""
+
+ blocks = OrderedDict()
+ key = None
+ section = ""
+ for line in Path(path).read_text(encoding="utf-8").splitlines():
+ if line.startswith("#") or line.startswith(">"):
+ # Any heading closes the block being read. Without this the last
+ # block of a section swallowed the next section's title lines.
+ if line.startswith("## ") and not line.startswith("### "):
+ section = line[3:].strip()
+ key = None
+ if line.startswith("## ") and not line.startswith("### "):
+ continue
+ heading = re.match(r"^#### `(.*)`", line)
+ if heading:
+ repeats = sum(
+ 1 for existing in blocks
+ if existing[0] == section and existing[1] == heading.group(1)
+ )
+ key = (section, heading.group(1), repeats)
+ blocks[key] = []
+ continue
+ if key is not None and line.strip():
+ blocks[key].append(line.strip())
+ return blocks
+
+
+def to_element_basis(expression):
+ """Rewrite both coordinate spellings of a poloidal bracket into one basis.
+
+ The element routine writes the same antisymmetric product either as
+ ``a_s*b_t - a_t*b_s`` or as ``xjac*(a_x*b_y - a_y*b_x)``, and it is not
+ always consistent between a residual and its own tangent. Expanding every
+ ``_s``/``_t`` derivative through the chain rule,
+
+ f_s = x_s*f_x + y_s*f_y, f_t = x_t*f_x + y_t*f_y,
+ xjac = x_s*y_t - x_t*y_s,
+
+ maps both spellings onto the same polynomial, so a residual that vanishes
+ in this basis means the two sides carry the same term.
+ """
+
+ x_s, x_t, y_s, y_t = sp.symbols("x_s x_t y_s y_t")
+ replacements = {sp.Symbol("xjac"): x_s * y_t - x_t * y_s}
+ for symbol in expression.free_symbols:
+ name = str(symbol)
+ if name in ("x_s", "x_t", "y_s", "y_t", "xjac") or name.startswith("var_"):
+ continue
+ if name.endswith("_s"):
+ replacements[symbol] = (
+ x_s * sp.Symbol(name[:-2] + "_x") + y_s * sp.Symbol(name[:-2] + "_y")
+ )
+ elif name.endswith("_t"):
+ replacements[symbol] = (
+ x_t * sp.Symbol(name[:-2] + "_x") + y_t * sp.Symbol(name[:-2] + "_y")
+ )
+ return sp.expand(expression.xreplace(replacements))
+
+
+def to_physical(expression):
+ """Render a residual with a single spelling for poloidal derivatives.
+
+ Specialising the element map to ``x_s = y_t = 1``, ``x_t = y_s = 0`` makes
+ ``f_s`` and ``f_t`` coincide with ``f_x`` and ``f_y`` and ``xjac`` with 1,
+ so a residual printed this way is readable and free of the (s,t)/(R,Z)
+ bookkeeping. The verdict is always taken in the full element basis; this
+ is only for display.
+ """
+
+ replacements = {sp.Symbol("xjac"): sp.Integer(1)}
+ for symbol in expression.free_symbols:
+ name = str(symbol)
+ if name in ("xjac",) or name.startswith("var_"):
+ continue
+ if name.endswith("_s"):
+ replacements[symbol] = sp.Symbol(name[:-2] + "_x")
+ elif name.endswith("_t"):
+ replacements[symbol] = sp.Symbol(name[:-2] + "_y")
+ return sp.expand(expression.xreplace(replacements))
+
+
+def check_alignment(source_path, generated_path, expected):
+ """Verify that blank report cells account for the multiset difference.
+
+ The reports are line aligned, so a blank generated cell marks a term the
+ element routine has and the generator does not. There must be exactly as
+ many of them as the block has source-only monomials, and likewise on the
+ other side. Any surplus is a pair the matcher failed to put on one line
+ although the report renders both the same way — a defect in the exporter
+ that the block verdicts cannot see, because the multiset difference is
+ computed from the file contents and cancels such a pair out.
+ """
+
+ source = Path(source_path).read_text(encoding="utf-8").splitlines()
+ generated = Path(generated_path).read_text(encoding="utf-8").splitlines()
+ section = None
+ key = None
+ seen = collections.Counter()
+ unpaired = collections.Counter()
+ for left, right in zip(source, generated):
+ if left.startswith("#") or left.startswith(">"):
+ if left.startswith("## ") and not left.startswith("### "):
+ section = left[3:].strip()
+ heading = re.match(r"^#### `(.*)`", left)
+ key = None
+ if heading:
+ name = heading.group(1)
+ key = (section, name, seen[(section, name)])
+ seen[(section, name)] += 1
+ continue
+ if key is None or not (left.strip() or right.strip()):
+ continue
+ if left.strip() and not right.strip():
+ unpaired[(key, "source")] += 1
+ elif right.strip() and not left.strip():
+ unpaired[(key, "generated")] += 1
+ surplus = collections.Counter()
+ for (key, side), count in unpaired.items():
+ # A negative difference means two terms that do differ were still
+ # placed on one line, by the coefficient or radius fallback. That is
+ # the intent; only a surplus of blanks is a failure to pair.
+ extra = count - expected.get((key, side), 0)
+ if extra > 0:
+ surplus[(key, side)] = extra
+ return surplus
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ reports = PROJECT_ROOT / "reports"
+ parser.add_argument("source", nargs="?", default=reports / "model600_fortran_terms.md")
+ parser.add_argument("generated", nargs="?", default=reports / "model600_generated_terms.md")
+ parser.add_argument(
+ "--quiet", action="store_true", help="print only the per-block verdicts",
+ )
+ args = parser.parse_args()
+
+ source = read_blocks(args.source)
+ generated = read_blocks(args.generated)
+
+ differing = 0
+ expected_blanks = {}
+ keys = list(OrderedDict.fromkeys(list(source) + list(generated)))
+ for key in keys:
+ source_only = Counter(source.get(key, ())) - Counter(generated.get(key, ()))
+ generated_only = Counter(generated.get(key, ())) - Counter(source.get(key, ()))
+ label = "{} {} #{}".format(*key)
+ expected_blanks[(key, "source")] = sum(source_only.values())
+ expected_blanks[(key, "generated")] = sum(generated_only.values())
+ if not source_only and not generated_only:
+ print("OK {} ({} terms)".format(label, len(source.get(key, ()))))
+ continue
+ def _sum(terms):
+ """Sum the parsable terms, listing the ones that are not.
+
+ A source line the exporter could not expand is copied into the
+ report verbatim; it is not a SymPy expression.
+ """
+
+ total = sp.S.Zero
+ unparsed = []
+ for term in terms:
+ # Parse every identifier as a plain symbol: several JOREK work
+ # variables (``zeta``, ``gamma``, ``beta``, ...) collide with
+ # SymPy's own function names.
+ locals_ = {
+ name: sp.Symbol(name)
+ for name in re.findall(r"[A-Za-z_][A-Za-z_0-9]*", term)
+ if name != "sqrt"
+ }
+ try:
+ total += sp.sympify(term, locals=locals_, evaluate=True)
+ except Exception:
+ unparsed.append(term)
+ return total, unparsed
+
+ source_total, source_unparsed = _sum(source_only.elements())
+ generated_total, generated_unparsed = _sum(generated_only.elements())
+ residual = sp.expand(source_total - generated_total)
+ # The verdict is taken in the element basis, where the two spellings
+ # of a poloidal bracket coincide; anything left there is a genuine
+ # difference.
+ residual_basis = to_element_basis(residual)
+ if residual != 0 and residual_basis == 0:
+ # The two sides carry the same physics; the element routine simply
+ # wrote a poloidal bracket in (s,t) where the generator wrote it
+ # in (R,Z), or the other way round.
+ print("SAME {}: {} line(s), identical up to the (s,t)/(R,Z) "
+ "spelling of a poloidal bracket".format(
+ label, sum(source_only.values()),
+ ))
+ continue
+ differing += 1
+ print("DIFF {}: source-only={} generated-only={} residual terms (element basis)={}{}".format(
+ label, sum(source_only.values()), sum(generated_only.values()),
+ 0 if residual_basis == 0 else len(sp.Add.make_args(residual_basis)),
+ "" if not (source_unparsed or generated_unparsed)
+ else " (unparsed: {})".format(
+ len(source_unparsed) + len(generated_unparsed)
+ ),
+ ))
+ if args.quiet:
+ continue
+ for term in sorted(source_only.elements()):
+ print(" source {}".format(term))
+ for term in sorted(generated_only.elements()):
+ print(" generated {}".format(term))
+ # Print the residual with one spelling for the poloidal derivatives,
+ # so that lines differing only in (s,t) versus (R,Z) cancel here as
+ # they do in the verdict above.
+ shown = to_physical(residual)
+ if shown == 0:
+ print(" residual: 0 in the displayed basis; the difference is "
+ "visible only in the element metric")
+ else:
+ print(" residual (source-generated, with f_s->f_x, f_t->f_y, "
+ "xjac->1):")
+ for term in sp.Add.make_args(shown):
+ print(" {}".format(term))
+ print("\nBlocks with differences: {} / {}".format(differing, len(keys)))
+ misaligned = check_alignment(args.source, args.generated, expected_blanks)
+ if misaligned:
+ print("MISALIGNED: {} report line(s) left unpaired although both "
+ "reports contain the term:".format(sum(misaligned.values())))
+ for (block, side), count in misaligned.items():
+ print(" {} {} #{} ({} side): {}".format(*block, side, count))
+ else:
+ print("Alignment check: blank cells match the multiset difference in "
+ "every block.")
+ return 1 if differing else 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/util/equation_codegen/examples/export_model600_terms.py b/util/equation_codegen/examples/export_model600_terms.py
new file mode 100644
index 0000000000..3e74fc589b
--- /dev/null
+++ b/util/equation_codegen/examples/export_model600_terms.py
@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Export aligned model-600 source/generated terms for a visual diff."""
+
+from pathlib import Path
+import argparse
+import sys
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+REPOSITORY_ROOT = PROJECT_ROOT.parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+
+from jorek_equations.model600_markdown import export_model600_markdown # noqa: E402
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--equation",
+ choices=("psi", "u", "zj", "w", "rho", "vpar", "rhoimp", "rhon", "Ti", "Te", "T"),
+ action="append",
+ help="export only this equation; repeat the option for multiple equations",
+ )
+ parser.add_argument(
+ "--include-neo",
+ action="store_true",
+ help="include NEO terms (omitted by default for the base-equation comparison)",
+ )
+ args = parser.parse_args()
+ reports = PROJECT_ROOT / "reports"
+ source, generated = export_model600_markdown(
+ REPOSITORY_ROOT / "models/model600/mod_elt_matrix_fft.f90",
+ reports / "model600_fortran_terms.md",
+ reports / "model600_generated_terms.md",
+ equations=args.equation,
+ include_neo=args.include_neo,
+ )
+ print("Wrote {}".format(source))
+ print("Wrote {}".format(generated))
+ if args.equation:
+ print("Selected equations: {}".format(", ".join(args.equation)))
+ print("Compare with: meld {} {}".format(source, generated))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/util/equation_codegen/examples/model600_docs.py b/util/equation_codegen/examples/model600_docs.py
new file mode 100644
index 0000000000..faaf4c7140
--- /dev/null
+++ b/util/equation_codegen/examples/model600_docs.py
@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Keep the model-600 weak-form documentation in step with model600.py.
+
+ model600_docs.py render # rewrite the generated regions of the page
+ model600_docs.py check # change nothing; exit 1 if the page is stale
+
+The equations are written in ``src/jorek_equations/model600.py``; the page
+``docs/physics/base_fluid_models/RMHD/weak_form.md`` is generated from them.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+REPOSITORY_ROOT = PROJECT_ROOT.parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+
+from jorek_equations.model600_docs import render_page # noqa: E402
+
+PAGE = REPOSITORY_ROOT / "docs" / "physics" / "base_fluid_models" / "RMHD" / "weak_form.md"
+SOURCE = PROJECT_ROOT / "src" / "jorek_equations" / "model600.py"
+REFERENCE = PROJECT_ROOT / "reference" / "model600_discrepancies.json"
+
+
+def rendered():
+ text = PAGE.read_text(encoding="utf-8")
+ return text, render_page(
+ text,
+ SOURCE.read_text(encoding="utf-8"),
+ REFERENCE.read_text(encoding="utf-8"),
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+ parser.add_argument("command", choices=("render", "check"))
+ arguments = parser.parse_args()
+ name = PAGE.relative_to(REPOSITORY_ROOT)
+ text, new = rendered()
+ if arguments.command == "render":
+ if new != text:
+ PAGE.write_text(new, encoding="utf-8")
+ print("Rewrote the generated regions of {}".format(name))
+ else:
+ print("{} is up to date".format(name))
+ return 0
+ if new != text:
+ print("STALE: {} does not match model600.py; run 'python3 examples/model600_docs.py "
+ "render'".format(name))
+ return 1
+ print("{} is in step with model600.py".format(name))
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/util/equation_codegen/final_test.py b/util/equation_codegen/final_test.py
new file mode 100755
index 0000000000..96bfa93d46
--- /dev/null
+++ b/util/equation_codegen/final_test.py
@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+"""Recompute every model-600 discrepancy and compare it with the frozen reference.
+
+The reference records, for every ``rhs_ij``/``amat`` assignment of the element
+routine, which monomials the element routine has and the generated
+linearization does not, and the other way round. Those two lists are the
+complete statement of what disagrees; everything else in the reports is
+identical on both sides.
+
+Run it with no arguments to check the current tree:
+
+ ./final_test.py
+
+It exports the reports into a temporary directory, so it never touches
+``reports/``. It exits 0 when the discrepancies match the reference and 1
+when they do not, printing what moved. ``--update`` rewrites the reference
+after an intended change; annotations of blocks that still exist are carried
+over.
+"""
+
+import argparse
+import hashlib
+import json
+import re
+import sys
+import tempfile
+from collections import Counter, OrderedDict
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+REPOSITORY_ROOT = PROJECT_ROOT.parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+sys.path.insert(0, str(PROJECT_ROOT / "examples"))
+
+from diff_model600_reports import read_blocks # noqa: E402
+from jorek_equations.model600_markdown import export_model600_markdown # noqa: E402
+
+REFERENCE = PROJECT_ROOT / "reference" / "model600_discrepancies.json"
+SOURCE = REPOSITORY_ROOT / "models" / "model600" / "mod_elt_matrix_fft.f90"
+FORMAT_VERSION = 1
+
+
+def block_name(key):
+ """Render a block key as the string used in the reference file."""
+
+ section, lhs, repeat = key
+ return "{} | {} | #{}".format(section, lhs, repeat)
+
+
+def _row_of(name):
+ """The equation row a block name belongs to, e.g. ``ti`` in amat(var_ti,..)."""
+
+ match = re.search(r"\(var_([a-z]+)", name.split("|")[1])
+ return match.group(1) if match else ""
+
+
+def measure(source_path=SOURCE, rows=None):
+ """Export the reports and return the discrepancies of every block."""
+
+ with tempfile.TemporaryDirectory() as folder:
+ folder = Path(folder)
+ source_report, generated_report = export_model600_markdown(
+ source_path, folder / "source.md", folder / "generated.md",
+ equations=rows,
+ )
+ source = read_blocks(source_report)
+ generated = read_blocks(generated_report)
+ blocks = OrderedDict()
+ for key in list(dict.fromkeys(list(source) + list(generated))):
+ source_only = Counter(source.get(key, ())) - Counter(generated.get(key, ()))
+ generated_only = Counter(generated.get(key, ())) - Counter(source.get(key, ()))
+ if not source_only and not generated_only:
+ continue
+ blocks[block_name(key)] = {
+ "source_only": sorted(source_only.elements()),
+ "generated_only": sorted(generated_only.elements()),
+ }
+ return blocks
+
+
+def build(blocks, source_path=SOURCE, annotations=None):
+ """Wrap the measured blocks with the metadata stored in the reference."""
+
+ annotations = annotations or {}
+ digest = hashlib.sha256(Path(source_path).read_bytes()).hexdigest()
+ document = OrderedDict()
+ document["format_version"] = FORMAT_VERSION
+ document["source"] = str(Path(source_path).relative_to(REPOSITORY_ROOT))
+ document["source_sha256"] = digest
+ document["block_count"] = len(blocks)
+ document["source_only_lines"] = sum(
+ len(item["source_only"]) for item in blocks.values()
+ )
+ document["generated_only_lines"] = sum(
+ len(item["generated_only"]) for item in blocks.values()
+ )
+ document["blocks"] = OrderedDict(
+ (name, OrderedDict((
+ ("findings", annotations.get(name, [])),
+ ("source_only", item["source_only"]),
+ ("generated_only", item["generated_only"]),
+ )))
+ for name, item in sorted(blocks.items())
+ )
+ return document
+
+
+def compare(reference, blocks):
+ """Return a list of human-readable differences, empty when they agree."""
+
+ stored = reference.get("blocks", {})
+ problems = []
+ for name in sorted(set(stored) | set(blocks)):
+ want = stored.get(name)
+ have = blocks.get(name)
+ if want is None:
+ problems.append(
+ "NEW BLOCK {}: {} source-only, {} generated-only".format(
+ name, len(have["source_only"]), len(have["generated_only"]),
+ )
+ )
+ problems.extend(" + " + t for t in have["source_only"])
+ problems.extend(" + " + t for t in have["generated_only"])
+ continue
+ if have is None:
+ problems.append(
+ "BLOCK NOW AGREES {} (the reference expects {} source-only and "
+ "{} generated-only)".format(
+ name, len(want["source_only"]), len(want["generated_only"]),
+ )
+ )
+ continue
+ for side in ("source_only", "generated_only"):
+ gone = Counter(want[side]) - Counter(have[side])
+ extra = Counter(have[side]) - Counter(want[side])
+ if not gone and not extra:
+ continue
+ problems.append("{} [{}] findings {}".format(
+ name, side, want.get("findings") or "unannotated",
+ ))
+ problems.extend(" - " + t for t in sorted(gone.elements()))
+ problems.extend(" + " + t for t in sorted(extra.elements()))
+ return problems
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+ parser.add_argument(
+ "--reference", type=Path, default=REFERENCE,
+ help="reference file to compare against (default: %(default)s)",
+ )
+ parser.add_argument(
+ "--source", type=Path, default=SOURCE,
+ help="element routine to read (default: the model-600 one)",
+ )
+ parser.add_argument(
+ "--equation", action="append", dest="rows", metavar="ROW",
+ help="check only this row; repeat for several. A partial run compares "
+ "only the blocks of the rows it exported, which is meant for "
+ "iterating, not for signing off a change",
+ )
+ parser.add_argument(
+ "--update", action="store_true",
+ help="rewrite the reference from the current tree",
+ )
+ arguments = parser.parse_args()
+
+ if arguments.rows and arguments.update:
+ parser.error("--update needs a full run; drop --equation")
+ print("Exporting model-600 terms and collecting discrepancies ...")
+ blocks = measure(arguments.source, arguments.rows)
+ source_only = sum(len(item["source_only"]) for item in blocks.values())
+ generated_only = sum(len(item["generated_only"]) for item in blocks.values())
+ print(" {} block(s) disagree: {} source-only and {} generated-only "
+ "monomials".format(len(blocks), source_only, generated_only))
+
+ if arguments.update:
+ previous = {}
+ if arguments.reference.exists():
+ previous = json.loads(arguments.reference.read_text(encoding="utf-8"))
+ annotations = {
+ name: item.get("findings", [])
+ for name, item in previous.get("blocks", {}).items()
+ }
+ document = build(blocks, arguments.source, annotations)
+ arguments.reference.parent.mkdir(parents=True, exist_ok=True)
+ arguments.reference.write_text(
+ json.dumps(document, indent=1) + "\n", encoding="utf-8",
+ )
+ missing = [n for n in document["blocks"] if not document["blocks"][n]["findings"]]
+ print("Wrote {}".format(arguments.reference))
+ if missing:
+ print(" {} block(s) carry no finding annotation:".format(len(missing)))
+ for name in missing:
+ print(" {}".format(name))
+ return 0
+
+ if not arguments.reference.exists():
+ print("No reference at {}; create one with --update.".format(
+ arguments.reference))
+ return 1
+ reference = json.loads(arguments.reference.read_text(encoding="utf-8"))
+
+ digest = hashlib.sha256(Path(arguments.source).read_bytes()).hexdigest()
+ if digest != reference.get("source_sha256"):
+ print("NOTE: {} has changed since the reference was written.".format(
+ reference.get("source")))
+ print(" reference sha256 {}".format(reference.get("source_sha256")))
+ print(" current sha256 {}".format(digest))
+
+ if arguments.rows:
+ wanted = {row.lower() for row in arguments.rows}
+ reference = dict(reference)
+ reference["blocks"] = {
+ name: item for name, item in reference["blocks"].items()
+ if _row_of(name) in wanted
+ }
+ print(" comparing the {} recorded block(s) of row(s) {}".format(
+ len(reference["blocks"]), ", ".join(sorted(wanted))))
+ problems = compare(reference, blocks)
+ if problems:
+ print("\nFINAL TEST: FAILED — {} block(s) differ from the "
+ "reference".format(sum(1 for p in problems if not p.startswith(" "))))
+ for line in problems:
+ print(" " + line)
+ print("\nA '-' line is a discrepancy the reference expects but the run no "
+ "longer produces; a '+' line is a new one.")
+ return 1
+ print("\nFINAL TEST: PASSED — every block reproduces the recorded "
+ "discrepancies exactly.")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/util/equation_codegen/pyproject.toml b/util/equation_codegen/pyproject.toml
new file mode 100644
index 0000000000..f9b653d221
--- /dev/null
+++ b/util/equation_codegen/pyproject.toml
@@ -0,0 +1,14 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "jorek-equation-codegen"
+version = "0.1.0"
+description = "Symbolic weak-form linearization for JOREK equations"
+requires-python = ">=3.9"
+dependencies = ["sympy>=1.12,<2"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
diff --git a/util/equation_codegen/reference/model600_discrepancies.json b/util/equation_codegen/reference/model600_discrepancies.json
new file mode 100644
index 0000000000..04c8f97ed9
--- /dev/null
+++ b/util/equation_codegen/reference/model600_discrepancies.json
@@ -0,0 +1,372 @@
+{
+ "format_version": 1,
+ "source": "models/model600/mod_elt_matrix_fft.f90",
+ "source_sha256": "d6a21cedf0026e97a0ef5def48277a65ab9c2a2a2c1e53e664f0c6d907512ca4",
+ "block_count": 24,
+ "source_only_lines": 8,
+ "generated_only_lines": 158,
+ "blocks": {
+ "Single-temperature (T) model | amat(var_psi,var_rhoimp) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-F0**2*v*rhoimp*T0_s*theta*tstep*alpha_e_bis*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-F0**2*v*rhoimp_s*T0*theta*tstep*alpha_e*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-F0**3*v*rhoimp*T0_p*xjac*theta*tstep*alpha_e_bis*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)",
+ "F0**2*v*rhoimp*T0_t*theta*tstep*alpha_e_bis*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "F0**2*v*rhoimp_t*T0*theta*tstep*alpha_e*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_psi,var_t) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-F0**2*v*rhoimp0*T*T0_s*theta*tstep*alpha_e_tri*ps0_t*tauIC/(2*F0**2*r0_corr + 2*ps0_x**2*r0_corr + 2*ps0_y**2*r0_corr)",
+ "-F0**2*v*rhoimp0*T_s*theta*tstep*alpha_e_bis*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-F0**2*v*rhoimp0_s*T*theta*tstep*alpha_e_bis*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-F0**3*v*rhoimp0*T*T0_p*xjac*theta*tstep*alpha_e_tri*tauIC/(2*BigR*F0**2*r0_corr + 2*BigR*ps0_x**2*r0_corr + 2*BigR*ps0_y**2*r0_corr)",
+ "-F0**3*v*rhoimp0_p*T*xjac*theta*tstep*alpha_e_bis*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)",
+ "F0**2*v*rhoimp0*T*T0_t*theta*tstep*alpha_e_tri*ps0_s*tauIC/(2*F0**2*r0_corr + 2*ps0_x**2*r0_corr + 2*ps0_y**2*r0_corr)",
+ "F0**2*v*rhoimp0*T_t*theta*tstep*alpha_e_bis*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "F0**2*v*rhoimp0_t*T*theta*tstep*alpha_e_bis*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_rho,var_psi) | #0": {
+ "findings": [
+ 4
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-BigR*v_x*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v_x*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v_y*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v_y*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_y**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v_x*psi_x*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v_y*psi_y*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_rho,var_rhoimp) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*BigR*v*rhoimp*T0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR*v*rhoimp_y*T0*xjac*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_rho,var_t) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*BigR*v*rhoimp0*T_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR*v*rhoimp0_y*T*xjac*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_u,var_rhoimp) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*BigR**4*v*u0_xy*rhoimp_y*T0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp*T0_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp_x*T0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp_xy*T0*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp_y*T0_x*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xy*rhoimp*T0_yy*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xy*rhoimp_yy*T0*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*w0_s*rhoimp*T0_t*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*w0_s*rhoimp_t*T0*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_xy*rhoimp_x*T0_x*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_x*u0_x*rhoimp*T0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_x*u0_x*rhoimp_y*T0*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_y*u0_y*rhoimp*T0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_y*u0_y*rhoimp_y*T0*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_xy*rhoimp*T0_xx*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_xy*rhoimp_xx*T0*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp*T0_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp_x*T0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp_xy*T0*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp_y*T0_x*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*w0_t*rhoimp*T0_s*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*w0_t*rhoimp_s*T0*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_u,var_t) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*BigR**4*v*u0_xy*rhoimp0_y*T_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp0*T_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp0_x*T_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp0_xy*T*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xx*rhoimp0_y*T_x*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xy*rhoimp0*T_yy*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*u0_xy*rhoimp0_yy*T*xjac*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*w0_s*rhoimp0*T_t*theta*tstep*alpha_i*tauIC",
+ "-BigR**4*v*w0_s*rhoimp0_t*T*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_xy*rhoimp0_x*T_x*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_x*u0_x*rhoimp0*T_y*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_x*u0_x*rhoimp0_y*T*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_y*u0_y*rhoimp0*T_y*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**3*v_y*u0_y*rhoimp0_y*T*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_xy*rhoimp0*T_xx*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_xy*rhoimp0_xx*T*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp0*T_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp0_x*T_y*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp0_xy*T*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*u0_yy*rhoimp0_y*T_x*xjac*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*w0_t*rhoimp0*T_s*theta*tstep*alpha_i*tauIC",
+ "BigR**4*v*w0_t*rhoimp0_s*T*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_vpar,var_psi) | #0": {
+ "findings": [
+ 4
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-BigR*v*vpar0_x*psi_x*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v*vpar0_y*psi_y*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_x*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_x*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_y*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_y*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_y**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_vpar,var_rho) | #0": {
+ "findings": [
+ 9
+ ],
+ "source_only": [
+ "BigR*v*vpar0_x*rho*xjac*theta*tstep*V_prof_pinch*ps0_x/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_y*rho*xjac*theta*tstep*V_prof_pinch*ps0_y/(sqrt(ps0_x**2 + ps0_y**2))"
+ ],
+ "generated_only": [
+ "-BigR*v*vpar0_x*rho*xjac*theta*tstep*V_prof_pinch*ps0_x/(sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v*vpar0_y*rho*xjac*theta*tstep*V_prof_pinch*ps0_y/(sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Single-temperature (T) model | amat(var_vpar,var_vpar) | #0": {
+ "findings": [
+ 9
+ ],
+ "source_only": [
+ "BigR*v*vpar_x*xjac*theta*tstep*V_prof_pinch*ps0_x*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar_y*xjac*theta*tstep*V_prof_pinch*ps0_y*r0/(sqrt(ps0_x**2 + ps0_y**2))"
+ ],
+ "generated_only": [
+ "-BigR*v*vpar_x*xjac*theta*tstep*V_prof_pinch*ps0_x*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v*vpar_y*xjac*theta*tstep*V_prof_pinch*ps0_y*r0/(sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Single-temperature (T) model | amat_n(var_psi,var_rhoimp) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-F0**3*v*rhoimp_p*T0*xjac*theta*tstep*alpha_e*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)"
+ ]
+ },
+ "Single-temperature (T) model | amat_n(var_psi,var_t) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-F0**3*v*rhoimp0*T_p*xjac*theta*tstep*alpha_e_bis*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_psi,var_rhoimp) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*F0**2*v*rhoimp*Te0_s*theta*tstep*alpha_e_bis*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-2*F0**2*v*rhoimp_s*Te0*theta*tstep*alpha_e*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-2*F0**3*v*rhoimp*Te0_p*xjac*theta*tstep*alpha_e_bis*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)",
+ "2*F0**2*v*rhoimp*Te0_t*theta*tstep*alpha_e_bis*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "2*F0**2*v*rhoimp_t*Te0*theta*tstep*alpha_e*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_psi,var_te) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*F0**2*v*rhoimp0*Te*Te0_s*theta*tstep*alpha_e_tri*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-2*F0**2*v*rhoimp0*Te_s*theta*tstep*alpha_e_bis*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-2*F0**2*v*rhoimp0_s*Te*theta*tstep*alpha_e_bis*ps0_t*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "-2*F0**3*v*rhoimp0*Te*Te0_p*xjac*theta*tstep*alpha_e_tri*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)",
+ "-2*F0**3*v*rhoimp0_p*Te*xjac*theta*tstep*alpha_e_bis*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)",
+ "2*F0**2*v*rhoimp0*Te*Te0_t*theta*tstep*alpha_e_tri*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "2*F0**2*v*rhoimp0*Te_t*theta*tstep*alpha_e_bis*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)",
+ "2*F0**2*v*rhoimp0_t*Te*theta*tstep*alpha_e_bis*ps0_s*tauIC/(F0**2*r0_corr + ps0_x**2*r0_corr + ps0_y**2*r0_corr)"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_rho,var_psi) | #0": {
+ "findings": [
+ 4
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-BigR*v_x*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v_x*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v_y*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v_y*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_y**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v_x*psi_x*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v_y*psi_y*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_rho,var_rhoimp) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-4*BigR*v*rhoimp*Ti0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-4*BigR*v*rhoimp_y*Ti0*xjac*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_rho,var_ti) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-4*BigR*v*rhoimp0*Ti_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-4*BigR*v*rhoimp0_y*Ti*xjac*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_u,var_rhoimp) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*BigR**4*v*u0_xx*rhoimp*Ti0_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xx*rhoimp_x*Ti0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xx*rhoimp_xy*Ti0*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xx*rhoimp_y*Ti0_x*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xy*rhoimp*Ti0_yy*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xy*rhoimp_yy*Ti0*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*w0_s*rhoimp*Ti0_t*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*w0_s*rhoimp_t*Ti0*theta*tstep*alpha_i*tauIC",
+ "-4*BigR**4*v*u0_xy*rhoimp_y*Ti0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_x*u0_x*rhoimp*Ti0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_x*u0_x*rhoimp_y*Ti0*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_y*u0_y*rhoimp*Ti0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_y*u0_y*rhoimp_y*Ti0*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_xy*rhoimp*Ti0_xx*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_xy*rhoimp_xx*Ti0*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp*Ti0_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp_x*Ti0_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp_xy*Ti0*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp_y*Ti0_x*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*w0_t*rhoimp*Ti0_s*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*w0_t*rhoimp_s*Ti0*theta*tstep*alpha_i*tauIC",
+ "4*BigR**4*v*u0_xy*rhoimp_x*Ti0_x*xjac*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_u,var_ti) | #0": {
+ "findings": [
+ 1
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*BigR**4*v*u0_xx*rhoimp0*Ti_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xx*rhoimp0_x*Ti_y*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xx*rhoimp0_xy*Ti*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xx*rhoimp0_y*Ti_x*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xy*rhoimp0*Ti_yy*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*u0_xy*rhoimp0_yy*Ti*xjac*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*w0_s*rhoimp0*Ti_t*theta*tstep*alpha_i*tauIC",
+ "-2*BigR**4*v*w0_s*rhoimp0_t*Ti*theta*tstep*alpha_i*tauIC",
+ "-4*BigR**4*v*u0_xy*rhoimp0_y*Ti_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_x*u0_x*rhoimp0*Ti_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_x*u0_x*rhoimp0_y*Ti*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_y*u0_y*rhoimp0*Ti_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**3*v_y*u0_y*rhoimp0_y*Ti*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_xy*rhoimp0*Ti_xx*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_xy*rhoimp0_xx*Ti*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp0*Ti_xy*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp0_x*Ti_y*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp0_xy*Ti*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*u0_yy*rhoimp0_y*Ti_x*xjac*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*w0_t*rhoimp0*Ti_s*theta*tstep*alpha_i*tauIC",
+ "2*BigR**4*v*w0_t*rhoimp0_s*Ti*theta*tstep*alpha_i*tauIC",
+ "4*BigR**4*v*u0_xy*rhoimp0_x*Ti_x*xjac*theta*tstep*alpha_i*tauIC"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_vpar,var_psi) | #0": {
+ "findings": [
+ 4
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-BigR*v*vpar0_x*psi_x*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v*vpar0_y*psi_y*xjac*theta*tstep*V_prof_pinch*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_x*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_x*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_y*psi_x*xjac*theta*tstep*V_prof_pinch*ps0_x*ps0_y*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_y*psi_y*xjac*theta*tstep*V_prof_pinch*ps0_y**2*r0/(ps0_x**2*sqrt(ps0_x**2 + ps0_y**2) + ps0_y**2*sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_vpar,var_rho) | #0": {
+ "findings": [
+ 9
+ ],
+ "source_only": [
+ "BigR*v*vpar0_x*rho*xjac*theta*tstep*V_prof_pinch*ps0_x/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar0_y*rho*xjac*theta*tstep*V_prof_pinch*ps0_y/(sqrt(ps0_x**2 + ps0_y**2))"
+ ],
+ "generated_only": [
+ "-BigR*v*vpar0_x*rho*xjac*theta*tstep*V_prof_pinch*ps0_x/(sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v*vpar0_y*rho*xjac*theta*tstep*V_prof_pinch*ps0_y/(sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat(var_vpar,var_vpar) | #0": {
+ "findings": [
+ 9
+ ],
+ "source_only": [
+ "BigR*v*vpar_x*xjac*theta*tstep*V_prof_pinch*ps0_x*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "BigR*v*vpar_y*xjac*theta*tstep*V_prof_pinch*ps0_y*r0/(sqrt(ps0_x**2 + ps0_y**2))"
+ ],
+ "generated_only": [
+ "-BigR*v*vpar_x*xjac*theta*tstep*V_prof_pinch*ps0_x*r0/(sqrt(ps0_x**2 + ps0_y**2))",
+ "-BigR*v*vpar_y*xjac*theta*tstep*V_prof_pinch*ps0_y*r0/(sqrt(ps0_x**2 + ps0_y**2))"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat_n(var_psi,var_rhoimp) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*F0**3*v*rhoimp_p*Te0*xjac*theta*tstep*alpha_e*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)"
+ ]
+ },
+ "Two-temperature (Ti/Te) model | amat_n(var_psi,var_te) | #0": {
+ "findings": [
+ 2
+ ],
+ "source_only": [],
+ "generated_only": [
+ "-2*F0**3*v*rhoimp0*Te_p*xjac*theta*tstep*alpha_e_bis*tauIC/(BigR*F0**2*r0_corr + BigR*ps0_x**2*r0_corr + BigR*ps0_y**2*r0_corr)"
+ ]
+ }
+ }
+}
diff --git a/util/equation_codegen/run_test.sh b/util/equation_codegen/run_test.sh
new file mode 100755
index 0000000000..698dc6772b
--- /dev/null
+++ b/util/equation_codegen/run_test.sh
@@ -0,0 +1,221 @@
+#!/bin/bash
+#
+# Regression test for the equation_codegen tool.
+#
+# Unlike the JOREK non-regression tests under reg_tests/ (which compile
+# binaries, run an MPI job and diff HDF5 restart files), this test is a pure
+# Python check: it recomputes every model-600 discrepancy between the
+# hand-written element routine (models/model600/mod_elt_matrix_fft.f90) and
+# the symbolic linearization (util/equation_codegen/src/jorek_equations),
+# and compares it against the frozen reference in
+# reference/model600_discrepancies.json. A change to either side that shifts
+# which terms disagree — a real regression, not just a reordering — makes
+# this test fail and print exactly which monomials are no longer accounted
+# for.
+#
+# Follows the same pass/fail convention as reg_tests/run_test.sh: exit 0 and
+# a line "Test 'equation_codegen' passed." on success, non-zero and a
+# detailed failure report otherwise. Meant to be run directly, or as its own
+# step in CI — it does not go through reg_tests/run_test.sh.
+#
+# The only runtime dependency is the ``sympy`` package; final_test.py adds
+# src/ to sys.path itself, so nothing needs to be "installed" as a package.
+#
+# Usage:
+# util/equation_codegen/run_test.sh # check the current tree
+# util/equation_codegen/run_test.sh --update # rewrite the reference
+# # after an intended change
+# util/equation_codegen/run_test.sh --parallel # check the 11 rows as
+# # separate background
+# # processes (one per
+# # equation), for a CI
+# # agent with several
+# # cores. Equivalent in
+# # coverage to a plain
+# # run; not combinable
+# # with --update or
+# # --equation.
+
+set -u
+
+ROWS="psi u rho vpar Ti Te T rhoimp rhon zj w "
+
+TESTNAME="equation_codegen"
+STARTDIR=$(readlink -f "$(dirname "$0")")
+cd "$STARTDIR" || exit 1
+
+PYTHON=python3
+
+# --- Prefer the system Python if it already has sympy; this is the common
+# case and avoids touching venv/pip at all.
+if ! $PYTHON -c "import sympy" >/dev/null 2>&1; then
+
+ # --- Try the environment-modules ``sympy`` package next (e.g. Lmod's
+ # "module load sympy/1.14.0-gfbf-2025b" on the ITER Bamboo agents),
+ # which needs no network access and no venv at all. ``module`` is a
+ # shell function normally set up by an interactive login shell, so a
+ # non-interactive script (like this one, run from a Bamboo task) may
+ # not have it yet; source Lmod's init script first if not. Lmod's own
+ # scripts are not written for ``set -u``, so relax it around this.
+ set +u
+ # A non-interactive, non-login shell (as Bamboo runs this) does not
+ # source /etc/profile.d, so MODULEPATH itself may be unset even before
+ # getting to the ``module`` function below; set it up the same way a
+ # login shell would.
+ if [ -z "$MODULEPATH" ]; then
+ for f in /etc/profile.d/*.sh; do
+ [ -r "$f" ] && . "$f" >/dev/null 2>&1
+ done
+ fi
+ if ! type module >/dev/null 2>&1; then
+ for init in \
+ "${LMOD_PKG:-}/init/bash" \
+ /usr/share/lmod/lmod/init/bash \
+ /etc/profile.d/lmod.sh \
+ /usr/share/Modules/init/bash
+ do
+ [ -n "$init" ] && [ -r "$init" ] && . "$init" && break
+ done
+ fi
+ if type module >/dev/null 2>&1; then
+ module load sympy/1.14.0-gfbf-2025b 2>/dev/null
+ if ! $PYTHON -c "import sympy" >/dev/null 2>&1; then
+ # Fall back to whatever sympy module this host happens to have,
+ # in case the pinned version above is not the one available here.
+ sympy_module=$(module -t avail sympy 2>&1 | grep -m1 '^sympy/')
+ [ -n "$sympy_module" ] && module load "$sympy_module" 2>/dev/null
+ fi
+ fi
+ set -u
+fi
+
+if ! $PYTHON -c "import sympy" >/dev/null 2>&1; then
+
+ # --- Otherwise, set up a virtual environment the first time this runs.
+ # --system-site-packages lets it fall back to whatever the system
+ # Python already provides (sympy included, if installed there).
+ if [ ! -x ".venv/bin/python" ]; then
+ echo "Setting up .venv for ${TESTNAME} ..."
+ $PYTHON -m venv --system-site-packages .venv || exit 1
+
+ # Some Python installs (e.g. RHEL's python3 without python3-pip) create
+ # a venv with no pip inside it. Bootstrap it explicitly rather than
+ # assume it is there.
+ if ! .venv/bin/python -m pip --version >/dev/null 2>&1; then
+ .venv/bin/python -m ensurepip --upgrade >/dev/null 2>&1
+ fi
+
+ if .venv/bin/python -m pip --version >/dev/null 2>&1; then
+ # A plain package install, not an editable/local one: this only
+ # needs sympy on sys.path, so there is nothing to build.
+ .venv/bin/python -m pip install --quiet 'sympy>=1.12,<2' || exit 1
+ elif ! .venv/bin/python -c "import sympy" >/dev/null 2>&1; then
+ echo "ERROR: no working pip in .venv and no system sympy available;"
+ echo " install sympy for ${PYTHON} manually and re-run."
+ exit 1
+ fi
+ fi
+
+ PYTHON=.venv/bin/python
+fi
+
+# --- Pull --parallel out of the arguments; everything else passes through
+# to final_test.py as before.
+parallel="no"
+args=()
+for arg in "$@"; do
+ if [ "$arg" = "--parallel" ]; then
+ parallel="yes"
+ else
+ args+=("$arg")
+ fi
+done
+
+# --- The weak-form documentation is generated from model600.py; fail early if
+# it has not been re-rendered after a change to the equations.
+if ! $PYTHON examples/model600_docs.py check; then
+ echo "Test '${TESTNAME}' FAILED."
+ exit 1
+fi
+
+if [ "$parallel" = "yes" ]; then
+ for arg in "${args[@]:-}"; do
+ case "$arg" in
+ --update)
+ echo "ERROR: --parallel cannot be combined with --update" \
+ "(concurrent writers would race on the reference file)."
+ exit 1
+ ;;
+ --equation)
+ echo "ERROR: --parallel already checks every row;" \
+ "drop --equation."
+ exit 1
+ ;;
+ esac
+ done
+ # --- Check every row as its own background process. Each row's blocks
+ # are independent of every other row's, so this finds exactly what a
+ # single full run would; it is just faster on a multi-core agent.
+ #
+ # Each row gets one CPU: at most NPROC rows run at a time (NPROC
+ # from the CPU affinity mask this process actually has, which on
+ # Linux already reflects a cgroup/container CPU quota, not just the
+ # host's total core count), and when ``taskset`` is available each
+ # one is pinned to its own core so it cannot spill onto a neighbour's.
+ # SymPy itself is single-threaded, but a couple of underlying C
+ # libraries thread on their own unless told not to, so that is
+ # disabled too.
+ NPROC=$(nproc 2>/dev/null || getconf _NPROCESSORS_ONLN 2>/dev/null || echo 1)
+ export OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
+ NUMEXPR_NUM_THREADS=1
+
+ logdir=$(mktemp -d)
+ trap 'rm -rf "$logdir"' EXIT
+ cpu=0
+ running=0
+ for row in $ROWS; do
+ (
+ if command -v taskset >/dev/null 2>&1; then
+ taskset -c "$((cpu % NPROC))" \
+ $PYTHON final_test.py --equation "$row" "${args[@]}" \
+ > "$logdir/$row.log" 2>&1
+ else
+ $PYTHON final_test.py --equation "$row" "${args[@]}" \
+ > "$logdir/$row.log" 2>&1
+ fi
+ echo $? > "$logdir/$row.status"
+ ) &
+ cpu=$((cpu + 1))
+ running=$((running + 1))
+ if [ "$running" -ge "$NPROC" ]; then
+ wait -n
+ running=$((running - 1))
+ fi
+ done
+ wait
+
+ status=0
+ for row in $ROWS; do
+ rc=$(cat "$logdir/$row.status" 2>/dev/null || echo 1)
+ if [ "$rc" != "0" ]; then
+ status=1
+ echo "--- ${row}: FAILED ---"
+ cat "$logdir/$row.log"
+ else
+ tail -n 1 "$logdir/$row.log"
+ fi
+ done
+else
+ # --- Run the regression check. final_test.py prints, on failure, exactly
+ # which monomials each disagreeing block is missing or has gained.
+ $PYTHON final_test.py "${args[@]}"
+ status=$?
+fi
+
+if [ $status -eq 0 ]; then
+ echo "Test '${TESTNAME}' passed."
+else
+ echo "Test '${TESTNAME}' FAILED."
+fi
+
+exit $status
diff --git a/util/equation_codegen/setup.cfg b/util/equation_codegen/setup.cfg
new file mode 100644
index 0000000000..705dd50891
--- /dev/null
+++ b/util/equation_codegen/setup.cfg
@@ -0,0 +1,15 @@
+[metadata]
+name = jorek-equation-codegen
+version = 0.1.0
+description = Symbolic weak-form linearization for JOREK equations
+
+[options]
+package_dir =
+ = src
+packages = find:
+python_requires = >=3.9
+install_requires =
+ sympy>=1.12,<2
+
+[options.packages.find]
+where = src
diff --git a/util/equation_codegen/src/jorek_equations/__init__.py b/util/equation_codegen/src/jorek_equations/__init__.py
new file mode 100644
index 0000000000..7427bf095f
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/__init__.py
@@ -0,0 +1,105 @@
+"""Public API for the JOREK weak-equation symbolic DSL."""
+
+from .channels import ToroidalChannels, split_toroidal_channels
+from .equations import ConstraintEquation, EvolutionEquation, LinearizedEquation
+from .exceptions import (
+ EquationDslError,
+ InvalidDeclarationError,
+ MissingDerivativeError,
+ UnsupportedExpressionError,
+)
+from .external import ExternalFunction, LinearizationPolicy, external_function
+from .fortran import FortranPrinter, fortran
+from .linearization import jacobian_blocks, resolve_current, variation
+from .operators import (
+ R,
+ Z,
+ phi,
+ SpatialDerivative,
+ PoloidalBracket,
+ bracket,
+ dR,
+ dZ,
+ derivative,
+ dot,
+ dphi,
+ poiss_bracket_st,
+ poiss_bracket,
+ laplacian,
+ ds,
+ dt,
+ expand_brackets,
+ expand_derivatives,
+ grad,
+ lower_brackets_to_element,
+)
+from .symbols import (
+ Field,
+ FieldRole,
+ FieldValue,
+ Frozen,
+ TestFunction,
+ coefficient,
+ delta,
+ field,
+ freeze,
+ frozen_abs,
+ previous_delta,
+ test_function,
+ trial,
+)
+from .fortran_source import FortranComparisonError
+
+__all__ = [
+ "ConstraintEquation",
+ "EquationDslError",
+ "EvolutionEquation",
+ "ExternalFunction",
+ "FortranPrinter",
+ "FortranComparisonError",
+ "Field",
+ "FieldRole",
+ "FieldValue",
+ "Frozen",
+ "InvalidDeclarationError",
+ "LinearizationPolicy",
+ "LinearizedEquation",
+ "MissingDerivativeError",
+ "R",
+ "SpatialDerivative",
+ "PoloidalBracket",
+ "TestFunction",
+ "ToroidalChannels",
+ "UnsupportedExpressionError",
+ "Z",
+ "bracket",
+ "coefficient",
+ "dR",
+ "dZ",
+ "delta",
+ "derivative",
+ "dot",
+ "dphi",
+ "poiss_bracket_st",
+ "poiss_bracket",
+ "laplacian",
+ "ds",
+ "dt",
+ "external_function",
+ "expand_brackets",
+ "expand_derivatives",
+ "field",
+ "freeze",
+ "frozen_abs",
+ "fortran",
+ "grad",
+ "jacobian_blocks",
+ "lower_brackets_to_element",
+ "phi",
+ "previous_delta",
+ "resolve_current",
+ "test_function",
+ "split_toroidal_channels",
+ "trial",
+ "variation",
+]
diff --git a/util/equation_codegen/src/jorek_equations/channels.py b/util/equation_codegen/src/jorek_equations/channels.py
new file mode 100644
index 0000000000..c035e3adcb
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/channels.py
@@ -0,0 +1,61 @@
+"""Classification by toroidal derivatives on test and trial functions."""
+
+from dataclasses import dataclass
+
+import sympy as sp
+
+from .operators import SpatialDerivative, expand_derivatives, phi
+from .symbols import FieldRole, FieldValue, TestFunction
+
+
+@dataclass(frozen=True)
+class ToroidalChannels:
+ """The four JOREK FFT assembly channels for one expression."""
+
+ p: sp.Expr = sp.S.Zero
+ n: sp.Expr = sp.S.Zero
+ k: sp.Expr = sp.S.Zero
+ kn: sp.Expr = sp.S.Zero
+
+
+def _term_channel(term):
+ test_order = 0
+ trial_order = 0
+ for item in term.atoms(SpatialDerivative):
+ if item.coordinate != phi:
+ continue
+ if item.expression.has(TestFunction):
+ test_order += 1
+ if any(
+ value.role is FieldRole.TRIAL
+ for value in item.expression.atoms(FieldValue)
+ ):
+ trial_order += 1
+
+ if test_order > 1 or trial_order > 1:
+ raise ValueError(
+ "Unsupported toroidal derivative order in term {}".format(term)
+ )
+ if test_order and trial_order:
+ return "kn"
+ if test_order:
+ return "k"
+ if trial_order:
+ return "n"
+ return "p"
+
+
+def split_toroidal_channels(expression) -> ToroidalChannels:
+ """Split an expanded sum into p, n, k, and kn contributions."""
+
+ # Apply the product rule before classifying a term. For example,
+ # ``dphi(T * delta_rho)`` must become ``T_phi*delta_rho`` (p channel)
+ # plus ``T*dphi(delta_rho)`` (n channel). Classifying the unevaluated
+ # derivative would incorrectly send both contributions to ``n``.
+ expression = expand_derivatives(expression)
+ groups = {"p": [], "n": [], "k": [], "kn": []}
+ for term in sp.Add.make_args(sp.expand(expression)):
+ groups[_term_channel(term)].append(term)
+ return ToroidalChannels(
+ **{name: sp.Add(*terms) for name, terms in groups.items()}
+ )
diff --git a/util/equation_codegen/src/jorek_equations/equations.py b/util/equation_codegen/src/jorek_equations/equations.py
new file mode 100644
index 0000000000..37393a0f2f
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/equations.py
@@ -0,0 +1,102 @@
+"""Weak equation containers and JOREK sign transformations."""
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, Mapping
+
+import sympy as sp
+
+from .linearization import resolve_current, variation
+from .symbols import Field, FieldRole, TestFunction
+
+
+@dataclass(frozen=True)
+class LinearizedEquation:
+ """Symbolic local RHS and AMAT blocks for one weak equation."""
+
+ name: str
+ rhs: sp.Expr
+ amat: Dict[Field, sp.Expr]
+ kind: str = "evolution"
+
+
+@dataclass(frozen=True)
+class EvolutionEquation:
+ """Weak evolution equation dA/dt = B."""
+
+ name: str
+ test: TestFunction
+ A: sp.Expr
+ B: sp.Expr
+ kind: str = "evolution"
+ # Optional field-specific residual expressions used when a legacy model
+ # defines a Jacobian tangent that intentionally differs from dB/dq.
+ amat_variation_overrides: Mapping[Field, sp.Expr] = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.test, TestFunction):
+ raise TypeError("test must be a TestFunction")
+ object.__setattr__(self, "A", sp.sympify(self.A))
+ object.__setattr__(self, "B", sp.sympify(self.B))
+ overrides = self.amat_variation_overrides or {}
+ object.__setattr__(
+ self,
+ "amat_variation_overrides",
+ {field: sp.sympify(expression) for field, expression in overrides.items()},
+ )
+
+ def linearize(self, *, fields: Iterable[Field], timestep, theta, zeta):
+ """Apply the model-199 evolution-equation sign convention."""
+
+ fields = tuple(fields)
+ timestep = sp.sympify(timestep)
+ theta = sp.sympify(theta)
+ zeta = sp.sympify(zeta)
+
+ history = sp.Add(
+ *(
+ variation(self.A, value, direction=FieldRole.PREVIOUS_DELTA)
+ for value in fields
+ )
+ )
+ rhs = sp.expand(timestep * resolve_current(self.B) + zeta * history)
+ amat = {
+ value: sp.expand(
+ (1 + zeta)
+ * variation(self.A, value, direction=FieldRole.TRIAL)
+ - theta
+ * timestep
+ * variation(
+ self.amat_variation_overrides.get(value, self.B),
+ value,
+ direction=FieldRole.TRIAL,
+ )
+ )
+ for value in fields
+ }
+ return LinearizedEquation(self.name, rhs, amat, self.kind)
+
+
+@dataclass(frozen=True)
+class ConstraintEquation:
+ """Weak algebraic constraint C(q) = 0."""
+
+ name: str
+ test: TestFunction
+ C: sp.Expr
+ kind: str = "static"
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.test, TestFunction):
+ raise TypeError("test must be a TestFunction")
+ object.__setattr__(self, "C", sp.sympify(self.C))
+
+ def linearize(self, *, fields: Iterable[Field]):
+ """Return RHS=-C and AMAT=D(C) for the Newton correction."""
+
+ fields = tuple(fields)
+ rhs = -resolve_current(self.C)
+ amat = {
+ value: variation(self.C, value, direction=FieldRole.TRIAL)
+ for value in fields
+ }
+ return LinearizedEquation(self.name, rhs, amat, self.kind)
diff --git a/util/equation_codegen/src/jorek_equations/exceptions.py b/util/equation_codegen/src/jorek_equations/exceptions.py
new file mode 100644
index 0000000000..e7386e4ceb
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/exceptions.py
@@ -0,0 +1,18 @@
+"""Exceptions raised by the JOREK equation DSL."""
+
+
+class EquationDslError(Exception):
+ """Base class for errors produced by the equation DSL."""
+
+
+class InvalidDeclarationError(EquationDslError, ValueError):
+ """A symbolic object was declared inconsistently."""
+
+
+class MissingDerivativeError(EquationDslError):
+ """An active external dependency has no supplied derivative."""
+
+
+class UnsupportedExpressionError(EquationDslError):
+ """An expression cannot be linearized safely."""
+
diff --git a/util/equation_codegen/src/jorek_equations/external.py b/util/equation_codegen/src/jorek_equations/external.py
new file mode 100644
index 0000000000..57c02aaed9
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/external.py
@@ -0,0 +1,117 @@
+"""Dependency declarations for externally evaluated functions."""
+
+from dataclasses import dataclass, field as dataclass_field
+from enum import Enum
+from typing import Dict, Mapping, Tuple, Union
+
+import sympy as sp
+
+from .exceptions import InvalidDeclarationError
+from .symbols import validate_name
+
+
+class LinearizationPolicy(str, Enum):
+ ACTIVE = "active"
+ FROZEN = "frozen"
+ PIECEWISE_ACTIVE = "piecewise_active"
+
+
+_FUNCTION_REGISTRY: Dict[object, "ExternalFunction"] = {}
+_NAME_REGISTRY: Dict[str, "ExternalFunction"] = {}
+
+
+@dataclass(frozen=True)
+class ExternalFunction:
+ """Declaration of a value supplied outside the symbolic equation layer."""
+
+ name: str
+ arguments: Tuple[str, ...]
+ derivatives: Mapping[str, str]
+ policy: LinearizationPolicy = LinearizationPolicy.ACTIVE
+ fortran_name: str = ""
+ _sympy_function: object = dataclass_field(init=False, repr=False, compare=False)
+
+ def __post_init__(self) -> None:
+ validate_name(self.name, "external-function")
+ arguments = tuple(self.arguments)
+ if len(set(arguments)) != len(arguments):
+ raise InvalidDeclarationError(
+ "External function {} has duplicate arguments".format(self.name)
+ )
+ for argument in arguments:
+ validate_name(argument, "external-function argument")
+
+ derivatives = dict(self.derivatives)
+ unknown = set(derivatives) - set(arguments)
+ if unknown:
+ raise InvalidDeclarationError(
+ "Derivatives for {} refer to unknown arguments: {}".format(
+ self.name, ", ".join(sorted(unknown))
+ )
+ )
+ for derivative_name in derivatives.values():
+ validate_name(derivative_name, "external derivative")
+
+ policy = LinearizationPolicy(self.policy)
+ fortran_name = self.fortran_name or self.name
+ validate_name(fortran_name, "Fortran external-function")
+
+ object.__setattr__(self, "arguments", arguments)
+ object.__setattr__(self, "derivatives", derivatives)
+ object.__setattr__(self, "policy", policy)
+ object.__setattr__(self, "fortran_name", fortran_name)
+ object.__setattr__(self, "_sympy_function", sp.Function(self.name))
+
+ def __call__(self, *values):
+ if len(values) != len(self.arguments):
+ raise InvalidDeclarationError(
+ "{} expects {} arguments ({}), received {}".format(
+ self.name,
+ len(self.arguments),
+ ", ".join(self.arguments),
+ len(values),
+ )
+ )
+ return self._sympy_function(*(sp.sympify(value) for value in values))
+
+ def derivative_call(self, argument: str, values):
+ """Return the symbolic value of one supplied partial derivative."""
+
+ return sp.Function(self.derivatives[argument])(*values)
+
+
+def external_function(
+ name: str,
+ *,
+ arguments: Tuple[str, ...],
+ derivatives: Mapping[str, str],
+ policy: Union[str, LinearizationPolicy] = LinearizationPolicy.ACTIVE,
+ fortran_name: str = ""
+) -> ExternalFunction:
+ """Declare and register an externally evaluated symbolic function."""
+
+ candidate = ExternalFunction(
+ name=name,
+ arguments=arguments,
+ derivatives=derivatives,
+ policy=LinearizationPolicy(policy),
+ fortran_name=fortran_name,
+ )
+ existing = _NAME_REGISTRY.get(name)
+ if existing is not None:
+ if existing != candidate:
+ raise InvalidDeclarationError(
+ "External function {!r} was redeclared inconsistently".format(name)
+ )
+ return existing
+
+ _NAME_REGISTRY[name] = candidate
+ _FUNCTION_REGISTRY[candidate._sympy_function] = candidate
+ return candidate
+
+
+def definition_for_call(expression):
+ """Return the registered declaration for a SymPy function call."""
+
+ return _FUNCTION_REGISTRY.get(expression.func)
+
diff --git a/util/equation_codegen/src/jorek_equations/fortran.py b/util/equation_codegen/src/jorek_equations/fortran.py
new file mode 100644
index 0000000000..ffdc29198a
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/fortran.py
@@ -0,0 +1,130 @@
+"""Fortran-style rendering of the symbolic DSL expressions."""
+
+import sympy as sp
+from sympy.printing.str import StrPrinter
+
+from .external import definition_for_call
+from .operators import R, Z, phi, s, t, PoloidalBracket, SpatialDerivative, expand_derivatives
+from .symbols import FieldRole, FieldValue, Frozen, TestFunction
+
+
+DEFAULT_FIELD_NAMES = {
+ "psi": ("ps0", "psi"),
+ "u": ("u0", "u"),
+ "j": ("zj0", "zj"),
+ "omega": ("w0", "w"),
+ "rho": ("r0", "rho"),
+ "T": ("T0", "T"),
+}
+
+COORDINATE_NAMES = {R: "BigR", Z: "Z", phi: "phi"}
+DERIVATIVE_SUFFIXES = {R: "x", Z: "y", phi: "p", s: "s", t: "t"}
+
+
+class FortranPrinter(StrPrinter):
+ """Render the current model-199 symbolic naming convention."""
+
+ def __init__(
+ self, *, field_names=None, coordinate_names=None, previous_names=None, **settings
+ ):
+ super().__init__(settings)
+ self.field_names = dict(DEFAULT_FIELD_NAMES)
+ if field_names:
+ self.field_names.update(field_names)
+ self.coordinate_names = dict(COORDINATE_NAMES)
+ if coordinate_names:
+ self.coordinate_names.update(coordinate_names)
+ self.previous_names = dict(previous_names or {})
+
+ def _print(self, expr, **kwargs):
+ # DSL atoms provide _sympystr methods for normal SymPy output. Intercept
+ # them here so the Fortran backend can apply its own naming convention.
+ if isinstance(expr, SpatialDerivative):
+ return self._print_SpatialDerivative(expr)
+ if isinstance(expr, PoloidalBracket):
+ return self._print_PoloidalBracket(expr)
+ if isinstance(expr, FieldValue):
+ return self._print_FieldValue(expr)
+ if isinstance(expr, Frozen):
+ return self._print_Frozen(expr)
+ if isinstance(expr, TestFunction):
+ return self._print_TestFunction(expr)
+ if isinstance(expr, sp.Function):
+ return self._print_Function(expr)
+ return super()._print(expr, **kwargs)
+
+ def _print_Symbol(self, expr):
+ return self.coordinate_names.get(expr, super()._print_Symbol(expr))
+
+ def _print_TestFunction(self, expr):
+ return expr.name
+
+ def _print_Frozen(self, expr):
+ # Frozen values are evaluated by the host before assembly. Do not
+ # expose the Python ``freeze(Abs(...))`` wrapper in generated Fortran.
+ return self._print(expr.expression)
+
+ def _print_Abs(self, expr):
+ # Model-199 stores abs(rho) and abs(T) in the current-state arrays;
+ # their DSL derivatives are intentionally disabled.
+ return self._print(expr.args[0])
+
+ def _print_FieldValue(self, expr):
+ current, trial = self.field_names.get(
+ expr.field_name, (expr.field_name + "0", expr.field_name)
+ )
+ if expr.role is FieldRole.CURRENT:
+ return current
+ if expr.role is FieldRole.TRIAL:
+ return trial
+ if expr.role is FieldRole.PREVIOUS_DELTA:
+ if expr.field_name in self.previous_names:
+ return self.previous_names[expr.field_name]
+ return "delta_{}_prev".format(trial)
+ if expr.role is FieldRole.DELTA:
+ return "delta_{}".format(trial)
+ return trial
+
+ def _print_SpatialDerivative(self, expr):
+ base, suffix = self._derivative_parts(expr)
+ return self._print(base) + "_" + suffix
+
+ def _print_PoloidalBracket(self, expr):
+ return "bracket({}, {})".format(
+ self._print(expr.left), self._print(expr.right)
+ )
+
+ def _derivative_parts(self, expr):
+ if isinstance(expr, SpatialDerivative):
+ base, suffix = self._derivative_parts(expr.expression)
+ suffix += DERIVATIVE_SUFFIXES[expr.coordinate]
+ # JOREK writes mixed derivatives with poloidal derivatives first
+ # and toroidal derivatives last (``u_xpp``, not ``u_ppx``).
+ order = {"x": 0, "y": 1, "s": 2, "t": 3, "p": 4}
+ suffix = "".join(sorted(suffix, key=lambda item: order[item]))
+ return base, suffix
+ return expr, ""
+
+ def _print_Function(self, expr):
+ external = definition_for_call(expr)
+ if external is not None:
+ return external.fortran_name
+ # Supplied derivative functions are represented as SymPy calls in the
+ # symbolic tree, but JOREK evaluates them into scalar variables.
+ for definition in self._external_definitions():
+ if expr.func.__name__ in definition.derivatives.values():
+ return expr.func.__name__
+ return super()._print_Function(expr)
+
+ def _external_definitions(self):
+ # The registry is intentionally queried lazily to keep this printer's
+ # public API independent of registry internals.
+ from .external import _NAME_REGISTRY
+
+ return tuple(_NAME_REGISTRY.values())
+
+
+def fortran(expression, **settings) -> str:
+ """Return a JOREK-style Fortran expression string."""
+
+ return FortranPrinter(**settings).doprint(expand_derivatives(expression))
diff --git a/util/equation_codegen/src/jorek_equations/fortran_source.py b/util/equation_codegen/src/jorek_equations/fortran_source.py
new file mode 100644
index 0000000000..7d687b430d
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/fortran_source.py
@@ -0,0 +1,744 @@
+"""Reading the JOREK element routine and comparing its terms with the DSL.
+
+This module owns step 1 and step 3 of the pipeline: it extracts the assembled
+``rhs_ij``/``amat`` expressions from a model's ``mod_elt_matrix_fft.f90``,
+rewrites the element routine's work variables into the names the symbolic DSL
+prints, parses the result into SymPy, and compares two such maps.
+"""
+
+from pathlib import Path
+import re
+from collections import Counter
+from typing import Dict, Iterable, Mapping
+
+import sympy as sp
+from sympy.parsing.sympy_parser import parse_expr
+
+from .channels import split_toroidal_channels
+from .fortran import fortran
+from .model199 import (
+ FIELDS,
+ T,
+ current_constraint_equation_3,
+ density_equation_5,
+ induction_equation_1,
+ j,
+ momentum_equation_2,
+ psi,
+ rho,
+ u,
+ vorticity_constraint_equation_4,
+ temperature_equation_6,
+ xjac,
+)
+from .operators import expand_brackets, lower_brackets_to_element
+from .symbols import coefficient
+
+
+EQUATION_1_ASSIGNMENTS = (
+ "rhs_ij_1",
+ "amat_11",
+ "amat_12",
+ "amat_12_n",
+ "amat_13",
+ "amat_16",
+)
+
+EQUATION_2_ASSIGNMENTS = (
+ "rhs_ij_2", "amat_21", "amat_22", "amat_23", "amat_23_n",
+ "amat_24", "amat_25", "amat_26",
+)
+EQUATION_3_ASSIGNMENTS = ("rhs_ij_3", "amat_31", "amat_33")
+EQUATION_4_ASSIGNMENTS = ("rhs_ij_4", "amat_42", "amat_44")
+EQUATION_5_ASSIGNMENTS = (
+ "rhs_ij_5", "rhs_ij_5_k", "amat_51", "amat_51_k", "amat_52",
+ "amat_55", "amat_55_k", "amat_55_n", "amat_55_kn",
+)
+EQUATION_6_ASSIGNMENTS = (
+ "rhs_ij_6", "rhs_ij_6_k", "amat_61", "amat_61_k", "amat_62",
+ "amat_63", "amat_66", "amat_66_k", "amat_66_n", "amat_66_kn",
+)
+
+
+class FortranComparisonError(AssertionError):
+ """Generated and source Fortran expressions differ."""
+
+
+def format_expression_terms(expression) -> str:
+ """Format an expression as one additive term per line.
+
+ The comparison itself remains symbolic; this is only a diagnostic
+ representation. Expanding the outer addition makes it much easier to
+ spot a missing sign or coefficient in the generated weak-form terms.
+ """
+
+ expanded = sp.expand(expression)
+ if expanded == 0:
+ return " 0"
+ terms = sp.Add.make_args(expanded)
+ lines = []
+ for term in terms:
+ if term.could_extract_minus_sign():
+ lines.append(" - {}".format(str(-term)))
+ else:
+ lines.append(" + {}".format(str(term)))
+ return "\n".join(lines)
+
+
+def format_assignment(name: str, expression, label: str = "") -> str:
+ """Return a readable, multiline assignment diagnostic."""
+
+ heading = "{}{}:".format(label + " " if label else "", name)
+ return heading + "\n" + format_expression_terms(expression)
+
+
+def format_term_difference(source, generated) -> str:
+ """Show unmatched additive terms without subtracting the expressions."""
+
+ source_terms = Counter(sp.Add.make_args(sp.expand(source)))
+ generated_terms = Counter(sp.Add.make_args(sp.expand(generated)))
+ source_only = list((source_terms - generated_terms).elements())
+ generated_only = list((generated_terms - source_terms).elements())
+ lines = []
+ if source_only:
+ lines.append(" source-only:")
+ lines.extend(" {}".format(_term_text(term)) for term in source_only)
+ if generated_only:
+ lines.append(" generated-only:")
+ lines.extend(" {}".format(_term_text(term)) for term in generated_only)
+ return "\n".join(lines) if lines else " (no unmatched terms)"
+
+
+def _term_text(term) -> str:
+ """Render one term while retaining its sign."""
+
+ if term.could_extract_minus_sign():
+ return "- {}".format(str(-term))
+ return "+ {}".format(str(term))
+
+
+def extract_fortran_assignments(path, names: Iterable[str]) -> Dict[str, str]:
+ """Extract continued scalar assignments from a Fortran source file."""
+
+ requested = tuple(names)
+ patterns = {
+ name: re.compile(r"^\s*{}\s*=\s*(.*)$".format(re.escape(name)), re.I)
+ for name in requested
+ }
+ lines = Path(path).read_text(encoding="utf-8").splitlines()
+ found = {}
+ index = 0
+ while index < len(lines):
+ uncommented = lines[index].split("!", 1)[0]
+ matched_name = None
+ matched_rhs = None
+ for name, pattern in patterns.items():
+ match = pattern.match(uncommented)
+ if match:
+ matched_name = name
+ matched_rhs = match.group(1)
+ break
+ if matched_name is None:
+ index += 1
+ continue
+
+ pieces = []
+ current = matched_rhs
+ while True:
+ current = current.strip()
+ continued = current.endswith("&")
+ pieces.append(current[:-1] if continued else current)
+ if not continued:
+ break
+ index += 1
+ if index >= len(lines):
+ raise ValueError("Unterminated assignment for {}".format(matched_name))
+ current = lines[index].split("!", 1)[0].lstrip()
+ if current.startswith("&"):
+ current = current[1:]
+ found[matched_name] = " ".join(piece.strip() for piece in pieces)
+ index += 1
+
+ missing = [name for name in requested if name not in found]
+ if missing:
+ raise ValueError(
+ "Assignments not found in {}: {}".format(path, ", ".join(missing))
+ )
+ return found
+
+
+_D_LITERAL = re.compile(r"(?i)(? str:
+ """Normalize the small Fortran-expression subset used by the comparator."""
+
+ result = expression.replace("&", " ")
+ result = re.sub(r"\bBigr\b", "BigR", result, flags=re.I)
+ result = re.sub(
+ r"current_source\s*\(\s*ms\s*,\s*mt\s*\)",
+ "current_source",
+ result,
+ flags=re.I,
+ )
+ for function_name in ("particle_source", "heat_source_i", "heat_source_e",
+ "heat_source"):
+ result = re.sub(
+ r"{}\s*\(\s*ms\s*,\s*mt\s*\)".format(function_name),
+ function_name,
+ result,
+ flags=re.I,
+ )
+ result = re.sub(
+ r"delta_g\s*\(\s*mp\s*,\s*1\s*,\s*ms\s*,\s*mt\s*\)",
+ "delta_g_1",
+ result,
+ flags=re.I,
+ )
+ result = re.sub(
+ r"delta_g\s*\(\s*mp\s*,\s*([56])\s*,\s*ms\s*,\s*mt\s*\)",
+ r"delta_g_\1",
+ result,
+ flags=re.I,
+ )
+ aliases = {
+ "D_par_local": "D_par",
+ "vv2": "(BigR**2*(u0_x**2+u0_y**2))",
+ "r0_hat": "(BigR**2*r0)",
+ "r0_x_hat": "(2*BigR*BigR_x*r0+BigR**2*r0_x)",
+ "r0_y_hat": "(BigR**2*r0_y)",
+ "rho_hat": "(BigR**2*rho)",
+ "rho_x_hat": "(2*BigR*BigR_x*rho+BigR**2*rho_x)",
+ "rho_y_hat": "(BigR**2*rho_y)",
+ "P0_s": "(r0_s*T0+r0*T0_s)",
+ "P0_t": "(r0_t*T0+r0*T0_t)",
+ "BB2": "((F0**2+ps0_x**2+ps0_y**2)/BigR**2)",
+ "BB2_psi": "(2*(psi_x*ps0_x+psi_y*ps0_y)/BigR**2)",
+ "Bgrad_rho_star": "((v_x*ps0_y-v_y*ps0_x)/BigR)",
+ "Bgrad_rho_k_star": "(F0*v_p/BigR**2)",
+ "Bgrad_rho": "((F0*r0_p/BigR+r0_x*ps0_y-r0_y*ps0_x)/BigR)",
+ "Bgrad_rho_star_psi": "((v_x*psi_y-v_y*psi_x)/BigR)",
+ "Bgrad_rho_psi": "((r0_x*psi_y-r0_y*psi_x)/BigR)",
+ "Bgrad_rho_rho": "((rho_x*ps0_y-rho_y*ps0_x)/BigR)",
+ "Bgrad_rho_rho_n": "(F0*rho_p/BigR**2)",
+ "Bgrad_T_star": "((v_x*ps0_y-v_y*ps0_x)/BigR)",
+ "Bgrad_T_k_star": "(F0*v_p/BigR**2)",
+ "Bgrad_T": "((F0*T0_p/BigR+T0_x*ps0_y-T0_y*ps0_x)/BigR)",
+ "Bgrad_T_star_psi": "((v_x*psi_y-v_y*psi_x)/BigR)",
+ "Bgrad_T_psi": "((T0_x*psi_y-T0_y*psi_x)/BigR)",
+ "Bgrad_T_T": "((T_x*ps0_y-T_y*ps0_x)/BigR)",
+ "Bgrad_T_T_n": "(F0*T_p/BigR**2)",
+ }
+ for name, replacement in aliases.items():
+ result = re.sub(r"\b{}\b".format(name), replacement, result, flags=re.I)
+ result = _D_LITERAL.sub(lambda match: match.group(1) + "e" + match.group(2), result)
+ return " ".join(result.split())
+
+
+def parse_fortran_expression(expression: str):
+ """Parse a normalized scalar Fortran expression into an exact SymPy tree."""
+
+ normalized = normalize_fortran_text(expression)
+ names = set(_IDENTIFIER.findall(normalized))
+ local_dict = {name: sp.Symbol(name) for name in names}
+ parsed = parse_expr(normalized, local_dict=local_dict, evaluate=True)
+ return sp.nsimplify(parsed)
+
+
+def generated_model199_equation1_assignments() -> Dict[str, sp.Expr]:
+ """Generate equation-1 terms in the same local convention as the source."""
+
+ equation = induction_equation_1()
+ result = equation.linearize(
+ fields=FIELDS,
+ timestep=coefficient("tstep"),
+ theta=coefficient("theta"),
+ zeta=coefficient("zeta"),
+ )
+
+ rhs_channels = split_toroidal_channels(result.rhs)
+ psi_channels = split_toroidal_channels(result.amat[psi])
+ u_channels = split_toroidal_channels(result.amat[u])
+ j_channels = split_toroidal_channels(result.amat[j])
+ T_channels = split_toroidal_channels(result.amat[T])
+
+ unsupported = {
+ "rhs": (rhs_channels.n, rhs_channels.k, rhs_channels.kn),
+ "psi": (psi_channels.n, psi_channels.k, psi_channels.kn),
+ "u": (u_channels.k, u_channels.kn),
+ "j": (j_channels.n, j_channels.k, j_channels.kn),
+ "T": (T_channels.n, T_channels.k, T_channels.kn),
+ }
+ for label, values in unsupported.items():
+ if any(value != 0 for value in values):
+ raise ValueError(
+ "Unexpected toroidal channel in equation-1 {} block".format(label)
+ )
+
+ generated = {
+ "rhs_ij_1": rhs_channels.p,
+ "amat_11": psi_channels.p,
+ "amat_12": u_channels.p,
+ "amat_12_n": u_channels.n,
+ "amat_13": j_channels.p,
+ "amat_16": T_channels.p,
+ }
+ return {
+ name: lower_brackets_to_element(expression, xjac)
+ for name, expression in generated.items()
+ }
+
+
+def generated_assignment_text() -> Dict[str, str]:
+ """Render generated terms with the identifiers used in model 199."""
+
+ return {
+ name: fortran(
+ expression,
+ previous_names={"psi": "delta_g(mp,1,ms,mt)"},
+ )
+ for name, expression in generated_model199_equation1_assignments().items()
+ }
+
+
+def generated_partial_model199_assignments() -> Dict[str, str]:
+ """Render generated assignments for the currently implemented equations."""
+
+ timestep = coefficient("tstep")
+ theta = coefficient("theta")
+ zeta = coefficient("zeta")
+ linearized = momentum_equation_2().linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ )
+ channels = split_toroidal_channels(linearized.rhs)
+ result_map = {"rhs_ij_2": channels.p}
+ for name, field, channel in (
+ ("amat_21", psi, "p"), ("amat_22", u, "p"),
+ ("amat_23", j, "p"), ("amat_23_n", j, "n"),
+ ("amat_24", FIELDS[3], "p"), ("amat_25", FIELDS[4], "p"),
+ ("amat_26", T, "p"),
+ ):
+ result_map[name] = getattr(split_toroidal_channels(linearized.amat[field]), channel)
+ for equation, rhs_name, block_fields in (
+ (current_constraint_equation_3(), "rhs_ij_3", (("amat_31", psi), ("amat_33", j))),
+ (vorticity_constraint_equation_4(), "rhs_ij_4", (("amat_42", u), ("amat_44", FIELDS[3]))),
+ ):
+ linearized = equation.linearize(fields=FIELDS)
+ result_map[rhs_name] = linearized.rhs
+ for assignment, field in block_fields:
+ result_map[assignment] = linearized.amat[field]
+ return {
+ name: fortran(value, previous_names={"u": "delta_u"})
+ for name, value in result_map.items()
+ }
+
+
+def generated_model199_transport_assignments() -> Dict[str, str]:
+ """Render generated RHS/AMAT assignments for equations 5 and 6."""
+
+ timestep = coefficient("tstep")
+ theta = coefficient("theta")
+ zeta = coefficient("zeta")
+ output = {}
+ for equation, fields, names in (
+ (density_equation_5(), FIELDS, EQUATION_5_ASSIGNMENTS),
+ (temperature_equation_6(), FIELDS, EQUATION_6_ASSIGNMENTS),
+ ):
+ linearized = equation.linearize(
+ fields=fields, timestep=timestep, theta=theta, zeta=zeta
+ )
+ rhs = split_toroidal_channels(linearized.rhs)
+ output[names[0]] = rhs.p
+ output[names[1]] = rhs.k
+ row_blocks = (("51", psi), ("52", u), ("55", rho)) if equation.name.endswith("density") else (("61", psi), ("62", u), ("63", j), ("66", T))
+ for block, field in row_blocks:
+ blocks = split_toroidal_channels(linearized.amat[field])
+ output["amat_{}".format(block)] = blocks.p
+ for channel in ("k", "n", "kn"):
+ if blocks.__getattribute__(channel) != 0:
+ output["amat_{}_{}".format(block, channel)] = getattr(blocks, channel)
+ previous_names = {
+ "rho": "delta_g(mp,5,ms,mt)",
+ "T": "delta_g(mp,6,ms,mt)",
+ }
+ return {
+ name: fortran(expand_brackets(value), previous_names=previous_names)
+ for name, value in output.items()
+ }
+
+
+def compare_assignment_maps(
+ source: Mapping[str, str],
+ generated: Mapping[str, str],
+ *,
+ set_BigR_x_one: bool = True,
+) -> None:
+ """Compare assignment maps and raise with symbolic differences.
+
+ ``set_BigR_x_one`` reflects the usual cylindrical-coordinate convention
+ and removes element-geometry ``BigR_x`` factors from this comparison.
+ """
+
+ messages = []
+ for name in source:
+ if name not in generated:
+ messages.append("{}: missing generated assignment".format(name))
+ continue
+ source_expr = parse_fortran_expression(source[name])
+ generated_expr = parse_fortran_expression(generated[name])
+ if set_BigR_x_one:
+ bigr_x = sp.Symbol("BigR_x")
+ source_expr = source_expr.subs(bigr_x, 1)
+ generated_expr = generated_expr.subs(bigr_x, 1)
+ difference = sp.simplify(generated_expr - source_expr)
+ if difference != 0:
+ messages.append(
+ "{name}:\n"
+ "{difference}".format(
+ name=name,
+ difference=format_term_difference(source_expr, generated_expr),
+ )
+ )
+ extra = sorted(set(generated) - set(source))
+ if extra:
+ messages.append("unexpected generated assignments: {}".format(", ".join(extra)))
+ if messages:
+ raise FortranComparisonError(
+ "Equation comparison failed:\n" + "\n".join(messages)
+ )
+
+
+def compare_model199_equation1(path, *, set_BigR_x_one: bool = True) -> None:
+ """Compare generated equation 1 with model199/mod_elt_matrix_fft.f90."""
+
+ source = extract_fortran_assignments(path, EQUATION_1_ASSIGNMENTS)
+ compare_assignment_maps(
+ source, generated_assignment_text(), set_BigR_x_one=set_BigR_x_one
+ )
+
+
+def normalize_model600_text(
+ expression, *, single_temperature=False, two_temperature=False
+):
+ # The element routine uses short background-density names, whereas the
+ # symbolic model uses the descriptive field names. They denote the same
+ # frozen quantities in the single-temperature momentum equation.
+ expression = re.sub(r"\brn0(?=\b|_)", "rhon0", expression, flags=re.I)
+ expression = re.sub(r"\brimp0(?=\b|_)", "rhoimp0", expression, flags=re.I)
+ # Convert the physical-coordinate cross product used by the source to its
+ # equivalent element-coordinate bracket before expanding aliases.
+ expression = re.sub(
+ r"ps0_x\s*\*\s*Pe0_y\s*-\s*ps0_y\s*\*\s*Pe0_x",
+ "(ps0_s*Pe0_t-ps0_t*Pe0_s)/xjac",
+ expression,
+ flags=re.I,
+ )
+ expression = re.sub(
+ r"factor\(\s*var_psi\s*,\s*([1-6])\s*\)",
+ "1",
+ expression,
+ flags=re.I,
+ )
+ expression = re.sub(
+ r"factor\(\s*var_(?:zj|w)\s*,\s*([1-6])\s*\)",
+ "1",
+ expression,
+ flags=re.I,
+ )
+ # ``delta_g(mp,var_X,ms,mt)`` is the previous Newton increment of field X.
+ # Give every field the same short spelling the DSL prints.
+ expression = re.sub(
+ r"delta_g\s*\(\s*mp\s*,\s*var_([a-z0-9_]+)\s*,\s*ms\s*,\s*mt\s*\)",
+ r"delta_\1",
+ expression,
+ flags=re.I,
+ )
+ expression = re.sub(
+ r"current_source\(\s*ms\s*,\s*mt\s*\)",
+ "current_source",
+ expression,
+ flags=re.I,
+ )
+ # NEO profile work arrays are scalar values at the element/toroidal
+ # point, despite being called as indexed Fortran functions.
+ for profile in ("amu_neo_prof", "aki_neo_prof"):
+ expression = re.sub(
+ r"\b{}\s*\(\s*ms\s*,\s*mt\s*\)".format(profile),
+ profile,
+ expression,
+ flags=re.I,
+ )
+ # The symbolic DSL prints the supplied density correction as a function
+ # call, while the element routine stores its value and derivative in the
+ # work variables ``r0_corr`` and ``dr0_corr_dn``.
+ expression = re.sub(
+ r"corr_neg_dens\(\s*r0\s*\)", "r0_corr", expression, flags=re.I
+ )
+ expression = re.sub(r"\bcorr_neg_dens\b", "r0_corr", expression, flags=re.I)
+ expression = re.sub(
+ r"dr0_corr_dn\(\s*r0\s*\)", "dr0_corr_dn", expression, flags=re.I
+ )
+ # The element routine's displayed tangent is written for the default
+ # correction branch where d(r0_corr)/d(rho)=1; compare that convention
+ # with the explicit DSL derivative call.
+ expression = re.sub(r"\bdr0_corr_dn\b", "1", expression)
+ if single_temperature:
+ # The one-temperature model evolves the total temperature. Do not
+ # identify either Ti0 or Te0 with T0 individually; only their sum
+ # (and the corresponding directional/spatial derivatives) is T0.
+ for suffix in ("", "_s", "_t", "_p", "_x", "_y", "_xx", "_yy", "_xy"):
+ expression = re.sub(
+ r"\bTi0{}\s*\+\s*Te0{}\b".format(suffix, suffix),
+ "T0{}".format(suffix), expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bTe0{}\s*\+\s*Ti0{}\b".format(suffix, suffix),
+ "T0{}".format(suffix), expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\balpha_imp_T\b", "(alpha_imp*T0)", expression, flags=re.I,
+ )
+ # ``alpha_e_T`` is the DSL print name of the closure value alpha_e*Te0.
+ # It is spelled the same way in both temperature models; the single
+ # temperature substitution Te0 = T0/2 is applied at the end.
+ expression = re.sub(
+ r"\balpha_e_T\b", "(alpha_e*Te0)", expression, flags=re.I,
+ )
+ pressure_temperature = "Te0"
+ ion_temperature = "Ti0"
+ # ``construct_pressure`` always includes impurity pressure when the
+ # impurity extension is active. Reports compare that full definition,
+ # rather than treating P0/Pi0/Pe0 as opaque element work variables.
+ # ``construct_pressure`` is written once for both temperature models: it
+ # always builds the species pressures from Ti0/Te0 and the per-species
+ # impurity coefficients. The one-temperature branch reaches it with
+ # Ti0=Te0=T0/2, which the substitutions at the end of this function
+ # apply. Expanding the aliases in the two-species basis therefore
+ # reproduces the element routine exactly in both branches.
+ ion_alpha = "alpha_i"
+ ion_density = "(r0+rhoimp0*alpha_i)"
+ electron_density = "(r0+rhoimp0*alpha_e)"
+ electron_density_tangent = "(r0+rhoimp0*alpha_e_bis)"
+
+ pi0 = "({}*{})".format(ion_density, ion_temperature)
+ pe0 = "({}*{})".format(electron_density, pressure_temperature)
+ pi_derivatives = {
+ "s": "((r0_s+rhoimp0_s*{a})*{t}+{d}*{t}_s)",
+ "t": "((r0_t+rhoimp0_t*{a})*{t}+{d}*{t}_t)",
+ "p": "((r0_p+rhoimp0_p*{a})*{t}+{d}*{t}_p)",
+ "x": "((r0_x+rhoimp0_x*{a})*{t}+{d}*{t}_x)",
+ "y": "((r0_y+rhoimp0_y*{a})*{t}+{d}*{t}_y)",
+ "xx": "((r0_xx+rhoimp0_xx*{a})*{t}+2*(r0_x+rhoimp0_x*{a})*{t}_x+{d}*{t}_xx)",
+ "yy": "((r0_yy+rhoimp0_yy*{a})*{t}+2*(r0_y+rhoimp0_y*{a})*{t}_y+{d}*{t}_yy)",
+ "xy": "((r0_xy+rhoimp0_xy*{a})*{t}+(r0_x+rhoimp0_x*{a})*{t}_y+(r0_y+rhoimp0_y*{a})*{t}_x+{d}*{t}_xy)",
+ }
+ pi_derivatives = {
+ key: value.format(a=ion_alpha, d=ion_density, t=ion_temperature)
+ for key, value in pi_derivatives.items()
+ }
+ pe_derivatives = {
+ "s": "((r0_s+rhoimp0_s*alpha_e)*{t}+{dt}*{t}_s)",
+ "t": "((r0_t+rhoimp0_t*alpha_e)*{t}+{dt}*{t}_t)",
+ "p": "((r0_p+rhoimp0_p*alpha_e)*{t}+{dt}*{t}_p)",
+ "x": "((r0_x+rhoimp0_x*alpha_e)*{t}+{dt}*{t}_x)",
+ "y": "((r0_y+rhoimp0_y*alpha_e)*{t}+{dt}*{t}_y)",
+ }
+ pe_derivatives = {
+ key: value.format(d=electron_density, dt=electron_density_tangent,
+ t=pressure_temperature)
+ for key, value in pe_derivatives.items()
+ }
+ p0 = "({}+{})".format(pi0, pe0)
+ p0_derivatives = {
+ key: "({}+{})".format(pi_derivatives[key], pe_derivatives[key])
+ for key in ("s", "t", "p")
+ }
+ # The implicit heating floor of the ion and electron energy equations is
+ # written inline with ``min`` and ``exp``. Name the two pieces so that
+ # the expression parser sees ordinary work values, and resolve the
+ # derivative of the exponential through its own value.
+ expression = re.sub(
+ r"\bmin\s*\(\s*Ti0\s*,\s*Tie_min_neg\s*\)", "Ti0_floor",
+ expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bexp\s*\(\s*\(\s*Ti0_floor\s*-\s*Tie_min_neg\s*\)\s*/\s*"
+ r"\(\s*0\.5(?:d0)?\s*\*\s*Tie_min_neg\s*\)\s*\)",
+ "Ti_floor_exp", expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bdTi_floor_exp\b", "(Ti_floor_exp/(0.5*Tie_min_neg))",
+ expression, flags=re.I,
+ )
+ expression = re.sub(r"\bdTi_floor\b", "1", expression, flags=re.I)
+ expression = re.sub(
+ r"\bmin\s*\(\s*Te0\s*,\s*Tie_min_neg\s*\)", "Te0_floor",
+ expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bexp\s*\(\s*\(\s*Te0_floor\s*-\s*Tie_min_neg\s*\)\s*/\s*"
+ r"\(\s*0\.5(?:d0)?\s*\*\s*Tie_min_neg\s*\)\s*\)",
+ "Te_floor_exp", expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bdTe_floor_exp\b", "(Te_floor_exp/(0.5*Tie_min_neg))",
+ expression, flags=re.I,
+ )
+ expression = re.sub(r"\bdTe_floor\b", "1", expression, flags=re.I)
+ expression = re.sub(
+ r"\bmin\s*\(\s*T0\s*,\s*T_min_neg\s*\)", "T0_floor",
+ expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bexp\s*\(\s*\(\s*T0_floor\s*-\s*T_min_neg\s*\)\s*/\s*"
+ r"\(\s*0\.5(?:d0)?\s*\*\s*T_min_neg\s*\)\s*\)",
+ "T_floor_exp", expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bdT_floor_exp\b", "(T_floor_exp/(0.5*T_min_neg))",
+ expression, flags=re.I,
+ )
+ expression = re.sub(r"\bdT_floor\b", "1", expression, flags=re.I)
+ # Corrected neutral density, mirroring ``corr_neg_dens``.
+ # ``rn0`` was renamed to ``rhon0`` at the top of this function, so the
+ # corrected neutral density must be spelled the same way here; otherwise
+ # the matcher compares ``rn0_corr`` with ``rhon0_corr`` while the report
+ # displays both as ``rhon0_corr``.
+ expression = re.sub(
+ r"corr_neg_dens_n\(\s*rhon0?\s*\)", "rhon0_corr", expression,
+ flags=re.I,
+ )
+ expression = re.sub(
+ r"\bcorr_neg_dens_n\b", "rhon0_corr", expression, flags=re.I,
+ )
+ expression = re.sub(r"\bdrn0_corr_dn\b", "1", expression, flags=re.I)
+ # Impurity negative-density correction, mirroring ``corr_neg_dens``.
+ expression = re.sub(
+ r"corr_neg_dens_imp\(\s*rhoimp0?\s*\)", "rhoimp0_corr",
+ expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\bcorr_neg_dens_imp\b", "rhoimp0_corr", expression, flags=re.I,
+ )
+ expression = re.sub(r"\bdrimp0_corr_dn\b", "1", expression, flags=re.I)
+ expression = re.sub(
+ r"\brimp0_corr\b", "rhoimp0_corr", expression, flags=re.I,
+ )
+ # ``sqrt`` is the only Fortran intrinsic appearing in the exported volume
+ # terms. Rewrite it as a power so the expression parser needs no function
+ # table.
+ expression = re.sub(
+ r"\bsqrt\s*\(([^()]*)\)", r"((\1)**(1/2))", expression, flags=re.I,
+ )
+ # Fortran is case insensitive, and the routine spells the adiabatic index
+ # both ``GAMMA`` and ``gamma`` — sometimes in a residual and its own
+ # tangent. Use one spelling.
+ expression = re.sub(r"\bgamma\b", "GAMMA", expression, flags=re.I)
+ # Fortran is case insensitive; the density and parallel-velocity blocks
+ # spell the parallel-velocity trial function ``Vpar`` while the DSL prints
+ # ``vpar``. ``Vpar0`` is covered by its own alias below.
+ expression = re.sub(
+ r"\bVpar(_[a-z]+)?\b", r"vpar\1", expression, flags=re.I,
+ )
+ aliases = {
+ "BB2": "((F0**2+ps0_x**2+ps0_y**2)/BigR**2)",
+ "psi_grad2": "(ps0_x**2+ps0_y**2)",
+ # Parallel-gradient work values of the density equation. The density
+ # ones are already covered by ``normalize_fortran_text``; the impurity
+ # ones must be expanded after ``rimp0`` has been renamed.
+ "Bgrad_rhoimp": "((F0*rhoimp0_p/BigR+rhoimp0_x*ps0_y-rhoimp0_y*ps0_x)/BigR)",
+ "Bgrad_rhoimp_psi": "((rhoimp0_x*psi_y-rhoimp0_y*psi_x)/BigR)",
+ "Bgrad_rhoimp_rhoimp": "((rhoimp_x*ps0_y-rhoimp_y*ps0_x)/BigR)",
+ "Bgrad_rhoimp_rhoimp_n": "(F0*rhoimp_p/BigR**2)",
+ # Parallel-velocity work values.
+ # Single-temperature parallel-gradient work values. ``Bgrad_T`` is
+ # already defined by ``normalize_fortran_text``; only the tangents
+ # are needed here.
+ "Bgrad_T_T": "((T_x*ps0_y-T_y*ps0_x)/BigR)",
+ "Bgrad_T_T_n": "(F0*T_p/BigR**2)",
+ # Ion-temperature parallel-gradient work values.
+ "Bgrad_Ti": "((F0*Ti0_p/BigR+Ti0_x*ps0_y-Ti0_y*ps0_x)/BigR)",
+ "Bgrad_Ti_psi": "((Ti0_x*psi_y-Ti0_y*psi_x)/BigR)",
+ "Bgrad_Ti_Ti": "((Ti_x*ps0_y-Ti_y*ps0_x)/BigR)",
+ "Bgrad_Ti_Ti_n": "(F0*Ti_p/BigR**2)",
+ # Electron-temperature parallel-gradient work values.
+ "Bgrad_Te": "((F0*Te0_p/BigR+Te0_x*ps0_y-Te0_y*ps0_x)/BigR)",
+ "Bgrad_Te_psi": "((Te0_x*psi_y-Te0_y*psi_x)/BigR)",
+ "Bgrad_Te_Te": "((Te_x*ps0_y-Te_y*ps0_x)/BigR)",
+ "Bgrad_Te_Te_n": "(F0*Te_p/BigR**2)",
+ "Bgrad_vpar": "((F0*vpar0_p/BigR+vpar0_x*ps0_y-vpar0_y*ps0_x)/BigR)",
+ "Bgrad_vpar_psi": "((vpar0_x*psi_y-vpar0_y*psi_x)/BigR)",
+ "Bgrad_vpar_vpar": "((vpar_x*ps0_y-vpar_y*ps0_x)/BigR)",
+ "Bgrad_vpar_vpar_n": "(F0*vpar_p/BigR**2)",
+ # Prescribed rotation profile: a flux function whose flux derivative
+ # is stored as ``dV_dpsi_source``.
+ "Vt0_x": "(dV_dpsi_source*ps0_x)",
+ "Vt0_y": "(dV_dpsi_source*ps0_y)",
+ "Vt_x_psi": "(dV_dpsi_source*psi_x)",
+ "Vt_y_psi": "(dV_dpsi_source*psi_y)",
+ # Fortran is case insensitive: amat_n(var_vpar,var_Ti) spells the
+ # background density ``R0``.
+ "R0": "r0",
+ "Btheta2": "((ps0_x**2+ps0_y**2)/BigR**2)",
+ "Btheta2_psi": "(2*(psi_x*ps0_x+psi_y*ps0_y)/BigR**2)",
+ "Vpar0": "vpar0",
+
+ "Pe0": pe0,
+ "Pe0_s": pe_derivatives["s"],
+ "Pe0_t": pe_derivatives["t"],
+ "Pe0_p": pe_derivatives["p"],
+ "Pe0_x": pe_derivatives["x"],
+ "Pe0_y": pe_derivatives["y"],
+ "Pi0": pi0,
+ "Pi0_s": pi_derivatives["s"],
+ "Pi0_t": pi_derivatives["t"],
+ "Pi0_y": pi_derivatives["y"],
+ # Cartesian derivatives used by the diamagnetic viscosity block.
+ # Keep these as explicit product rules so source and DSL expressions
+ # are expanded to the same monomials by the report matcher.
+ "Pi0_x": pi_derivatives["x"],
+ "Pi0_xx": pi_derivatives["xx"],
+ "Pi0_yy": pi_derivatives["yy"],
+ "Pi0_xy": pi_derivatives["xy"],
+ "P0": p0,
+ "P0_s": p0_derivatives["s"],
+ "P0_t": p0_derivatives["t"],
+ "P0_p": p0_derivatives["p"],
+ }
+ for name, replacement in aliases.items():
+ expression = re.sub(
+ r"\b{}\b".format(name), replacement, expression, flags=re.I
+ )
+ if single_temperature:
+ # The element routine sets Ti0 = Te0 = T0/2 in the one-temperature
+ # branch (the evolved T0 is the total temperature), so a species
+ # temperature reaching this point is half of the evolved one.
+ expression = re.sub(r"\bTe0([_a-z]*)\b", r"(T0\1/2)", expression)
+ expression = re.sub(r"\bTi0([_a-z]*)\b", r"(T0\1/2)", expression)
+ # The one-temperature impurity closure stores the means of the
+ # per-species coefficients: alpha_imp = (alpha_i+alpha_e)/2,
+ # alpha_imp_bis = (alpha_i+alpha_e_bis)/2 (alpha_i does not depend on
+ # temperature, so it is its own "bis"), and alpha_imp_tri =
+ # alpha_e_tri/4, the extra quarter coming from dTe0/dT0 = 1/2 applied
+ # twice. Express everything in the two-species basis so that
+ # assignments shared between the branches and assignments written for
+ # one branch only use the same symbols.
+ expression = re.sub(
+ r"\balpha_imp_tri\b", "(alpha_e_tri/4)", expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\balpha_imp_bis\b", "((alpha_i+alpha_e_bis)/2)",
+ expression, flags=re.I,
+ )
+ expression = re.sub(
+ r"\balpha_imp\b", "((alpha_i+alpha_e)/2)", expression, flags=re.I,
+ )
+ # W_dia_Ti is assembled from Pi0_*_Ti, which the element routine
+ # builds from the *trial* temperature basis function without the
+ # dTi0/dT0 = 1/2 chain factor. It is therefore twice the derivative
+ # of W_dia with respect to the evolved one-temperature field.
+ expression = re.sub(
+ r"\bW_dia_Ti\b", "(2*W_dia_T)", expression, flags=re.I,
+ )
+ return expression
diff --git a/util/equation_codegen/src/jorek_equations/latex_render.py b/util/equation_codegen/src/jorek_equations/latex_render.py
new file mode 100644
index 0000000000..ee71233286
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/latex_render.py
@@ -0,0 +1,690 @@
+"""Render Python weak-form expressions (as ``ast`` nodes) in LaTeX.
+
+The equations are written in Python (``model600.py``); this module turns their
+source into display math for the documentation. It does not import SymPy and
+never evaluates the equations: it reads them as syntax, so a helper call such
+as ``_B_dot_grad(rho)`` stays a named operator instead of being expanded.
+"""
+
+import ast
+import copy
+import json
+import re
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+
+class RenderError(ValueError):
+ """An expression that cannot be rendered."""
+
+# ---------------------------------------------------------------------------
+# LaTeX printing
+# ---------------------------------------------------------------------------
+
+ATOM, POW, MUL, ADD, LOW = 5, 4, 3, 2, 1
+
+GREEK = {
+ "alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta",
+ "kappa", "lambda", "mu", "nu", "xi", "pi", "rho", "sigma", "tau", "phi",
+ "chi", "psi", "omega", "Gamma", "Delta", "Theta", "Lambda", "Xi", "Pi",
+ "Sigma", "Phi", "Psi", "Omega",
+}
+
+
+def _name_part(part):
+ match = re.fullmatch(r"([A-Za-z]+)(\d*)", part)
+ if not match:
+ return r"\mathrm{%s}" % part
+ letters, digits = match.groups()
+ if letters in GREEK:
+ body = "\\" + letters
+ elif len(letters) == 1:
+ body = letters
+ else:
+ body = r"\mathrm{%s}" % letters
+ return body + ("_{%s}" % digits if digits else "")
+
+
+def default_symbol(name):
+ """LaTeX for a name not in the notation table: ``D_par_local -> D_{par,local}``."""
+
+ parts = [p for p in name.split("_") if p]
+ if not parts:
+ return r"\mathrm{%s}" % name.replace("_", r"\_")
+ head = _name_part(parts[0])
+ if len(parts) == 1:
+ return head
+ subscript = ",".join(_name_part(p) for p in parts[1:])
+ if "_{" in head:
+ return r"{%s}_{%s}" % (head, subscript)
+ return r"%s_{%s}" % (head, subscript)
+
+
+PLACEHOLDER = re.compile(r"\{(\w+)(!?)\}")
+
+DERIVATIVES = {"dR": "R", "dZ": "Z", "dphi": r"\phi", "ds": "s", "dt": "t"}
+
+BUILTIN_TEMPLATES = {
+ "grad": (r"\nabla_{\mathrm{pol}} {0}", POW),
+ "dot": (r"{0!}\cdot{1!}", MUL),
+ "laplacian": (r"\nabla^2_{\mathrm{pol}} {0}", POW),
+ "poiss_bracket": (r"[{0!},{1!}]", ATOM),
+ # poiss_bracket_st(a, b) = J [a,b]; divided by xjac it is shown as [a,b]^{st}
+ "poiss_bracket_st": (r"\mathcal{J}\,[{0!},{1!}]^{st}", MUL),
+ "bracket": (r"\mathcal{J}\,[{0!},{1!}]", MUL),
+ "__bracket_st": (r"[{0!},{1!}]^{st}", ATOM),
+ "__bracket": (r"[{0!},{1!}]", ATOM),
+ "freeze": (r"\overline{{0!}}", ATOM),
+ "__dt": (r"\partial_t {0}", ATOM),
+}
+
+
+@dataclass
+class Helper:
+ name: str
+ template: str
+ params: List[str]
+ keyword_only: List[str]
+ defaults: Dict[str, ast.expr]
+ display: Dict[str, str]
+
+
+def _fill(template, values):
+ """Substitute ``{key}``/``{key!}`` for the keys in ``values``; other braces are LaTeX."""
+
+ def replace(match):
+ key, raw = match.group(1), match.group(2)
+ if key not in values:
+ return match.group(0)
+ latex, prec = values[key]
+ if not raw and prec < ATOM:
+ return r"\left(%s\right)" % latex
+ return latex
+ return PLACEHOLDER.sub(replace, template)
+
+
+class LatexPrinter:
+ """Render Python expressions and statements as LaTeX.
+
+ ``values`` holds the local definitions of the function being rendered
+ (name -> ``ast`` expression), used to see through ``psi_gradient[0]``;
+ ``aliases`` maps a local ``functools.partial`` name to the helper it wraps.
+ """
+
+ def __init__(self, notation, helpers=None, local=None, terms=r"\Big(\textstyle\sum_k b_k\Big)",
+ values=None, aliases=None):
+ self.notation = notation
+ self.helpers = helpers or {}
+ self.local = local or {}
+ self.terms = terms
+ self.values = values or {}
+ self.aliases = aliases or {}
+
+ def with_local(self, extra):
+ return LatexPrinter(self.notation, self.helpers, dict(self.local, **extra), self.terms,
+ self.values, self.aliases)
+
+ # -- names ---------------------------------------------------------------
+ def symbol(self, name):
+ if name in self.local:
+ return self.local[name]
+ if name in self.notation and "{0}" not in self.notation[name]:
+ return self.notation[name]
+ if name in BUILTIN_TEMPLATES:
+ return _fill(BUILTIN_TEMPLATES[name][0], {"0": (r"\cdot", ATOM), "1": (r"\cdot", ATOM)})
+ return default_symbol(name)
+
+ # -- expressions -----------------------------------------------------------
+ def expr(self, node) -> Tuple[str, int]:
+ method = getattr(self, "_" + type(node).__name__, None)
+ if method is None:
+ raise RenderError("cannot render {} ({})".format(
+ type(node).__name__, ast.unparse(node)))
+ return method(node)
+
+ def latex(self, node):
+ return self.expr(node)[0]
+
+ def wrap(self, node, minimum):
+ latex, prec = self.expr(node)
+ if prec < minimum:
+ return r"\left(%s\right)" % latex
+ return latex
+
+ def _Constant(self, node):
+ if isinstance(node.value, str):
+ return r"\text{%s}" % node.value.replace("_", r"\_"), ATOM
+ return repr(node.value), ATOM
+
+ def _Name(self, node):
+ if node.id == "terms":
+ return self.terms, ATOM
+ return self.symbol(node.id), ATOM
+
+ def _Attribute(self, node):
+ return self.symbol(node.attr), ATOM
+
+ def _UnaryOp(self, node):
+ sign = {ast.USub: "-", ast.UAdd: "+"}.get(type(node.op))
+ if sign is None:
+ raise RenderError("cannot render " + ast.unparse(node))
+ return sign + self.wrap(node.operand, MUL), ADD
+
+ def _BinOp(self, node):
+ op = type(node.op)
+ if op is ast.Add:
+ return "%s + %s" % (self.wrap(node.left, ADD), self.wrap(node.right, MUL if _is_signed(node.right) else ADD)), ADD
+ if op is ast.Sub:
+ return "%s - %s" % (self.wrap(node.left, ADD), self.wrap(node.right, MUL)), ADD
+ if op in (ast.Mult, ast.Div):
+ sign, unsigned = _pull_sign(node)
+ if sign == "-":
+ return "-" + self.wrap(unsigned, MUL), ADD
+ node = unsigned
+ if op in (ast.Mult, ast.Div):
+ numerator, denominator = _cancel_jacobian(*_num_den(node))
+ if not denominator:
+ return self.product(numerator), MUL
+ bottom = self.latex(denominator[0]) if len(denominator) == 1 else self.product(denominator)
+ number = 1
+ for factor in numerator:
+ if _is_number(factor):
+ number *= factor.value
+ rest = [f for f in numerator if not _is_number(f)]
+ if len(rest) == 1 and self.expr(rest[0])[1] <= ADD:
+ return r"\frac{%s}{%s}\left(%s\right)" % (repr(number), bottom, self.latex(rest[0])), MUL
+ if len(rest) <= 1:
+ top = self.latex(rest[0]) if rest and number == 1 else self.product(numerator)
+ return r"\frac{%s}{%s}" % (top, bottom), MUL
+ # keep the factors at full size: (number / denominator) * factors
+ return r"\frac{%s}{%s}\,%s" % (repr(number), bottom, self.product(rest)), MUL
+ if op is ast.Pow:
+ return "%s^{%s}" % (self.wrap(node.left, ATOM), self.latex(node.right)), POW
+ raise RenderError("cannot render " + ast.unparse(node))
+
+ def product(self, factors):
+ """Juxtapose factors, numbers first; bracket dot products that are followed by more."""
+
+ number = 1
+ for factor in factors:
+ if _is_number(factor):
+ number *= factor.value
+ rest = [f for f in factors if not _is_number(f)]
+ rest = [f for f in rest if _is_rational(f)] + [f for f in rest if not _is_rational(f)]
+ parts = [repr(number)] if number != 1 or not rest else []
+ for position, factor in enumerate(rest):
+ latex = self.wrap(factor, MUL)
+ if (re.search(r"\\cdot(?![a-z])", latex) and len(rest) > 1
+ and not latex.startswith(r"\left(")):
+ latex = r"\left(%s\right)" % latex
+ parts.append(latex)
+ return r"\,".join(parts)
+
+ def _IfExp(self, node):
+ return (
+ r"\begin{cases} %s & \text{if } %s \\ %s & \text{otherwise} \end{cases}"
+ % (self.latex(node.body), _condition(node.test), self.latex(node.orelse))
+ ), ATOM
+
+ def _Lambda(self, node):
+ params = [a.arg for a in node.args.args]
+ inner = self.with_local({p: p for p in params})
+ return r"(%s) \mapsto %s" % (", ".join(params), inner.latex(node.body)), LOW
+
+ def _Tuple(self, node):
+ return r"\left(%s\right)" % ", ".join(self.latex(e) for e in node.elts), ATOM
+
+ def _Subscript(self, node):
+ value, index = node.value, node.slice
+ if isinstance(value, ast.Name) and isinstance(self.values.get(value.id), ast.Call):
+ value = self.values[value.id]
+ if (isinstance(value, ast.Call) and _call_name(value.func) == "grad" and len(value.args) == 1
+ and isinstance(index, ast.Constant) and index.value in (0, 1)):
+ return r"\partial_{%s} %s" % ("RZ"[index.value], self.wrap(value.args[0], ATOM)), POW
+ return "{%s}_{%s}" % (self.wrap(node.value, ATOM), self.latex(node.slice)), ATOM
+
+ def _GeneratorExp(self, node):
+ """``(f(a, b) for a, b in zip(A, B))``: the vector ``f(A, B)``."""
+
+ (loop,) = node.generators
+ if (loop.ifs or not isinstance(loop.iter, ast.Call) or _call_name(loop.iter.func) != "zip"
+ or not isinstance(loop.target, ast.Tuple)):
+ raise RenderError("cannot render " + ast.unparse(node))
+ names = [target.id for target in loop.target.elts]
+ inner = self.with_local({n: self.latex(a) for n, a in zip(names, loop.iter.args)})
+ return inner.expr(node.elt)
+
+ def _Call(self, node):
+ name = _call_name(node.func)
+ name = self.aliases.get(name, name)
+ args = node.args
+ if name == "tuple" and len(args) == 1 and isinstance(args[0], ast.GeneratorExp):
+ return self.expr(args[0])
+ if name in DERIVATIVES and len(args) == 1:
+ chain = [DERIVATIVES[name]]
+ inner = args[0]
+ while (isinstance(inner, ast.Call) and _call_name(inner.func) in DERIVATIVES
+ and len(inner.args) == 1):
+ chain.append(DERIVATIVES[_call_name(inner.func)])
+ inner = inner.args[0]
+ ops = "".join(r"\partial_{%s}" % c for c in chain)
+ return "%s %s" % (ops, self.wrap(inner, ATOM)), POW
+ if name == "Rational" and len(args) == 2:
+ return r"\frac{%s}{%s}" % (self.latex(args[0]), self.latex(args[1])), ATOM
+ if name == "sqrt" and len(args) == 1:
+ return r"\sqrt{%s}" % self.latex(args[0]), ATOM
+ if name in ("coefficient", "test_function") and args and isinstance(args[0], ast.Constant):
+ return self.symbol(args[0].value), ATOM
+ if name in BUILTIN_TEMPLATES:
+ template, prec = BUILTIN_TEMPLATES[name]
+ return _fill(template, {str(i): self.expr(a) for i, a in enumerate(args)}), prec
+ if name in self.helpers:
+ helper = self.helpers[name]
+ names = set(helper.params + helper.keyword_only)
+ shown = [k for k, _ in PLACEHOLDER.findall(helper.template) if k in names]
+ atomic = not shown or helper.template.rstrip().endswith(")")
+ return self._helper_call(helper, node), ATOM if atomic else POW
+ values = [self.latex(a) for a in args]
+ if name in self.notation and "{0}" in self.notation[name]:
+ return _fill(self.notation[name], {str(i): self.expr(a) for i, a in enumerate(args)}), POW
+ head = self.symbol(name)
+ if not values:
+ return head, ATOM
+ return r"%s(%s)" % (head, ", ".join(values)), ATOM
+
+ def _helper_call(self, helper, node):
+ values = {}
+ for param, arg in zip(helper.params, node.args):
+ values[param] = arg
+ for keyword in node.keywords:
+ values[keyword.arg] = keyword.value
+ names = set(helper.params + helper.keyword_only)
+ used = {key for key, _ in PLACEHOLDER.findall(helper.template) if key in names}
+ for param in values:
+ if param not in used and param not in helper.keyword_only:
+ raise RenderError(
+ "{}(...) is called with {!r}, which its template {!r} does not show; "
+ "add {{{}}} to the template".format(helper.name, param, helper.template, param))
+ filled = {}
+ for key in used:
+ if key in values:
+ filled[key] = self.expr(values[key])
+ elif key in helper.defaults:
+ filled[key] = self.expr(helper.defaults[key])
+ else:
+ raise RenderError("{}(...) is called without {!r}".format(helper.name, key))
+ return _fill(helper.template, filled)
+
+ # -- statements ------------------------------------------------------------
+ def rows(self, statements, indent=""):
+ rows = []
+ for statement in statements:
+ if isinstance(statement, ast.Expr) and isinstance(statement.value, ast.Constant):
+ continue
+ if isinstance(statement, ast.Raise) or (
+ isinstance(statement, ast.If) and not statement.orelse
+ and all(isinstance(s, ast.Raise) for s in statement.body)):
+ continue # guards are not part of the equation
+ if (isinstance(statement, ast.Assign) and isinstance(statement.value, ast.Call)
+ and _call_name(statement.value.func) == "partial"):
+ continue # an alias of a helper, resolved where it is called
+ if isinstance(statement, ast.Assign):
+ targets = " = ".join(self.latex(t) for t in statement.targets)
+ value = self.latex(statement.value)
+ if value != targets:
+ rows.append(r"%s%s &= %s" % (indent, targets, value))
+ elif isinstance(statement, ast.AugAssign) and isinstance(statement.op, (ast.Add, ast.Sub)):
+ op = "+" if isinstance(statement.op, ast.Add) else "-"
+ rows.append(r"%s%s &\mathrel{%s}= %s" % (
+ indent, self.latex(statement.target), op, self.latex(statement.value)))
+ elif isinstance(statement, ast.If):
+ rows.append(r"&%s\text{if } %s\text{:}" % (indent, _condition(statement.test)))
+ rows.extend(self.rows(statement.body, indent + r"\quad "))
+ if statement.orelse:
+ rows.append(r"&%s\text{otherwise:}" % indent)
+ rows.extend(self.rows(statement.orelse, indent + r"\quad "))
+ elif isinstance(statement, ast.Return):
+ rows.append(r"%s &= %s" % (indent, self.latex(statement.value)))
+ else:
+ raise RenderError("cannot render statement " + ast.unparse(statement))
+ return rows
+
+
+def _call_name(func):
+ if isinstance(func, ast.Name):
+ return func.id
+ if isinstance(func, ast.Attribute):
+ return func.attr
+ return ""
+
+
+def _is_signed(node):
+ return isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd))
+
+
+def _is_number(node):
+ return isinstance(node, ast.Constant) and isinstance(node.value, (int, float)) \
+ and not isinstance(node.value, bool)
+
+
+def _is_rational(node):
+ return isinstance(node, ast.Call) and _call_name(node.func) == "Rational"
+
+
+def _pull_sign(node):
+ """Move a unary sign on the leftmost factor of a product to the front."""
+
+ if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
+ return ("-" if isinstance(node.op, ast.USub) else "+"), node.operand
+ if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Mult, ast.Div)):
+ sign, left = _pull_sign(node.left)
+ if left is not node.left:
+ return sign, ast.BinOp(left=left, op=node.op, right=node.right)
+ return "+", node
+
+
+def _num_den(node):
+ """Numerator and denominator factors of a product/quotient chain."""
+
+ if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Mult):
+ left, right = _num_den(node.left), _num_den(node.right)
+ return left[0] + right[0], left[1] + right[1]
+ if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Div):
+ left, right = _num_den(node.left), _num_den(node.right)
+ return left[0] + right[1], left[1] + right[0]
+ return [node], []
+
+
+BRACKETS_ST = {"poiss_bracket_st": "__bracket_st", "bracket": "__bracket"}
+
+
+def _cancel_jacobian(numerator, denominator):
+ """Pair each element-coordinate bracket with a ``xjac`` of the denominator."""
+
+ numerator, denominator = list(numerator), list(denominator)
+ if any(_call_name(getattr(f, "func", None)) in BRACKETS_ST for f in numerator if isinstance(f, ast.Call)):
+ # the volume weight dV = R*xjac in a denominator also supplies a xjac
+ expanded = []
+ for factor in denominator:
+ if isinstance(factor, ast.Name) and factor.id == "dV":
+ expanded += [ast.Name(id="R", ctx=ast.Load()), ast.Name(id="xjac", ctx=ast.Load())]
+ else:
+ expanded.append(factor)
+ denominator = expanded
+ for index, factor in enumerate(numerator):
+ name = _call_name(factor.func) if isinstance(factor, ast.Call) else ""
+ if name not in BRACKETS_ST:
+ continue
+ jacobian = next((i for i, d in enumerate(denominator)
+ if isinstance(d, ast.Name) and d.id == "xjac"), None)
+ if jacobian is None:
+ continue
+ del denominator[jacobian]
+ numerator[index] = ast.Call(func=ast.Name(id=BRACKETS_ST[name], ctx=ast.Load()),
+ args=factor.args, keywords=[])
+ return numerator, denominator
+
+
+def _rebuild(numerator, denominator):
+ def chain(factors):
+ node = factors[0]
+ for factor in factors[1:]:
+ node = ast.BinOp(left=node, op=ast.Mult(), right=factor)
+ return node
+ top = chain(numerator) if numerator else ast.Constant(value=1)
+ return ast.BinOp(left=top, op=ast.Div(), right=chain(denominator)) if denominator else top
+
+
+def _without_volume_element(node, name="dV"):
+ """``node / dV`` for an expression with a ``dV`` factor in its numerator."""
+
+ numerator, denominator = _num_den(node)
+ for index, factor in enumerate(numerator):
+ if isinstance(factor, ast.Name) and factor.id == name:
+ return _rebuild(numerator[:index] + numerator[index + 1:], denominator)
+ raise RenderError("expected a factor {} in {}".format(name, ast.unparse(node)))
+
+
+def _factors(node):
+ if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Mult):
+ return _factors(node.left) + _factors(node.right)
+ return [node]
+
+
+def _condition(test):
+ # GitHub's Markdown-to-KaTeX pipeline unescapes "\_" back to "_" before the
+ # math is parsed, which then errors inside \texttt{...} (text mode, where a
+ # bare "_" is invalid). Render the flag name with spaces instead of a
+ # backslash escape, so no "\_" sequence reaches the page at all.
+ return r"\texttt{%s}" % ast.unparse(test).replace("_", " ")
+
+
+def _signed_terms(node):
+ """Split a sum into ``(sign, term)`` pairs, keeping the written order."""
+
+ if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub)):
+ left = _signed_terms(node.left)
+ right = _signed_terms(node.right)
+ if isinstance(node.op, ast.Sub):
+ right = [("-" if s == "+" else "+", t) for s, t in right]
+ return left + right
+ if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
+ sign = "-" if isinstance(node.op, ast.USub) else "+"
+ return [(sign if s == "+" else ("+" if sign == "-" else "-"), t)
+ for s, t in _signed_terms(node.operand)] if _is_sum(node.operand) else [(sign, node.operand)]
+ sign, unsigned = _pull_sign(node)
+ if unsigned is not node:
+ return [(sign, unsigned)]
+ return [("+", node)]
+
+
+def _is_sum(node):
+ return isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub))
+
+
+def _aligned(rows):
+ if len(rows) == 1 and "&" not in rows[0].replace(r"\&", ""):
+ return rows[0]
+ return "\\begin{aligned}\n" + " \\\\\n".join(rows) + "\n\\end{aligned}"
+
+
+# ---------------------------------------------------------------------------
+# Helpers and blocks
+# ---------------------------------------------------------------------------
+
+
+LONG_SUM = 4
+LONG_LINE = 160 # characters of LaTeX that still read well on one line
+
+
+def _definition_rows(printer, lhs, value):
+ """``lhs = value``, one term per row when the value is a long sum."""
+
+ single = r"%s &= %s" % (lhs, printer.latex(value))
+ if len(single) <= LONG_LINE:
+ return [single]
+ terms = _signed_terms(value)
+ if len(terms) >= LONG_SUM:
+ rows = [r"%s &= %s" % (lhs, _signed(printer, terms[0], first=True))]
+ rows += [r"&\quad %s" % _signed(printer, term) for term in terms[1:]]
+ return rows
+ sign, unsigned = _pull_sign(value)
+ numerator, denominator = _num_den(unsigned)
+ if not denominator and len(numerator) > 1:
+ inner = _signed_terms(numerator[-1])
+ if len(inner) >= 3:
+ factor = ("-" if sign == "-" else "") + printer.product(numerator[:-1])
+ rows = [r"%s &= %s\Big(%s" % (lhs, factor, _signed(printer, inner[0], first=True))]
+ rows += [r"&\qquad %s" % _signed(printer, term) for term in inner[1:]]
+ rows[-1] += r"\Big)"
+ return rows
+ return [r"%s &= %s" % (lhs, printer.latex(value))]
+
+
+def _signed(printer, signed_term, first=False):
+ sign, term = signed_term
+ latex = printer.wrap(term, MUL)
+ if first:
+ return ("-" if sign == "-" else "") + latex
+ return "%s %s" % (sign, latex)
+
+
+def _is_variant(statement):
+ """An ``if flag: return ...`` spelling variant of a helper's result."""
+
+ return (isinstance(statement, ast.If) and not statement.orelse
+ and all(isinstance(s, ast.Return) for s in statement.body))
+
+
+SENTINEL = "\x00terms\x00"
+FIELD_NAMES = {"psi", "u", "j", "omega", "rho", "T", "Ti", "Te", "vpar", "rhon", "rhoimp"}
+
+
+class _TimeDerivative:
+ """Write out the time derivative of a mass term.
+
+ ``freeze(x)`` marks a factor that the time derivative does not act on, so
+ a mass term with frozen factors is expanded by the product rule over its
+ remaining factors. A term without frozen factors is kept as
+ ``partial_t(term)``.
+ """
+
+ def __init__(self, values, helpers):
+ self.values = values
+ self.helpers = helpers
+
+ def depends(self, node, seen=()):
+ if isinstance(node, ast.Call) and _call_name(node.func) == "freeze":
+ return False
+ if isinstance(node, ast.Name):
+ if node.id in FIELD_NAMES:
+ return True
+ if node.id in self.values and node.id not in seen:
+ return self.depends(self.values[node.id], seen + (node.id,))
+ return False
+ if isinstance(node, ast.Call) and _call_name(node.func) in self.helpers:
+ helper = self.helpers[_call_name(node.func)]
+ given = set(helper.params[:len(node.args)]) | {k.arg for k in node.keywords}
+ if any(self.depends(d, seen) for p, d in helper.defaults.items()
+ if p not in given and p not in helper.keyword_only):
+ return True
+ return any(self.depends(child, seen) for child in ast.iter_child_nodes(node))
+
+ @staticmethod
+ def has_freeze(node):
+ return any(isinstance(n, ast.Call) and _call_name(n.func) == "freeze" for n in ast.walk(node))
+
+ @staticmethod
+ def thaw(node):
+ class Thaw(ast.NodeTransformer):
+ def visit_Call(self, call):
+ self.generic_visit(call)
+ if _call_name(call.func) == "freeze" and len(call.args) == 1:
+ return call.args[0]
+ return call
+ import copy
+ return Thaw().visit(copy.deepcopy(node))
+
+ LINEAR = {"grad", "dR", "dZ", "dphi", "laplacian"}
+
+ def mark(self, node):
+ """``partial_t node``, moved inside linear operators of a single time-dependent argument."""
+
+ if isinstance(node, ast.Name) and node.id in self.values and node.id not in FIELD_NAMES:
+ value = self.values[node.id]
+ if isinstance(value, ast.Call) and _call_name(value.func) in self.LINEAR | {"dot"}:
+ return self.mark(value)
+ if isinstance(node, ast.Call) and not node.keywords:
+ name = _call_name(node.func)
+ timed = [i for i, a in enumerate(node.args) if self.depends(a)]
+ if (name in self.LINEAR or name == "dot") and len(timed) == 1:
+ args = list(node.args)
+ args[timed[0]] = self.mark(args[timed[0]])
+ return ast.Call(func=node.func, args=args, keywords=[])
+ return ast.Call(func=ast.Name(id="__dt", ctx=ast.Load()), args=[node], keywords=[])
+
+ def derive(self, node):
+ if not self.depends(node):
+ return None
+ if not self.has_freeze(node):
+ return self.mark(node)
+ if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub)):
+ left, right = self.derive(node.left), self.derive(node.right)
+ if right is None:
+ return left
+ if left is None:
+ return right if isinstance(node.op, ast.Add) else ast.UnaryOp(op=ast.USub(), operand=right)
+ return ast.BinOp(left=left, op=node.op, right=right)
+ if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
+ inner = self.derive(node.operand)
+ return None if inner is None else ast.UnaryOp(op=node.op, operand=inner)
+ if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Mult, ast.Div)):
+ numerator, denominator = _num_den(node)
+ if any(self.depends(d) for d in denominator):
+ return self.mark(self.thaw(node))
+ result = None
+ for index, factor in enumerate(numerator):
+ derived = self.derive(factor)
+ if derived is None:
+ continue
+ others = [self.thaw(f) for f in numerator[:index] + numerator[index + 1:]]
+ term = _rebuild(others + [derived], denominator)
+ result = term if result is None else ast.BinOp(left=result, op=ast.Add(), right=term)
+ return result
+ return self.mark(self.thaw(node))
+
+
+
+POLICY_TEXT = {
+ "active": "differentiated",
+ "piecewise_active": "differentiated (current branch)",
+ "frozen": "not differentiated",
+}
+
+
+def functions_table(symbols_source):
+ """A Markdown table of the ``external_function`` declarations of a module."""
+
+ rows = ["| DSL | Arguments | Fortran value | Derivatives supplied | Jacobian |",
+ "|---|---|---|---|---|"]
+ for statement in ast.parse(symbols_source).body:
+ if not (isinstance(statement, ast.Assign) and isinstance(statement.value, ast.Call)
+ and _call_name(statement.value.func) == "external_function"):
+ continue
+ call = statement.value
+ keywords = {k.arg: ast.literal_eval(k.value) for k in call.keywords}
+ name = statement.targets[0].id
+ fortran = keywords.get("fortran_name") or ast.literal_eval(call.args[0])
+ arguments = ", ".join("`%s`" % a for a in keywords.get("arguments", ())) or "none"
+ derivatives = ", ".join("`%s` for `%s`" % (d, a) for a, d in keywords.get("derivatives", {}).items())
+ policy = POLICY_TEXT[keywords.get("policy", "active")]
+ rows.append("| `%s` | %s | `%s` | %s | %s |" % (name, arguments, fortran, derivatives or "none", policy))
+ return rows
+
+
+FINDINGS_URL = ("https://github.com/iterorganization/JOREK/blob/develop/"
+ "util/equation_codegen/JOREK_FINDINGS.md")
+
+
+def findings_table(reference_source):
+ """A Markdown table of the findings recorded in the checker's reference."""
+
+ import json
+
+ blocks = json.loads(reference_source)["blocks"]
+ by_finding = {}
+ for name, block in blocks.items():
+ lhs = name.split(" | ")[1]
+ for finding in block["findings"]:
+ kinds = by_finding.setdefault(finding, ({}, {}))
+ target = kinds[0] if lhs.startswith("rhs") else kinds[1]
+ target[lhs] = target.get(lhs, 0) + len(block["source_only"]) + len(block["generated_only"])
+ rows = ["| Finding | Residual blocks | Jacobian blocks |", "|---:|---|---|"]
+ for finding in sorted(by_finding):
+ residual, jacobian = by_finding[finding]
+ rows.append("| [%d](%s) | %s | %s |" % (
+ finding, FINDINGS_URL,
+ ", ".join("`%s`" % b for b in sorted(residual)) or "none",
+ ", ".join("`%s`" % b for b in sorted(jacobian)) or "none"))
+ return rows
+
diff --git a/util/equation_codegen/src/jorek_equations/linearization.py b/util/equation_codegen/src/jorek_equations/linearization.py
new file mode 100644
index 0000000000..f0a24b5901
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/linearization.py
@@ -0,0 +1,154 @@
+"""Directional linearization of JOREK DSL expressions."""
+
+from functools import lru_cache
+
+import sympy as sp
+from sympy.core.function import AppliedUndef, ArgumentIndexError
+
+from .exceptions import MissingDerivativeError, UnsupportedExpressionError
+from .external import LinearizationPolicy, definition_for_call
+from .operators import PoloidalBracket, SpatialDerivative, bracket, derivative
+from .symbols import Field, FieldRole, FieldValue, Frozen, TestFunction
+
+
+def _is_zero(expression) -> bool:
+ return expression == 0
+
+
+def resolve_current(expression):
+ """Resolve abstract equation fields to current-state values."""
+
+ expression = sp.sympify(expression)
+ replacements = {
+ value: FieldValue(value.field_name, FieldRole.CURRENT)
+ for value in expression.atoms(FieldValue)
+ if value.role is FieldRole.ABSTRACT
+ }
+ return expression.xreplace(replacements)
+
+
+def variation(expression, wrt: Field, *, direction=FieldRole.DELTA):
+ """Compute the directional variation of an expression for one field."""
+
+ if not isinstance(wrt, Field):
+ raise TypeError("wrt must be a Field declaration")
+ direction = FieldRole(direction)
+ if direction is FieldRole.CURRENT:
+ raise ValueError("A variation direction cannot have role 'current'")
+ # Equation definitions contain abstract fields. At the tangent stage,
+ # non-varied factors are evaluated at the current state.
+ expression = resolve_current(expression)
+
+ @lru_cache(maxsize=None)
+ def visit(node):
+ if isinstance(node, Frozen):
+ return sp.S.Zero
+
+ if isinstance(node, FieldValue):
+ if node.role is FieldRole.CURRENT and node.field_name == wrt.name:
+ return wrt.value(direction)
+ return sp.S.Zero
+
+ if isinstance(node, (TestFunction, sp.Symbol, sp.Number)):
+ return sp.S.Zero
+
+ if isinstance(node, SpatialDerivative):
+ inner = visit(node.expression)
+ if _is_zero(inner):
+ return sp.S.Zero
+ return derivative(inner, node.coordinate)
+
+ if isinstance(node, PoloidalBracket):
+ left_variation = visit(node.left)
+ right_variation = visit(node.right)
+ result = sp.S.Zero
+ if not _is_zero(left_variation):
+ result += bracket(left_variation, node.right)
+ if not _is_zero(right_variation):
+ result += bracket(node.left, right_variation)
+ return result
+
+ external = definition_for_call(node)
+ if external is not None:
+ if external.policy is LinearizationPolicy.FROZEN:
+ return sp.S.Zero
+ result = sp.S.Zero
+ for argument_name, argument_value in zip(external.arguments, node.args):
+ argument_variation = visit(argument_value)
+ if _is_zero(argument_variation):
+ continue
+ if argument_name not in external.derivatives:
+ raise MissingDerivativeError(
+ "External function {} is active in argument {!r}, but no "
+ "derivative was supplied".format(external.name, argument_name)
+ )
+ result += (
+ external.derivative_call(argument_name, node.args)
+ * argument_variation
+ )
+ return result
+
+ if isinstance(node, AppliedUndef):
+ raise UnsupportedExpressionError(
+ "Undefined function {} has no external-function declaration".format(
+ node.func
+ )
+ )
+
+ if isinstance(node, sp.Add):
+ return sp.Add(*(visit(argument) for argument in node.args))
+
+ if isinstance(node, sp.Mul):
+ terms = []
+ for index, argument in enumerate(node.args):
+ argument_variation = visit(argument)
+ if _is_zero(argument_variation):
+ continue
+ factors = list(node.args)
+ factors[index] = argument_variation
+ terms.append(sp.Mul(*factors))
+ return sp.Add(*terms)
+
+ if isinstance(node, sp.Pow):
+ base, exponent = node.args
+ base_variation = visit(base)
+ exponent_variation = visit(exponent)
+ result = sp.S.Zero
+ if not _is_zero(base_variation):
+ result += exponent * base ** (exponent - 1) * base_variation
+ if not _is_zero(exponent_variation):
+ result += node * exponent_variation * sp.log(base)
+ return result
+
+ if isinstance(node, sp.Function):
+ result = sp.S.Zero
+ for index, argument in enumerate(node.args, start=1):
+ argument_variation = visit(argument)
+ if _is_zero(argument_variation):
+ continue
+ try:
+ partial = node.fdiff(index)
+ except (ArgumentIndexError, NotImplementedError) as exc:
+ raise UnsupportedExpressionError(
+ "Cannot differentiate function {} safely".format(node.func)
+ ) from exc
+ result += partial * argument_variation
+ return result
+
+ if not node.args:
+ return sp.S.Zero
+
+ raise UnsupportedExpressionError(
+ "Unsupported expression node {} in {}".format(type(node).__name__, node)
+ )
+
+ return sp.expand(visit(expression))
+
+
+def jacobian_blocks(expression, fields):
+ """Return one trial-function directional derivative per field."""
+
+ return {
+ value: variation(expression, value, direction=FieldRole.TRIAL)
+ for value in fields
+ }
diff --git a/util/equation_codegen/src/jorek_equations/model199.py b/util/equation_codegen/src/jorek_equations/model199.py
new file mode 100644
index 0000000000..1cfa958556
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/model199.py
@@ -0,0 +1,173 @@
+"""Canonical fields and weak equations for JOREK model 199."""
+
+import sympy as sp
+
+from .equations import ConstraintEquation, EvolutionEquation
+from .external import external_function
+from .operators import R, bracket, dR, dZ, dphi, dot, laplacian, poiss_bracket_st, grad
+from .symbols import coefficient, field, freeze, test_function
+
+
+psi = field("psi", fortran_current="ps0", fortran_trial="psi")
+u = field("u", fortran_current="u0", fortran_trial="u")
+j = field("j", fortran_current="zj0", fortran_trial="zj")
+omega = field("omega", fortran_current="w0", fortran_trial="w")
+rho = field("rho", fortran_current="r0", fortran_trial="rho")
+T = field("T", fortran_current="T0", fortran_trial="T")
+
+FIELDS = (psi, u, j, omega, rho, T)
+
+
+# The constitutive interface is deliberately explicit: the value is evaluated
+# by the host application and the tangent is named for generated Fortran.
+eta = external_function(
+ "eta",
+ arguments=("T",),
+ derivatives={"T": "deta_dT"},
+ policy="piecewise_active",
+ fortran_name="eta_T",
+)
+
+xjac = coefficient("xjac")
+eps_cyl = coefficient("eps_cyl")
+F0 = coefficient("F0")
+eta_num = coefficient("eta_num")
+current_source = coefficient("current_source")
+GAMMA = coefficient("GAMMA")
+gamma = coefficient("gamma")
+visco_num = coefficient("visco_num")
+D_par = coefficient("D_par")
+ZK_par = coefficient("ZK_par")
+D_prof = coefficient("D_prof")
+ZK_prof = coefficient("ZK_prof")
+particle_source = coefficient("particle_source")
+heat_source = coefficient("heat_source")
+visco = external_function(
+ "visco", arguments=("T",), derivatives={"T": "dvisco_dT"},
+ policy="piecewise_active", fortran_name="visco_T"
+)
+eta_ohmic = external_function(
+ "eta_ohmic", arguments=("T",), derivatives={"T": "deta_dT_ohm"},
+ policy="piecewise_active", fortran_name="eta_T_ohm"
+)
+
+
+def _rho_hat():
+ return R**2 * rho
+
+
+def _rho_hat_time():
+ # The momentum time coefficient is frozen in the Fortran assembly.
+ return R**2 * freeze(rho)
+
+
+def _pressure():
+ return rho * T
+
+
+def _vv2():
+ """Squared poloidal velocity used by the momentum equation."""
+
+ return R**2 * (dR(u)**2 + dZ(u)**2)
+
+
+def _parallel_gradient(value, flux):
+ """Full parallel-gradient factor in the physical-coordinate form."""
+
+ return bracket(value, flux) / R + F0 * dphi(value) / R**2
+
+
+def _parallel_norm():
+ return (F0**2 + dR(psi)**2 + dZ(psi)**2) / R**2
+
+
+def induction_equation_1():
+ """Return model-199 equation 1 in its already-integrated weak form.
+
+ The expression is written in physical (R, Z, phi) derivatives. The
+ factor ``xjac`` is explicit because the existing element routine performs
+ the poloidal mapping and Gaussian weight outside the equation terms.
+ """
+
+ v = test_function("v")
+ B = (
+ v * eta(T) * (j - current_source) / R * xjac
+ + v * xjac * bracket(psi, u)
+ - v * eps_cyl * F0 / R * dphi(u) * xjac
+ + eta_num * xjac * dot(grad(v), grad(j))
+ )
+ A = v * psi / R * xjac
+ return EvolutionEquation("model199_induction", v, A, B)
+
+
+def momentum_equation_2():
+ """Return model-199 equation 2 in its integrated weak form."""
+
+ v = test_function("v")
+ rho_hat = _rho_hat()
+ rho_hat_time = _rho_hat_time()
+ pressure = _pressure()
+ B = (
+ -sp.Rational(1, 2) * _vv2()
+ * (dR(v) * dZ(rho_hat) - dZ(v) * dR(rho_hat)) * xjac
+ - rho_hat * R**2 * omega * poiss_bracket_st(v, u)
+ + v * poiss_bracket_st(psi, j)
+ - visco(T) * R * dot(grad(v), grad(omega)) * xjac
+ - v * eps_cyl * F0 / R * dphi(j) * xjac
+ + R**2 * poiss_bracket_st(v, pressure)
+ - visco_num * laplacian(v) * laplacian(omega) * xjac
+ )
+ A = -R * rho_hat_time * dot(grad(v), grad(u)) * xjac
+ return EvolutionEquation("model199_momentum", v, A, B)
+
+
+def current_constraint_equation_3():
+ """Return model-199 equation 3, a static current constraint."""
+
+ v = test_function("v")
+ C = (dot(grad(v), grad(psi)) + v * j) / R * xjac
+ return ConstraintEquation("model199_current_constraint", v, C, kind="static")
+
+
+def vorticity_constraint_equation_4():
+ """Return model-199 equation 4, a static vorticity constraint."""
+
+ v = test_function("v")
+ C = (dot(grad(v), grad(u)) + v * omega) * R * xjac
+ return ConstraintEquation("model199_vorticity_constraint", v, C, kind="static")
+
+
+def density_equation_5():
+ """Return model-199 equation 5 (density transport)."""
+
+ v = test_function("v")
+ rho_hat = R**2 * rho
+ B = (
+ v * R * particle_source * xjac
+ + v * R**2 * poiss_bracket_st(rho, u)
+ + 2 * v * R * rho * dZ(u) * xjac
+ - (D_par - D_prof) * R / _parallel_norm()
+ * _parallel_gradient(v, psi) * _parallel_gradient(rho, psi) * xjac
+ - D_prof * R * dot(grad(v), grad(rho)) * xjac
+ - D_prof * R * eps_cyl**2 / R**2 * dphi(v) * dphi(rho) * xjac
+ )
+ A = v * R * rho * xjac
+ return EvolutionEquation("model199_density", v, A, B)
+
+
+def temperature_equation_6():
+ """Return model-199 equation 6 (temperature transport)."""
+
+ v = test_function("v")
+ B = (
+ v * R * heat_source * xjac
+ + v * R**2 * poiss_bracket_st(T, u)
+ + 2 * (GAMMA - 1) * v * R * T * dZ(u) * xjac
+ - (ZK_par - ZK_prof) * R / _parallel_norm()
+ * _parallel_gradient(v, psi) * _parallel_gradient(T, psi) * xjac
+ - ZK_prof * R * dot(grad(v), grad(T)) * xjac
+ - ZK_prof * R / R**2 * dphi(v) * dphi(T) * xjac
+ + v * (gamma - 1) * eta_ohmic(T) * (j / R)**2 * R * xjac
+ )
+ A = v * R * T * xjac
+ return EvolutionEquation("model199_temperature", v, A, B)
diff --git a/util/equation_codegen/src/jorek_equations/model600.py b/util/equation_codegen/src/jorek_equations/model600.py
new file mode 100644
index 0000000000..76888267bf
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/model600.py
@@ -0,0 +1,1052 @@
+"""Model-600 weak equations, extending the model-199 DSL definitions."""
+
+import functools
+
+import sympy as sp
+
+from .equations import ConstraintEquation, EvolutionEquation
+from .external import external_function
+from .model199 import FIELDS as MODEL199_FIELDS
+from .model199 import GAMMA, gamma, j, omega, psi, rho, T, u, xjac
+from .operators import R, dR, dZ, dphi, ds, dt, poiss_bracket_st, poiss_bracket, laplacian, dot, grad
+from .symbols import coefficient, field, freeze, test_function
+
+
+# Model-600 extension fields. The model199 declarations are reused so their
+# current/trial and Fortran naming conventions remain identical.
+# -------------------------------------------------------------------------
+# Fields
+# -------------------------------------------------------------------------
+
+vpar = field("vpar", fortran_current="vpar0", fortran_trial="vpar")
+Ti = field("Ti", fortran_current="Ti0", fortran_trial="Ti")
+Te = field("Te", fortran_current="Te0", fortran_trial="Te")
+rhon = field("rhon", fortran_current="rn0", fortran_trial="rhon")
+rhoimp = field("rhoimp", fortran_current="rimp0", fortran_trial="rhoimp")
+
+FIELDS = MODEL199_FIELDS + (vpar, Ti, Te, rhon, rhoimp)
+
+# -------------------------------------------------------------------------
+# Scalar work values
+# -------------------------------------------------------------------------
+# Values the element routine computes once per quadrature point and
+# freezes in the Newton tangent.
+Jb = coefficient("Jb")
+aux_jre_ind = coefficient("aux_jre_ind")
+tauIC = coefficient("tauIC")
+tstep = coefficient("tstep")
+r0_corr = coefficient("r0_corr")
+F0 = coefficient("F0")
+factor_psi = tuple(coefficient("factor_psi_{}".format(index)) for index in range(1, 7))
+visco_fact_old = coefficient("visco_fact_old")
+visco_fact_new = coefficient("visco_fact_new")
+dvisco_dT = coefficient("dvisco_dT")
+fact_conservative_u = coefficient("fact_conservative_u")
+delta_n_convection = coefficient("delta_n_convection")
+tgnum_u = coefficient("tgnum_u")
+amu_neo_prof = coefficient("amu_neo_prof")
+aki_neo_prof = coefficient("aki_neo_prof")
+Btheta2 = coefficient("Btheta2")
+epsil = coefficient("epsil")
+alpha_e = coefficient("alpha_e")
+alpha_i = coefficient("alpha_i")
+alpha_imp = coefficient("alpha_imp")
+Sion_T = coefficient("Sion_T")
+Srec_T = coefficient("Srec_T")
+particle_source = coefficient("particle_source")
+source_pellet = coefficient("source_pellet")
+source_bg_drift = coefficient("source_bg_drift")
+source_imp_drift = coefficient("source_imp_drift")
+aux_P_par_re = coefficient("aux_P_par_re")
+aux_P_perp_re = coefficient("aux_P_perp_re")
+aux_divPIR_perp = coefficient("aux_divPIR_perp")
+aux_divPIZ_perp = coefficient("aux_divPIZ_perp")
+D_par_local = coefficient("D_par_local")
+D_par_local_imp = coefficient("D_par_local_imp")
+D_par_sc_num = coefficient("D_par_sc_num")
+D_par_imp_sc_num = coefficient("D_par_imp_sc_num")
+D_perp_num_psin = coefficient("D_perp_num_psin")
+D_prof = coefficient("D_prof")
+D_prof_imp = coefficient("D_prof_imp")
+tau_sc = coefficient("tau_sc")
+tgnum_rho = coefficient("tgnum_rho")
+V_prof_pinch = coefficient("V_prof_pinch")
+aux_rho0 = coefficient("aux_rho0")
+visco_par = coefficient("visco_par")
+visco_par_sc_num = coefficient("visco_par_sc_num")
+visco_par_num = coefficient("visco_par_num")
+visco_par_par = coefficient("visco_par_par")
+tgnum_vpar = coefficient("tgnum_vpar")
+aux_mom_par0 = coefficient("aux_mom_par0")
+dV_dpsi_source = coefficient("dV_dpsi_source")
+Dn_perp_num = coefficient("Dn_perp_num")
+Dn0x = coefficient("Dn0x")
+Dn0y = coefficient("Dn0y")
+Dn0p = coefficient("Dn0p")
+source_neutral_drift = coefficient("source_neutral_drift")
+tgnum_rhoimp = coefficient("tgnum_rhoimp")
+heat_source_i = coefficient("heat_source_i")
+ZK_i_perp_num_psin = coefficient("ZK_i_perp_num_psin")
+tgnum_Ti = coefficient("tgnum_Ti")
+visco_par_heating = coefficient("visco_par_heating")
+implicit_heat_source = coefficient("implicit_heat_source")
+Tie_min_neg = coefficient("Tie_min_neg")
+aux_E0_Ti = coefficient("aux_E0_Ti")
+heat_source_e = coefficient("heat_source_e")
+ZK_e_perp_num_psin = coefficient("ZK_e_perp_num_psin")
+tgnum_Te = coefficient("tgnum_Te")
+ksi_ion_norm = coefficient("ksi_ion_norm")
+aux_jre = coefficient("aux_jre")
+aux_E0_Te = coefficient("aux_E0_Te")
+power_dens_teleport_ju = coefficient("power_dens_teleport_ju")
+heat_source_total = coefficient("heat_source")
+ZK_perp_num_psin = coefficient("ZK_perp_num_psin")
+tgnum_T = coefficient("tgnum_T")
+T_min_neg = coefficient("T_min_neg")
+aux_E0 = coefficient("aux_E0")
+
+# -------------------------------------------------------------------------
+# State-dependent work values
+# -------------------------------------------------------------------------
+# Quantities the element routine supplies together with the derivatives
+# the Jacobian needs; declaring the interface lets the generator
+# differentiate them by name.
+eta = external_function(
+ "eta600",
+ arguments=("T", "rho", "rhoimp"),
+ derivatives={"T": "deta_dT", "rho": "deta_dr0", "rhoimp": "deta_drimp0",},
+ policy="piecewise_active",
+ fortran_name="eta_T",
+)
+
+eta_num_T = external_function(
+ "eta_num600",
+ arguments=("T",),
+ derivatives={"T": "deta_num_dT"},
+ policy="piecewise_active",
+ fortran_name="eta_num_T",
+)
+
+visco = external_function(
+ "visco600",
+ arguments=("T",),
+ derivatives={"T": "dvisco_dT"},
+ policy="piecewise_active",
+ fortran_name="visco_T",
+)
+
+visco_num = external_function(
+ "visco_num600",
+ arguments=("T",),
+ # The element routine computes dvisco_num_dT but deliberately does not
+ # place that derivative in the momentum AMAT temperature columns.
+ derivatives={},
+ policy="frozen",
+ fortran_name="visco_num_T",
+)
+
+alpha_i_state = external_function(
+ "alpha_i600", arguments=(), derivatives={}, policy="frozen",
+ fortran_name="alpha_i",
+)
+
+dvisco_state = external_function(
+ "dvisco_state600", arguments=("T",),
+ derivatives={"T": "d2visco_dT2"},
+ policy="piecewise_active", fortran_name="dvisco_dT",
+)
+
+Sion_rate = external_function(
+ "Sion_rate600", arguments=("T",), derivatives={"T": "dSion_dT"},
+ policy="piecewise_active", fortran_name="Sion_T",
+)
+
+Srec_rate = external_function(
+ "Srec_rate600", arguments=("T",), derivatives={"T": "dSrec_dT"},
+ policy="piecewise_active", fortran_name="Srec_T",
+)
+
+alpha_e_state = external_function(
+ "alpha_e600", arguments=("T",), derivatives={"T": "dalpha_e_dT"},
+ policy="piecewise_active", fortran_name="alpha_e",
+)
+
+alpha_e_bis_state = external_function(
+ "alpha_e_bis", arguments=("T",), derivatives={"T": "alpha_e_tri"},
+ policy="piecewise_active", fortran_name="alpha_e_bis",
+)
+
+alpha_e_temperature = external_function(
+ "alpha_e_temperature600", arguments=("T",), derivatives={"T": "alpha_e_bis"},
+ policy="piecewise_active", fortran_name="alpha_e_T",
+)
+
+alpha_imp_bis_state = external_function(
+ "alpha_imp_bis", arguments=("T",), derivatives={"T": "alpha_imp_tri"},
+ policy="piecewise_active", fortran_name="alpha_imp_bis",
+)
+
+alpha_imp_temperature = external_function(
+ "alpha_imp_temperature600", arguments=("T",),
+ derivatives={"T": "alpha_imp_bis"},
+ policy="piecewise_active", fortran_name="alpha_imp_T",
+)
+
+corr_neg_dens = external_function(
+ "corr_neg_dens",
+ arguments=("rho",),
+ derivatives={"rho": "dr0_corr_dn"},
+ policy="piecewise_active",
+ fortran_name="corr_neg_dens",
+)
+
+W_dia_single = external_function(
+ "W_dia_single600",
+ arguments=("rho", "T"),
+ derivatives={"rho": "W_dia_rho", "T": "W_dia_T"},
+ policy="piecewise_active",
+ fortran_name="W_dia",
+)
+
+W_dia_two = external_function(
+ "W_dia_two600",
+ arguments=("rho", "Ti"),
+ derivatives={"rho": "W_dia_rho", "Ti": "W_dia_Ti",},
+ policy="piecewise_active",
+ fortran_name="W_dia",
+)
+
+ZKi_par = external_function(
+ "ZKi_par600", arguments=("Ti",), derivatives={"Ti": "dZKi_par_dT"},
+ policy="piecewise_active", fortran_name="ZKi_par_T",
+)
+
+ZKi_perp = external_function(
+ "ZKi_prof600", arguments=("rho",), derivatives={"rho": "dZKi_prof_drho"},
+ policy="piecewise_active", fortran_name="ZKi_prof",
+)
+
+visco_heating = external_function(
+ "visco_heating600", arguments=("Te",),
+ derivatives={"Te": "dvisco_dT_heating"},
+ policy="piecewise_active", fortran_name="visco_T_heating",
+)
+
+Ti_e_exchange = external_function(
+ "dTi_e600", arguments=("Ti", "Te", "rho", "rhoimp"),
+ derivatives={"Ti": "ddTi_e_dTi", "Te": "ddTi_e_dTe", "rho": "ddTi_e_drho", "rhoimp": "ddTi_e_drhoimp"},
+ policy="piecewise_active", fortran_name="dTi_e",
+)
+
+Ti_floor = external_function(
+ "Ti_floor600", arguments=("Ti",), derivatives={"Ti": "dTi_floor"},
+ policy="piecewise_active", fortran_name="Ti0_floor",
+)
+
+Ti_floor_exp = external_function(
+ "Ti_floor_exp600", arguments=("Ti",),
+ derivatives={"Ti": "dTi_floor_exp"},
+ policy="piecewise_active", fortran_name="Ti_floor_exp",
+)
+
+corr_neg_dens_imp = external_function(
+ "corr_neg_dens_imp", arguments=("rhoimp",),
+ derivatives={"rhoimp": "drimp0_corr_dn"},
+ policy="piecewise_active", fortran_name="corr_neg_dens_imp",
+)
+
+E_ion_bg_state = external_function(
+ "E_ion_bg600", arguments=(), derivatives={}, policy="frozen",
+ fortran_name="E_ion_bg",
+)
+
+ZKe_par = external_function(
+ "ZKe_par600", arguments=("Te",), derivatives={"Te": "dZKe_par_dT"},
+ policy="piecewise_active", fortran_name="ZKe_par_T",
+)
+
+ZKe_perp = external_function(
+ "ZKe_prof600", arguments=("rho",), derivatives={"rho": "dZKe_prof_drho"},
+ policy="piecewise_active", fortran_name="ZKe_prof",
+)
+
+eta_ohm_e = external_function(
+ "eta_ohm600", arguments=("Te", "rho", "rhoimp"),
+ derivatives={"Te": "deta_dT_ohm", "rho": "deta_dr0_ohm", "rhoimp": "deta_drimp0_ohm"},
+ policy="piecewise_active", fortran_name="eta_T_ohm",
+)
+
+dE_ion_dT_state = external_function(
+ "dE_ion_dT", arguments=("Te",), derivatives={}, policy="frozen",
+)
+
+E_ion_state = external_function(
+ "E_ion600", arguments=("Te",), derivatives={"Te": "dE_ion_dT"},
+ policy="piecewise_active", fortran_name="E_ion",
+)
+
+LradDrays = external_function(
+ "LradDrays600", arguments=("Te",), derivatives={"Te": "dLradDrays_dT"},
+ policy="piecewise_active", fortran_name="LradDrays_T",
+)
+
+LradDcont = external_function(
+ "LradDcont600", arguments=("Te",),
+ derivatives={"Te": "dLradDcont_dT_corr"},
+ policy="piecewise_active", fortran_name="LradDcont_corr",
+)
+
+frad_bg_state = external_function(
+ "frad_bg600", arguments=("Te",), derivatives={"Te": "dfrad_bg_dT"},
+ policy="piecewise_active", fortran_name="frad_bg",
+)
+
+Lrad_state = external_function(
+ "Lrad600", arguments=("Te",), derivatives={"Te": "dLrad_dT"},
+ policy="piecewise_active", fortran_name="Lrad",
+)
+
+Te_i_exchange = external_function(
+ "dTe_i600", arguments=("Ti", "Te", "rho", "rhoimp"),
+ derivatives={"Ti": "ddTe_i_dTi", "Te": "ddTe_i_dTe", "rho": "ddTe_i_drho", "rhoimp": "ddTe_i_drhoimp"},
+ policy="piecewise_active", fortran_name="dTe_i",
+)
+
+Te_floor = external_function(
+ "Te_floor600", arguments=("Te",), derivatives={"Te": "dTe_floor"},
+ policy="piecewise_active", fortran_name="Te0_floor",
+)
+
+Te_floor_exp = external_function(
+ "Te_floor_exp600", arguments=("Te",),
+ derivatives={"Te": "dTe_floor_exp"},
+ policy="piecewise_active", fortran_name="Te_floor_exp",
+)
+
+corr_neg_dens_n = external_function(
+ "corr_neg_dens_n", arguments=("rhon",),
+ derivatives={"rhon": "drn0_corr_dn"},
+ policy="piecewise_active", fortran_name="corr_neg_dens_n",
+)
+
+ZK_par = external_function(
+ "ZK_par600", arguments=("T",), derivatives={"T": "dZK_par_dT"},
+ policy="piecewise_active", fortran_name="ZK_par_T",
+)
+
+ZK_perp = external_function(
+ "ZK_prof600", arguments=("rho",), derivatives={"rho": "dZK_prof_drho"},
+ policy="piecewise_active", fortran_name="ZK_prof",
+)
+
+T_floor = external_function(
+ "T_floor600", arguments=("T",), derivatives={"T": "dT_floor"},
+ policy="piecewise_active", fortran_name="T0_floor",
+)
+
+T_floor_exp = external_function(
+ "T_floor_exp600", arguments=("T",), derivatives={"T": "dT_floor_exp"},
+ policy="piecewise_active", fortran_name="T_floor_exp",
+)
+
+# -------------------------------------------------------------------------
+# Operators and shared term groups
+# -------------------------------------------------------------------------
+def _rho_hat(value=rho):
+ return R**2 * value
+
+def _B2(flux):
+ """Total magnetic field squared (B^2)"""
+ return (F0**2 + grad(flux)[0]**2 + grad(flux)[1]**2) / R**2
+
+def _B2_pol(flux):
+ """Poloidal magnetic field squared (Bpol^2)"""
+ return (grad(flux)[0]**2 + grad(flux)[1]**2) / R**2
+
+def _B_dot_grad(value, flux=psi, *, st_form=False):
+ """Return B.grad(value). st_form computes the Poiss bracket in st coordinates"""
+
+ if st_form:
+ return (F0 / R * dphi(value) + poiss_bracket_st(value, flux) / xjac) / R
+
+ return (F0 / R * dphi(value) + poiss_bracket(value, flux)) / R
+
+def _diffusion_tot_intg_by_parts(test, value):
+ """Total diffusion operator, including its toroidal part, integrated by parts."""
+
+ return dot(grad(test), grad(value)) + dphi(test) * dphi(value) / R**2
+
+def _u_convection(density, *, bracket=poiss_bracket_st):
+ """Advection and compression of a particle density by u flow."""
+
+ return R * bracket(density, u) / xjac + 2 * density * dZ(u)
+
+def _parallel_convection(quantity, vpar=vpar):
+ """Convection of a quantity density along the field by ``vpar``."""
+
+ return (
+ - vpar * _B_dot_grad(quantity, st_form=True)
+ - quantity * _B_dot_grad(vpar, st_form=True)
+ )
+
+def _dens_tgnum_intg_by_parts(v, density):
+ """Taylor-Galerkin stabilization of a particle-density equation."""
+
+ return (
+ -sp.Rational(1, 4) * R**2 * poiss_bracket(density, u) * poiss_bracket(v, u) * tstep
+ -sp.Rational(1, 4) * vpar**2 * _B_dot_grad(density) * _B_dot_grad(v) * tstep
+ )
+
+def _par_diff_intg_by_parts(v, density):
+ """Parallel diffusion term along the field integrated by parts."""
+
+ return - _B_dot_grad(v) * _B_dot_grad(density) / _B2(psi)
+
+def _diamagnetic_pressure(*, with_TiTe=False):
+ """Ion pressure ``Pi = (rho + alpha_i*rhoimp)*Ti`` of ``construct_pressure``.
+
+ The one-temperature model evolves the total temperature and supplies
+ Ti0 = T0/2, so the ion temperature is half of the evolved field there.
+ ``alpha_i`` carries no temperature dependence in either model.
+ """
+
+ ion_temperature = Ti if with_TiTe else T / 2
+ return (rho + rhoimp * alpha_i_state()) * ion_temperature
+
+def _diamagnetic_viscosity(*, with_TiTe=False):
+ """The element routine's ``W_dia``.
+
+ It is built from the ion pressure alone, so it carries no explicit
+ electron-temperature dependence.
+ """
+
+ return W_dia_two(rho, Ti) if with_TiTe else W_dia_single(rho, T)
+
+
+def _press_u_convection(pressure, *, bracket):
+ """Convection and compression of a pressure by the u (ExB) flow."""
+
+ return (R * bracket(pressure, u) / xjac + 2 * GAMMA * pressure * dZ(u))
+
+
+def _press_parallel_convection(pressure):
+ """Convection and compression of a pressure along the field by ``vpar``."""
+
+ return (
+ - vpar * _B_dot_grad(pressure, st_form=True)
+ - GAMMA * pressure * _B_dot_grad(vpar, st_form=True)
+ )
+
+
+def _energy_tgnum(v, pressure, factor):
+ """Taylor-Galerkin stabilization of an energy equation. `."""
+
+ return (
+ -factor * sp.Rational(1, 4) * R**2 * poiss_bracket(pressure, u) * poiss_bracket(v, u) * tstep
+ - factor * sp.Rational(1, 4) * vpar**2 * _B_dot_grad(pressure) * _B_dot_grad(v) * tstep
+ )
+
+
+def _heating_floor(v, exponential, floor, minimum):
+ """Implicit heating that keeps a temperature away from its floor. `."""
+
+ return implicit_heat_source * (gamma - 1) * v * (sp.Rational(1,2) * minimum * (1 + exponential) - floor)
+
+
+def _released_kinetic_energy(T_or_Te):
+ """Particle sources whose kinetic energy is released into the ions."""
+
+ return (
+ (rho + alpha_e_state(T_or_Te) * rhoimp) * rhon * Sion_rate(T_or_Te)
+ + particle_source + source_pellet + source_bg_drift + source_imp_drift
+ )
+
+
+def _parallel_viscous_heating(v):
+ """Heating from the parallel (``vpar``) viscosity. """
+
+ grad_vpar = grad(vpar)
+ return (GAMMA - 1) * visco_par_heating * (v * dot(grad_vpar, grad_vpar) + vpar * dot(grad(v), grad_vpar))
+
+
+def _u_viscous_heating(v, heating):
+ """Heating from the perpendicular (poloidal ``u`` flow) viscosity."""
+
+ return (
+ - (GAMMA - 1) * v * heating * R**2 * visco_fact_old * dot(grad(u), grad(omega))
+ - (GAMMA - 1) * v * heating * 2 * R * visco_fact_new * omega * dR(u)
+ - (GAMMA - 1) * v * heating * visco_fact_new * (dR(u) * dR(dphi(dphi(u))) + dZ(u) * dZ(dphi(dphi(u))))
+ )
+
+
+def _kinetic_coupling(v):
+ """Energy and momentum handed over by the kinetic neutral/impurity model"""
+
+ return (
+ (gamma - 1) * sp.Rational(1, 2) * v * aux_rho0 * vpar**2 * _B2(psi)
+ - (gamma - 1) * v * aux_mom_par0 * vpar
+ )
+
+
+def _ionization_energy_transport(v, energy, temperature):
+ """Transport and diffusive flux of the ionization potential energy."""
+
+ d_par_tot = D_par_local + D_par_sc_num * tau_sc - D_prof
+ d_par_imp_tot = D_par_local_imp + D_par_imp_sc_num * tau_sc - D_prof_imp
+ bb2 = _B2(psi)
+ return (GAMMA - 1) * (
+ + v * R * poiss_bracket_st(energy, u) / xjac
+ + 2 * v * energy * dZ(u)
+ - v * F0 * vpar * dphi(energy) / R**2
+ - v * vpar * poiss_bracket_st(energy, psi) / (R * xjac)
+ - v * energy * _B_dot_grad(vpar, st_form=True)
+ - E_ion_state(temperature) * d_par_imp_tot / bb2 * _B_dot_grad(v) * _B_dot_grad(rhoimp)
+ - E_ion_state(temperature) * D_prof_imp * _diffusion_tot_intg_by_parts(v, rhoimp)
+ - E_ion_bg_state() * d_par_tot / bb2 * _B_dot_grad(v) * _B_dot_grad(rho - rhoimp)
+ - E_ion_bg_state() * D_prof * _diffusion_tot_intg_by_parts(v, rho - rhoimp)
+ )
+
+
+# -------------------------------------------------------------------------
+# Equations
+# -------------------------------------------------------------------------
+# In the order of the element routine.
+def induction_equation_1( with_TiTe: bool = False):
+ """Weak form of the induction equation for var_psi."""
+
+ v = test_function("v")
+ T_or_Te = Te if with_TiTe else T
+ Te_gen = Te if with_TiTe else T / 2
+ # Electron pressure
+ Pe = rho * Te_gen + rhoimp * alpha_e_temperature(Te_gen)
+ rho_corr = corr_neg_dens(rho)
+ dV = R * xjac
+
+ # Time derivative
+ A = v * psi / R**2 * dV
+
+ # RHS
+ B = (
+ # -B.grad u
+ + v * poiss_bracket_st(psi, u) / (R * xjac) - v * F0 / R**2 * dphi(u)
+
+ # eta*j
+ + v * eta(T_or_Te, rho, rhoimp)*(j - coefficient("current_source") - Jb) / R**2
+
+ # hyper-resistivity
+ + eta_num_T(T_or_Te) * dot(grad(v), grad(j)) / R
+
+ # Runaway current coupling, -eta*j_RE
+ -v * eta(T_or_Te, rho, rhoimp) * aux_jre_ind / R**2
+
+ # B.grad Pe diamagnetic term (no impurity contribution)
+ - v * 2 * tauIC / (rho_corr * _B2(psi)) * F0**2 / R**3 * poiss_bracket_st(psi, Pe) / xjac
+ + v * 2 * tauIC / (rho_corr * _B2(psi)) * F0**3 / R**4 * dphi(Pe)
+
+ ) * dV
+
+ return EvolutionEquation("model600_induction", v, A, B)
+
+def momentum_equation_2(*, with_TiTe=False, include_neo=False):
+ """Model-600 perpendicular momentum equation (``var_u``). """
+
+ v = test_function("v")
+ T_or_Te = Te if with_TiTe else T
+ alpha_e_value = alpha_e_state(T_or_Te)
+ sion_rate = Sion_rate(T_or_Te)
+ srec_rate = Srec_rate(T_or_Te)
+ dV = R * xjac
+
+ # In the two-temperature branch p0 is the sum of ion and electron
+ # pressures; the single-temperature branch uses the unified T field.
+ if with_TiTe:
+ pressure = rho * (Ti + Te) + rhoimp * (alpha_i_state() * Ti + alpha_e_temperature(Te))
+ else:
+ pressure = rho * T + rhoimp * alpha_imp_temperature(T)
+
+ rho_hat = _rho_hat()
+ grad_v = grad(v)
+ grad_u = grad(u)
+ grad_omega = grad(omega)
+ lap_v = laplacian(v)
+ lap_omega = laplacian(omega)
+ pi = _diamagnetic_pressure(with_TiTe=with_TiTe)
+ W_dia = _diamagnetic_viscosity(with_TiTe=with_TiTe)
+ ne = rho + alpha_e_value * rhoimp
+ neutral_sources = ne * rhon * sion_rate - ne * (rho - rhoimp) * srec_rate
+
+ # Time derivative
+ A = R**2 * dV * (
+ - corr_neg_dens(freeze(rho)) * dot(grad_v, grad_u)
+ - fact_conservative_u * rho * dot(grad_v, grad(freeze(u))) # Not sure one should freeze u
+ )
+
+ # The RHS, in the order the element routine builds it.
+ B = (
+ # Perpendicular inertia (poloidal Jacobian of the kinetic energy).
+ - R * sp.Rational(1, 2) * (dR(u)**2 + dZ(u)**2) * poiss_bracket(v, rho_hat)
+
+ # Vorticity advection.
+ - R**3 * rho * omega * poiss_bracket_st(v, u) / xjac
+
+ # jxB force term. which here is B.gra j
+ + v * poiss_bracket_st(psi, j) / dV
+ - v * F0 * dphi(j) / R**2
+
+ # Pressure advection.
+ + R * poiss_bracket_st(v, pressure) / xjac
+
+ # Perpendicular and toroidal viscosity.
+ - visco(T_or_Te) * R**2 * visco_fact_old * dot(grad_v, grad_omega)
+ - 2 * visco(T_or_Te) * R * visco_fact_new * omega * dR(v)
+ - visco(T_or_Te) * visco_fact_new * (dR(v) * dR(dphi(dphi(u))) + dZ(v) * dZ(dphi(dphi(u))))
+ - visco_num(T_or_Te) * lap_v * lap_omega / R
+
+ # Diamagnetic pressure advection and diamagnetic viscosity.
+ - v * tauIC * 2 * R**3 * poiss_bracket_st(pi, omega) / xjac
+ - tauIC * 2 * R**2 * dZ(pi) * dot(grad_v, grad_u)
+ - v * tauIC * 2 * R**3 * (dR(dZ(u)) * (dR(dR(pi)) - dZ(dZ(pi))) - dR(dZ(pi)) * (dR(dR(u)) - dZ(dZ(u))))
+ + dvisco_state(T_or_Te) * W_dia * dot(grad(Ti if with_TiTe else T / 2), grad_v)
+ + visco(T_or_Te) * W_dia * lap_v
+
+ # Conservative form of the momentum equation.
+ + fact_conservative_u * dot(grad_v, grad_u) * (
+ - R * poiss_bracket(rho_hat, u)
+ + F0 * (rho * dphi(vpar) + vpar * dphi(rho))
+ + R * rho * poiss_bracket(vpar, psi)
+ + R * vpar * poiss_bracket(rho, psi)
+ )
+
+ # Runaway/auxiliary pressure contributions
+ - v * (aux_P_par_re + aux_P_perp_re)
+ + R * (-aux_divPIR_perp * dZ(v) + aux_divPIZ_perp * dR(v))
+
+ # Momentum carried by the neutral and particle sources
+ + R**2 * dot(grad_v, grad_u) * (
+ (1 - delta_n_convection) * neutral_sources
+ + (1 - fact_conservative_u) * (particle_source + source_pellet + source_bg_drift + source_imp_drift)
+ )
+
+ # Taylor-Galerkin stabilization.
+ - tgnum_u * sp.Rational(1, 4) * R**4 * rho* poiss_bracket(omega, u) * poiss_bracket(v, u) * tstep
+ - tgnum_u * sp.Rational(1, 4) * R**2 * omega * fact_conservative_u * poiss_bracket(rho_hat, u) * poiss_bracket(v, u) * tstep
+
+ ) * dV
+
+ # Neo-classical terms still missing!
+
+ return EvolutionEquation("model600_momentum", v, A, B, kind="evolution", amat_variation_overrides={})
+
+def current_constraint_equation_zj():
+ """Model-600 current-definition equation for ``var_zj``."""
+
+ v = test_function("v")
+ C = (dot(grad(v), grad(psi)) + v * j) / R * xjac
+ return ConstraintEquation("model600_current_constraint", v, C, kind="static")
+
+def vorticity_constraint_equation_w():
+ """Model-600 vorticity-definition equation for ``var_w``."""
+
+ v = test_function("v")
+ C = (dot(grad(v), grad(u)) + v * omega) * R * xjac
+ return ConstraintEquation("model600_vorticity_constraint", v, C, kind="static")
+
+def density_equation_rho(*, with_TiTe=False):
+ """Model-600 density equation (``var_rho``)."""
+
+ v = test_function("v")
+ T_or_Te = Te if with_TiTe else T
+ rho_main = rho - rhoimp
+ ne = rho + alpha_e_state(T_or_Te) * rhoimp
+ D_par_tot = D_par_local + D_par_sc_num * tau_sc
+ D_par_imp_tot = D_par_local_imp + D_par_imp_sc_num * tau_sc
+ psi_gradient = grad(psi)
+ dV = R * xjac
+
+ # Time derivative
+ A = v * rho * dV
+
+ # Other terms and RHS
+ B = (
+ # particle sources
+ + v * (particle_source+source_pellet+source_bg_drift+source_imp_drift+aux_rho0)
+
+ # sources/sinks due to neutrals
+ + v * ne * rhon * Sion_rate(T_or_Te)
+ - v * ne * rho_main * Srec_rate(T_or_Te)
+
+ # convection/compression by u variable (ExB)
+ + v * _u_convection(rho)
+
+ # parallel convection/compression by vpar
+ + v * _parallel_convection(rho)
+
+ # parallel and perpendicular diffusion of main ions
+ + (D_par_tot-D_prof) * _par_diff_intg_by_parts(v, rho_main)
+ - D_prof * _diffusion_tot_intg_by_parts(v, rho_main)
+
+ # parallel and perpendicular diffusion of impurities (total mass)
+ + (D_par_imp_tot-D_prof_imp) * _par_diff_intg_by_parts(v, rhoimp)
+ - D_prof_imp * _diffusion_tot_intg_by_parts(v, rhoimp)
+
+ # numerical stabilization
+ + tgnum_rho * _dens_tgnum_intg_by_parts(v, rho)
+ - D_perp_num_psin * laplacian(v) * laplacian(rho)
+
+ # Diamagnetic drift
+ + v * 2 * tauIC * 2 * dZ(_diamagnetic_pressure(with_TiTe=with_TiTe))
+
+ # Pinch term, not sure about this one here
+ - V_prof_pinch / sp.sqrt(psi_gradient[0]**2 + psi_gradient[1]**2) * dot(grad(v), psi_gradient) * rho
+
+ ) * dV
+ return EvolutionEquation("model600_density", v, A, B)
+
+def parallel_velocity_equation_vpar(*, with_TiTe=False, st_form=True):
+ """Model-600 parallel velocity equation (``var_vpar``)."""
+
+ v = test_function("v")
+ T_or_Te = Te if with_TiTe else T
+ bb2 = _B2(psi)
+ dV = R * xjac
+ rho_hat = _rho_hat()
+ psi_gradient = grad(psi)
+ ne = rho + alpha_e_state(T_or_Te) * rhoimp
+ pressure = rho * (Ti + Te) if with_TiTe else rho * T
+ if with_TiTe:
+ pressure += rhoimp * (alpha_i_state() * Ti + alpha_e_temperature(Te))
+ else:
+ pressure += rhoimp * alpha_imp_temperature(T)
+
+ visco_par_eff = visco_par + visco_par_sc_num * tau_sc
+ source_dens_tot = particle_source + source_pellet + source_bg_drift + source_imp_drift
+
+ Bdot_grad = functools.partial(_B_dot_grad, st_form=st_form)
+ # Only the flux is varied in the prescribed rotation profile, so
+ # ``grad(Vt) = dV_dpsi_source*grad(psi)`` is the faithful form.
+ rotation_shear = tuple(
+ parallel - dV_dpsi_source * flux
+ for parallel, flux in zip(grad(vpar), grad(psi))
+ )
+
+ # Time derivative
+ # The B \cdot d (rho vpar B)
+ A = v * (
+ + corr_neg_dens(freeze(rho)) * vpar * _B2(freeze(psi))
+ + corr_neg_dens(freeze(rho)) * freeze(vpar) * _B2_pol(psi) * sp.Rational(1, 2)
+ + fact_conservative_u * rho * freeze(vpar) * _B2(freeze(psi))
+ ) * dV
+
+ # The RHS and others
+ B = (
+ # Parallel pressure gradient.
+ - v * _B_dot_grad(pressure, st_form=True)
+
+ # Parallel advection of the kinetic energy, 0.5*v_par**2*B**2.
+ + sp.Rational(1, 2) * vpar**2 * bb2 * (rho * Bdot_grad(v) + v * Bdot_grad(rho))
+
+ # Term to obtain conservative form of momentum equation
+ + fact_conservative_u * v * vpar * bb2 * (poiss_bracket(rho_hat, u)/R - _B_dot_grad(rho * vpar))
+
+ # Numerical and physical parallel viscosities.
+ - visco_par_num * laplacian(v) * laplacian(vpar)
+ - visco_par_par * F0**2 / (R**2 * bb2) * _B_dot_grad(vpar) * _B_dot_grad(v)
+ - visco_par_eff * dot(grad(v), rotation_shear)
+
+ # External momentum sources
+ + v * aux_mom_par0
+
+ # Momentum carried by the particle sources; not active in conservative form
+ - v * source_dens_tot * vpar * bb2 * (1 - fact_conservative_u)
+ - v * aux_rho0 * vpar * bb2 * (1 - fact_conservative_u)
+
+ # This one should probably be multiplied by (1 - fact_conservative_u
+ + (1 - delta_n_convection) * v * ne * vpar * bb2 * ((rho - rhoimp) * Srec_rate(T_or_Te) - rhon * Sion_rate(T_or_Te))
+
+ # Taylor-Galerkin stabilization.
+ - tgnum_vpar * sp.Rational(1, 4) * rho * vpar**2 * bb2* _B_dot_grad(vpar, st_form=True) * _B_dot_grad(v, st_form=True) * tstep
+ - tgnum_vpar * sp.Rational(1, 4) * v * vpar**2 * bb2 * (1 - fact_conservative_u) * _B_dot_grad(vpar, st_form=True) * _B_dot_grad(rho, st_form=True) * tstep
+ - tgnum_vpar * sp.Rational(1, 4) * vpar**3 * bb2 * fact_conservative_u * _B_dot_grad(rho, st_form=True) * _B_dot_grad(v, st_form=True) * tstep
+
+ # Not sure this pinch term should be here
+ + V_prof_pinch / sp.sqrt(psi_gradient[0]**2 + psi_gradient[1]**2) * dot(psi_gradient, grad(vpar)) * rho * v
+
+ ) * dV
+
+ return EvolutionEquation("model600_parallel_velocity", v, A, B)
+
+def impurity_density_equation_rhoimp(*, with_TiTe=False):
+ """Model-600 impurity-density equation (``var_rhoimp``)"""
+
+ v = test_function("v")
+ D_par_imp_tot = D_par_local_imp + D_par_imp_sc_num * tau_sc
+ dV = R * xjac
+
+ # Time derivative
+ A = v * rhoimp * dV
+
+ # Other terms and RHS
+ B = (
+ # particle sources
+ + v * source_imp_drift
+
+ # convection/compression by u variable (ExB)
+ + v * _u_convection(rhoimp)
+
+ # parallel convection/compression by vpar
+ + v * _parallel_convection(rhoimp)
+
+ # parallel and perpendicular diffusion
+ + (D_par_imp_tot-D_prof_imp) * _par_diff_intg_by_parts(v, rhoimp)
+ - D_prof_imp * _diffusion_tot_intg_by_parts(v, rhoimp)
+
+ # numerical stabilization
+ + tgnum_rhoimp * _dens_tgnum_intg_by_parts(v, rhoimp)
+ - Dn_perp_num * laplacian(v) * laplacian(rhoimp)
+ ) * dV
+
+ return EvolutionEquation("model600_impurity_density", v, A, B)
+
+def neutral_density_equation_rhon(*, with_TiTe=False):
+ """Model-600 fluid-neutral density equation (``var_rhon``)."""
+
+ v = test_function("v")
+ T_or_Te = Te if with_TiTe else T
+ ne = rho + alpha_e_state(T_or_Te) * rhoimp
+ rho_main = rho - rhoimp
+ dV = R * xjac
+
+ # Time derivative
+ A = v * rhon * dV
+
+ # Other terms and RHS
+ B = (
+ # anisotropic diffusion
+ - (Dn0x * dR(v) * dR(rhon) + Dn0y * dZ(v) * dZ(rhon))
+ - Dn0p * dphi(v) * dphi(rhon) / R**2
+
+ # convection/compression with the plasma flow
+ + delta_n_convection * v * (_u_convection(rhon) + _parallel_convection(rhon))
+
+ # ionization and recombination with the main/impurity ions
+ - v * ne * rhon * Sion_rate(T_or_Te)
+ + v * ne * rho_main * Srec_rate(T_or_Te)
+
+ # neutral source
+ + v * source_neutral_drift
+
+ # numerical stabilization
+ - Dn_perp_num * laplacian(v) * laplacian(rhon)
+ ) * dV
+
+ return EvolutionEquation("model600_neutral_density", v, A, B)
+
+def ion_energy_equation_Ti(*, st_form=True):
+ """Model-600 ion energy equation (``var_Ti``).
+
+ ``st_form`` selects the spelling of the poloidal advection bracket.
+ """
+
+ v = test_function("v")
+ bracket = poiss_bracket_st if st_form else (lambda a, b: poiss_bracket(a, b) * xjac)
+ dV = R * xjac
+ ion_density = rho + alpha_i_state() * rhoimp
+ ion_pressure = ion_density * Ti
+
+ # Time derivative
+ A = v * (
+ corr_neg_dens(rho) + alpha_i_state() * corr_neg_dens_imp(rhoimp)
+ ) * Ti * dV
+
+ # RHS, many terms are common to T/Ti/Te (so many operators are reused to avoid duplications)
+ B = (
+ # convection/compression by u (ExB)
+ + v * _press_u_convection(ion_pressure, bracket=bracket)
+
+ # parallel convection/compression by vpar
+ + v * _press_parallel_convection(ion_pressure)
+
+ # parallel heat conduction
+ + (ZKi_par(Ti) - ZKi_perp(rho)) * _par_diff_intg_by_parts(v, Ti)
+
+ # perpendicular (total) heat conduction
+ - ZKi_perp(rho) * _diffusion_tot_intg_by_parts(v, Ti)
+
+ # heat sources/sinks
+ + v * (heat_source_i + aux_E0_Ti)
+
+ # energy exchange with electrons
+ + v * Ti_e_exchange(Ti, Te, rho, rhoimp)
+
+ # neutral recombination sink
+ - v * Ti * corr_neg_dens(rho)**2 * Srec_rate(Te)
+
+ # kinetic energy released by particle sources (friction heating)
+ + v * (GAMMA - 1) / 2 * (vpar**2 * _B2(psi) + R**2 * (dR(u)**2 + dZ(u)**2)) * _released_kinetic_energy(Te)
+
+ # parallel viscous heating
+ + _parallel_viscous_heating(v)
+
+ # perpendicular (u) viscous heating
+ + _u_viscous_heating(v, visco_heating(Te))
+
+ # energy and momentum from the kinetic neutral/impurity model
+ + _kinetic_coupling(v)
+
+ # numerical stabilization
+ + _energy_tgnum(v, ion_pressure, tgnum_Ti)
+ - ZK_i_perp_num_psin * laplacian(v) * laplacian(Ti)
+
+ # implicit heat source to avoid negative temperatures
+ + _heating_floor(v, Ti_floor_exp(Ti), Ti_floor(Ti), Tie_min_neg)
+
+ ) * dV
+ return EvolutionEquation("model600_ion_energy", v, A, B)
+
+def electron_energy_equation_Te(*, st_form=True):
+ """Model-600 electron energy equation (``var_Te``).
+
+ ``st_form`` applies to the electron-pressure advection only; the
+ ionization-energy advection is always written in element coordinates.
+ """
+
+ v = test_function("v")
+ bracket = poiss_bracket_st if st_form else (lambda a, b: poiss_bracket(a, b) * xjac)
+ dV = R * xjac
+ electron_pressure = rho * Te + rhoimp * alpha_e_temperature(Te)
+ electron_density = corr_neg_dens(rho) + alpha_e_state(Te) * corr_neg_dens_imp(rhoimp)
+ ionization_energy = E_ion_state(Te) * rhoimp + E_ion_bg_state() * (rho - rhoimp)
+
+ # Time derivative
+ A = v * (
+ corr_neg_dens(rho) * Te + corr_neg_dens_imp(rhoimp) * alpha_e_temperature(Te)
+ + (GAMMA - 1) * ionization_energy
+ ) * dV
+
+ # RHS, many terms are common to T/Ti/Te (so many operators are reused to avoid duplications)
+ B = (
+ # convection/compression by u (ExB)
+ + v * _press_u_convection(electron_pressure, bracket=bracket)
+
+ # parallel convection/compression by vpar
+ + v * _press_parallel_convection(electron_pressure)
+
+ # parallel heat conduction
+ + (ZKe_par(Te) - ZKe_perp(rho)) * _par_diff_intg_by_parts(v, Te)
+
+ # perpendicular (total) heat conduction
+ - ZKe_perp(rho) * _diffusion_tot_intg_by_parts(v, Te)
+
+ # heat sources/sinks
+ + v * ( heat_source_e + aux_E0_Te + power_dens_teleport_ju)
+
+ # Ionization sink due to neutrals
+ - v * ksi_ion_norm * (rho + alpha_e_state(Te) * rhoimp) * rhon * Sion_rate(Te)
+
+ # energy exchange with ions
+ + v * Te_i_exchange(Ti, Te, rho, rhoimp)
+
+ # ohmic heating
+ + v * (GAMMA - 1) * eta_ohm_e(Te, rho, rhoimp) * ((j - aux_jre) / R)**2
+
+ # neutral (line) radiation
+ - v * electron_density * corr_neg_dens_n(rhon) * LradDrays(Te)
+
+ # background radiation
+ - v * electron_density * (corr_neg_dens(rho) - corr_neg_dens_imp(rhoimp)) * LradDcont(Te)
+ - v * electron_density * frad_bg_state(Te)
+
+ # impurity radiation
+ - v * electron_density * corr_neg_dens_imp(rhoimp) * Lrad_state(Te)
+
+ # ionization potential transport
+ + _ionization_energy_transport(v, ionization_energy, Te)
+
+ # numerical stabilization
+ + _energy_tgnum(v, electron_pressure, tgnum_Te)
+ - ZK_e_perp_num_psin * laplacian(v) * laplacian(Te)
+
+ # implicit heat source to avoid negative temperatures
+ + _heating_floor(v, Te_floor_exp(Te), Te_floor(Te), Tie_min_neg)
+
+ ) * dV
+ return EvolutionEquation("model600_electron_energy", v, A, B)
+
+def total_energy_equation_T(*, st_form=True):
+ """Model-600 single-temperature energy equation (``var_T``)"""
+
+ v = test_function("v")
+ bracket = poiss_bracket_st if st_form else (lambda a, b: poiss_bracket(a, b) * xjac)
+ dV = R * xjac
+ pressure = rho * T + rhoimp * alpha_imp_temperature(T)
+ electron_density = corr_neg_dens(rho) + alpha_e_state(T) * corr_neg_dens_imp(rhoimp)
+ ionization_energy = E_ion_state(T) * rhoimp + E_ion_bg_state() * (rho - rhoimp)
+
+ # Time derivative
+ A = v * (
+ corr_neg_dens(rho) * T + corr_neg_dens_imp(rhoimp) * alpha_imp_temperature(T)
+ + (GAMMA - 1) * ionization_energy
+ ) * dV
+
+ # RHS, many terms are common to T/Ti/Te (so many operators are reused to avoid duplications)
+ B = (
+ # convection/compression by u (ExB)
+ + v * _press_u_convection(pressure, bracket=bracket)
+
+ # parallel convection/compression by vpar
+ + v * _press_parallel_convection(pressure)
+
+ # parallel heat conduction
+ + (ZK_par(T) - ZK_perp(rho)) * _par_diff_intg_by_parts(v, T)
+
+ # perpendicular (total) heat conduction
+ - ZK_perp(rho) * _diffusion_tot_intg_by_parts(v, T)
+
+ # heat sources/sinks
+ + v * (heat_source_total + aux_E0 + power_dens_teleport_ju)
+
+ # Ionization sink due to neutrals
+ - v * ksi_ion_norm * (rho + alpha_e_state(T) * rhoimp) * rhon * Sion_rate(T)
+
+ # recombination sink due to neutrals
+ - (GAMMA - 1) * v * sp.Rational(1, 2) * T * corr_neg_dens(rho)**2 * Srec_rate(T)
+
+ # ohmic heating
+ + v * (GAMMA - 1) * eta_ohm_e(T, rho, rhoimp) * ((j - aux_jre) / R)**2
+
+ # neutral (line) radiation
+ - v * electron_density * corr_neg_dens_n(rhon) * LradDrays(T)
+
+ # background radiation
+ - v * electron_density * (corr_neg_dens(rho) - corr_neg_dens_imp(rhoimp)) * LradDcont(T)
+ - v * electron_density * frad_bg_state(T)
+
+ # impurity radiation
+ - v * electron_density * corr_neg_dens_imp(rhoimp) * Lrad_state(T)
+
+ # ionization potential transport
+ + _ionization_energy_transport(v, ionization_energy, T)
+
+ # kinetic energy released by particle sources (friction heating)
+ + v * (GAMMA - 1) / 2 * (vpar**2 * _B2(psi) + R**2 * (dR(u)**2 + dZ(u)**2)) * _released_kinetic_energy(T)
+
+ # parallel viscous heating
+ + _parallel_viscous_heating(v)
+
+ # perpendicular (u) viscous heating
+ + _u_viscous_heating(v, visco_heating(T))
+
+ # energy and momentum from the kinetic neutral/impurity model
+ + _kinetic_coupling(v)
+
+ # numerical stabilization
+ + _energy_tgnum(v, pressure, tgnum_T)
+ - ZK_perp_num_psin * laplacian(v) * laplacian(T)
+
+ # implicit heat source to avoid negative temperatures
+ + _heating_floor(v, T_floor_exp(T), T_floor(T), T_min_neg)
+
+ ) * dV
+ return EvolutionEquation("model600_total_energy", v, A, B)
+
+
diff --git a/util/equation_codegen/src/jorek_equations/model600_docs.py b/util/equation_codegen/src/jorek_equations/model600_docs.py
new file mode 100644
index 0000000000..36b80b24dc
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/model600_docs.py
@@ -0,0 +1,409 @@
+"""Generate the model-600 weak-form documentation from ``model600.py``.
+
+``model600.py`` is the only place the equations are written. This module
+reads its source (with ``ast``, without importing or evaluating it) and writes
+the generated regions of ``docs/physics/base_fluid_models/RMHD/weak_form.md``:
+the full equations, the notation tables, the shared operators, the
+per-equation term groups, and two tables built from the checker's data. The
+prose around them is hand-written and kept as it is.
+
+A generated region is delimited by::
+
+
+ ...
+
+"""
+
+import ast
+import re
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+
+from . import model600_notation as notation
+from .latex_render import (
+ ATOM,
+ MUL,
+ SENTINEL,
+ Helper,
+ LatexPrinter,
+ RenderError,
+ _TimeDerivative,
+ _aligned,
+ _call_name,
+ _definition_rows,
+ _fill,
+ _is_variant,
+ _num_den,
+ _rebuild,
+ _signed_terms,
+ functions_table,
+ findings_table,
+)
+
+SOURCE_NAME = "util/equation_codegen/src/jorek_equations/model600.py"
+BEGIN = re.compile(r"^\s*$")
+END = ""
+MATH = "$$"
+
+
+# ---------------------------------------------------------------------------
+# Reading model600.py
+# ---------------------------------------------------------------------------
+
+@dataclass
+class TermGroup:
+ title: str
+ terms: List[Tuple[str, ast.expr]]
+
+
+@dataclass
+class EquationSource:
+ name: str
+ title: str
+ anchor: str
+ prose: str
+ definitions: List[ast.stmt]
+ values: Dict[str, ast.expr]
+ aliases: Dict[str, str]
+ mass: Optional[ast.expr]
+ rhs_name: str
+ rhs: Optional[ast.expr]
+ groups: List[TermGroup] = field(default_factory=list)
+ notes: List[str] = field(default_factory=list)
+
+
+def _title(text):
+ text = text.strip().rstrip(".").strip()
+ return text[:1].upper() + text[1:]
+
+
+def _comment_block(lines, lineno):
+ """The comment lines directly above line ``lineno`` (1-based), top first."""
+
+ comments, index = [], lineno - 2
+ while index >= 0 and lines[index].strip().startswith("#"):
+ comments.insert(0, lines[index].strip().lstrip("#").strip())
+ index -= 1
+ return comments
+
+
+def _split_groups(inner_lines):
+ """Split the lines inside ``B = ( ... )`` into ``(title, code)`` groups at comments."""
+
+ groups, title, code, depth = [], None, [], 0
+ for line in inner_lines:
+ stripped = line.strip()
+ if depth == 0 and stripped.startswith("#"):
+ comment = stripped.lstrip("#").strip()
+ if code:
+ groups.append((title, code))
+ title, code = comment, []
+ else:
+ title = comment if title is None else title + " " + comment
+ continue
+ if depth == 0 and not stripped:
+ if code:
+ groups.append((title, code))
+ title, code = None, []
+ continue
+ code.append(line)
+ cleaned = re.sub(r"#.*", "", line)
+ depth += cleaned.count("(") - cleaned.count(")")
+ if code:
+ groups.append((title, code))
+ return [(_title(t) if t else "Other terms", "\n".join(c)) for t, c in groups]
+
+
+def _local_names(statements):
+ values, aliases = {}, {}
+ for statement in ast.walk(ast.Module(body=list(statements), type_ignores=[])):
+ if isinstance(statement, ast.Assign) and len(statement.targets) == 1 \
+ and isinstance(statement.targets[0], ast.Name):
+ name, value = statement.targets[0].id, statement.value
+ if isinstance(value, ast.Call) and _call_name(value.func) == "partial" and value.args:
+ aliases[name] = _call_name(value.args[0])
+ else:
+ values.setdefault(name, value)
+ return values, aliases
+
+
+def read_equation(function, lines, title, anchor, prose):
+ body = function.body
+ if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant):
+ body = body[1:]
+ values, aliases = _local_names(body)
+ equation = EquationSource(
+ name=function.name, title=title, anchor=anchor, prose=prose,
+ definitions=[], values=values, aliases=aliases, mass=None, rhs_name="", rhs=None,
+ )
+ for statement in body:
+ target = (statement.targets[0].id if isinstance(statement, ast.Assign)
+ and isinstance(statement.targets[0], ast.Name) else None)
+ if isinstance(statement, ast.Return):
+ equation.notes = _comment_block(lines, _first_code_line_above(lines, statement.lineno))
+ continue
+ if target == "A":
+ equation.mass = statement.value
+ elif target in ("B", "C"):
+ equation.rhs_name = target
+ equation.rhs = statement.value
+ text = lines[statement.lineno - 1:statement.end_lineno]
+ if target == "B":
+ if text[0].strip() != "B = (" or not text[-1].strip().startswith(")"):
+ raise RenderError("{}: write the right-hand side as 'B = (' ... ') * dV'".format(
+ function.name))
+ for title_, code in _split_groups(text[1:-1]):
+ expression = ast.parse("(\n" + _dedent(code) + "\n)", mode="eval").body
+ equation.groups.append(TermGroup(title_, _signed_terms(expression)))
+ else:
+ equation.definitions.append(statement)
+ if equation.rhs is None:
+ raise RenderError("{} defines neither B nor C".format(function.name))
+ return equation
+
+
+def _first_code_line_above(lines, lineno):
+ """Skip the blank lines above ``lineno`` so a comment block separated by one is found."""
+
+ while lineno > 1 and not lines[lineno - 2].strip():
+ lineno -= 1
+ return lineno
+
+
+def _dedent(code):
+ import textwrap
+ return textwrap.dedent(code).strip("\n")
+
+
+def read_helpers(tree):
+ helpers, functions = {}, {}
+ for node in tree.body:
+ if isinstance(node, ast.FunctionDef) and node.name in notation.HELPERS:
+ info = notation.HELPERS[node.name]
+ arguments = node.args
+ params = [a.arg for a in arguments.args]
+ defaults = dict(zip(params[len(params) - len(arguments.defaults):], arguments.defaults))
+ for arg, default in zip(arguments.kwonlyargs, arguments.kw_defaults):
+ if default is not None:
+ defaults[arg.arg] = default
+ helpers[node.name] = Helper(
+ name=node.name, template=info["template"], params=params,
+ keyword_only=[a.arg for a in arguments.kwonlyargs],
+ defaults=defaults, display=dict(info["display"]),
+ )
+ functions[node.name] = node
+ missing = [name for name in notation.HELPERS if name not in helpers]
+ if missing:
+ raise RenderError("helpers documented but not in model600.py: " + ", ".join(missing))
+ undocumented = [n.name for n in tree.body if isinstance(n, ast.FunctionDef)
+ and n.name.startswith("_") and n.name not in notation.HELPERS]
+ if undocumented:
+ raise RenderError("add these helpers to model600_notation.HELPERS: " + ", ".join(undocumented))
+ return helpers, functions
+
+
+def read_source(source):
+ tree = ast.parse(source)
+ lines = source.split("\n")
+ helpers, helper_functions = read_helpers(tree)
+ functions = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)}
+ equations = []
+ for name, title, anchor, prose in notation.EQUATIONS:
+ if name not in functions:
+ raise RenderError("equation {} is documented but not in model600.py".format(name))
+ equations.append(read_equation(functions[name], lines, title, anchor, prose))
+ public = [n for n in functions if not n.startswith("_")]
+ missing = [n for n in public if n not in {e.name for e in equations}]
+ if missing:
+ raise RenderError("add these equations to model600_notation.EQUATIONS: " + ", ".join(missing))
+ return helpers, helper_functions, equations
+
+
+# ---------------------------------------------------------------------------
+# Rendering
+# ---------------------------------------------------------------------------
+
+def symbol_map():
+ table = {}
+ table.update(notation.FIELDS)
+ table.update(notation.SYMBOLS)
+ table.update(notation.FUNCTIONS)
+ return table
+
+
+def _without_weight(node):
+ """``node / (R*xjac)``: drop the volume weight ``dV = R*xjac`` from a term."""
+
+ numerator, denominator = _num_den(node)
+ names = [f.id if isinstance(f, ast.Name) else None for f in numerator]
+ if "dV" in names:
+ del numerator[names.index("dV")]
+ return _rebuild(numerator, denominator)
+ if "xjac" not in names:
+ raise RenderError("no volume weight in " + ast.unparse(node))
+ del numerator[names.index("xjac")]
+ names = [f.id if isinstance(f, ast.Name) else None for f in numerator]
+ if "R" in names:
+ del numerator[names.index("R")]
+ else:
+ below = [f.id if isinstance(f, ast.Name) else None for f in denominator]
+ if "R" in below:
+ denominator[below.index("R")] = ast.BinOp(
+ left=ast.Name(id="R", ctx=ast.Load()), op=ast.Pow(), right=ast.Constant(value=2))
+ else:
+ denominator.append(ast.Name(id="R", ctx=ast.Load()))
+ return _rebuild(numerator, denominator)
+
+
+def _math(latex):
+ return [MATH] + latex.split("\n") + [MATH]
+
+
+def render_summary(equation, symbols, helpers):
+ printer = LatexPrinter(symbols, helpers, values=equation.values, aliases=equation.aliases)
+ if equation.rhs_name == "B":
+ rows = ["%s %s" % ("" if i == 0 and sign == "+" else sign, printer.wrap(term, MUL))
+ for i, (sign, term) in enumerate(t for g in equation.groups for t in g.terms)]
+ mass = _without_weight(equation.mass)
+ if _TimeDerivative.has_freeze(mass):
+ lhs = printer.latex(_TimeDerivative(equation.values, helpers).derive(mass))
+ else:
+ lhs = r"\frac{\partial}{\partial t}\Big(%s\Big)" % printer.latex(mass)
+ else:
+ lhs, rows = "0", [printer.latex(_without_weight(equation.rhs))]
+ rows[0] = r"%s \;=\; &%s" % (lhs, rows[0])
+ rows[1:] = ["&" + row for row in rows[1:]]
+ return "\\begin{aligned}\n" + " \\\\\n".join(rows) + "\n\\end{aligned}"
+
+
+def render_helper(function, helper, symbols, helpers):
+ printer = LatexPrinter(symbols, helpers)
+ local = {}
+ for param in helper.params + helper.keyword_only:
+ if param in helper.display:
+ local[param] = helper.display[param]
+ elif param in helper.defaults and param not in helper.keyword_only:
+ local[param] = printer.latex(helper.defaults[param])
+ inner = printer.with_local(local)
+ lhs = _fill(helper.template, {p: (local.get(p, inner.symbol(p)), ATOM)
+ for p in helper.params + helper.keyword_only})
+ body = function.body
+ if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant):
+ body = body[1:]
+ body = [s for s in body if not _is_variant(s)]
+ if not isinstance(body[-1], ast.Return):
+ raise RenderError("helper {} must end with a return".format(function.name))
+ rows = inner.rows(body[:-1])
+ rows.extend(_definition_rows(inner, lhs, body[-1].value))
+ return _aligned(rows)
+
+
+def region_equations(equations, symbols, helpers):
+ out = []
+ for equation in equations:
+ out += ['' % equation.anchor, "", "### " + equation.title, ""]
+ out += _math(render_summary(equation, symbols, helpers))
+ out += ["", "[Definitions and term groups](#details-%s)" % equation.anchor, ""]
+ return out
+
+
+def region_operators(helper_functions, symbols, helpers):
+ out, group = [], None
+ for name, info in notation.HELPERS.items():
+ if info["group"] != group:
+ group = info["group"]
+ out += ["### " + group, ""]
+ out += ["#### " + info["title"], ""]
+ if info["description"]:
+ out += [info["description"], ""]
+ out += _math(render_helper(helper_functions[name], helpers[name], symbols, helpers))
+ out += ["", "Source: `%s`." % name, ""]
+ return out
+
+
+def region_details(equations, symbols, helpers):
+ out = []
+ for equation in equations:
+ printer = LatexPrinter(symbols, helpers, values=equation.values, aliases=equation.aliases)
+ out += ['' % equation.anchor, "", "### " + equation.title, ""]
+ if equation.prose:
+ out += [equation.prose, ""]
+ out += ["Source: `%s` in `%s`." % (equation.name, SOURCE_NAME), ""]
+ definitions = printer.rows(equation.definitions)
+ if definitions:
+ out += ["#### Definitions", ""] + _math(_aligned(definitions)) + [""]
+ if equation.mass is not None:
+ out += ["#### Time derivative", ""]
+ out += _math(r"A = %s" % printer.latex(equation.mass)) + [""]
+ if equation.rhs_name == "C":
+ out += ["#### Constraint", ""] + _math(r"C = %s" % printer.latex(equation.rhs)) + [""]
+ continue
+ out += ["#### Right-hand side", ""]
+ out += _math(r"B = %s" % LatexPrinter(symbols, helpers).latex(
+ _replace_sum(equation.rhs))) + [""]
+ for group in equation.groups:
+ out += ["##### " + group.title, ""]
+ rows = ["&%s %s" % (sign, printer.wrap(term, MUL)) for sign, term in group.terms]
+ out += _math("\\begin{aligned}\n" + " \\\\\n".join(rows) + "\n\\end{aligned}") + [""]
+ if equation.notes:
+ out += ["**Note:** " + " ".join(equation.notes), ""]
+ return out
+
+
+def _replace_sum(node):
+ """``B = (sum) * dV`` shown as ``(sum_k b_k) dV``."""
+
+ return ast.BinOp(left=ast.Name(id="terms", ctx=ast.Load()), op=node.op, right=node.right)
+
+
+def region_notation():
+ out = ["### Work values", "", "| Python | Symbol |", "|---|---|"]
+ out += ["| `%s` | $%s$ |" % (n, s) for n, s in notation.SYMBOLS.items()]
+ out += ["", "### State-dependent functions", "",
+ r"A function of the state is shown with its arguments, e.g. $S_{ion}(T)$;",
+ r"$x^{\mathrm{c}}$ is a density corrected for negative values.", "",
+ "| Python | Symbol |", "|---|---|"]
+ out += ["| `%s` | $%s$ |" % (n, s) for n, s in notation.FUNCTIONS.items()]
+ return out
+
+
+def render_regions(source, reference):
+ helpers, helper_functions, equations = read_source(source)
+ symbols = symbol_map()
+ return {
+ "equations": region_equations(equations, symbols, helpers),
+ "notation": region_notation(),
+ "linearization": functions_table(source),
+ "findings": findings_table(reference),
+ "operators": region_operators(helper_functions, symbols, helpers),
+ "details": region_details(equations, symbols, helpers),
+ }
+
+
+def render_page(text, source, reference):
+ """Return ``text`` with every generated region rewritten."""
+
+ regions = render_regions(source, reference)
+ lines = text.split("\n")
+ output, index, seen = [], 0, set()
+ while index < len(lines):
+ match = BEGIN.match(lines[index])
+ if not match:
+ output.append(lines[index])
+ index += 1
+ continue
+ name = match.group(1)
+ if name not in regions:
+ raise RenderError("unknown generated region {!r}".format(name))
+ end = index + 1
+ while end < len(lines) and lines[end].strip() != END.format(name):
+ end += 1
+ if end == len(lines):
+ raise RenderError("generated region {!r} is not closed".format(name))
+ output += [lines[index], ""] + regions[name] + ["", lines[end]]
+ seen.add(name)
+ index = end + 1
+ missing = set(regions) - seen
+ if missing:
+ raise RenderError("the page has no region for: " + ", ".join(sorted(missing)))
+ return "\n".join(output)
diff --git a/util/equation_codegen/src/jorek_equations/model600_markdown.py b/util/equation_codegen/src/jorek_equations/model600_markdown.py
new file mode 100644
index 0000000000..aefab40d76
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/model600_markdown.py
@@ -0,0 +1,1143 @@
+"""Meld-friendly Markdown exports of model-600 source and generated terms."""
+
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+import re
+
+import sympy as sp
+
+from .equations import LinearizedEquation
+from .linearization import resolve_current, variation
+from .fortran import fortran
+from .model600 import (
+ FIELDS,
+ density_equation_rho,
+ impurity_density_equation_rhoimp,
+ neutral_density_equation_rhon,
+ electron_energy_equation_Te,
+ ion_energy_equation_Ti,
+ total_energy_equation_T,
+ parallel_velocity_equation_vpar,
+ T,
+ Te,
+ current_constraint_equation_zj,
+ induction_equation_1,
+ j,
+ momentum_equation_2,
+ omega,
+ psi,
+ rho,
+ u,
+ vorticity_constraint_equation_w,
+)
+from .fortran_source import normalize_model600_text
+from .operators import SpatialDerivative, expand_derivatives, phi
+from .fortran_source import parse_fortran_expression
+from .symbols import FieldRole, FieldValue, TestFunction, coefficient
+
+
+ROWS = ("psi", "u", "zj", "w", "rho", "vpar", "rhoimp", "rhon", "ti", "te", "t")
+FIELD_NAMES = {
+ field: name
+ for field, name in zip(FIELDS, ("psi", "u", "zj", "w", "rho", "T", "vpar", "Ti", "Te", "rhon", "rhoimp"))
+}
+ASSIGNMENT_RE = re.compile(
+ r"(?P(?:rhs_ij(?:_k)?|amat(?:_n|_k|_kn|_nn)?)\s*\(\s*"
+ r"var_(?Ppsi|u|zj|w|rho|rhoimp|rhon|vpar|Ti|Te|T)\b[^=]*?\))\s*=\s*(?P.*)$",
+ re.I,
+)
+
+
+@dataclass(frozen=True)
+class SourceAssignment:
+ lhs: str
+ row: str
+ line: int
+ expression: str
+ temperature_model: str = "both"
+
+
+def _clean_lhs(lhs):
+ return re.sub(r"\s+", "", lhs).lower()
+
+
+def _read_source_assignments(path):
+ """Read every relevant assignment, retaining repetitions and file order."""
+
+ lines = Path(path).read_text(encoding="utf-8").splitlines()
+ result = []
+ temperature_model = "both"
+ # Stack only conditionals controlled by ``with_TiTe``. It is necessary
+ # because model600 contains further ordinary IF blocks inside each
+ # temperature branch, and a bare ``endif`` must not accidentally change
+ # the active temperature convention.
+ temperature_stack = []
+ velocity_profile_branch = None
+ index = 0
+ while index < len(lines):
+ raw = lines[index]
+ code = raw.split("!", 1)[0]
+ lowered = raw.lower()
+ if re.search(r"\bif\s*\(\s*with_tite\s*\)\s*then", lowered):
+ temperature_stack.append([temperature_model, 0])
+ temperature_model = "two"
+ elif temperature_stack and re.search(r"\bif\s*\([^)]*\)\s*then", lowered):
+ temperature_stack[-1][1] += 1
+ elif (
+ temperature_stack
+ and temperature_stack[-1][1] == 0
+ and re.match(r"^\s*else\b(?!\s*if)", lowered)
+ ):
+ temperature_model = "single"
+ elif temperature_stack and re.search(r"\bend\s*if\b", lowered):
+ if temperature_stack[-1][1]:
+ temperature_stack[-1][1] -= 1
+ else:
+ temperature_model = temperature_stack.pop()[0]
+ # ``normalized_velocity_profile`` defaults to .true.; export that
+ # branch and skip the ``else`` alternative. These blocks contain no
+ # nested conditionals, so a single flag is enough.
+ if re.search(r"\bif\s*\(\s*normalized_velocity_profile\s*\)\s*then", lowered):
+ velocity_profile_branch = True
+ index += 1
+ continue
+ if velocity_profile_branch is not None:
+ if re.match(r"^\s*else\b(?!\s*if)", lowered):
+ velocity_profile_branch = False
+ index += 1
+ continue
+ if re.search(r"\bend\s*if\b", lowered):
+ velocity_profile_branch = None
+ index += 1
+ continue
+ if not velocity_profile_branch:
+ index += 1
+ continue
+ match = ASSIGNMENT_RE.search(code)
+ if match is None:
+ index += 1
+ continue
+ lhs = _clean_lhs(match.group("lhs"))
+ pieces = []
+ current = match.group("rhs")
+ while True:
+ current = current.strip()
+ continued = current.endswith("&")
+ pieces.append(current[:-1].strip() if continued else current)
+ if not continued:
+ break
+ index += 1
+ if index >= len(lines):
+ raise ValueError("Unterminated assignment at line {}".format(index + 1))
+ # JOREK often places a blank/comment line between continued
+ # terms. It is not the end of the Fortran statement.
+ current = lines[index].split("!", 1)[0].strip()
+ while not current:
+ index += 1
+ if index >= len(lines):
+ raise ValueError("Unterminated assignment at line {}".format(index + 1))
+ current = lines[index].split("!", 1)[0].strip()
+ if current.startswith("&"):
+ current = current[1:]
+ expression = " ".join(piece for piece in pieces if piece)
+ # Repeated extension assignments have the form ``a = a + extension``.
+ # Extension assignments commonly use either ``lhs = lhs + term`` or
+ # ``lhs = lhs - term``. Remove only the repeated lhs and preserve the
+ # following sign so the extension remains a source term.
+ expression = re.sub(
+ r"^{}\s*".format(re.escape(lhs)), "",
+ _clean_assignment_refs(expression), flags=re.I,
+ )
+ result.append(
+ SourceAssignment(
+ lhs, match.group("row").lower(), index + 2 - len(pieces),
+ expression, temperature_model,
+ )
+ )
+ index += 1
+ return result
+
+
+def _clean_assignment_refs(expression):
+ """Canonicalize only assignment references, leaving term spelling intact."""
+
+ return re.sub(
+ r"(?:rhs_ij(?:_k)?|amat(?:_n|_k|_kn|_nn)?)\s*\([^)]*\)",
+ lambda match: _clean_lhs(match.group(0)),
+ expression,
+ flags=re.I,
+ )
+
+
+def _split_top_level(expression):
+ """Split an expression at ordered, outermost additive signs."""
+
+ terms = []
+ start = 0
+ depth = 0
+ for index, char in enumerate(expression):
+ if char == "(":
+ depth += 1
+ elif char == ")":
+ depth = max(0, depth - 1)
+ elif char in "+-" and depth == 0 and index > start:
+ previous = expression[index - 1]
+ if previous not in "eEdD":
+ terms.append(expression[start:index].strip())
+ start = index
+ tail = expression[start:].strip()
+ if tail:
+ terms.append(tail)
+ # ``lhs = lhs + &`` continued onto a line that itself starts with ``+``
+ # leaves a piece that is nothing but a sign. It is not a term.
+ return [term for term in terms if term.strip(" +-")]
+
+
+def _normalized_text(text, *, single_temperature=False, two_temperature=False):
+ text = re.sub(
+ r"factor\(\s*var_[a-z0-9_]+\s*,\s*[0-9]+\s*\)", "1", text,
+ flags=re.I,
+ )
+ text = re.sub(r"\bBigR_x\b", "1", text, flags=re.I)
+ result = normalize_model600_text(
+ text,
+ single_temperature=single_temperature,
+ two_temperature=two_temperature,
+ )
+ return re.sub(r"\bBigR_x\b", "1", result, flags=re.I)
+
+
+def _parsed_terms(text, *, temperature_model=None):
+ """Return expanded monomials, or the original text when unsupported."""
+
+ variants = []
+ variants_to_try = {
+ "single": ((True, False),),
+ "two": ((False, True),),
+ None: ((False, False), (True, False), (False, True)),
+ }[temperature_model]
+ for single, two in variants_to_try:
+ try:
+ parsed = parse_fortran_expression(
+ _normalized_text(text, single_temperature=single, two_temperature=two)
+ )
+ # ``normalize_fortran_text`` expands several element-routine
+ # aliases (notably ``r0_x_hat``) after the textual normalization
+ # above. Those expansions can re-introduce ``BigR_x``. The
+ # report follows the comparison convention BigR_x=1, so apply it
+ # to the parsed tree as well as to the input text.
+ parsed = parsed.subs(sp.Symbol("BigR_x"), sp.Integer(1))
+ except Exception:
+ continue
+ terms = tuple(sp.Add.make_args(sp.expand(parsed)))
+ if terms not in variants:
+ variants.append(terms)
+ return variants
+
+
+def _channel(term):
+ test_order = 0
+ trial_order = 0
+ for derivative in term.atoms(SpatialDerivative):
+ if derivative.coordinate != phi:
+ continue
+ if derivative.expression.has(TestFunction):
+ test_order += 1
+ if any(
+ value.role is FieldRole.TRIAL
+ for value in derivative.expression.atoms(FieldValue)
+ ):
+ trial_order += 1
+ if test_order == 0 and trial_order == 0:
+ return ""
+ if test_order == 0 and trial_order == 1:
+ return "_n"
+ if test_order == 1 and trial_order == 0:
+ return "_k"
+ if test_order == 1 and trial_order == 1:
+ return "_kn"
+ if test_order == 0 and trial_order == 2:
+ return "_nn"
+ raise ValueError("Unsupported toroidal channel in {}".format(term))
+
+
+def _add_linearized(
+ pools, row, linearized: LinearizedEquation, *, previous_names=None,
+ include_rhs=True, include_amat=True, include_fields=None,
+):
+ if include_rhs:
+ for term in sp.Add.make_args(sp.expand(expand_derivatives(linearized.rhs))):
+ if term == 0:
+ continue
+ channel = _channel(term)
+ lhs = "rhs_ij{}(var_{})".format(channel, row)
+ pools[_clean_lhs(lhs)].append(fortran(term, previous_names=previous_names))
+ if include_amat:
+ expressions = linearized.amat.items()
+ else:
+ expressions = ()
+ if include_fields is not None:
+ include_fields = set(include_fields)
+ expressions = (
+ (field, expression)
+ for field, expression in expressions
+ if field in include_fields
+ )
+ for field, expression in expressions:
+ column = FIELD_NAMES[field]
+ for term in sp.Add.make_args(sp.expand(expand_derivatives(expression))):
+ if term == 0:
+ continue
+ channel = _channel(term)
+ lhs = "amat{}(var_{},var_{})".format(channel, row, column)
+ pools[_clean_lhs(lhs)].append(fortran(term, previous_names=previous_names))
+
+
+def _raw_evolution_rhs(equation, fields, timestep, zeta):
+ """Build RHS without SymPy's global additive expansion."""
+
+ history = sp.Add(
+ *(variation(equation.A, value, direction=FieldRole.PREVIOUS_DELTA) for value in fields)
+ )
+ terms = [
+ timestep * resolve_current(term)
+ for term in sp.Add.make_args(equation.B)
+ ]
+ terms.extend(zeta * term for term in sp.Add.make_args(history))
+ return terms
+
+
+def _add_raw_rhs(pools, row, expression, *, previous_names=None):
+ for term in expression:
+ if term == 0:
+ continue
+ channel = _channel(term)
+ lhs = "rhs_ij{}(var_{})".format(channel, row)
+ pools[_clean_lhs(lhs)].append(fortran(term, previous_names=previous_names))
+
+
+def _generated_pools(
+ rows=ROWS, *, include_neo=True, with_tite=None,
+ st_form=True,
+):
+ """Generate implemented model-600 terms for the requested rows."""
+
+ rows = set(rows)
+ pools = defaultdict(list)
+ timestep = coefficient("tstep")
+ theta = coefficient("theta")
+ zeta = coefficient("zeta")
+ if "psi" in rows:
+ # Both the residual and every tangent column follow the temperature
+ # model of the section being exported. The two branches differ in
+ # more than naming: the one-temperature electron pressure is built
+ # from Te0 = T0/2.
+ psi_with_tite = bool(with_tite)
+ equation = induction_equation_1(with_TiTe=psi_with_tite)
+ result = equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ )
+ _add_linearized(
+ pools, "psi", result,
+ previous_names={"psi": "delta_g(mp,var_psi,ms,mt)"},
+ )
+ if "rho" in rows:
+ equation = density_equation_rho(with_TiTe=bool(with_tite))
+ _add_linearized(
+ pools, "rho",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={"rho": "delta_rho_g"},
+ )
+ if "vpar" in rows:
+ equation = parallel_velocity_equation_vpar(
+ with_TiTe=bool(with_tite),
+ st_form=st_form,
+ )
+ _add_linearized(
+ pools, "vpar",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={
+ "vpar": "delta_g(mp,var_vpar,ms,mt)",
+ "rho": "delta_rho_g",
+ "psi": "delta_ps",
+ },
+ )
+ if "rhoimp" in rows:
+ equation = impurity_density_equation_rhoimp(with_TiTe=bool(with_tite))
+ _add_linearized(
+ pools, "rhoimp",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={"rhoimp": "delta_g(mp,var_rhoimp,ms,mt)"},
+ )
+ if "rhon" in rows:
+ equation = neutral_density_equation_rhon(with_TiTe=bool(with_tite))
+ _add_linearized(
+ pools, "rhon",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={"rhon": "delta_g(mp,var_rhon,ms,mt)"},
+ )
+ # The ion and electron energy equations exist only in the two-temperature
+ # branch of the element routine.
+ if "ti" in rows and with_tite:
+ equation = ion_energy_equation_Ti(
+ st_form=st_form,
+ )
+ _add_linearized(
+ pools, "ti",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={
+ "Ti": "delta_g(mp,var_Ti,ms,mt)",
+ "rho": "delta_rho_g",
+ "rhoimp": "delta_g(mp,var_rhoimp,ms,mt)",
+ },
+ )
+ if "te" in rows and with_tite:
+ equation = electron_energy_equation_Te(
+ st_form=st_form,
+ )
+ _add_linearized(
+ pools, "te",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={
+ "Te": "delta_g(mp,var_Te,ms,mt)",
+ "rho": "delta_rho_g",
+ "rhoimp": "delta_g(mp,var_rhoimp,ms,mt)",
+ },
+ )
+ # The single-temperature energy equation exists only in the other branch.
+ if "t" in rows and not with_tite:
+ equation = total_energy_equation_T(st_form=st_form)
+ _add_linearized(
+ pools, "t",
+ equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ ),
+ previous_names={
+ "T": "delta_g(mp,var_T,ms,mt)",
+ "rho": "delta_rho_g",
+ "rhoimp": "delta_g(mp,var_rhoimp,ms,mt)",
+ },
+ )
+ if "zj" in rows:
+ _add_linearized(
+ pools, "zj", current_constraint_equation_zj().linearize(fields=FIELDS)
+ )
+ if "w" in rows:
+ _add_linearized(
+ pools, "w", vorticity_constraint_equation_w().linearize(fields=FIELDS)
+ )
+ if "u" in rows:
+ temperature_branches = (False, True) if with_tite is None else (with_tite,)
+ for with_tite in temperature_branches:
+ equation = momentum_equation_2(
+ with_TiTe=with_tite,
+ include_neo=include_neo,
+ )
+ # The source routine contains both the one- and two-temperature
+ # RHS branches. Retain both raw expressions so terms involving
+ # ``Ti0``/``Te0`` can be aligned as well as the single-T terms.
+ _add_raw_rhs(
+ pools, "u", _raw_evolution_rhs(equation, FIELDS, timestep, zeta),
+ previous_names={"u": "delta_u", "rho": "delta_rho_g"},
+ )
+ result = equation.linearize(
+ fields=FIELDS, timestep=timestep, theta=theta, zeta=zeta
+ )
+ _add_linearized(
+ pools, "u", result,
+ previous_names={"u": "delta_u", "rho": "delta_rho_g"},
+ include_rhs=False,
+ )
+ # For the momentum report retain repeated terms: the Fortran routine has
+ # separate base/extension assignments whose identical monomials must be
+ # available more than once for source-row alignment. Other rows retain
+ # the historical duplicate suppression.
+ if rows == {"u"}:
+ return {lhs: list(terms) for lhs, terms in pools.items()}
+ return {
+ lhs: list(dict.fromkeys(terms))
+ for lhs, terms in pools.items()
+ }
+
+
+
+
+
+def _expanded_ordered_monomials(text, *, temperature_model=None):
+ """Expand generated text using the same temperature convention as source."""
+
+ variants = _parsed_terms(text, temperature_model=temperature_model)
+ if not variants:
+ return []
+ terms = variants[0]
+ common_symbols = set.intersection(
+ *(set(term.free_symbols) for term in terms)
+ ) if terms else set()
+ return sorted(
+ terms,
+ key=lambda term: _source_term_order(text, term, common_symbols),
+ )
+
+
+def _monomial_structure(term):
+ """Return the coefficient-free part of a monomial."""
+
+ return sp.cancel(term).as_coeff_Mul()[1]
+
+
+def _geometric_structure(term):
+ """Monomial structure with the cylindrical radius factored out.
+
+ The element routine and the generator sometimes disagree by a power of
+ ``BigR`` alone. Keying on the rest of the monomial lets such a pair be
+ shown on one report line instead of drifting apart, which is how a wrong
+ power of the radius becomes visible in Meld.
+ """
+
+ structure = _monomial_structure(term)
+ radius = sp.Symbol("BigR")
+ exponent = sp.degree(sp.together(structure).as_numer_denom()[0], radius)
+ try:
+ return sp.cancel(structure / radius**exponent)
+ except Exception:
+ return structure
+
+
+def _scaling_structure(term):
+ """Monomial structure with every scalar scaling factor removed.
+
+ On top of the cylindrical radius this drops the implicitness factor
+ ``theta``. Both are pure scalings of an otherwise identical term, so a
+ pair that differs only in them belongs on one report line; that is how a
+ tangent written without its ``theta`` becomes visible in Meld.
+ """
+
+ structure = _geometric_structure(term)
+ theta = sp.Symbol("theta")
+ exponent = sp.degree(sp.together(structure).as_numer_denom()[0], theta)
+ try:
+ return sp.cancel(structure / theta**exponent)
+ except Exception:
+ return structure
+
+
+def _combine_source_duplicates(source_monomials, generated_by_key):
+ """Sum source monomials the generated side reports only once.
+
+ A Fortran assignment writes a Jacobian as a sum of separately factored
+ outer terms, and two of those outer terms can expand onto the same
+ monomial. The product rule in the generator, by contrast, always returns
+ fully combined monomials, so the same contribution appears there with the
+ summed coefficient. The example in model 600 is the Taylor-Galerkin
+ ``tgnum_u`` tangent, where ``0.25*(w0_x*u_y - w0_y*u_x)*(v_x*u0_y - ...)``
+ and ``0.25*(w0_x*u0_y - w0_y*u0_x)*(v_x*u_y - ...)`` both contain
+ ``v_x*u0_y*u_y*w0_x`` and therefore combine to ``0.5``.
+
+ Source monomials are combined only when the generated block really has
+ fewer copies of that monomial structure; otherwise the one-slot-per-outer
+ term layout of the report is preserved.
+ """
+
+ generated_structures = defaultdict(int)
+ for terms in generated_by_key.values():
+ for term in terms:
+ generated_structures[_monomial_structure(term)] += 1
+
+ positions = defaultdict(list)
+ for index, (_, source_term, _) in enumerate(source_monomials):
+ if source_term is None:
+ continue
+ positions[_monomial_structure(source_term)].append(index)
+
+ dropped = set()
+ combined = {}
+ for structure, indices in positions.items():
+ if len(indices) <= max(generated_structures.get(structure, 0), 1):
+ continue
+ total = sp.Add(*(source_monomials[index][1] for index in indices))
+ combined[indices[0]] = sp.expand(total)
+ dropped.update(indices[1:])
+
+ result = []
+ for index, entry in enumerate(source_monomials):
+ if index in dropped:
+ continue
+ if index in combined:
+ if combined[index] == 0:
+ continue
+ assignment, _, source_raw = entry
+ entry = (assignment, combined[index], source_raw)
+ result.append(entry)
+ return result
+
+
+def _source_monomial_slots(assignments, generated_terms, *, temperature_model=None):
+ """Expand both sides and align every generated monomial to source order."""
+
+ source_monomials = []
+ # The NEO psi tangent is written in JOREK as four additive variations,
+ # while SymPy combines equal monomials when differentiating the compact
+ # residual. Combine the source block first as well; otherwise a source
+ # coefficient such as ``1`` is compared with the algebraically equivalent
+ # combined generated coefficient ``3`` on a different source line.
+ source_assignments = assignments
+ if (
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_psi)"
+ and any("amu_neo_prof" in item.expression for item in assignments)
+ ):
+ combined = " + ".join(item.expression for item in assignments)
+ source_assignments = (
+ SourceAssignment(assignments[0].lhs, assignments[0].row,
+ assignments[0].line, combined),
+ )
+ combine_neo_psi = (
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_psi)"
+ and any("amu_neo_prof" in item.expression for item in assignments)
+ )
+ for assignment in source_assignments:
+ pieces = ([assignment.expression] if combine_neo_psi
+ else _split_top_level(assignment.expression))
+ for piece in pieces:
+ variants = _parsed_terms(piece, temperature_model=temperature_model)
+ use_two_temperature = (
+ assignment.lhs in {
+ "amat(var_psi,var_psi)",
+ "amat(var_psi,var_rho)",
+ "amat_n(var_psi,var_rho)",
+ }
+ or "var_te" in assignment.lhs
+ )
+ expanded = []
+ if variants:
+ if use_two_temperature:
+ # Explicit Te0 source terms already identify the branch;
+ # abstract Pe0 terms require the two-temperature alias.
+ variant_index = (
+ 0 if "te0" in piece.lower() else len(variants) - 1
+ )
+ else:
+ variant_index = 0
+ expanded = list(variants[variant_index])
+ common_symbols = set.intersection(
+ *(set(term.free_symbols) for term in expanded)
+ ) if expanded else set()
+ expanded.sort(
+ key=lambda term: _source_term_order(
+ piece, term, common_symbols
+ )
+ )
+ if not expanded:
+ # Keep unsupported source expressions (currently NEO terms)
+ # visible in the report instead of silently dropping them.
+ source_monomials.append((assignment, None, piece))
+ continue
+ for term in expanded:
+ if term == 0:
+ continue
+ source_monomials.append((assignment, term, None))
+
+ generated_by_key = defaultdict(list)
+ generated_by_structure = defaultdict(list)
+ generated_by_radius = defaultdict(list)
+ generated_by_scaling = defaultdict(list)
+ # For a source NEO block, keep only generated NEO terms in this local
+ # alignment. The complete generated pool also contains the ordinary
+ # momentum terms, which must not appear as unrelated generated-only rows
+ # in the NEO subsection.
+ neo_block = bool(
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_psi)"
+ and any("amu_neo_prof" in item.expression for item in assignments)
+ )
+ pool_terms = (
+ [text for text in generated_terms if "amu_neo_prof" in text]
+ if neo_block else generated_terms
+ )
+ for text in pool_terms:
+ for term in _expanded_ordered_monomials(
+ text, temperature_model=temperature_model,
+ ):
+ # The residual history in the element routine uses the physical
+ # background density ``r0``; only the velocity tangent uses the
+ # corrected coefficient ``r0_corr``. The DSL shares one A form,
+ # so apply that source convention when matching RHS monomials.
+ if assignments and assignments[0].lhs == "rhs_ij(var_u)":
+ term = term.xreplace({sp.Symbol("r0_corr"): sp.Symbol("r0")})
+ # JOREK stores W_dia_rho as a logarithmic density derivative
+ # (the explicit trial-density factor is absorbed in the work
+ # variable), whereas the DSL variation exposes that factor.
+ if (
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_rho)"
+ and term.has(sp.Symbol("W_dia_rho"))
+ and term.has(sp.Symbol("rho"))
+ ):
+ term = sp.cancel(term / sp.Symbol("rho"))
+ if (
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_ti)"
+ and term.has(sp.Symbol("W_dia_Ti"))
+ and term.has(sp.Symbol("Ti"))
+ ):
+ term = sp.cancel(term / sp.Symbol("Ti"))
+ if (
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_t)"
+ and term.has(sp.Symbol("W_dia_T"))
+ and term.has(sp.Symbol("T"))
+ ):
+ term = sp.cancel(term / sp.Symbol("T"))
+ if (
+ assignments
+ and assignments[0].lhs == "amat(var_u,var_rhoimp)"
+ and term.has(sp.Symbol("alpha_e_bis"))
+ and term.has(sp.Symbol("Te0"))
+ ):
+ # alpha_e_bis is supplied as d(alpha_e*Te)/dTe, so the work
+ # variable already contains the explicit background Te0.
+ term = sp.cancel(term / sp.Symbol("Te0"))
+ generated_by_key[_canonical_monomial(term)].append(term)
+ generated_by_structure[_monomial_structure(term)].append(term)
+ generated_by_radius[_geometric_structure(term)].append(term)
+ generated_by_scaling[_scaling_structure(term)].append(term)
+
+ source_monomials = _combine_source_duplicates(
+ source_monomials, generated_by_key
+ )
+
+ def _take(pool, key):
+ """Pop the first unconsumed term registered under ``key``."""
+
+ while pool[key]:
+ term = pool[key].pop(0)
+ if id(term) in consumed:
+ continue
+ consumed.add(id(term))
+ return term
+ return None
+
+ # Match in three passes rather than term by term. A single pass lets an
+ # early source monomial with no exact partner consume, through one of the
+ # relaxed keys, a generated monomial that a later source monomial matches
+ # exactly; the pair then drifts apart in the report even though the block
+ # agrees. Exact matches are therefore all resolved first, then the ones
+ # that differ only by a coefficient, then the ones that differ by a power
+ # of the cylindrical radius as well.
+ consumed = set()
+ entries = []
+ for assignment, source_term, source_raw in source_monomials:
+ if source_term is None:
+ entries.append([assignment, source_raw, "", None])
+ continue
+ entries.append([
+ assignment, _canonical_monomial(source_term), "", source_term,
+ ])
+ for pool, key_of in (
+ (generated_by_key, _canonical_monomial),
+ (generated_by_structure, _monomial_structure),
+ (generated_by_radius, _geometric_structure),
+ (generated_by_scaling, _scaling_structure),
+ ):
+ for entry in entries:
+ source_term = entry[3]
+ if source_term is None or entry[2]:
+ continue
+ match = _take(pool, key_of(source_term))
+ if match is not None:
+ entry[2] = _canonical_monomial(match)
+ slots = [(assignment, source, generated)
+ for assignment, source, generated, _ in entries]
+
+ # Do not hide generated-only terms. They are emitted after the aligned
+ # source rows with an empty source column, which makes sign or convention
+ # inconsistencies visible in Meld.
+ synthetic = SourceAssignment(
+ assignments[0].lhs, assignments[0].row, 0, ""
+ ) if assignments else None
+ if synthetic is not None:
+ emitted = set()
+ for values in generated_by_key.values():
+ for term in values:
+ if id(term) in consumed:
+ continue
+ rendered = _canonical_monomial(term)
+ if rendered in emitted:
+ continue
+ emitted.add(rendered)
+ slots.append((synthetic, "", rendered))
+ return slots
+
+
+def _unmatched_generated_blocks(rows, source_lhs, generated, temperature_model):
+ """Emit generated assignments the source routine does not write at all.
+
+ A Jacobian column that is missing from the element routine entirely has no
+ source assignment to align against, so it would otherwise disappear from
+ the report instead of showing up as an unimplemented block.
+ """
+
+ slots = []
+ for lhs, terms in generated.items():
+ if lhs in source_lhs or not terms:
+ continue
+ row_match = re.search(r"\(var_(psi|u|zj|w|rho|rhoimp|vpar|ti|te|t)\b", lhs)
+ if row_match is None or row_match.group(1) not in rows:
+ continue
+ synthetic = SourceAssignment(lhs, row_match.group(1), 0, "")
+ emitted = set()
+ for text in terms:
+ for term in _expanded_ordered_monomials(
+ text, temperature_model=temperature_model,
+ ):
+ rendered = _canonical_monomial(term)
+ if rendered in emitted:
+ continue
+ emitted.add(rendered)
+ slots.append((synthetic, "", rendered))
+ return slots
+
+
+def _source_term_order(source_piece, term, common_symbols=()):
+ """Estimate textual order for expanded terms inside one source term."""
+
+ text = source_piece.lower()
+ positions = []
+ common_symbols = set(common_symbols)
+ for symbol in term.free_symbols:
+ if symbol in common_symbols:
+ continue
+ name = str(symbol).lower()
+ # A symbol appearing once is usually the factor that distinguishes this
+ # expanded monomial (zj0/current_source/Jb in the induction example).
+ if len(re.findall(r"\b{}\b".format(re.escape(name)), text)) == 1:
+ positions.append(text.find(name))
+ return min(positions) if positions else len(text)
+
+
+def _factor_order(factor):
+ """Use a stable JOREK-oriented order for multiplicative factors."""
+
+ base = factor
+ while getattr(base, "is_Pow", False):
+ base = base.base
+ name = str(base)
+ prefixes = (
+ "F0", "BigR", "Z", "v", "u", "psi", "zj", "w", "rho",
+ "T", "Ti", "Te", "xjac", "theta", "tstep", "zeta",
+ )
+ prefix_rank = next(
+ (index for index, prefix in enumerate(prefixes) if name.startswith(prefix)),
+ len(prefixes),
+ )
+ return prefix_rank, name, str(factor)
+
+
+def _canonical_monomial(term):
+ """Render one monomial with the same factor order in both reports."""
+
+ # SymPy may leave a common numeric factor in both the numerator and an
+ # expanded NEO denominator (for example ``-2/(2*D)``). Reduce rational
+ # factors before ordering products so algebraically identical source and
+ # generated monomials receive the same key.
+ term = sp.cancel(term)
+ coefficient = sp.S.One
+ factors = []
+ for factor in sp.Mul.make_args(term):
+ if factor.is_Number:
+ coefficient *= factor
+ else:
+ factors.append(factor)
+ factors.sort(key=_factor_order)
+ numerator = [factor for factor in factors if not (factor.is_Pow and factor.exp.is_negative)]
+ denominator = [factor.base ** (-factor.exp) for factor in factors if factor.is_Pow and factor.exp.is_negative]
+ pieces = []
+ if coefficient not in (1, -1) or (not numerator and not denominator):
+ pieces.append(sp.sstr(coefficient))
+ pieces.extend(sp.sstr(factor) for factor in numerator)
+ result = "*".join(pieces) if pieces else "1"
+ for factor in denominator:
+ rendered = sp.sstr(factor)
+ if not factor.is_Atom:
+ rendered = "({})".format(rendered)
+ result += "/{}".format(rendered)
+ if coefficient == -1 and (numerator or denominator):
+ result = "-{}".format(result)
+ return result
+
+
+def _canonical_display(text, *, temperature_model=None):
+ """Normalize multiplication ordering without changing term grouping.
+
+ The temperature convention must be named rather than picked by position:
+ ``_parsed_terms`` drops variants that coincide, so an expression without
+ two-temperature aliases would otherwise be displayed through whichever
+ convention happened to survive.
+ """
+
+ variants = _parsed_terms(text, temperature_model=temperature_model)
+ if not variants:
+ return text
+ terms = variants[0]
+ if len(terms) == 1 and terms[0] == 0:
+ return ""
+ return " + ".join(_canonical_monomial(term) for term in terms)
+
+
+def _slot_mismatches(slots):
+ """Count report lines that are blank on one side."""
+
+ return sum(1 for _, source, generated in slots if not source or not generated)
+
+
+def _best_slots(lhs_assignments, pools, temperature_model):
+ """Align against whichever generated spelling fits the source block.
+
+ The element routine writes a poloidal bracket sometimes as
+ ``a_s*b_t - a_t*b_s`` and sometimes as ``xjac*(a_x*b_y - a_y*b_x)``, and
+ it picks differently for a residual and for one of its own tangents. The
+ two spellings are identical, but they expand into different monomials, so
+ aligning them line by line requires using the spelling the block at hand
+ happens to use.
+ """
+
+ best = None
+ for pool in pools:
+ slots = _source_monomial_slots(
+ lhs_assignments, pool, temperature_model=temperature_model,
+ )
+ score = _slot_mismatches(slots)
+ if best is None or score < best[0]:
+ best = (score, slots)
+ if score == 0:
+ break
+ return best[1]
+
+
+def _source_slots(assignments, generated, *, temperature_model=None,
+ generated_alternative=None):
+ """Align generated monomials to source monomials in source order."""
+
+ alternative = generated_alternative or {}
+
+ slots = []
+ used = defaultdict(set)
+ positional_consumed = defaultdict(int)
+ # For a model-600 momentum report, use the same monomial alignment for
+ # every RHS/AMAT block. Each assignment is matched independently because
+ # the source contains separate columns and toroidal channels. This keeps
+ # the generated line in the same row as its source monomial and leaves a
+ # genuinely empty line when that contribution is not implemented.
+ if assignments and all(assignment.row == "u" for assignment in assignments):
+ def _remove_branch_duplicates(pool):
+ """Drop exact terms repeated only because both T branches run.
+
+ Terms containing an explicit background temperature belong to a
+ branch and must remain separate. Terms without ``T0``, ``Ti0``
+ or ``Te0`` are branch-independent (base, magnetic, and most
+ geometric terms); retaining two identical copies made the report
+ look as if the generator had duplicated physics.
+ """
+ seen = set()
+ result = []
+ for term in pool:
+ lowered = term.lower()
+ branch_specific = any(
+ marker in lowered for marker in ("t0", "ti0", "te0")
+ )
+ if not branch_specific:
+ if term in seen:
+ continue
+ seen.add(term)
+ result.append(term)
+ return result
+
+ by_lhs = defaultdict(list)
+ for assignment in assignments:
+ by_lhs[assignment.lhs].append(assignment)
+ slots = []
+ for lhs, lhs_assignments in by_lhs.items():
+ pool = generated.get(lhs, ())
+ neo_assignments = [
+ item for item in lhs_assignments
+ if "amu_neo_prof" in item.expression
+ ]
+ regular_assignments = [
+ item for item in lhs_assignments
+ if "amu_neo_prof" not in item.expression
+ ]
+ if regular_assignments:
+ slots.extend(
+ _source_monomial_slots(
+ regular_assignments,
+ _remove_branch_duplicates(
+ [term for term in pool if "amu_neo_prof" not in term]
+ ), temperature_model=temperature_model,
+ )
+ )
+ if neo_assignments:
+ slots.extend(
+ _source_monomial_slots(
+ neo_assignments,
+ [term for term in pool if "amu_neo_prof" in term],
+ temperature_model=temperature_model,
+ )
+ )
+ slots.extend(_unmatched_generated_blocks(
+ {"u"}, set(by_lhs), generated, temperature_model,
+ ))
+ return slots
+ # Export the remaining equations (notably PSI) at monomial granularity as
+ # well. The older path below aligned whole Fortran additive pieces, which
+ # made the PSI report look as though monomials were missing even when the
+ # generated expression contained them.
+ if assignments:
+ by_lhs = defaultdict(list)
+ for assignment in assignments:
+ by_lhs[assignment.lhs].append(assignment)
+ slots = [
+ slot
+ for lhs, lhs_assignments in by_lhs.items()
+ for slot in _best_slots(
+ lhs_assignments,
+ [generated.get(lhs, ())]
+ + ([alternative[lhs]] if lhs in alternative else []),
+ temperature_model,
+ )
+ ]
+ slots.extend(_unmatched_generated_blocks(
+ {assignment.row for assignment in assignments},
+ set(by_lhs), generated, temperature_model,
+ ))
+ return slots
+ # ``assignments`` is empty only when the requested row has no source
+ # assignment in this temperature branch; the generator produces nothing
+ # for it either, so there is nothing to align.
+ return []
+
+
+def _render_report(
+ slots, side, rows=ROWS, *, include_neo=False, temperature_model=None,
+):
+ section = {
+ "single": "Single-temperature (T) model",
+ "two": "Two-temperature (Ti/Te) model",
+ }.get(temperature_model)
+ lines = [
+ "## {}".format(section) if section else "# Model 600 equation terms",
+ "",
+ "> Rows currently exported: {}.".format(
+ ", ".join("`{}`".format(row) for row in rows)
+ ),
+ (
+ "> NEO RHS and AMAT terms are included when the corresponding source branch is enabled."
+ if include_neo
+ else "> NEO RHS and AMAT terms are omitted from this comparison."
+ ),
+ "",
+ ]
+ for row in rows:
+ lines.extend(("### var_{}".format(row) if section else "## var_{}".format(row), ""))
+ row_slots = [slot for slot in slots if slot[0].row == row]
+ lhs_order = list(dict.fromkeys(slot[0].lhs for slot in row_slots))
+ for lhs in lhs_order:
+ lines.extend(("#### `{}`".format(lhs) if section else "### `{}`".format(lhs), ""))
+ for assignment, source, generated in row_slots:
+ if assignment.lhs != lhs:
+ continue
+ value = source if side == "source" else generated
+ value = " ".join(value.split())
+ value = _canonical_display(
+ value, temperature_model=temperature_model,
+ )
+ # Keep exactly one physical line per aligned term. An
+ # unavailable term is intentionally an actually empty line,
+ # which makes its position immediately visible in Meld.
+ lines.append(value)
+ lines.append("")
+ return "\n".join(lines) + "\n"
+
+
+def export_model600_markdown(
+ source_path, source_output, generated_output, equations=None,
+ include_neo=False,
+):
+ """Write aligned source/generated reports and return their paths.
+
+ ``equations`` may contain any of ``psi``, ``u``, ``zj`` and ``w``. If it
+ is omitted, all four rows are exported. NEO terms are omitted by default;
+ pass ``include_neo=True`` to include them.
+ """
+
+ # The Fortran spells the temperature rows ``var_Ti``/``var_Te``; the
+ # exporter keys everything on the lower-case assignment name.
+ rows = tuple(row.lower() for row in (equations or ROWS))
+ invalid = set(rows) - set(ROWS)
+ if invalid:
+ raise ValueError("Unknown model-600 equation(s): {}".format(", ".join(sorted(invalid))))
+ all_assignments = _read_source_assignments(source_path)
+ source_sections = []
+ generated_sections = []
+ for temperature_model, with_tite in (("single", False), ("two", True)):
+ source_assignments = [
+ assignment for assignment in all_assignments
+ if assignment.row in rows
+ and assignment.temperature_model in ("both", temperature_model)
+ and (include_neo or "amu_neo_prof" not in assignment.expression)
+ ]
+ generated = _generated_pools(
+ rows, include_neo=include_neo,
+ with_tite=with_tite,
+ )
+ generated_alternative = {}
+ if {"vpar", "ti", "te", "t"} & set(rows):
+ # The parallel-velocity row is the one whose source assignments
+ # disagree among themselves about how to spell a poloidal
+ # bracket, so build the other spelling as well.
+ alternative_rows = tuple(
+ row for row in ("vpar", "ti", "te", "t") if row in rows
+ )
+ alternative_pools = _generated_pools(
+ alternative_rows,
+ include_neo=include_neo, with_tite=with_tite,
+ st_form=False,
+ )
+ generated_alternative = dict(alternative_pools)
+ slots = _source_slots(
+ source_assignments, generated, temperature_model=temperature_model,
+ generated_alternative=generated_alternative,
+ )
+ source_sections.append(
+ _render_report(
+ slots, "source", rows, include_neo=include_neo,
+ temperature_model=temperature_model,
+ )
+ )
+ generated_sections.append(
+ _render_report(
+ slots, "generated", rows, include_neo=include_neo,
+ temperature_model=temperature_model,
+ )
+ )
+ source_output = Path(source_output)
+ generated_output = Path(generated_output)
+ source_output.parent.mkdir(parents=True, exist_ok=True)
+ generated_output.parent.mkdir(parents=True, exist_ok=True)
+ header = "# Model 600 equation terms\n\n"
+ source_output.write_text(header + "\n".join(source_sections), encoding="utf-8")
+ generated_output.write_text(header + "\n".join(generated_sections), encoding="utf-8")
+ return source_output, generated_output
diff --git a/util/equation_codegen/src/jorek_equations/model600_notation.py b/util/equation_codegen/src/jorek_equations/model600_notation.py
new file mode 100644
index 0000000000..de1775a0a8
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/model600_notation.py
@@ -0,0 +1,350 @@
+"""LaTeX notation for the model-600 documentation.
+
+Rendering data only: the equations themselves are in ``model600.py``. This
+module says how each name used there is written in
+``docs/physics/base_fluid_models/RMHD/weak_form.md``; names that are not
+listed are rendered from their spelling (``D_par_local -> D_{par,local}``).
+"""
+
+# The evolved fields, in the order of the variables table.
+FIELDS = {
+ 'psi': '\\psi',
+ 'u': 'u',
+ 'j': 'j',
+ 'omega': '\\omega',
+ 'rho': '\\rho',
+ 'rhoimp': '\\rho_{imp}',
+ 'rhon': '\\rho_n',
+ 'vpar': 'v_\\parallel',
+ 'T': 'T',
+ 'Ti': 'T_i',
+ 'Te': 'T_e',
+}
+
+# Work values, coefficients and local names of the equations.
+SYMBOLS = {
+ 'T_or_Te': '\\check{T}_e',
+ 'Te_gen': '\\hat{T}_e',
+ 'ion_temperature': '\\tilde{T}_i',
+ 'R': 'R',
+ 'xjac': '\\mathcal{J}',
+ 'dV': 'w_V',
+ 'F0': 'F_0',
+ 'tauIC': '\\tau_{IC}',
+ 'tstep': '\\Delta t',
+ 'GAMMA': '\\Gamma',
+ 'gamma': '\\Gamma',
+ 'ne': 'n_e',
+ 'rho_main': '\\rho_{main}',
+ 'rho_corr': '\\rho^{\\mathrm{c}}',
+ 'rho_hat': '\\hat{\\rho}',
+ 'bb2': 'B^2',
+ 'Pe': 'p_e',
+ 'pressure': 'p',
+ 'pi': 'p_i',
+ 'W_dia': 'W_{dia}',
+ 'ion_density': 'n_i',
+ 'ion_pressure': 'p_i',
+ 'electron_pressure': 'p_e',
+ 'electron_density': 'n_e^{\\mathrm{c}}',
+ 'ionization_energy': 'W_{ion}',
+ 'visco_par_eff': '\\mu_\\parallel^{\\mathrm{eff}}',
+ 'source_dens_tot': 'S_n',
+ 'D_par_tot': 'D_\\parallel^{\\mathrm{tot}}',
+ 'D_par_imp_tot': 'D_{\\parallel,imp}^{\\mathrm{tot}}',
+ 'd_par_tot': '\\hat{D}_\\parallel',
+ 'd_par_imp_tot': '\\hat{D}_{\\parallel,imp}',
+ 'alpha_e_value': '\\alpha_e',
+ 'sion_rate': 'S_{ion}',
+ 'srec_rate': 'S_{rec}',
+ 'grad_v': '\\nabla_{\\mathrm{pol}} v',
+ 'grad_u': '\\nabla_{\\mathrm{pol}} u',
+ 'grad_omega': '\\nabla_{\\mathrm{pol}} \\omega',
+ 'grad_vpar': '\\nabla_{\\mathrm{pol}} v_\\parallel',
+ 'lap_v': '\\nabla^2_{\\mathrm{pol}} v',
+ 'lap_omega': '\\nabla^2_{\\mathrm{pol}} \\omega',
+ 'neutral_sources': 'S_{neut}',
+ 'current_source': 'j_{src}',
+ 'Jb': 'j_b',
+ 'visco_par': '\\mu_\\parallel',
+ 'visco_par_num': '\\mu_{\\parallel,num}',
+ 'visco_par_par': '\\mu_{\\parallel\\parallel}',
+ 'visco_par_sc_num': '\\mu_{\\parallel,sc}',
+ 'D_prof': 'D_\\perp',
+ 'D_prof_imp': 'D_{\\perp,imp}',
+ 'D_par_local': 'D_\\parallel',
+ 'D_par_local_imp': 'D_{\\parallel,imp}',
+ 'D_par_sc_num': 'D_{\\parallel,sc}',
+ 'D_par_imp_sc_num': 'D_{\\parallel,imp,sc}',
+ 'tau_sc': '\\tau_{sc}',
+ 'V_prof_pinch': 'V_{pinch}',
+ 'fact_conservative_u': 'f_{cons}',
+ 'delta_n_convection': '\\delta_{n}',
+ 'dV_dpsi_source': '\\frac{\\mathrm{d}V_{rot}}{\\mathrm{d}\\psi}',
+ 'particle_source': 'S_p',
+ 'source_pellet': 'S_{pellet}',
+ 'source_bg_drift': 'S_{bg,drift}',
+ 'source_imp_drift': 'S_{imp,drift}',
+ 'aux_rho0': 'S_\\rho^{aux}',
+ 'aux_mom_par0': 'M_\\parallel^{aux}',
+ 'heat_source_total': 'H',
+ 'heat_source_i': 'H_i',
+ 'heat_source_e': 'H_e',
+ 'aux_E0': 'H^{aux}',
+ 'aux_E0_Ti': 'H_i^{aux}',
+ 'aux_E0_Te': 'H_e^{aux}',
+ 'power_dens_teleport_ju': 'P_{teleport}',
+ 'ksi_ion_norm': '\\xi_{ion}',
+ 'implicit_heat_source': 'c_{\\mathrm{floor}}',
+ 'T_min_neg': 'T_{\\mathrm{min}}',
+ 'Tie_min_neg': 'T_{\\mathrm{min}}',
+ 'aux_jre': 'j_{RE}',
+ 'aux_jre_ind': 'j_{RE}^{ind}',
+ 'aux_P_par_re': 'P_\\parallel^{RE}',
+ 'aux_P_perp_re': 'P_\\perp^{RE}',
+ 'aux_divPIR_perp': '\\Pi_R^{RE}',
+ 'aux_divPIZ_perp': '\\Pi_Z^{RE}',
+ 'visco_fact_old': 'f_\\mu^{old}',
+ 'visco_fact_new': 'f_\\mu^{new}',
+ 'visco_par_heating': '\\mu_{\\parallel,heat}',
+ 'D_perp_num_psin': 'D_{\\perp,num}',
+ 'Dn0x': 'D_{n,R}',
+ 'Dn0y': 'D_{n,Z}',
+ 'Dn0p': 'D_{n,\\phi}',
+ 'source_neutral_drift': 'S_n^{drift}',
+ 'Dn_perp_num': 'D_{\\perp,num}^{n}',
+ 'ZK_perp_num_psin': '\\kappa_{\\perp,num}',
+ 'ZK_i_perp_num_psin': '\\kappa_{\\perp,num}^{i}',
+ 'ZK_e_perp_num_psin': '\\kappa_{\\perp,num}^{e}',
+ 'tgnum_u': 'c_{TG}^{u}',
+ 'tgnum_rho': 'c_{TG}^{\\rho}',
+ 'tgnum_vpar': 'c_{TG}^{v_\\parallel}',
+ 'tgnum_rhoimp': 'c_{TG}^{\\rho_{imp}}',
+ 'tgnum_T': 'c_{TG}^{T}',
+ 'tgnum_Ti': 'c_{TG}^{T_i}',
+ 'tgnum_Te': 'c_{TG}^{T_e}',
+ 'psi_gradient': '\\nabla_{\\mathrm{pol}} \\psi',
+ 'rotation_shear': '\\nabla_{\\mathrm{pol}}(v_\\parallel - V_{rot})',
+}
+
+# State-dependent functions. ``{0}`` makes the symbol a template for the
+# argument; otherwise the arguments follow in parentheses.
+FUNCTIONS = {
+ 'eta': '\\eta',
+ 'eta_num_T': '\\eta_{num}',
+ 'visco': '\\mu_\\perp',
+ 'visco_num': '\\mu_{num}',
+ 'alpha_i_state': '\\alpha_i',
+ 'dvisco_state': "\\mu_\\perp'",
+ 'Sion_rate': 'S_{ion}',
+ 'Srec_rate': 'S_{rec}',
+ 'alpha_e_state': '\\alpha_e',
+ 'alpha_e_bis_state': "\\alpha_e'",
+ 'alpha_e_temperature': 'A_e',
+ 'alpha_imp_bis_state': "\\alpha_{imp}'",
+ 'alpha_imp_temperature': 'A_{imp}',
+ 'corr_neg_dens': '{0}^{\\mathrm{c}}',
+ 'corr_neg_dens_imp': '{0}^{\\mathrm{c}}',
+ 'corr_neg_dens_n': '{0}^{\\mathrm{c}}',
+ 'W_dia_single': 'W_{dia}',
+ 'W_dia_two': 'W_{dia}',
+ 'ZKi_par': '\\kappa_{\\parallel,i}',
+ 'ZKi_perp': '\\kappa_{\\perp,i}',
+ 'visco_heating': '\\mu_{heat}',
+ 'Ti_e_exchange': 'Q_{ie}',
+ 'Ti_floor': 'T_i^{\\mathrm{floor}}',
+ 'Ti_floor_exp': '\\epsilon_i^{\\mathrm{floor}}',
+ 'E_ion_bg_state': 'E_{ion}^{bg}',
+ 'E_ion_state': 'E_{ion}',
+ 'ZKe_par': '\\kappa_{\\parallel,e}',
+ 'ZKe_perp': '\\kappa_{\\perp,e}',
+ 'eta_ohm_e': '\\eta_{ohm}',
+ 'LradDrays': 'L_{rays}',
+ 'LradDcont': 'L_{cont}',
+ 'frad_bg_state': 'f_{rad}^{bg}',
+ 'Lrad_state': 'L_{rad}',
+ 'Te_i_exchange': 'Q_{ei}',
+ 'Te_floor': 'T_e^{\\mathrm{floor}}',
+ 'Te_floor_exp': '\\epsilon_e^{\\mathrm{floor}}',
+ 'ZK_par': '\\kappa_\\parallel',
+ 'ZK_perp': '\\kappa_\\perp',
+ 'T_floor': 'T^{\\mathrm{floor}}',
+ 'T_floor_exp': '\\epsilon^{\\mathrm{floor}}',
+}
+
+# Shared operators (the ``_helper`` functions of model600.py), grouped as on
+# the page. ``template`` is how a call is written, with ``{param}``
+# placeholders (``{param!}`` inserts the argument without parentheses);
+# ``display`` names a parameter in the operator's own definition.
+HELPERS = {
+ '_B2': dict(
+ group='Magnetic field',
+ title='Magnetic field strength',
+ description='The square of the total field, with $\\mathbf{B} = F_0\\,\\nabla\\phi + \\nabla\\psi\\times\\nabla\\phi$.',
+ template='B^2({flux!})',
+ display={'flux': '\\psi'},
+ ),
+ '_B2_pol': dict(
+ group='Magnetic field',
+ title='Poloidal magnetic field strength',
+ description='The square of the poloidal field, $\\lvert\\nabla_{\\mathrm{pol}}\\psi\\times\\nabla\\phi\\rvert^2$. Since $F_0$ is constant, $\\partial_t B^2 = \\partial_t B^2_{pol}$.',
+ template='B^2_{pol}({flux!})',
+ display={'flux': '\\psi'},
+ ),
+ '_B_dot_grad': dict(
+ group='Magnetic field',
+ title='Parallel gradient',
+ description='The derivative along the field. With `st_form` the bracket is written in element coordinates, $[f,\\psi]^{st}$, which is the same quantity.',
+ template='\\mathbf{B}\\cdot\\nabla {value}',
+ display={'value': 'f'},
+ ),
+ '_rho_hat': dict(
+ group='Magnetic field',
+ title='Weighted density',
+ description='The density weighted by $R^2$, as it enters the momentum equation.',
+ template='\\hat{\\rho}',
+ display={'value': '\\rho'},
+ ),
+ '_u_convection': dict(
+ group='Convection',
+ title='Convection by the ExB flow',
+ description='Advection and compression of a density by the $E\\times B$ flow of stream function $u$.',
+ template='\\mathcal{C}_u({density!})',
+ display={'density': 'n'},
+ ),
+ '_parallel_convection': dict(
+ group='Convection',
+ title='Parallel convection',
+ description='Advection and compression of a density by the parallel flow $v_\\parallel$.',
+ template='\\mathcal{C}_\\parallel({quantity!})',
+ display={'quantity': 'n'},
+ ),
+ '_press_u_convection': dict(
+ group='Convection',
+ title='Pressure convection by the ExB flow',
+ description='The same for a pressure: the compression carries the factor $\\Gamma$.',
+ template='\\mathcal{C}_u^{p}({pressure!})',
+ display={'pressure': 'p', 'bracket': '[\\cdot,\\cdot]'},
+ ),
+ '_press_parallel_convection': dict(
+ group='Convection',
+ title='Parallel pressure convection',
+ description='The same for a pressure along the field.',
+ template='\\mathcal{C}_\\parallel^{p}({pressure!})',
+ display={'pressure': 'p'},
+ ),
+ '_par_diff_intg_by_parts': dict(
+ group='Diffusion',
+ title='Parallel diffusion',
+ description='Diffusion along the magnetic field, after integration by parts.',
+ template='\\mathcal{D}_\\parallel({v!}, {density!})',
+ display={'density': 'n'},
+ ),
+ '_diffusion_tot_intg_by_parts': dict(
+ group='Diffusion',
+ title='Total diffusion',
+ description='Isotropic diffusion, including the toroidal derivative, after integration by parts.',
+ template='\\mathcal{D}_{tot}({test!}, {value!})',
+ display={'test': 'v', 'value': 'f'},
+ ),
+ '_dens_tgnum_intg_by_parts': dict(
+ group='Numerical stabilization',
+ title='Taylor-Galerkin stabilization of a density',
+ description='Taylor-Galerkin stabilization of the convection by the $E\\times B$ and parallel flows.',
+ template='\\mathcal{T}_n({v!}, {density!})',
+ display={'density': 'n'},
+ ),
+ '_energy_tgnum': dict(
+ group='Numerical stabilization',
+ title='Taylor-Galerkin stabilization of an energy equation',
+ description='The same for a pressure, with the coefficient $c$ of the equation.',
+ template='\\mathcal{T}_p({v!}, {pressure!}; {factor!})',
+ display={'pressure': 'p', 'factor': 'c'},
+ ),
+ '_heating_floor': dict(
+ group='Numerical stabilization',
+ title='Heating at the temperature floor',
+ description='Implicit heat source that keeps a temperature away from its floor $T_{\\mathrm{min}}$. $\\epsilon$ and $T^{\\mathrm{floor}}$ are the floor functions of the element routine (`Ti_floor_exp`, `Ti0_floor` and their electron and single-temperature counterparts).',
+ template='\\mathcal{H}({v!}, {exponential!}, {floor!}, {minimum!})',
+ display={'exponential': '\\epsilon', 'floor': 'T^{\\mathrm{floor}}', 'minimum': 'T_{\\mathrm{min}}'},
+ ),
+ '_diamagnetic_pressure': dict(
+ group='Diamagnetic terms',
+ title='Diamagnetic ion pressure',
+ description='The ion pressure of the diamagnetic terms, as built in `construct_pressure`. The single-temperature model evolves the total temperature, so the ion temperature is $T/2$ there. $\\alpha_i$ does not depend on the temperature.',
+ template='p_i^{dia}',
+ display={},
+ ),
+ '_diamagnetic_viscosity': dict(
+ group='Diamagnetic terms',
+ title='Diamagnetic viscosity',
+ description="The element routine's `W_dia`. It is built from the ion pressure alone, so it has no explicit electron-temperature dependence.",
+ template='W_{dia}',
+ display={},
+ ),
+ '_released_kinetic_energy': dict(
+ group='Sources and heating',
+ title='Particle sources releasing kinetic energy',
+ description='Particle sources whose kinetic energy is released into the ions: ionization of neutrals and the external sources.',
+ template='S_{kin}({T_or_Te!})',
+ display={'T_or_Te': 'T'},
+ ),
+ '_parallel_viscous_heating': dict(
+ group='Sources and heating',
+ title='Parallel viscous heating',
+ description='Heating by the parallel viscosity.',
+ template='Q_{\\mu\\parallel}({v!})',
+ display={},
+ ),
+ '_u_viscous_heating': dict(
+ group='Sources and heating',
+ title='Perpendicular viscous heating',
+ description='Heating by the perpendicular viscosity of the $u$ flow, with the heating viscosity $\\mu_{heat}$.',
+ template='Q_{\\mu\\perp}({v!}, {heating!})',
+ display={'heating': '\\mu_{heat}'},
+ ),
+ '_kinetic_coupling': dict(
+ group='Sources and heating',
+ title='Coupling to the kinetic neutral and impurity model',
+ description='Energy and momentum handed over by the kinetic neutral and impurity model.',
+ template='Q_{kin}({v!})',
+ display={},
+ ),
+ '_ionization_energy_transport': dict(
+ group='Sources and heating',
+ title='Transport of the ionization potential energy',
+ description='The ionization potential energy $W$ is convected like a density, and diffuses with the impurity and main-ion densities that carry it.',
+ template='\\mathcal{Q}_{ion}({v!}, {energy!}, {temperature!})',
+ display={'energy': 'W', 'temperature': 'T'},
+ ),
+}
+
+# The equations, in the order of the page: function in model600.py, title,
+# anchor suffix, and an optional paragraph shown above the equation's terms.
+EQUATIONS = (
+ ("induction_equation_1", "Induction equation (`var_psi`)", "psi", ""),
+ ("momentum_equation_2", "Perpendicular momentum equation (`var_u`)", "u", ""),
+ ("current_constraint_equation_zj", "Current definition (`var_zj`)", "zj", ""),
+ ("vorticity_constraint_equation_w", "Vorticity definition (`var_w`)", "w", ""),
+ ("density_equation_rho", "Density equation (`var_rho`)", "rho", ""),
+ ("parallel_velocity_equation_vpar", "Parallel velocity equation (`var_vpar`)", "vpar",
+ r"The equation is the projection of the momentum equation on $\mathbf{B}$. "
+ r"With the parallel flow $v_\parallel\mathbf{B}$ and "
+ r"$\mathbf{B}\cdot\partial_t\mathbf{B} = \frac{1}{2}\partial_t B^2 = \frac{1}{2}\partial_t B^2_{pol}$, "
+ r"$\mathbf{B}\cdot\rho\,\partial_t(v_\parallel\mathbf{B}) = \rho\,\big(B^2\,\partial_t v_\parallel "
+ r"+ \frac{1}{2}v_\parallel\,\partial_t B^2_{pol}\big)$: the time derivative acts separately on "
+ r"$v_\parallel$ and on $B^2_{pol}$, and is not the derivative of $\rho\,v_\parallel B^2$. "
+ r"In the conservative form ($f_{cons} = 1$) it adds $v_\parallel B^2\,\partial_t\rho$. "
+ r"The element routine uses the corrected density $\rho^{\mathrm{c}}$ in the first two terms."),
+ ("impurity_density_equation_rhoimp", "Impurity density equation (`var_rhoimp`)", "rhoimp",
+ r"`with_TiTe` has no effect here; it is accepted so that every density equation has the "
+ r"same interface."),
+ ("neutral_density_equation_rhon", "Neutral density equation (`var_rhon`)", "rhon",
+ r"Fluid neutrals are diffused with an anisotropic diffusivity, convected with the plasma flow "
+ r"when $\delta_n = 1$, ionized and recombined with the same rates as in the density "
+ r"equation, and fed by a prescribed source."),
+ ("ion_energy_equation_Ti", "Ion energy equation (`var_Ti`)", "Ti", ""),
+ ("electron_energy_equation_Te", "Electron energy equation (`var_Te`)", "Te", ""),
+ ("total_energy_equation_T", "Single-temperature energy equation (`var_T`)", "T", ""),
+)
+
diff --git a/util/equation_codegen/src/jorek_equations/operators.py b/util/equation_codegen/src/jorek_equations/operators.py
new file mode 100644
index 0000000000..aed62f5d31
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/operators.py
@@ -0,0 +1,210 @@
+"""Structured spatial operators for the equation DSL."""
+
+from typing import Sequence, Tuple
+
+import sympy as sp
+
+from .external import definition_for_call
+
+from .exceptions import InvalidDeclarationError
+from .symbols import Field, FieldValue, Frozen, TestFunction
+
+
+R = sp.Symbol("R")
+Z = sp.Symbol("Z")
+phi = sp.Symbol("phi")
+s = sp.Symbol("s")
+t = sp.Symbol("t")
+
+
+class SpatialDerivative(sp.Expr):
+ """An unevaluated derivative with respect to a physical coordinate."""
+
+ is_commutative = True
+
+ def __new__(cls, expression, coordinate):
+ expression = sp.sympify(expression)
+ coordinate = sp.sympify(coordinate)
+ if coordinate not in (R, Z, phi, s, t):
+ raise InvalidDeclarationError(
+ "Unsupported physical coordinate {!r}".format(coordinate)
+ )
+ return sp.Expr.__new__(cls, expression, coordinate)
+
+ @property
+ def expression(self):
+ return self.args[0]
+
+ @property
+ def coordinate(self):
+ return self.args[1]
+
+ def _sympystr(self, printer) -> str:
+ return "d{}({})".format(
+ printer._print(self.coordinate), printer._print(self.expression)
+ )
+
+ def _latex(self, printer) -> str:
+ return r"\frac{{\partial {}}}{{\partial {}}}".format(
+ printer._print(self.expression), printer._print(self.coordinate)
+ )
+
+
+def derivative(expression, coordinate) -> SpatialDerivative:
+ return SpatialDerivative(expression, coordinate)
+
+
+def dR(expression) -> SpatialDerivative:
+ return derivative(expression, R)
+
+
+def dZ(expression) -> SpatialDerivative:
+ return derivative(expression, Z)
+
+
+def dphi(expression) -> SpatialDerivative:
+ return derivative(expression, phi)
+
+
+def ds(expression) -> SpatialDerivative:
+ return derivative(expression, s)
+
+
+def dt(expression) -> SpatialDerivative:
+ return derivative(expression, t)
+
+
+def grad(expression) -> Tuple[SpatialDerivative, SpatialDerivative]:
+ """Return the two poloidal gradient components."""
+
+ return dR(expression), dZ(expression)
+
+
+def dot(left: Sequence, right: Sequence):
+ if len(left) != len(right):
+ raise InvalidDeclarationError("Dot-product operands have different sizes")
+ return sp.Add(*(sp.sympify(a) * sp.sympify(b) for a, b in zip(left, right)))
+
+
+class PoloidalBracket(sp.Expr):
+ """Structured physical-coordinate bracket [a,b]."""
+
+ is_commutative = True
+
+ def __new__(cls, left, right):
+ return sp.Expr.__new__(cls, sp.sympify(left), sp.sympify(right))
+
+ @property
+ def left(self):
+ return self.args[0]
+
+ @property
+ def right(self):
+ return self.args[1]
+
+ def _sympystr(self, printer) -> str:
+ return "bracket({}, {})".format(
+ printer._print(self.left), printer._print(self.right)
+ )
+
+
+def bracket(left, right):
+ """Return the structured poloidal bracket [a,b]."""
+
+ return PoloidalBracket(left, right)
+
+
+def poiss_bracket_st(left, right):
+ """Return ``left_s right_t - left_t right_s`` in element coordinates."""
+
+ return ds(left) * dt(right) - dt(left) * ds(right)
+
+
+def poiss_bracket(left, right):
+ """Return ``left_R right_Z - left_Z right_R`` in physical coordinates."""
+
+ return dR(left) * dZ(right) - dZ(left) * dR(right)
+
+
+def laplacian(value):
+ """Return the axisymmetric (R,Z) Laplacian of ``value``."""
+
+ return dR(dR(value)) + dZ(dZ(value)) + dR(value) / R
+
+
+def expand_brackets(expression):
+ """Expand physical brackets into R/Z derivatives."""
+
+ expression = sp.sympify(expression)
+ replacements = {
+ item: dR(item.left) * dZ(item.right) - dZ(item.left) * dR(item.right)
+ for item in expression.atoms(PoloidalBracket)
+ }
+ return expression.xreplace(replacements)
+
+
+def expand_derivatives(expression):
+ """Apply the product rule to structured spatial derivatives."""
+
+ expression = sp.sympify(expression)
+
+ def expand(node):
+ if isinstance(node, (FieldValue, Frozen, TestFunction)):
+ return node
+ if isinstance(node, SpatialDerivative):
+ inner = expand(node.expression)
+ coordinate = node.coordinate
+ if inner.is_Number:
+ # A numeric factor (such as the 1/2 in the one-temperature
+ # ion temperature Ti0 = T0/2) is constant in space. Without
+ # this the product rule below emits a spurious derivative of
+ # that number.
+ return sp.S.Zero
+ if isinstance(inner, sp.Add):
+ return sp.Add(*(expand(SpatialDerivative(term, coordinate))
+ for term in inner.args))
+ if isinstance(inner, sp.Mul):
+ terms = []
+ for index, factor in enumerate(inner.args):
+ factors = list(inner.args)
+ factors[index] = expand(SpatialDerivative(factor, coordinate))
+ terms.append(sp.Mul(*factors))
+ return sp.Add(*terms)
+ if isinstance(inner, sp.Pow) and inner.exp.is_integer:
+ return inner.exp * inner.base ** (inner.exp - 1) * expand(
+ SpatialDerivative(inner.base, coordinate)
+ )
+ external = definition_for_call(inner)
+ if external is not None:
+ return sp.Add(*(
+ external.derivative_call(argument_name, inner.args)
+ * expand(SpatialDerivative(argument_value, coordinate))
+ for argument_name, argument_value in zip(
+ external.arguments, inner.args
+ )
+ if argument_name in external.derivatives
+ ))
+ if inner in (R, Z, phi, s, t):
+ return sp.S.One if inner == coordinate else sp.S.Zero
+ return SpatialDerivative(inner, coordinate)
+ if not node.args:
+ return node
+ return node.func(*[expand(argument) for argument in node.args])
+
+ return sp.expand(expand(expression))
+
+
+def lower_brackets_to_element(expression, jacobian):
+ """Rewrite [a,b] as (a_s b_t-a_t b_s)/jacobian."""
+
+ expression = sp.sympify(expression)
+ jacobian = sp.sympify(jacobian)
+ replacements = {
+ item: (
+ ds(item.left) * dt(item.right)
+ - dt(item.left) * ds(item.right)
+ )
+ / jacobian
+ for item in expression.atoms(PoloidalBracket)
+ }
+ return sp.expand(expression.xreplace(replacements))
diff --git a/util/equation_codegen/src/jorek_equations/symbols.py b/util/equation_codegen/src/jorek_equations/symbols.py
new file mode 100644
index 0000000000..a5bd41e9fd
--- /dev/null
+++ b/util/equation_codegen/src/jorek_equations/symbols.py
@@ -0,0 +1,238 @@
+"""Core symbolic objects used by the JOREK equation DSL."""
+
+from dataclasses import dataclass
+from enum import Enum
+import re
+from typing import Optional, Union
+
+import sympy as sp
+
+from .exceptions import InvalidDeclarationError
+
+
+_VALID_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def validate_name(name: str, kind: str) -> None:
+ if not isinstance(name, str) or not _VALID_NAME.match(name):
+ raise InvalidDeclarationError(
+ "Invalid {} name {!r}; expected an identifier".format(kind, name)
+ )
+
+
+class FieldRole(str, Enum):
+ """Role played by a field value during linearization."""
+
+ CURRENT = "current"
+ ABSTRACT = "abstract"
+ DELTA = "delta"
+ PREVIOUS_DELTA = "previous_delta"
+ TRIAL = "trial"
+
+
+class FieldValue(sp.Expr):
+ """One physical field in a specific linearization role."""
+
+ is_commutative = True
+ is_Atom = True
+
+ def __new__(cls, field_name: str, role: Union[str, FieldRole]):
+ validate_name(field_name, "field")
+ role_value = FieldRole(role).value
+ return sp.Expr.__new__(
+ cls, sp.Symbol(field_name), sp.Symbol(role_value)
+ )
+
+ @property
+ def field_name(self) -> str:
+ return str(self.args[0])
+
+ @property
+ def role(self) -> FieldRole:
+ return FieldRole(str(self.args[1]))
+
+ def _sympystr(self, printer) -> str:
+ if self.role is FieldRole.ABSTRACT:
+ return self.field_name
+ if self.role is FieldRole.CURRENT:
+ return "{}0".format(self.field_name)
+ if self.role is FieldRole.DELTA:
+ return "delta_{}".format(self.field_name)
+ if self.role is FieldRole.PREVIOUS_DELTA:
+ return "delta_{}_prev".format(self.field_name)
+ return "{}_trial".format(self.field_name)
+
+ def _latex(self, printer) -> str:
+ name = printer._print(sp.Symbol(self.field_name))
+ if self.role is FieldRole.ABSTRACT:
+ return name
+ if self.role is FieldRole.CURRENT:
+ return "{}_0".format(name)
+ if self.role is FieldRole.DELTA:
+ return r"\delta {}".format(name)
+ if self.role is FieldRole.PREVIOUS_DELTA:
+ return r"\delta {}^{{n-1}}".format(name)
+ return r"\varphi_{{{}}}".format(name)
+
+
+@dataclass(frozen=True)
+class Field:
+ """Declaration of one evolved physical field."""
+
+ name: str
+ fortran_current: Optional[str] = None
+ fortran_trial: Optional[str] = None
+
+ def __post_init__(self) -> None:
+ validate_name(self.name, "field")
+ if self.fortran_current is not None:
+ validate_name(self.fortran_current, "Fortran current-value")
+ if self.fortran_trial is not None:
+ validate_name(self.fortran_trial, "Fortran trial-value")
+
+ @property
+ def current(self) -> FieldValue:
+ return FieldValue(self.name, FieldRole.CURRENT)
+
+ @property
+ def symbol(self) -> FieldValue:
+ """Return the abstract symbol used in equation definitions."""
+
+ return FieldValue(self.name, FieldRole.ABSTRACT)
+
+ @property
+ def increment(self) -> FieldValue:
+ return FieldValue(self.name, FieldRole.DELTA)
+
+ @property
+ def previous_increment(self) -> FieldValue:
+ return FieldValue(self.name, FieldRole.PREVIOUS_DELTA)
+
+ @property
+ def trial(self) -> FieldValue:
+ return FieldValue(self.name, FieldRole.TRIAL)
+
+ def value(self, role: Union[str, FieldRole]) -> FieldValue:
+ return FieldValue(self.name, role)
+
+ def _sympy_(self):
+ """Allow a bare field declaration in SymPy expressions."""
+
+ return self.symbol
+
+ # Delegate arithmetic to the abstract symbolic field, so equations can be
+ # written as ``v * eta(T) * j / R``.
+ def __add__(self, other):
+ return self.symbol + other
+
+ def __radd__(self, other):
+ return other + self.symbol
+
+ def __sub__(self, other):
+ return self.symbol - other
+
+ def __rsub__(self, other):
+ return other - self.symbol
+
+ def __mul__(self, other):
+ return self.symbol * other
+
+ def __rmul__(self, other):
+ return other * self.symbol
+
+ def __truediv__(self, other):
+ return self.symbol / other
+
+ def __rtruediv__(self, other):
+ return other / self.symbol
+
+ def __pow__(self, other):
+ return self.symbol ** other
+
+ def __neg__(self):
+ return -self.symbol
+
+
+class TestFunction(sp.Expr):
+ """Symbolic Galerkin test function."""
+
+ is_commutative = True
+ is_Atom = True
+
+ def __new__(cls, name: str):
+ validate_name(name, "test-function")
+ return sp.Expr.__new__(cls, sp.Symbol(name))
+
+ @property
+ def name(self) -> str:
+ return str(self.args[0])
+
+ def _sympystr(self, printer) -> str:
+ return self.name
+
+ def _latex(self, printer) -> str:
+ return printer._print(sp.Symbol(self.name))
+
+
+class Frozen(sp.Expr):
+ """Expression evaluated at the current state but excluded from variation."""
+
+ is_commutative = True
+
+ def __new__(cls, expression):
+ return sp.Expr.__new__(cls, sp.sympify(expression))
+
+ @property
+ def expression(self):
+ return self.args[0]
+
+ def _sympystr(self, printer) -> str:
+ return "freeze({})".format(printer._print(self.expression))
+
+
+def field(
+ name: str,
+ *,
+ fortran_current: Optional[str] = None,
+ fortran_trial: Optional[str] = None
+) -> Field:
+ """Declare an evolved physical field."""
+
+ return Field(name, fortran_current, fortran_trial)
+
+
+def test_function(name: str = "v") -> TestFunction:
+ """Declare a symbolic weak-form test function."""
+
+ return TestFunction(name)
+
+
+def coefficient(name: str, **assumptions) -> sp.Symbol:
+ """Declare a state-independent scalar coefficient."""
+
+ validate_name(name, "coefficient")
+ return sp.Symbol(name, **assumptions)
+
+
+def delta(value: Field) -> FieldValue:
+ return value.increment
+
+
+def previous_delta(value: Field) -> FieldValue:
+ return value.previous_increment
+
+
+def trial(value: Field) -> FieldValue:
+ return value.trial
+
+
+def freeze(expression) -> Frozen:
+ """Mark a value as frozen during directional linearization."""
+
+ return Frozen(expression)
+
+
+def frozen_abs(expression) -> Frozen:
+ """Return an absolute value evaluated at the current state only."""
+
+ return freeze(sp.Abs(sp.sympify(expression)))
diff --git a/util/equation_codegen/tests/test_docs.py b/util/equation_codegen/tests/test_docs.py
new file mode 100644
index 0000000000..b252f2a8a8
--- /dev/null
+++ b/util/equation_codegen/tests/test_docs.py
@@ -0,0 +1,56 @@
+"""The model-600 documentation generated from model600.py."""
+
+import ast
+import sys
+import unittest
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+sys.path.insert(0, str(PROJECT_ROOT / "examples"))
+
+from jorek_equations.latex_render import LatexPrinter # noqa: E402
+from jorek_equations.model600_docs import read_source, render_page # noqa: E402
+import model600_docs as tool # noqa: E402
+
+
+def latex(source, **options):
+ return LatexPrinter({}, **options).latex(ast.parse(source, mode="eval").body)
+
+
+class PrinterTest(unittest.TestCase):
+ def test_quotient_keeps_its_factors_at_full_size(self):
+ self.assertEqual(latex("2 * a / b * c / d"), r"\frac{2}{b\,d}\,a\,c")
+
+ def test_element_bracket_absorbs_the_jacobian(self):
+ self.assertEqual(latex("R * poiss_bracket_st(a, u) / xjac"), r"R\,[a,u]^{st}")
+ self.assertEqual(latex("v * poiss_bracket_st(a, u) / dV"), r"\frac{1}{R}\,v\,[a,u]^{st}")
+
+ def test_gradient_component_through_a_local(self):
+ values = {"g": ast.parse("grad(psi)", mode="eval").body}
+ self.assertEqual(latex("g[1]", values=values), r"\partial_{Z} \psi")
+
+ def test_partial_alias_is_rendered_as_its_helper(self):
+ self.assertEqual(latex("f(a)", aliases={"f": "dR"}), r"\partial_{R} a")
+
+ def test_zip_comprehension_is_a_vector(self):
+ self.assertEqual(
+ latex("tuple(a - k * b for a, b in zip(grad(x), grad(y)))"),
+ r"\nabla_{\mathrm{pol}} x - k\,\nabla_{\mathrm{pol}} y")
+
+
+class Model600PageTest(unittest.TestCase):
+ """weak_form.md is generated from model600.py and must follow it."""
+
+ def test_every_equation_and_helper_is_documented(self):
+ helpers, _, equations = read_source(tool.SOURCE.read_text(encoding="utf-8"))
+ self.assertIn("neutral_density_equation_rhon", {e.name for e in equations})
+ self.assertIn("_B_dot_grad", helpers)
+
+ def test_page_is_up_to_date(self):
+ text, new = tool.rendered()
+ self.assertEqual(new, text, "run: python3 examples/model600_docs.py render")
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/util/equation_codegen/tests/test_dsl.py b/util/equation_codegen/tests/test_dsl.py
new file mode 100644
index 0000000000..8ee71e9e08
--- /dev/null
+++ b/util/equation_codegen/tests/test_dsl.py
@@ -0,0 +1,260 @@
+"""The symbolic DSL: fields and roles, operators, externals, sign conventions."""
+
+import sys
+import unittest
+from pathlib import Path
+
+import sympy as sp
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+
+from jorek_equations import ( # noqa: E402
+ ConstraintEquation,
+ EvolutionEquation,
+ FieldRole,
+ MissingDerivativeError,
+ bracket,
+ coefficient,
+ dR,
+ dZ,
+ dphi,
+ poiss_bracket_st,
+ expand_brackets,
+ expand_derivatives,
+ external_function,
+ field,
+ fortran,
+ freeze,
+ grad,
+ dot,
+ test_function,
+ variation,
+)
+from jorek_equations.model199 import FIELDS, eta, j, psi, u
+
+
+
+class SymbolsAndOperatorsTest(unittest.TestCase):
+ def test_model199_field_order_and_fortran_names(self):
+ self.assertEqual(
+ [value.name for value in FIELDS],
+ ["psi", "u", "j", "omega", "rho", "T"],
+ )
+ self.assertEqual(FIELDS[2].fortran_current, "zj0")
+ self.assertEqual(FIELDS[3].fortran_trial, "w")
+
+ def test_roles_are_distinct(self):
+ psi = field("psi_role_test")
+ self.assertEqual(psi.current.role, FieldRole.CURRENT)
+ self.assertEqual(psi.increment.role, FieldRole.DELTA)
+ self.assertEqual(psi.previous_increment.role, FieldRole.PREVIOUS_DELTA)
+ self.assertEqual(psi.trial.role, FieldRole.TRIAL)
+ self.assertNotEqual(psi.current, psi.trial)
+
+ def test_bare_field_is_abstract_until_linearization(self):
+ psi = field("psi_abstract_test")
+ u = field("u_abstract_test")
+ expression = psi * u
+
+ self.assertEqual(variation(expression, psi), psi.increment * u.current)
+ self.assertEqual(variation(expression, u), psi.current * u.increment)
+
+ def test_derivative_variation_commutes(self):
+ psi = field("psi_derivative_test")
+ actual = variation(dR(psi.current), psi)
+ self.assertEqual(actual, dR(psi.increment))
+
+ def test_bracket_product_rule(self):
+ psi = field("psi_bracket_test")
+ u = field("u_bracket_test")
+ expression = bracket(psi.current, u.current)
+
+ self.assertEqual(
+ variation(expression, psi, direction=FieldRole.TRIAL),
+ bracket(psi.trial, u.current),
+ )
+ self.assertEqual(
+ variation(expression, u, direction=FieldRole.TRIAL),
+ bracket(psi.current, u.trial),
+ )
+
+ def test_gradient_dot_product_rule(self):
+ u = field("u_gradient_test")
+ expression = sum(component**2 for component in grad(u.current))
+ expected = 2 * (
+ dR(u.current) * dR(u.trial)
+ + dZ(u.current) * dZ(u.trial)
+ )
+ actual = variation(expression, u, direction=FieldRole.TRIAL)
+ self.assertEqual(sp.expand(actual - expected), 0)
+
+ def test_frozen_expression_has_zero_variation(self):
+ rho = field("rho_frozen_test")
+ u = field("u_frozen_test")
+ expression = freeze(rho.current) * dR(u.current)
+
+ self.assertEqual(variation(expression, rho), 0)
+ self.assertEqual(
+ variation(expression, u),
+ freeze(rho.current) * dR(u.increment),
+ )
+
+ def test_test_function_and_coefficients_do_not_vary(self):
+ psi = field("psi_test_function_test")
+ v = test_function("v_test_function_test")
+ Rloc = coefficient("R_test_function_test")
+ expression = v * psi.current / Rloc
+
+ expected = v * psi.trial / Rloc
+ actual = variation(expression, psi, direction=FieldRole.TRIAL)
+ self.assertEqual(actual, expected)
+
+
+
+class ExternalFunctionsTest(unittest.TestCase):
+ def test_active_function_uses_supplied_derivative(self):
+ T = field("T_external_active_test")
+ j = field("j_external_active_test")
+ eta = external_function(
+ "eta_external_active_test",
+ arguments=("T",),
+ derivatives={"T": "deta_dT_external_active_test"},
+ )
+ expression = eta(T.current) * j.current
+
+ expected_T = (
+ sp.Function("deta_dT_external_active_test")(T.current)
+ * T.trial
+ * j.current
+ )
+ expected_j = eta(T.current) * j.trial
+
+ self.assertEqual(
+ variation(expression, T, direction=FieldRole.TRIAL), expected_T
+ )
+ self.assertEqual(
+ variation(expression, j, direction=FieldRole.TRIAL), expected_j
+ )
+
+ def test_piecewise_active_uses_branch_derivative_interface(self):
+ T = field("T_external_piecewise_test")
+ eta = external_function(
+ "eta_external_piecewise_test",
+ arguments=("T",),
+ derivatives={"T": "deta_dT_external_piecewise_test"},
+ policy="piecewise_active",
+ )
+
+ expected = (
+ sp.Function("deta_dT_external_piecewise_test")(T.current)
+ * T.increment
+ )
+ self.assertEqual(variation(eta(T.current), T), expected)
+
+ def test_frozen_external_function_does_not_propagate_dependency(self):
+ psi = field("psi_external_frozen_test")
+ Dperp = external_function(
+ "Dperp_external_frozen_test",
+ arguments=("psi_norm",),
+ derivatives={},
+ policy="frozen",
+ )
+
+ self.assertEqual(variation(Dperp(psi.current), psi), 0)
+
+ def test_missing_active_derivative_is_an_error(self):
+ psi = field("psi_external_missing_test")
+ profile = external_function(
+ "profile_external_missing_test",
+ arguments=("psi",),
+ derivatives={},
+ policy="active",
+ )
+
+ with self.assertRaises(MissingDerivativeError):
+ variation(profile(psi.current), psi)
+
+
+
+class EquationConventionTest(unittest.TestCase):
+ def test_evolution_rhs_and_amat_signs(self):
+ psi = field("psi_evolution_test")
+ j = field("j_evolution_test")
+ T = field("T_evolution_test")
+ v = test_function("v_evolution_test")
+ Rloc = coefficient("R_evolution_test")
+ timestep = coefficient("tstep_evolution_test")
+ theta = coefficient("theta_evolution_test")
+ zeta = coefficient("zeta_evolution_test")
+ eta = external_function(
+ "eta_evolution_test",
+ arguments=("T",),
+ derivatives={"T": "deta_dT_evolution_test"},
+ )
+
+ A = v * psi.current / Rloc
+ B = v * eta(T.current) * j.current / Rloc
+ result = EvolutionEquation("induction_test", v, A, B).linearize(
+ fields=(psi, j, T),
+ timestep=timestep,
+ theta=theta,
+ zeta=zeta,
+ )
+
+ expected_rhs = (
+ timestep * B
+ + zeta * v * psi.previous_increment / Rloc
+ )
+ expected_psi = (1 + zeta) * v * psi.trial / Rloc
+ expected_j = -theta * timestep * v * eta(T.current) * j.trial / Rloc
+ expected_T = (
+ -theta
+ * timestep
+ * v
+ * sp.Function("deta_dT_evolution_test")(T.current)
+ * T.trial
+ * j.current
+ / Rloc
+ )
+
+ self.assertEqual(sp.expand(result.rhs - expected_rhs), 0)
+ self.assertEqual(sp.expand(result.amat[psi] - expected_psi), 0)
+ self.assertEqual(sp.expand(result.amat[j] - expected_j), 0)
+ self.assertEqual(sp.expand(result.amat[T] - expected_T), 0)
+
+ def test_constraint_uses_newton_residual_sign(self):
+ psi = field("psi_constraint_test")
+ j = field("j_constraint_test")
+ v = test_function("v_constraint_test")
+ Rloc = coefficient("R_constraint_test")
+ C = (dR(v) * dR(psi.current) + v * j.current) / Rloc
+
+ result = ConstraintEquation("current_definition_test", v, C).linearize(
+ fields=(psi, j)
+ )
+
+ self.assertEqual(result.rhs, -C)
+ self.assertEqual(result.kind, "static")
+ self.assertEqual(result.amat[psi], dR(v) * dR(psi.trial) / Rloc)
+ self.assertEqual(result.amat[j], v * j.trial / Rloc)
+
+
+
+class FortranPrinterTest(unittest.TestCase):
+ def test_model199_field_and_derivative_names(self):
+ self.assertEqual(fortran(dR(j.current)), "zj0_x")
+ self.assertEqual(fortran(dphi(u.trial)), "u_p")
+ self.assertEqual(fortran(eta(psi.current)), "eta_T")
+
+ def test_previous_increment_can_use_assembly_array_name(self):
+ self.assertEqual(
+ fortran(
+ psi.previous_increment,
+ previous_names={"psi": "delta_g(mp,1,ms,mt)"},
+ ),
+ "delta_g(mp,1,ms,mt)",
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/util/equation_codegen/tests/test_final_test.py b/util/equation_codegen/tests/test_final_test.py
new file mode 100644
index 0000000000..bc56cbd623
--- /dev/null
+++ b/util/equation_codegen/tests/test_final_test.py
@@ -0,0 +1,94 @@
+"""Checks for the discrepancy benchmark's comparison logic."""
+
+import sys
+import unittest
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+sys.path.insert(0, str(PROJECT_ROOT / "examples"))
+
+from final_test import REFERENCE, block_name, compare # noqa: E402
+
+
+def reference(blocks):
+ return {"blocks": blocks}
+
+
+class CompareTest(unittest.TestCase):
+ def test_identical_discrepancies_pass(self):
+ block = {"findings": [7], "source_only": ["a*b"], "generated_only": ["c*d"]}
+ self.assertEqual(
+ compare(reference({"x": block}), {"x": dict(block)}), [],
+ )
+
+ def test_a_discrepancy_that_disappeared_is_reported(self):
+ stored = {"findings": [7], "source_only": ["a*b"], "generated_only": []}
+ current = {"source_only": [], "generated_only": []}
+ problems = compare(reference({"x": stored}), {"x": current})
+ self.assertTrue(any("- a*b" in line for line in problems))
+
+ def test_a_new_discrepancy_is_reported(self):
+ stored = {"findings": [7], "source_only": [], "generated_only": []}
+ current = {"source_only": ["new*term"], "generated_only": []}
+ problems = compare(reference({"x": stored}), {"x": current})
+ self.assertTrue(any("+ new*term" in line for line in problems))
+
+ def test_a_block_that_now_agrees_is_reported(self):
+ stored = {"findings": [7], "source_only": ["a*b"], "generated_only": []}
+ problems = compare(reference({"x": stored}), {})
+ self.assertTrue(any("BLOCK NOW AGREES" in line for line in problems))
+
+ def test_a_block_that_started_to_disagree_is_reported(self):
+ current = {"source_only": ["a*b"], "generated_only": []}
+ problems = compare(reference({}), {"x": current})
+ self.assertTrue(any("NEW BLOCK" in line for line in problems))
+
+ def test_repeated_monomials_are_compared_with_multiplicity(self):
+ stored = {"findings": [], "source_only": ["a*b", "a*b"], "generated_only": []}
+ current = {"source_only": ["a*b"], "generated_only": []}
+ self.assertTrue(compare(reference({"x": stored}), {"x": current}))
+
+ def test_block_name_is_stable(self):
+ self.assertEqual(
+ block_name(("Two-temperature (Ti/Te) model", "amat(var_u,var_u)", 0)),
+ "Two-temperature (Ti/Te) model | amat(var_u,var_u) | #0",
+ )
+
+
+class ReferenceTest(unittest.TestCase):
+ """The stored reference must stay well formed and fully annotated."""
+
+ def test_reference_is_complete(self):
+ import json
+
+ document = json.loads(REFERENCE.read_text(encoding="utf-8"))
+ blocks = document["blocks"]
+ self.assertEqual(document["block_count"], len(blocks))
+ self.assertEqual(
+ document["source_only_lines"],
+ sum(len(item["source_only"]) for item in blocks.values()),
+ )
+ self.assertEqual(
+ document["generated_only_lines"],
+ sum(len(item["generated_only"]) for item in blocks.values()),
+ )
+ for name, item in blocks.items():
+ with self.subTest(block=name):
+ self.assertTrue(
+ item["source_only"] or item["generated_only"],
+ "a recorded block must actually disagree",
+ )
+ self.assertTrue(
+ item["findings"],
+ "every recorded block must name the finding it belongs to",
+ )
+ self.assertEqual(item["source_only"], sorted(item["source_only"]))
+ self.assertEqual(
+ item["generated_only"], sorted(item["generated_only"]),
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/util/equation_codegen/tests/test_model199.py b/util/equation_codegen/tests/test_model199.py
new file mode 100644
index 0000000000..ac47451a96
--- /dev/null
+++ b/util/equation_codegen/tests/test_model199.py
@@ -0,0 +1,187 @@
+"""Model 199: read the element routine, generate the weak form, compare."""
+
+import unittest
+from pathlib import Path
+
+import sympy as sp
+
+from jorek_equations import (
+ coefficient,
+ dR,
+ dZ,
+ dphi,
+ expand_brackets,
+ split_toroidal_channels,
+ variation,
+)
+from jorek_equations.model199 import (
+ FIELDS,
+ T,
+ F0,
+ eta,
+ induction_equation_1,
+ j,
+ omega,
+ psi,
+ rho,
+ u,
+)
+from jorek_equations.fortran_source import (
+ EQUATION_1_ASSIGNMENTS,
+ FortranComparisonError,
+ compare_assignment_maps,
+ compare_model199_equation1,
+ extract_fortran_assignments,
+)
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+REPOSITORY_ROOT = PROJECT_ROOT.parents[1]
+FORTRAN_FILE = REPOSITORY_ROOT / "models/model199/mod_elt_matrix_fft.f90"
+
+
+class SourceComparisonTest(unittest.TestCase):
+ def test_extracts_all_equation_1_assignments(self):
+ assignments = extract_fortran_assignments(
+ FORTRAN_FILE, EQUATION_1_ASSIGNMENTS
+ )
+ self.assertEqual(set(assignments), set(EQUATION_1_ASSIGNMENTS))
+ self.assertIn("eta_T", assignments["rhs_ij_1"])
+ self.assertIn("deta_dT", assignments["amat_16"])
+
+ def test_real_model199_equation_1_matches(self):
+ compare_model199_equation1(FORTRAN_FILE)
+
+ def test_toroidal_trial_term_is_in_n_channel(self):
+ result = induction_equation_1().linearize(
+ fields=FIELDS,
+ timestep=coefficient("tstep"),
+ theta=coefficient("theta"),
+ zeta=coefficient("zeta"),
+ )
+ channels = split_toroidal_channels(result.amat[u])
+ self.assertNotEqual(channels.p, 0)
+ self.assertNotEqual(channels.n, 0)
+ self.assertEqual(channels.k, 0)
+ self.assertEqual(channels.kn, 0)
+
+ def test_mismatch_raises_and_reports_difference(self):
+ with self.assertRaises(FortranComparisonError) as raised:
+ compare_assignment_maps(
+ {"amat_test": "x + y"},
+ {"amat_test": "x - y"},
+ )
+ message = str(raised.exception)
+ self.assertIn("amat_test", message)
+ self.assertIn("source-only:", message)
+ self.assertIn("generated-only:", message)
+ self.assertIn(" + y", message)
+ self.assertIn(" - y", message)
+
+
+if __name__ == "__main__":
+ unittest.main()
+
+class Model199Equation1Test(unittest.TestCase):
+ def setUp(self):
+ self.timestep = coefficient("tstep_model199_eq1_test")
+ self.theta = coefficient("theta_model199_eq1_test")
+ self.zeta = coefficient("zeta_model199_eq1_test")
+ self.xjac = coefficient("xjac")
+ self.R = sp.Symbol("R")
+
+ def test_rhs_matches_fortran_equation_1_terms(self):
+ equation = induction_equation_1()
+ result = equation.linearize(
+ fields=(psi, u, j, T),
+ timestep=self.timestep,
+ theta=self.theta,
+ zeta=self.zeta,
+ )
+ v = equation.test
+ current_source = sp.Symbol("current_source")
+ eta_num = sp.Symbol("eta_num")
+ eps_cyl = sp.Symbol("eps_cyl")
+
+ expected_B = (
+ v * eta(T.current) * (j.current - current_source) / self.R * self.xjac
+ + v * self.xjac * (
+ dR(psi.current) * dZ(u.current)
+ - dZ(psi.current) * dR(u.current)
+ )
+ - v * eps_cyl * F0 / self.R * dphi(u.current) * self.xjac
+ + eta_num * self.xjac * (
+ dR(v) * dR(j.current) + dZ(v) * dZ(j.current)
+ )
+ )
+ expected = self.timestep * expected_B + self.zeta * (
+ v * psi.previous_increment / self.R * self.xjac
+ )
+ self.assertEqual(sp.expand(expand_brackets(result.rhs) - expected), 0)
+
+ def test_amat_blocks_match_fortran_equation_1_terms(self):
+ equation = induction_equation_1()
+ result = equation.linearize(
+ fields=(psi, u, j, T),
+ timestep=self.timestep,
+ theta=self.theta,
+ zeta=self.zeta,
+ )
+ v = equation.test
+ eta_num = sp.Symbol("eta_num")
+ eps_cyl = sp.Symbol("eps_cyl")
+ current_source = sp.Symbol("current_source")
+
+ expected_psi = (
+ (1 + self.zeta) * v * psi.trial / self.R * self.xjac
+ - self.theta * self.timestep * v * self.xjac * (
+ dR(psi.trial) * dZ(u.current)
+ - dZ(psi.trial) * dR(u.current)
+ )
+ )
+ expected_u = (
+ -self.theta * self.timestep * v * self.xjac * (
+ dR(psi.current) * dZ(u.trial)
+ - dZ(psi.current) * dR(u.trial)
+ )
+ + self.theta * self.timestep * v * eps_cyl * F0 / self.R
+ * dphi(u.trial) * self.xjac
+ )
+ expected_j = -self.theta * self.timestep * (
+ eta_num * self.xjac * (
+ dR(v) * dR(j.trial) + dZ(v) * dZ(j.trial)
+ )
+ + eta(T.current) * v * j.trial / self.R * self.xjac
+ )
+ expected_T = -self.theta * self.timestep * (
+ sp.Function("deta_dT")(T.current) * T.trial * v
+ * (j.current - current_source) / self.R * self.xjac
+ )
+
+ self.assertEqual(
+ sp.expand(expand_brackets(result.amat[psi]) - expected_psi), 0
+ )
+ self.assertEqual(
+ sp.expand(expand_brackets(result.amat[u]) - expected_u), 0
+ )
+ self.assertEqual(sp.expand(result.amat[j] - expected_j), 0)
+ self.assertEqual(sp.expand(result.amat[T] - expected_T), 0)
+
+ def test_unrelated_model_fields_have_zero_blocks(self):
+ equation = induction_equation_1()
+ result = equation.linearize(
+ fields=FIELDS,
+ timestep=self.timestep,
+ theta=self.theta,
+ zeta=self.zeta,
+ )
+ self.assertEqual(result.amat[omega], 0)
+ self.assertEqual(result.amat[rho], 0)
+
+
+if __name__ == "__main__":
+ unittest.main()
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/util/equation_codegen/tests/test_report.py b/util/equation_codegen/tests/test_report.py
new file mode 100644
index 0000000000..7eff7b5377
--- /dev/null
+++ b/util/equation_codegen/tests/test_report.py
@@ -0,0 +1,115 @@
+"""Report generation: term alignment, coordinate spellings, rendering."""
+
+import sys
+import unittest
+from pathlib import Path
+
+import sympy as sp
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+sys.path.insert(0, str(PROJECT_ROOT / "examples"))
+
+from diff_model600_reports import to_element_basis, to_physical # noqa: E402
+from jorek_equations.model600_markdown import ( # noqa: E402
+ SourceAssignment,
+ _canonical_display,
+ _render_report,
+ _source_monomial_slots,
+ _split_top_level,
+)
+
+
+class CoordinateSpellingTest(unittest.TestCase):
+ """The element routine writes poloidal brackets both ways."""
+
+ def _both_spellings(self):
+ a_s, a_t, b_s, b_t = sp.symbols("a_s a_t b_s b_t")
+ a_x, a_y, b_x, b_y = sp.symbols("a_x a_y b_x b_y")
+ xjac = sp.Symbol("xjac")
+ return a_s * b_t - a_t * b_s, xjac * (a_x * b_y - a_y * b_x)
+
+ def test_element_basis_identifies_the_two_spellings(self):
+ element, physical = self._both_spellings()
+ self.assertEqual(to_element_basis(element - physical), 0)
+
+ def test_display_basis_identifies_the_two_spellings(self):
+ element, physical = self._both_spellings()
+ self.assertEqual(to_physical(element - physical), 0)
+
+ def test_a_genuine_difference_survives(self):
+ element, physical = self._both_spellings()
+ self.assertNotEqual(to_element_basis(element - 2 * physical), 0)
+
+ def test_variable_names_are_not_mangled(self):
+ # ``var_t`` is an assignment name, not a derivative, and ``Sion_T``
+ # ends in an upper-case T.
+ expression = sp.Symbol("var_t") * sp.Symbol("Sion_T")
+ self.assertEqual(to_element_basis(expression), expression)
+ self.assertEqual(to_physical(expression), expression)
+
+
+class CoefficientAlignmentTest(unittest.TestCase):
+ """A coefficient mismatch must stay on one report line."""
+
+ def test_same_structure_different_coefficient_is_paired(self):
+ assignment = SourceAssignment("amat(var_u,var_u)", "u", 1, "2*a*b")
+ slots = _source_monomial_slots(assignment_list(assignment), ["a*b"])
+ self.assertEqual(len(slots), 1)
+ _, source, generated = slots[0]
+ self.assertEqual(source, "2*a*b")
+ self.assertEqual(generated, "a*b")
+
+
+def assignment_list(assignment):
+ return [assignment]
+
+
+
+class MatchingOrderTest(unittest.TestCase):
+ """Exact matches must be resolved before the relaxed ones."""
+
+ def test_exact_match_is_not_stolen_by_a_relaxed_one(self):
+ # ``2*a*b`` has no exact partner and would, in a single pass, consume
+ # ``a*b`` through the coefficient fallback, leaving the source line
+ # that matches ``a*b`` exactly with an empty cell.
+ assignment = SourceAssignment(
+ "amat(var_u,var_u)", "u", 1, "2*a*b + a*b",
+ )
+ slots = _source_monomial_slots([assignment], ["a*b", "3*a*b"])
+ rendered = [(source, generated) for _, source, generated in slots]
+ self.assertEqual(len(rendered), 2)
+ self.assertIn(("a*b", "a*b"), rendered)
+ self.assertIn(("2*a*b", "3*a*b"), rendered)
+ self.assertTrue(all(source and generated for source, generated in rendered))
+
+
+
+class RenderTest(unittest.TestCase):
+ """Report rendering keeps one physical line per aligned term."""
+
+ def test_top_level_split_keeps_parenthesized_sum_together(self):
+ self.assertEqual(
+ _split_top_level("a + b*(c-d) - e"),
+ ["a", "+ b*(c-d)", "- e"],
+ )
+
+ def test_missing_generated_term_is_an_empty_aligned_line(self):
+ assignment = SourceAssignment("rhs_ij(var_psi)", "psi", 1, "a")
+ slots = [(assignment, "a", "")]
+ source = _render_report(slots, "source").splitlines()
+ generated = _render_report(slots, "generated").splitlines()
+ self.assertEqual(len(source), len(generated))
+ self.assertIn("a", source)
+ self.assertEqual(generated[source.index("a")], "")
+
+
+ def test_factor_order_is_shared_between_source_and_generated_forms(self):
+ self.assertEqual(
+ _canonical_display("F0 / BigR * v * u_p * xjac * theta * tstep"),
+ _canonical_display("F0*theta*tstep*xjac*u_p*v/BigR"),
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()