AI4Plasma Documentation

Welcome to the AI4Plasma documentation! AI4Plasma is a python- and pytorch-based library for physics-informed and operator learning, especially designed for plasma simulation. As the world’s first AI package for plasma simulation, AI4Plasma has been determined from its inception to make it easier for plasma researchers to use AI tools, thereby enhancing the efficiency and even the accuracy of plasma simulation. Although there are several excellent predecessors, such as DeepXDE, AI4Plasma is (and will always be) the AI algorithm package that understands plasma the most and the plasma algorithm package that understands AI the most.

Attention: The documentation is prepared with the help of generative AI due to the limitd time. It will be checked and improved in the next version

Getting Started

• 1. Installation Guide
  • 1.1. Prerequisites
  • 1.2. Install from PyPI
  • 1.3. Install from Source
  • 1.4. Install by Conda
  • 1.5. Build documentation locally
• 2. Quick Start
  • 2.1. Running an Example
  • 2.2. Monitor training with TensorBoard
• 3. Basic Concepts
  • 3.1. Core Architecture
  • 3.2. Physics-Informed Machine Learning (PIML)
  • 3.3. Neural Operators
  • 3.4. Equation Terms and Loss Components
  • 3.5. Visualization and Monitoring
  • 3.6. Training Utilities
  • 3.7. Plasma Physics Applications
  • 3.8. Typical Workflow
  • 3.9. References
  • 3.10. Next Steps
• 4. Configuration
  • 4.1. Floating-point precision: REAL
  • 4.2. Compute device: DEVICE
  • 4.3. Common pitfalls

User Guide

• 1. Physics-Informed Machine Learning (PIML)
  • 1.1. What is PIML?
  • 1.2. Core abstractions
  • 1.3. Geometry and sampling
  • 1.4. Autograd derivatives
  • 1.5. Loss design and weighting
  • 1.6. Metrics and evaluation
  • 1.7. Worked example: 1D steady arc (CS-PINN)
  • 1.8. Practical tips
  • 1.9. Common pitfalls
  • 1.10. Next steps
• 2. Neural Operators
  • 2.1. What is an Operator?
  • 2.2. DeepONet
  • 2.3. DeepCSNet
  • 2.4. Losses and Metrics
  • 2.5. Minimal working example (DeepONet)
  • 2.6. Common pitfalls
  • 2.7. Where to go next
• 3. Training, Logging, Checkpoints
  • 3.1. Training components
  • 3.2. TensorBoard
  • 3.3. Checkpointing
  • 3.4. Reproducibility
  • 3.5. Practical tips for PINN training
  • 3.6. Example: basic training skeleton
  • 3.7. Troubleshooting
  • 3.8. Checkpoint resumption workflow
  • 3.9. Next steps
• 4. Utilities
  • 4.1. Reproducibility
  • 4.2. Device utilities
  • 4.3. Data conversion
  • 4.4. Timing
  • 4.5. I/O helpers
  • 4.6. Math helpers
  • 4.7. Physical constants
• 5. Plasma Models
  • 5.1. Arc discharge models (ai4plasma.plasma.arc)
  • 5.2. Plasma properties and interpolation (ai4plasma.plasma.prop)
  • 5.3. Where this fits in the workflow

Examples

• 1. Physics-Informed Machine Learning (PIML)
  • 1.1. Overview
  • 1.2. Standard PINN
  • 1.3. CS-PINN (Coefficient-Subnet PINN)
  • 1.4. RK-PINN (Runge-Kutta PINN)
  • 1.5. Meta-PINN
  • 1.6. Best Practices
  • 1.7. Performance Benchmarks
  • 1.8. Visualization and Monitoring
  • 1.9. Advanced Topics
  • 1.10. Troubleshooting Guide
  • 1.11. Further Reading
  • 1.12. Quick Start Guide
• 2. Operator Learning
  • 2.1. Overview
  • 2.2. Common Setup
  • 2.3. DeepONet (Deep Operator Network)
  • 2.4. DeepCSNet (Deep Coefficient-Subnet Network)
  • 2.5. Best Practices
  • 2.6. Performance Benchmarks
  • 2.7. Troubleshooting
  • 2.8. References
  • 2.9. Further Reading
• 3. Plasma Simulation
  • 3.1. Overview
  • 3.2. Arc Discharge Simulations
  • 3.3. Examples
  • 3.4. Best Practices
  • 3.5. Performance Benchmarks
  • 3.6. Troubleshooting Guide
  • 3.7. Advanced Topics
  • 3.8. Applications
  • 3.9. References
  • 3.10. Quick Start Guide

Core Modules

• 1. Core Module
  • 1.1. Model
  • 1.2. Network
• 2. Physics-Informed Machine Learning (PIML)
  • 2.1. Geometry
  • 2.2. PINN
  • 2.3. CS-PINN (Conservative PINN)
  • 2.4. RK-PINN (Runge-Kutta PINN)
  • 2.5. Meta-PINN (Meta-Learning PINN)
  • 2.6. NAS-PINN (Neural Architecture Search for PINN)
• 3. Operator Learning Module
  • 3.1. DeepONet
  • 3.2. DeepCSNet
• 4. Plasma Physics Module
  • 4.1. Properties
  • 4.2. Arc Discharge
• 5. Utilities
  • 5.1. Common
  • 5.2. Math
  • 5.3. Device
  • 5.4. I/O

Development

• 1. Contributing

Overview

AI4Plasma provides a flexible framework for:

• Operator: A neural network architecture for learning nonlinear operators

• PIML: Physics-Informed Neural Networks for solving PDEs

• Utilities: Common tools for data handling, device management, and mathematical operations

• Core Models: Base classes and network architectures for neural network models

Key Features

✨ Modular Design: Clean separation of concerns with dedicated modules for operators, PINNs, and utilities

⚡ GPU Support: Automatic GPU/CPU device management for efficient computation

🔬 Physics-Informed: Built-in support for physics-informed learning approaches

📊 Flexible Data Handling: Support for various data formats and batch processing strategies

Quick Example

import torch
import numpy as np
from ai4plasma.operator.deeponet import DeepONet, DeepONetModel
from ai4plasma.core.network import FNN

# Create branch and trunk networks
branch_net = FNN(layers=[1, 64, 64, 64])
trunk_net = FNN(layers=[1, 64, 64, 64])

# Initialize DeepONet
network = DeepONet(branch_net, trunk_net)
model = DeepONetModel(network=network)

# Prepare training data
model.prepare_train_data(branch_inputs, trunk_inputs, targets, batch_size=32)

# Train the model
model.train(num_epochs=100, lr=1e-3)

# Make predictions
predictions = model.predict(test_branch, test_trunk)

Package Structure

ai4plasma/
├── core/          # Base models and network architectures
├── operator/      # Neural operator models (DeepONet, etc.)
├── piml/          # Physics-Informed Machine Learning (PINN, CS-PINN, RK-PINN, ...)
└── utils/         # Utility functions and common tools

Installation

See Installation Guide for detailed setup instructions.

Getting Help

• 📖 Read the Getting Started Guide

• 🔍 Check the API Reference

• 📚 Browse Examples

• 💬 Open an issue on GitHub

Citation

If you use AI4Plasma in your research, please cite:

@software{ai4plasma2026,
  author = {Zhong, Linlin},
  title = {AI4Plasma: An AI Library for Plasma Physics Simulation},
  year = {2026},
  url = {https://github.com/mathboylinlin/ai4plasma}
}

License

This project is licensed under the MIT License - see the LICENSE file for details.

---

Last Updated: January 14, 2026