Modeling Charged Multilayer Devices with Defects Using JAX

Multilayer van der Waals devices, such as twisted WSe₂ bilayers, have emerged as versatile platforms for realizing exotic correlated states, including atomic and molecular Wigner crystals and excitonic crystals. While integer lattice fillings have been extensively studied through both experiments and theory, the role of defects remains less explored. A major challenge lies in the lack of efficient and user-friendly computational tools capable of handling long-range Coulomb interactions in such systems.

In this work, we introduce an open-source computational framework built on the JAX automatic differentiation and just-in-time compilation ecosystem to model charged multilayer systems with vacancies and interstitials. The tool enables flexible simulations of multilayer electron configurations within variable simulation cells and supports optimization of the structural degrees of freedom. We benchmark the efficiency and robustness of various optimization algorithms from the machine learning community, providing insights that help bridge the gap between the two disciplines. Our results demonstrate that this approach significantly streamlines the study of defect-driven physics in 2D materials, providing an accessible platform to explore future multilayer quantum devices.