Read ZMap+ grids to NumPy arrays + metadata

zmapgrid.py

# Copyright (c) 2024 dwt | 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 enum import Enum, auto
from pathlib import Path
from typing import Callable, NamedTuple
import sys

import numpy as np
import numpy.typing as npt


class ZMapGrid(NamedTuple):
    xmin: float
    xmax: float
    ymin: float
    ymax: float
    rows: int
    cols: int
    grid: npt.NDArray[np.float64]

    def plot(self, pixel_is_point: bool = False) -> None:
        import matplotlib.pyplot as plt
        if pixel_is_point:
            extents = self.extents_pixel_is_point()
        else:
            extents = self.extents_pixel_is_area()
        plt.imshow(self.grid, extent=extents, origin='upper')
        plt.colorbar()
        plt.show()

    # (x, y)
    def pixel_size_pixel_is_area(self) -> tuple[float, float]:
        xsize = (self.xmax - self.xmin) / self.cols
        ysize = (self.ymax - self.ymin) / self.rows
        return xsize, ysize

    # (x, y)
    def pixel_size_pixel_is_point(self) -> tuple[float, float]:
        xsize = (self.xmax - self.xmin) / (self.cols - 1)
        ysize = (self.ymax - self.ymin) / (self.rows - 1)
        return xsize, ysize

    # (xmin, xmax, ymin, ymax)
    def extents_pixel_is_area(self) -> tuple[float, float, float, float]:
        return self.xmin, self.xmax, self.ymin, self.ymax

    # (xmin, xmax, ymin, ymax)
    def extents_pixel_is_point(self) -> tuple[float, float, float, float]:
        xsize, ysize = self.pixel_size_pixel_is_point()
        return (
            self.xmin - 0.5 * xsize,
            self.xmax + 0.5 * xsize,
            self.ymin - 0.5 * ysize,
            self.ymax + 0.5 * ysize
        )


class ParseState(Enum):
    BEGIN = auto()
    HEADER1 = auto()
    HEADER2 = auto()
    HEADER3 = auto()
    HEADEREND = auto()
    DATA = auto()

def load_zmap_grid(path: Path, check_nodes: bool = True) -> ZMapGrid:
    nodes = 0
    missing = ''
    parse_float: Callable[[str], float] = float
    rows = 0
    cols = 0
    xmin = np.nan
    xmax = np.nan
    ymin = np.nan
    ymax = np.nan
    state = ParseState.BEGIN
    expect_lens: list[int] = []
    cyc = 0
    data: list[float] = []

    with open(path, 'r') as f:
        # handle tagging with line #s, comment lines and line endings
        lines = (
            (i + 1, line.rstrip())
            for i, line in enumerate(f)
            if not line.startswith('!')
        )

        for num, line in lines:
            # tokenize by , or whitespace depdending on file section
            if state == ParseState.DATA:
                toks = line.split()
            else:
                toks = [w.strip() for w in line.split(',')]

            match state:
                case ParseState.BEGIN:
                    match toks:
                        case [udf, 'GRID', s_nodes]:
                            if not udf.startswith('@'):
                                raise ValueError(
                                  f'invalid grid file {path}:\n'
                                  + f'line {num} does not begin with "@"'
                                )
                            nodes = int(s_nodes)
                            state = ParseState.HEADER1
                        case _:
                            raise ValueError(
                                f'invalid grid file {path}:\n'
                                + f'line {num} is not "@..., GRID, N"'
                            )

                case ParseState.HEADER1:
                    if len(toks) != 5:
                        raise ValueError(
                            f'invalid grid file {path}:\n'
                            + f'line {num} does not contain 5 entries'
                        )
                    _width = int(toks[0])
                    missing = toks[1]
                    if missing == '':
                        missing = toks[2]
                    elif toks[2] != '':
                        raise ValueError(
                            f'invalid grid file {path}:\n'
                            + f'line {num} specifies missing number '
                            + 'AND text'
                        )
                    ndec = int(toks[3])
                    parse_float = lambda s: float(s) if '.' in s else float(s[:len(s) - ndec] + '.' + s[-ndec:])
                    _startcol = int(toks[4])
                    state = ParseState.HEADER2

                case ParseState.HEADER2:
                    if len(toks) != 6:
                        raise ValueError(
                            f'invalid grid file {path}:\n'
                            + f'line {num} does not contain 6 entries'
                        )
                    rows = int(toks[0])
                    cols = int(toks[1])
                    xmin = float(toks[2])
                    xmax = float(toks[3])
                    ymin = float(toks[4])
                    ymax = float(toks[5])
                    expect_lens = [nodes] * (rows // nodes)
                    if rows % nodes != 0:
                        expect_lens += [rows % nodes]
                    state = ParseState.HEADER3

                case ParseState.HEADER3:
                    match toks:
                        case ['0.0', '0.0', '0.0']:
                            state = ParseState.HEADEREND
                        case _:
                            raise ValueError(
                                f'invalid grid file {path}:\n'
                                + f'line {num} was not "0.0, 0.0, 0.0"'
                            )

                case ParseState.HEADEREND:
                    match toks:
                        case ['@']:
                            state = ParseState.DATA
                        case _:
                            raise ValueError(
                                f'invalid grid file {path}:\n'
                                + f'line {num} was not "@"'
                            )

                case ParseState.DATA:
                    if check_nodes and len(toks) != expect_lens[cyc]:
                        raise ValueError(
                            f'invalid grid file {path}:\n'
                            + f'line {num} had {len(toks)} entries;'
                            + f' expected {expect_lens[cyc]}'
                        )
                    for tok in toks:
                        if tok == missing:
                            data.append(np.nan)
                        else:
                            data.append(parse_float(tok))
                    cyc = (cyc + 1) % len(expect_lens)

    grid = np.array(data, dtype=np.float64, order='F'
      ).reshape((rows, cols), order='F')

    return ZMapGrid(
        xmin=xmin,
        xmax=xmax,
        ymin=ymin,
        ymax=ymax,
        rows=rows,
        cols=cols,
        grid=grid
    )