OpenFUSIONToolkit.TokaMaker.pulse_design Namespace Reference

Detailed Description

TokaMaker + TORAX Coupled Pulse Design and Simulation Workflow (TokaMaker_TORAX/tmtx).

TokaMaker (tm) = free-boundary equilibrium solver, part of OFT TORAX (tx) = Python Jax transport solver (developed by Google DeepMind) https://github.com/google-deepmind/torax

TokaMaker_TORAX workflow couples the two codes for pulse planning, predictive kinetic equilibria, and other integrated modeling applications.

Authors

Freddie Sheehan and John Lhota

Date

May 2026

Classes
class	MyEncoder
	JSON Encoder Object to store simulation results. More...
class	TokaMaker_TORAX
	TokaMaker + TORAX Coupled Pulse Simulation Code. More...

Functions
	_coil_aturn_clim (tt)
	Return (min, max) colormap limits in A-turns from _coil_bounds (A/turn * turns).
	_coil_net_turns (tt, cname)
	Return net_turns for a coil set name, defaulting to 1.
	_draw_equil_from_dir (ax, tt, loop, idx, t_now, equil_dir)
	Draw equilibrium from pre-rendered PNG in equil_dir.
	_draw_info (ax, tt, loop, run_name)
	_draw_profiles_movie (axes, tt, idx)
	Draw the 6 radial profile panels for a movie frame.
	_draw_scalars_movie (axes, tt, times, t_now, flux_con_tm, flux_con_tx)
	Draw the 6 scalar time-series panels for a movie frame.
	_embed_video_jupyter (video_path)
	Embed an MP4 video in a Jupyter notebook.
	_encode_video_from_dir (frame_dir, out_path, fps=2)
	Encode frames from a directory into an MP4 file.
	_in_jupyter ()
	Return True if running inside a Jupyter notebook.
	_make_temp_dir_viz ()
	Create temp directory for visualization temporary files: RAM-backed on Linux, OS default elsewhere.
	_prof (state_dict, idx)
	Return (x, y) arrays from a profile state dict, or (None, None).
	_render_equil_frames (tt, loop, equil_dir)
	Render equilibrium plots from stored equilibrium snapshots to PNG files.
	_render_frame (tt, loop, idx, t_now, times, flux_con_tm, flux_con_tx, out_path, run_name, equil_dir)
	Create a single movie frame.
	_save_or_display (fig, save_path=None, display=True)
	Save figure and/or display it, then close.
	_style (ax)
	Apply consistent styling to a matplotlib Axes.
	_trace_lcfs_for_evolution_plot (equil)
	Try successive normalized psi surfaces until tracing returns a non-empty contour.
	_tx_profile_at_rho (tt, var_name, rho_val, rho_coord='rho_norm', scale=1.0)
	Return (times, values) for a profile variable at a fixed rho, full resolution.
	_tx_scalar (tt, var_name, scale=1.0)
	Return (times, values) arrays from data_tree.scalars at full TORAX resolution.
	_x_points_active (tt, i, t=None)
	Return True when X-point targets should be applied at timestep index i.
	log_redirect_setup ()
	Step 1/3 of setup to redirect noisy outputs to log file.
	make_movie (tt, save_path=None, display=True, speed_factor=1.0, loop=None, notebook_mode=None)
	Create pulse movie from simulation data.
	plot_coils (tt, save_path=None, display=True)
	Plot coil current traces in MA-turns with limit bands.
	plot_equil_interactive (tt, loop=None, notebook_mode=None, save_path=None)
	Equilibrium viewer — widget slider in notebook, saved PNGs otherwise.
	plot_lcfs_evolution (tt, save_path=None, display=True, one_plot=False)
	Plot time evolution of the last closed flux surface for each phase.
	plot_profile_evolution (tt, save_path=None, display=True, one_plot=False)
	Plot profile evolution by pulse phase by default, or as one figure when one_plot=True.
	plot_profiles_interactive (tt)
	Interactive profile viewer with ipywidgets slider to scrub through timesteps.
	plot_scalars (tt, save_path=None, display=True)
	Plot 4x3 grid of time-series scalars.
	plot_tx_relax_profiles (tt, data_tree, *, stage, t_final_relax, user_ref_curves=None, save_path=None, display=False)
	Plot psi, n_e, T_e, T_i vs rho_norm after a short TORAX relax.
	profile_plot (tt, i, t, save_path=None, display=True)
	Detailed profile comparison at a single timestep.
	redirect_outputs_to_log (filename)
	Step 2/3 of setup to redirect noisy outputs to log file.
	summary (tt)
	Print and return key physics figures of merit from the simulation.
	tm_diagnostic_plot (tt, i, t, level_attempts, solve_succeeded, save_path=None, display=True)
	TokaMaker input/output diagnostic plot for a single timestep.
	tm_loop_summary_plot (tt, loop_level_log, save_path=None, display=True)
	Summary figure showing per-timestep GS solve outcomes.

