diff --git a/docs/conf.py b/docs/conf.py index 665dc4a5..c6095044 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -56,7 +56,11 @@ def skip_aliases(app, what, name, obj, skip, options): """Skip documentation for classes that are exported from multiple modules.""" # For names that are defined in a private sub-module and aliased into a # public package, hide the definition. - if name.startswith("nitypes.time._") or name.startswith("nitypes.waveform._"): + if ( + name.startswith("nitypes.complex._") + or name.startswith("nitypes.time._") + or name.startswith("nitypes.waveform._") + ): skip = True return skip diff --git a/src/nitypes/_arguments.py b/src/nitypes/_arguments.py index 5afbe94a..ba2f3d1a 100644 --- a/src/nitypes/_arguments.py +++ b/src/nitypes/_arguments.py @@ -6,7 +6,12 @@ import numpy as np import numpy.typing as npt -from nitypes._exceptions import invalid_arg_type, invalid_arg_value, unsupported_arg +from nitypes._exceptions import ( + invalid_arg_type, + invalid_arg_value, + unsupported_arg, + unsupported_dtype, +) def arg_to_float( @@ -144,14 +149,7 @@ def validate_dtype(dtype: npt.DTypeLike, supported_dtypes: tuple[npt.DTypeLike, if not isinstance(dtype, (type, np.dtype)): dtype = np.dtype(dtype) if not np.isdtype(dtype, supported_dtypes): - # Remove duplicate names because distinct types (e.g. int vs. long) may have the same name - # ("int32"). - supported_dtype_names = {np.dtype(d).name: None for d in supported_dtypes}.keys() - raise TypeError( - "The requested data type is not supported.\n\n" - f"Data type: {np.dtype(dtype)}\n" - f"Supported data types: {', '.join(supported_dtype_names)}" - ) + raise unsupported_dtype("requested data type", dtype, supported_dtypes) def validate_unsupported_arg(arg_description: str, value: object) -> None: diff --git a/src/nitypes/_exceptions.py b/src/nitypes/_exceptions.py index 890bf6e2..784ad1ac 100644 --- a/src/nitypes/_exceptions.py +++ b/src/nitypes/_exceptions.py @@ -3,6 +3,9 @@ import reprlib import sys +import numpy as np +import numpy.typing as npt + def add_note(exception: Exception, note: str) -> None: """Add a note to an exception. @@ -66,6 +69,20 @@ def unsupported_arg(arg_description: str, value: object) -> ValueError: ) +def unsupported_dtype( + arg_description: str, dtype: npt.DTypeLike, supported_dtypes: tuple[npt.DTypeLike, ...] +) -> TypeError: + """Create a TypeError for an unsupported dtype.""" + # Remove duplicate names because distinct types (e.g. int vs. long) may have the same name + # ("int32"). + supported_dtype_names = {np.dtype(d).name: None for d in supported_dtypes}.keys() + return TypeError( + f"The {arg_description} is not supported.\n\n" + f"Data type: {np.dtype(dtype)}\n" + f"Supported data types: {', '.join(supported_dtype_names)}" + ) + + # English-specific hack. This is why we prefer "Key: value" for localizable errors. TODO: consider # moving the full strings into a string table instead of building them out of English noun phrases. def _a(noun: str) -> str: diff --git a/src/nitypes/complex/__init__.py b/src/nitypes/complex/__init__.py new file mode 100644 index 00000000..14e9fcdd --- /dev/null +++ b/src/nitypes/complex/__init__.py @@ -0,0 +1,129 @@ +"""Complex number data types for NI Python APIs. + +================ +Complex Integers +================ + +Some NI driver APIs (such as NI-FGEN, NI-SCOPE, NI-RFSA, and NI-RFSG) use complex numbers to +represent I/Q data. Python and NumPy have native support for complex floating-point numbers, but +not complex integers, so the :mod:`nitypes.complex` submodule provides a NumPy representation of +complex integers. + +:any:`ComplexInt32DType` is a NumPy structured data type object representing a complex integer with +16-bit ``real`` and ``imag`` fields. This structured data type has the same memory layout as the +``NIComplexI16`` C struct used by NI driver APIs. + +For more information about NumPy structured data types, see the +:ref:`NumPy documentation on structured arrays `. + +.. note:: + In ``NIComplexI16``, the number 16 refers to the number of bits in each field. In + :any:`ComplexInt32DType`, the number 32 refers to the total number of bits, following the + precedent set by NumPy's other complex types. For example, :any:`numpy.complex128` contains + 64-bit ``real`` and ``imag`` fields. + +Constructing arrays of complex integers +--------------------------------------- + +You can construct an array of complex integers from a sequence of tuples using :func:`numpy.array`: + +>>> import numpy as np +>>> np.array([(1, 2), (3, 4)], dtype=ComplexInt32DType) +array([(1, 2), (3, 4)], dtype=[('real', '>> np.zeros(3, dtype=ComplexInt32DType) +array([(0, 0), (0, 0), (0, 0)], dtype=[('real', '>> x = np.array([(1, 2), (3, 4), (5, 6)], dtype=ComplexInt32DType) +>>> x[0] +np.void((1, 2), dtype=[('real', '>> x[1] +np.void((3, 4), dtype=[('real', '>> x[0][0] +np.int16(1) +>>> x[0][1] +np.int16(2) + +You can also index by the field names ``real`` and ``imag``: + +>>> x[0]['real'] +np.int16(1) +>>> x[0]['imag'] +np.int16(2) + +Or you can index the entire array by the field names ``real`` and ``imag``: + +>>> x['real'] +array([1, 3, 5], dtype=int16) +>>> x['imag'] +array([2, 4, 6], dtype=int16) + +Arrays of complex integers support slicing and negative indices like any other array: + +>>> x[0:2] +array([(1, 2), (3, 4)], dtype=[('real', '>> x[1:] +array([(3, 4), (5, 6)], dtype=[('real', '>> x[-1] +np.void((5, 6), dtype=[('real', '>> x[0].item() +(1, 2) +>>> [y.item() for y in x] +[(1, 2), (3, 4), (5, 6)] + +To convert NumPy arrays between between different complex number data types, use the +:func:`convert_complex` function: + +>>> convert_complex(np.complex128, x) +array([1.+2.j, 3.+4.j, 5.+6.j]) +>>> convert_complex(ComplexInt32DType, np.array([1.23+4.56j])) +array([(1, 4)], dtype=[('real', '>> convert_complex(np.complex128, x[0]) +np.complex128(1+2j) +>>> convert_complex(ComplexInt32DType, np.complex128(3+4j)) +np.void((3, 4), dtype=[('real', ' np.ndarray[_Shape, np.dtype[_ScalarType]]: ... + + +@overload +def convert_complex( + requested_dtype: npt.DTypeLike, value: np.ndarray[_Shape, Any] +) -> np.ndarray[_Shape, Any]: ... + + +# https://numpy.org/doc/2.2/reference/typing.html#d-arrays +# "While thus not strictly correct, all operations are that can potentially perform a 0D-array -> +# scalar cast are currently annotated as exclusively returning an ndarray." +@overload +def convert_complex( + requested_dtype: type[_ScalarType] | np.dtype[_ScalarType], + value: np.generic[Any], +) -> np.ndarray[tuple[()], np.dtype[_ScalarType]]: ... + + +@overload +def convert_complex( + requested_dtype: npt.DTypeLike, + value: np.generic[Any], +) -> np.ndarray[tuple[()], Any]: ... + + +def convert_complex( + requested_dtype: npt.DTypeLike, value: np.ndarray[_Shape, Any] | np.generic[Any] +) -> np.ndarray[_Shape, Any]: + """Convert a NumPy array or scalar of complex numbers to the specified dtype. + + Args: + requested_dtype: The NumPy data type to convert to. Supported data types: + :any:`numpy.complex64`, :any:`numpy.complex128`, :any:`ComplexInt32DType`. + value: The NumPy array or scalar to convert. + + Returns: + The value converted to the specified dtype. + """ + validate_dtype(requested_dtype, _COMPLEX_DTYPES) + if requested_dtype == value.dtype: + return cast(np.ndarray[_Shape, Any], value) + elif requested_dtype == ComplexInt32DType or value.dtype == ComplexInt32DType: + # ndarray.view on scalars requires the source and destination types to have the same size, + # so reshape the scalar into an 1-element array before converting and index it afterwards. + # shape == () means this is either a scalar (np.generic) or a 0-dimension array, but mypy + # doesn't know that. + if value.shape == (): + return cast( + np.ndarray[_Shape, Any], + _convert_complexint32_array(requested_dtype, value.reshape(1))[0], + ) + else: + return _convert_complexint32_array( + requested_dtype, cast(np.ndarray[_Shape, Any], value) + ) + else: + return value.astype(requested_dtype) + + +def _convert_complexint32_array( + requested_dtype: npt.DTypeLike | type[_ScalarType] | np.dtype[_ScalarType], + value: np.ndarray[_Shape, Any], +) -> np.ndarray[_Shape, np.dtype[_ScalarType]]: + if not isinstance(requested_dtype, np.dtype): + requested_dtype = np.dtype(requested_dtype) + + requested_field_dtype = _FIELD_DTYPE.get(requested_dtype) + if requested_field_dtype is None: + raise unsupported_dtype("requested data type", requested_dtype, _COMPLEX_DTYPES) + + value_field_dtype = _FIELD_DTYPE.get(value.dtype) + if value_field_dtype is None: + raise unsupported_dtype("array data type", value.dtype, _COMPLEX_DTYPES) + + return value.view(value_field_dtype).astype(requested_field_dtype).view(requested_dtype) diff --git a/src/nitypes/complex/_dtypes.py b/src/nitypes/complex/_dtypes.py new file mode 100644 index 00000000..a07e58f7 --- /dev/null +++ b/src/nitypes/complex/_dtypes.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import numpy as np +from typing_extensions import TypeAlias + +ComplexInt32Base: TypeAlias = np.void +"""Type alias for the base type of :any:`ComplexInt32DType`, which is :any:`numpy.void`.""" + +ComplexInt32DType = np.dtype((ComplexInt32Base, [("real", np.int16), ("imag", np.int16)])) +"""NumPy structured data type for a complex integer with 16-bit ``real`` and ``imag`` fields.""" diff --git a/tests/unit/complex/__init__.py b/tests/unit/complex/__init__.py new file mode 100644 index 00000000..3986ed07 --- /dev/null +++ b/tests/unit/complex/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the nitypes.complex package.""" diff --git a/tests/unit/complex/test_conversion.py b/tests/unit/complex/test_conversion.py new file mode 100644 index 00000000..1e783631 --- /dev/null +++ b/tests/unit/complex/test_conversion.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import numpy.typing as npt +from typing_extensions import assert_type + +from nitypes.complex import ComplexInt32Base, ComplexInt32DType, convert_complex + + +############################################################################### +# convert arrays +############################################################################### +def test___complexint32_array_to_complex64_array___convert_complex___converts_array() -> None: + value_in = np.array([(1, 2), (3, -4), (-5, 6), (-7, -8)], ComplexInt32DType) + + value_out = convert_complex(np.complex64, value_in) + + assert_type(value_out, npt.NDArray[np.complex64]) + assert isinstance(value_out, np.ndarray) and value_out.dtype == np.complex64 + assert list(value_out) == [1 + 2j, 3 - 4j, -5 + 6j, -7 - 8j] + + +def test___complexint32_array_to_complex128_array___convert_complex___converts_array() -> None: + value_in = np.array([(1, 2), (3, -4), (-5, 6), (-7, -8)], ComplexInt32DType) + + value_out = convert_complex(np.complex128, value_in) + + assert_type(value_out, npt.NDArray[np.complex128]) + assert isinstance(value_out, np.ndarray) and value_out.dtype == np.complex128 + assert list(value_out) == [1 + 2j, 3 - 4j, -5 + 6j, -7 - 8j] + + +def test___complexint32_array_to_complexint32_array___convert_complex___returns_original_array() -> ( + None +): + value_in = np.array([(1, 2), (3, -4), (-5, 6), (-7, -8)], ComplexInt32DType) + + value_out = convert_complex(ComplexInt32DType, value_in) + + assert_type(value_out, npt.NDArray[ComplexInt32Base]) + assert value_out is value_in + + +def test___complex64_array_to_complexint32_array___convert_complex___converts_array() -> None: + value_in = np.array([1 + 2j, 3 - 4j, -5 + 6j, -7 - 8j], np.complex64) + + value_out = convert_complex(ComplexInt32DType, value_in) + + assert_type(value_out, npt.NDArray[ComplexInt32Base]) + assert isinstance(value_out, np.ndarray) and value_out.dtype == ComplexInt32DType + assert [x.item() for x in value_out] == [(1, 2), (3, -4), (-5, 6), (-7, -8)] + + +def test___complex64_array_to_complex64_array___convert_complex___returns_original_array() -> None: + value_in = np.array([1 + 2j, 3 - 4j, -5 + 6j, -7 - 8j], np.complex64) + + value_out = convert_complex(np.complex64, value_in) + + assert_type(value_out, npt.NDArray[np.complex64]) + assert value_out is value_in + + +def test___complex64_array_to_complex128_array___convert_complex___converts_array() -> None: + value_in = np.array([1.23 + 4.56j, 6.78 - 9.01j], np.complex64) + + value_out = convert_complex(np.complex128, value_in) + + assert_type(value_out, npt.NDArray[np.complex128]) + assert isinstance(value_out, np.ndarray) and value_out.dtype == np.complex128 + # complex64 bruises the numbers (e.g. np.complex128(1.2300000190734863+4.559999942779541j)) so + # round to 3 decimal places. + assert list(np.round(value_out, 3)) == [1.23 + 4.56j, 6.78 - 9.01j] + + +def test___complex128_array_to_complexint32_array___convert_complex___converts_array() -> None: + value_in = np.array([1 + 2j, 3 - 4j, -5 + 6j, -7 - 8j], np.complex128) + + value_out = convert_complex(ComplexInt32DType, value_in) + + assert_type(value_out, npt.NDArray[ComplexInt32Base]) + assert isinstance(value_out, np.ndarray) and value_out.dtype == ComplexInt32DType + assert [x.item() for x in value_out] == [(1, 2), (3, -4), (-5, 6), (-7, -8)] + + +def test___complex128_array_to_complex64_array___convert_complex___converts_array() -> None: + value_in = np.array([1.23 + 4.56j, 6.78 - 9.01j], np.complex128) + + value_out = convert_complex(np.complex64, value_in) + + assert_type(value_out, npt.NDArray[np.complex64]) + assert isinstance(value_out, np.ndarray) and value_out.dtype == np.complex64 + assert list(value_out) == [1.23 + 4.56j, 6.78 - 9.01j] + + +def test___complex128_array_to_complex128_array___convert_complex___returns_original_array() -> ( + None +): + value_in = np.array([1.23 + 4.56j, 6.78 - 9.01j], np.complex128) + + value_out = convert_complex(np.complex128, value_in) + + assert_type(value_out, npt.NDArray[np.complex128]) + assert value_out is value_in + + +def test___2d_complexint32_array_to_complex128_array___convert_complex___preserves_shape() -> None: + value_in = np.array( + [[(1, 2), (3, -4)], [(-5, 6), (-7, -8)], [(9, 10), (11, 12)]], ComplexInt32DType + ) + + value_out = convert_complex(np.complex128, value_in) + + # Use npt.NDArray because np.array() can't infer the array shape when type checking. + assert_type(value_out, npt.NDArray[np.complex128]) + assert isinstance(value_out, np.ndarray) and value_out.shape == (3, 2) + assert [list(x) for x in value_out] == [ + [1 + 2j, 3 - 4j], + [-5 + 6j, -7 - 8j], + [9 + 10j, 11 + 12j], + ] + + +def test___2d_complex64_array_to_complex128_array___convert_complex___preserves_shape() -> None: + value_in = np.array([[1 + 2j, 3 - 4j], [-5 + 6j, -7 - 8j], [9 + 10j, 11 + 12j]], np.complex64) + + value_out = convert_complex(np.complex128, value_in) + + # Use npt.NDArray because np.array() can't infer the array shape when type checking. + assert_type(value_out, npt.NDArray[np.complex128]) + assert isinstance(value_out, np.ndarray) and value_out.shape == (3, 2) + assert [list(x) for x in value_out] == [ + [1 + 2j, 3 - 4j], + [-5 + 6j, -7 - 8j], + [9 + 10j, 11 + 12j], + ] + + +def test___arrays_with_static_shape___convert_complex___preserves_static_and_runtime_shape() -> ( + None +): + # np.zeros() can infer the array shape when type checking because it takes a shape argument. + value_in_1d = np.zeros(3, np.complex64) + value_in_2d = np.zeros((4, 5), np.complex128) + value_in_3d = np.zeros((6, 7, 8), ComplexInt32DType) + + value_out_1d = convert_complex(np.complex128, value_in_1d) + value_out_2d = convert_complex(ComplexInt32DType, value_in_2d) + value_out_3d = convert_complex(np.complex64, value_in_3d) + + assert_type(value_out_1d, np.ndarray[tuple[int], np.dtype[np.complex128]]) + assert_type(value_out_2d, np.ndarray[tuple[int, int], np.dtype[ComplexInt32Base]]) + assert_type(value_out_3d, np.ndarray[tuple[int, int, int], np.dtype[np.complex64]]) + assert isinstance(value_out_1d, np.ndarray) and value_out_1d.shape == (3,) + assert isinstance(value_out_2d, np.ndarray) and value_out_2d.shape == (4, 5) + assert isinstance(value_out_3d, np.ndarray) and value_out_3d.shape == (6, 7, 8) + + +############################################################################### +# convert scalars +############################################################################### +def test___complexint32_scalar_to_complex128_scalar___convert_complex___converts_scalar() -> None: + value_in = np.array([(1, 2)], ComplexInt32DType)[0] + assert_type(value_in, Any) # ¯\_(ツ)_/¯ + + value_out = convert_complex(np.complex128, value_in) + + assert_type(value_out, np.ndarray[Any, Any]) # This is less than ideal. + assert isinstance(value_out, np.complex128) + assert value_out == (1 + 2j) + + +def test___complex128_scalar_to_complexint32_scalar___convert_complex___converts_scalar() -> None: + value_in = np.complex128(1 + 2j) + + value_out = convert_complex(ComplexInt32DType, value_in) + + assert_type(value_out, np.ndarray[tuple[()], np.dtype[ComplexInt32Base]]) + assert isinstance(value_out, ComplexInt32Base) + assert value_out.item() == (1, 2) + + +def test___complexint32_scalar_to_complexint32_scalar___convert_complex___returns_original_scalar() -> ( + None +): + value_in = np.array([(1, 2)], ComplexInt32DType)[0] + assert_type(value_in, Any) # ¯\_(ツ)_/¯ + + value_out = convert_complex(ComplexInt32DType, value_in) + + assert_type(value_out, np.ndarray[Any, Any]) # This is less than ideal. + assert value_out is value_in + + +def test___complex64_scalar_to_complex128_scalar___convert_complex___converts_scalar() -> None: + value_in = np.complex64(1.23 + 4.56j) + + value_out = convert_complex(np.complex128, value_in) + + assert_type(value_out, np.ndarray[tuple[()], np.dtype[np.complex128]]) + assert isinstance(value_out, np.complex128) + # complex64 bruises the numbers (e.g. np.complex128(1.2300000190734863+4.559999942779541j)) so + # round to 3 decimal places. + assert np.round(value_out, 3) == (1.23 + 4.56j) diff --git a/tests/unit/complex/test_dtypes.py b/tests/unit/complex/test_dtypes.py new file mode 100644 index 00000000..f9e4df2e --- /dev/null +++ b/tests/unit/complex/test_dtypes.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import numpy.typing as npt +import pytest +from typing_extensions import assert_type + +from nitypes.complex import ComplexInt32Base, ComplexInt32DType + + +def test___complexint32_dtype___np_array___constructs_array_with_dtype() -> None: + value = np.array([(1, 2), (3, -4), (-5, 6), (-7, -8)], ComplexInt32DType) + + assert_type(value, npt.NDArray[ComplexInt32Base]) + assert isinstance(value, np.ndarray) and value.dtype == ComplexInt32DType + assert [x.item() for x in value] == [(1, 2), (3, -4), (-5, 6), (-7, -8)] + + +def test___complexint32_dtype___np_zeros___constructs_array_with_dtype() -> None: + value = np.zeros(3, ComplexInt32DType) + + assert_type(value, np.ndarray[tuple[int], np.dtype[ComplexInt32Base]]) + assert isinstance(value, np.ndarray) and value.dtype == ComplexInt32DType + assert [x.item() for x in value] == [(0, 0), (0, 0), (0, 0)] + + +def test___complexint32_array___index___returns_complexint32_scalar() -> None: + array = np.array([(1, 2), (3, -4)], ComplexInt32DType) + + value = array[1] + + assert_type(value, Any) # ¯\_(ツ)_/¯ + assert isinstance(value, ComplexInt32Base) # alias for np.void + assert value["real"] == 3 + assert value["imag"] == -4 + + +def test___complexint32_arrays___add___raises_type_error() -> None: + left = np.array([(1, 2), (3, -4)], ComplexInt32DType) + right = np.array([(-5, 6), (-7, -8)], ComplexInt32DType) + + with pytest.raises(TypeError): + _ = left + right # type: ignore[operator] + + +def test___complexint32_array_and_int16_array___add___raises_type_error() -> None: + left = np.array([(1, 2), (3, -4)], ComplexInt32DType) + right = np.array([5, -6], np.int16) + + with pytest.raises(TypeError): + _ = left + right # type: ignore[operator]