The Hasegawa-Wakatani drift-wave model

KAPPA = 1.0

rung "The drift wave" {
    #| In the limit $C \to \infty$ the electrons follow the potential at once and $n = \phi$.
    #| Subtracting the two equations then leaves one, the Hasegawa-Mima equation
    #| $\partial_t(\nabla^2\phi - \phi) = \kappa\,\partial_y\phi$ (linear part). A wave
    #| $e^{i(k_x x + k_y y - \omega t)}$ solves it when
    #|
    #| $$\omega = \frac{\kappa k_y}{1 + k_x^2 + k_y^2}.$$
    #|
    #| The wave moves only in $y$, along the surfaces of constant density. The 1 in the
    #| denominator comes from $n = \phi$; the $k^2$ comes from the ion polarization drift,
    #| which is stronger for small eddies. So the frequency first rises with $k_y$, peaks at
    #| $k_y = 1$ (for $k_x = 0$), and then falls.

    #| $\omega_{HM} = \kappa k_y / (1 + k_x^2 + k_y^2)$: the frequency of the drift wave when the electrons follow the potential exactly.
    def omega_hm(kx, ky) {
        return KAPPA * ky / (1.0 + kx * kx + ky * ky)
    }
    w1 = omega_hm(0.0, 1.0)
    show w1
    show omega_hm(0.3, 0.7)
    test "the drift wave at k = (0, 1) has frequency 1/2" {
        assert_near(w1, 0.5, 1.0e-13)
    }
    test "at k = (0.3, 0.7) it is 0.7 / 1.58 = 0.44303797468354" {
        # 1.58 is 1 + k^2
        assert_near(1.0 + 0.3 * 0.3 + 0.7 * 0.7, 1.58, 1.0e-15)
        assert_near(omega_hm(0.3, 0.7), 0.4430379746835443, 1.0e-13)
    }
    test "past k_y = 1 the frequency falls: 0.4 at k_y = 2" {
        assert_near(omega_hm(0.0, 2.0), 0.4, 1.0e-13)
    }
}

rung "The growth rate is an eigenvalue" {
    #| At finite $C$ the density lags the potential, and the lag lets the wave draw
    #| energy from the gradient. For one Fourier mode, $\phi = -\zeta/k^2$ and $\partial_y \to ik_y$,
    #| and the linear equations become $\frac{d}{dt}(\zeta, n) = M(\zeta, n)$ with the $2 \times 2$
    #| matrix under the title. Its eigenvalues are $\lambda = \gamma - i\omega$: the real part is
    #| the growth rate, and minus the imaginary part is the frequency.
    #|
    #| The trace of $M$ is real and its determinant, $i C\kappa k_y / k^2$, is imaginary, so the
    #| quadratic formula needs the square root of a complex number. Viv has complex
    #| arithmetic, and `sqrt` of a complex value is the principal root, whose real part is
    #| not negative; that picks the eigenvalue that grows faster.
    #|
    #| At $C = 1$ and $k = (0, 1)$: $\gamma = 0.098684$ and $\omega = 0.45509$, below the
    #| drift-wave 0.5. Scanning $k_y$, the fastest growth at $C = 1$ is $\gamma = 0.10687$ near
    #| $k_y = 1.30$.

    #| The faster-growing eigenvalue $\lambda = \frac12\left(\operatorname{tr} M + \sqrt{(\operatorname{tr} M)^2 - 4\det M}\right)$ of the linearized mode.
    def hw_lambda(kx, ky, c) {
        k2 = kx * kx + ky * ky
        tr = -c / k2 - c
        det = complex(0.0, c * KAPPA * ky / k2)
        return 0.5 * (tr + sqrt(tr * tr - 4.0 * det))
    }
    #| $\gamma = \operatorname{Re}\lambda$: how fast the mode grows.
    def growth(kx, ky, c) {
        return re(hw_lambda(kx, ky, c))
    }
    #| $\omega = -\operatorname{Im}\lambda$: the mode's frequency.
    def freq(kx, ky, c) {
        return -im(hw_lambda(kx, ky, c))
    }
    show growth(0.0, 1.0, 1.0)
    show freq(0.0, 1.0, 1.0)
    test "at C = 1, k = (0, 1): gamma = 0.09868411347 and omega = 0.45508986056 (numpy eigvals)" {
        assert_near(growth(0.0, 1.0, 1.0), 0.09868411346781, 1.0e-12)
        assert_near(freq(0.0, 1.0, 1.0), 0.4550898605622274, 1.0e-12)
    }
    test "at C = 0.5, k = (0.3, 0.7): gamma = 0.10155874337 and omega = 0.38554401937 (numpy eigvals)" {
        assert_near(growth(0.3, 0.7, 0.5), 0.10155874336669901, 1.0e-12)
        assert_near(freq(0.3, 0.7, 0.5), 0.3855440193686499, 1.0e-12)
    }

    # k_y from 0.05 to 3 in steps of 0.01
    N_KY = 296
    kys = linspace(0.05, 3.0, N_KY)
    g_c1 = zeros(N_KY)
    w_c1 = zeros(N_KY)
    w_hm = zeros(N_KY)
    g_best = 0.0
    ky_best = 0.0
    for i to N_KY {
        g_c1[i] = growth(0.0, kys[i], 1.0)
        w_c1[i] = freq(0.0, kys[i], 1.0)
        w_hm[i] = omega_hm(0.0, kys[i])
        if g_c1[i] > g_best {
            g_best = g_c1[i]
            ky_best = kys[i]
        }
    }
    show g_best
    show ky_best
    test "the fastest growth at C = 1 is 0.106870 near k_y = 1.30 (numpy scan with step 0.0005)" {
        assert_near(g_best, 0.10687017183708586, 1.0e-6)
        assert_near(ky_best, 1.3035, 0.01)
    }
    figure("One mode of Hasegawa-Wakatani, $k_x = 0$, $C = 1$", "$k_y$", "rate")
    plot(kys, g_c1, "growth rate $\\gamma$")
    plot(kys, w_c1, "frequency $\\omega$")
    plot(kys, w_hm, "drift wave, $\\kappa k_y/(1 + k^2)$")
    style("drift wave, $\\kappa k_y/(1 + k^2)$", "dash")
}

