簡介
在宇宙學研究中，在計算機上進行 數值模擬 [1] 是個重要手段，這是宇宙學家用於分析理解宇宙測繪收集的大量數據的基本工具。此類模擬在計算上的花銷通常特別高昂，並且往往需要在大型超級計算機上離線運行。但是，假如我們可以以超快的速度執行此類模擬級，並實現與機器學習組件集成到一個統一框架內，會怎樣？這就是全新的 N 體宇宙學模擬代碼 FlowPM [2] 要完成的目標。在本文中，我們將向您展示如何在 TensorFlow 中模擬自己的微型宇宙，並展望這將為宇宙學帶來的新前景。
圖 1：（藍色）2dFGRS 測繪計劃中觀察到的宇宙結構。（紅色）千年 N 體模擬中生成的對應結構
N 體宇宙模擬
宇宙學的 N 體模擬展示出宇宙從大爆炸中誕生到至今的發展演化過程。宇宙學家根據物質的初始分佈將大量粒子分佈在一個方框內，然後模擬這些粒子隨時間推移及在引力作用下的運動，從而完成模擬。而產生的最終粒子分佈可再現宇宙的大尺度結構。例如，這些粒子可以反映當今星系和星系團的分佈狀況。
在分析這些結構的形成和演化之前，需要回答關於宇宙的一些基本問題。例如，宇宙中有多少物質？宇宙膨脹的速度有多快？驅動宇宙膨脹的暗物質和暗能量的性質是什麼？
現代宇宙學測繪使用最強大的望遠鏡在最廣闊的天空區域範圍內觀測，並追溯到數十億年前，最終繪製出這些結構。將 N 體模擬得出的預測與這些測繪的結果進行匹配，有助於我們回答上述基本問題，進而增進我們對宇宙的誕生和演化的理解。
TensorFlow 宇宙 N 體解算器：FlowPM
數十年來，N 體模擬一直是宇宙學界的主要工具。受益於機器學習和統計推斷領域的最新發展，我們引入了 FlowPM ，宇宙 N 體模擬的純 TensorFlow 實現。我們提供 Google Colab notebook 來進行此類模擬實驗。備註：目前，我們的代碼使用 TF1 編寫而成（未來我們會將研究更新至 TF2）。
- FlowPMhttps://github.com/modichirag/flowpm
- Google Colab notebookhttps://colab.research.google.com/drive/1zc1nvra1HF01mTtpOlRigj-UBVzYgFqq#scrollTo=pvtJPFS3LfYv
設置好模擬參數（如框長和網格大小）後，執行模擬的代碼片段將非常簡單，並已包含在下文的完整代碼中。此代碼片段將生成觀測到的大尺度結構，如圖 2 所示。
<code># Generate Gaussian initial conditions for the matter distribution
initial_conditions = flowpm.linear_field(N, L, ipklin, batch_size=batch)
# Sample particles, i.e. generate the initial displacement and velocity
state = flowpm.lpt_init(initial_conditions, a0=a0)
# Evolve particles from initial state down to present time with N-Body simulations
final_state = flowpm.nbody(state, stages, N)
# Visualize final density field i.e interpolate the particles to a grid
Final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions), final_state[0])
# Execute the graph!
with tf.Session() as sess:
ic, istate, fstate, sim = sess.run([initial_conditions, state, final_state, final_field])
/<code>
圖 2：（左）N 體模擬開始時，物質在宇宙中的初始分佈。（右）模擬最終快照中的物質最終分佈。圖上清晰可見具有塌縮暗暈、長纖維和空洞的大尺度結構
那麼，在 TensorFlow 中進行此類模擬有哪些好處？該框架的優點主要可分為兩大類：
1.分析和推理 ：TensorFlow 中的模擬具有宇宙學家以前沒有使用的特性 — 可微分性。該功能為科學家開啟了全新分析工具的大門，例如，開發基於模擬的高效推理技術。此外，通過此功能，我們還能夠量化針對各種輸入參數的最終觀測結果響應。在我們位於 伯克利宇宙物理中心 （BCCP，加州大學伯克利分校）的研究小組中，我們很希望根據目前觀測到的大尺度結構來進行反推 (時光倒流) 並重建 初始條件 [4] 。這就涉及到需要解決數 百萬 個維度中的高度非線性優化問題，這隻有使用具備可微分性的模擬（如 FlowPM）才能實現。下方視頻是此重建的圖示。
- 伯克利宇宙物理中心http://bccp.berkeley.edu/
所示為從最終暗物質宇宙場（右）中的大尺度結構重建宇宙初始條件（左）的動態
2.混合物理/深度學習模擬 ：該模擬支持開發混合前向模型，在該模型中，我們可將深度學習組件用作 N 體模擬的一部分。在現有計算資源條件下，無法對所有組件進行高精度模擬。根據科學研究目標的不同，當前模擬通常會在不同要素之間進行取捨：如在宇宙中觀測到的各種長度尺度，觀測到的星系質量範圍，以及形成這些星系的各種物理過程。
但是，現在，針對具體的科學要求，我們可以使用深度學習替代模型，通過與基礎 N 體模擬自然對接的方式納入這些元素，進而在這種取捨之上做出改進。我們位於 BCCP 的小組開發的此類混合模擬示例 [3] 如圖 3 所示。假設我們要模擬宇宙中的氣體動量密度。這些可觀察量當前使用代價高昂的水力模擬進行模擬，但現在可改用 FlowPM 端到端模型來生成模擬結果，成本為原來的 1/1000。
圖 3：混合模擬示例，其中我們用 2 層非線性變換（網絡）來補充 PM 模擬的暗物質輸出（左），以模擬宇宙中的氣體動量密度（中心）。將其與使用成本高達 1000 倍的水力模擬（右）所模擬的真實結果進行對比
使用 Mesh TensorFlow 進行分佈式模擬
所有這些模擬的原理看起來都很簡單，但是我們還面臨另一個挑戰；我們的宇宙 浩瀚無邊 ！為準確匹配當前和將來測繪的觀測結果，我們將需要同時模擬數十億個粒子的演化。這給 N 體宇宙學模擬帶來兩種挑戰：
1.持續模擬數十億個粒子的演化在計算上代價高昂 。另外，要估算所有粒子之間的重力，我們將需要在模擬中計算所有粒子對。這使得計算量 按 N 2 級的規模增長，因此也無法實現。幸運的是，還有一些近似算法讓計算易於駕馭。我們在 FlowPM 中採用的一種算法是 Particle-Mesh (PM) 方法 。在 PM 方法中，為估計重力，我們在大小為 N g 的規則網格上離散空間，然後使用高度優化的 3D 快 速傅立葉變換計算整個空間上的力。這可使計算成本從 N 2 g 減少至 N g log(N g ) 。
2.儘管對這些算法進行了優化，但數十億個粒子也會讓這些模擬非常佔用內存 。因此，需要網格至少為 1024x1024x1024 的有效尺寸的模擬無法在單個 GPU 上顯示。因此，我們需要 模型並行 框架以開發大尺寸模擬，而這正是 Mesh TensorFlow 的用武之地。
- Mesh TensorFlow https://github.com/tensorflow/mesh
Mesh TensorFlow 框架讓我們能輕鬆地用分佈式張量描述模擬，並在幕後持續跟蹤設備間的分佈式梯度和內存通信。通過在 Mesh TensorFlow 中編寫 N 體解算器 ，我們可以將這些超級計算機上的龐大模擬量分佈在多臺設備之間。在此種模擬中，每個處理和網格組件在每個時間步下都會演化出不同的空間區域。使用相同的模擬代碼，我們可以同時在雲 TPU 上演化 128 個網格尺寸為 128x128x128 的獨立宇宙，或者在 NERSC 等國家計算機設施的 64 個 GPU 上演化單個網格尺寸為 1024x1024x1024 的宇宙。除支持大型模擬之外，模型並行框架還允許我們通過在多個處理器之間分配計算來加速中間尺寸的模擬。如下方圖 4 所示，我們可以看到，FlowPM 模擬的平均速度是當前可微分 python 模擬 FastPM 的 40 倍。
- N 體解算器https://github.com/modichirag/flowpm/tree/mesh
圖 4：我們在 FastPM（在 Cori Haswell 內核上運行的基於 CPU 的 python 代碼）和 FlowPM（在 Cori GPU 上運行的基於 GPU 的 Mesh TensorFlow 代碼）的 2563 網格 PM 模擬中，將時間縮放與處理器數量進行 1 個時間步的比較
展望
三十多年來，我們的宇宙數值模擬已成為大尺度結構宇宙學的骨幹技術。藉助 FlowPM，我們正在邁出第一步，將此類模擬與深度學習組件集成到一個統一框架內，同時保持對基本現象的精準物理認知。在宇宙學中，這兩者的結合為開發全新分析工具以及在迄今難以解決的領域中實現建模開啟大門。這些是活躍的研究領域，隨著新一代宇宙學測繪的臨近而變得越來越緊迫，因為屆時需要觀測在十年交替時宇宙中即將顯現的上千萬個物體。藉助 Mesh TensorFlow 的模型並行框架，物理建模和機器學習相結合的可能性變大，並且我們希望使用 FlowPM 開發的組件分析和計算工具也能使使宇宙學以外的其他學科的大尺寸科學應用受益。
閱讀更多 sandag 的文章
