Skip to content

Commit

Permalink
Improvements to transformations (#3013)
Browse files Browse the repository at this point in the history 
* 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.
  • Loading branch information
@groutr
groutr committed Apr 8, 2024
1 parent 90d75b0 commit 7fb4442
Show file tree
Hide file tree
Showing 2 changed files with 98 additions and 59 deletions.
76 changes: 46 additions & 30 deletions rasterio/_transform.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ from contextlib import ExitStack
import logging
import warnings

import numpy as np

from rasterio._err import GDALError
from rasterio._err cimport exc_wrap_pointer
from rasterio.errors import NotGeoreferencedWarning, TransformWarning
Expand Down Expand Up @@ -132,7 +134,7 @@ cdef class RPCTransformerBase:
Parameters
----------
xs, ys, zs : list
xs, ys, zs : ndarray
List of coordinates to be transformed. May be either pixel/line/height or
lon/lat/height)
transform_direction : TransformDirection
Expand All @@ -151,7 +153,7 @@ cdef class RPCTransformerBase:
Returns
-------
tuple of list
tuple of ndarray
Notes
-----
Expand All @@ -161,45 +163,52 @@ cdef class RPCTransformerBase:
if self._transformer == NULL:
raise ValueError("Unexpected NULL transformer")

cdef int i
cdef Py_ssize_t i
cdef int n
cdef double *x = NULL
cdef double *y = NULL
cdef double *z = NULL
cdef int bDstToSrc = transform_direction
cdef int src_count
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))
y = <double *>CPLMalloc(n * sizeof(double))
z = <double *>CPLMalloc(n * sizeof(double))
panSuccess = <int *>CPLMalloc(n * sizeof(int))

for i in range(n):
x[i] = xs[i]
y[i] = ys[i]
z[i] = zs[i]
for i, xyz in enumerate(bi):
x[i] = <double>xyz[0]
y[i] = <double>xyz[1]
z[i] = <double>xyz[2]

