Bayesian ensembles for inferring exoplanetary atmospheres - Oxford AI4Science Lab

Team

Adam D. Cobb (University of Oxford)
Michael D. Himes (University of Central Florida)
Frank Soboczenski (Kings College London)
Simone Zorzan (Luxembourg Institute of Science and Technology)
Molly D. O’Beirne (University of Pittsburgh)
Atılım Güneş Baydin (University of Oxford)
Yarin Gal (University of Oxford)
Shawn D. Domagal-Goldman (NASA Goddard Space Flight Center)
Giada N. Arney (NASA Goddard Space Flight Center)
Daniel Angerhausen (University of Bern)

Abstract

Machine learning is now used in many areas of astrophysics, from detecting exoplanets in Kepler transit signals to removing telescope systematics. Recent work demonstrated the potential of using machine learning algorithms for atmospheric retrieval by implementing a random forest to perform retrievals in seconds that are consistent with the traditional, computationally-expensive nested-sampling retrieval method. We expand upon their approach by presenting a new machine learning model, plan-net, based on an ensemble of Bayesian neural networks that yields more accurate inferences than the random forest for the same data set of synthetic transmission spectra. We demonstrate that an ensemble provides greater accuracy and more robust uncertainties than a single model. In addition to being the first to use Bayesian neural networks for atmospheric retrieval, we also introduce a new loss function for Bayesian neural networks that learns correlations between the model outputs. Importantly, we show that designing machine learning models to explicitly incorporate domain-specific knowledge both improves performance and provides additional insight by inferring the covariance of the retrieved atmospheric parameters. We apply plan-net to the Hubble Space Telescope Wide Field Camera 3 transmission spectrum for WASP-12b and retrieve an isothermal temperature and water abundance consistent with the literature. We highlight that our method is flexible and can be expanded to higher-resolution spectra and a larger number of atmospheric parameters.

Press Releases and Media Coverage

Nov 2019: NASA Takes a Cue From Silicon Valley to Hatch Artificial Intelligence Technologies, NASA front page

Video

Publications

Cobb, Adam D., Michael D. Himes, Frank Soboczenski, Simone Zorzan, Molly D. O’Beirne, Atılım Güneş Baydin, Yarin Gal, Shawn D. Domagal-Goldman, Giada N. Arney, and Daniel Angerhausen. 2019. “An Ensemble of Bayesian Neural Networks for Exoplanetary Atmospheric Retrieval.” The Astronomical Journal 158 (1). doi:10.3847/1538-3881/ab2390.
- Abstract
- Bib
- PDF
Machine learning is now used in many areas of astrophysics, from detecting exoplanets in Kepler transit signals to removing telescope systematics. Recent work demonstrated the potential of using machine learning algorithms for atmospheric retrieval by implementing a random forest to perform retrievals in seconds that are consistent with the traditional, computationally-expensive nested-sampling retrieval method. We expand upon their approach by presenting a new machine learning model, plan-net, based on an ensemble of Bayesian neural networks that yields more accurate inferences than the random forest for the same data set of synthetic transmission spectra. We demonstrate that an ensemble provides greater accuracy and more robust uncertainties than a single model. In addition to being the first to use Bayesian neural networks for atmospheric retrieval, we also introduce a new loss function for Bayesian neural networks that learns correlations between the model outputs. Importantly, we show that designing machine learning models to explicitly incorporate domain-specific knowledge both improves performance and provides additional insight by inferring the covariance of the retrieved atmospheric parameters. We apply plan-net to the Hubble Space Telescope Wide Field Camera 3 transmission spectrum for WASP-12b and retrieve an isothermal temperature and water abundance consistent with the literature. We highlight that our method is flexible and can be expanded to higher-resolution spectra and a larger number of atmospheric parameters.
```
@article{cobb-2019-ensemble,
  title = {An Ensemble of Bayesian Neural Networks for Exoplanetary Atmospheric Retrieval},
  author = {Cobb, Adam D. and Himes, Michael D. and Soboczenski, Frank and Zorzan, Simone and O’Beirne, Molly D. and Baydin, Atılım Güneş and Gal, Yarin and Domagal-Goldman, Shawn D. and Arney, Giada N. and Angerhausen, Daniel},
  journal = {The Astronomical Journal},
  volume = {158},
  number = {1},
  year = {2019},
  doi = {10.3847/1538-3881/ab2390},
  url = {https://doi.org/10.3847/1538-3881/ab2390}
}
```

