-
Neural General Circulation Models for Weather and Climate
Authors:
Dmitrii Kochkov,
Janni Yuval,
Ian Langmore,
Peter Norgaard,
Jamie Smith,
Griffin Mooers,
Milan Klöwer,
James Lottes,
Stephan Rasp,
Peter Düben,
Sam Hatfield,
Peter Battaglia,
Alvaro Sanchez-Gonzalez,
Matthew Willson,
Michael P. Brenner,
Stephan Hoyer
Abstract:
General circulation models (GCMs) are the foundation of weather and climate prediction. GCMs are physics-based simulators which combine a numerical solver for large-scale dynamics with tuned representations for small-scale processes such as cloud formation. Recently, machine learning (ML) models trained on reanalysis data achieved comparable or better skill than GCMs for deterministic weather fore…
▽ More
General circulation models (GCMs) are the foundation of weather and climate prediction. GCMs are physics-based simulators which combine a numerical solver for large-scale dynamics with tuned representations for small-scale processes such as cloud formation. Recently, machine learning (ML) models trained on reanalysis data achieved comparable or better skill than GCMs for deterministic weather forecasting. However, these models have not demonstrated improved ensemble forecasts, or shown sufficient stability for long-term weather and climate simulations. Here we present the first GCM that combines a differentiable solver for atmospheric dynamics with ML components, and show that it can generate forecasts of deterministic weather, ensemble weather and climate on par with the best ML and physics-based methods. NeuralGCM is competitive with ML models for 1-10 day forecasts, and with the European Centre for Medium-Range Weather Forecasts ensemble prediction for 1-15 day forecasts. With prescribed sea surface temperature, NeuralGCM can accurately track climate metrics such as global mean temperature for multiple decades, and climate forecasts with 140 km resolution exhibit emergent phenomena such as realistic frequency and trajectories of tropical cyclones. For both weather and climate, our approach offers orders of magnitude computational savings over conventional GCMs. Our results show that end-to-end deep learning is compatible with tasks performed by conventional GCMs, and can enhance the large-scale physical simulations that are essential for understanding and predicting the Earth system.
△ Less
Submitted 7 March, 2024; v1 submitted 13 November, 2023;
originally announced November 2023.
-
Learning to correct spectral methods for simulating turbulent flows
Authors:
Gideon Dresdner,
Dmitrii Kochkov,
Peter Norgaard,
Leonardo Zepeda-Núñez,
Jamie A. Smith,
Michael P. Brenner,
Stephan Hoyer
Abstract:
Despite their ubiquity throughout science and engineering, only a handful of partial differential equations (PDEs) have analytical, or closed-form solutions. This motivates a vast amount of classical work on numerical simulation of PDEs and more recently, a whirlwind of research into data-driven techniques leveraging machine learning (ML). A recent line of work indicates that a hybrid of classical…
▽ More
Despite their ubiquity throughout science and engineering, only a handful of partial differential equations (PDEs) have analytical, or closed-form solutions. This motivates a vast amount of classical work on numerical simulation of PDEs and more recently, a whirlwind of research into data-driven techniques leveraging machine learning (ML). A recent line of work indicates that a hybrid of classical numerical techniques and machine learning can offer significant improvements over either approach alone. In this work, we show that the choice of the numerical scheme is crucial when incorporating physics-based priors. We build upon Fourier-based spectral methods, which are known to be more efficient than other numerical schemes for simulating PDEs with smooth and periodic solutions. Specifically, we develop ML-augmented spectral solvers for three common PDEs of fluid dynamics. Our models are more accurate (2-4x) than standard spectral solvers at the same resolution but have longer overall runtimes (~2x), due to the additional runtime cost of the neural network component. We also demonstrate a handful of key design principles for combining machine learning and numerical methods for solving PDEs.
△ Less
Submitted 25 June, 2023; v1 submitted 1 July, 2022;
originally announced July 2022.
-
Learned Coarse Models for Efficient Turbulence Simulation
Authors:
Kimberly Stachenfeld,
Drummond B. Fielding,
Dmitrii Kochkov,
Miles Cranmer,
Tobias Pfaff,
Jonathan Godwin,
Can Cui,
Shirley Ho,
Peter Battaglia,
Alvaro Sanchez-Gonzalez
Abstract:
Turbulence simulation with classical numerical solvers requires high-resolution grids to accurately resolve dynamics. Here we train learned simulators at low spatial and temporal resolutions to capture turbulent dynamics generated at high resolution. We show that our proposed model can simulate turbulent dynamics more accurately than classical numerical solvers at the comparably low resolutions ac…
▽ More
Turbulence simulation with classical numerical solvers requires high-resolution grids to accurately resolve dynamics. Here we train learned simulators at low spatial and temporal resolutions to capture turbulent dynamics generated at high resolution. We show that our proposed model can simulate turbulent dynamics more accurately than classical numerical solvers at the comparably low resolutions across various scientifically relevant metrics. Our model is trained end-to-end from data and is capable of learning a range of challenging chaotic and turbulent dynamics at low resolution, including trajectories generated by the state-of-the-art Athena++ engine. We show that our simpler, general-purpose architecture outperforms various more specialized, turbulence-specific architectures from the learned turbulence simulation literature. In general, we see that learned simulators yield unstable trajectories; however, we show that tuning training noise and temporal downsampling solves this problem. We also find that while generalization beyond the training distribution is a challenge for learned models, training noise, added loss constraints, and dataset augmentation can help. Broadly, we conclude that our learned simulator outperforms traditional solvers run on coarser grids, and emphasize that simple design choices can offer stability and robust generalization.
△ Less
Submitted 22 April, 2022; v1 submitted 30 December, 2021;
originally announced December 2021.
-
Learning ground states of quantum Hamiltonians with graph networks
Authors:
Dmitrii Kochkov,
Tobias Pfaff,
Alvaro Sanchez-Gonzalez,
Peter Battaglia,
Bryan K. Clark
Abstract:
Solving for the lowest energy eigenstate of the many-body Schrodinger equation is a cornerstone problem that hinders understanding of a variety of quantum phenomena. The difficulty arises from the exponential nature of the Hilbert space which casts the governing equations as an eigenvalue problem of exponentially large, structured matrices. Variational methods approach this problem by searching fo…
▽ More
Solving for the lowest energy eigenstate of the many-body Schrodinger equation is a cornerstone problem that hinders understanding of a variety of quantum phenomena. The difficulty arises from the exponential nature of the Hilbert space which casts the governing equations as an eigenvalue problem of exponentially large, structured matrices. Variational methods approach this problem by searching for the best approximation within a lower-dimensional variational manifold. In this work we use graph neural networks to define a structured variational manifold and optimize its parameters to find high quality approximations of the lowest energy solutions on a diverse set of Heisenberg Hamiltonians. Using graph networks we learn distributed representations that by construction respect underlying physical symmetries of the problem and generalize to problems of larger size. Our approach achieves state-of-the-art results on a set of quantum many-body benchmark problems and works well on problems whose solutions are not positive-definite. The discussed techniques hold promise of being a useful tool for studying quantum many-body systems and providing insights into optimization and implicit modeling of exponentially-sized objects.
△ Less
Submitted 12 October, 2021;
originally announced October 2021.
-
Variational Data Assimilation with a Learned Inverse Observation Operator
Authors:
Thomas Frerix,
Dmitrii Kochkov,
Jamie A. Smith,
Daniel Cremers,
Michael P. Brenner,
Stephan Hoyer
Abstract:
Variational data assimilation optimizes for an initial state of a dynamical system such that its evolution fits observational data. The physical model can subsequently be evolved into the future to make predictions. This principle is a cornerstone of large scale forecasting applications such as numerical weather prediction. As such, it is implemented in current operational systems of weather forec…
▽ More
Variational data assimilation optimizes for an initial state of a dynamical system such that its evolution fits observational data. The physical model can subsequently be evolved into the future to make predictions. This principle is a cornerstone of large scale forecasting applications such as numerical weather prediction. As such, it is implemented in current operational systems of weather forecasting agencies across the globe. However, finding a good initial state poses a difficult optimization problem in part due to the non-invertible relationship between physical states and their corresponding observations. We learn a mapping from observational data to physical states and show how it can be used to improve optimizability. We employ this mapping in two ways: to better initialize the non-convex optimization problem, and to reformulate the objective function in better behaved physics space instead of observation space. Our experimental results for the Lorenz96 model and a two-dimensional turbulent fluid flow demonstrate that this procedure significantly improves forecast quality for chaotic systems.
△ Less
Submitted 20 May, 2021; v1 submitted 22 February, 2021;
originally announced February 2021.
-
Machine learning accelerated computational fluid dynamics
Authors:
Dmitrii Kochkov,
Jamie A. Smith,
Ayya Alieva,
Qing Wang,
Michael P. Brenner,
Stephan Hoyer
Abstract:
Numerical simulation of fluids plays an essential role in modeling many physical phenomena, such as weather, climate, aerodynamics and plasma physics. Fluids are well described by the Navier-Stokes equations, but solving these equations at scale remains daunting, limited by the computational cost of resolving the smallest spatiotemporal features. This leads to unfavorable trade-offs between accura…
▽ More
Numerical simulation of fluids plays an essential role in modeling many physical phenomena, such as weather, climate, aerodynamics and plasma physics. Fluids are well described by the Navier-Stokes equations, but solving these equations at scale remains daunting, limited by the computational cost of resolving the smallest spatiotemporal features. This leads to unfavorable trade-offs between accuracy and tractability. Here we use end-to-end deep learning to improve approximations inside computational fluid dynamics for modeling two-dimensional turbulent flows. For both direct numerical simulation of turbulence and large eddy simulation, our results are as accurate as baseline solvers with 8-10x finer resolution in each spatial dimension, resulting in 40-80x fold computational speedups. Our method remains stable during long simulations, and generalizes to forcing functions and Reynolds numbers outside of the flows where it is trained, in contrast to black box machine learning approaches. Our approach exemplifies how scientific computing can leverage machine learning and hardware accelerators to improve simulations without sacrificing accuracy or generalization.
△ Less
Submitted 28 January, 2021;
originally announced February 2021.
-
Learned discretizations for passive scalar advection in a 2-D turbulent flow
Authors:
Jiawei Zhuang,
Dmitrii Kochkov,
Yohai Bar-Sinai,
Michael P. Brenner,
Stephan Hoyer
Abstract:
The computational cost of fluid simulations increases rapidly with grid resolution. This has given a hard limit on the ability of simulations to accurately resolve small scale features of complex flows. Here we use a machine learning approach to learn a numerical discretization that retains high accuracy even when the solution is under-resolved with classical methods. We apply this approach to pas…
▽ More
The computational cost of fluid simulations increases rapidly with grid resolution. This has given a hard limit on the ability of simulations to accurately resolve small scale features of complex flows. Here we use a machine learning approach to learn a numerical discretization that retains high accuracy even when the solution is under-resolved with classical methods. We apply this approach to passive scalar advection in a two-dimensional turbulent flow. The method maintains the same accuracy as traditional high-order flux-limited advection solvers, while using 4x lower grid resolution in each dimension. The machine learning component is tightly integrated with traditional finite-volume schemes and can be trained via an end-to-end differentiable programming framework. The solver can achieve near-peak hardware utilization on CPUs and accelerators via convolutional filters. Code is available at https://github.com/google-research/data-driven-pdes.
△ Less
Submitted 5 November, 2020; v1 submitted 11 April, 2020;
originally announced April 2020.
-
Variational optimization in the AI era: Computational Graph States and Supervised Wave-function Optimization
Authors:
Dmitrii Kochkov,
Bryan K. Clark
Abstract:
Representing a target quantum state by a compact, efficient variational wave-function is an important approach to the quantum many-body problem. In this approach, the main challenges include the design of a suitable variational ansatz and optimization of its parameters. In this work, we address both of these challenges. First, we define the variational class of Computational Graph States (CGS) whi…
▽ More
Representing a target quantum state by a compact, efficient variational wave-function is an important approach to the quantum many-body problem. In this approach, the main challenges include the design of a suitable variational ansatz and optimization of its parameters. In this work, we address both of these challenges. First, we define the variational class of Computational Graph States (CGS) which gives a uniform framework for describing all computable variational ansatz. Secondly, we develop a novel optimization scheme, supervised wave-function optimization (SWO), which systematically improves the optimized wave-function by drawing on ideas from supervised learning. While SWO can be used independently of CGS, utilizing them together provides a flexible framework for the rapid design, prototyping and optimization of variational wave-functions. We demonstrate CGS and SWO by optimizing for the ground state wave-function of 1D and 2D Heisenberg models on nine different variational architectures including architectures not previously used to represent quantum many-body wave-functions and find they are energetically competitive to other approaches. One interesting application of this architectural exploration is that we show that fully convolution neural network wave-functions can be optimized for one system size and, using identical parameters, produce accurate energies for a range of system sizes. We expect these methods to increase the rate of discovery of novel variational ansatz and bring further insights to the quantum many body problem.
△ Less
Submitted 29 November, 2018;
originally announced November 2018.
-
Classical phase diagram of the stuffed honeycomb lattice
Authors:
Jyotisman Sahoo,
Dmitrii Kochkov,
Bryan K. Clark,
Rebecca Flint
Abstract:
We investigate the classical phase diagram of the stuffed honeycomb Heisenberg lattice, which consists of a honeycomb lattice with a superimposed triangular lattice formed by sites at the center of each hexagon. This lattice encompasses and interpolates between the honeycomb, triangular and dice lattices, preserving the hexagonal symmetry while expanding the phase space for potential spin liquids.…
▽ More
We investigate the classical phase diagram of the stuffed honeycomb Heisenberg lattice, which consists of a honeycomb lattice with a superimposed triangular lattice formed by sites at the center of each hexagon. This lattice encompasses and interpolates between the honeycomb, triangular and dice lattices, preserving the hexagonal symmetry while expanding the phase space for potential spin liquids. We use a combination of iterative minimization, classical Monte Carlo and analytical techniques to determine the complete ground state phase diagram. It is quite rich, with a variety of non-coplanar and non-collinear phases not found in the previously studied limits. In particular, our analysis reveals the triangular lattice critical point to be a multicritical point with two new phases vanishing via second order transitions at the critical point. We analyze these phases within linear spin wave theory and discuss consequences for the S = 1/2 spin liquid.
△ Less
Submitted 20 September, 2018; v1 submitted 24 May, 2018;
originally announced May 2018.
-
The mother of all states of the kagome quantum antiferromagnet
Authors:
Hitesh J. Changlani,
Dmitrii Kochkov,
Krishna Kumar,
Bryan K. Clark,
Eduardo Fradkin
Abstract:
Frustrated quantum magnets are a central theme in condensed matter physics due to the richness of their phase diagrams. They support a panoply of phases including various ordered states and topological phases. Yet, this problem has defied a solution for a long time due to the lack of controlled approximations which make it difficult to distinguish between competing phases. Here we report the disco…
▽ More
Frustrated quantum magnets are a central theme in condensed matter physics due to the richness of their phase diagrams. They support a panoply of phases including various ordered states and topological phases. Yet, this problem has defied a solution for a long time due to the lack of controlled approximations which make it difficult to distinguish between competing phases. Here we report the discovery of a special quantum macroscopically degenerate point in the XXZ model on the spin-1/2 kagome quantum antiferromagnet for the ratio of Ising to antiferromagnetic transverse coupling Jz/J=-1/2. This point is proximate to many competing phases explaining the source of the complexity of the phase diagram. We identify five phases near this point including both spin-liquid and broken-symmetry phases and give evidence that the kagome Heisenberg antiferromagnet is close to a transition between two phases.
△ Less
Submitted 15 March, 2018; v1 submitted 14 March, 2017;
originally announced March 2017.
-
Deviation from the dipole-ice model in the new spinel spin-ice candidate, MgEr$_2$Se$_4$
Authors:
D. Reig-i-Plessis,
S. V. Geldern,
A. A. Aczel,
D. Kochkov,
B. K. Clark,
G. J. MacDougall
Abstract:
In spin ice research, small variations in structure or interactions drive a multitude of different behaviors, yet the collection of known materials relies heavily on the `227' pyrochlore structure. Here, we present thermodynamic, structural and inelastic neutron scattering data on a new spin-ice material, MgEr$_2$Se$_4$, which contributes to the relatively under-explored family of rare-earth spine…
▽ More
In spin ice research, small variations in structure or interactions drive a multitude of different behaviors, yet the collection of known materials relies heavily on the `227' pyrochlore structure. Here, we present thermodynamic, structural and inelastic neutron scattering data on a new spin-ice material, MgEr$_2$Se$_4$, which contributes to the relatively under-explored family of rare-earth spinel chalcogenides. X-ray and neutron diffraction confirm a normal spinel structure, and places Er$^{3+}$ moments on an ideal pyrochlore sublattice. Measurement of crystal electric field excitations with inelastic neutron scattering confirms that the moments have perfect Ising character, and further identifies the ground state Kramers doublet as having dipolar-octupolar form with a significant multipolar character. Heat capacity and magnetic neutron diffuse scattering have ice-like features, but are inconsistent with Monte Carlo simulations of the nearest-neighbor and next-nearest-neighbor dipolar spin-ice (DSI) models. A significant remnant entropy is observed as $T\rightarrow0$~K, but again falls short of the full Pauling expectation for DSI, unless significant disorder is added. We show that these observations are fully in-line with what is recently reported for CdEr$_2$Se$_4$, and point to the importance of quantum fluctuations in these materials.
△ Less
Submitted 29 April, 2019; v1 submitted 13 March, 2017;
originally announced March 2017.
