wrf-python/test/misc/wps.py

# Hastily made script to read WPS intermediate files
from __future__ import print_function
import struct

import numpy as np

# Number of bytes used at the start and end of a fortran record to
# indicate the record size
SIZE_BYTES = 4


class WPSData(object):
    def __init__(self, ifv=None, hdate=None, xfcst=None, map_source=None,
                 field=None, units=None, desc=None, xlvl=None, nx=None,
                 ny=None, iproj=None, startloc=None, startlat=None,
                 startlon=None, deltalat=None, deltalon=None, nlats=None,
                 dx=None, dy=None, xlonc=None, truelat1=None, truelat2=None,
                 earth_radius=None, is_wind_earth_rel=None, slab=None):

        self.ifv = ifv
        self.hdate = hdate
        self.xfcst = xfcst
        self.map_source = map_source
        self.field = field
        self.units = units
        self.desc = desc
        self.xlvl = xlvl
        self.nx = nx
        self.ny = ny
        self.iproj = iproj
        self.startloc = startloc
        self.startlat = startlat
        self.startlon = startlon
        self.deltalat = deltalat
        self.deltalon = deltalon
        self.nlats = nlats
        self.dx = dx
        self.dy = dy
        self.xlonc = xlonc
        self.truelat1 = truelat1
        self.truelat2 = truelat2
        self.earth_radius = earth_radius
        self.is_wind_earth_rel = is_wind_earth_rel
        self.slab = slab


def _parse_record1(data):
    result = {}
    result["ifv"] = struct.unpack(">i", data)

    return result


def _parse_record2(data):
    result = {}
    parsed = struct.unpack(">24sf32s9s25s46sfiii", data)
    result["hdate"] = parsed[0]
    result["xfcst"] = parsed[1]
    result["map_source"] = parsed[2]
    result["field"] = parsed[3]
    result["units"] = parsed[4]
    result["desc"] = parsed[5]
    result["xlvl"] = parsed[6]
    result["nx"] = parsed[7]
    result["ny"] = parsed[8]
    result["iproj"] = parsed[9]

    return result


def _parse_record3(data, iproj):
    result = {}
    if iproj == 0:
        fmt = ">8sfffff"
        parsed = struct.unpack(fmt, data)
        result["startloc"] = parsed[0]
        result["startlat"] = parsed[1]
        result["startlon"] = parsed[2]
        result["deltalat"] = parsed[3]
        result["deltalon"] = parsed[4]
        result["earth_radius"] = parsed[5]
    elif iproj == 1:
        fmt = ">8sffffff"
        parsed = struct.unpack(fmt, data)
        result["startloc"] = parsed[0]
        result["startlat"] = parsed[1]
        result["startlon"] = parsed[2]
        result["dx"] = parsed[3]
        result["dy"] = parsed[4]
        result["truelat1"] = parsed[5]
        result["earth_radius"] = parsed[6]
    elif iproj == 3:
        fmt = ">8sffffffff"
        parsed = struct.unpack(fmt, data)
        result["startloc"] = parsed[0]
        result["startlat"] = parsed[1]
        result["startlon"] = parsed[2]
        result["dx"] = parsed[3]
        result["dy"] = parsed[4]
        result["xlonc"] = parsed[5]
        result["truelat1"] = parsed[6]
        result["truelat2"] = parsed[7]
        result["earth_radius"] = parsed[8]
    elif iproj == 4:
        fmt = ">8sfffff"
        parsed = struct.unpack(fmt, data)
        result["startloc"] = parsed[0]
        result["startlat"] = parsed[1]
        result["startlon"] = parsed[2]
        result["nlats"] = parsed[3]
        result["deltalon"] = parsed[4]
        result["earth_radius"] = parsed[5]
    elif iproj == 5:
        fmt = ">8sfffffff"
        parsed = struct.unpack(fmt, data)
        result["startloc"] = parsed[0]
        result["startlat"] = parsed[1]
        result["startlon"] = parsed[2]
        result["dx"] = parsed[3]
        result["dy"] = parsed[4]
        result["xlonc"] = parsed[5]
        result["truelat1"] = parsed[6]
        result["earth_radius"] = parsed[7]

    return result


def _parse_record4(data):
    result = {}
    result["is_wind_earth_rel"] = struct.unpack(">i", data)

    return result


def _parse_record5(data, nx, ny):
    result = {}

    size = nx * ny
    fmt = ">{}f".format(size)
    parsed = struct.unpack(fmt, data)

    arr = np.array(parsed, dtype=np.float32)
    result["slab"] = arr.reshape((nx, ny), order="F")

    return result


_PARSE_MAP = {0: _parse_record1,
              1: _parse_record2,
              2: _parse_record3,
              3: _parse_record4,
              4: _parse_record5}


def parse_record(record_idx, data, iproj=None, nx=None, ny=None):

    if record_idx == 0 or record_idx == 1 or record_idx == 3:
        return _PARSE_MAP[record_idx](data)
    elif record_idx == 2:
        return _PARSE_MAP[record_idx](data, iproj)
    elif record_idx == 4:
        return _PARSE_MAP[record_idx](data, nx, ny)


def read_wps(wps_file, field=None):
    with open(wps_file, "rb") as f:
        wps_params = {}
        record_list = []
        raw_data = f.read()

        record_size_idx = 0
        end_of_file_idx = len(raw_data) - 1

        while True:
            iproj = None
            nx = None
            ny = None
            keep_record = True
            for record_idx in range(5):
                # Each record has the size (in SIZE_BYTES bytes) at the
                # start and end of each record.  This might be compiler
                # dependent though, so this might need to be modified.
                # Also, the WPS files are stored big endian.

                record_size = struct.unpack(
                    ">i",
                    raw_data[record_size_idx: record_size_idx + SIZE_BYTES])
                record_start = record_size_idx + SIZE_BYTES
                record_end = record_start + record_size[0]
                record_data = raw_data[record_start:record_end]

                parsed_record = parse_record(record_idx, record_data, iproj,
                                             nx, ny)

                try:
                    field_name = parsed_record["field"].strip()
                except KeyError:
                    pass
                else:
                    if field is not None:
                        if field_name != field:
                            keep_record = False

                try:
                    iproj = parsed_record["iproj"]
                except KeyError:
                    pass

                try:
                    nx = parsed_record["nx"]
                except KeyError:
                    pass

                try:
                    ny = parsed_record["ny"]
                except KeyError:
                    pass

                wps_params.update(parsed_record)

                record_size_idx = record_end + SIZE_BYTES

            if keep_record:
                record_list.append(WPSData(**wps_params))

            # Repeat for all record slabs
            if record_end + SIZE_BYTES > end_of_file_idx:
                break

    return record_list