Soboczenski, Frank, Michael D. Himes, Molly D. O’Beirne, Simone Zorzan, Atılım Güneş Baydin, Adam D. Cobb, Yarin Gal, Daniel Angerhausen, Massimo Mascaro, Giada N. Arney, and Shawn D. Domagal-Goldman. 2018. “Bayesian Deep Learning for Exoplanet Atmospheric Retrieval.” In Third Workshop on Bayesian Deep Learning (NeurIPS 2018), Montreal, Canada.
- Abstract
- Bib
- PDF
Over the past decade, the study of extrasolar planets has evolved rapidly from plain detection and identification to comprehensive categorization and characterization of exoplanet systems and their atmospheres. Atmospheric retrieval, the inverse modeling technique used to determine an exoplanetary atmosphere’s temperature structure and composition from an observed spectrum, is both time-consuming and compute-intensive, requiring complex algorithms that compare thousands to millions of atmospheric models to the observational data to find the most probable values and associated uncertainties for each model parameter. For rocky, terrestrial planets, the retrieved atmospheric composition can give insight into the surface fluxes of gaseous species necessary to maintain the stability of that atmosphere, which may in turn provide insight into the geological and/or biological processes active on the planet. These atmospheres contain many molecules, some of them biosignatures, spectral fingerprints indicative of biological activity, which will become observable with the next generation of telescopes. Runtimes of traditional retrieval models scale with the number of model parameters, so as more molecular species are considered, runtimes can become prohibitively long. Recent advances in machine learning (ML) and computer vision offer new ways to reduce the time to perform a retrieval by orders of magnitude, given a sufficient data set to train with. Here we present an ML-based retrieval framework called Intelligent exoplaNet Atmospheric RetrievAl (INARA) that consists of a Bayesian deep learning model for retrieval and a data set of 3,000,000 synthetic rocky exoplanetary spectra generated using the NASA Planetary Spectrum Generator. Our work represents the first ML retrieval model for rocky, terrestrial exoplanets and the first synthetic data set of terrestrial spectra generated at this scale.
```
@inproceedings{soboczenski-2018-bayesian-exoplanet,
  title = {Bayesian Deep Learning for Exoplanet Atmospheric Retrieval},
  author = {Soboczenski, Frank and Himes, Michael D. and O'Beirne, Molly D. and Zorzan, Simone and Baydin, Atılım Güneş and Cobb, Adam D. and Gal, Yarin and Angerhausen, Daniel and Mascaro, Massimo and Arney, Giada N. and Domagal-Goldman, Shawn D.},
  booktitle = {Third workshop on Bayesian Deep Learning (NeurIPS 2018), Montreal, Canada},
  year = {2018}
}
```

