Ideal ballooning solver

This example runs the ideal ballooning solver in pyrokinetics for a range of \(\hat{s}\) and \(\beta\) and plots this data

import matplotlib.pyplot as plt
import numpy as np

from pyrokinetics import Pyro, template_dir
from pyrokinetics.diagnostics import Diagnostics

nshat = 15
nbprime = 20
shat = np.linspace(0.0, 2, nshat)
bprime = np.linspace(0.0, -0.5, nbprime)

pyro = Pyro(gk_file=template_dir / "input.gs2")

gamma = np.empty((nshat, nbprime))

for i_s, s in enumerate(shat):
    for i_b, b in enumerate(bprime):
        pyro.local_geometry.shat = s
        pyro.local_geometry.beta_prime = b
        diag = Diagnostics(pyro)
        gamma[i_s, i_b] = diag.ideal_ballooning_solver()


fig, ax = plt.subplots(1, 1, sharex=True, figsize=(8, 7))

cs = ax.contourf(abs(bprime), shat, gamma)
cs_0 = ax.contour(
    abs(bprime),
    shat,
    gamma,
    colors="w",
    linewidths=(2,),
    levels=[
        0.0,
    ],
)

ax.clabel(cs_0, fmt="%.1e", colors="w", fontsize=22)
ax.set_ylabel(r"$\hat{s}$")
ax.set_xlabel(r"$|\beta'|$")
fig.colorbar(cs)
ax.set_title("IBM growth rate")
fig.tight_layout()
plt.show()

This would generate the following figure

Note this works for any pyro object independant of the LocalGeometry chosen.

The algorithm used to solve the ideal ballooning equation was taken from ideal-ballooning-solver