Breakthrough

Berkeley Lab researchers have developed a new machine learning model, the spatiotemporal Fourier transformer (StFT), that can accurately predict the long-term behavior of complex physical systems. These systems — such as turbulent plasma in fusion reactors and large-scale fluid flows — involve processes that occur at multiple spatial and temporal scales, constantly interacting with one another. StFT outperforms state-of-the-art methods by explicitly capturing these multi-scale interactions through a hierarchical architecture that combines spatial and frequency-domain information. This design enables it to model how large, slow-changing patterns influence small, fast-changing ones — and vice versa. A built-in error-correction mechanism improves stability and provides meaningful uncertainty estimates. 

This capability is especially important for building digital twins — virtual replicas of real-world systems — that can be used to test scenarios, optimize designs, and inform operational decisions. With StFT, scientists can create more accurate and reliable digital twins for applications ranging from atmospheric modeling to the design and control of fusion power plants.

Background

Modeling and simulating the long-term dynamics of complex physical systems present significant challenges in science and engineering. These systems involve tightly coupled processes that occur across vastly different spatial and temporal scales. Accurately predicting their behavior over time is critical for understanding phenomena such as turbulence, which poses major challenges in various fields of physics and engineering. For example, in aerospace engineering, understanding turbulent airflow around aircraft is essential for optimizing design and ensuring flight safety. In mechanical systems, turbulence in fluid flows can affect performance and efficiency in applications such as pumping, cooling, and manufacturing processes.

The project lead, Zhe Bai of Berkeley Lab, highlights the difficulties in modeling these systems, noting that they often exhibit multiscale and highly intricate dynamics. Conventional methods, like reduced-order modeling and projection-based techniques, may struggle to effectively capture these dynamics. Meanwhile, while neural operators offer promise for short-term predictions, they tend to suffer from long-range instabilities. The StFT model aims to address these limitations by balancing efficiency, accuracy, and stability, which is essential for creating reliable digital twins for design, optimization, and operational control in fields such as fusion energy.

Breakdown 

• Multi-Scale Architecture: Developed a structured hierarchy of StFT blocks, each dedicated to learning dynamics at specific spatial scales. This design explicitly captures interactions across macro- and micro-scales, enabling the model to effectively learn how large, slowly varying patterns influence small, rapidly changing structures. This comprehensive approach enhances the model’s predictive capacity, allowing it to accurately capture complex dynamic behaviors over long timeframes.
  
• Dual-Path Representation: Incorporated a dual-path architecture that integrates frequency-domain (Fourier) and spatio-temporal representations, providing a richer and more complete view of system dynamics than existing neural operator models.
  
• Uncertainty-Aware Forecasting: Introduced a generative residual correction mechanism, implemented through flow matching, to iteratively refine predictions and quantify intrinsic uncertainties. This improves both accuracy and stability in long-term, autoregressive forecasts.
  
• Smooth Spatial Representation: Designed an overlapping tokenizer-detokenizer scheme that shares boundaries between adjacent spatial regions, improving spatial continuity and reducing artifacts — unwanted irregularities that can occur at the edges of the segmented areas — in the model’s outputs.
  
• Evaluation and Results: Tested on benchmark datasets in plasma, fluid, and atmospheric dynamics, StFT consistently outperformed state-of-the-art machine learning methods in accuracy, stability, and uncertainty calibration.
  
• Future Improvements: Planned enhancements include parallelizing computations across multi-scale blocks and extending the method to handle irregular geometries, broadening its applicability to more complex scientific and engineering problems.

Co-authors

Da Long, Shandian Zhe, Samuel Williams, Leonid Oliker, Zhe Bai

Publications

Spatio-temporal Fourier Transformer (StFT) for Long-term Dynamics Prediction

Funding

Department of Energy, Scientific Discovery Through Advanced Computing (SciDAC) program