Here are two possible solutions using both gdal and numpy. The first solution consists of looping through each pair of coordinates, getting the corresponding raster pixel and then, extracting its value along with their n neighbours. The second solution consists of a vectorized (and faster) version of the first solution.
For the sake of this example, assume you want to extract precipitation values (and their 24 neighbours) for populated places in the world. A Shapefile with (7343) populated places of the world can be found on Natural Earth's 1:10m Cultural Vectors and a zip file containing 12 precipitation GeoTIFF files (one for each month) can be found on WorldClim's Historical climate data. Here I'll be using January's GeoTIFF of the prec 2.5m dataset, which has 4320 rows by 8640 columns.
Solution 1: looping through each pair of coordinate and extracting values one step at a time
First, open the raster using gdal, get the GeoTransform, get the NoData value and read the raster as a numpy array.
from osgeo import gdal
ds = gdal.Open('wc2.1_2.5m_prec_01.tif', 0)
ox, pw, xskew, oy, yskew, ph = ds.GetGeoTransform()
nd_value = ds.GetRasterBand(1).GetNoDataValue()
arr = ds.ReadAsArray()
del ds
Second, define the window size taking into account the number of neighbours and pad the array at each side with NoData values to handle the cases when the coordinates lie at or close to the edge of the raster.
import numpy as np
window_size = (5, 5) # 25 cells, 24 neighbours and window center
padding_y = (2, 2) # 2 rows above and 2 rows below
padding_X = (2, 2) # 2 columns to the left and 2 columns to the right
padded_arr = np.pad(arr, pad_width=(padding_y, padding_x), mode='constant', constant_values=nd_value)
Then, open the Shapefile using ogr and get the x and y coordinates of all the features.
from osgeo import ogr
lyr = ds.GetLayer()
coords = [(feat.geometry().GetX(), feat.geometry().GetY()) for feat in lyr]
coords = np.array(coords)
x = coords.T[0]
y = coords.T[1]
del ds, lyr
Before looping through each pair of coordinates, I'll write a helper function to convert real world coordinates (WGS84 in this case) to array coordinates (i.e. positive integer indices).
import math
def get_index(x: float, y: float, ox: float, oy: float, pw: float, ph: float) -> tuple:
"""
Gets the row (i) and column (j) indices in an NumPy 2D array for a given
pair of coordinates.
Parameters
----------
x : float
x (longitude) coordinate
y : float
y (latitude) coordinate
ox : float
Raster x origin (minimum x coordinate)
oy : float
Raster y origin (maximum y coordinate)
pw : float
Raster pixel width
ph : float
Raster pixel height
Returns
-------
Two-element tuple with the column and row indices.
Notes
-----
This function is based on: https://gis.stackexchange.com/a/92015/86131.
Both x and y coordinates must be within the raster boundaries. Otherwise,
the index will not correspond to the actual values or will be out of
bounds.
"""
# make sure pixel height is positive
ph = abs(ph)
i = math.floor((oy-y) / ph)
j = math.floor((x-ox) / pw)
return i, j
Finally, define an offset based on the number of cells you want to extract at each side of the pixel center (which in this case would be 2), loop through each pair of coordinates, convert them to array coordinates and then extract the values.
offset = 2
for x_coord, y_coord in zip(x, y):
# get index
i, j = get_index(x_coord, y_coord, ox, oy, pw, ph)
# get pixel value and its 24 neighbours
values = padded_arr[i-offset:i+offset+1, j-offset:j+offset+1]
Solution 2: getting a rolling window and extracting values at once
The first thing to do is to write a function that will get a 4D (the first two dimensions correspond to the shape of the original array and the last two dimensions correspond to the shape of the window) numpy array with a 5x5 window for each pixel in the original array. To do this, I'll use the numpy.lib.stride_tricks.as_strided function. It is worth mentioning that this function returns a view on the original array rather than a new array.
def rolling_window(arr: np.ndarray, window_size: tuple = (3, 3)) -> np.ndarray:
"""
Gets a view with a window of a specific size for each element in arr.
Parameters
----------
arr : np.ndarray
NumPy 2D array.
window_size : tuple
Tuple with the number of rows and columns for the window. Both values
have to be positive (i.e. greater than zero) and they cannot exceed
arr dimensions.
Returns
-------
NumPy 4D array
Notes
-----
This function has been slightly adapted from the one presented on:
https://rigtorp.se/2011/01/01/rolling-statistics-numpy.html.
It is advised to read the notes on the numpy.lib.stride_tricks.as_strided
function, which can be found on:
https://docs.scipy.org/doc/numpy-1.17.0/reference/generated/numpy.lib.stride_tricks.as_strided.html
"""
# validate window size
err1 = 'window size must be postive'
err2 = 'window size exceeds input array dimensions'
assert window_size[0] > 0 and window_size[1] > 0, err1
assert window_size[0] <= arr.shape[0] and window_size[1] <= arr.shape[1], err2
# calculate output array's shape
y_size = (arr.shape[0] - window_size[0]) + 1
x_size = (arr.shape[1] - window_size[1]) + 1
shape = (y_size, x_size) + window_size
# define strides
strides = arr.strides * 2
return np.lib.stride_tricks.as_strided(arr, shape, strides, writeable=False)
I'll also write a vectorized version of the get_index function so it takes an array of x coordinates and an array of y coordinates as input, and returns a tuple with two arrays containing the indices.
def get_indices(x: np.ndarray, y: np.ndarray, ox: float, oy: float,
pw: float, ph: float) -> tuple:
"""
Gets the row (i) and column (j) indices in an NumPy 2D array for a given
set of coordinates.
Parameters
----------
x : np.ndarray
NumPy 1D array containing the x (longitude) coordinates.
y : np.ndarray
NumPy 1D array containing the y (latitude) coordinates.
ox : float
Raster x origin (minimum x coordinate)
oy : float
Raster y origin (maximum y coordinate)
pw : float
Raster pixel width
ph : float
Raster pixel height
Returns
-------
Two-element tuple with the column and row indices.
Notes
-----
This function is based on: https://gis.stackexchange.com/a/92015/86131.
All x and y coordinates must be within the raster boundaries. Otherwise,
indices will not correspond to the actual values or will be out of bounds.
"""
# make sure pixel height is positive
ph = abs(ph)
i = np.floor((oy-y) / ph).astype('int')
j = np.floor((x-ox) / pw).astype('int')
return i, j
Now it is just a matter of getting the windows for each pixel (using padded_arr to handle edge cases; pun intended), getting the indices for all the coordinates and then indexing the view with the windows to get the values.
windows = rolling_window(padded_arr, window_size=window_size)
idx = get_indices(x, y, ox, oy, pw, ph)
values = windows[idx]
If you take a look at values, you'll see that it is a numpy 3D array with shape 7343 by 5 by 5. This means there is a 2D 5 x 5 (window size) array for each point in the Shapefile.
Benchmarking
To compare the execution time of the proposed solutions, I wrote a wrapper function for each one and then used IPython's %timeit built-in magic command.
Here is the wrapper function for solution 1:
def extract_n_neighbours(padded_arr, x, y, ox, oy, pw, ph):
offset = 2
for x_coord, y_coord in zip(x, y):
i, j = get_index(x_coord, y_coord, ox, oy, pw, ph)
values = padded_arr[i-offset:i+offset+1, j-offset:j+offset+1]
Here is the wrapper function for solution 2:
def extract_n_neighbours_vectorized(padded_arr, x, y, ox, oy, pw, ph):
windows = rolling_window(padded_arr, window_size=window_size)
idx = get_indices(x, y, ox, oy, pw, ph)
values = windows[idx]
And here are the results:
In[2]: %timeit -n 1000 extract_n_neighbours(padded_arr, x, y, ox, oy, pw, ph)
21.7 ms ± 1.62 ms per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In[3]: %timeit -n 1000 extract_n_neighbours_vectorized(padded_arr, x, y, ox, oy, pw, ph)
2.37 ms ± 156 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) # almost ten times faster