Variables
tuple	_LCFS_EVO_PSI_TRY = (1.0, 0.99, 0.98, 0.97)
int	_NBI_W_TO_MA = 1/16e6
dict	BASE_CONFIG
str	COLOR_TM = 'steelblue'
str	COLOR_TX = 'crimson'
list	COLORS_MULTI
int	DEFAULT_LOOP0_TX_FACE_POINTS = 51
int	DIAG_FS = 13
tuple	EQDSK_SAVE_NR_NZ_SEQUENCE = (100, 150, 200, 250, 300, 350, 400, 450, 500)
float	GRID_ALPHA = 0.2
int	INFO_FS = 13
int	LABEL_FS = 11
float	LCFS_WEIGHT = 100.0
int	LEGEND_FS = 9
str	LS_PRI = '-'
str	LS_SEC = '--'
float	LW = 1.6
int	MK_SZ = 3
str	MK_TM = '.'
int	MOVIE_DPI = 200
	MOVIE_FIG_H
	MOVIE_FIG_W
float	mu_0 = 4.0 * np.pi * 1e-7
int	N_PSI = 1000
float	RELAX_FIXED_DT = 0.01
int	TICK_FS = 10
int	TITLE_FS = 13
str	VLINE_COLOR = 'black'
str	VLINE_LS = ':'
float	VLINE_LW = 1.0

Function Documentation

_coil_aturn_clim()

_coil_aturn_clim ( tt )

protected

Return (min, max) colormap limits in A-turns from _coil_bounds (A/turn * turns).

_coil_net_turns()

_coil_net_turns ( tt, cname )

protected

Return net_turns for a coil set name, defaulting to 1.

_draw_equil_from_dir()

_draw_equil_from_dir ( ax, tt, loop, idx, t_now, equil_dir )

protected

Draw equilibrium from pre-rendered PNG in equil_dir.

_draw_info()

_draw_info ( ax, tt, loop, run_name )

protected

_draw_profiles_movie()

_draw_profiles_movie ( axes, tt, idx )

protected

Draw the 6 radial profile panels for a movie frame.

_draw_scalars_movie()

_draw_scalars_movie ( axes, tt, times, t_now, flux_con_tm, flux_con_tx )

protected

Draw the 6 scalar time-series panels for a movie frame.

_embed_video_jupyter()

_embed_video_jupyter ( video_path )

protected

Embed an MP4 video in a Jupyter notebook.

_encode_video_from_dir()

_encode_video_from_dir ( frame_dir, out_path, fps = 2 )

protected

Encode frames from a directory into an MP4 file.

_in_jupyter()

_in_jupyter ( )

protected

Return True if running inside a Jupyter notebook.

_make_temp_dir_viz()

_make_temp_dir_viz ( )

protected

Create temp directory for visualization temporary files: RAM-backed on Linux, OS default elsewhere.

_prof()

_prof ( state_dict, idx )

protected

Return (x, y) arrays from a profile state dict, or (None, None).

_render_equil_frames()

_render_equil_frames ( tt, loop, equil_dir )

protected

Render equilibrium plots from stored equilibrium snapshots to PNG files.

_render_frame()

_render_frame ( tt, loop, idx, t_now, times, flux_con_tm, flux_con_tx, out_path, run_name, equil_dir )

protected

Create a single movie frame.

_save_or_display()

_save_or_display ( fig, save_path = None, display = True )

protected

Save figure and/or display it, then close.

_style()

_style ( ax )

protected

Apply consistent styling to a matplotlib Axes.

_trace_lcfs_for_evolution_plot()

_trace_lcfs_for_evolution_plot ( equil )

protected

Try successive normalized psi surfaces until tracing returns a non-empty contour.

_tx_profile_at_rho()

_tx_profile_at_rho ( tt, var_name, rho_val, rho_coord = 'rho_norm', scale = 1.0 )

protected

Return (times, values) for a profile variable at a fixed rho, full resolution.

_tx_scalar()

_tx_scalar ( tt, var_name, scale = 1.0 )

protected

Return (times, values) arrays from data_tree.scalars at full TORAX resolution.

_x_points_active()

_x_points_active ( tt, i, t = None )

protected

Return True when X-point targets should be applied at timestep index i.

log_redirect_setup()

log_redirect_setup

(

)

Step 1/3 of setup to redirect noisy outputs to log file.

Performs the initial, minimal logging setup.

• Removes any handlers pre-configured by libraries.
• Sets the root logger's level to capture all desired messages.
• Adds a console handler for critical errors only.

make_movie()

make_movie	(		tt,
			save_path = None,
			display = True,
			speed_factor = 1.0,
			loop = None,
			notebook_mode = None )

Create pulse movie from simulation data.

Renders equilibrium plots from stored equilibrium snapshots, generates composite frames in a temp directory, encodes to MP4, then cleans up.

Parameters

tt : TokaMaker_TORAX Fully-populated TokaMaker_TORAX object (after fly()). save_path : str, optional Path to save the MP4 file. If None, uses a default in cwd. display : bool If True and running in Jupyter, embed the video in the notebook. speed_factor : float Playback speed relative to real time. loop : int, optional Loop to visualize. Default: last completed loop. notebook_mode : bool or None True → embed video in notebook after saving. False → save to file only, do not embed. None → auto-detect (embed if in Jupyter and display=True).

plot_coils()

