Full Plasma (Core + Edge profiles)¶

import numpy as np
import ultraplot as uplt
from matplotlib.colors import SymLogNorm
from rich import print as rprint
from rich.table import Table

from cherab.core.math import sample3d, sample3d_grid
from cherab.imas.datasets import iter_jintrac
from cherab.imas.plasma import load_plasma

# Set dark background for plots
uplt.rc.style = "dark_background"

uplt.rc["title.borderwidth"] = 0

plasma = load_plasma(path, "r")

# Organize plasma species by symbol and charge
species = {}
for s in plasma.composition:
    symbol = s.element.symbol
    if symbol not in species:
        species[symbol] = []
    species[symbol].append(s)

# Sort species of each element by charge
for symbol in species:
    species[symbol].sort(key=lambda x: x.charge)

# Print table of plasma species
table = Table(title="Plasma Species")
table.add_column("Element", style="cyan", justify="right")
table.add_column("Charge")
for symbol in species:
    for i, s in enumerate(species[symbol]):
        if i == 0:
            table.add_row(symbol, f"{s.charge:>2d}+")
        else:
            table.add_row("", f"{s.charge:>2d}+")
rprint(table)

# Sample electron density
r_pts, _, z_pts, n_e = sample3d(
    plasma.electron_distribution.density,
    (R_MIN, R_MAX, n_r),
    (0, 0, 1),
    (Z_MIN, Z_MAX, n_z),
)
n_e = n_e.squeeze()

# Plot electron density
fig, ax = uplt.subplots()
ax = plot_quantity(
    ax,
    n_e,
    extent,
    logscale=True,
)
ax.format(
    title="$n_\\mathrm{e}$ [/m³]",
)

# Sample electron temperature
te_plasma = sample3d_grid(
    plasma.electron_distribution.effective_temperature,
    r_pts,
    [0],
    z_pts,
).squeeze()

# Plot electron temperature
fig, ax = uplt.subplots()
ax = plot_quantity(
    ax,
    te_plasma,
    extent,
    logscale=True,
)
ax.format(
    title="$T_\\mathrm{e}$ [eV]",
)

def species_label(species) -> str:
    """Construct element symbol with charge label"""
    if species.charge == 0:
        label = f"{species.element.symbol}$^{{0}}$"
    elif species.charge == 1:
        label = f"{species.element.symbol}$^{{+}}$"
    else:
        label = f"{species.element.symbol}$^{{{species.charge}+}}$"
    return label


# === Density Plot ===
fig, axs = uplt.subplots(
    nrows=2,
    ncols=2,
)
axs.format(suptitle="Hydrogenic Species Density [/m³]")
for i_row, s_list in enumerate([species["D"], species["T"]]):
    for i_col, s in enumerate(s_list):
        # Sample density
        density = sample3d_grid(
            s.distribution.density,
            r_pts,
            [0],
            z_pts,
        ).squeeze()

        # Plot density
        ax = plot_quantity(
            axs[i_row, i_col],
            density,
            extent,
            logscale=True,
        )
        ax.format(urtitle=species_label(s))


# === Temperature Plot ===
fig, axs = uplt.subplots(
    nrows=2,
    ncols=2,
)
axs.format(suptitle="Hydrogenic Species Temperature [eV]")
for i_row, s_list in enumerate([species["D"], species["T"]]):
    for i_col, s in enumerate(s_list):
        # Sample temperature
        temperature = sample3d_grid(
            s.distribution.effective_temperature,
            r_pts,
            [0],
            z_pts,
        ).squeeze()

        # Plot deutrium temperature
        ax = plot_quantity(
            axs[i_row, i_col],
            temperature,
            extent,
            logscale=True,
        )
        ax.format(urtitle=species_label(s))

for s_list in [species["He"], species["Ne"], species["W"]]:
    # Each element gets its own figure
    fig, axs = uplt.subplots(
        ncols=len(s_list),
    )
    axs.format(suptitle=f"{s_list[0].element.name.capitalize()} Density [/m³]")

    for i_col, s in enumerate(s_list):
        # Sample density
        density = sample3d_grid(
            s.distribution.density,
            r_pts,
            [0],
            z_pts,
        ).squeeze()

        # Make non-physical negative densities zero
        density[density <= 0] = 0.0

        # Plot density
        ax = plot_quantity(
            axs[i_col],
            density,
            extent,
            logscale=True,
        )
        ax.format(urtitle=species_label(s))