ENH: Function Reverse Arithmetic Priority

rocketpy/mathutils/function.py

@@ -24,6 +24,9 @@
     extrapolation, plotting and algebra.
     """
 
+    # Arithmetic priority
+    __array_ufunc__ = None
+
     def __init__(
         self,
         source,
@@ -1837,7 +1840,7 @@
                 return Function(lambda x: (self.get_value(x) + other(x)))
         # If other is Float except...
         except AttributeError:
-            if isinstance(other, (float, int, complex)):
+            if isinstance(other, (float, int, complex, np.ndarray)):
                 # Check if Function object source is array or callable
                 if isinstance(self.source, np.ndarray):
                     # Operate on grid values
@@ -1967,7 +1970,7 @@
                 return Function(lambda x: (self.get_value(x) * other(x)))
         # If other is Float except...
         except AttributeError:
-            if isinstance(other, (float, int, complex)):
+            if isinstance(other, (float, int, complex, np.ndarray)):
                 # Check if Function object source is array or callable
                 if isinstance(self.source, np.ndarray):
                     # Operate on grid values
@@ -2056,7 +2059,7 @@
                 return Function(lambda x: (self.get_value_opt(x) / other(x)))
         # If other is Float except...
         except AttributeError:
-            if isinstance(other, (float, int, complex)):
+            if isinstance(other, (float, int, complex, np.ndarray)):
                 # Check if Function object source is array or callable
                 if isinstance(self.source, np.ndarray):
                     # Operate on grid values
@@ -2095,7 +2098,7 @@
             A Function object which gives the result of other(x)/self(x).
         """
         # Check if Function object source is array and other is float
-        if isinstance(other, (float, int, complex)):
+        if isinstance(other, (float, int, complex, np.ndarray)):
             if isinstance(self.source, np.ndarray):
                 # Operate on grid values
                 ys = other / self.y_array
@@ -2163,7 +2166,7 @@
                 return Function(lambda x: (self.get_value_opt(x) ** other(x)))
         # If other is Float except...
         except AttributeError:
-            if isinstance(other, (float, int, complex)):
+            if isinstance(other, (float, int, complex, np.ndarray)):
                 # Check if Function object source is array or callable
                 if isinstance(self.source, np.ndarray):
                     # Operate on grid values
@@ -2202,7 +2205,7 @@
             A Function object which gives the result of other(x)**self(x).
         """
         # Check if Function object source is array and other is float
-        if isinstance(other, (float, int, complex)):
+        if isinstance(other, (float, int, complex, np.ndarray)):
             if isinstance(self.source, np.ndarray):
                 # Operate on grid values
                 ys = other**self.y_array

tests/test_function.py

@@ -390,3 +390,19 @@ def test_shepard_interpolation(x, y, z_expected):
     func = Function(source=source, inputs=["x", "y"], outputs=["z"])
     z = func(x, y)
     assert np.isclose(z, z_expected, atol=1e-8).all()
+
+
+def test_arithmetic_priority():
+    """Test the arithmetic priority of the Function class, specially
+    comparing to the numpy array operations.
+    """
+    func_lambda = Function(lambda x: x**2)
+    func_array = Function([(0, 0), (1, 1), (2, 4)])
+    array = np.array([1])
+
+    assert isinstance(func_lambda + func_array, Function)
+    assert isinstance(func_array + func_lambda, Function)
+    assert isinstance(func_lambda + array, Function)
+    assert isinstance(array + func_lambda, Function)
+    assert isinstance(func_array + array, Function)
+    assert isinstance(array + func_array, Function)