numba: Add support for selecting multiple grades (#348)

clifford/numba/_multivector.py

@@ -285,27 +285,40 @@ def impl(a):
 
 
 @overload_call(MultiVectorType)
-def ga_call(self, arg):
-    # grade projection
-    grades = self.layout_type.obj._basis_blade_order.grades
-    if isinstance(arg, types.IntegerLiteral):
-        # Optimized case where the mask can be computed at compile-time.
-        # using `nonzero` makes the resulting array smaller.
-        inds, = (grades == arg.literal_value).nonzero()
-        def impl(self, arg):
-            mv = self.layout.MultiVector(np.zeros_like(self.value))
-            mv.value[inds] = self.value[inds]
-            return mv
-        return impl
-    elif isinstance(arg, types.Integer):
-        # runtime grade selection - should be less common
-        def impl(self, arg):
-            # probably faster to not call nonzero here
-            inds = grades == arg
-            mv = self.layout.MultiVector(np.zeros_like(self.value))
-            mv.value[inds] = self.value[inds]
-            return mv
-        return impl
+def ga_call(self, *args):
+    # a numba quirk means that varargs can end up passed in two different ways
+    if len(args) == 1 and isinstance(args[0], (types.StarArgTuple, types.StarArgUniTuple)):
+        args = args[0].types
+
+    # grade selection
+    if len(args) > 0:
+        # grade projection
+        grades = self.layout_type.obj._basis_blade_order.grades
+        if all(isinstance(arg, types.IntegerLiteral) for arg in args):
+            # Optimized case where the mask can be computed at compile-time.
+            inds = (grades == args[0].literal_value)
+            for arg in args[1:]:
+                inds |= (grades == arg.literal_value)
+            # using `nonzero` makes the resulting array smaller.
+            inds = inds.nonzero()
+            def impl(self, *args):
+                mv = self.layout.MultiVector(np.zeros_like(self.value))
+                mv.value[inds] = self.value[inds]
+                return mv
+            return impl
+        elif all(isinstance(arg, types.Integer) for arg in args):
+            # runtime grade selection - should be less common. This includes
+            # the case where only some grades are known at compile-time.
+            def impl(self, *args):
+                # probably faster to not call nonzero here
+                inds = (grades == args[0])
+                # can't use `for arg in args` here due to numba/numba#5372
+                for i in range(1, len(args)):
+                    inds |= (grades == args[i])
+                mv = self.layout.MultiVector(np.zeros_like(self.value))
+                mv.value[inds] = self.value[inds]
+                return mv
+            return impl
 
 
 @numba.extending.overload_method(MultiVectorType, 'mag2')

clifford/test/test_numba_extensions.py

@@ -127,7 +127,21 @@ def grade_func(a):
     def grade_func(a, grade):
         return a(grade)
 
-    @pytest.mark.parametrize('func', [literal_grade_func, grade_func])
+    # vararg versions
+    def literal_grade_func_star(a, *grades):
+        @numba.njit
+        def grade_func(a):
+            return a(*grades)
+        return grade_func(a)
+
+    @numba.njit
+    def grade_func_star(a, *grades):
+        return a(*grades)
+
+    @pytest.mark.parametrize('func', [
+        literal_grade_func, grade_func,
+        literal_grade_func_star, grade_func_star,
+    ])
     def test_grade_projection(self, func):
         a = 1 + e1 + (e1^e2) + (e1^e2^e3) + (e1^e2^e3^e4) + (e1^e2^e3^e4^e5)
 
@@ -137,3 +151,12 @@ def test_grade_projection(self, func):
         assert func(a, 3) == e1^e2^e3
         assert func(a, 4) == e1^e2^e3^e4
         assert func(a, 5) == e1^e2^e3^e4^e5
+
+    @pytest.mark.parametrize('func', [
+        literal_grade_func_star, grade_func_star,
+    ])
+    def test_multiple_grade_projection(self, func):
+        a = 1 + e1 + (e1^e2) + (e1^e2^e3) + (e1^e2^e3^e4) + (e1^e2^e3^e4^e5)
+
+        assert func(a, 0, 1) == 1 + e1
+        assert func(a, 0, 1, 2, 3, 4, 5) == a