Research

I work on deep learning for the physical sciences, with a particular focus on weather and climate modeling and on accelerating scientific simulations. My collaborators and I have published in Nature, Science, PNAS, Physical Review Letters and NeurIPS.

For a complete list, see my Google Scholar profile (40,000+ citations, h-index 29).

AI for weather and climate

I conceived and led the Neural General Circulation Models project at Google Research. NeuralGCM is the first AI-based model to improve on traditional physics-based 15-day weather forecasts and atmosphere-only climate simulations. It combines a differentiable atmospheric dynamical core written in JAX with learned parameterizations for sub-grid physics, and is trained end-to-end on reanalysis data.

The work was published in Nature (2024) with 16 co-authors, and received press coverage in Bloomberg and MIT Technology Review, among others. The open-source model has since been used to produce state-of-the-art monsoon-onset forecasts sent to 38 million farmers in India, in collaboration with the University of Chicago and the Indian Ministry of Agriculture.

On the technical side, I personally implemented many of the key components, including model and data parallelism scaling to 256 TPUs.

AI for fluids

For several years I led a research program on accelerating fluid simulations using deep learning and Google TPUs. The program produced six peer-reviewed publications (two in PNAS) and over 2,000 citations.

AI for physics

I have published research applying machine learning to a range of problems across the physical sciences, including drug discovery, microscopy, nanophotonics, quantum chemistry, structural engineering, flood modeling and fundamental physics.

Fundamental AI research

At Google, I collaborated with researchers at DeepMind and Google Brain and published in top machine learning conferences.

Tools for scientific computing at scale

Alongside specific research projects, I have contributed to general-purpose tools for working with large scientific datasets, including Xarray, NumPy and JAX. See Software for the open source projects.

Earlier: quantum dynamics

Before joining industry I completed a Ph.D. in theoretical physics at UC Berkeley with Birgitta Whaley, working on quantum dynamics in photosynthetic light harvesting (2008–2013).