Bootstrap

Python - 从WRF输出文件中提取变量(如经向风、比湿和温度),进行插值处理,并将结果保存到一个NetCDF文件

该代码的作用是从WRF输出文件中提取变量(如经向风、比湿和温度),进行插值处理,并将结果保存到一个NetCDF文件中。

  1. 将重复的代码抽取到单独的函数中,以提高代码的重用性和可读性。
  2. 统一了变量命名风格,使用了下划线命名法。
  3. 添加了对文件的打开和关闭,以及错误处理机制。
  4. 添加了详细的文档字符串,描述了每个函数的作用、参数和返回值。
  5. 使用了 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()

;