NAMEPreProcessor.py

import os
import sys

import cftime
import iris.util
from iris.cube import CubeList, Cube


def save_cube(cube: iris.cube.Cube,
              sanitised_field_name: str,
              year_string: str,
              month_string: str,
              output_path: str):

    os.makedirs(output_path, exist_ok = True)

    file_name = "{}-no_part-{}_{}.nc".format(sanitised_field_name, year_string, month_string)
    nc_file_path = os.path.join(output_path, file_name)

    print("writing {}".format(nc_file_path))
    iris.save(cube, nc_file_path)

    # nc_zip_file_path = "{}.zip".format(nc_file_path)
    # print("zipping to {}".format(nc_zip_file_path))
    # with ZipFile(nc_zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zip:
    #     zip.write(nc_file_path, basename(nc_file_path))
    #
    # print("deleting {}".format(nc_file_path))
    # os.remove(nc_file_path)


def process_met_office_NAME(name_file_wildcard: str, out_path: str):

    cube_list = iris.load_raw(name_file_wildcard)

    # get all the field names from the cubes as a unique set
    field_names: set = set()

    for cube in cube_list:
        field_names.add(cube.name())

    # write a cube for each calendar month for each field type
    for field_name in field_names:

        field_name: str = field_name
        print(field_name)
        cube: Cube = cube_list.extract(iris.Constraint(name = field_name))

        for month in range(1, 13):

            month_constraint = iris.Constraint(time = lambda cell: cell.point.month == month)

            cubes_for_month: CubeList = cube.extract(month_constraint)

            if len(cubes_for_month) != 0:
                cube_for_month: Cube = cubes_for_month.merge_cube()

                # hack, the 'cell_methods' field is NonType and this gets the save method of iris upset,
                # so setting to an empty Tuple
                cube_for_month.cell_methods = ()

                time_coord = cube_for_month.coord('time')
                sample_date = time_coord.units.num2date(time_coord.points[0])
                year_string = cftime.datetime.strftime(sample_date, "%Y")
                month_string = cftime.datetime.strftime(sample_date, "%m")
                # year_string = cftime.real_datetime.strftime(sample_date, "%Y")
                # month_string = cftime.real_datetime.strftime(sample_date, "%m")

                # get rid of funny characters, for filename
                sanitised_field_name = field_name.replace("(", "")
                sanitised_field_name = sanitised_field_name.replace(")", "")
                sanitised_field_name = sanitised_field_name.replace("/", "_")

                save_cube(cube_for_month,
                          sanitised_field_name,
                          year_string,
                          month_string,
                          out_path)


if __name__ == '__main__':

    """
    example usage:
    /media/scratch/lb584_scratch/projects/ews/input/eth/MET_20200707/WR_EnvSuit_Met_Ethiopia_20200707/Met_Data_*.txt
    /media/scratch/lb584_scratch/projects/ews/output/name_processing/another_dir
    """

    _name_file_wildcard = sys.argv[1]
    _out_path = sys.argv[2]

    process_met_office_NAME(_name_file_wildcard, _out_path)