try:
err = GDALRPCTransform(self._transformer, bDstToSrc, n, x, y, z, panSuccess)
if err == GDALError.failure:
warnings.warn(
"Could not transform points using RPCs.",
TransformWarning)
res_xs = [0] * n
res_ys = [0] * n
res_zs = [0] * n
res_xs = np.zeros(n)
res_ys = np.zeros(n)
#res_zs = np.zeros(n)
rxview = res_xs
ryview = res_ys
checked = False
for i in range(n):
# 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(
"One or more points could not be transformed using RPCs.",
TransformWarning)
checked = True
res_xs[i] = x[i]
res_ys[i] = y[i]
res_zs[i] = z[i]
rxview[i] = x[i]
ryview[i] = y[i]
#res_zs[i] = z[i]
finally:
CPLFree(x)
CPLFree(y)
Expand Down Expand Up @@ -285,7 +294,7 @@ cdef class GCPTransformerBase:
Parameters
----------
xs, ys, zs : list
xs, ys, zs : ndarray
List of coordinates to be transformed. May be either pixel/line/height or
lon/lat/height)
transform_direction : TransformDirection
Expand All @@ -304,29 +313,34 @@ cdef class GCPTransformerBase:
Returns
-------
tuple of list
tuple of ndarray
"""
if self._transformer == NULL:
raise ValueError("Unexpected NULL transformer")

cdef int i
cdef Py_ssize_t i
cdef int n
cdef double *x = NULL
cdef double *y = NULL
cdef double *z = NULL
cdef int bDstToSrc = transform_direction
cdef int src_count
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))
y = <double *>CPLMalloc(n * sizeof(double))
z = <double *>CPLMalloc(n * sizeof(double))
panSuccess = <int *>CPLMalloc(n * sizeof(int))

for i in range(n):
x[i] = xs[i]
y[i] = ys[i]
z[i] = zs[i]
for i, xyz in enumerate(bi):
x[i] = <double>xyz[0]
y[i] = <double>xyz[1]
z[i] = <double>xyz[2]

try:
if self._tps:
Expand All @@ -337,20 +351,22 @@ cdef class GCPTransformerBase:
warnings.warn(
"Could not transform points using GCPs.",
TransformWarning)
res_xs = [0] * n
res_ys = [0] * n
res_zs = [0] * n
res_xs = np.zeros(n)
res_ys = np.zeros(n)
#res_zs = np.zeros(n)
rxview = res_xs
ryview = res_ys
checked = False
for i in range(n):
# 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(
"One or more points could not be transformed using GCPs.",
TransformWarning)
checked = True
res_xs[i] = x[i]
res_ys[i] = y[i]
res_zs[i] = z[i]
rxview[i] = x[i]
ryview[i] = y[i]
#res_zs[i] = z[i]
finally:
CPLFree(x)
CPLFree(y)
Expand Down
81 changes: 52 additions & 29 deletions rasterio/transform.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@

from contextlib import ExitStack
from functools import partial
import math
import numpy as np
import warnings
from numbers import Number

from affine import Affine

Expand Down Expand Up @@ -86,7 +86,7 @@ def index(
x,
y,
z=None,
op=math.floor,
op=np.floor,
precision=None,
transform_method=TransformMethod.affine,
**rpc_options
Expand Down Expand Up @@ -251,7 +251,7 @@ def xy(transform, rows, cols, zs=None, offset='center', **rpc_options):
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).
Parameters
Expand Down Expand Up @@ -330,21 +330,31 @@ def _ensure_arr_input(xs, ys, zs=None):
TransformError
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:
xs, ys, zs = np.broadcast_arrays(xs, ys, 0 if zs is None else zs)
broadcasted = np.broadcast(xs, ys, zs)
if broadcasted.ndim != 1:
raise TransformError(
"Input coordinates must be broadcastable to a 1d array"
)
except ValueError as error:
raise TransformError(
"Input coordinates must be broadcastable to a 1d array"
) from error
return np.atleast_1d(xs), np.atleast_1d(ys), np.atleast_1d(zs)
raise TransformError() from error

return xs, ys, zs

def __enter__(self):
return self

def __exit__(self, *args):
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.
Parameters
Expand Down Expand Up @@ -378,18 +388,30 @@ def rowcol(self, xs, ys, zs=None, op=math.floor, precision=None):
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)

try:
new_cols, new_rows = self._transform(
xs, ys, zs, transform_direction=TransformDirection.reverse
)

if len(new_rows) == 1 and not AS_ARR:
return (op(new_rows[0]), op(new_cols[0]))
is_op_ufunc = isinstance(op, np.ufunc)
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:
return ([op(r) for r in new_rows], [op(c) for c in new_cols])
return new_rows, new_cols
except TypeError:
raise TransformError("Invalid inputs")

Expand Down Expand Up @@ -419,7 +441,7 @@ def xy(self, rows, cols, zs=None, offset='center'):
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)

if offset == 'center':
Expand All @@ -445,10 +467,11 @@ def xy(self, rows, cols, zs=None, offset='center'):
new_xs, new_ys = self._transform(
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:
return (list(new_xs), list(new_ys))
return new_xs.tolist(), new_ys.tolist()
except TypeError:
raise TransformError("Invalid inputs")

Expand Down Expand Up @@ -480,22 +503,22 @@ def __init__(self, affine_transform):
if not isinstance(affine_transform, Affine):
raise ValueError("Not an affine transform")
self._transformer = affine_transform
self._transform_arr = np.empty((3, 3), dtype='float64')

def _transform(self, xs, ys, zs, transform_direction):
transform = self._transform_arr
if transform_direction is TransformDirection.forward:
transform = self._transformer
transform.flat[:] = self._transformer
elif transform_direction is TransformDirection.reverse:
transform = ~self._transformer

is_arr = True if type(xs) in [list, tuple] else False
if is_arr:
a, b, c, d, e, f, _, _, _ = transform
transform_matrix = np.array([[a, b, c], [d, e, f]])
input_matrix = np.array([xs, ys, np.ones(len(xs))])
output_matrix = np.dot(transform_matrix, input_matrix)
return (list(output_matrix[0]), list(output_matrix[1]))
else:
return transform * (xs, ys)
transform.flat[:] = ~self._transformer

bi = np.broadcast(xs, ys)
input_matrix = np.empty((3, bi.size))
input_matrix[0] = bi.iters[0]
input_matrix[1] = bi.iters[1]
input_matrix[2] = 1
transform.dot(input_matrix, out=input_matrix)
return input_matrix[0], input_matrix[1]


def __repr__(self):
Expand Down

0 comments on commit 7fb4442

Please sign in to comment.