基于Dask并行加速的射電干涉成像網(wǎng)格化方法實(shí)現(xiàn)*

李?yuàn)檴?，駱開達(dá)，衛(wèi)守林,2，戴 偉,2，梁 波,2

(1. 昆明理工大學(xué)信息工程與自動(dòng)化學(xué)院，云南 昆明 650500; 2. 云南省計(jì)算機(jī)應(yīng)用技術(shù)重點(diǎn)實(shí)驗(yàn)室，云南 昆明 650500)

射電干涉陣列得到的是非均勻采樣的可見度數(shù)據(jù)，在應(yīng)用快速傅里葉變換對可見度數(shù)據(jù)進(jìn)行成像前，需使用網(wǎng)格化方法將采樣數(shù)據(jù)重采樣到一個(gè)均勻劃分的網(wǎng)格上。當(dāng)前網(wǎng)格化主要使用基于卷積的方法，卷積網(wǎng)格化過程的實(shí)質(zhì)是矩陣相乘，當(dāng)數(shù)據(jù)量較大時(shí)，網(wǎng)格化計(jì)算非常耗時(shí)。

近年來，天文學(xué)家在提高可見度數(shù)據(jù)網(wǎng)格化算法性能方面做了很多研究。其中W-projection算法是目前廣泛使用的網(wǎng)格化方法，由于該算法僅校準(zhǔn)W項(xiàng)，沒有校準(zhǔn)方向相關(guān)效應(yīng)的A項(xiàng)，當(dāng)天線彼此相距較遠(yuǎn)，W項(xiàng)的尺寸可能變得很大，算法效率低下且內(nèi)存占用率高[1]。通過將每個(gè)可見度數(shù)據(jù)的w值投影到鄰近w平面的W-stacking算法，可以顯著提高網(wǎng)格化性能，但是需要耗費(fèi)額外的內(nèi)存[2]。如果考慮方向相關(guān)效應(yīng)，網(wǎng)格化的計(jì)算難度將進(jìn)一步增加，同時(shí)修正方向相關(guān)效應(yīng)A項(xiàng)和W項(xiàng)被稱為AW-projection網(wǎng)格化算法[3]。在數(shù)值分析領(lǐng)域，文[4]提出基于 “半圓指數(shù)” 卷積核的非均勻快速傅里葉變換庫(Non-uniform Fast Fourier Transform, NUFFT)，將快速傅里葉變換推廣到離散化的網(wǎng)格數(shù)據(jù)中。首次將非均勻傅里葉變換應(yīng)用到射電天文中的Nifty-gridder算法，采用共享內(nèi)存和多線程技術(shù)，進(jìn)一步優(yōu)化W-stacking算法。

綜上所述，網(wǎng)格化算法的改進(jìn)和細(xì)化都需要計(jì)算更多的卷積核，卷積計(jì)算占據(jù)網(wǎng)格化算法開銷的主要部分。雖然采用多核中央處理器和圖形處理器[5-7]可以實(shí)現(xiàn)并行計(jì)算，提高算法性能，但基于Python實(shí)現(xiàn)的上述網(wǎng)格化方法主要局限于NumPy多維數(shù)組計(jì)算，難以適應(yīng)海量數(shù)據(jù)和實(shí)時(shí)處理的需求。近年來數(shù)組Dask.Array的提出，為超大矩陣的數(shù)值計(jì)算開辟了新途徑。文[8]采用Dask并行框架，配合Pipeline技術(shù)測試LOFAR(Low Frequency Analysis Recording)數(shù)據(jù)集[9]，使得原本需要11 h才能完成整個(gè)成像流程的串行化代碼縮短至8 min，大大減少了干涉成像時(shí)間。本文提出基于Dask并行加速的射電干涉可見度數(shù)據(jù)卷積網(wǎng)格化方法，在并行計(jì)算的基礎(chǔ)上兼顧系統(tǒng)的彈性縮放，主要特點(diǎn)是以Dask.Array矩陣分塊存儲(chǔ)和計(jì)算為核心，封裝Nifty-gridder卷積網(wǎng)格化算法提供的Python接口，采取數(shù)據(jù)分塊和延遲計(jì)算，提高了數(shù)值計(jì)算效率，配合Dask的分布式調(diào)度策略，實(shí)現(xiàn)了網(wǎng)格化算法從單機(jī)到集群的遷移。

1 網(wǎng)格化

射電干涉測量方程闡明了可見度數(shù)據(jù)V是天空亮度分布B的傅里葉變換，數(shù)學(xué)表達(dá)式為

