Added the start of support for N-d arrays of multivectors (#343)

Previously only 1d arrays of multivectors were supported.

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

clifford/_mvarray.py

@@ -1,57 +1,96 @@
-import numpy as np
+import functools
+import operator
+from typing import Tuple
 
+import numpy as np
 
 from clifford.io import write_ga_file, read_ga_file  # noqa: F401
+from ._layout import Layout
 from ._multivector import MultiVector
 
 dual_array = np.vectorize(MultiVector.dual)
 normal_array = np.vectorize(MultiVector.normal)
 call_array = np.vectorize(MultiVector.__call__)
 
 
+def _interrogate_nested_mvs(input_array) -> Tuple[Tuple[int, ...], Layout, np.dtype]:
+    """
+    Calculates the shape of the nested input_array, and gets the associated layout.
+    Stops descending when it encounters a MultiVector.
+    """
+    if not isinstance(input_array, MultiVector):
+        nested_shape, layout, dtype = _interrogate_nested_mvs(input_array[0])
+        return (len(input_array), *nested_shape), layout, dtype
+    else:
+        return (), input_array.layout, input_array.value.dtype
+
+
+def _index_nested_iterable(input_iterable, index):
+    """
+    Given a nested iterable, input_iterable, return the element given by the
+    1d index iterable
+    """
+    return functools.reduce(operator.getitem, index, input_iterable)
+
+
+@functools.lru_cache(None)
+def _get_vectorized_value_func(dtype):
+    return np.vectorize(lambda x: x.value, otypes=[dtype], signature='()->(n)')
+
+
 class MVArray(np.ndarray):
-    '''
+    """
     MultiVector Array
-    '''
-
+    """
     def __new__(cls, input_array):
-        obj = np.empty(len(input_array), dtype=object)
-        obj[:] = input_array
-        obj = obj.view(cls)
-        return obj
+
+        input_shape, layout, dtype = _interrogate_nested_mvs(input_array)
+        # copying this across elementwise is necessary to prevent numpy recursing into the multivector coefficients
+        obj = np.empty(input_shape, dtype=object)
+        for index in np.ndindex(input_shape):
+            obj[index] = _index_nested_iterable(input_array, index)
+
+        self = obj.view(cls)
+        return self
 
     def __array_finalize__(self, obj):
         if obj is None:
             return
 
+    def _get_first_element(self):
+        return self[(0,) * self.ndim]
+
     @property
     def value(self):
         """
         Return an np array of the values of multivectors
         """
-        return np.array([mv.value for mv in self])
+        value_dtype = self._get_first_element().value.dtype
+        return _get_vectorized_value_func(value_dtype)(self)
 
     @staticmethod
     def from_value_array(layout, value_array):
         """
         Constructs an array of mvs from a value array
         """
-        v_new_mv = np.vectorize(lambda v: MultiVector(layout, v), otypes=[MVArray], signature='(n)->()')
+        v_new_mv = np.vectorize(lambda v: layout.MultiVector(v), otypes=[object], signature='(n)->()')
         return MVArray(v_new_mv(value_array))
 
     def save(self, filename, compression=True, transpose=False,
              sparse=False, support=False, compression_opts=1):
         """
         Saves the array to a ga file
         """
-        write_ga_file(filename, self.value, self[0].layout.metric, self[0].layout.basis_names,
+        first_element = self._get_first_element()
+        write_ga_file(filename, self.value, first_element.layout.metric,
+                      first_element.layout.basis_names,
                       compression=compression, transpose=transpose,
                       sparse=sparse, support=support, compression_opts=compression_opts)
 
     def sum(self):
-        '''
+        """
         sum elements of this MVArray
-        '''
+        """
         out = self[0]
         for k in self[1:]:
             out += k

clifford/test/test_clifford.py

@@ -196,6 +196,59 @@ def test_add_float64(self, g3):
 
         assert 1 + e1 == e1 + np.float64(1)
 
+    def _random_value_array(self, layout, Nrows, Ncolumns):
+        value_array = np.zeros((Nrows, Ncolumns, layout.gaDims))
+        for i in range(Nrows):
+            for j in range(Ncolumns):
+                value_array[i, j, :] = layout.randomMV().value
+        return value_array
+
+    def test_2d_mv_array(self, g3):
+        layout, blades = g3, g3.blades
+        Nrows = 2
+        Ncolumns = 3
+        value_array_a = self._random_value_array(g3, Nrows, Ncolumns)
+        value_array_b = self._random_value_array(g3, Nrows, Ncolumns)
+
+        mv_array_a = MVArray.from_value_array(layout, value_array_a)
+        assert mv_array_a.shape == (Nrows, Ncolumns)
+        mv_array_b = MVArray.from_value_array(layout, value_array_b)
+        assert mv_array_b.shape == (Nrows, Ncolumns)
+
+        # check properties of the array are preserved (no need to check both a and b)
+        np.testing.assert_array_equal(mv_array_a.value, value_array_a)
+        assert mv_array_a.value.dtype == value_array_a.dtype
+        assert type(mv_array_a.value) == type(value_array_a)
+
+        # Check addition
+        mv_array_sum = mv_array_a + mv_array_b
+        array_sum = value_array_a + value_array_b
+        np.testing.assert_array_equal(mv_array_sum.value, array_sum)
+
+        # Check elementwise gp
+        mv_array_gp = mv_array_a * mv_array_b
+        value_array_gp = np.zeros((Nrows, Ncolumns, layout.gaDims))
+        for i in range(Nrows):
+            for j in range(Ncolumns):
+                value_array_gp[i, j, :] = layout.gmt_func(value_array_a[i, j, :], value_array_b[i, j, :])
+        np.testing.assert_array_equal(mv_array_gp.value, value_array_gp)
+
+        # Check elementwise op
+        mv_array_op = mv_array_a ^ mv_array_b
+        value_array_op = np.zeros((Nrows, Ncolumns, layout.gaDims))
+        for i in range(Nrows):
+            for j in range(Ncolumns):
+                value_array_op[i, j, :] = layout.omt_func(value_array_a[i, j, :], value_array_b[i, j, :])
+        np.testing.assert_array_equal(mv_array_op.value, value_array_op)
+
+        # Check elementwise ip
+        mv_array_ip = mv_array_a | mv_array_b
+        value_array_ip = np.zeros((Nrows, Ncolumns, layout.gaDims))
+        for i in range(Nrows):
+            for j in range(Ncolumns):
+                value_array_ip[i, j, :] = layout.imt_func(value_array_a[i, j, :], value_array_b[i, j, :])
+        np.testing.assert_array_equal(mv_array_ip.value, value_array_ip)
+
     def test_array_control(self, g3):
         '''
         test methods to take control addition from numpy arrays
@@ -242,11 +295,11 @@ def test_array_overload(self, algebra):
 
         normed_array = test_array.normal()
         other_array = np.array([t.normal().value for t in test_array])
-        np.testing.assert_almost_equal(normed_array.value, other_array)
+        np.testing.assert_array_equal(normed_array.value, other_array)
 
         dual_array = test_array.dual()
         other_array_2 = np.array([t.dual().value for t in test_array])
-        np.testing.assert_almost_equal(dual_array.value, other_array_2)
+        np.testing.assert_array_equal(dual_array.value, other_array_2)
 
     def test_comparison_operators(self, g3):
         layout, blades = g3, g3.blades