extract_anl_var (grib2)ΒΆ

#!/usr/bin/env python

# Extract 20CRv3 analysis variable from the full output

import os
import sys
import subprocess
import tempfile
import datetime

# Version and month
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--startyear", help="Year run started", type=int, required=True)
parser.add_argument("--year", help="Year to extract", type=int, required=True)
parser.add_argument("--month", help="Month to extract", type=int, required=True)
parser.add_argument("--version", help="Month to extract", type=int, default=451)
parser.add_argument("--var", help="Variable to extract", type=str, required=True)
parser.add_argument(
    "--level",
    help="Pressure level (hPa) to extract at",
    type=str,
    default=None,
    required=False,
    choices=[
        "1000",
        "975",
        "950",
        "925",
        "900",
        "850",
        "800",
        "750",
        "700",
        "650",
        "600",
        "550",
        "500",
        "450",
        "400",
        "350",
        "300",
        "250",
        "200",
        "150",
        "100",
        "70",
        "50",
        "30",
        "20",
        "10",
        "5",
        "1",
    ],
)
parser.add_argument(
    "--ilevel",
    help="Isentropic level (K) to extract at",
    type=str,
    default=None,
    required=False,
    choices=["300", "330", "350"],
)
parser.add_argument(
    "--height",
    help="Height (m) to extract at",
    type=str,
    default=None,
    required=False,
    choices=[
        "2",
        "10",
        "12",
        "20",
        "30",
        "50",
        "80",
        "100",
        "150",
        "200",
        "250",
        "300",
        "500",
    ],
)

args = parser.parse_args()

# Set default heights, where appropriate
if args.var == "air.2m":
    args.level = None
    args.height = "2"
    args.var = "TMP"
if args.var == "uwnd.10m":
    args.level = None
    args.height = "10"
    args.var = "UGRD"
if args.var == "vwnd.10m":
    args.level = None
    args.height = "10"
    args.var = "VGRD"


# Make an output file name
def opfile(var, level, height, ilevel):
    var = var.lower()
    if var == "prmsl":
        return "prmsl"
    if var == "air.sfc":
        return "air.sfc"
    if var == "air.2m":
        return "air.2m"
    if var == "tmp" and height is not None and height == 2:
        return "air.2m"
    if var == "uwnd" or var == "ugrd":
        if level is not None:
            return "uwnd.%dmb" % int(level)
        elif ilevel is not None:
            return "uwnd.%dK" % int(ilevel)
        elif height is not None:
            return "uwnd.%dm" % int(height)
        else:
            raise ValueError("Either height, level, or ilevel must be specified")
    if var == "vwnd" or var == "vgrd":
        if level is not None:
            return "vwnd.%dmb" % int(level)
        elif ilevel is not None:
            return "vwnd.%dK" % int(ilevel)
        elif height is not None:
            return "vwnd.%dm" % int(height)
        else:
            raise ValueError("Either height, level, or ilevel must be specified")
    if var == "icec":
        return "icec"

    # Default - specify grib2 var directly
    if level is not None:
        return "%s.%dmb" % (var.lower(), int(level))
    elif ilevel is not None:
        return "%s.%dK" % (var.lower(), int(ilevel))
    elif height is not None:
        return "%s.%dm" % (var.lower(), int(height))
    else:
        raise ValueError("Either height, level, or ilevel must be specified")


# Make a search string
def search_string(var, level, height, ilevel):
    if var == "prmsl":
        return "PRMSL"
    if var == "air.sfc":
        return "TMP:surface"
    if var == "uwnd" or var == "ugrd":
        if level is not None:
            return "UGRD:%d mb" % int(level)
        elif ilevel is not None:
            return "UGRD:%d K" % int(ilevel)
        elif height is not None:
            return "UGRD:%d m above ground" % int(height)
        else:
            raise ValueError("Either height, level, or ilevel must be specified")
    if var == "vwnd" or var == "vgrd":
        if level is not None:
            return "VGRD:%d mb" % int(level)
        elif ilevel is not None:
            return "VGRD:%d K" % int(ilevel)
        elif height is not None:
            return "VGRD:%d m above ground" % int(height)
        else:
            raise ValueError("Either height, level, or ilevel must be specified")
    if var == "icec":
        return "ICEC"

    # Default - specify grib2 var directly
    if level is not None:
        return "%s:%d mb" % (var.upper(), int(level))
    elif ilevel is not None:
        return "%s:%d K" % (var.upper(), int(ilevel))
    elif height is not None:
        return "%s:%d m above ground" % (var.upper(), int(height))
    else:
        raise ValueError("Either height, level, or ilevel must be specified")


# Where to find the grib (and obs) files retrieved from hsi
working_directory = "%s/20CRv3.working/ensda_%04d/%04d/%02d" % (
    os.getenv("SCRATCH"),
    args.startyear,
    args.year,
    args.month,
)
if not os.path.isdir(working_directory):
    os.makedirs(working_directory)
# Where to put the final output files for this month
final_directory = "%s/20CRv3.final/version_%1d.%1d.%1d/%04d/%02d" % (
    os.getenv("SCRATCH"),
    int(args.version / 100),
    int((args.version % 100) / 10),
    int(args.version % 10),
    args.year,
    args.month,
)
if not os.path.isdir(final_directory):
    os.makedirs(final_directory)


# Don't repeat pre-existing extractions
fn = "%s/%s.nc4" % (
    final_directory,
    opfile(args.var, args.level, args.height, args.ilevel),
)
if os.path.isfile(fn):
    raise Exception("Already done")

# Temporary file for staging extracted data
tfile = tempfile.NamedTemporaryFile(delete=False)

current_day = datetime.datetime(args.year, args.month, 1, 0)
while current_day.month == args.month:
    # Extract grids every 3 hours and concatenate to output
    for hour in range(0, 24, 3):
        for member in range(1, 81):
            an_file_name = "%s/pgrbanl_%04d%02d%02d%02d_mem%03d.grb2" % (
                working_directory,
                current_day.year,
                current_day.month,
                current_day.day,
                hour,
                member,
            )
            if not os.path.exists(an_file_name):
                raise Exception("Missing data %s" % an_file_name)

            proc = subprocess.Popen(
                "wgrib2 %s -match '%s' -grib %s; cat %s >> %s/%s.grb2"
                % (
                    an_file_name,
                    search_string(args.var, args.level, args.height, args.ilevel),
                    tfile.name,
                    tfile.name,
                    final_directory,
                    opfile(args.var, args.level, args.height, args.ilevel),
                ),
                shell=True,
            )
            (out, err) = proc.communicate()
            if out is not None or err is not None:
                raise StandardError(
                    "Failed to extract %s from %s"
                    % (
                        opfile(args.var, args.level, args.height, args.ilevel),
                        an_file_name,
                    )
                )

    current_day = current_day + datetime.timedelta(days=1)
os.remove(tfile.name)

# Convert to netCDF
proc = subprocess.Popen(
    "ncl_convert2nc %s.grb2 -i %s -o %s -L -nc4c -cl 5"
    % (
        opfile(args.var, args.level, args.height, args.ilevel),
        final_directory,
        final_directory,
    ),
    shell=True,
)
(out, err) = proc.communicate()
if out is not None or err is not None:
    raise StandardError(
        "Failed to convert %s to netCDF"
        % opfile(args.var, args.level, args.height, args.ilevel)
    )