Kpq=e-2πi[upql+vpqm+wpq(n-1)],

(1)

(2)

其中，(l,m,n)為觀測方向的余弦坐標(biāo)；(upq,vpq,wpq)為天線p和q組成的基線坐標(biāo)；G和A項(xiàng)分別是瓊斯矩陣參數(shù)化的方向無關(guān)和方向相關(guān)效應(yīng)。在小場近似的條件下，指數(shù)中的wpq(n-1)趨向于0，可見度和天空亮度近似為二維傅里葉變換關(guān)系。由于基線uv軌跡的不規(guī)則性，可見度數(shù)據(jù)并非等間隔離散采樣，直接對干涉測量方程進(jìn)行傅里葉反演的計(jì)算代價(jià)是非常昂貴的。為了應(yīng)用快速傅里葉變換算法成像，可見度數(shù)據(jù)必須重新采樣到規(guī)則化的笛卡爾網(wǎng)格中。

射電干涉成像流程如圖1。不同的光譜頻率(即圖像通道)測量所得的可見度數(shù)據(jù)可以獨(dú)立處理。一個(gè)圖像通道對應(yīng)一個(gè)或多個(gè)數(shù)據(jù)通道。成像通常從空白的天空模型開始迭代，經(jīng)過網(wǎng)格化和傅里葉逆變換運(yùn)算，一個(gè)或多個(gè)明亮的源可能掩蓋周圍微弱的光源，使用CLEAN算法提取明亮點(diǎn)源到天空模型中。與網(wǎng)格化相反的過程是對天空模型進(jìn)行快速傅里葉變換，即從天空模型中計(jì)算可見度，稱為去網(wǎng)格化(Degridding)。測量可見度減去模型可見度數(shù)據(jù)是為了進(jìn)一步提取微弱光源。重復(fù)網(wǎng)格化和去網(wǎng)格化，直到天空模型收斂。

圖1 射電干涉成像流程Fig.1 The imaging pipeline of radio interferometry

干涉成像是射電天文數(shù)據(jù)處理的關(guān)鍵步驟。簡單地將可見度數(shù)據(jù)插值到鄰近的網(wǎng)格會(huì)導(dǎo)致嚴(yán)重的偽影，特別是圖像混疊。為了抑制圖像偽影的副作用，通常采用可見度數(shù)據(jù)與網(wǎng)格化函數(shù)(卷積核)進(jìn)行卷積，然后再重采樣到網(wǎng)格中，這一過程可以產(chǎn)生抗重疊效果。由于卷積核的窗函數(shù)特性，邊界處的圖像裁剪誤差比中心位置高出幾個(gè)數(shù)量級(jí)，產(chǎn)生較大的臟圖。臟圖與修正函數(shù)相乘抵消卷積核產(chǎn)生的誤差，從而獲得正確的光通量，該修正函數(shù)通常是卷積核的傅里葉逆變換的倒數(shù)。相比W-stacking算法，Nifty-gridder為提高卷積核的計(jì)算精度做了以下改進(jìn)：(1)沿著w軸，對所有的可見度數(shù)據(jù)網(wǎng)格化到三維空間，而不是將每個(gè)可見度數(shù)據(jù)的w值投影到鄰近的w平面；(2)改進(jìn)后的修正函數(shù)使臟圖的快速傅里葉變換和離散傅里葉變換之間的差異最小化，因此獲得更高精度的臟圖[10]。

2 卷積網(wǎng)格化實(shí)現(xiàn)

2.1 測量集的并行讀取和分塊

相比于NumPy.ndarray數(shù)組，Dask.Array具有如下優(yōu)勢：(1)支持將超大數(shù)組分割成許多個(gè)NumPy.ndarray小數(shù)組；(2)采用阻塞算法能對大于內(nèi)存的數(shù)組進(jìn)行多核并行計(jì)算。此外我們利用Xarray實(shí)現(xiàn)矩陣的一致性分塊(chunksize)，相關(guān)字段的數(shù)據(jù)可以輕易轉(zhuǎn)化為Dask.Array類型。對于單(多)個(gè)測量集文件，統(tǒng)一將路徑信息放入列表對象，使用Dask.Bag分布式加載，然后按照測量集中的FIELD_ID和DATA_DESC_ID字段分組，實(shí)現(xiàn)并行加載。在本實(shí)驗(yàn)中整個(gè)數(shù)據(jù)集劃為4個(gè)子數(shù)據(jù)集(0_0， 0_1， 0_2和0_3)。以子集0_1為例，Xarray數(shù)據(jù)集定義的部分相關(guān)實(shí)驗(yàn)數(shù)據(jù)如表1。

