inst_eddies.py.in

#!/usr/bin/env python3

"""This script creates the directory {SHPC_dir}/Slice_{my_slice},
containing a time slice of instantaneous eddies. If
{SHPC_dir}/Slice_{my_slice} already exists then it is replaced.

"""

import shutil
import nco
import os
import subprocess
import shapefile
import sys
import time
import csv


def process_1_file(
    bbox,
    nco_instance,
    nc_file,
    inst_eddies_exe,
    d,
    my_slice,
    writer,
    SHPC_dir,
    config_nml_string,
    periodic,
):
    t0 = time.perf_counter()

    if bbox:
        xmin, xmax, ymin, ymax = bbox

        try:
            nco_instance.ncks(
                nc_file,
                output="cropped.nc",
                options=[
                    f"--dimension=longitude,{xmin},{xmax}",
                    f"--dimension=latitude,{ymin},{ymax}",
                ],
            )
        except nco.NCOException as err:
            print(err.stdout)
            raise

        nc_file = "cropped.nc"

    nco_instance.ncpdq(nc_file, output="unpacked.nc", options=["--unpack"])
    t1 = time.perf_counter()
    elapsed_NCO = t1 - t0
    t0 = t1
    subprocess.run(
        [inst_eddies_exe, SHPC_dir, "unpacked.nc"],
        check=True,
        input=f"&dates_nml date = {d}, slice = {my_slice}/\n"
        + config_nml_string
        + f"&input_ssh_nml periodic = {periodic}/\n",
        text=True,
    )
    elapsed_Fortran = time.perf_counter() - t0
    writer.writerow([d, elapsed_NCO, elapsed_Fortran])


def loop_inst_eddies(
    SHPC_dir, files, bbox, config_nml_string, d, my_slice, periodic
):
    """`files` is an iterator of file paths. The first file must exist,
    others can be missing. bbox is a tuple of four real numbers. `d`
    is a date index (integer value) for the first date. my_slice is
    the slice index (integer value) in the SHPC.

    """

    inst_eddies_exe = "@CMAKE_CURRENT_BINARY_DIR@/inst_eddies"

    if not os.access(inst_eddies_exe, os.X_OK):
        sys.exit(
            "inst_eddies.loop_inst_eddies:"
            + inst_eddies_exe
            + " not found or not executable"
        )

    nco_instance = nco.Nco()

    if os.access(f"{SHPC_dir}/Slice_{my_slice}", os.F_OK):
        shutil.rmtree(f"{SHPC_dir}/Slice_{my_slice}")

    for orient in ["Cyclones", "Anticyclones"]:
        os.makedirs(f"{SHPC_dir}/Slice_{my_slice}/{orient}")

    perf_report = open(
        f"{SHPC_dir}/Slice_{my_slice}/perf_report.csv", "w", newline=""
    )
    writer = csv.writer(perf_report, lineterminator="\n")
    writer.writerow(["", "s", "s"])  # units
    writer.writerow(["date", "elapsed NCO", "elapsed Fortran"])  # long names

    # Short names (without blanks):
    writer.writerow(["date", "elapsed_NCO", "elapsed_Fortran"])

    # First file is out of the loop because it must be present:
    nc_file = next(files)
    print("inst_eddies.loop_inst_eddies: Date:", d)
    print("inst_eddies.loop_inst_eddies: Input file:", nc_file, flush=True)
    process_1_file(
        bbox,
        nco_instance,
        nc_file,
        inst_eddies_exe,
        d,
        my_slice,
        writer,
        SHPC_dir,
        config_nml_string,
        periodic,
    )
    d += 1

    for nc_file in files:
        print("inst_eddies.loop_inst_eddies: Date:", d)
        print("inst_eddies.loop_inst_eddies: Input file:", nc_file, flush=True)

        if os.access(nc_file, os.F_OK):
            process_1_file(
                bbox,
                nco_instance,
                nc_file,
                inst_eddies_exe,
                d,
                my_slice,
                writer,
                SHPC_dir,
                config_nml_string,
                periodic,
            )
        else:
            print("inst_eddies.loop_inst_eddies: Missing file:", nc_file)

            for orient in ["Cyclones", "Anticyclones"]:
                with open(
                    f"{SHPC_dir}/Slice_{my_slice}/{orient}/ishape_last.txt", "a"
                ) as ishape_last, shapefile.Reader(
                    f"{SHPC_dir}/Slice_{my_slice}/{orient}/" "extremum"
                ) as shpr:
                    n_shapes = len(shpr)
                    ishape_last.write(str(n_shapes - 1) + "\n")

        d += 1

    perf_report.close()
    os.remove("unpacked.nc")

    if bbox:
        os.remove("cropped.nc")

    if my_slice == 0 and not os.access(f"{SHPC_dir}/n_slices.txt", os.F_OK):
        with open(f"{SHPC_dir}/n_slices.txt", "w") as f:
            f.write("1\n")


if __name__ == "__main__":
    import argparse

    argparser = argparse.ArgumentParser()
    argparser.add_argument("SHPC_dir")
    argparser.add_argument(
        "file",
        nargs="+",
        help="NetCDF file containing SSH and velocity " "at a single date",
    )
    argparser.add_argument(
        "-b",
        "--bbox",
        nargs=4,
        type=float,
        metavar=("xmin", "xmax", "ymin", "ymax"),
    )
    args = argparser.parse_args()

    if not os.access(args.file[0], os.F_OK):
        sys.exit(args.file[0] + " not found")

    if os.access("inst_eddies_nml.txt", os.R_OK):
        print("Will use inst_eddies_nml.txt.")
    else:
        sys.exit('"inst_eddies_nml.txt" not found -- Aborting')

    with open("inst_eddies_nml.txt") as f_obj:
        config_nml_string = f_obj.read()

    d = input("Enter first date index (integer value): ")
    d = int(d)
    my_slice = input("Enter slice index (integer value): ")
    my_slice = int(my_slice)
    periodic = input("periodic (t/f)? ")
    loop_inst_eddies(
        args.SHPC_dir,
        iter(args.file),
        args.bbox,
        config_nml_string,
        d,
        my_slice,
        periodic,
    )