Numericarray atom (#1512)

Initial feature for #1511: implement NumericArray atom.

This is sufficient to support the Plot3D/Graphics3D use case as those
will only require NumericArrays generated and consumed in Python.

---------

Co-authored-by: Juan Mauricio Matera <[email protected]>
Co-authored-by: R. Bernstein <[email protected]>

Author: 3 people

mathics/core/atoms.py

@@ -7,13 +7,15 @@
 from typing import Any, Dict, Generic, Optional, Tuple, TypeVar, Union
 
 import mpmath
+import numpy
 import sympy
 from sympy.core import numbers as sympy_numbers
 
 from mathics.core.element import BoxElementMixin, ImmutableValueMixin
 from mathics.core.keycomparable import (
     BASIC_ATOM_BYTEARRAY_ELT_ORDER,
     BASIC_ATOM_NUMBER_ELT_ORDER,
+    BASIC_ATOM_NUMERICARRAY_ELT_ORDER,
     BASIC_ATOM_STRING_ELT_ORDER,
 )
 from mathics.core.number import (
@@ -1098,6 +1100,117 @@ def is_zero(self) -> bool:
 }
 
 
+#
+# NumericArray
+#
+
+NUMERIC_ARRAY_TYPE_MAP = {
+    "UnsignedInteger8": numpy.dtype("uint8"),
+    "UnsignedInteger16": numpy.dtype("uint16"),
+    "UnsignedInteger32": numpy.dtype("uint32"),
+    "UnsignedInteger64": numpy.dtype("uint64"),
+    "Integer8": numpy.dtype("int8"),
+    "Integer16": numpy.dtype("int16"),
+    "Integer32": numpy.dtype("int32"),
+    "Integer64": numpy.dtype("int64"),
+    "Real32": numpy.dtype("float32"),
+    "Real64": numpy.dtype("float64"),
+    "ComplexReal32": numpy.dtype("complex64"),
+    "ComplexReal64": numpy.dtype("complex128"),
+}
+
+NUMERIC_ARRAY_DTYPE_TO_NAME = {
+    dtype: name for name, dtype in NUMERIC_ARRAY_TYPE_MAP.items()
+}
+
+
+class NumericArray(Atom, ImmutableValueMixin):
+    """
+    NumericArray provides compact storage and efficient access for machine-precision numeric arrays,
+    backed by NumPy arrays.
+    """
+
+    class_head_name = "NumericArray"
+
+    def __init__(self, value, dtype=None):
+        # compute value
+        if not isinstance(value, numpy.ndarray):
+            value = numpy.asarray(value, dtype=dtype)
+        elif dtype is not None:
+            value = value.astype(dtype)
+        self.value = value
+
+        # check type
+        self._type_name = NUMERIC_ARRAY_DTYPE_TO_NAME.get(self.value.dtype, None)
+        if not self._type_name:
+            allowed = ", ".join(str(dtype) for dtype in NUMERIC_ARRAY_TYPE_MAP.values())
+            message = f"Argument 'value' must be one of {allowed}; is {str(self.value.dtype)}."
+            raise ValueError(message)
+
+        # summary and hash
+        shape_string = "×".join(str(dim) for dim in self.value.shape) or "0"
+        self._summary_string = f"{self._type_name}, {shape_string}"
+        self._hash = None
+
+    def __hash__(self):
+        if not self._hash:
+            self._hash = hash(("NumericArray", self.value.shape, id(self.value)))
+        return self._hash
+
+    def __str__(self) -> str:
+        return f"NumericArray[{self._summary_string}]"
+
+    def atom_to_boxes(self, f, evaluation):
+        return String(f"<{self._summary_string}>")
+
+    def do_copy(self) -> "NumericArray":
+        return NumericArray(self.value.copy())
+
+    def default_format(self, evaluation, form) -> str:
+        return f"NumericArray[<{self._summary_string}>]"
+
+    @property
+    def items(self) -> tuple:
+        from mathics.core.convert.python import from_python
+
+        if len(self.value.shape) == 1:
+            return tuple(from_python(item.item()) for item in self.value)
+        else:
+            return tuple(NumericArray(array) for array in self.value)
+
+    @property
+    def element_order(self) -> tuple:
+        return (
+            BASIC_ATOM_NUMERICARRAY_ELT_ORDER,
+            self.value.shape,
+            self.value.dtype,
+            id(self.value),
+        )
+
+    @property
+    def pattern_precedence(self) -> tuple:
+        return super().pattern_precedence
+
+    def sameQ(self, rhs) -> bool:
+        return isinstance(rhs, NumericArray) and numpy.array_equal(
+            self.value, rhs.value
+        )
+
+    def to_sympy(self, **kwargs) -> None:
+        return None
+
+    # TODO: this returns a list instead of np.ndarray in keeping with
+    # idea that to_python should return only "native" Python types.
+    # Keep an eye on this because there is a slight risk that code may
+    # naively call to_python and cause a performance issue due to
+    # the cost of converting to a nested list structure for a large array.
+    def to_python(self, *args, **kwargs) -> list:
+        return self.value.tolist()
+
+    def user_hash(self, update) -> None:
+        update(self.value.tobytes())
+
+
 class String(Atom, BoxElementMixin):
     value: str
     class_head_name = "System`String"

mathics/core/convert/python.py

@@ -5,7 +5,9 @@
 
 from typing import Any
 
-from mathics.core.atoms import Complex, Integer, Rational, Real, String
+import numpy
+
+from mathics.core.atoms import Complex, Integer, NumericArray, Rational, Real, String
 from mathics.core.number import get_type
 from mathics.core.symbols import (
     BaseElement,
@@ -113,5 +115,7 @@ def from_python(arg: Any) -> BaseElement:
         from mathics.builtin.binary.bytearray import ByteArray
 
         return Expression(SymbolByteArray, ByteArray(arg))
+    elif isinstance(arg, numpy.ndarray):
+        return NumericArray(arg)
     else:
         raise NotImplementedError

mathics/core/keycomparable.py

@@ -279,6 +279,7 @@ def __ne__(self, other) -> bool:
 BASIC_ATOM_STRING_ELT_ORDER = 0x01
 BASIC_ATOM_BYTEARRAY_ELT_ORDER = 0x02
 LITERAL_EXPRESSION_ELT_ORDER = 0x03
+BASIC_ATOM_NUMERICARRAY_ELT_ORDER = 0x04
 
 BASIC_NUMERIC_EXPRESSION_ELT_ORDER = 0x12
 GENERAL_NUMERIC_EXPRESSION_ELT_ORDER = 0x13

test/builtin/test_binary.py

@@ -421,3 +421,38 @@ def test_ByteOrdering(str_expr, str_expected, fail_msg):
         hold_expected=True,
         failure_message=fail_msg,
     )
+
+
+@pytest.mark.skip(reason="NumericArray[] builtin not written yet.")
+@pytest.mark.parametrize(
+    ("str_expr", "str_expected"),
+    [
+        ("NumericArray[{{1,2},{3,4}}]", "<Integer64, 2×2>"),
+        ("ToString[NumericArray[{{1,2},{3,4}}]]", "<Integer64, 2×2>"),
+        ("Head[NumericArray[{1,2}]]", "NumericArray"),
+        ("AtomQ[NumericArray[{1,2}]]", "True"),
+        ("First[NumericArray[{1,2,3}]]", "1"),
+        ("First[NumericArray[{{1,2}, {3,4}}]]", "<Integer64, 2>"),
+        ("Last[NumericArray[{1,2,3}]]", "3"),
+        ("Last[NumericArray[{{1,2}, {3,4}}]]", "<Integer64, 2>"),
+        ("Normal[NumericArray[{{1,2}, {3,4}}]]", "{{1, 2}, {3, 4}}"),
+    ],
+)
+def test_basics(str_expr, str_expected):
+    check_evaluation(str_expr, str_expected, hold_expected=True)
+
+
+@pytest.mark.skip(reason="NumericArray[] builtin not written yet.")
+def test_type_conversion():
+    # Move below imports to the top when we've implementated NumericArray[]
+    from test.helper import evaluate
+
+    import numpy as np
+
+    from mathics.core.atoms import NumericArray
+
+    expr = evaluate("NumericArray[{1,2}]")
+    assert isinstance(expr, NumericArray)
+    assert expr.value.dtype == np.int64
+    expr = evaluate('NumericArray[{1,2}, "ComplexReal32"]')
+    assert expr.value.dtype == np.complex64

test/core/test_atoms.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 
-import sys
+import numpy as np
 
 import mathics.core.atoms as atoms
 import mathics.core.systemsymbols as system_symbols
@@ -12,11 +12,13 @@
     Integer2,
     MachineReal,
     MachineReal0,
+    NumericArray,
     Rational,
     RationalOneHalf,
     Real,
     String,
 )
+from mathics.core.convert.python import from_python
 from mathics.core.definitions import Definitions
 from mathics.core.evaluation import Evaluation
 from mathics.core.expression import Expression
@@ -238,3 +240,29 @@ def test_mixed_object_canonicalization():
         Complex(Rational(1, 0), Integer(0)), # 3
     )
     # fmt: on
+
+
+#
+# NumericArray tests
+#
+
+
+def test_numericarray_atom_preserves_array_reference():
+    array = np.array([1, 2, 3], dtype=np.int64)
+    atom = NumericArray(array)
+    assert atom.value is array, "NumericArray.value should be a NumPy array"
+
+
+def test_numericarray_atom_preserves_equality():
+    array = np.array([1, 2, 3], dtype=np.int64)
+    atom = NumericArray(array, dtype=np.float64)
+    np.testing.assert_array_equal(atom.value, array)
+
+
+def test_numericarray_expression_from_python_array():
+    array = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+    atom = from_python(array)
+    assert isinstance(
+        atom, NumericArray
+    ), "from_python() conversion of a NumPy Array should yield a NumericArray"
+    assert atom.value is array