plot_coils	(		tt,
			save_path = None,
			display = True )

Plot coil current traces in MA-turns with limit bands.

plot_equil_interactive()

plot_equil_interactive	(		tt,
			loop = None,
			notebook_mode = None,
			save_path = None )

Equilibrium viewer — widget slider in notebook, saved PNGs otherwise.

Parameters

notebook_mode : bool or None True → ipywidgets slider (Jupyter). False → save PNG files to save_path. None → auto-detect (widget if in Jupyter, save files otherwise). save_path : str, optional Directory for PNG files when notebook_mode=False. Defaults to ./equil_loop<N>/ in the current directory.

plot_lcfs_evolution()

plot_lcfs_evolution	(		tt,
			save_path = None,
			display = True,
			one_plot = False )

Plot time evolution of the last closed flux surface for each phase.

Produces phase-split figures by default (rampup, flattop, rampdown), or one combined figure when one_plot=True.

plot_profile_evolution()

plot_profile_evolution	(		tt,
			save_path = None,
			display = True,
			one_plot = False )

Plot profile evolution by pulse phase by default, or as one figure when one_plot=True.

plot_profiles_interactive()

plot_profiles_interactive

(

tt

)

Interactive profile viewer with ipywidgets slider to scrub through timesteps.

plot_scalars()

plot_scalars	(		tt,
			save_path = None,
			display = True )

Plot 4x3 grid of time-series scalars.

plot_tx_relax_profiles()

plot_tx_relax_profiles	(		tt,
			data_tree,
		*	,
			stage,
			t_final_relax,
			user_ref_curves = None,
			save_path = None,
			display = False )

Plot psi, n_e, T_e, T_i vs rho_norm after a short TORAX relax.

Shows main-run history (tt._relax_mainrun_profile_history), user kinetic profiles from flattened profile_conditions, init EQDSK psi at t_initial when psi is not prescribed (initial_psi_mode='geometry'), and the profile at t_final_relax. Used when TokaMaker_TORAX.fly(..., output_mode='debug').

profile_plot()

profile_plot	(		tt,
			i,
			t,
			save_path = None,
			display = True )

Detailed profile comparison at a single timestep.

redirect_outputs_to_log()

redirect_outputs_to_log

(

filename

)

Step 2/3 of setup to redirect noisy outputs to log file.

A context manager to temporarily redirect stdout and stderr to a file.

Parameters

filename

Name of log file (self._log_file)

summary()

summary

(

tt

)

Print and return key physics figures of merit from the simulation.

tm_diagnostic_plot()

tm_diagnostic_plot	(		tt,
			i,
			t,
			level_attempts,
			solve_succeeded,
			save_path = None,
			display = True )

TokaMaker input/output diagnostic plot for a single timestep.

tm_loop_summary_plot()

tm_loop_summary_plot	(		tt,
			loop_level_log,
			save_path = None,
			display = True )

Summary figure showing per-timestep GS solve outcomes.

Variable Documentation

_LCFS_EVO_PSI_TRY

tuple _LCFS_EVO_PSI_TRY = (1.0, 0.99, 0.98, 0.97)

protected

_NBI_W_TO_MA

int _NBI_W_TO_MA = 1/16e6

protected

BASE_CONFIG

dict BASE_CONFIG

COLOR_TM

str COLOR_TM = 'steelblue'

COLOR_TX

str COLOR_TX = 'crimson'

COLORS_MULTI

list COLORS_MULTI

Initial value:

=  [
    'darkorange', 'forestgreen', 'mediumpurple',
    'goldenrod', 'deeppink', 'teal', 'sienna',
]

DEFAULT_LOOP0_TX_FACE_POINTS

int DEFAULT_LOOP0_TX_FACE_POINTS = 51

DIAG_FS

int DIAG_FS = 13

EQDSK_SAVE_NR_NZ_SEQUENCE

tuple EQDSK_SAVE_NR_NZ_SEQUENCE = (100, 150, 200, 250, 300, 350, 400, 450, 500)

GRID_ALPHA

float GRID_ALPHA = 0.2

INFO_FS

int INFO_FS = 13

LABEL_FS

int LABEL_FS = 11

LCFS_WEIGHT

float LCFS_WEIGHT = 100.0

LEGEND_FS

int LEGEND_FS = 9

LS_PRI

str LS_PRI = '-'

LS_SEC

str LS_SEC = '--'

LW

float LW = 1.6

MK_SZ

int MK_SZ = 3

MK_TM

str MK_TM = '.'

MOVIE_DPI

int MOVIE_DPI = 200

MOVIE_FIG_H

MOVIE_FIG_H

MOVIE_FIG_W

MOVIE_FIG_W

mu_0

float mu_0 = 4.0 * np.pi * 1e-7

N_PSI

int N_PSI = 1000

RELAX_FIXED_DT

float RELAX_FIXED_DT = 0.01

TICK_FS

int TICK_FS = 10

TITLE_FS

int TITLE_FS = 13

VLINE_COLOR

str VLINE_COLOR = 'black'

VLINE_LS

str VLINE_LS = ':'

VLINE_LW

float VLINE_LW = 1.0