Soboczenski, Frank, Michael D. Himes, Molly D. O’Beirne, Simone Zorzan, Atılım Güneş Baydin, Adam D. Cobb, Yarin Gal, Daniel Angerhausen, Massimo Mascaro, Geronimo Villanueva, Shawn D. Domagal-Goldman, and Giada N. Arney. 2020. “INARA: A Bayesian Deep Learning Framework for Exoplanet Atmospheric Retrieval.” In Second AI and Data Science Workshop for Earth and Space Sciences, Jet Propulsion Laboratory (NASA JPL), Pasadena, CA, United States, March 24–26, 2020. https://datascience.jpl.nasa.gov/aiworkshop.
- Abstract
- Bib
Determining an exoplanet’s atmospheric properties from an observed spectrum (atmospheric retrieval) is a time-consuming and compute-intensive inverse modeling technique. They require complex algorithms that generate many atmospheric models and compare their simulated spectra to the observational data to find the most probable values and associated uncertainties for each model parameter. Retrieval may be the first method to find extraterrestrial life by remotely detecting biosignatures, atmospheric species indicative of biological activity. The work presented here is a result of the NASA Frontier Development Lab Astrobiology Team II. We present an ML-based retrieval framework called Intelligent exoplaNet Atmospheric RetrievAl (INARA) that consists of a Bayesian deep learning model for retrieval and a data set of 3,000,000 synthetic rocky exoplanetary spectra generated using approximately 2,000 high-end VMs and instances of the NASA Planetary Spectrum Generator (PSG). The generated dataset encompasses spectra based on a given planetary system model, where we consider F-, G-, K-, and M-type main sequence stars. Observations are simulated using an instrument model of the Large UltraViolet/Optical/InfraRed Surveyor (LUVOIR). Our work represents the first ML retrieval framework for rocky, terrestrial exoplanets and the first synthetic data set of terrestrial spectra generated at this scale.
```
@inproceedings{soboczenski-2020-inara,
  title = {{INARA}: A {Bayesian} Deep Learning Framework for Exoplanet Atmospheric Retrieval},
  author = {Soboczenski, Frank and Himes, Michael D. and O’Beirne, Molly D. and Zorzan, Simone and Baydin, Atılım Güneş and Cobb, Adam D. and Gal, Yarin and Angerhausen, Daniel and Mascaro, Massimo and Villanueva, Geronimo and Domagal-Goldman, Shawn D. and Arney, Giada N.},
  booktitle = {Second AI and Data Science Workshop for Earth and Space Sciences, Jet Propulsion Laboratory (NASA JPL), Pasadena, CA, United States, March 24--26, 2020},
  year = {2020},
  url = {https://datascience.jpl.nasa.gov/aiworkshop}
}
```
Himes, Michael D., Adam D. Cobb, Frank Soboczenski, Simone Zorzan, Molly D. O’Beirne, Atılım Güneş Baydin, Yarin Gal, Daniel Angerhausen, Shawn D. Domagal-Goldman, and Giada N. Arney. 2020. “Machine Learning Retrieval of Jovian and Terrestrial Atmospheres.” In American Astronomical Society Meeting #235, Id. 343.01. Bulletin of the American Astronomical Society, Vol. 52, No. 1. https://ui.adsabs.harvard.edu/abs/2020AAS...23534301H/abstract.
- Abstract
- Bib
Machine learning approaches to atmospheric retrieval offer results comparable to traditional numerical approaches in just seconds, compared to hundreds of compute hours. This opens the possibility for fully-3D retrievals to execute in times comparable to traditional approaches. Recently, we developed plan-net, an ensemble of Bayesian neural networks for atmospheric retrieval; we trained plan-net on synthetic Wide Field Camera 3 (WFC3) hot-Jupiter transmission spectra, applied it to the WFC3 spectrum of WASP-12b, and found results consistent with the literature. Here, we present updates to plan-net and expand its application to our 28-parameter data set of simulated LUVOIR spectra of terrestrial exoplanets generated using the NASA Planetary Spectrum Generator. By including both dense dropout and convolutional layers, we find a significant improvement in accuracy. MH and FS acknowledge the support of NVIDIA Corporation for the donation of the Titan Xp GPUs used for this research. AC is sponsored by the AIMS-CDT and EPSRC. AGB is funded by Lawrence Berkeley National Lab and EPSRC/MURI grant EP/N019474/1.
```
@inproceedings{himes-2020-machine,
  title = {Machine Learning Retrieval of Jovian and Terrestrial Atmospheres},
  author = {Himes, Michael D. and Cobb, Adam D. and Soboczenski, Frank and Zorzan, Simone and O’Beirne, Molly D. and Baydin, Atılım Güneş and Gal, Yarin and Angerhausen, Daniel and Domagal-Goldman, Shawn D. and Arney, Giada N.},
  booktitle = {American Astronomical Society meeting \#235, id. 343.01. Bulletin of the American Astronomical Society, Vol. 52, No. 1},
  year = {2020},
  url = {https://ui.adsabs.harvard.edu/abs/2020AAS...23534301H/abstract}
}
```
Himes, M., A. Cobb, A. Baydin, F. Soboczenski, S. Zorzan, M. O’Beirne, G.N. Arney, S. Domagal-Goldman, D. Angerhausen, and Y. Gal. 2019. “Exoplanetary Atmospheric Retrieval via Bayesian Machine Learning.” In American Astronomical Society Meeting on Extreme Solar Systems IV, Reykjavik, Iceland, August 19–23, 2019. https://sites.northwestern.edu/iceland2019/.
- Abstract
- Bib
Atmospheric retrieval, the inverse modeling technique whereby atmospheric properties are inferred from observations, is computationally expensive and time consuming. Recently, machine learning (ML) approaches to atmospheric retrieval have been shown to provide results consistent with traditional approaches in just seconds to minutes. We introduce plan-net, the first ensemble of Bayesian neural networks for atmospheric retrieval. Our novel likelihood function captures parameter correlations, improving uncertainty estimations over standard likelihood functions common in ML. We replicate the results of Marquez-Neila et al. (2018), and we demonstrate plan-net’s improvement in accuracy over their random forest regression tree when applied to their synthetic data set of hot Jupiter WFC3 transmission spectra. We apply a trained plan-net ensemble to the transmission spectrum of WASP-12b and find results generally consistent with the literature. We also apply plan-net to our data set of over 3 million synthetic terrestrial exoplanet spectra generated using the NASA Planetary Spectrum Generator.
```
@inproceedings{himes-2019-exoplanetary,
  title = {Exoplanetary Atmospheric Retrieval via Bayesian Machine Learning},
  author = {Himes, M. and Cobb, A. and Baydin, A. and Soboczenski, F. and Zorzan, S. and O'Beirne, M. and Arney, G.N. and Domagal-Goldman, S. and Angerhausen, D. and Gal, Y.},
  booktitle = {American Astronomical Society Meeting on Extreme Solar Systems IV, Reykjavik, Iceland, August 19--23, 2019},
  year = {2019},
  url = {https://sites.northwestern.edu/iceland2019/}
}
```
O’Beirne, Molly D., Michael D. Himes, Frank Soboczenski, Simone Zorzan, Adam Cobb, Atılım Güneş Baydin, Yarin Gal, Daniel Angerhausen, Massimo Mascaro, Giada N. Arney, and Shawn D. Domagal-Goldman. 2019. “INARA: A Machine Learning Retrieval Framework with a Data Set of 3 Million Simulated Exoplanet Atmospheric Spectra.” In Astrobiology Science Conference (AbSciCon 2019), Bellevue, Washington, June 24–28, 2019. https://agu.confex.com/agu/abscicon19/meetingapp.cgi/Paper/481266.
- Abstract
- Bib
Traditional approaches for determining the atmospheres of exoplanets from telescopic spectral data (i.e., atmospheric retrievals) involve time-consuming and compute-intensive Bayesian sampling methods, requiring a compromise between physical and chemical realism and overall computational feasibility. For rocky, terrestrial exoplanets, the retrieved atmospheric composition can give insight into the surface fluxes of gaseous species necessary to maintain the stability of that atmosphere, which may in turn provide insight into the geological and/or biological processes active on the planet. Machine learning (ML) offers a feasible and reliable approach to expedite the process of atmospheric retrievals; however, ML models require a large data set to train on. Here we present a data set of 3,000,000 simulated atmospheric spectra of rocky, terrestrial exoplanets generated across a broad parameter space of stellar and planetary properties, including 12 molecular species relevant for determining extant life. We then introduce INARA (Intelligent exoplaNet Atmospheric RetrievAl), our ML-based atmospheric retrieval framework. In a matter of seconds, INARA is capable of retrieving accurate concentrations of 12 molecular atmospheric constituents when given an observed spectrum. Our work represents the first large-scale simulated spectral data set and first atmospheric retrieval ML model for rocky, terrestrial exoplanets.
```
@inproceedings{obeirne-2019-inara,
  title = {INARA: A Machine Learning Retrieval Framework with a Data Set of 3 Million Simulated Exoplanet Atmospheric Spectra},
  author = {O’Beirne, Molly D. and Himes, Michael D. and Soboczenski, Frank and Zorzan, Simone and Cobb, Adam and Baydin, Atılım Güneş and Gal, Yarin and Angerhausen, Daniel and Mascaro, Massimo and Arney, Giada N. and Domagal-Goldman, Shawn D.},
  booktitle = {Astrobiology Science Conference (AbSciCon 2019), Bellevue, Washington, June 24--28, 2019},
  year = {2019},
  url = {https://agu.confex.com/agu/abscicon19/meetingapp.cgi/Paper/481266}
}
```

Acknowledgments

We thank Chloe Fisher for making the data set from Marquez-Neila et al. (2018) publicly available on GitHub upon request. Adam D. Cobb is sponsored by the AIMS CDT (http://aims.robots.ox.ac.uk) and the EPSRC (https://www.epsrc.ac.uk). Frank Soboczenski gratefully acknowledges the support of NVIDIA Corporation with the donation of the Titan Xp GPU used for this research (GPU No 900-1G611-2530-000). A.G. Baydin is funded by Lawrence Berkeley National Lab and EPSRC/MURI grant EP/N019474/1. We thank NASA FDL (http://www.frontierdevelopmentlab.org/) and SETI (https://www.seti.org) for making this collaboration possible.

NEXTEtalumis: Scientific simulators as probabilistic programs