source: palm/trunk/SCRIPTS/cut_empty_timesteps.py @ 4841

#!/usr/bin/env python3

import subprocess
import os
import sys
import netCDF4
from netCDF4 import Dataset, stringtochar
import numpy as np


def cut(input_path, output_path):
    """Cycle through all netCDF files within input_path and remove all
    empty time levels. Output is saved in output_path"""

    if not os.path.isdir(input_path):
        raise OSError("Error: input path does not exist: {0}".format(input_path))
    if not os.path.isdir(output_path):
        os.makedirs(output_path)

    list_input_data = \
        [file for file in os.listdir(input_path) if os.path.splitext(file)[1] == '.nc']

    for nc_file in list_input_data:
        file_in = input_path + "/" + nc_file
        ncfile = Dataset(file_in, "r", format="NETCDF4", encoding='ascii')
        time_dim = ncfile.dimensions["time"]
        time_var = ncfile.variables["time"][:]
        time_mask = ~np.ma.getmaskarray(time_var)

        print("process file ", file_in)
        if not time_dim.isunlimited() and any(time_mask):
            start_index = np.where(time_mask)[0][0]
            end_index = np.where(time_mask)[0][-1]
            file_out = output_path + "/" + nc_file

            ncks_command = "ncks -O -d time,{0},{1} {2} {3}".format(
                start_index,
                end_index,
                file_in,
                file_out,
                )

            ncks_output = subprocess.run(
                ncks_command,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                shell=True
                )

            if ncks_output.returncode != 0:
                raise RuntimeError("Error: {0}".format(ncks_output.stdout))


if __name__ == '__main__':
    arg_list = sys.argv
    if arg_list.__len__() == 3:
        input_path = arg_list[1]
        output_path = arg_list[2]
        cut(input_path, output_path)
    else:
        sys.exit('Arguments are missing! ' +
                 'Usage:\n cut_empty_timesteps.py <input_path> <output_path>')