import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from metpy.calc import divergence
from metpy.units import units
# 1. 读取NetCDF文件
def read_netcdf(file_path):
"""读取并处理NetCDF文件"""
ds = xr.open_dataset(file_path)
print("文件变量:", list(ds.data_vars))
# 检查常见变量名并重命名
var_mapping = {
'uwnd': 'u', 'u_wind': 'u',
'vwnd': 'v', 'v_wind': 'v',
'shum': 'q', 'qair': 'q', 'hus': 'q'
}
for old_name, new_name in var_mapping.items():
if old_name in ds.data_vars:
ds = ds.rename({old_name: new_name})
return ds
# 2. 计算水汽通量
def calculate_flux(ds, level_idx=0):
"""计算水汽通量分量"""
# 检查必要变量是否存在
required_vars = ['u', 'v', 'q']
missing_vars = [var for var in required_vars if var not in ds.data_vars]
if missing_vars:
raise ValueError(f"缺少必要变量: {missing_vars}")
# 添加单位 (metpy要求)
u = ds['u'] * units('m/s')
v = ds['v'] * units('m/s')
q = ds['q'] * units('kg/kg')
# 计算水汽通量分量
qu = u * q # x方向水汽通量 (kg/(m·s))
qv = v * q # y方向水汽通量 (kg/(m·s))
return qu, qv
# 3. 计算水汽通量散度
def calculate_flux_divergence(qu, qv):
"""计算水汽通量散度"""
# 使用metpy计算散度
div = divergence(qu, qv) * 1e7 # 缩放为常用单位 (10^-7 kg/(m²·s))
return div
# 4. 绘图函数
def plot_moisture_flux(qu, qv, div, time_idx=0, level_idx=0):
"""绘制水汽通量和水汽通量散度图"""
fig = plt.figure(figsize=(16, 12))
# 创建地图投影
proj = ccrs.PlateCarree()
ax = fig.add_subplot(1, 1, 1, projection=proj)
# 添加地图要素
ax.add_feature(cfeature.COASTLINE, linewidth=0.8)
ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.5)
ax.add_feature(cfeature.OCEAN, color='lightblue', alpha=0.3)
ax.add_feature(cfeature.LAND, color='beige', alpha=0.2)
ax.gridlines(draw_labels=True, linestyle='--', alpha=0.5)
# 提取指定时间和高度的数据
try:
# 尝试使用level索引
qu_slice = qu.isel(time=time_idx, level=level_idx)
div_slice = div.isel(time=time_idx, level=level_idx)
except:
# 如果没有level维度,尝试pressure或isobaric
try:
qu_slice = qu.isel(time=time_idx, pressure=level_idx)
div_slice = div.isel(time=time_idx, pressure=level_idx)
except:
# 最后尝试直接使用time索引
qu_slice = qu.isel(time=time_idx)
div_slice = div.isel(time=time_idx)
qv_slice = qv.isel(time=time_idx) # 假设qv与qu有相同维度
# 绘制水汽通量散度填色图
contour = ax.contourf(
qu_slice.lon,
qu_slice.lat,
div_slice.values,
levels=np.linspace(-5, 5, 21),
cmap='RdBu',
extend='both',
transform=proj
)
cbar = fig.colorbar(contour, ax=ax, shrink=0.7)
cbar.set_label('水汽通量散度 (10$^{-7}$ kg/(m$^2$·s))', fontsize=12)
# 绘制水汽通量箭头(降低密度)
skip = max(1, len(qu_slice.lon) // 30) # 自适应密度
# 确保网格是规则的
if hasattr(qu_slice, 'lon') and hasattr(qu_slice, 'lat'):
lons, lats = qu_slice.lon, qu_slice.lat
else:
lons, lats = np.meshgrid(qu_slice.longitude, qu_slice.latitude)
Q = ax.quiver(
lons[::skip, ::skip],
lats[::skip, ::skip],
qu_slice.values[::skip, ::skip],
qv_slice.values[::skip, ::skip],
scale=300,
color='k',
transform=proj
)
# 添加箭头比例尺
ax.quiverkey(Q, 0.85, 0.05, 100,
'100 kg/(m·s)',
labelpos='E',
coordinates='figure')
# 设置标题
time_str = str(qu_slice.time.values)[:19] if 'time' in qu_slice.coords else "未知时间"
level_value = ""
if 'level' in qu_slice.coords:
level_value = f" | 高度: {qu_slice.level.values} hPa"
elif 'pressure' in qu_slice.coords:
level_value = f" | 高度: {qu_slice.pressure.values} hPa"
elif 'isobaric' in qu_slice.coords:
level_value = f" | 高度: {qu_slice.isobaric.values} hPa"
ax.set_title(f'水汽通量和水汽通量散度\n时间: {time_str}{level_value}',
fontsize=16)
plt.tight_layout()
plt.savefig('moisture_flux_divergence.png', dpi=300, bbox_inches='tight')
print("图像已保存为: moisture_flux_divergence.png")
plt.show()
# 主程序
if __name__ == "__main__":
# 使用您提供的文件路径
file_path = r'C:/Users/25636/Desktop/bi/ziji/zhengduanfx/first/ziliao/200307.nc'
print(f"正在处理文件: {file_path}")
try:
# 处理数据
ds = read_netcdf(file_path)
qu, qv = calculate_flux(ds)
div = calculate_flux_divergence(qu, qv)
# 绘制图形(默认使用第一个时间和高度层)
plot_moisture_flux(qu, qv, div)
except Exception as e:
print(f"处理过程中发生错误: {str(e)}")
print("请检查以下可能原因:")
print("1. 文件路径是否正确")
print("2. 文件是否包含u/v风分量和比湿(q)数据")
print("3. 尝试使用不同的变量名映射")
print("4. 数据维度是否正确"),帮我修改这个代码,让他输出两张图