用 AI 编程工具做射电天文 Radio Astronomy with AI Coding Tools

"""生成用于测试的 HI 21cm 谱线数据立方体（合成测试数据）。"""
import sys
import numpy as np
from astropy.io import fits
from astropy.wcs import WCS


def make_cube(outpath="hi_cube.fits", nx=256, ny=256, nv=128, seed=42):
    w = WCS(naxis=3)
    w.wcs.ctype = ["RA---SIN", "DEC--SIN", "VELO-LSR"]
    w.wcs.cunit = ["deg", "deg", "m/s"]
    w.wcs.crval = [83.633, 22.014, -5000.0]   # 云心，中心速度 -5 km/s
    w.wcs.crpix = [nx / 2, ny / 2, nv / 2]
    w.wcs.cdelt = [-0.01, 0.01, 1000.0]       # 0.01 deg/px，1 km/s/channel
    w.wcs.restfrq = 1420.40575177e6           # HI 21cm 静止频率 (Hz)

    yy, xx = np.mgrid[0:ny, 0:nx]
    r = np.hypot(xx - nx / 2, yy - ny / 2)
    vv = np.arange(nv)[:, None, None]
    gauss_sky = np.exp(-(r ** 2) / (2 * 30.0 ** 2))
    gauss_v = np.exp(-((vv - nv / 2) ** 2) / (2 * 4.0 ** 2))
    cube = 20.0 * gauss_v * gauss_sky          # 峰值 20 Jy/beam
    rng = np.random.default_rng(seed)
    cube = cube + rng.normal(0, 0.1, cube.shape)  # 0.1 Jy/beam 噪声

    hdr = w.to_header()
    hdr["BUNIT"] = "Jy/beam"
    fits.writeto(outpath, cube, hdr, overwrite=True)
    print(f"已生成 {outpath}，形状 {cube.shape}")


if __name__ == "__main__":
    make_cube(sys.argv[1] if len(sys.argv) > 1 else "hi_cube.fits")

"""Step 0：探索数据 — 打印 FITS header 关键键。"""
import sys
from astropy.io import fits

KEYS = ["NAXIS1", "NAXIS2", "NAXIS3",
        "CTYPE1", "CTYPE2", "CTYPE3",
        "CUNIT1", "CUNIT2", "CUNIT3",
        "CDELT1", "CDELT2", "CDELT3",
        "CRVAL1", "CRVAL2", "CRVAL3",
        "CRPIX1", "CRPIX2", "CRPIX3",
        "BUNIT", "RESTFRQ", "VELREF"]


def explore(path="hi_cube.fits"):
    hdu = fits.open(path)[0]
    hdr = hdu.header
    print("文件 / file :", path)
    # 注意：numpy 数组轴与 FITS NAXIS 相反 → [速度, DEC, RA]
    print("shape (numpy, [速度, DEC, RA]):", hdu.data.shape)
    print("dtype:", hdu.data.dtype)
    for k in KEYS:
        print(f"  {k:8s} = {hdr.get(k)}")


if __name__ == "__main__":
    explore(sys.argv[1] if len(sys.argv) > 1 else "hi_cube.fits")

"""Step 1：噪声掩膜 + 矩图（spectral-cube + astropy）。"""
import sys
import numpy as np
from astropy import units as u
import spectral_cube as sc


def make_moments(path="hi_cube.fits"):
    cube = sc.SpectralCube.read(path)
    # 统一速度轴为 km/s（radio 约定，参考系随 FITS 头，通常 LSRK）
    cube = cube.with_spectral_unit(u.km / u.s, velocity_convention="radio")

    # 用 MAD 估计噪声（对离群值稳健）
    data = cube.filled_data[:].value
    mad = np.nanmedian(np.abs(data - np.nanmedian(data))) / 0.6745
    print(f"MAD 噪声估计 / noise estimate: {mad:.4f} Jy/beam")

    # 3σ 掩膜（自动排除 NaN；比较需带单位）
    mask = cube > (3 * mad) * cube.unit
    frac = mask.include().mean()
    print(f"掩膜覆盖比例 / mask coverage: {frac:.1%}")

    masked = cube.with_mask(mask)
    mom0 = masked.moment0()          # Jy/beam·km/s  积分强度
    mom1 = masked.moment1()          # km/s          流量加权速度
    mom2 = masked.moment2()          # km/s          速度弥散

    mom0.write("mom0.fits", overwrite=True)
    mom1.write("mom1.fits", overwrite=True)
    mom2.write("mom2.fits", overwrite=True)
    print("已写出 / wrote: mom0.fits, mom1.fits, mom2.fits")
    print("M0 峰值 / peak:", np.nanmax(mom0.value), mom0.unit)
    print("M1 范围 / range:", np.nanmin(mom1.value), "-", np.nanmax(mom1.value), mom1.unit)
    print("M2 范围 / range:", np.nanmin(mom2.value), "-", np.nanmax(mom2.value), mom2.unit)


