Improvements to transformations (#3013)

* Use ndarrays instead of lists.

* Make fast path for ndarrays.

* Use ndarrays instead of lists.

In the future, see if we can pass the underlying array pointer into GDAL directly without having to make a copy.

* Allow efficient use of ufuncs.

* Optimizations to AffineTransformer

1. Pre-allocate transform array. If the user creates their own AffineTransformer, they can avoid reallocating the transform array.
2. Faster construction of input_matrix
3. Faster, inplace dot product

Since _ensure_arr_input is called before _transform, the typecheck is redundant. _transform method should never be called directly by the user.

* Update xy method.

* Reorder condition to avoid array access if possible.

* Broadcast input arrays.

* Only check that input arrays satisfy broadcasting requirements.

Avoid allocating new arrays to store broadcasted results.

* Broadcast input arrays.

* Prefer np.floor over math.floor

* Remove unused import.
This commit is contained in:
Ryan Grout 2024-04-08 15:11:01 -06:00 committed by GitHub
parent 90d75b0d38
commit 7fb4442c40
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 97 additions and 58 deletions

View File

@ -8,6 +8,8 @@ from contextlib import ExitStack
import logging import logging
import warnings import warnings
import numpy as np
from rasterio._err import GDALError from rasterio._err import GDALError
from rasterio._err cimport exc_wrap_pointer from rasterio._err cimport exc_wrap_pointer
from rasterio.errors import NotGeoreferencedWarning, TransformWarning from rasterio.errors import NotGeoreferencedWarning, TransformWarning
@ -132,7 +134,7 @@ cdef class RPCTransformerBase:
Parameters Parameters
---------- ----------
xs, ys, zs : list xs, ys, zs : ndarray
List of coordinates to be transformed. May be either pixel/line/height or List of coordinates to be transformed. May be either pixel/line/height or
lon/lat/height) lon/lat/height)
transform_direction : TransformDirection transform_direction : TransformDirection
@ -151,7 +153,7 @@ cdef class RPCTransformerBase:
Returns Returns
------- -------
tuple of list tuple of ndarray
Notes Notes
----- -----
@ -161,7 +163,8 @@ cdef class RPCTransformerBase:
if self._transformer == NULL: if self._transformer == NULL:
raise ValueError("Unexpected NULL transformer") raise ValueError("Unexpected NULL transformer")
cdef int i cdef Py_ssize_t i
cdef int n
cdef double *x = NULL cdef double *x = NULL
cdef double *y = NULL cdef double *y = NULL
cdef double *z = NULL cdef double *z = NULL
@ -169,16 +172,20 @@ cdef class RPCTransformerBase:
cdef int src_count cdef int src_count
cdef int *panSuccess = NULL cdef int *panSuccess = NULL
n = len(xs) cdef double[:] rxview, ryview, rzview
cdef tuple xyz
bi = np.broadcast(xs, ys, zs)
n = bi.size
x = <double *>CPLMalloc(n * sizeof(double)) x = <double *>CPLMalloc(n * sizeof(double))
y = <double *>CPLMalloc(n * sizeof(double)) y = <double *>CPLMalloc(n * sizeof(double))
z = <double *>CPLMalloc(n * sizeof(double)) z = <double *>CPLMalloc(n * sizeof(double))
panSuccess = <int *>CPLMalloc(n * sizeof(int)) panSuccess = <int *>CPLMalloc(n * sizeof(int))
for i in range(n): for i, xyz in enumerate(bi):
x[i] = xs[i] x[i] = <double>xyz[0]
y[i] = ys[i] y[i] = <double>xyz[1]
z[i] = zs[i] z[i] = <double>xyz[2]
try: try:
err = GDALRPCTransform(self._transformer, bDstToSrc, n, x, y, z, panSuccess) err = GDALRPCTransform(self._transformer, bDstToSrc, n, x, y, z, panSuccess)
@ -186,20 +193,22 @@ cdef class RPCTransformerBase:
warnings.warn( warnings.warn(
"Could not transform points using RPCs.", "Could not transform points using RPCs.",
TransformWarning) TransformWarning)
res_xs = [0] * n res_xs = np.zeros(n)
res_ys = [0] * n res_ys = np.zeros(n)
res_zs = [0] * n #res_zs = np.zeros(n)
rxview = res_xs
ryview = res_ys
checked = False checked = False
for i in range(n): for i in range(n):
# GDALRPCTransformer may return a success overall despite individual points failing. Warn once. # GDALRPCTransformer may return a success overall despite individual points failing. Warn once.
if not panSuccess[i] and not checked: if not checked and not panSuccess[i]:
warnings.warn( warnings.warn(
"One or more points could not be transformed using RPCs.", "One or more points could not be transformed using RPCs.",
TransformWarning) TransformWarning)
checked = True checked = True
res_xs[i] = x[i] rxview[i] = x[i]
res_ys[i] = y[i] ryview[i] = y[i]
res_zs[i] = z[i] #res_zs[i] = z[i]
finally: finally:
CPLFree(x) CPLFree(x)
CPLFree(y) CPLFree(y)
@ -285,7 +294,7 @@ cdef class GCPTransformerBase:
Parameters Parameters
---------- ----------
xs, ys, zs : list xs, ys, zs : ndarray
List of coordinates to be transformed. May be either pixel/line/height or List of coordinates to be transformed. May be either pixel/line/height or
lon/lat/height) lon/lat/height)
transform_direction : TransformDirection transform_direction : TransformDirection
@ -304,12 +313,13 @@ cdef class GCPTransformerBase:
Returns Returns
------- -------
tuple of list tuple of ndarray
""" """
if self._transformer == NULL: if self._transformer == NULL:
raise ValueError("Unexpected NULL transformer") raise ValueError("Unexpected NULL transformer")
cdef int i cdef Py_ssize_t i
cdef int n
cdef double *x = NULL cdef double *x = NULL
cdef double *y = NULL cdef double *y = NULL
cdef double *z = NULL cdef double *z = NULL
@ -317,16 +327,20 @@ cdef class GCPTransformerBase:
cdef int src_count cdef int src_count
cdef int *panSuccess = NULL cdef int *panSuccess = NULL
n = len(xs) cdef double[:] rxview, ryview, rzview
cdef tuple xyz
bi = np.broadcast(xs, ys, zs)
n = bi.size
x = <double *>CPLMalloc(n * sizeof(double)) x = <double *>CPLMalloc(n * sizeof(double))
y = <double *>CPLMalloc(n * sizeof(double)) y = <double *>CPLMalloc(n * sizeof(double))
z = <double *>CPLMalloc(n * sizeof(double)) z = <double *>CPLMalloc(n * sizeof(double))
panSuccess = <int *>CPLMalloc(n * sizeof(int)) panSuccess = <int *>CPLMalloc(n * sizeof(int))
for i in range(n): for i, xyz in enumerate(bi):
x[i] = xs[i] x[i] = <double>xyz[0]
y[i] = ys[i] y[i] = <double>xyz[1]
z[i] = zs[i] z[i] = <double>xyz[2]
try: try:
if self._tps: if self._tps:
@ -337,20 +351,22 @@ cdef class GCPTransformerBase:
warnings.warn( warnings.warn(
"Could not transform points using GCPs.", "Could not transform points using GCPs.",
TransformWarning) TransformWarning)
res_xs = [0] * n res_xs = np.zeros(n)
res_ys = [0] * n res_ys = np.zeros(n)
res_zs = [0] * n #res_zs = np.zeros(n)
rxview = res_xs
ryview = res_ys
checked = False checked = False
for i in range(n): for i in range(n):
# GDALGCPTransformer or GDALTPSTransformer may return a success overall despite individual points failing. Warn once. # GDALGCPTransformer or GDALTPSTransformer may return a success overall despite individual points failing. Warn once.
if not panSuccess[i] and not checked: if not checked and not panSuccess[i]:
warnings.warn( warnings.warn(
"One or more points could not be transformed using GCPs.", "One or more points could not be transformed using GCPs.",
TransformWarning) TransformWarning)
checked = True checked = True
res_xs[i] = x[i] rxview[i] = x[i]
res_ys[i] = y[i] ryview[i] = y[i]
res_zs[i] = z[i] #res_zs[i] = z[i]
finally: finally:
CPLFree(x) CPLFree(x)
CPLFree(y) CPLFree(y)

