大數據文摘出品
來源：tensorflowblog
編譯：睡不著的iris、劉俊寰
去年6月，俄羅斯天文學家兼碼農小哥開發了一款名叫Space Engine的宇宙模擬遊戲，目前，該遊戲已經收到2639份測評，絕大部分都是好評。
不過，Space Engine還只是在遊戲層面對宇宙的一種呈現，試想以下，如果用數值模擬的基礎工具，從觀測的海量宇宙數據進行分析……
停！看到這兒也大概能猜到，這樣的模擬計算肯定代價不菲，而且通常需要在大規模的超級算計算機上離線運行。
但是，若是提升模擬計算速度，將其與機器學習組件整合到統一框架內，會有什麼不同呢？
這便是新的N體宇宙模擬代碼FlowPM在做的。
沒錯，文摘菌今天給大家展示的就是，如何使用TensorFlow模擬自己的小宇宙，並說明為什麼它對宇宙學有著激動人心的前景展望，搓搓小手跟進了~
圖1：（藍色）通過2dFGRS測量的宇宙結構，（紅色）使用N體模擬生成的千禧年宇宙結構。
N體宇宙模擬
簡單來說，N體模擬的方法是為了解釋宇宙演化中的非線性演化問題，科學家們提出的一套研究宇宙學的數值模擬方法。
在宇宙學領域，N體模擬還原了宇宙是如何從誕生之初的大爆炸演變至今的。
如果我們按照物質初始分布，將大量的粒子放置在一個盒子中，由於引力的牽引，粒子會不斷移動。而最終的粒子狀態，就是對浩瀚宇宙結構的再現。例如，從這個分布中，你可以看到星系和星系群是如何分布的。
這些結構的形成和演化取決於一些與宇宙相關的基礎問題，諸如宇宙中有多少物質？宇宙膨脹有多快？導致宇宙膨脹的黑物質和暗能量有哪些特性？
現代的宇宙學研究用強大的望遠鏡覆蓋了遼闊的觀測區域，繪製出了這些結構，時間甚至可以追溯到幾十億年前。通過將N體模擬的預測與觀測數據進行比較，可以幫助我們解答前述基礎問題，從而提高對宇宙的誕生和演化的理解。
FlowPM：一個TensorFlow的宇宙N體求解程序
隨著幾十年研究的更迭，N體模擬已是宇宙學界的主要研究工具。
得益於機器學習和統計推理領域的最新進展，一個完全基於TensorFlow實現的N體宇宙模擬程序FlowPM應運而生了。
在這裡，我們使用Google Colab notebook對其進行實驗，同時，以下主要提供TF1的編寫代碼。
首先需要設置好邊長、網格大小和執行代碼片段等模擬參數，隨後的執行模擬代碼片段就非常簡單了，完整的代碼如下所示，如此就能生成下圖所示的大型結構。
# Generate Gaussian initial conditions for the matter distributioninitial_conditions = flowpm.linear_field(N, L, ipklin, batch_size=batch)# Sample particles, i.e. generate the initial displacement and velocitystate = flowpm.lpt_init(initial_conditions, a0=a0)# Evolve particles from initial state down to present time with N-Body simulationsfinal_state = flowpm.nbody(state, stages, N)# Visualize final density field i.e interpolate the particles to a gridFinal_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])
圖2 ：（左）N體模擬啟動時，宇宙中物質的初始分布；（右）在模擬的最終快照，物質的最終分布。可以清楚地看到大型結構，具有暈圈，長絲和空洞塌陷。
或許不少人會疑惑，TensorFlow真有這麼香？文摘菌把話撂這，還這真有這麼香，TensorFlow的優勢主要體現在以下兩個方面。
首先是分析和推理，TensorFlow模擬為宇宙學家提供了前所未有的獨特功能，也就是我們所說的可微性（differentiability）。這個功能為科學家提供一種新的分析工具，例如開發基於模擬的有效推理技術。
不僅如此，這個功能還支持量化各類輸入參數的調整對最終觀察的影響。在伯克利宇宙中心，研究者們對時光倒流和重構初始狀態十分感興趣，這也有助於更好地了解現在的宇宙結構，但這是有前提的，必須先解決一個百萬維度的高度非線性優化問題，即模擬過程得可微，而這，FlowPM就能做到，重構後的圖如下所示。
從最終的黑物質的大型宇宙結構（右邊）如何重建宇宙初始狀態（左邊）
另一個就是混合物理/深度學習模擬，在N體模擬運用深度學習組件時，你會需要開發一個混合前向模型。但是，利用當前的計算資源，不可能高度準確地同時模擬所有的組件，只能依據科研目標，在不同元素之間進行取捨，例如宇宙可觀察的深度、星系可見的範圍和形成星系的多維物理過程等。
然而現在，我們可以用深度學習替代模型，將這些元素自然地整合進N體模擬中。在BCCP，科學家們開發了一個這樣的混合模擬，如下圖所示。假如我們要模擬宇宙中的氣體動量密度，目前，這類觀察採用混合模擬實現代價非常大，但使用一個端到端的FlowPM模型，成本至少可以降低1000倍。
一個混合模擬的示例，用兩層的非線性變換（網絡）補充PM模擬的暗物質輸出（左），以模擬宇宙（中心）的氣體動量密度。將使用昂貴1000倍的混合模擬所模擬與真相進行比較（右）
採用Mesh TensorFlow分布式框架仿真
雖然原理看上去很簡單，但實際的挑戰還是很大的，畢竟宇宙真的很龐大！
為了準確匹配當前和未來的觀測，我們要同時模擬數以億計的粒子的變化。
要實現這個，N體宇宙模擬主要有以下兩個難點：
連續演化數十億個粒子，計算成本非常高。要估算所有粒子之間的重力，需要計算所有的粒子對，當計算規模擴展至N2，無法得出結果。但幸運的是，可以在FlowPM中使用一種叫做particle-mesh（PM）的方法，即把一個Ng規格的網格離散空間，使用高度優化的3D快速傅立葉變換計算整個空間的力，這將計算成本從N2g降到了Nglog(Ng)。儘管算法進行了優化，但模擬數以萬計的粒子也非常佔用內存。結果是一個有效尺寸為1024x1024x1024的最小的模擬網格，都無法再單個GPU上顯示。因此，我們需要一個模型並行框架來支撐大規模模擬，這也正是Mesh TensorFlow的用武之地。
Mesh TensorFlow框架可以輕鬆地使用分布式張量來表示模擬，並持續追蹤分布式梯度和設備之間的存儲通信，通過Mesh TensorFlow編寫N體求解程序，可以將巨大的仿真量跨多個超級計算機進行分布。
在這種模擬中，每一步操作和網格組件在一個時間步長可以演化出不同的空間區域。在雲端TPU部署相同的模擬代碼，可以模擬128個128x128x128網格的獨立宇宙，或者部署在國家計算機（NERSC）的64個GPU上模擬一個1024x1024x1024的宇宙。
除了支持大規模仿真，模型並行框架允許通過將計算任務分配給多個處理器來加速中間尺寸的模擬，如下圖所示，FlowPM的平均模擬速度比當前的可微python模擬器FastPM快了40倍不止。
使用FastPM（在Cori Haswell內核上運行基於CPU的python代碼）和FlowPM（在Cori GPU上運行基於GPU的Mesh TensorFlow代碼）模擬2563網格PM，比較1個步長內時間縮放與處理器數量
展望
過去的三十多年，數值模擬是大型宇宙的研究基石。現在，有了FlowPM，研究即將邁出新的一步，將模擬器集成在包含深度學習組件的統一框架內，同時保持對潛在現象的準確物理理解。
在宇宙學中，這種結合不僅促生了新穎的分析工具，在研究體系中加入深度模型攻克棘手的歷史難題。這些活躍的研究領域，加速催生下一代宇宙研究，在未來的十年之際可以觀察宇宙中成千上萬的個體。
多虧了Mesh TensorFlow模型並行框架，使得物理建模和機器學習的融合得以實現，希望成分分析和FlowPM計算工具的發展，讓宇宙學以外的領域也得以受益。
相關報導：
https://blog.tensorflow.org/2020/03/simulating-universe-in-tensorflow.html