DataCube Write Functions

writers.py

"""
Useful functions for Datacube users

Not used internally, those should go in `utils.py`
"""
import rasterio
import numpy as np

DEFAULT_PROFILE = {
    'blockxsize': 256,
    'blockysize': 256,
    'compress': 'lzw',
    'driver': 'GTiff',
    'interleave': 'band',
    'nodata': 0.0,
    'photometric': 'RGBA',
    'tiled': True}


def write_geotiff(filename, dataset, time_index=None, profile_override=None):
    """
    Write an xarray dataset to a geotiff

    :param filename: Output filename
    :attr dataset: xarray dataset containing multiple bands to write to file
    :attr time_index: time index to write to file
    :attr profile_override: option dict, overrides rasterio file creation options.
    """
    profile_override = profile_override or {}

    dtypes = {val.dtype for val in dataset.data_vars.values()}
    assert len(dtypes) == 1  # Check for multiple dtypes

    profile = DEFAULT_PROFILE.copy()
    profile.update({
        'width': dataset.dims[dataset.crs.dimensions[1]],
        'height': dataset.dims[dataset.crs.dimensions[0]],
        'affine': dataset.affine,
        'crs': dataset.crs.crs_str,
        'count': len(dataset.data_vars),
        'dtype': str(dtypes.pop())
    })
    profile.update(profile_override)

    with rasterio.open(str(filename), 'w', **profile) as dest:
        for bandnum, data in enumerate(dataset.data_vars.values(), start=1):
            dest.write(data.isel(time=time_index).data, bandnum)


def write_xr_dataset_to_disk(filename, dataset, **profile_override):
    """
    Write an xarray dataset to a geotiff.

    Expects a single time value, with one or more data variables. Each data variable
    will be written to a separate band.

    For a multi-time dataset, use ``sel()`` or ``isel()`` to select a single time value.

    :param filename: Output filename
    :attr dataset: xarray dataset containing multiple bands to write to file
    :attr profile_override: option dict, overrides rasterio file creation options.
    """
    dtypes = {val.dtype for val in dataset.data_vars.values()}
    assert len(dtypes) == 1  # Check for multiple dtypes

    profile = {
        'width': dataset.dims[dataset.crs.dimensions[1]],
        'height': dataset.dims[dataset.crs.dimensions[0]],
        'transform': dataset.affine,
        'crs': dataset.crs.crs_str,
        'count': len(dataset.data_vars),
        'dtype': str(dtypes.pop())
    }
    profile.update(profile_override)

    with rasterio.open(str(filename), 'w', **profile) as dest:
        for bandnum, data in enumerate(dataset.data_vars.values(), start=1):
            dest.write(data.data, bandnum)

import rasterio
def write_multi_time_dataarray(filename, dataarray, **profile_override):
    profile = {
        'width': len(dataarray[dataarray.crs.dimensions[1]]),
        'height': len(dataarray[dataarray.crs.dimensions[0]]),
        'transform': dataarray.affine,
        'crs': dataarray.crs.crs_str,
        'count': len(dataarray.time),
        'dtype': str(dataarray.dtype)
    }
    profile.update(profile_override)

    with rasterio.open(str(filename), 'w', **profile) as dest:
        for time_idx in range(len(dataarray.time)):
            bandnum = time_idx + 1
            dest.write(dataarray.isel(time=time_idx).data, bandnum)

write_multi_time_dataarray('/short/v10/dra547/test2.bin', data.red, driver='ENVI')

def ga_pq_fuser(dest, src):
    """
    Fuse two Geoscience Australia Pixel Quality ndarrays
    
    Takes a **pessimistic** approach to merging, if either array value shows cloud or cloud shadow, the output
    array will also show 

    To be used as a `fuse_func` when loaded `grouped` data, for example when grouping
    by solar day to avoid duplicate data from scene overlaps.
    """
    valid_bit = 8
    valid_val = (1 << valid_bit)

    no_data_dest_mask = ~(dest & valid_val).astype(bool)
    np.copyto(dest, src, where=no_data_dest_mask)

    both_data_mask = (valid_val & dest & src).astype(bool)
    np.copyto(dest, src & dest, where=both_data_mask)