该代码的作用是从WRF输出文件中提取变量(如经向风、比湿和温度),进行插值处理,并将结果保存到一个NetCDF文件中。
- 将重复的代码抽取到单独的函数中,以提高代码的重用性和可读性。
- 统一了变量命名风格,使用了下划线命名法。
- 添加了对文件的打开和关闭,以及错误处理机制。
- 添加了详细的文档字符串,描述了每个函数的作用、参数和返回值。
- 使用了 main() 函数来组织主要的执行逻辑。
import xarray as xr
import numpy as np
import pandas as pd
import glob
import os
from netCDF4 import Dataset
from wrf import getvar, interplevel, latlon_coords
from metpy.units import units
import metpy.calc as mpcalc
def open_dataset(path):
"""
Open a WRF output dataset.
Args:
path (str): Path to WRF output file.
Returns:
xr.Dataset: WRF output dataset.
"""
return Dataset(path)
def extract_variables(ncfile):
"""
Extract variables from WRF output dataset.
Args:
ncfile (xr.Dataset): WRF output dataset.
Returns:
tuple: Tuple containing time stamp, v-wind component, specific humidity, and temperature.
"""
time = str(getvar(ncfile, "times").data)[:19]
p = getvar(ncfile, "pressure")
ua = getvar(ncfile, "ua", units="m s-1")
va = getvar(ncfile, "va", units="m s-1")
sh = mpcalc.specific_humidity_from_mixing_ratio(getvar(ncfile, 'QVAPOR'))
tem = getvar(ncfile, 'temp')
return time, ua, va, sh, tem,p
def interpolate_variables(ua, va, sh, tem,p):
"""
Interpolate variables.
Args:
time (str): Time stamp.
ua (xr.DataArray): U-wind component.
va (xr.DataArray): V-wind component.
sh (xr.DataArray): Specific humidity.
tem (xr.DataArray): Temperature.
Returns:
tuple: Tuple containing interpolated v-wind component, specific humidity, and temperature.
"""
lats, lons = latlon_coords(ua)
lon = lons.data
lon[lon < 0] += 360
level = np.array([100,125,150,175,200,225,250,300,350,400,450,500,550,600,650,700,750,
775,800,825,850,875,900,925,950,975,1000])
plevs = level[::-1]
v_interp = interplevel(va, p, plevs)
sh_interp = interplevel(sh, p, plevs)
tem_interp = interplevel(tem, p, plevs)
return v_interp.data, sh_interp.data, tem_interp.data
def save_netcdf(save_nc_path, time_list, v_list, sh_list, temp_list, lat, lon):
"""
Save variables to a NetCDF file.
Args:
save_nc_path (str): Path to save the NetCDF file.
time_list (list): List of time stamps.
v_list (list): List of v-wind component arrays.
sh_list (list): List of specific humidity arrays.
temp_list (list): List of temperature arrays.
lat (xr.DataArray): Latitude coordinates.
lon (xr.DataArray): Longitude coordinates.
"""
da_nc = xr.Dataset(
data_vars=dict(
v = (['time', 'level', 'lat', 'lon'], v_list),
sh = (['time', 'level', 'lat', 'lon'], sh_list),
tem = (['time', 'level', 'lat', 'lon'], temp_list),
),
coords={
'time': time_list,
'level': np.flip(np.array([100,125,150,175,200,225,250,300,350,400,450,500,550,600,
650,700,750,775,800,825,850,875,900,925,950,975,1000])),
'lat': lat.data,
'lon': lon.data,
},
attrs=dict(
description="destaggered v-wind component, temperature, and water vapor",
units='m/s, K, gram/kilogram',
),
)
da_nc.to_netcdf(save_nc_path)
print('NetCDF has been output')
def cal_dxdy(file):
ncfile = Dataset(file)
P = getvar(ncfile, "pressure")
lats, lons = latlon_coords(P)
lon = lons[0]
lon[lon<=0]=lon[lon<=0]+360
lat = lats[:,0]
dx, dy = mpcalc.lat_lon_grid_deltas(lon.data, lat.data)
return lon,lat,dx,dy
def main():
case = 'case22'
data_path = f'/WRFV3/test/em_real/output_experiment/{case}'
save_nc_path = f'/WRFV3/test/em_real/v_sh_temp_{case}.nc'
files = sorted(glob.glob(os.path.join(data_path, 'wrfout*')))
time_list, v_list, sh_list, temp_list = [], [], [], []
lon,lat,_,_ = cal_dxdy(files[0])
for file in files:
print(file)
ncfile = open_dataset(file)
time, ua, va, sh, tem,p = extract_variables(ncfile)
v_interp, sh_interp, tem_interp = interpolate_variables(ua, va, sh, tem,p)
time_list.append(pd.to_datetime(time))
v_list.append(v_interp)
sh_list.append(sh_interp)
temp_list.append(tem_interp)
save_netcdf(save_nc_path, time_list, v_list, sh_list, temp_list, lat, lon)
if __name__ == "__main__":
main()