基于Python的SKA-RASCIL中的天線增益校準(zhǔn)實現(xiàn)*

黎 靜，曹 忠，鄧 輝，王 鋒，梅 盈 ，戴 偉

(1. 廣州大學(xué)物理與材料科學(xué)學(xué)院/天體物理中心，廣東 廣州 510006；2. 昆明理工大學(xué)云南省計算機技術(shù)應(yīng)用重點實驗室，云南 昆明 650500)

宇宙再電離、脈沖星、尋找外星智慧生命體等科學(xué)研究是當(dāng)前天文領(lǐng)域的研究熱點，要實現(xiàn)這些研究目標(biāo)，迫切需要射電望遠(yuǎn)鏡具有更高的靈敏度和分辨率。國際大科學(xué)工程——平方公里陣列是人類有史以來建造的最大射電望遠(yuǎn)鏡，信號接收總面積約1 km2，具有高靈敏度，高時間、高空間、高頻率分辨率，高動態(tài)范圍的特點，它的建成將給天文學(xué)研究帶來革命性的影響。

校準(zhǔn)是射電干涉測量的重要組成部分，實際觀測過程中，電離層、溫度、地面環(huán)境變化，電子元件的噪聲和背景溫度等造成天線增益的變化[1-2]。天線增益校準(zhǔn)的準(zhǔn)確性是影響復(fù)可見度數(shù)據(jù)準(zhǔn)確性的主要因素[3]。針對極為微弱且錯綜復(fù)雜的各類系統(tǒng)與環(huán)境效應(yīng)的校準(zhǔn)是困擾當(dāng)前平方公里陣列科學(xué)研究的關(guān)鍵因素，為了提高成像精度和動態(tài)范圍，保證成像的動態(tài)范圍達(dá)到106以上[4]，天線增益的校準(zhǔn)成為需要深入研究的課題。

射電天文仿真校準(zhǔn)成像軟件庫是射電天文數(shù)據(jù)處理專家為平方公里陣列研制的數(shù)據(jù)處理軟件，開發(fā)語言為Python，有良好的可讀性，利于后續(xù)計算機專家對算法進(jìn)行并行實現(xiàn)和部署。為符合RASCIL的編程規(guī)定，本文基于Python語言，實現(xiàn)了經(jīng)典的天線增益校準(zhǔn)算法Antsol，并對其進(jìn)行了性能優(yōu)化。

1 天線增益校準(zhǔn)算法

在過去的幾十年中，針對射電干涉陣成像動態(tài)范圍受校準(zhǔn)精度制約的問題，科學(xué)家們給予了相當(dāng)大的關(guān)注，并在增益校準(zhǔn)算法上取得了一定的成果[5-9]。文[5-6]討論分析了甚大陣(Very Large Array, VLA)數(shù)據(jù)校準(zhǔn)過程中的誤差，給出了利用最小二乘法求解天線增益的經(jīng)典算法——Antsol算法，該算法對可見度數(shù)據(jù)的異常值比較敏感。文[7]提出了一種魯棒性更好的求解增益的方法，與Antsol相比，這種方法受異常值的影響較小，同時性能與Antsol相當(dāng)。文[8]基于最小二乘法提出了4種從協(xié)方差矩陣求解天線增益的算法，并用韋斯特博克綜合孔徑射電望遠(yuǎn)鏡(Westerbork Synthesis Radio Telescope, WSRT)的實驗數(shù)據(jù)進(jìn)行了仿真測試。文[9]提出了一種較新的算法StEFCal(Statistically Efficient And Fast Calibration)，該算法在低頻陣列(Low Frequency Array, LOFAR)數(shù)據(jù)校準(zhǔn)管線的幾個階段得到了應(yīng)用。Antsol算法至今仍在大型米波射電望遠(yuǎn)鏡(Giant Metrewave Radio Telescope, GMRT)和通用天文軟件包(Common Astronomy Software Applications, CASA)中使用，是天文圖像處理系統(tǒng)(Astronomical Image Processing System, AIPS)的默認(rèn)算法，因此SKA-RASCIL必須有該算法的高性能實現(xiàn)。

2 天線增益校準(zhǔn)算法的基本原理

在綜合孔徑成像中，不同天線接收的信號通過乘法電路和時間平均電路進(jìn)行合成并輸出。乘法電路和時間平均電路組成相關(guān)器，干涉陣的每個相關(guān)器的輸出為天線信號和系統(tǒng)噪聲的總和。對于天線p和q組成的天線對，相關(guān)器的輸出為

(1)

gk=|gk|e-iφk，

(2)

其中，|gk|為天線接收信號的幅值；φ為信號基于天線的相位。因此，p，q天線對的復(fù)增益為

(3)

(4)

(5)

增益校準(zhǔn)的殘差rpq為

(6)

為了提高校準(zhǔn)的精度，要求殘差收斂得盡可能小，需要對(5)式進(jìn)行多次迭代求解。

3 Antsol算法的實現(xiàn)

