comparing two methods for "is near" algorithm

is_near_comparison.py

import numpy
import math
from osgeo import gdal, osr
import pygeoprocessing
import timeit

FLOAT_NODATA = -1
UINT8_NODATA = 255


def is_near1(input_path, radius, distance_path, out_path):
    # Calculate the distance from each pixel to the nearest '1' pixel
    pygeoprocessing.distance_transform_edt(
        (input_path, 1),
        distance_path
    )

    def lte_threshold_op(array, threshold):
        out = numpy.full(array.shape, 0, dtype=numpy.uint8)
        valid_mask = ~numpy.isclose(array, FLOAT_NODATA)
        out[valid_mask] = array[valid_mask] <= threshold
        return out

    # Threshold that to a binary array so '1' means it's within the radius
    pygeoprocessing.raster_calculator(
        [(distance_path, 1), (radius, 'raw')],
        lte_threshold_op,
        out_path,
        gdal.GDT_Byte,
        UINT8_NODATA)


def is_near2(input_path, radius, kernel_path, count_path, output_path):
    """
    Args:
        input_path:
        radius
    """
    search_kernel = make_search_kernel(input_path, radius)
    # create and open the output raster and save the kernel to it
    # set up a spatial reference and raster properties
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(3857)
    projection_wkt = srs.ExportToWkt()

    pygeoprocessing.numpy_array_to_raster(
        search_kernel,
        UINT8_NODATA,
        (20, -20),
        (-10300000, 5610000),
        projection_wkt,
        kernel_path)

    pygeoprocessing.convolve_2d(
        (input_path, 1),
        (kernel_path, 1),
        count_path,
        # pixels with nodata or off the edge of the raster won't count
        ignore_nodata_and_edges=True,
        # output will have nodata where ratio_path has nodata
        mask_nodata=True,
        target_datatype=gdal.GDT_UInt16,
        target_nodata=FLOAT_NODATA)

    def gte_threshold_op(array, threshold):
        out = numpy.full(array.shape, 0, dtype=numpy.uint8)
        valid_mask = ~numpy.isclose(array, FLOAT_NODATA)
        out[valid_mask] = array[valid_mask] >= threshold
        return out

    pygeoprocessing.raster_calculator(
        [(count_path, 1), (1, 'raw')],
        gte_threshold_op,
        output_path,
        gdal.GDT_Byte,
        UINT8_NODATA)


def make_search_kernel(raster_path, radius):
    """Make a search kernel for a raster that marks pixels within a radius.
    Save the search kernel to a raster for use with pygeoprocessing.convolve_2d

    Args:
        raster_path (str): path to a raster to make kernel for
        radius (float): distance around each pixel's centerpoint to search
            in raster coordinate system units

    Returns:
        2D boolean numpy.ndarray. '1' pixels are within ``radius`` of the
        center pixel, measured centerpoint-to-centerpoint. '0' pixels are
        outside the radius. The array dimensions are as small as possible
        while still including the entire radius.
    """
    raster_info = pygeoprocessing.get_raster_info(raster_path)
    pixel_radius = radius / raster_info['pixel_size'][0]
    pixel_margin = math.floor(pixel_radius)
    # the search kernel is just large enough to contain all pixels that
    # *could* be within the radius of the center pixel
    search_kernel_shape = tuple([pixel_margin * 2 + 1] * 2)
    # arrays of the column index and row index of each pixel
    col_indices, row_indices = numpy.indices(search_kernel_shape)
    # adjust them so that (0, 0) is the center pixel
    col_indices -= pixel_margin
    row_indices -= pixel_margin
    # hypotenuse_i = sqrt(col_indices_i**2 + row_indices_i**2) for each pixel i
    hypotenuse = numpy.hypot(col_indices, row_indices)
    # boolean kernel where 1=pixel centerpoint is within the radius of the
    # center pixel's centerpoint
    search_kernel = numpy.array(hypotenuse <= pixel_radius, dtype=numpy.uint8)
    return search_kernel


radius = 30
input_array = numpy.random.choice([0, 1], (1000, 1000), p=[0.9, 0.1]).astype(numpy.uint8)
input_path = '/Users/emily/Documents/test_is_near_input.tif'
distance_path = '/Users/emily/Documents/test_is_near_distance.tif'
kernel_path = '/Users/emily/Documents/test_is_near_kernel.tif'
count_path = '/Users/emily/Documents/test_is_near_count.tif'
output_path_1 = '/Users/emily/Documents/test_is_near_output1.tif'
output_path_2 = '/Users/emily/Documents/test_is_near_output2.tif'

# set up a spatial reference and raster properties
srs = osr.SpatialReference()
srs.ImportFromEPSG(3857)
projection_wkt = srs.ExportToWkt()
origin = (-10300000, 5610000)
pixel_size = (20, -20)

pygeoprocessing.numpy_array_to_raster(input_array, UINT8_NODATA, pixel_size,
    origin, projection_wkt, input_path)

is_near1(input_path, radius / pixel_size[0], distance_path, output_path_1)
is_near2(input_path, radius, kernel_path, count_path, output_path_2)

result_1 = pygeoprocessing.raster_to_numpy_array(output_path_1)
result_2 = pygeoprocessing.raster_to_numpy_array(output_path_2)

numpy.testing.assert_equal(result_1, result_2)


def call_is_near1():
    is_near1(
        input_path,
        radius / pixel_size[0],
        distance_path,
        output_path_1)


def call_is_near2():
    is_near2(
        input_path,
        radius,
        kernel_path,
        count_path,
        output_path_2)


duration = timeit.timeit(call_is_near1, number=1000)
print('Method 1:', duration)

duration = timeit.timeit(call_is_near2, number=1000)
print('Method 2:', duration)