ENH: simplifying and optimizing the function, implementing tests.

Author: Lucas-Prates

rocketpy/mathutils/piecewise_function.py

@@ -40,62 +40,25 @@ def __new__(
             The number of points in which the piecewise function will be
             evaluated to create a base function. The default value is 100.
         """
+        cls.__validate__source(source)
         if inputs is None:
             inputs = ["Scalar"]
         if outputs is None:
             outputs = ["Scalar"]
-        # Check if source is a dictionary
-        if not isinstance(source, dict):
-            raise TypeError("source must be a dictionary")
-        # Check if all keys are tuples
-        for key in source.keys():
-            if not isinstance(key, tuple):
-                raise TypeError("keys of source must be tuples")
-        # Check if all domains are disjoint
-        for key1 in source.keys():
-            for key2 in source.keys():
-                if key1 != key2:
-                    if key1[0] < key2[1] and key1[1] > key2[0]:
-                        raise ValueError("domains must be disjoint")
-
-        # Crate Function
-        def calc_output(func, inputs):
-            """Receives a list of inputs value and a function, populates another
-            list with the results corresponding to the same results.
-
-            Parameters
-            ----------
-            func : Function
-                The Function object to be
-            inputs : list, tuple, np.array
-                The array of points to applied the func to.
-
-            Examples
-            --------
-            >>> inputs = [0, 1, 2, 3, 4, 5]
-            >>> def func(x):
-            ...     return x*10
-            >>> calc_output(func, inputs)
-            [0, 10, 20, 30, 40, 50]
-
-            Notes
-            -----
-            In the future, consider using the built-in map function from python.
-            """
-            output = np.zeros(len(inputs))
-            for j, value in enumerate(inputs):
-                output[j] = func.get_value_opt(value)
-            return output
 
         input_data = []
         output_data = []
-        for key in sorted(source.keys()):
-            i = np.linspace(key[0], key[1], datapoints)
-            i = i[~np.isin(i, input_data)]
-            input_data = np.concatenate((input_data, i))
+        for interval in sorted(source.keys()):
+            grid = np.linspace(interval[0], interval[1], datapoints)
+
+            # since intervals are disjoint and sorted, we only need to check
+            # if the first point is already included
+            if interval[0] in input_data:
+                grid = np.delete(grid, 0)
+            input_data = np.concatenate((input_data, grid))
 
-            f = Function(source[key])
-            output_data = np.concatenate((output_data, calc_output(f, i)))
+            f = Function(source[interval])
+            output_data = np.concatenate((output_data, f(grid)))
 
         return Function(
             np.concatenate(([input_data], [output_data])).T,
@@ -104,3 +67,27 @@ def calc_output(func, inputs):
             interpolation=interpolation,
             extrapolation=extrapolation,
         )
+
+    @staticmethod
+    def __validate__source(source):
+        """Validates that source is dictionary with non-overlapping
+        intervals
+
+        Parameters
+        ----------
+        source : dict
+            A dictionary of Function objects, where the keys are the domains.
+        """
+        # Check if source is a dictionary
+        if not isinstance(source, dict):
+            raise TypeError("source must be a dictionary")
+        # Check if all keys are tuples
+        for key in source.keys():
+            if not isinstance(key, tuple):
+                raise TypeError("keys of source must be tuples")
+        # Check if all domains are disjoint
+        for interval1 in source.keys():
+            for interval2 in source.keys():
+                if interval1 != interval2:
+                    if interval1[0] < interval2[1] and interval1[1] > interval2[0]:
+                        raise ValueError("domains must be disjoint")

tests/unit/test_piecewise_function.py

@@ -0,0 +1,35 @@
+import pytest
+
+from rocketpy import PiecewiseFunction
+
+
+@pytest.mark.parametrize(
+    "source",
+    [
+        ((0, 4), lambda x: x),
+        {"0-4": lambda x: x},
+        {(0, 4): lambda x: x, (3, 5): lambda x: 2 * x},
+    ],
+)
+def test_invalid_source(source):
+    """Test an error is raised when the source parameter is invalid"""
+    with pytest.raises((TypeError, ValueError)):
+        PiecewiseFunction(source)
+
+
+@pytest.mark.parametrize(
+    "source",
+    [
+        {(-1, 0): lambda x: -x, (0, 1): lambda x: x},
+        {
+            (0, 1): lambda x: x,
+            (1, 2): lambda x: 1,
+            (2, 3): lambda x: 3 - x,
+        },
+    ],
+)
+@pytest.mark.parametrize("inputs", [None, "X"])
+@pytest.mark.parametrize("outputs", [None, "Y"])
+def test_new(source, inputs, outputs):
+    """Test if PiecewiseFunction.__new__ runs correctly"""
+    PiecewiseFunction(source, inputs, outputs)