Predicting Wave Reflection and Transmission in Heterogeneous Media via Fourier Operator-Based Transformer Modeling

Zhe Bai; Hans Johansen

cs.LGMar 31, 2026

Predicting Wave Reflection and Transmission in Heterogeneous Media via Fourier Operator-Based Transformer Modeling

Authors: Zhe Bai, Hans Johansen

Organizations: AMCR Division, Lawrence Berkeley National Lab, Berkeley, USA

Abstract

We develop a machine learning (ML) surrogate model to approximate solutions to Maxwell's equations in one dimension, focusing on scenarios involving a material interface that reflects and transmits electro-magnetic waves. Derived from high-fidelity Finite Volume (FV) simulations, our training data includes variations of the initial conditions, as well as variations in one material's speed of light, allowing for the model to learn a range of wave-material interaction behaviors. The ML model autoregressively learns both the physical and frequency embeddings in a vision transformer-based framework. By incorporating Fourier transforms in the latent space, the wave number spectra of the solutions aligns closely with the simulation data. Prediction errors exhibit an approximately linear growth over time with a sharp increase at the material interface. Test results show that the ML solution has adequate relative errors below $10\%$ in over $75$ time step rollouts, despite the presence of the discontinuity and unknown material properties.

Figures & tables

Figure 1: Autoregressive modeling for forecasting long-range dynamics with a series of time snapshots of

\tilde{s}_{t}=[s_{t},s_{t-1},\dots,s_{t-m+1}]

at each step rolling out for prediction of

s_{t+1}

Figure 2: Dual-path architecture: Fourier transformer. Overlapping tokenizer is applied prior to the transformations into embeddings of $f_{t}$ and $e_{t}$ . After the two embeddings are merged, an overlapping detokenizer is employed to reconstruct the spatial field, mapping the latent representation back to the original physical domain.

Panel of Figure 3 — Figure 3: Comparison between the simulated (top), ML predicted (middle), and residual (bottom) results for the electric field in Case 2, at $t=50$ (left column), and $t=100$ (right column).

Hidden dimension	$[256,512]$
Patch size	$[33,66,132]$
Network depth	$[6,10,12]$
Overlap tokenizer	$[0,1]$

Table I: Hyperparameter search range

Figure 4: Simulated and predicted solutions for Case 1, at two time steps $t=100$ and $t=200$ . The dotted line at $x_{j}=0.5$ is the material interface where the speed of light changes.

Figure 5: Fourier domain content of the predicted solution for Case 1, at step $t=100$ (left), when the wave packet is crossing the material boundary, and $t=200$ (right) when it has moved away after reflecting/transmitting.

Figure 6: Case 1 predicted solution relative error over time steps $t=[0,200]$ , with the mean value (dark red line) and standard deviation (shaded region) across all test cases. Time step $t=100$ is approximately when the wave pulse is crossing the material boundary.

Figure 7: Simulated and predicted solutions for Case 2, with the initial condition and two time steps $t=50$ and $t=100$ . The material to the right of the interface (dotted line) has a random speed of light that must be learned from training data.

Figure 8: Fourier domain content of the predicted solution for Case 2, at step $t=50$ (left) and $t=100$ (right), as it is moving away from the interface after reflecting/transmitting.

Figure 9: Case 2 predicted solution relative error for $t\in[0,100]$ , mean value (dark red line) and standard deviation (shaded region) across all test data. The error grows approximately linearly in time.

Problem	Model(s)	$L_{2}$ error
Case 1	${f}_{t}+{e}_{t}$	0.0506
Case 1	${e_{t}}$	0.3945
Case 2	${f_{t}}+{e_{t}}$	0.0827
Case 2	${e_{t}}$	0.1819

Table II: Effect of the frequency path.

Problem	Setting(s)	$L_{2}$ error
Case 1	$r=4$	0.0825
	$r=8$	0.0506
	$r=16$	0.4795
Case 2	$r=4$	0.1191
	$r=8$	0.0827
	$r=16$	0.3099

Table III: Ablation results on the Fourier mode truncation.

PDE Surrogate Modeling

PDE surrogate modeling builds fast approximations of partial differential equation solutions for simulation, prediction, and design.

See all 322 papers

Learning Metamaterial Eigenmodes with Wavelet-Encoded Fourier Neural Operators
Sep 8, 2026Han Zhang, Alexander Ogren, Cynthia Rudin +2Fourier Neural Operator PDE Surrogate Modeling
Department of Mechanical Engineering and Materials Science, Duke University, Durham, NC, USA · Department of Mechanical Engineering, California Institute of Technology, Pasadena, CA, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA
Neural Surrogate Modeling of Ultrasonic Wave Propagation in Heterogeneous Media
Oct 9, 2026Burla Nur Korkmaz, Petro Moilanen, Ari Salmi +2Neural Surrogate Modeling PDE Surrogate Modeling
Department of Computer Science, University of Helsinki, Finland
Windowed Fourier Propagator: A Frequency-Local Neural Operator for Wave Equations in Inhomogeneous Media
Mar 15, 2026Yiyang Cai, Zixuan Qiu, Yunlu Shu +4PDE Operator Learning Neural Operators
School of Mathematical Sciences, Fudan University, Shanghai, 200433, China · Shanghai Center for Mathematical Sciences, Fudan University, Shanghai, 200438, China · Center for Applied Mathematics, Fudan University, Shanghai, 200438, China

Sep 8, 2026J/K move · Enter open · S save

Learning Metamaterial Eigenmodes with Wavelet-Encoded Fourier Neural Operators

Han Zhang, Alexander Ogren, Cynthia Rudin, Johann Guilleminot+1

Department of Mechanical Engineering and Materials Science, Duke University, Durham, NC, USA · Department of Mechanical Engineering, California Institute of Technology, Pasadena, CA, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA

Fourier Neural Operator PDE Surrogate Modeling

Open paper AskPDF arXiv

Machine learning surrogates based on neural operators have shown broad applicability in solving forward PDE problems. However, eigenvalue problems, in which an eigenparameter and one of several valid eigenmodes must be simultaneously solved, remain difficult because standard operator learning formulations assume a unique input-output map. This work demonstrates that Fourier Neural Operators (FNOs), combined with wavelet-based encodings of PDE inputs, can learn and predict multiple eigenmodes of the elastic wave equation, corresponding to deformation modes of acoustic waves propagating through arbitrary metamaterial geometries. We provide a mechanistic explanation and experimental evidence for why wavelet encodings are well matched to the dual spatial-spectral structure of the FNO, enabling deterministic mode selection on both continuous-valued and binary-valued geometries within a single model, and for why prediction accuracy varies with geometric discontinuities. For metamaterial design, the resulting surrogate accelerates the simulation stage of the design cycle by three orders of magnitude relative to finite element analysis on a consumer-grade CPU, while preserving high fidelity. These results also carry broader implications for designing input encodings in other multi-mode PDE solvers based on spectral neural operators.

Abstract

PDE Surrogate Modeling

Explore similar work

Learning Metamaterial Eigenmodes with Wavelet-Encoded Fourier Neural Operators

Neural Surrogate Modeling of Ultrasonic Wave Propagation in Heterogeneous Media

Windowed Fourier Propagator: A Frequency-Local Neural Operator for Wave Equations in Inhomogeneous Media