vapor 2

"""
wrfvap

Script to pre-process WRF data to create a file that should work welll with the Vapor 
visialuzation software

- Combines multiple netcdf files into a single file

- Crops the boundary region of the model.  This reduces file size and eliminates unrealistic
  data values that can occur in these areas

- Gets some useful aditional variables that are useful when explorine 3D volumetric data

Notes:

- Requires that the nco program ncrcat is installed

- Processing may take several minutes

"""
import xarray as xa
import numpy as np
from glob import glob
import os
import sys
from netCDF4 import Dataset
import wrf

# Settings
ncrcat = "ncrcat"

def crop_wrf(fpattern, fout, n, max_z):
    """
    Combines and crops a sequence of WRF files definede by fpattern (e.g. wrfout_d01_*). The cropping buffer
    is given by n and should be at least 5, but 10 or more is recomended for larger domains
    """
    file = glob(fpattern)[0]

    ds = xa.open_dataset(file,mode="r")
    ny = ds["XLAT"].shape[1]
    nx = ds["XLAT"].shape[2]

    command = f'{ncrcat} -x -v EL_PBL,DTAUX3D,DTAUY3D  -d south_north,{n},{ny-n} -d south_north_stag,{n},{ny-n+1} -d west_east,{n},{nx-n} -d west_east_stag,{n},{nx-n+1} -d bottom_top,0,{max_z} -d bottom_top_stag,0,{max_z+1} {fpattern} {fout}'

    print(command)
    os.system( command )

def append_ds( v, fnc ):
    """
    Appends the data is dataset v to an existing netCDF file fnc
    """
    ftmp = fnc+'.tmp'   # Temporary file

    v.to_netcdf(ftmp,mode="w", compute=True)

    nc_command = "ncks -A %s %s" % (ftmp, fnc)
    print(nc_command)
    os.system(nc_command)
    
    os.system('rm -f ' + ftmp)

# MAIN PROGRAM

#**************************************************************************
# User defined inpit and output files
#**************************************************************************
#
fpattern = "wrfout_d04_2023-04-26_00:00:00"
fnc      = "wrf_vapor_d04.nc"

# Combine and crop input data
print('Combine input files and crop')
crop_wrf(fpattern, fnc, 10, 25)

# Open the combined dataset
print('Compute new variables')
ds = xa.open_dataset(fnc,mode="r")
data = {}

# Precipitation rate
#data["PRATE"] = ds["RAINNC"].differentiate("Time").compute()
da = ds["RAINNC"].diff(dim='Time')
da = xa.concat( [ds["RAINNC"].isel(Time=0)*0, da], dim="Time")
da = da.transpose("Time",...)

print(ds["RAINNC"])
print('prate')
print(da)
data["PRATE"] = da

# Wind velocity
u = wrf.getvar( Dataset(fnc), 'ua', timeidx=wrf.ALL_TIMES)
v = wrf.getvar( Dataset(fnc), 'ua', timeidx=wrf.ALL_TIMES)
data["VEL"] = np.sqrt(u*u+v*v)
data["VEL"] = data["VEL"].assign_coords( XTIME=data["PRATE"].coords["XTIME"] )

# Water vapor flux
qv = ds["QVAPOR"]
data["QFLX"] = data["VEL"] * qv

# Relative humidity
data["RH"] = wrf.getvar( Dataset(fnc), 'rh', timeidx=wrf.ALL_TIMES)
data["RH"] = data["RH"].assign_coords( XTIME=data["PRATE"].coords["XTIME"] )

# Temperatura C
data["TC"] = wrf.getvar( Dataset(fnc), 'tc', timeidx=wrf.ALL_TIMES)
data["TC"] = data["TC"].assign_coords( XTIME=data["PRATE"].coords["XTIME"] )

ds.close()

# Clean some metadata variables from wrf-python that are not compatible with NetCDF
dsnew = xa.Dataset( data )
del dsnew["RH"].attrs["projection"]
del dsnew["TC"].attrs["projection"]

# Append the new variables to the file
print('Write output file')
append_ds( dsnew, fnc )