postprocess

#!/usr/bin/env python
# This function processes results printed as tables to each directory
import os
import numpy as np
import xarray as xr
import glob
import sys
# Initial stuff
pattern = sys.argv[-1] # pattern
files = sorted(glob.glob(pattern)) # list of files
if len(files)==0:
    print(f'No files with pattern {pattern} found.'); exit() # exit script
directory = '/'.join(os.path.abspath(__file__).split('/')[:-1]) # path of current file
prefix = os.path.abspath(files[0]).split('/')[-2].split('sb_')[-1] # directory
# Put everything into single numpy array
# Optionally iterate through all experiments
print(f'Post-processing files with pattern {pattern}.')
expname = os.path.abspath(pattern).split('/')[-2] # the experiment name i.e. subdirectory
filetype = pattern.split('/')[-1].split('*')[0] # file prefix
initdata = np.loadtxt(files[0], skiprows=1)
z = initdata[:,0] # height
p = initdata[:,1] # pressure
times, datas = [], []
for file in files:
    data = np.loadtxt(file, usecols=(2,3,5,6,7), skiprows=1) # exclude z/p and direct beam
    datas.append(data[:,:,None]) # includes flux down, flux up, flux divergence, heating, and temperature
    times.append(int(file.split(filetype)[-1].split('.txt')[0])) # the observation timestep
data = np.concatenate(datas, axis=-1) # final axis is now time dimension
# Create NetCDF file with XArray from that
# Syntax is to say dataset['name'] = (<dim list>, <data>, [<attr dictionary>]) or equivalently
# when declaring dataset, with Dataset({'var1': above_tuple_1, 'var2': above_tuple_2})
names = ['fdown', 'fup', 'dfdz', 'dtempdz', 'temp']
longs = ['downward flux', 'upward flux', 'flux divergence', 'heating rate', 'temperature']
units = ['W/m2', 'W/m2', 'W/m3', 'K/day', 'K']
dataset = xr.Dataset(coords={
        'level': ('level', np.arange(data.shape[0]), {'long_name':'level number'}),
        'time': ('time', times, {'long_name':'time', 'units':'hours'}),
        }) # initialize with just coordinates
dataset['z'] = ('level', z, {'long_name':'altitude', 'units':'km'})
dataset['p'] = ('level', p, {'long_name':'pressure', 'units':'mb'}) # add height and pressure
for i,name in enumerate(names):
    scale = 1e-3 if 'divergence' in name else 1
    dataset[name] = (('level','time'), scale*data[:,i,:], {'long_name':longs[i], 'units':units[i]}) # create height by time data
print(f'Saving to \"{directory}/results/{prefix}_{filetype}.nc\"...')
dataset.to_netcdf(f'{directory}/results/{prefix}_{filetype}.nc') # save in parent directory