rung "Stiff coupling gives back the drift wave" {
    #| Rung 1 took $C \to \infty$ by hand. The eigenvalue should agree, and expanding the
    #| small root of $\lambda^2 - \operatorname{tr}M\,\lambda + \det M = 0$ in powers of $1/C$ says how fast.
    #| With $\operatorname{tr}M = -C(1+k^2)/k^2$, the first term is $\det M/\operatorname{tr}M = -i\omega_{HM}$,
    #| the drift wave. The next two terms are
    #|
    #| $$\gamma \approx \frac{\omega_{HM}^2\, k^2}{C(1 + k^2)}, \qquad 1 - \frac{\omega}{\omega_{HM}} \approx \frac{2\,\omega_{HM}^2\, k^4}{C^2 (1 + k^2)^2}.$$
    #|
    #| The growth rate vanishes like $1/C$ and the frequency error like $1/C^2$. At $k = (0, 1)$
    #| both coefficients are $1/8$, and the cell checks $C\gamma$ and $C^2(1 - \omega/\omega_{HM})$
    #| at $C = 1000$. The growth needs the lag: electrons that keep up exactly give a wave that
    #| neither grows nor decays.

    C_STIFF = 1000.0
    k2_ref = 1.0
    c_gamma = C_STIFF * growth(0.0, 1.0, C_STIFF)
    c2_dw = C_STIFF * C_STIFF * (1.0 - freq(0.0, 1.0, C_STIFF) / w1)
    show c_gamma
    show c2_dw
    test "C gamma approaches omega_HM^2 k^2 / (1 + k^2) = 1/8 (0.12499996 at C = 1000)" {
        assert_near(c_gamma, w1 * w1 * k2_ref / (1.0 + k2_ref), 1.0e-6)
    }
    test "C^2 (1 - omega / omega_HM) approaches 2 omega_HM^2 k^4 / (1 + k^2)^2 = 1/8" {
        assert_near(c2_dw, 2.0 * w1 * w1 * k2_ref * k2_ref / ((1.0 + k2_ref) * (1.0 + k2_ref)), 1.0e-6)
    }
}