if __name__ == "__main__":
    make_moments(sys.argv[1] if len(sys.argv) > 1 else "hi_cube.fits")

"""Step 2：CASA 交叉验证 — 在 CASA 中运行: casa --nogui -c casa_check.py"""
# --- 1. 导入 ---
importfits(fitsimage="hi_cube.fits", imagename="hi_cube.image")

# --- 2. 先看统计，确定掩膜阈值（3σ ≈ 0.3 Jy/beam，请按实际噪声设定）---
imstat(imagename="hi_cube.image")

# --- 3. 阈值掩膜 + 矩 ---
immoments(imagename="hi_cube.image", moments=[0, 1, 2],
          mask='"hi_cube.image" > 0.3',        # 示意阈值，请按实际噪声设定
          outfile="hi_moments.image")

# --- 4. 导出为 FITS 以便与 mom0/1/2.fits 对比 ---
exportfits(imagename="hi_moments.image", fitsimage="hi_moments_casa.fits")

"""Step 3：出版级科学图 — 三面板矩图 + 云心谱线。"""
import os
import numpy as np
# 把 matplotlib 配置/字体缓存放到脚本目录（避免写入用户目录失败）
os.environ.setdefault(
    "MPLCONFIGDIR",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), ".mplconfig"),
)
import matplotlib
matplotlib.use("Agg")          # 无界面环境
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy.wcs import WCS


def plot_moments():
    fig = plt.figure(figsize=(14, 4.6))
    files = ["mom0.fits", "mom1.fits", "mom2.fits"]
    titles = ["Moment 0 · Integrated Intensity",
              "Moment 1 · Velocity",
              "Moment 2 · Dispersion"]
    for i, (f, t) in enumerate(zip(files, titles)):
        hdu = fits.open(f)[0]
        ax = fig.add_subplot(1, 3, i + 1, projection=WCS(hdu.header))
        im = ax.imshow(hdu.data, origin="lower", cmap="inferno")
        ax.set_title(t, fontsize=10)
        ax.coords[0].set_axislabel("RA (J2000)")
        ax.coords[1].set_axislabel("Dec (J2000)")
        cb = fig.colorbar(im, ax=ax, fraction=0.046)
        cb.set_label(hdu.header.get("BUNIT", ""))
    plt.tight_layout()
    plt.savefig("moments.png", dpi=200)
    print("已保存 / saved: moments.png")


def plot_spectrum(ra=83.633, dec=22.014):
    """抽取云心位置的 HI 谱线（世界坐标 → 像素坐标）。"""
    hdu = fits.open("hi_cube.fits")[0]
    w = WCS(hdu.header)
    px, py, _ = w.all_world2pix(ra, dec, 0.0, 0)
    spectrum = hdu.data[:, int(round(float(py))), int(round(float(px)))]
    v = w.all_pix2world(0, 0, np.arange(hdu.data.shape[0]), 0)[2] / 1000.0

    plt.figure(figsize=(6, 3.5))
    plt.plot(v, spectrum)
    plt.xlabel("V_LSR (km/s)")
    plt.ylabel("Flux (Jy/beam)")
    plt.title("HI spectrum at cloud center")
    plt.tight_layout()
    plt.savefig("spectrum.png", dpi=200)
    print("已保存 / saved: spectrum.png")


if __name__ == "__main__":
    plot_moments()
    plot_spectrum()