PIC Simulation Agentic Workflow: Automated Safe Drift Optimization Under the JAX-in-Cell Framework

The open-source PIC_AGENTIC_WORKFLOW project from the University of Washington Plasma Research Group innovatively integrates large language model (LLM)-driven agents into Particle-in-Cell (PIC) computing. Built on the JAX-in-Cell framework, this project enables automated parameter optimization and safe drift control (including momentum/energy/charge conservation drifts) for plasma physics simulations. It supports three human-machine collaboration modes—supervised, automatic, and hybrid—significantly improving simulation efficiency and reliability, with applications in nuclear fusion, accelerator physics, space physics, and other fields.

The Particle-in-Cell (PIC) method is a core numerical tool for plasma physics research, widely used in nuclear fusion, space physics, accelerator physics, and other fields. It captures complex dynamic effects by tracking the motion of charged particles and solving Maxwell's equations in a coupled manner. However, PIC simulations face challenges such as high computational cost, difficult parameter tuning, and numerical instability: traditional manual parameter tuning is inefficient, and automated optimization requires domain experts to define constraints and safety boundaries.

JAX-in-Cell Base Framework

JAX-in-Cell is a modern PIC implementation based on JAX, with advantages such as automatic differentiation, JIT compilation (performance close to handwritten CUDA), function transformations (vmap/pmap support for batch processing parallelism), and pure function design.

Hierarchical Agent Architecture

• Planning Agent: Formulates high-level strategies and decomposes optimization goals into subtasks (e.g., parameter search strategies, evaluation metric definition).
• Execution Agent: Calls JAX-in-Cell to run simulations, collects metrics, adjusts parameters, and maintains a queue of parameter suggestions.
• Safety Agent: Monitors parameter suggestions and simulation runs, detects violations of conservation laws for physical quantities, numerical instability, etc., and intervenes.
• Evaluation Agent: Analyzes simulation results, calculates objective function values, identifies drift patterns, and generates feedback reports.

Constraint Classification

• Hard Constraints: Absolutely non-violable (e.g., Courant condition, non-negative particle count, finite field quantities, memory limits). Violation leads to immediate simulation termination and parameter rollback.
• Soft Constraints: Recommended to follow (e.g., upper limit of energy drift rate, momentum conservation accuracy). Added to the objective function as penalty terms for trade-off.
• Meta Constraints: Constraints on the optimization process (e.g., maximum number of iterations, parameter change amplitude limits, exploration vs. exploitation balance).

Safety Intervention Mechanism

Hierarchical response: Early warning (record anomalies) → Intervention (pause parameter adjustment and retry) → Termination (abandon parameter combination) → Rollback (restore safe state and re-plan). All events are recorded in the safety log.

Typical Optimization Cycle

1. Initialization: Load basic configuration, set optimization goals and constraints;
2. Parameter Suggestion: Planning Agent generates candidate parameter sets;
3. Safety Review: Safety Agent pre-reviews and filters out violations;
4. Simulation Execution: Execution Agent calls JAX-in-Cell to run simulations;
5. Real-time Monitoring: Safety Agent monitors status and intervenes;
6. Result Evaluation: Evaluation Agent analyzes outputs and calculates objective function values;
7. Feedback Learning: Agents update their states and adjust strategies;
8. Convergence Check: Check if the goal is achieved or iteration limit is reached.

Human-Machine Collaboration Modes

• Supervised Mode: Key decision points require human confirmation (suitable for exploratory research);
• Automatic Mode: Agents run autonomously and only report results and anomalies (suitable for batch tasks);
• Hybrid Mode: Routine decisions are handled automatically; complex situations request guidance (balances efficiency and reliability).

• Nuclear Fusion Research: Optimize wave-particle interactions, heating efficiency, and transport barrier conditions in tokamaks/stellarators;
• Accelerator Physics: Design and optimize beam focusing stability, wakefield effect suppression, and brightness;
• Space Physics: Simulate reconnection processes and particle acceleration mechanisms in solar wind-magnetosphere interactions;
• Basic Plasma Physics: Study classic problems such as two-stream instability, Landau damping, and wave-wave interactions.

Current Limitations

1. LLM Cost: Frequent API calls incur high expenses;
2. Domain Knowledge Dependence: Prompt engineering and constraint definition require PIC expertise;
3. Convergence Guarantee: Complex problems are prone to local optima;
4. Debugging Complexity: Tracking the behavior of multi-agent systems is difficult.

Future Directions

• Reinforcement Learning: Replace LLMs to reduce costs;
• Transfer Learning: Transfer optimization experience across PIC problems;
• Uncertainty Quantification: Evaluate the statistical significance of optimization results;
• Visualization Tools: Develop monitoring and debugging interfaces.

PIC_AGENTIC_WORKFLOW is a cutting-edge exploration of the integration of scientific computing and AI, establishing a "safe exploration" paradigm that balances automation and reliability. The project is hosted by the University of Washington Plasma Physics Group, and contributions are welcome:

• Application cases for new physical problems;
• Improvements to optimization algorithms;
• Expansion of constraint rules;
• Development of visualization tools.