Other Modules

This section provides documentation for additional modules in GymTorax:

Logging System

Comprehensive logging configuration for debugging and monitoring.

gymtorax.logger.setup_logging(level=logging.WARNING, log_file=None, suppress_external=True)[source]

Setup global logging configuration with optional external library suppression.

When using DEBUG level, external libraries (JAX, TORAX, etc.) can generate overwhelming amounts of log messages. This function allows to debug the gymtorax code while keeping external libraries at WARNING level or higher.

Parameters:

  • level (int) – Logging level for gymtorax modules (e.g., logging.DEBUG, logging.INFO, logging.WARNING, …).

  • log_file (str or None) – If provided, logs will also be written to this file.

  • suppress_external (bool) – If True and level=DEBUG, suppress verbose output from external libraries (JAX, TORAX, TensorFlow, etc.) by setting them to WARNING level. Default: True.

Example

>>> # Debug gymtorax only, suppress external libraries
>>> setup_logging(level=logging.DEBUG)
>>>
>>> # Debug everything including external libraries
>>> setup_logging(level=logging.DEBUG, suppress_external=False)
>>>
>>> # Normal usage (unchanged behavior)
>>> setup_logging(level=logging.WARNING)

Logging Configuration

from gymtorax.logger import setup_logging
import logging

# Configure comprehensive logging
setup_logging(
    level=logging.DEBUG,
    log_file="simulation.log",
    suppress_external=True  # Quiet JAX/TORAX messages
)

# Environment will use configured logging
env = gym.make('gymtorax/IterHybrid-v0', log_level="debug")

Custom Log Levels

# Different log levels for different components
setup_logging(level=logging.INFO)

# Set specific loggers
logging.getLogger("gymtorax.torax_wrapper").setLevel(logging.DEBUG)
logging.getLogger("gymtorax.envs").setLevel(logging.WARNING)

Visualization and Rendering

Real-time visualization system with TORAX integration for plasma simulation monitoring and video recording.

Real-time visualization and rendering system for GymTORAX plasma simulations.

This module provides a comprehensive visualization framework for TORAX plasma simulation data within the Gymnasium environment ecosystem. It combines TORAX native plotting capabilities with Gymnasium rendering pipeline to support both live visualization and video recording of simulation episodes.

The system is designed to handle real-time visualization during RL training while maintaining full compatibility with TORAX plot configurations and styling conventions. It supports multiple rendering modes and provides optimized performance for video generation.

gymtorax.rendering.process_plot_config(plot_config: str | plotruns_lib.FigureProperties) → plotruns_lib.FigureProperties[source]

Load and validate TORAX plot configuration from string identifier or object.

This function handles the loading of TORAX plot configurations, which define the layout, variables, and styling for plasma simulation visualizations. It supports both string-based configuration names (loaded from TORAX default configurations) and direct FigureProperties objects.

Parameters:

plot_config (str or plotruns_lib.FigureProperties) – Either a string identifier for a default TORAX plot configuration (e.g., "default", "simple", …) or a pre-configured FigureProperties object.

Returns:

Validated plot configuration object

containing axes definitions, layout properties, and styling options.

Return type:

plotruns_lib.FigureProperties

Raises:

  • TypeError – If plot_config is neither a string nor FigureProperties instance.

  • ImportError – If the specified plot configuration module is not found.

  • AttributeError – If the configuration module lacks the PLOT_CONFIG attribute.

Note

String-based configurations are loaded from the TORAX package at torax.plotting.configs.{plot_config}_plot_config.PLOT_CONFIG.

class gymtorax.rendering.Plotter(plot_config: torax._src.plotting.plotruns_lib.FigureProperties, render_mode: str)[source]

Bases: object

Real-time plotter using TORAX plot style and axis conventions.

This class provides real-time visualization capabilities for TORAX plasma simulation data with support for both spatial profiles and time series plots.

The plotter can operate in different render modes and supports live visualization during simulation runs. It automatically handles figure layout, axis scaling, and legend management based on the provided configuration.

