Skip to content

Commit

Permalink
Add the private _to_numpy function to convert anything to a numpy array
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Nov 5, 2024
1 parent f91475d commit 0f8d82f
Showing 1 changed file with 47 additions and 21 deletions.
68 changes: 47 additions & 21 deletions pygmt/clib/conversion.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,52 @@ def dataarray_to_matrix(
return matrix, region, inc


def _to_numpy(data: Any) -> np.ndarray:
"""
Convert an array-like object to a C contiguous numpy array.
The function aims to convert any array-like objects (e.g., Python lists or tuples,
NumPy arrays with various dtypes, pandas.Series with NumPy/Pandas/PyArrow dtypes,
PyArrow arrays with various dtypes) to a numpy array.
The function is internally used in the ``vectors_to_arrays`` function, which is
responsible for converting a sequence of vectors to a list of C contiguous numpy
arrays. Thus, the function uses the :numpy:func:`numpy.ascontiguousarray` function
rather than the :numpy:func:`numpy.asarray`/:numpy::func:`numpy.asanyarray`
functions, to ensure the returned numpy array is C contiguous.
Parameters
----------
data
The array-like object to convert.
Returns
-------
array
The C contiguous numpy array.
"""
# Mapping of unsupported dtypes to the expected numpy dtype.
dtypes: dict[str, type] = {
"date32[day][pyarrow]": np.datetime64,
"date64[ms][pyarrow]": np.datetime64,
}

if (
hasattr(data, "isna")
and data.isna().any()
and Version(pd.__version__) < Version("2.2")
):
# Workaround for dealing with pd.NA with pandas < 2.2.
# Bug report at: https://github.com/GenericMappingTools/pygmt/issues/2844
# Following SPEC0, pandas 2.1 will be dropped in 2025 Q3, so it's likely
# we can remove the workaround in PyGMT v0.17.0.
array = np.ascontiguousarray(data.astype(float))
else:
vec_dtype = str(getattr(data, "dtype", ""))
array = np.ascontiguousarray(data, dtype=dtypes.get(vec_dtype))
return array


def vectors_to_arrays(vectors: Sequence[Any]) -> list[np.ndarray]:
"""
Convert 1-D vectors (scalars, lists, or array-like) to C contiguous 1-D arrays.
Expand Down Expand Up @@ -171,27 +217,7 @@ def vectors_to_arrays(vectors: Sequence[Any]) -> list[np.ndarray]:
>>> all(i.ndim == 1 for i in arrays)
True
"""
dtypes = {
"date32[day][pyarrow]": np.datetime64,
"date64[ms][pyarrow]": np.datetime64,
}
arrays = []
for vector in vectors:
if (
hasattr(vector, "isna")
and vector.isna().any()
and Version(pd.__version__) < Version("2.2")
):
# Workaround for dealing with pd.NA with pandas < 2.2.
# Bug report at: https://github.com/GenericMappingTools/pygmt/issues/2844
# Following SPEC0, pandas 2.1 will be dropped in 2025 Q3, so it's likely
# we can remove the workaround in PyGMT v0.17.0.
array = np.ascontiguousarray(vector.astype(float))
else:
vec_dtype = str(getattr(vector, "dtype", ""))
array = np.ascontiguousarray(vector, dtype=dtypes.get(vec_dtype))
arrays.append(array)
return arrays
return [_to_numpy(vector) for vector in vectors]


def sequence_to_ctypes_array(
Expand Down

0 comments on commit 0f8d82f

Please sign in to comment.