WRF | 后处理 | 读取最近站点的变量数据

在WRF后处理过程中，读取输出的wrfout数据，同时找到需要站点的数据，将其提取出来

# -*- coding: utf-8 -*-
"""
Created on Thu Jan 25 15:04:59 2024

@author: Jianpu

@blog:  https://blog.csdn.net/weixin_44237337?spm=1000.2115.3001.5343

@email: xianpu.ji@hhu.edu.cn

@introduction: keep learning although slowly
"""


import matplotlib.ticker as mticker
import xarray as xr
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import numpy as np
import cartopy.crs as crs
from wrf import (getvar, interplevel, to_np, latlon_coords, get_cartopy,ll_to_xy,
                 cartopy_xlim, cartopy_ylim)
import pandas as pd
import matplotlib.patches as mpatches
import argparse




def read_station_file(file_path):
    """
    Reads station information from a text file.

    Each line of the file should contain: station_name latitude longitude
    """
    names = []
    lats  = []
    lons  = []
    
    with open(file_path, 'r') as file:
        for line in file:
            values = line.strip().split()
            station_name = values[0]
            latitude = float(values[1])
            longitude = float(values[2])
            names.append(station_name)
            lats.append( latitude)
            lons.append( longitude)

    return names,lats,lons





def interpolate_wrf_data(wrf_file_path, variables, lat_station, lon_station, station_file_path):
    # 读取WRF文件
    wrf_data = Dataset(wrf_file_path)
    # 获取时间坐标
    p = getvar(wrf_data, "pressure")
    
    interpolated_data = {}

    for variable in variables:

        # 获取WRF变量数据
        wrf_variable = getvar(wrf_data, str(variable), timeidx=0)[:]

        print(f"Variable: {variable}, Shape: {wrf_variable.shape}")

        if station_file_path is None:
            
            
            # 对单个站点或多个站点进行插值
            if len(lat_station) == 1:

                lat_point, lon_point = lat_station, lon_station
                point = ll_to_xy(wrf_data, lat_point, lon_point, timeidx=0, meta=False)

                x, y = point[0,:].data, point[1,:].data


                if len(wrf_variable.shape) == 2:
                    
                    interpolated_data[variable] = wrf_variable[y, x].data
                    
                elif (len(wrf_variable.shape) == 3) and (wrf_variable.shape[0] == len(p)):

                    wrf_variable_interp = interplevel(wrf_variable, p, 850)
                    
                    interpolated_data[variable] = wrf_variable_interp[y, x][:].data

                elif (len(wrf_variable.shape) == 3) and (wrf_variable.shape[0] == 2):
                    
                    interpolated_data[variable] = wrf_variable[:, y, x][:].data

                else:
                    print(f"Unhandled case for Variable: {variable}")

            elif len(lat_station) >= 2:  # 多站点
                lat_point, lon_point = lat_station, lon_station
                point = ll_to_xy(wrf_data, lat_point, lon_point, timeidx=0, meta=False)
                x, y = point[0, :], point[1, :]
                if len(wrf_variable.shape) == 2:
                    interpolated_data[variable] = wrf_variable[y, x].data
                elif (len(wrf_variable.shape) == 3) and (wrf_variable.shape[0] == len(p)):
                    wrf_variable = interplevel(wrf_variable, p, 850,)
                    interpolated_data[variable] = wrf_variable[y, x].data
                elif (len(wrf_variable.shape) == 3) and (wrf_variable.shape[0] == 2):
                    interpolated_data[variable] = wrf_variable[:, y, x][:].data
                else:
                    print(f"Unhandled case for Variable: {variable}")
                    
        else:


            # Read station information from the file
            names, lats, lons = read_station_file(station_file_path)

            lat_point, lon_point = lats, lons
            point = ll_to_xy(wrf_data, lat_point, lon_point, timeidx=0, meta=False)
            x, y = point[0, :], point[1, :]
            if len(wrf_variable.shape) == 2:
                interpolated_data[variable] = wrf_variable[y, x].data
            elif (len(wrf_variable.shape) == 3) and (wrf_variable.shape[0] == len(p)):
                wrf_variable = interplevel(wrf_variable, p, 850,)
                interpolated_data[variable] = wrf_variable[y, x].data
            elif (len(wrf_variable.shape) == 3) and (wrf_variable.shape[0] == 2):
                interpolated_data[variable] = wrf_variable[:, y, x][:].data
            else:
                print(f"Unhandled case for Variable: {variable}")

    print("Interpolated Data:", interpolated_data)
    return interpolated_data



def parse_args():
    

    parser = argparse.ArgumentParser(description="Interpolate WRF data at specified station coordinates.")
    parser.add_argument("-wrf_file_path", type=str, help="Path to the WRF file.")
    parser.add_argument("-station_file_path", type=str, help="Path to the station coordinates file.")
    parser.add_argument("-lat_station", type=float, nargs='+', help="List of latitudes for station coordinates.")
    parser.add_argument("-lon_station", type=float, nargs='+', help="List of longitudes for station coordinates.")
    
    return parser.parse_args()

if __name__ == "__main__":
    
    variables = ['RAINNC','RAINC', 'U10', 'V10', 'ctt', 'rh','wspd_wdir10']
    args = parse_args()
    get_data = interpolate_wrf_data(args.wrf_file_path, variables, args.lat_station, args.lon_station, args.station_file_path)
    print(get_data)


### example 
# single 
# python your_script.py -wrf_file_path path/to/wrf_file.nc -lat_station 30.0 -lon_station 40.0
# multiple
# python your_script.py -wrf_file_path path/to/wrf_file.nc -lat_station 30.0 31.0 32.0 -lon_station 40.0 41.0 42.0