Variables:

  • plot_config (plotruns_lib.FigureProperties) – Configuration for plot layout and styling

  • render_mode (str) – Rendering mode, either "human" for interactive display or "rgb_array" for video recording

  • fig (matplotlib.figure.Figure) – Main figure object

  • axes (list[matplotlib.axes.Axes]) – List of matplotlib axes objects

  • lines (list[matplotlib.lines.Line2D]) – List of matplotlib line objects for plotting

Parameters:

  • plot_config (plotruns_lib.FigureProperties)

  • render_mode (str)

__init__(plot_config: torax._src.plotting.plotruns_lib.FigureProperties, render_mode: str)[source]

Initialize the real-time plotter with configuration and display settings.

Sets up the plotter with the specified configuration and creates the matplotlib figure structure.

Automatically configures font scaling based on render mode and plot layout to ensure optimal readability in different visualization contexts.

Parameters:

  • plot_config (plotruns_lib.FigureProperties) – Configuration object defining the plot layout, variables to display, axis properties, and styling options. This determines which plasma variables are plotted and how they are arranged.

  • render_mode (str) – Rendering mode that affects font scaling and backend selection. Supported values are "human" for interactive display, "rgb_array" for video recording.

Note

The figure and axes are created immediately during initialization based on the plot_config. Data histories are initialized as empty and will be populated through subsequent update() calls. Font scaling is automatically applied for multi-row layouts to maintain readability.

reset()[source]

Reset the plotter to its initial state without closing the figure.

This method clears all accumulated data histories and resets all plot lines to empty state. The figure remains open and ready for new data. If in human render mode, the display is refreshed to show the cleared state.

Note

This method is useful for starting a new simulation run while keeping the same plotter instance and figure window.

update(current_state: DataTree, t: float)[source]

Update the visualization with new simulation data from TORAX output.

This method processes the latest simulation state and updates all plot lines with new data points.

The method performs data validation, coordinate transformation, and line updates while maintaining TORAX native plotting conventions and styling.

Parameters:

  • current_state (xarray.DataTree) – Complete simulation output from TORAX, containing profiles and scalars datasets with plasma variables (temperatures, densities, currents, etc.).

  • t (float) – Current simulation time in seconds. Used for time series data accumulation and display labeling.

Raises:

ValueError – If required variables specified in plot_config are missing from the current_state DataTree.

Note

On first call, this method initializes all plot lines and formatting. Subsequent calls efficiently update existing lines with new data points. The method automatically handles data unit conversions as defined in the TORAX plotting system.

render_frame(t: float | None = None)[source]

Render and display a single frame for interactive visualization.

This method completes the visualization pipeline by updating axis limits and legends , and refreshing the display for interactive viewing. It is designed for "human" render mode where users observe the simulation in real-time through matplotlib windows.

Parameters:

t (float) – Current simulation time in seconds to display in the figure title. Formats as “t = {value:.3f}”.

close()[source]

Close the visualization figure and release all associated resources.

This method properly closes the matplotlib figure and releases memory resources associated with the visualization. It should be called when the plotter is no longer needed.

render_rgb_array(t: float | None = None) → ndarray[source]

Render the current visualization state as RGB array for video recording.

This method captures the current plot state as a RGB image suitable for video generation and programmatic processing. It performs the same axis scaling and legend management as render_frame() but outputs a numpy array instead of displaying interactively.

Parameters:

t (float) – Current simulation time in seconds to display in the figure title. Formats as "t = {value:.3f}".

Returns:

RGB image array with shape (height, width, 3) and dtype

uint8. Values are in the range [0, 255] representing RGB color channels. The array can be directly used by video encoding libraries or saved as image files.

Return type:

numpy.ndarray

Real-time Visualization

The visualization system provides live plotting during simulation episodes with support for both interactive display and video recording:

import gymnasium as gym
from gymnasium.wrappers import RecordVideo

# Interactive visualization
env = gym.make('gymtorax/IterHybrid-v0', render_mode="human")
obs, info = env.reset()

terminated = False
while not terminated:
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    env.render()  # Live matplotlib display

    if terminated or truncated:
        break

TORAX Plot Configuration

Uses standard TORAX plot configurations, or create custom configurations and pass them to the plot_config argument of the environment:

from gymtorax.rendering import Plotter, process_plot_config

# Custom configuration for specific variables
env = gym.make('gymtorax/IterHybrid-v0',
               render_mode="human",
               plot_config="simple")