WIP: reverse engineering (parts of) the Petra .GRD format

petra_grid_viewer.py

# Copyright (c) 2023 Derrick W. Turk / terminus, LLC

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


from datetime import datetime, timedelta
from enum import Enum
import sys
import struct
import warnings

import numpy as np
import numpy.lib.stride_tricks as nplst
from numpy.typing import NDArray

from matplotlib.tri import Triangulation # type: ignore
import matplotlib.pyplot as plt # type: ignore

from typing import NamedTuple


_DELPHI_DATE_ORIGIN = datetime(1899, 12, 30)


class UnitOfMeasure(Enum):
    Feet = 0
    Meters = 1


class Grid(NamedTuple):
    version: int
    name: str

    xmin: float
    xmax: float
    ymin: float
    ymax: float
    xstep: float
    ystep: float
    zmin: float
    zmax: float
    xyunits: UnitOfMeasure
    zunits: UnitOfMeasure

    created_date: datetime
    source_data: str
    unknown_metadata: str
    projection: str
    datum: str

    # we know these exist, but not necessarily what they "mean"!
    gridmethod: int
    projection_code: int

    # same...
    cm: float
    rlat: float

    # shape (rows, cols)
    grid: NDArray[np.float64]

    def plot(self) -> None:
        plt.imshow(self.grid,
          extent=(self.xmin, self.xmax, self.ymin, self.ymax),
          origin='lower')
        plt.title(f'{self.name}')
        plt.colorbar()
        plt.show()


class TriangulatedGrid(NamedTuple):
    version: int
    name: str

    # these are present, but meaningless for triangulated grids
    # (I think they reflect the original "as-specified" grid prior to
    #   triangulation)
    original_size: int
    original_rows: int
    original_cols: int

    xmin: float
    xmax: float
    ymin: float
    ymax: float
    xstep: float
    ystep: float
    zmin: float
    zmax: float
    xyunits: UnitOfMeasure
    zunits: UnitOfMeasure

    created_date: datetime
    source_data: str
    unknown_metadata: str
    projection: str
    datum: str

    # we know these exist, but not necessarily what they "mean"!
    gridmethod: int
    projection_code: int

    # same...
    cm: float
    rlat: float

    # an array of triangles, each with 3 vertices having 3 dimensions (x, y, z)
    # shape (n_triangles, 3, 3)
    triangles: NDArray[np.float64]

    def plot(self) -> None:
        zs = np.ravel(self.triangles[:, :, 2])
        fig, ax = plt.subplots()
        ax.set_aspect('equal')
        tc = ax.tripcolor(self._triangulation(), zs)
        # TODO: adjust aspect ratio?
        ax.set_title(f'{self.name}')
        fig.colorbar(tc)
        plt.show()

    def _triangulation(self) -> Triangulation:
        xs = np.ravel(self.triangles[:, :, 0])
        ys = np.ravel(self.triangles[:, :, 1])
        ixs = np.reshape(np.indices(xs.shape)[0], (self.triangles.shape[0], 3))
        return Triangulation(xs, ys, ixs)


def _asciz_string(buf: bytes) -> str:
    return buf[:buf.index(b'\x00')].decode('ascii')


def _delphi_datetime(days: float) -> datetime:
    return _DELPHI_DATE_ORIGIN + timedelta(days=days)


def grid_from_buffer(buf: bytes) -> Grid | TriangulatedGrid:
    (ver, name, sz, xmin, xmax, ymin, ymax,
      xstep, ystep, zmin, zmax) = struct.unpack_from('<I81sI8d', buf, 0x0)
    cm, rlat = struct.unpack_from('<2d', buf, 0xb9)
    (rows, cols, meth, proj_code, xyunits) = struct.unpack_from(
      '<5I', buf, 0x3fd)
    zunits, = struct.unpack_from('<I', buf, 0x429)
    n_triangles, = struct.unpack_from('<I', buf, 0x431)

    if ver != 2:
        warnings.warn(f'unknown version(?): {ver}')

    if rows * cols != sz:
        raise ValueError(
          f'total size {sz} != {rows} x {cols}')

    if (np.abs(xmin + (cols - 1) * xstep - xmax) / xmax > 0.0001):
        raise ValueError(
          f'invalid x spec: {xmin} to {xmax} by {xstep} but {cols} cols')

    if (np.abs(ymin + (rows - 1) * ystep - ymax) / ymax > 0.0001):
        raise ValueError(
          f'invalid y spec: {ymin} to {ymax} by {ystep} but {rows} rows')

    if n_triangles == 0 and (len(buf) - 0x119c) / 8 != sz:
        raise ValueError(
          f"data can't begin at offset 0x119c and have size {sz}")

    if n_triangles > 0 and (len(buf) - 0x119c) / 8 / 9 != n_triangles:
        raise ValueError(
          f"data can't begin at offset 0x119c and have {n_triangles} triangles")

    (src,) = struct.unpack_from('<246s', buf, 0x5b9)
    (days,) = struct.unpack_from('<d', buf, 0xe1)
    (unk, proj, datum) = struct.unpack_from('<2009s65s195s', buf, 0x8bf)

    if n_triangles == 0:
        grid = np.frombuffer(buf,
          dtype=np.dtype(float).newbyteorder('<'), offset=0x119c).copy()
        grid[grid == 1e30] = np.nan
        grid = grid.reshape(rows, cols)
        return Grid(
            ver,
            _asciz_string(name),
            xmin,
            xmax,
            ymin,
            ymax,
            xstep,
            ystep,
            zmin,
            zmax,
            UnitOfMeasure(xyunits),
            UnitOfMeasure(zunits),
            _delphi_datetime(days),
            _asciz_string(src),
            _asciz_string(unk),
            _asciz_string(proj),
            _asciz_string(datum),
            meth,
            proj_code,
            cm,
            rlat,
            grid,
        )
    else:
        raw = np.frombuffer(buf,
          dtype=np.dtype(float).newbyteorder('<'), offset=0x119c)
        tris = nplst.as_strided(raw, (n_triangles, 3, 3), (72, 8, 24)).copy()
        tris[tris == 1e30] = np.nan
        return TriangulatedGrid(
            ver,
            _asciz_string(name),
            sz,
            rows,
            cols,
            xmin,
            xmax,
            ymin,
            ymax,
            xstep,
            ystep,
            zmin,
            zmax,
            UnitOfMeasure(xyunits),
            UnitOfMeasure(zunits),
            _delphi_datetime(days),
            _asciz_string(src),
            _asciz_string(unk),
            _asciz_string(proj),
            _asciz_string(datum),
            meth,
            proj_code,
            cm,
            rlat,
            tris,
        )


def main(argv: list[str]) -> int:
    if len(argv) <= 1:
        print(f'Usage: {argv[0] if argv else "grid_viewer"} grd-files',
          file=sys.stderr)
        return 2

    failed = False
    for grid_path in argv[1:]:
        with open(grid_path, 'rb') as f:
            buf = f.read()
            try:
                grid = grid_from_buffer(buf)
                print(grid)
                grid.plot()
            except ValueError as e:
                print(f'{grid_path}: {e}', file=sys.stderr)
                failed = True

    if failed:
        return 1
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv))


__all__ = [
    'Grid',
    'TriangulatedGrid',
    'grid_from_buffer',
]