Antsol算法已經(jīng)有Fortran和C語言的實現(xiàn)版本。本文參考美國甚大陣的Fortran程序(Thomspon經(jīng)典書籍的附錄1[5])，并使用Python語言對Antsol算法進(jìn)行高性能實現(xiàn)。同時，由于SKA-SDHP采用了DASK(https://www.dask.org)作為多任務(wù)分布計算框架 ，而DASK庫可以將計算擴展到多個內(nèi)核甚至多個機器，在天線增益校準(zhǔn)軟件模塊開發(fā)時，我們不需要在模塊內(nèi)使用多任務(wù)并行計算，同時，為了使軟件具有更好的通用性，也不采用圖形處理器加速。

3.1 主要流程

Antsol算法的實現(xiàn)流程如圖1，主要包括5個步驟：(1)輸入可見度數(shù)據(jù)以及各種參數(shù)，包括觀測可見度、模型可見度、迭代次數(shù)、容差等；(2)對觀測可見度和模型可見度進(jìn)行擬合，求解增益表；(3)求解單個天線的增益；(4)計算殘差；(5)收斂條件檢測，當(dāng)達(dá)到迭代次數(shù)或者前后兩次求解的增益變化小于預(yù)設(shè)值時，結(jié)束計算。

圖1 Antsol算法實現(xiàn)流程圖Fig.1 The flow chart of the antsol algorithm implementation

3.2 相關(guān)函數(shù)

主要的函數(shù)見表1，其中函數(shù)1用來求解增益表(包含增益、權(quán)重、殘差等數(shù)據(jù)變量)；函數(shù)2到4用來迭代求解不同極化數(shù)下所有天線的增益；函數(shù)5到7用來更新天線增益，分別被函數(shù)2、3、4調(diào)用；函數(shù)8和9用來求解殘差，分別被函數(shù)5和函數(shù)7調(diào)用。

參考甚大陣的Fortran程序，我們用Python語言實現(xiàn)了整個算法。以殘值計算的函數(shù)solution_resudial_matrix為例， 圖2給出了用基本的Python語言實現(xiàn)的殘值計算流程，用了5重循環(huán)實現(xiàn)多天線的多通道計算，最后返回殘值，參數(shù)gain表示增益數(shù)組，x表示點源等效可見度數(shù)組,xwt表示等效點源可見度權(quán)重數(shù)組。

4 程序優(yōu)化

顯然，圖2的代碼執(zhí)行效率并不高。開發(fā)環(huán)境PyCharm中的Profile顯示，圖2中的第7行到第16行的5重循環(huán)耗時最長，其次是第17行、18行的花式索引也消耗了大量時間。

表1 算法實現(xiàn)的相關(guān)函數(shù)Table 1 The functions of the Antsol algorithm

圖2 殘值計算函數(shù)solution_resudial_matrix()Fig.2 Python code for the solution_resudial_matrix()

4.1 代碼優(yōu)化

經(jīng)過多次嘗試，本文最終采用Numpy的愛因斯坦求和約定(Einstein Summation Convention)對這段5重循環(huán)代碼進(jìn)行優(yōu)化。愛因斯坦求和約定又稱愛因斯坦標(biāo)記法，是大科學(xué)家愛因斯坦于 1916 年提出的一種標(biāo)記約定，Numpy, PyTorch, TensorFlow都實現(xiàn)了該標(biāo)記法，相應(yīng)的函數(shù)名稱為 einsum()。einsum()函數(shù)能快速靈活地實現(xiàn)矩陣計算，如矩陣轉(zhuǎn)置、矩陣乘法、求跡、張量乘法、數(shù)組求和等。針對第17行、18行的花式索引，我們使用Numpy的putmask()函數(shù)實現(xiàn)數(shù)組的部分替換。圖3為最終的代碼，可以看出，代碼變得更為簡潔優(yōu)雅。

圖3 改進(jìn)的殘值計算函數(shù)solution_residual_matrix()Fig.3 Optimized code for the solution_residual_matrix()

4.2 性能測試

我們使用RASCIL對SKA1-LOW進(jìn)行了模擬觀測，然后對所實現(xiàn)的Antsol算法的性能進(jìn)行了測試。測試服務(wù)器環(huán)境為CentOS 7.8操作系統(tǒng)，Intel E5 2620 V4 CPU，128 G內(nèi)存。表2列出了優(yōu)化前后幾個主要函數(shù)的運行時間，可以看出，優(yōu)化以后的程序執(zhí)行時間大幅度縮短，其中殘值矩陣計算函數(shù)Solution_residual_matrix()的計算時間約為之前的0.40%，而最快的殘值矢量計算函數(shù)Solution_residual_vector()只有之前的0.20%。這是因為在用Python實現(xiàn)Antsol算法的過程中，在確保結(jié)果正確的同時，調(diào)用了Numpy的einsum()函數(shù)和putmask()函數(shù)，而Numpy是基于C/C++實現(xiàn)，因此程序的運行效率得到了大幅提升。

表2 優(yōu)化前后各主要函數(shù)計算耗時比較Table 2 The comparison of the execution time between the original and the optimized code

5 結(jié) 論

本文分析了天線增益校準(zhǔn)算法Antsol的基本原理，并利用Numpy的愛因斯坦求和約定函數(shù)對Antsol算法進(jìn)行了高性能實現(xiàn)，所實現(xiàn)的Antsol算法滿足平方公里陣列數(shù)據(jù)處理軟件 RASCIL的性能要求，并已經(jīng)集成到RASCIL中(https://gitlab.com/ska-telescope/external/rascil/-/blob/master/rascil/processing_components/calibration/solvers.py)，不僅為當(dāng)前平方公里陣列數(shù)據(jù)處理提供了支撐，也為未來數(shù)據(jù)處理的性能優(yōu)化提供了算法參考。