李姍姍,駱開達(dá),衛(wèi)守林,2,戴 偉,2,梁 波,2
(1. 昆明理工大學(xué)信息工程與自動化學(xué)院,云南 昆明 650500; 2. 云南省計算機(jī)應(yīng)用技術(shù)重點(diǎn)實(shí)驗(yàn)室,云南 昆明 650500)
射電干涉陣列得到的是非均勻采樣的可見度數(shù)據(jù),在應(yīng)用快速傅里葉變換對可見度數(shù)據(jù)進(jìn)行成像前,需使用網(wǎng)格化方法將采樣數(shù)據(jù)重采樣到一個均勻劃分的網(wǎng)格上。當(dāng)前網(wǎng)格化主要使用基于卷積的方法,卷積網(wǎng)格化過程的實(shí)質(zhì)是矩陣相乘,當(dāng)數(shù)據(jù)量較大時,網(wǎng)格化計算非常耗時。
近年來,天文學(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]。通過將每個可見度數(shù)據(jù)的w值投影到鄰近w平面的W-stacking算法,可以顯著提高網(wǎng)格化性能,但是需要耗費(fèi)額外的內(nèi)存[2]。如果考慮方向相關(guān)效應(yīng),網(wǎng)格化的計算難度將進(jìn)一步增加,同時修正方向相關(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ù)網(wǎng)格化算法開銷的主要部分。雖然采用多核中央處理器和圖形處理器[5-7]可以實(shí)現(xiàn)并行計算,提高算法性能,但基于Python實(shí)現(xiàn)的上述網(wǎng)格化方法主要局限于NumPy多維數(shù)組計算,難以適應(yīng)海量數(shù)據(jù)和實(shí)時處理的需求。近年來數(shù)組Dask.Array的提出,為超大矩陣的數(shù)值計算開辟了新途徑。文[8]采用Dask并行框架,配合Pipeline技術(shù)測試LOFAR(Low Frequency Analysis Recording)數(shù)據(jù)集[9],使得原本需要11 h才能完成整個成像流程的串行化代碼縮短至8 min,大大減少了干涉成像時間。本文提出基于Dask并行加速的射電干涉可見度數(shù)據(jù)卷積網(wǎng)格化方法,在并行計算的基礎(chǔ)上兼顧系統(tǒng)的彈性縮放,主要特點(diǎn)是以Dask.Array矩陣分塊存儲和計算為核心,封裝Nifty-gridder卷積網(wǎng)格化算法提供的Python接口,采取數(shù)據(jù)分塊和延遲計算,提高了數(shù)值計算效率,配合Dask的分布式調(diào)度策略,實(shí)現(xiàn)了網(wǎng)格化算法從單機(jī)到集群的遷移。
射電干涉測量方程闡明了可見度數(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)行傅里葉反演的計算代價是非常昂貴的。為了應(yīng)用快速傅里葉變換算法成像,可見度數(shù)據(jù)必須重新采樣到規(guī)則化的笛卡爾網(wǎng)格中。
射電干涉成像流程如圖1。不同的光譜頻率(即圖像通道)測量所得的可見度數(shù)據(jù)可以獨(dú)立處理。一個圖像通道對應(yīng)一個或多個數(shù)據(jù)通道。成像通常從空白的天空模型開始迭代,經(jīng)過網(wǎng)格化和傅里葉逆變換運(yùn)算,一個或多個明亮的源可能掩蓋周圍微弱的光源,使用CLEAN算法提取明亮點(diǎn)源到天空模型中。與網(wǎng)格化相反的過程是對天空模型進(jìn)行快速傅里葉變換,即從天空模型中計算可見度,稱為去網(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)格會導(dǎo)致嚴(yán)重的偽影,特別是圖像混疊。為了抑制圖像偽影的副作用,通常采用可見度數(shù)據(jù)與網(wǎng)格化函數(shù)(卷積核)進(jìn)行卷積,然后再重采樣到網(wǎng)格中,這一過程可以產(chǎn)生抗重疊效果。由于卷積核的窗函數(shù)特性,邊界處的圖像裁剪誤差比中心位置高出幾個數(shù)量級,產(chǎn)生較大的臟圖。臟圖與修正函數(shù)相乘抵消卷積核產(chǎn)生的誤差,從而獲得正確的光通量,該修正函數(shù)通常是卷積核的傅里葉逆變換的倒數(shù)。相比W-stacking算法,Nifty-gridder為提高卷積核的計算精度做了以下改進(jìn):(1)沿著w軸,對所有的可見度數(shù)據(jù)網(wǎng)格化到三維空間,而不是將每個可見度數(shù)據(jù)的w值投影到鄰近的w平面;(2)改進(jìn)后的修正函數(shù)使臟圖的快速傅里葉變換和離散傅里葉變換之間的差異最小化,因此獲得更高精度的臟圖[10]。
相比于NumPy.ndarray數(shù)組,Dask.Array具有如下優(yōu)勢:(1)支持將超大數(shù)組分割成許多個NumPy.ndarray小數(shù)組;(2)采用阻塞算法能對大于內(nèi)存的數(shù)組進(jìn)行多核并行計算。此外我們利用Xarray實(shí)現(xiàn)矩陣的一致性分塊(chunksize),相關(guān)字段的數(shù)據(jù)可以輕易轉(zhuǎn)化為Dask.Array類型。對于單(多)個測量集文件,統(tǒng)一將路徑信息放入列表對象,使用Dask.Bag分布式加載,然后按照測量集中的FIELD_ID和DATA_DESC_ID字段分組,實(shí)現(xiàn)并行加載。在本實(shí)驗(yàn)中整個數(shù)據(jù)集劃為4個子數(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
分布式計算是海量數(shù)據(jù)處理的有效途徑,Dask并行計算框架提供了靈活多變的分布式調(diào)度方式。由于Dask任務(wù)調(diào)度方式和用戶自定義的算法相解耦,用戶只需切換調(diào)度方式,便可以使算法在單(多)機(jī)以多進(jìn)程的方式彈性擴(kuò)展,但需要根據(jù)算法的特點(diǎn),選擇合理的任務(wù)調(diào)度方式,以獲取最佳的計算性能。本文使用最為復(fù)雜的dask.distributed調(diào)度方式在兩臺機(jī)器節(jié)點(diǎn)執(zhí)行Nifty-gridder網(wǎng)格化算法。任務(wù)的調(diào)度算法(圖2(a))采用多進(jìn)程的執(zhí)行方式:多個測量集文件的物理分離有利于使用多進(jìn)程并行讀取數(shù)據(jù)集。測量集文件通常包含多個射電源(即多個Sub-dataset),基于子數(shù)據(jù)集的任務(wù)調(diào)度更進(jìn)一步細(xì)粒度化整個Nifty-gridder網(wǎng)格化的并行流程。算法使用高階函數(shù)Dask.Array.blockwise封裝和調(diào)用Nifty-gridder的Python接口,實(shí)現(xiàn)了基于子塊的并行計算以及協(xié)調(diào)子塊的縮聚和拼接操作(圖2(b))。為降低Dask.Array在進(jìn)程之間的傳輸成本,數(shù)值計算采用多線程的執(zhí)行方式計算臟圖。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個采樣平面,并均勻分布到w軸(從w0~wNw-1)。
(2)沿w軸將可見度數(shù)據(jù)網(wǎng)格化到w平面。
(3)初始化Nx×Ny的零矩陣I,對每一個w=wi平面有:
①將每個w平面再進(jìn)行uv網(wǎng)格化,然后執(zhí)行二維傅里葉逆變換;
③將上述結(jié)果累加到矩陣I中。
(4)修正函數(shù)乘以矩陣I,得到最終的臟圖。
實(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)境為兩臺高性能服務(wù)器Intel Xeon CPU E5-2660 v4CPU@3.4 GHz處理器(56核),512 GB隨機(jī)存取存儲器(Random Access Memory, RAM)。本文使用CASA 5.6.2(Common Astronomy Software Applications)進(jìn)行數(shù)據(jù)結(jié)果的驗(yàn)證。
以4個子數(shù)據(jù)集為例,chunksize設(shè)置為20 000行,經(jīng)網(wǎng)格化處理生成臟圖,使用可見度數(shù)據(jù)的行數(shù)度量數(shù)據(jù)集的體積。在同一軟硬件環(huán)境下,比較Dask.Array和NumPy版本Nifty-gridder算法的運(yùn)行時間(單位: s),實(shí)驗(yàn)結(jié)果如表2。
由表2可知,基于Dask.Array改進(jìn)的Nifty-gridder算法,其中央處理器的時間均大于進(jìn)程實(shí)際占用時間,說明對于計算密集型問題,使用多核計算并行效果顯著,可以明顯降低程序的運(yùn)行時間。以0_0和0_1數(shù)據(jù)集的對比分析為例:即使將可見度數(shù)據(jù)體積增大10.5倍(≈413 696/39 274),相應(yīng)的執(zhí)行時間僅增加1.85倍(≈4.95/2.68),且加速比進(jìn)一步提高。然而Dask.Array是在NumPy的基礎(chǔ)上增加了一層復(fù)雜的設(shè)計,對于較小的數(shù)據(jù)體積(0_1數(shù)據(jù)集約占40MB),NumPy可能是正確的選擇,這恰恰說明Dask.Array適宜處理超大型矩陣的數(shù)值計算。
Dask允許跨集群提交Python函數(shù),以實(shí)現(xiàn)并行任務(wù),從而生成大量可以監(jiān)視、控制和計算的異步調(diào)用對象。對于復(fù)雜的程序處理流程,動態(tài)的可視化監(jiān)控有助于了解算法的性能瓶頸,實(shí)驗(yàn)執(zhí)行過程中的實(shí)時性能監(jiān)控如圖3。
圖3 網(wǎng)格化流程中任務(wù)流的實(shí)時狀態(tài)Fig.3 The real-time state of task stream in the gridding process
為了說明Dask并行框架的優(yōu)越性,通過增加測量集的輸入量和限定每臺機(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í)行時間(見表2)。主要因?yàn)镈ask.Array數(shù)組采取分塊加載和延遲計算,尚不具備計算條件的子塊會駐留磁盤,以節(jié)約系統(tǒng)資源,而滿足計算條件的子塊送入內(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),算法能正確識別射電源的分布位置。
圖5 CASA和實(shí)驗(yàn)結(jié)果的臟圖對比Fig.5 Comparison of dirty image from CASA and experimental result
高性能分布式并行計算已經(jīng)成為射電干涉成像過程中應(yīng)對高分辨率和大視場干涉陣列產(chǎn)生的海量數(shù)據(jù)的必要方法。可見度數(shù)據(jù)的網(wǎng)格化和去網(wǎng)格化是成像的重要組成部分,網(wǎng)格化并行加速無疑對提高整個成像速度有重要意義。本文使用開源的Dask分布式計算框架,結(jié)合Nifty-gridder實(shí)現(xiàn)了測量集的分布式加載和并行網(wǎng)格化加速過程,充分發(fā)揮集群的規(guī)?;瘍?yōu)勢,提高了多核中央處理器的利用率。干涉成像過程包含多個復(fù)雜的處理流程,都涉及矩陣的數(shù)值計算,而Dask.Array可以高效地處理多維超大矩陣的數(shù)值計算,因此,下一步的工作考慮基于Dask實(shí)現(xiàn)去網(wǎng)格化、校準(zhǔn)、成像等算法的并行加速。