#kuang@DEVP /nas2/cmaqruns/2022fcst/grid45/cctm.fcst/daily
#$ cat join_NCs.py
#!/cluster/miniconda/envs/unresp/bin/python
import os, sys
import netCDF4
import numpy as np
from dtconvertor import dt2jul, jul2dt
nf=len(sys.argv)-1
fnames=[sys.argv[i+1] for i in range(nf)]
cdates, tflags=[],[]
for n in range(nf-1):
nc = netCDF4.Dataset(fnames[n],'r')
t0,t1=n*24,(n+1)*24
if n==0: #first time declarations
V=[list(filter(lambda x:nc.variables[x].ndim==j, [i for i in nc.variables])) for j in [1,2,3,4]]
nt,nlay,nrow,ncol=nc.variables[V[3][0]].shape
nv=len(V[3])
var=np.zeros(shape=(nv,(nf-1)*24,nlay,nrow,ncol))
for iv in range(nv): #store the variables and current time/flags of files
var[iv,t0:t1,:,:,:]=nc.variables[V[3][iv]][:,:,:,:]
cdates.append(jul2dt([nc.CDATE,nc.CTIME]))
tflags.append([jul2dt(nc['TFLAG'][t,0,:]) for t in range(24)])
nt=24 #time interplation period, in hours
for n in range(nf-2):
delt=(cdates[n+1]-cdates[n]).total_seconds()
if 0<=delt<3600.:continue # if conc are from same batch, skip interpolation
t0,t1=(n+1)*24-nt,(n+1)*24
var0=np.zeros(shape=(nv,nt,nlay,nrow,ncol))
for t in range(nt): # transition in nt hours
var0[:,t,:,:,:]=var[:,t0+t,:,:,:]*(nt-t)/nt + var[:,t1,:,:,:]*t/nt
var[:,t0:t1,:,:,:]=var0[:,:,:,:,:]
os.system('cp '+fnames[-2]+' '+fnames[-1])
nc = netCDF4.Dataset(fnames[-1],'r+')
#expand the nc file
for n in range(nf-1):
t0,t1=n*24,(n+1)*24
for t in range(t0,t1):
nc.variables['TFLAG'][t,0,:]=dt2jul(tflags[n][t-t0])
var0=np.zeros(shape=nc.variables['TFLAG'].shape)
var0[:,:,:]=np.array(nc.variables['TFLAG'][:,0,:])[:,None,:]
nc.variables['TFLAG'][:,:,:]=var0[:,:,:]
for iv in range(nv):
nc.variables[V[3][iv]][:,:,:,:]=var[iv,:,:,:,:]
nc.close()