AI4Plasma

AI4Plasma is a Python- and PyTorch-based library for physics-informed machine learning and operator learning, specifically designed for plasma physics simulation. As the world's first AI package dedicated to plasma simulation, AI4Plasma bridges the gap between cutting-edge AI techniques and plasma physics research, making it easier for plasma researchers to leverage AI tools to enhance both the efficiency and accuracy of their simulations.

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. We welcome your testing and use, and also invite you to join the AI4Plasma development team.

✨ Key Features

• Physics-Informed Machine Learning (PIML)

  • 🧠 PINNs: Classic Physics-Informed Neural Networks for solving PDEs
  • 🎯 CS-PINNs: Coefficient-Subnet PINNs optimized for solving plasma equations with variable coefficients
  • ⏰ RK-PINNs: Runge-Kutta PINNs for time-dependent problems
  • 🔄 Meta-PINNs: Meta-learning approach for rapid adaptation across different physics scenarios
  • 🔍 NAS-PINNs: Neural Architecture Search for automatic PINN architecture optimization

• Operator Learning

  • 🌐 DeepONet: Deep Operator Networks for learning solution operators
  • 🔬 DeepCSNet: Deep Operator Networks for predicting cross sections

• Plasma Physics Models

  • ⚡ Arc Plasma Simulation: Steady and transient arc plasma simulations (1-D in current version)
  • 📊 Plasma Properties: Built-in plasma property calculations and estimations

• Comprehensive Utilities

  • 📐 Geometry Tools: 1D, 2D, and 3D domain definitions with flexible boundary conditions
  • 🎨 Visualization: TensorBoard integration and custom plotting callbacks
  • 🚀 Auto-differentiation: Efficient gradient computation for complex PDEs
  • 💾 I/O Management: Easy model saving/loading and checkpoint management
  • 🖥️ GPU Support: Automatic device detection and optimization

📦 Installation

Install from PyPI

The easiest way to install AI4Plasma is via pip:

pip install ai4plasma

Or upgrade to the latest version:

pip install --upgrade ai4plasma

Install from Source

For the latest development version:

git clone https://github.com/ai4plasma/ai4plasma.git
cd ai4plasma
pip install -e .

Install with Conda

If you prefer using Conda/Mamba:

conda create -n ai4plasma python=3.12
conda activate ai4plasma
pip install -e .

Prerequisites

• Python >= 3.10
• PyTorch >= 2.0
• NumPy, SciPy, Pandas, Matplotlib
• TensorBoard, Imageio (for visualization)
• FiPy (for traditional numerical methods comparison)
• Shapely (for geometry processing)
• Huggingface-hub (for downloading dataset from Hugging Face)

🚀 Quick Start

Example 1: Solving a Simple PDE with PINN

Solve the 1D ODE: d²u/dx² = -sin(x) with boundary conditions u(0) = u(1) = 0

import torch
import torch.nn as nn
from ai4plasma.config import REAL
from ai4plasma.piml.pinn import PINN
from ai4plasma.piml.geo import Geo1D
from ai4plasma.core.network import FNN
from ai4plasma.utils.math import df_dX, calc_relative_l2_err
from ai4plasma.utils.common import set_seed, numpy2torch

# Set seed for reproducibility
set_seed(2026)

# Define custom PINN class
class SimplePINN(PINN):
    def __init__(self, network):
        self.geo = Geo1D([0.0, 1.0])
        super().__init__(network)
    
    @staticmethod
    def _pde_residual(network, x):
        """Compute residual: d²u/dx² + sin(x)"""
        u = network(x)
        u_x = df_dX(u, x)
        u_xx = df_dX(u_x, x)
        return u_xx + torch.sin(x)
    
    @staticmethod
    def _bc_residual(network, x):
        """Dirichlet BC: u(x) = 0"""
        return network(x)
    
    def _define_loss_terms(self):
        """Set up domain and boundary loss terms"""
        x_domain = self.geo.sample_domain(100, mode='uniform')
        x_bc = self.geo.sample_boundary()
        
        self.add_equation('Domain', self._pde_residual, weight=1.0, data=x_domain)
        self.add_equation('Left BC', self._bc_residual, weight=10.0, data=x_bc[0])
        self.add_equation('Right BC', self._bc_residual, weight=10.0, data=x_bc[1])

# Create and train PINN
network = FNN([1, 64, 64, 64, 1])
pinn = SimplePINN(network)
pinn.set_loss_func(nn.MSELoss())

pinn.train(num_epochs=5000, lr=1e-3, print_loss=True, print_loss_freq=500)