View File

@ -2,9 +2,9 @@
from contextlib import ExitStack from contextlib import ExitStack
from functools import partial from functools import partial
import math
import numpy as np import numpy as np
import warnings import warnings
from numbers import Number
from affine import Affine from affine import Affine
@ -86,7 +86,7 @@ class TransformMethodsMixin:
x, x,
y, y,
z=None, z=None,
op=math.floor, op=np.floor,
precision=None, precision=None,
transform_method=TransformMethod.affine, transform_method=TransformMethod.affine,
**rpc_options **rpc_options
@ -251,7 +251,7 @@ def xy(transform, rows, cols, zs=None, offset='center', **rpc_options):
return transformer.xy(rows, cols, zs=zs, offset=offset) return transformer.xy(rows, cols, zs=zs, offset=offset)
def rowcol(transform, xs, ys, zs=None, op=math.floor, precision=None, **rpc_options): def rowcol(transform, xs, ys, zs=None, op=np.floor, precision=None, **rpc_options):
"""Get rows and cols of the pixels containing (x, y). """Get rows and cols of the pixels containing (x, y).
Parameters Parameters
@ -330,13 +330,23 @@ class TransformerBase:
TransformError TransformError
If input coordinates are not all of the same length If input coordinates are not all of the same length
""" """
xs = np.atleast_1d(xs)
ys = np.atleast_1d(ys)
if zs is not None:
zs = np.atleast_1d(zs)
else:
zs = np.zeros(1)
try: try:
xs, ys, zs = np.broadcast_arrays(xs, ys, 0 if zs is None else zs) broadcasted = np.broadcast(xs, ys, zs)
except ValueError as error: if broadcasted.ndim != 1:
raise TransformError( raise TransformError(
"Input coordinates must be broadcastable to a 1d array" "Input coordinates must be broadcastable to a 1d array"
) from error )
return np.atleast_1d(xs), np.atleast_1d(ys), np.atleast_1d(zs) except ValueError as error:
raise TransformError() from error
return xs, ys, zs
def __enter__(self): def __enter__(self):
return self return self
@ -344,7 +354,7 @@ class TransformerBase:
def __exit__(self, *args): def __exit__(self, *args):
pass pass
def rowcol(self, xs, ys, zs=None, op=math.floor, precision=None): def rowcol(self, xs, ys, zs=None, op=np.floor, precision=None):
"""Get rows and cols coordinates given geographic coordinates. """Get rows and cols coordinates given geographic coordinates.
Parameters Parameters
@ -378,7 +388,7 @@ class TransformerBase:
RasterioDeprecationWarning, RasterioDeprecationWarning,
) )
AS_ARR = True if hasattr(xs, "__iter__") else False IS_SCALAR = isinstance(xs, Number) and isinstance(ys, Number)
xs, ys, zs = self._ensure_arr_input(xs, ys, zs=zs) xs, ys, zs = self._ensure_arr_input(xs, ys, zs=zs)
try: try:
@ -386,10 +396,22 @@ class TransformerBase:
xs, ys, zs, transform_direction=TransformDirection.reverse xs, ys, zs, transform_direction=TransformDirection.reverse
) )
if len(new_rows) == 1 and not AS_ARR: is_op_ufunc = isinstance(op, np.ufunc)
return (op(new_rows[0]), op(new_cols[0])) if is_op_ufunc:
op(new_rows, out=new_rows)
op(new_cols, out=new_cols)
new_rows = new_rows.tolist()
new_cols = new_cols.tolist()
if not is_op_ufunc:
new_rows = list(map(op, new_rows))
new_cols = list(map(op, new_cols))
if IS_SCALAR:
return new_rows[0], new_cols[0]
else: else:
return ([op(r) for r in new_rows], [op(c) for c in new_cols]) return new_rows, new_cols
except TypeError: except TypeError:
raise TransformError("Invalid inputs") raise TransformError("Invalid inputs")
@ -419,7 +441,7 @@ class TransformerBase:
tuple of float or list of float tuple of float or list of float
""" """
AS_ARR = True if hasattr(rows, "__iter__") else False IS_SCALAR = isinstance(rows, Number) and isinstance(cols, Number)
rows, cols, zs = self._ensure_arr_input(rows, cols, zs=zs) rows, cols, zs = self._ensure_arr_input(rows, cols, zs=zs)
if offset == 'center': if offset == 'center':
@ -445,10 +467,11 @@ class TransformerBase:
new_xs, new_ys = self._transform( new_xs, new_ys = self._transform(
offset_cols, offset_rows, zs, transform_direction=TransformDirection.forward offset_cols, offset_rows, zs, transform_direction=TransformDirection.forward
) )
if len(new_xs) == 1 and not AS_ARR:
return (new_xs[0], new_ys[0]) if IS_SCALAR:
return new_xs[0], new_ys[0]
else: else:
return (list(new_xs), list(new_ys)) return new_xs.tolist(), new_ys.tolist()
except TypeError: except TypeError:
raise TransformError("Invalid inputs") raise TransformError("Invalid inputs")
@ -480,22 +503,22 @@ class AffineTransformer(TransformerBase):
if not isinstance(affine_transform, Affine): if not isinstance(affine_transform, Affine):
raise ValueError("Not an affine transform") raise ValueError("Not an affine transform")
self._transformer = affine_transform self._transformer = affine_transform
self._transform_arr = np.empty((3, 3), dtype='float64')
def _transform(self, xs, ys, zs, transform_direction): def _transform(self, xs, ys, zs, transform_direction):
transform = self._transform_arr
if transform_direction is TransformDirection.forward: if transform_direction is TransformDirection.forward:
transform = self._transformer transform.flat[:] = self._transformer
elif transform_direction is TransformDirection.reverse: elif transform_direction is TransformDirection.reverse:
transform = ~self._transformer transform.flat[:] = ~self._transformer
is_arr = True if type(xs) in [list, tuple] else False bi = np.broadcast(xs, ys)
if is_arr: input_matrix = np.empty((3, bi.size))
a, b, c, d, e, f, _, _, _ = transform input_matrix[0] = bi.iters[0]
transform_matrix = np.array([[a, b, c], [d, e, f]]) input_matrix[1] = bi.iters[1]
input_matrix = np.array([xs, ys, np.ones(len(xs))]) input_matrix[2] = 1
output_matrix = np.dot(transform_matrix, input_matrix) transform.dot(input_matrix, out=input_matrix)
return (list(output_matrix[0]), list(output_matrix[1])) return input_matrix[0], input_matrix[1]
else:
return transform * (xs, ys)
def __repr__(self): def __repr__(self):