zmap2xyz.py

"""
Converts one or more zmap grids to ascii x,y,z tab-delimited files.  Exported
files will be saved to the current directory with a ".xyz" extension and the
same file name as the original.  Null grid values will not be included unless
the "--with_nulls" option is specified.

Usage:
    zmap2xyz.py [options] INPUT_FILES...

Options:
    -h --help       Show this message
    --with_nulls    Include null grid points in the output
"""
import os
import numpy as np
from docopt import docopt

def main(args):
    for filename in args['INPUT_FILES']:
        process_file(filename, args['--with_nulls'])

def process_file(filename, include_nulls=False):
    x, y, z = grid2xyz(filename)
    if include_nulls:
        x, y, z = [item.ravel() for item in [x, y, z]]
    else:
        x, y, z = [item.compressed() for item in [x, y, z]]
    base, ext = os.path.splitext(filename)
    np.savetxt(base + '.xyz', np.vstack([x, y, z]).T)

def grid2xyz(filename):
    grid = ZmapGridReader(filename)
    # Note the reversed y-order!! (Grid is flipped up/down)
    y, x = np.mgrid[grid.ymax:grid.ymin:1j*grid.nrows,
                    grid.xmin:grid.xmax:1j*grid.ncols]
    x = np.ma.masked_array(x, grid.data.mask)
    y = np.ma.masked_array(y, grid.data.mask)
    return x, y, grid.data


class ZmapGridReader(object):
    def __init__(self, filename):
        self.filename = filename
        with open(self.filename, 'r') as infile:
            values = self._parse_header(infile)
        null, nrows, ncols, xmin, xmax, ymin, ymax = values
        self.null = float(null)
        self.nrows, self.ncols = int(nrows), int(ncols)
        self.xmin, self.xmax = float(xmin), float(xmax)
        self.ymin, self.ymax = float(ymin), float(ymax)
        self.data = self._read_values()

    def _skip_comments(self, infile):
        while True:
            last_position = infile.tell()
            if not infile.readline()[0].startswith('!'):
                break
        infile.seek(last_position)

    def _parse_header(self, infile):
        def line(infile):
            return infile.readline().replace(',', ' ').strip().split()
        infile.seek(0)
        self._skip_comments(infile)
        header = infile.readline()
        width, null = line(infile)[:2]
        nrows, ncols, xmin, xmax, ymin, ymax = line(infile)[:6] # Trailing ,
        return null, nrows, ncols, xmin, xmax, ymin, ymax

    def _read_values(self):
        def value_iter(infile):
            for line in infile:
                line = line.strip().split()
                for item in line:
                    yield float(item)

        with open(self.filename, 'r') as infile:
            for line in infile:
                if line.strip() == '@':
                    break
            data = np.fromiter(value_iter(infile), dtype=float)
        data = data.reshape(self.ncols, self.nrows).T
        return np.ma.masked_equal(data, self.null)

if __name__ == '__main__':
    main(docopt(__doc__))