pyart.io.output_to_geotiff — pyart-mch 0.4.1 documentation

# Source code for pyart.io.output_to_geotiff

    """
    pyart.io.write_grid_geotiff
    ===========================

    Write a Py-ART Grid object to a GeoTIFF file.

    .. autosummary::
        :toctree: generated/

        write_grid_geotiff
        _get_rgb_values
        _create_sld

    """

    import os
    import shutil

    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.colors as colors
    from ..exceptions import MissingOptionalDependency
    try:
        from osgeo import gdal
        IMPORT_FLAG = True
    except ImportError:
        IMPORT_FLAG = False

    def write_grid_geotiff(grid, filename, field, rgb=False, level=None,
                           cmap='viridis', vmin=0, vmax=75, color_levels=None,
                           warp=False, sld=False, use_doublequotes=False):
        """
        Write a Py-ART Grid object to a GeoTIFF file.

        The GeoTIFF can be the standard Azimuthal Equidistant projection used
        in Py-ART, or a lat/lon projection on a WGS84 sphere. The latter is
        typically more usable in web mapping applications. The GeoTIFF can
        contain a single float-point raster band, or three RGB byte raster bands.
        The former will require an SLD file for colorful display using standard
        GIS or web mapping software, while the latter will show colors
        "out-of-the-box" but lack actual data values. The function also can
        output an SLD file based on the user-specified inputs. User can specify
        the 2D vertical level to be output. If this is not specified, a 2D
        composite is created. User also can specify the field to output.

        This function requires GDAL Python libraries to be installed. These are
        available via conda; e.g., 'conda install gdal'

        Parameters
        ----------
        grid : pyart.core.Grid object
            Grid object to write to file.
        filename : str
            Filename for the GeoTIFF.
        field : str
            Field name to output to file.

        Other Parameters
        ----------------
        rbg : bool, optional
            True - Output 3-band RGB GeoTIFF

            False - Output single-channel, float-valued GeoTIFF. For display,
                    likely will need an SLD file to provide a color table.

        level : int or None, optional
            Index for z-axis plane to output. None gives composite values
            (i.e., max in each vertical column).
        cmap : str or matplotlib.colors.Colormap object, optional
            Colormap to use for RGB output or SLD file.
        vmin : int or float, optional
            Minimum value to color for RGB output or SLD file.
        vmax : int or float, optional
            Maximum value to color for RGB output or SLD file.
        color_levels : int or None, optional
            Number of color levels in cmap. Useful for categorical colormaps
            with steps << 255 (e.g., hydrometeor ID).
        warp : bool, optional
            True - Use gdalwarp (called from command line using os.system)
                   to warp to a lat/lon WGS84 grid.

            False - No warping will be performed. Output will be Az. Equidistant.

        sld : bool, optional
            True - Create a Style Layer Descriptor file (SLD) mapped to vmin/vmax
                   and cmap. File is named same as output TIFF, except for .sld
                   extension.

            False - Don't do this.

        use_doublequotes : bool, optional
            True - Use double quotes in the gdalwarp call (requires warp=True),
                   which may help if that command is producing and error like:
                   'Translating source or target SRS failed'.

            False - Use single quotes instead.

        """
        if not IMPORT_FLAG:
            raise MissingOptionalDependency(
                'GDAL not detected, GeoTIFF output failure!')

        if field not in grid.fields.keys():
            raise KeyError('Failed -', field, 'field not found in Grid object.')

        # Determine whether filename template already contains a suffix
        # If not, append an appropriate one.
        if '.' not in filename:
            name = filename
            end = 'tif'
            ofile = name + "." + end
        else:
            ofile = filename
        nz, ny, nx = grid.fields[field]['data'].shape
        dist = max(grid.x['data'])
        rangestep = grid.x['data'][1] - grid.x['data'][2]
        lat = grid.origin_latitude['data'][0]
        lon = grid.origin_longitude['data'][0]
        # Check if masked array; if so, fill missing data
        filled = np.ma.filled(grid.fields[field]['data'], fill_value=-32768)
        if level is None:
            data = np.amax(filled, 0)
        else:
            data = filled[level]
        data = data.astype(float)
        data[data == -32768] = np.nan

        iproj = 'PROJCS["unnamed",GEOGCS["WGS 84",DATUM["unknown",' + \
            'SPHEROID["WGS84",6378137,298.257223563]],' + \
            'PRIMEM["Greenwich",0],' + \
            'UNIT["degree",0.0174532925199433]],' + \
            'PROJECTION["Azimuthal_Equidistant"],' + \
            'PARAMETER["latitude_of_center",' + str(lat) + '],' + \
            'PARAMETER["longitude_of_center",' + str(lon) + '],' + \
            'PARAMETER["false_easting",0],' + \
            'PARAMETER["false_northing",0],' + \
            'UNIT["metre",1,AUTHORITY["EPSG","9001"]]]'
        out_driver = gdal.GetDriverByName("GTiff")

        # Output dataset depends on rgb flag
        if not rgb:
            # Single-channel, floating-point output
            dst_options = ['COMPRESS=LZW', 'ALPHA=YES']
            dst_ds = out_driver.Create(
                ofile, data.shape[1], data.shape[0], 1,
                gdal.GDT_Float32, dst_options)
        else:
            # Assign data RGB levels based on value relative to vmax/vmin
            rarr, garr, barr = _get_rgb_values(
                data, vmin, vmax, color_levels, cmap)
            dst_ds = out_driver.Create(ofile, data.shape[1],
                                       data.shape[0], 3, gdal.GDT_Byte)

        # Common Projection and GeoTransform
        dst_ds.SetGeoTransform([-dist, -rangestep, 0, dist, 0, rangestep])
        dst_ds.SetProjection(iproj)

        # Final output depends on rgb flag
        if not rgb:
            dst_ds.GetRasterBand(1).WriteArray(data[::-1, :])
        else:
            dst_ds.GetRasterBand(1).WriteArray(rarr[::-1, :])
            dst_ds.GetRasterBand(2).WriteArray(garr[::-1, :])
            dst_ds.GetRasterBand(3).WriteArray(barr[::-1, :])
        dst_ds.FlushCache()
        dst_ds = None

        if sld:
            _create_sld(cmap, vmin, vmax, ofile, color_levels)

        if warp:
            # Warps TIFF to lat/lon WGS84 projection that is more useful
            # for web mapping applications. Likely changes array shape.
            if use_doublequotes:
                os.system('gdalwarp -q -t_srs \"+proj=longlat +ellps=WGS84 ' +
                          '+datum=WGS84 +no_defs\" ' + ofile + ' ' +
                          ofile + '_tmp.tif')
            else:
                os.system('gdalwarp -q -t_srs \'+proj=longlat +ellps=WGS84 ' +
                          '+datum=WGS84 +no_defs\' ' + ofile + ' ' +
                          ofile + '_tmp.tif')
            shutil.move(ofile+'_tmp.tif', ofile)

    def _get_rgb_values(data, vmin, vmax, color_levels, cmap):
        """
        Get RGB values for later output to GeoTIFF, given a 2D data field,
        display min/max and color table info. Missing data get numpy.nan.
        Only called if rgb is True in write_grid_geotiff.

        Parameters
        ----------
        data : numpy.ndarray object, dtype int or float
            Two-dimensional data array.
        vmin : int or float
            Minimum value to color for RGB output or SLD file.
        vmax : int or float
            Maximum value to color for RGB output or SLD file.
        color_levels : int
            Number of color levels in cmap. Useful for categorical colormaps
            with steps << 255 (e.g., hydrometeor ID).
        cmap : str or matplotlib.colors.Colormap object, optional
            Colormap to use for RGB output or SLD file.

        Returns
        -------
        rarr : numpy.ndarray object, dtype int
            Red channel indices (range = 0-255).
        barr : numpy.ndarray object, dtype int
            Blue channel indices (range = 0-255).
        garr : numpy.ndarray object, dtype int
            Green channel indices (range = 0-255).

        """
        frac = (data - vmin) / np.float(vmax-vmin)
        if color_levels is None:
            color_levels = 255
        index = (frac * color_levels).ravel()
        # Out-of-bounds values will be lowest/highest colors
        index[index < 0] = 0
        index[index > 255] = 255
        rarr = []
        garr = []
        barr = []
        cmap = plt.cm.get_cmap(cmap)
        for val in index:
            if not np.isnan(val):
                ind = np.int(np.round(val))
                r, g, b, t = cmap(ind)
                rarr.append(np.int(np.round(r * 255)))
                garr.append(np.int(np.round(g * 255)))
                barr.append(np.int(np.round(b * 255)))
            else:
                rarr.append(np.nan)
                garr.append(np.nan)
                barr.append(np.nan)
        rarr = np.reshape(rarr, data.shape)
        garr = np.reshape(garr, data.shape)
        barr = np.reshape(barr, data.shape)
        return rarr, garr, barr

    def _create_sld(cmap, vmin, vmax, filename, color_levels=None):
        """
        Develop a Style Layer Descriptor file given a color table and
        user-specified min/max files. Output color info to that file.
        Only called if sld is True in write_grid_geotiff.

        Parameters
        ----------
        cmap : str or matplotlib.colors.Colormap object, optional
            Colormap to use for RGB output or SLD file.
        vmin : int or float
            Minimum value to color for RGB output or SLD file.
        vmax : int or float
            Maximum value to color for RGB output or SLD file.
        filename : str
            Template for SLD filename. The suffix (presumably .tif or .tiff)
            is removed and replaced with .sld. Thus, if provided a filename
            radar_reflectivity.tif, the output SLD file will be called
            radar_reflectivity.sld.

        Other Parameters
        ----------------
        color_levels : int or None, optional
            Number of color levels in cmap. Useful for categorical colormaps
            with steps << 255 (e.g., hydrometeor ID).

        """
        cmap = plt.cm.get_cmap(cmap)
        if color_levels is None:
            color_levels = 255
        name, _ = filename.split('.')
        ofile = name + '.sld'
        fileobj = open(ofile, 'w')

        header = """<?xml version="1.0" encoding="UTF-8"?>
    <sld:StyledLayerDescriptor xmlns="http://www.opengis.net/sld" xmlns:sld="http://www.opengis.net/sld" xmlns:ogc="http://www.opengis.net/ogc" xmlns:gml="http://www.opengis.net/gml" version="1.0.0">
        <sld:UserLayer>
            <sld:LayerFeatureConstraints>
                <sld:FeatureTypeConstraint/>
            </sld:LayerFeatureConstraints>
            <sld:UserStyle>
                <sld:Name>""" + str(name) + """</sld:Name>
                <sld:FeatureTypeStyle>
                    <sld:Name>name</sld:Name>
                    <sld:FeatureTypeName>Feature</sld:FeatureTypeName>
                    <sld:Rule>
                        <sld:RasterSymbolizer>
                            <sld:ColorMap>"""
        fileobj.write(header)

        for i in np.arange(color_levels + 1):
            val = i * (vmax - vmin) / (color_levels) + vmin  # color_levels + 1
            rgbt = cmap(i)
            if i == 0:
                op = 0.0
            else:
                op = 1.0
            hexval = colors.rgb2hex(rgbt[0:3])
            wstr = '\n                            <sld:ColorMapEntry color=\"' + \
                str(hexval).upper() + '\" opacity=\"' + str(op) + \
                '\" quantity=\"' + str(val) + '\"/>'
            fileobj.write(wstr)

        footer = """
                            </sld:ColorMap>
                        </sld:RasterSymbolizer>
                    </sld:Rule>
                </sld:FeatureTypeStyle>
            </sld:UserStyle>
        </sld:UserLayer>
    </sld:StyledLayerDescriptor>
    """
        fileobj.write(footer)
        fileobj.close()