表1 Xarray數(shù)據(jù)集定義的部分相關(guān)實(shí)驗(yàn)數(shù)據(jù)Table 1 Xarray dataset definitions for some related experiment data

2.2 網(wǎng)格化方法的并行實(shí)現(xiàn)

分布式計(jì)算是海量數(shù)據(jù)處理的有效途徑，Dask并行計(jì)算框架提供了靈活多變的分布式調(diào)度方式。由于Dask任務(wù)調(diào)度方式和用戶自定義的算法相解耦，用戶只需切換調(diào)度方式，便可以使算法在單(多)機(jī)以多進(jìn)程的方式彈性擴(kuò)展，但需要根據(jù)算法的特點(diǎn)，選擇合理的任務(wù)調(diào)度方式，以獲取最佳的計(jì)算性能。本文使用最為復(fù)雜的dask.distributed調(diào)度方式在兩臺(tái)機(jī)器節(jié)點(diǎn)執(zhí)行Nifty-gridder網(wǎng)格化算法。任務(wù)的調(diào)度算法(圖2(a))采用多進(jìn)程的執(zhí)行方式：多個(gè)測量集文件的物理分離有利于使用多進(jìn)程并行讀取數(shù)據(jù)集。測量集文件通常包含多個(gè)射電源(即多個(gè)Sub-dataset)，基于子數(shù)據(jù)集的任務(wù)調(diào)度更進(jìn)一步細(xì)粒度化整個(gè)Nifty-gridder網(wǎng)格化的并行流程。算法使用高階函數(shù)Dask.Array.blockwise封裝和調(diào)用Nifty-gridder的Python接口，實(shí)現(xiàn)了基于子塊的并行計(jì)算以及協(xié)調(diào)子塊的縮聚和拼接操作(圖2(b))。為降低Dask.Array在進(jìn)程之間的傳輸成本，數(shù)值計(jì)算采用多線程的執(zhí)行方式計(jì)算臟圖。Nifty-gridder算法的執(zhí)行過程如下：

圖2 (a)分布式任務(wù)調(diào)度；(b)Nifty-gridder網(wǎng)格化算法Fig.2 (a) Distributed task scheduling; (b) the Nifty-gridder algorithm

(1)沿w軸確定Nw個(gè)采樣平面，并均勻分布到w軸(從w0～wNw-1)。

(2)沿w軸將可見度數(shù)據(jù)網(wǎng)格化到w平面。

(3)初始化Nx×Ny的零矩陣I，對每一個(gè)w=wi平面有：

①將每個(gè)w平面再進(jìn)行uv網(wǎng)格化，然后執(zhí)行二維傅里葉逆變換；

③將上述結(jié)果累加到矩陣I中。

(4)修正函數(shù)乘以矩陣I，得到最終的臟圖。

3 實(shí)驗(yàn)結(jié)果和分析

3.1 實(shí)驗(yàn)環(huán)境

實(shí)驗(yàn)的數(shù)據(jù)集來源于2010年8月23日甚大型Karl G. Jansky干涉陣列對超新星遺跡G055.7 + 3.4長達(dá)8 h的觀測。該陣列采用D型配置，觀測頻率范圍為1～2 GHz，覆蓋所有可用的L波段。實(shí)驗(yàn)的硬件環(huán)境為兩臺(tái)高性能服務(wù)器Intel Xeon CPU E5-2660 v4CPU@3.4 GHz處理器(56核)，512 GB隨機(jī)存取存儲(chǔ)器(Random Access Memory, RAM)。本文使用CASA 5.6.2(Common Astronomy Software Applications)進(jìn)行數(shù)據(jù)結(jié)果的驗(yàn)證。

3.2 Dask并行加速和實(shí)驗(yàn)結(jié)果

以4個(gè)子數(shù)據(jù)集為例，chunksize設(shè)置為20 000行，經(jīng)網(wǎng)格化處理生成臟圖，使用可見度數(shù)據(jù)的行數(shù)度量數(shù)據(jù)集的體積。在同一軟硬件環(huán)境下，比較Dask.Array和NumPy版本Nifty-gridder算法的運(yùn)行時(shí)間(單位： s)，實(shí)驗(yàn)結(jié)果如表2。