# Predict on evaluation points
import numpy as np
x_eval = np.linspace(0, 1, 200, dtype=REAL()).reshape(-1, 1)
u_true = np.sin(x_eval) - x_eval * np.sin(1)  # Analytical solution for comparison
x_eval_tensor = numpy2torch(x_eval)
u_pred = pinn.predict(x_eval_tensor).detach().numpy()

print(f"Relative L2 error: {calc_relative_l2_err(u_true, u_pred):.4e}")

Example 2: Learning Operators with DeepONet

Learn the solution operator for -Δu(x) = f(x) where f(x) = v·π²·sin(πx) and u(x) = v·sin(πx)

import numpy as np
from ai4plasma.core.network import FNN
from ai4plasma.operator.deeponet import DeepONet, DeepONetModel
from ai4plasma.utils.common import set_seed, numpy2torch
from ai4plasma.utils.device import check_gpu
from ai4plasma.utils.math import calc_relative_l2_err
from ai4plasma.config import DEVICE, REAL

# Set seed for reproducibility
set_seed(2026)

## Set device ##
if check_gpu(print_required=True):
    DEVICE.set_device(0)  # Using cuda:0
else:
    DEVICE.set_device(-1) # Using cpu
print(DEVICE)

# Define branch and trunk networks
branch_net = FNN([1, 10, 10, 10, 10])
trunk_net = FNN([1, 10, 10, 10, 10])

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

# Prepare training data
# Branch input: parameter v (A samples)
v = np.array([[2, 4, 6, 8, 10]], dtype=REAL()).reshape((-1, 1))
# Trunk input: spatial coordinate x (B points)
x = np.linspace(-1, 1, 40, endpoint=True, dtype=REAL()).reshape((-1, 1))
# Target output: solution u (A × B)
u = v * np.sin(np.pi * x.T)

# Convert to tensors and prepare data
v, x, u = numpy2torch(v), numpy2torch(x), numpy2torch(u)
model.prepare_train_data(v, x, u)

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

# Test on unseen parameter
vv = np.array([[5.5]], dtype=REAL())
xx = np.linspace(-1, 1, 30, endpoint=True, dtype=REAL()).reshape((-1, 1))
uu = vv * np.sin(np.pi * xx.T)  # True solution

# Predict
vv, xx = numpy2torch(vv), numpy2torch(xx)
u_pred = model.predict(vv, xx).cpu().detach().numpy()

# Evaluate
l2_err = calc_relative_l2_err(uu, u_pred)
print(f"Relative L2 error: {l2_err:.4e}")

📚 Module Overview

Core Modules (ai4plasma.core)

• model.py: Base model class with training utilities and checkpoint management
• network.py: Neural network architectures (FNN, CNN, ResNet, etc.)

Physics-Informed ML (ai4plasma.piml)

• geo.py: Geometry definitions and domain sampling
• pinn.py: Standard Physics-Informed Neural Networks
• cs_pinn.py: Coefficient-Subnet PINNs (CS-PINN) for plasma applications
• meta_pinn.py: Meta-learning PINNs (Meta-PINN) for multi-task scenarios
• rk_pinn.py: Runge-Kutta PINNs (RK-PINN) for temporal problems
• nas_pinn.py: Neural Architecture Search for PINNs (NAS-PINN)

Operator Learning (ai4plasma.operator)

• deeponet.py: Deep Operator Networks implementation
• deepcsnet.py: Deep Operator Networks for cross section prediction

Plasma Physics (ai4plasma.plasma)

• arc.py: Arc plasma models and solvers
• prop.py: Plasma property calculations

Utilities (ai4plasma.utils)

• common.py: Common utilities (seed setting, timers, etc.)
• device.py: GPU/CPU device management
• io.py: File I/O and checkpoint utilities
• math.py: Mathematical utilities (automatic differentiation, operators)

📖 Documentation

Full documentation is available at https://ai4plasma.readthedocs.io

• API Reference: Detailed documentation of all modules and classes
• User Guide: Step-by-step tutorials for getting started
• Examples: Comprehensive examples covering various use cases

🎯 Example Applications

The app/ directory contains numerous ready-to-run examples:

Physics-Informed Neural Networks

• 1D/2D Poisson Equation: app/piml/pinn/solve_1d_pinn.py
• 2D Rectangular Domain: app/piml/pinn/solve_2d_rect_pinn.py
• Polynomial Domain: app/piml/pinn/solve_2d_poly_pinn.py

Plasma Simulations

