GOES16_to_COG.py

"""geos: netcdf_to_tiff."""

import os

import numpy

from netCDF4 import Dataset

from affine import Affine

import rasterio
from rasterio.io import MemoryFile
from rasterio.warp import calculate_default_transform, reproject
from rasterio.transform import array_bounds

from rio_cogeo.cogeo import cog_translate
from rio_cogeo.profiles import cog_profiles


# From https://github.com/Solcast/netcdf-tiff
def get_goes_transform(resolution):
    """Return GOES16 geotransform for corresponding resolution."""
    if resolution.startswith("0.5km"):
        return Affine.from_gdal(
            -5434895.081637931,
            501.0043288718853,
            0,
            -5434894.837566491,
            0,
            501.0043288718853,
        )
    elif resolution.startswith("1km"):
        return Affine.from_gdal(
            -5434894.954752678,
            1002.0086577437705,
            0,
            -5434894.964451744,
            0,
            1002.0086577437705,
        )
    elif resolution.startswith("2km"):
        return Affine.from_gdal(
            -5434894.700982173,
            2004.0173154875410,
            0,
            -5434895.218222249,
            0,
            2004.0173154875410,
        )
    else:
        raise Exception(f"Invalid resolution: {resolution}")


def convert(
    netcdf_file,
    cog_path,
    sds="Rad",
    cogeo_profile="deflate",
    cogeo_options={"blockxsize": 128, "blockysize": 128, "predictor": 2, "zlevel": 7},
    autoclean=True,
):
    """Convert a netcdf to COG."""
    ds = Dataset(netcdf_file, "r")

    proj_name = ds.variables[sds].getncattr("grid_mapping")
    proj_info = ds.variables[proj_name]

    # From https://github.com/Solcast/netcdf-tiff
    proj_string = " ".join([
        "+proj=geos",  # The projection name is very important this means geostationary
        "+lon_0={0}".format(proj_info.longitude_of_projection_origin),
        "+h={0}".format(proj_info.perspective_point_height),
        "+a={0}".format(proj_info.semi_major_axis),
        "+b={0}".format(proj_info.semi_minor_axis),
        "+units={0}".format('m'),
        "+sweep={0}".format(proj_info.sweep_angle_axis),
        "+no_defs"
    ])
    crs = rasterio.crs.CRS.from_string(proj_string)
    transform = get_goes_transform(ds.getncattr("spatial_resolution"))

    tags = {
        k: ds.variables["Rad"].getncattr(k) for k in ds.variables["Rad"].ncattrs()
    }

    nodata = ds.variables[sds]._FillValue
    dtype = ds.variables[sds].dtype
    extracted_data = numpy.flip(ds.variables[sds][:].data.astype(dtype), axis=0)
    height, width = extracted_data.shape
    ds.close()

    if autoclean:
        os.remove(netcdf_file)

    config = dict(
        GDAL_NUM_THREADS="ALL_CPUS",
        GDAL_TIFF_OVR_BLOCKSIZE="128",
    )
    src_bounds = array_bounds(height, width, transform)
    dst_transform, dst_width, dst_height = calculate_default_transform(
        crs, "epsg:4326", width, height, *src_bounds
    )

    meta = dict(
        driver="GTiff",
        dtype=dtype,
        count=1,
        height=dst_height,
        width=dst_width,
        crs="epsg:4326",
        transform=dst_transform,
        nodata=nodata,
        tiled=True,
        compress="deflate",
        blockxsize=512,
        blockysize=512,
    )

    with rasterio.Env(**config):
        with MemoryFile() as memfile:
            with memfile.open(**meta) as mem:
                reproject(
                    extracted_data,
                    rasterio.band(mem, 1),
                    src_transform=transform,
                    src_crs=crs,
                    src_nodata=nodata,
                    resampling=rasterio.enums.Resampling.nearest,
                    num_threads=8,
                )
                mem.update_tags(**tags)
                extracted_data = None

                profile = cog_profiles.get(cogeo_profile)
                profile.update(cogeo_options)
                cog_translate(
                    mem,
                    cog_path,
                    profile,
                    config=config,
                    in_memory=True,
                    # allow_intermediate_compression=True,
                    quiet=True,
                )