Aligning a global raster, including pixel area-aware conversion to/from population pixel counts

convert.py

import numpy
import pygeoprocessing
from osgeo import gdal

FLOAT32_NODATA = float(numpy.finfo(numpy.float32).min)


def _convert_to_from_density(source_raster_path, target_raster_path,
                             direction='to_density'):
    """Convert a raster to/from counts/pixel and counts/unit area.

    Args:
        source_raster_path (string): The path to a raster containing units that
            need to be converted.
        target_raster_path (string): The path to where the target raster with
            converted units should be stored.
        direction='to_density' (string): The direction of unit conversion.  If
            'to_density', then the units of ``source_raster_path`` must be in
            counts per pixel and will be converted to counts per square meter.
            If 'from_density', then the units of ``source_raster_path`` must be
            in counts per square meter and will be converted to counts per
            pixel.

    Returns:
        ``None``
    """
    source_raster_info = pygeoprocessing.get_raster_info(source_raster_path)
    source_nodata = source_raster_info['nodata'][0]

    # Calculate the per-pixel area based on the latitude.
    _, miny, _, maxy = source_raster_info['bounding_box']
    pixel_area_in_m2_by_latitude = (
        pygeoprocessing._create_latitude_m2_area_column(
            miny, maxy, source_raster_info['raster_size'][1]))

    def _convert(array, pixel_area):
        out_array = numpy.full(array.shape, FLOAT32_NODATA,
                               dtype=numpy.float32)
        valid_mask = slice(None)
        if source_nodata is not None:
            valid_mask = ~numpy.isclose(array, source_nodata)

        if direction == 'to_density':
            out_array[valid_mask] = array[valid_mask] / pixel_area[valid_mask]
        elif direction == 'from_density':
            out_array[valid_mask] = array[valid_mask] * pixel_area[valid_mask]
        else:
            raise AssertionError(f'Invalid direction: {direction}')
        return out_array

    pygeoprocessing.raster_calculator(
        [(source_raster_path, 1), pixel_area_in_m2_by_latitude],
        _convert, target_raster_path, gdal.GDT_Float32, FLOAT32_NODATA)


if __name__ == '__main__':
    population_raster_path = 'population.tif'
    pop_per_sq_m = 'population_per_sq_m.tif'
    _convert_to_from_density(population_raster_path, pop_per_sq_m,
                             'to_density')

    # call align_and_resize_raster_stack here using pop_per_sq_m as a source
    # path.
    aligned_pop_per_sq_m = 'aligned_pop_per_sq_m.tif'
    pygeoprocessing.align_and_resize_raster_stack(
        base_raster_path_list=[pop_per_sq_m],
        target_raster_path_list=[aligned_pop_per_sq_m],
        resample_method_list=['bilinear'],
        target_pixel_size=('YOUR PIXEL SIZE HERE'),
        bounding_box_mode='intersection'
    )

    # Once pop-per-sq-m has been aligned with the stack, convert it back to
    # pop-per-pixel.
    pop_per_pixel = 'aligned_population_per_pixel.tif'
    _convert_to_from_density(aligned_pop_per_sq_m, pop_per_pixel,
                             'from_density')