由表2可知，基于Dask.Array改進(jìn)的Nifty-gridder算法，其中央處理器的時(shí)間均大于進(jìn)程實(shí)際占用時(shí)間，說明對于計(jì)算密集型問題，使用多核計(jì)算并行效果顯著，可以明顯降低程序的運(yùn)行時(shí)間。以0_0和0_1數(shù)據(jù)集的對比分析為例：即使將可見度數(shù)據(jù)體積增大10.5倍(≈413 696/39 274)，相應(yīng)的執(zhí)行時(shí)間僅增加1.85倍(≈4.95/2.68)，且加速比進(jìn)一步提高。然而Dask.Array是在NumPy的基礎(chǔ)上增加了一層復(fù)雜的設(shè)計(jì)，對于較小的數(shù)據(jù)體積(0_1數(shù)據(jù)集約占40MB)，NumPy可能是正確的選擇，這恰恰說明Dask.Array適宜處理超大型矩陣的數(shù)值計(jì)算。

Dask允許跨集群提交Python函數(shù)，以實(shí)現(xiàn)并行任務(wù)，從而生成大量可以監(jiān)視、控制和計(jì)算的異步調(diào)用對象。對于復(fù)雜的程序處理流程，動(dòng)態(tài)的可視化監(jiān)控有助于了解算法的性能瓶頸，實(shí)驗(yàn)執(zhí)行過程中的實(shí)時(shí)性能監(jiān)控如圖3。

圖3 網(wǎng)格化流程中任務(wù)流的實(shí)時(shí)狀態(tài)Fig.3 The real-time state of task stream in the gridding process

為了說明Dask并行框架的優(yōu)越性，通過增加測量集的輸入量和限定每臺(tái)機(jī)器內(nèi)存占有量并保持實(shí)驗(yàn)環(huán)境一致。從系統(tǒng)資源利用率角度分析并比較基于Dask.Array和NumPy的Nifty-gridder算法性能。由圖4可知，資源利用率的峰值和平均值相比NumPy版本，Dask.Array類型的網(wǎng)格化算法中央處理器利用率和內(nèi)存占有率更低，卻能獲得更快的網(wǎng)格化執(zhí)行時(shí)間(見表2)。主要因?yàn)镈ask.Array數(shù)組采取分塊加載和延遲計(jì)算，尚不具備計(jì)算條件的子塊會(huì)駐留磁盤，以節(jié)約系統(tǒng)資源，而滿足計(jì)算條件的子塊送入內(nèi)存并行執(zhí)行，相反，NumPy數(shù)組必須全部加載到內(nèi)存，導(dǎo)致較高內(nèi)存的持有率。

圖4 網(wǎng)格化流程中中央處理器和內(nèi)存的使用情況對比(Dask. Array vs. NumPy)Fig.4 CPU and memory usage in the gridding process compared Dask. Array to NumPy

為了進(jìn)一步驗(yàn)證代碼的正確性，我們使用標(biāo)準(zhǔn)的CASA軟件對該數(shù)據(jù)集成像，生成的臟圖(圖5(a))與實(shí)驗(yàn)結(jié)果(圖5(b))進(jìn)行對比，兩幅灰度圖中的灰白色點(diǎn)代表觀測源。由圖5可以發(fā)現(xiàn)，算法能正確識(shí)別射電源的分布位置。

圖5 CASA和實(shí)驗(yàn)結(jié)果的臟圖對比Fig.5 Comparison of dirty image from CASA and experimental result

4 總 結(jié)

高性能分布式并行計(jì)算已經(jīng)成為射電干涉成像過程中應(yīng)對高分辨率和大視場干涉陣列產(chǎn)生的海量數(shù)據(jù)的必要方法。可見度數(shù)據(jù)的網(wǎng)格化和去網(wǎng)格化是成像的重要組成部分，網(wǎng)格化并行加速無疑對提高整個(gè)成像速度有重要意義。本文使用開源的Dask分布式計(jì)算框架，結(jié)合Nifty-gridder實(shí)現(xiàn)了測量集的分布式加載和并行網(wǎng)格化加速過程，充分發(fā)揮集群的規(guī)?；瘍?yōu)勢，提高了多核中央處理器的利用率。干涉成像過程包含多個(gè)復(fù)雜的處理流程，都涉及矩陣的數(shù)值計(jì)算，而Dask.Array可以高效地處理多維超大矩陣的數(shù)值計(jì)算，因此，下一步的工作考慮基于Dask實(shí)現(xiàn)去網(wǎng)格化、校準(zhǔn)、成像等算法的并行加速。