From 96b12634d95a0fd2a262bcf5245a942cb4077e69 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Fri, 11 Oct 2024 14:29:44 +0800 Subject: [PATCH] clib.conversion: Add type hints and improve docstrings for dataarray_to_matrix/vectors_to_arrays/array_to_datetime (#3496) --- pygmt/clib/conversion.py | 94 ++++++++++++++++++++-------------------- 1 file changed, 47 insertions(+), 47 deletions(-) diff --git a/pygmt/clib/conversion.py b/pygmt/clib/conversion.py index ee80ec1a187..07f8756adcd 100644 --- a/pygmt/clib/conversion.py +++ b/pygmt/clib/conversion.py @@ -5,41 +5,45 @@ import ctypes as ctp import warnings from collections.abc import Sequence +from typing import Any import numpy as np import pandas as pd +import xarray as xr from packaging.version import Version from pygmt.exceptions import GMTInvalidInput -def dataarray_to_matrix(grid): +def dataarray_to_matrix( + grid: xr.DataArray, +) -> tuple[np.ndarray, list[float], list[float]]: """ - Transform an xarray.DataArray into a data 2-D array and metadata. + Transform an xarray.DataArray into a 2-D numpy array and metadata. - Use this to extract the underlying numpy array of data and the region and - increment for the grid. + Use this to extract the underlying numpy array of data and the region and increment + for the grid. - Only allows grids with two dimensions and constant grid spacing (GMT - doesn't allow variable grid spacing). If the latitude and/or longitude - increments of the input grid are negative, the output matrix will be - sorted by the DataArray coordinates to yield positive increments. + Only allows grids with two dimensions and constant grid spacings (GMT doesn't allow + variable grid spacings). If the latitude and/or longitude increments of the input + grid are negative, the output matrix will be sorted by the DataArray coordinates to + yield positive increments. - If the underlying data array is not C contiguous, for example if it's a - slice of a larger grid, a copy will need to be generated. + If the underlying data array is not C contiguous, for example, if it's a slice of a + larger grid, a copy will need to be generated. Parameters ---------- - grid : xarray.DataArray - The input grid as a DataArray instance. Information is retrieved from - the coordinate arrays, not from headers. + grid + The input grid as a DataArray instance. Information is retrieved from the + coordinate arrays, not from headers. Returns ------- - matrix : 2-D array + matrix The 2-D array of data from the grid. - region : list + region The West, East, South, North boundaries of the grid. - inc : list + inc The grid spacing in East-West and North-South, respectively. Raises @@ -64,8 +68,8 @@ def dataarray_to_matrix(grid): (180, 360) >>> matrix.flags.c_contiguous True - >>> # Using a slice of the grid, the matrix will be copied to guarantee - >>> # that it's C-contiguous in memory. The increment should be unchanged. + >>> # Using a slice of the grid, the matrix will be copied to guarantee that it's + >>> # C-contiguous in memory. The increment should be unchanged. >>> matrix, region, inc = dataarray_to_matrix(grid[10:41, 30:101]) >>> matrix.flags.c_contiguous True @@ -75,7 +79,7 @@ def dataarray_to_matrix(grid): [-150.0, -79.0, -80.0, -49.0] >>> print(inc) [1.0, 1.0] - >>> # but not if only taking every other grid point. + >>> # The increment should change accordingly if taking every other grid point. >>> matrix, region, inc = dataarray_to_matrix(grid[10:41:2, 30:101:2]) >>> matrix.flags.c_contiguous True @@ -87,21 +91,19 @@ def dataarray_to_matrix(grid): [2.0, 2.0] """ if len(grid.dims) != 2: - raise GMTInvalidInput( - f"Invalid number of grid dimensions '{len(grid.dims)}'. Must be 2." - ) + msg = f"Invalid number of grid dimensions 'len({grid.dims})'. Must be 2." + raise GMTInvalidInput(msg) + # Extract region and inc from the grid - region = [] - inc = [] - # Reverse the dims because it is rows, columns ordered. In geographic - # grids, this would be North-South, East-West. GMT's region and inc are - # East-West, North-South. + region, inc = [], [] + # Reverse the dims because it is rows, columns ordered. In geographic grids, this + # would be North-South, East-West. GMT's region and inc are East-West, North-South. for dim in grid.dims[::-1]: coord = grid.coords[dim].to_numpy() - coord_incs = coord[1:] - coord[0:-1] + coord_incs = coord[1:] - coord[:-1] coord_inc = coord_incs[0] if not np.allclose(coord_incs, coord_inc): - # calculate the increment if irregular spacing is found + # Calculate the increment if irregular spacing is found. coord_inc = (coord[-1] - coord[0]) / (coord.size - 1) msg = ( f"Grid may have irregular spacing in the '{dim}' dimension, " @@ -110,9 +112,8 @@ def dataarray_to_matrix(grid): ) warnings.warn(msg, category=RuntimeWarning, stacklevel=2) if coord_inc == 0: - raise GMTInvalidInput( - f"Grid has a zero increment in the '{dim}' dimension." - ) + msg = f"Grid has a zero increment in the '{dim}' dimension." + raise GMTInvalidInput(msg) region.extend( [ coord.min() - coord_inc / 2 * grid.gmt.registration, @@ -131,26 +132,25 @@ def dataarray_to_matrix(grid): return matrix, region, inc -def vectors_to_arrays(vectors): +def vectors_to_arrays(vectors: Sequence[Any]) -> list[np.ndarray]: """ - Convert 1-D vectors (lists, arrays, or pandas.Series) to C contiguous 1-D arrays. + Convert 1-D vectors (scalars, lists, or array-like) to C contiguous 1-D arrays. - Arrays must be in C contiguous order for us to pass their memory pointers - to GMT. If any are not, convert them to C order (which requires copying the - memory). This usually happens when vectors are columns of a 2-D array or - have been sliced. + Arrays must be in C contiguous order for us to pass their memory pointers to GMT. + If any are not, convert them to C order (which requires copying the memory). This + usually happens when vectors are columns of a 2-D array or have been sliced. - If a vector is a list or pandas.Series, get the underlying numpy array. + The returned arrays are guaranteed to be C contiguous and at least 1-D. Parameters ---------- - vectors : list of lists, 1-D arrays, or pandas.Series + vectors The vectors that must be converted. Returns ------- - arrays : list of 1-D arrays - The converted numpy arrays + arrays + List of converted numpy arrays. Examples -------- @@ -307,16 +307,15 @@ def strings_to_ctypes_array(strings: Sequence[str]) -> ctp.Array: return (ctp.c_char_p * len(strings))(*[s.encode() for s in strings]) -def array_to_datetime(array): +def array_to_datetime(array: Sequence[Any]) -> np.ndarray: """ Convert a 1-D datetime array from various types into numpy.datetime64. - If the input array is not in legal datetime formats, raise a ValueError - exception. + If the input array is not in legal datetime formats, raise a ValueError exception. Parameters ---------- - array : list or 1-D array + array The input datetime array in various formats. Supported types: @@ -328,7 +327,8 @@ def array_to_datetime(array): Returns ------- - array : 1-D datetime array in numpy.datetime64 + array + 1-D datetime array in numpy.datetime64. Raises ------