• Steady Arc by FVM: app/plasma/arc/solve_1d_arc_steady.py
• Transient Arc by FVM: app/plasma/arc/solve_1d_arc_transient_explicit.py
• Steady Arc by CS-PINN: app/piml/cs_pinn/solve_1d_arc_steady_cs_pinn.py
• Transient Arc (with radial velocity) by CS-PINN: app/piml/cs_pinn/solve_1d_arc_transient_cs_pinn.py
• Transient Arc (without radial velocity) by CS-PINN: app/piml/cs_pinn/solve_1d_arc_transient_noV_cs_pinn.py
• Corona Discharge by RK-PINN: app/piml/rk_pinn/solve_1d_corona_rk_pinn.py
• Resume Training: app/piml/cs_pinn/resume_1d_arc_transient_cs_pinn.py

Operator Learning

• DeepONet for Poisson: app/operator/deeponet/solve_1d_poisson.py
• 2D CNN-based DeepONet: app/operator/deeponet/solve_2d_poisson_cnn.py
• DeepCSNet for Cross-Sections: app/operator/deepcsnet/predict_total_ionxsec.py

Meta-Learning & NAS

• Meta-PINN: app/piml/meta_pinn/solve_1d_arc_steady_meta_pinn.py
• NAS-PINN: app/piml/nas_pinn/search_pinn_2d_poisson.py

📊 Project Structure

ai4plasma/
├── ai4plasma/              # Main package
│   ├── core/              # Core model and network architectures
│   ├── piml/              # Physics-informed machine learning
│   ├── operator/          # Operator learning methods
│   ├── plasma/            # Plasma physics models
│   └── utils/             # Utility functions
├── app/                   # Example applications
│   ├── piml/              # PIML examples
│   ├── operator/          # Operator learning examples
│   └── plasma/            # Plasma simulation examples
├── docs/                  # Documentation source files

🤝 Contributing

We welcome contributions from the community! Whether you're fixing bugs, adding new features, improving documentation, or sharing examples, your help is appreciated.

How to Contribute

1. Fork the repository on GitHub
2. Create a new branch for your feature or bugfix

   git checkout -b feature/your-feature-name
   

3. Make your changes and commit with clear messages
4. Submit a pull request with a detailed description

Code Style

• Follow PEP 8 guidelines & "Numpy" style comment
• Use type hints where applicable
• Write docstrings for all public functions and classes
• Keep functions focused and modular

📝 Citation

If you use AI4Plasma in your research, please cite:

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

For specific methods, please also cite the relevant papers:

• PINNs: M. Raissi, P. Perdikaris, and G. E. Karniadakis, "Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations," Journal of Computational Physics 378, 686 (2019).

• CS-PINN: L. Zhong, B. Wu, and Y. Wang, "Low-temperature plasma simulation based on physics-informed neural networks: Frameworks and preliminary applications," Physics of Fluids 34, 087116 (2022)

• RK-PINN: L. Zhong, B. Wu, and Y. Wang, "Low-temperature plasma simulation based on physics-informed neural networks: Frameworks and preliminary applications," Physics of Fluids 34, 087116 (2022)

• Meta-PINN: L. Zhong, B. Wu, and Y. Wang, "Accelerating physics-informed neural network based 1D arc simulation by meta learning," Journal of Physics D: Applied Physics 56, 074006 (2023).

• NAS-PINN: Y. Wang, and L. Zhong, "NAS-PINN: Neural architecture search-guided physics-informed neural network for solving PDEs," Journal of Computational Physics 496, 112603 (2024).

• DeepONet: L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis, "Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators," Nature Machine Intelligence 3, 218 (2021).

• DeepCSNet: Y. Wang, and L. Zhong, "DeepCSNet: a deep learning method for predicting electron-impact doubly differential ionization cross sections," Plasma Sources Science and Technology 33, 105012 (2024).

📄 License

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

🙏 Acknowledgments

• Thanks to all contributors who have helped shape AI4Plasma
• Inspired by DeepXDE and other physics-informed ML and operator learning frameworks
• Built on top of PyTorch and the broader Python scientific computing ecosystem
• Special thanks to the plasma physics and machine learning communities

📬 Contact

• Author: Linlin Zhong
• Email: linlin@seu.edu.cn
• Homepage: https://github.com/mathboylinlin/ai4plasma
• Documentation: https://ai4plasma.readthedocs.io

🌟 Star History

If you find AI4Plasma useful, please consider giving it a star ⭐ on GitHub!

---

Made with ❤️ by the AI4Plasma Team