Skip to content

Commit 96b1263

Browse files
authored
clib.conversion: Add type hints and improve docstrings for dataarray_to_matrix/vectors_to_arrays/array_to_datetime (#3496)
1 parent 90c1262 commit 96b1263

File tree

1 file changed

+47
-47
lines changed

1 file changed

+47
-47
lines changed

pygmt/clib/conversion.py

+47-47
Original file line numberDiff line numberDiff line change
@@ -5,41 +5,45 @@
55
import ctypes as ctp
66
import warnings
77
from collections.abc import Sequence
8+
from typing import Any
89

910
import numpy as np
1011
import pandas as pd
12+
import xarray as xr
1113
from packaging.version import Version
1214
from pygmt.exceptions import GMTInvalidInput
1315

1416

15-
def dataarray_to_matrix(grid):
17+
def dataarray_to_matrix(
18+
grid: xr.DataArray,
19+
) -> tuple[np.ndarray, list[float], list[float]]:
1620
"""
17-
Transform an xarray.DataArray into a data 2-D array and metadata.
21+
Transform an xarray.DataArray into a 2-D numpy array and metadata.
1822
19-
Use this to extract the underlying numpy array of data and the region and
20-
increment for the grid.
23+
Use this to extract the underlying numpy array of data and the region and increment
24+
for the grid.
2125
22-
Only allows grids with two dimensions and constant grid spacing (GMT
23-
doesn't allow variable grid spacing). If the latitude and/or longitude
24-
increments of the input grid are negative, the output matrix will be
25-
sorted by the DataArray coordinates to yield positive increments.
26+
Only allows grids with two dimensions and constant grid spacings (GMT doesn't allow
27+
variable grid spacings). If the latitude and/or longitude increments of the input
28+
grid are negative, the output matrix will be sorted by the DataArray coordinates to
29+
yield positive increments.
2630
27-
If the underlying data array is not C contiguous, for example if it's a
28-
slice of a larger grid, a copy will need to be generated.
31+
If the underlying data array is not C contiguous, for example, if it's a slice of a
32+
larger grid, a copy will need to be generated.
2933
3034
Parameters
3135
----------
32-
grid : xarray.DataArray
33-
The input grid as a DataArray instance. Information is retrieved from
34-
the coordinate arrays, not from headers.
36+
grid
37+
The input grid as a DataArray instance. Information is retrieved from the
38+
coordinate arrays, not from headers.
3539
3640
Returns
3741
-------
38-
matrix : 2-D array
42+
matrix
3943
The 2-D array of data from the grid.
40-
region : list
44+
region
4145
The West, East, South, North boundaries of the grid.
42-
inc : list
46+
inc
4347
The grid spacing in East-West and North-South, respectively.
4448
4549
Raises
@@ -64,8 +68,8 @@ def dataarray_to_matrix(grid):
6468
(180, 360)
6569
>>> matrix.flags.c_contiguous
6670
True
67-
>>> # Using a slice of the grid, the matrix will be copied to guarantee
68-
>>> # that it's C-contiguous in memory. The increment should be unchanged.
71+
>>> # Using a slice of the grid, the matrix will be copied to guarantee that it's
72+
>>> # C-contiguous in memory. The increment should be unchanged.
6973
>>> matrix, region, inc = dataarray_to_matrix(grid[10:41, 30:101])
7074
>>> matrix.flags.c_contiguous
7175
True
@@ -75,7 +79,7 @@ def dataarray_to_matrix(grid):
7579
[-150.0, -79.0, -80.0, -49.0]
7680
>>> print(inc)
7781
[1.0, 1.0]
78-
>>> # but not if only taking every other grid point.
82+
>>> # The increment should change accordingly if taking every other grid point.
7983
>>> matrix, region, inc = dataarray_to_matrix(grid[10:41:2, 30:101:2])
8084
>>> matrix.flags.c_contiguous
8185
True
@@ -87,21 +91,19 @@ def dataarray_to_matrix(grid):
8791
[2.0, 2.0]
8892
"""
8993
if len(grid.dims) != 2:
90-
raise GMTInvalidInput(
91-
f"Invalid number of grid dimensions '{len(grid.dims)}'. Must be 2."
92-
)
94+
msg = f"Invalid number of grid dimensions 'len({grid.dims})'. Must be 2."
95+
raise GMTInvalidInput(msg)
96+
9397
# Extract region and inc from the grid
94-
region = []
95-
inc = []
96-
# Reverse the dims because it is rows, columns ordered. In geographic
97-
# grids, this would be North-South, East-West. GMT's region and inc are
98-
# East-West, North-South.
98+
region, inc = [], []
99+
# Reverse the dims because it is rows, columns ordered. In geographic grids, this
100+
# would be North-South, East-West. GMT's region and inc are East-West, North-South.
99101
for dim in grid.dims[::-1]:
100102
coord = grid.coords[dim].to_numpy()
101-
coord_incs = coord[1:] - coord[0:-1]
103+
coord_incs = coord[1:] - coord[:-1]
102104
coord_inc = coord_incs[0]
103105
if not np.allclose(coord_incs, coord_inc):
104-
# calculate the increment if irregular spacing is found
106+
# Calculate the increment if irregular spacing is found.
105107
coord_inc = (coord[-1] - coord[0]) / (coord.size - 1)
106108
msg = (
107109
f"Grid may have irregular spacing in the '{dim}' dimension, "
@@ -110,9 +112,8 @@ def dataarray_to_matrix(grid):
110112
)
111113
warnings.warn(msg, category=RuntimeWarning, stacklevel=2)
112114
if coord_inc == 0:
113-
raise GMTInvalidInput(
114-
f"Grid has a zero increment in the '{dim}' dimension."
115-
)
115+
msg = f"Grid has a zero increment in the '{dim}' dimension."
116+
raise GMTInvalidInput(msg)
116117
region.extend(
117118
[
118119
coord.min() - coord_inc / 2 * grid.gmt.registration,
@@ -131,26 +132,25 @@ def dataarray_to_matrix(grid):
131132
return matrix, region, inc
132133

133134

134-
def vectors_to_arrays(vectors):
135+
def vectors_to_arrays(vectors: Sequence[Any]) -> list[np.ndarray]:
135136
"""
136-
Convert 1-D vectors (lists, arrays, or pandas.Series) to C contiguous 1-D arrays.
137+
Convert 1-D vectors (scalars, lists, or array-like) to C contiguous 1-D arrays.
137138
138-
Arrays must be in C contiguous order for us to pass their memory pointers
139-
to GMT. If any are not, convert them to C order (which requires copying the
140-
memory). This usually happens when vectors are columns of a 2-D array or
141-
have been sliced.
139+
Arrays must be in C contiguous order for us to pass their memory pointers to GMT.
140+
If any are not, convert them to C order (which requires copying the memory). This
141+
usually happens when vectors are columns of a 2-D array or have been sliced.
142142
143-
If a vector is a list or pandas.Series, get the underlying numpy array.
143+
The returned arrays are guaranteed to be C contiguous and at least 1-D.
144144
145145
Parameters
146146
----------
147-
vectors : list of lists, 1-D arrays, or pandas.Series
147+
vectors
148148
The vectors that must be converted.
149149
150150
Returns
151151
-------
152-
arrays : list of 1-D arrays
153-
The converted numpy arrays
152+
arrays
153+
List of converted numpy arrays.
154154
155155
Examples
156156
--------
@@ -307,16 +307,15 @@ def strings_to_ctypes_array(strings: Sequence[str]) -> ctp.Array:
307307
return (ctp.c_char_p * len(strings))(*[s.encode() for s in strings])
308308

309309

310-
def array_to_datetime(array):
310+
def array_to_datetime(array: Sequence[Any]) -> np.ndarray:
311311
"""
312312
Convert a 1-D datetime array from various types into numpy.datetime64.
313313
314-
If the input array is not in legal datetime formats, raise a ValueError
315-
exception.
314+
If the input array is not in legal datetime formats, raise a ValueError exception.
316315
317316
Parameters
318317
----------
319-
array : list or 1-D array
318+
array
320319
The input datetime array in various formats.
321320
322321
Supported types:
@@ -328,7 +327,8 @@ def array_to_datetime(array):
328327
329328
Returns
330329
-------
331-
array : 1-D datetime array in numpy.datetime64
330+
array
331+
1-D datetime array in numpy.datetime64.
332332
333333
Raises
334334
------

0 commit comments

Comments
 (0)