rung "A zonal flow relaxes to k^2 / (1 + k^2) of its start" {
    #| A zonal flow is a mode with $k_y = 0$: a band of $E \times B$ flow in the $y$ direction,
    #| constant along $y$. For it the drive $\kappa\,\partial_y\phi$ is zero and only the coupling
    #| acts, $\partial_t\zeta = \partial_t n = C(-\zeta/k^2 - n)$. Here $M$ is singular. One eigenvalue is
    #| $-C(1 + 1/k^2)$; the other is 0, because $\zeta - n$ never changes. So a band started with
    #| $\zeta = 1$, $n = 0$ does not decay away: it relaxes to
    #|
    #| $$\frac{\zeta(t)}{\zeta(0)} = \frac{k^2}{1 + k^2} + \frac{e^{-C(1 + 1/k^2)t}}{1 + k^2},$$
    #|
    #| which is $0.2$ at $k_x = 0.5$. Wide bands (small $k$) keep little, narrow ones keep
    #| most. The cell integrates the two equations with fourth-order Runge-Kutta to
    #| $t = 2$, where the exact ratio is $0.2000363$.

    #| Integrate a zonal mode ($k_y = 0$) with coupling $c$ from $\zeta = 1$, $n = 0$ to time $t_{end}$
    #| in `nstep` steps of fourth-order Runge-Kutta, storing $\zeta$ at `len(hist)` evenly spaced times,
    #| the first of them $t = 0$.
    def zonal_history(kz, c, t_end, nstep, hist: Arr) {
        every = nstep / (len(hist) - 1.0)
        h = t_end / nstep
        k2 = kz * kz
        z = 1.0
        n = 0.0
        hist[0] = z
        for i to nstep {
            a1 = c * (-z / k2 - n)
            a2 = c * (-(z + 0.5 * h * a1) / k2 - (n + 0.5 * h * a1))
            a3 = c * (-(z + 0.5 * h * a2) / k2 - (n + 0.5 * h * a2))
            a4 = c * (-(z + h * a3) / k2 - (n + h * a3))
            dz = h / 6.0 * (a1 + 2.0 * a2 + 2.0 * a3 + a4)
            z = z + dz
            n = n + dz
            if mod(i + 1.0, every) == 0.0 {
                hist[(i + 1.0) / every] = z
            }
        }
        return z
    }
    KZ = 0.5
    T_END = 2.0
    N_RK = 4000
    # the figure's samples: the start, then every 40th step
    N_SAMPLE = 101
    zon_classic = zeros(N_SAMPLE)
    zr = zonal_history(KZ, 1.0, T_END, N_RK, zon_classic)
    lam_z = 1.0 + 1.0 / (KZ * KZ)
    zr_exact = KZ * KZ / (1.0 + KZ * KZ) + exp(-lam_z * T_END) / (1.0 + KZ * KZ)
    show zr
    show zr_exact
    test "Runge-Kutta gives the exact k^2/(1+k^2) + e^(-C(1+1/k^2)t)/(1+k^2) at t = 2 to 1e-12" {
        assert_near(zr, zr_exact, 1.0e-12)
    }
    test "and scipy's matrix exponential gives 0.20003631994381 for the same band" {
        assert_near(zr, 0.20003631994381002, 1.0e-12)
    }

    t_rk = linspace(0.0, T_END, N_SAMPLE)
    resid = zeros(N_SAMPLE)
    for i to N_SAMPLE {
        resid[i] = KZ * KZ / (1.0 + KZ * KZ)
    }
    figure("A zonal band, $k_x = 0.5$, $C = 1$", "t", "zeta(t) / zeta(0)", 0.0, 2.0, 0.0, 1.05)
    plot(t_rk, zon_classic, "classical Hasegawa-Wakatani")
    plot(t_rk, resid, "residual $k^2/(1 + k^2)$")
    style("residual $k^2/(1 + k^2)$", "dash")
}

rung "The modified model leaves the band alone" {
    #| In a tokamak the coupling term comes from electrons moving along field lines, and
    #| a zonal flow is constant on a whole flux surface, so there is nothing along the field
    #| line for the electrons to smooth out. The modified Hasegawa-Wakatani model
    #| (Numata, Ball and Dewar, 2007) writes this in: the coupling acts only on the
    #| non-zonal parts, $C(\tilde\phi - \tilde n)$, where the tilde removes the average over $y$.
    #| For a $k_y = 0$ mode that makes the coupling exactly zero, and every non-zonal mode is
    #| unchanged.
    #|
    #| With the same Runge-Kutta loop the zonal band now stays at its starting value, where
    #| rung 4's fell to $0.2$. The band has no linear sink left; in the full equations it
    #| changes through the nonlinear bracket: the Reynolds stress of the drift waves around it.

    #| The coupling a mode sees in the modified model: $C$ for $k_y \ne 0$, zero for a zonal mode.
    def modified_coupling(ky, c) {
        if ky == 0.0 {
            return 0.0
        }
        return c
    }
    zon_mod = zeros(N_SAMPLE)
    zr_mod = zonal_history(KZ, modified_coupling(0.0, 1.0), T_END, N_RK, zon_mod)
    show zr_mod
    test "in the modified model the zonal band keeps its full amplitude" {
        assert_near(zr_mod, 1.0, 0.0)
    }
    test "a non-zonal mode is unchanged: at k = (0, 1) the growth rate is still 0.09868411347" {
        assert_near(growth(0.0, 1.0, modified_coupling(1.0, 1.0)), 0.09868411346781, 1.0e-12)
    }
    figure("The same band in the two models", "t", "zeta(t) / zeta(0)", 0.0, 2.0, 0.0, 1.05)
    plot(t_rk, zon_classic, "classical")
    plot(t_rk, zon_mod, "modified")
    plot(t_rk, resid, "residual $k^2/(1 + k^2)$")
    style("residual $k^2/(1 + k^2)$", "dash")
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors