Merge pull request #519 from RocketPy-Team/enh/optional-discretize-mutation

ENH: Argument for Optional Mutation on Function Discretize

Author: Gui-FernandesBR

CHANGELOG.md

@@ -32,7 +32,7 @@ straightforward as possible.
 
 ### Added
 
--
+- ENH: Argument for Optional Mutation on Function Discretize [#519](https://github.com/RocketPy-Team/RocketPy/pull/519)
 
 ### Changed
 

rocketpy/mathutils/function.py

@@ -10,6 +10,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from copy import deepcopy
 from scipy import integrate, linalg, optimize
 
 try:
@@ -515,12 +516,16 @@ def set_discrete(
         interpolation="spline",
         extrapolation="constant",
         one_by_one=True,
+        mutate_self=True,
     ):
-        """This method transforms function defined Functions into list
-        defined Functions. It evaluates the function at certain points
-        (sampling range) and stores the results in a list, which is converted
-        into a Function and then returned. The original Function object is
-        replaced by the new one.
+        """This method discretizes a 1-D or 2-D Function by evaluating it at
+        certain points (sampling range) and storing the results in a list,
+        which is converted into a Function and then returned. By default, the
+        original Function object is replaced by the new one, which can be
+        changed by the attribute `mutate_self`.
+
+        This method is specially useful to change a dataset sampling or to
+        convert a Function defined by a callable into a list based Function.
 
         Parameters
         ----------
@@ -544,18 +549,32 @@ def set_discrete(
         one_by_one : boolean, optional
             If True, evaluate Function in each sample point separately. If
             False, evaluates Function in vectorized form. Default is True.
+        mutate_self : boolean, optional
+            If True, the original Function object source will be replaced by
+            the new one. If False, the original Function object source will
+            remain unchanged, and the new one is simply returned.
+            Default is True.
 
         Returns
         -------
         self : Function
+
+        Notes
+        -----
+        1. This method performs by default in place replacement of the original
+        Function object source. This can be changed by the attribute `mutate_self`.
+
+        2. Currently, this method only supports 1-D and 2-D Functions.
         """
-        if self.__dom_dim__ == 1:
+        func = deepcopy(self) if not mutate_self else self
+
+        if func.__dom_dim__ == 1:
             xs = np.linspace(lower, upper, samples)
-            ys = self.get_value(xs.tolist()) if one_by_one else self.get_value(xs)
-            self.set_source(np.concatenate(([xs], [ys])).transpose())
-            self.set_interpolation(interpolation)
-            self.set_extrapolation(extrapolation)
-        elif self.__dom_dim__ == 2:
+            ys = func.get_value(xs.tolist()) if one_by_one else func.get_value(xs)
+            func.set_source(np.concatenate(([xs], [ys])).transpose())
+            func.set_interpolation(interpolation)
+            func.set_extrapolation(extrapolation)
+        elif func.__dom_dim__ == 2:
             lower = 2 * [lower] if isinstance(lower, (int, float)) else lower
             upper = 2 * [upper] if isinstance(upper, (int, float)) else upper
             sam = 2 * [samples] if isinstance(samples, (int, float)) else samples
@@ -564,22 +583,29 @@ def set_discrete(
             ys = np.linspace(lower[1], upper[1], sam[1])
             xs, ys = np.array(np.meshgrid(xs, ys)).reshape(2, xs.size * ys.size)
             # Evaluate function at all mesh nodes and convert it to matrix
-            zs = np.array(self.get_value(xs, ys))
-            self.__interpolation__ = "shepard"
-            self.__extrapolation__ = "natural"
-            self.set_source(np.concatenate(([xs], [ys], [zs])).transpose())
-        return self
+            zs = np.array(func.get_value(xs, ys))
+            func.set_source(np.concatenate(([xs], [ys], [zs])).transpose())
+            func.__interpolation__ = "shepard"
+            func.__extrapolation__ = "natural"
+        else:
+            raise ValueError(
+                "Discretization is only supported for 1-D and 2-D Functions."
+            )
+        return func
 
     def set_discrete_based_on_model(
-        self, model_function, one_by_one=True, keep_self=True
+        self, model_function, one_by_one=True, keep_self=True, mutate_self=True
     ):
-        """This method transforms the domain of Function instance into a list of
-        discrete points based on the domain of a model Function instance. It
-        does so by retrieving the domain, domain name, interpolation method and
-        extrapolation method of the model Function instance. It then evaluates
-        the original Function instance in all points of the retrieved domain to
-        generate the list of discrete points that will be used for interpolation
-        when this Function is called.
+        """This method transforms the domain of a 1-D or 2-D Function instance
+        into a list of discrete points based on the domain of a model Function
+        instance. It does so by retrieving the domain, domain name,
+        interpolation method and extrapolation method of the model Function
+        instance. It then evaluates the original Function instance in all
+        points of the retrieved domain to generate the list of discrete points
+        that will be used for interpolation when this Function is called.
+
+        By default, the original Function object is replaced by the new one,
+        which can be changed by the attribute `mutate_self`.
 
         Parameters
         ----------
@@ -589,15 +615,17 @@ def set_discrete_based_on_model(
             Must be a Function whose source attribute is a list (i.e. a list
             based Function instance). Must have the same domain dimension as the
             Function to be discretized.
-
         one_by_one : boolean, optional
             If True, evaluate Function in each sample point separately. If
             False, evaluates Function in vectorized form. Default is True.
-
-        keepSelf : boolean, optional
+        keep_self : boolean, optional
             If True, the original Function interpolation and extrapolation
             methods will be kept. If False, those are substituted by the ones
             from the model Function. Default is True.
+        mutate_self : boolean, optional
+            If True, the original Function object source will be replaced by
+            the new one. If False, the original Function object source will
+            remain unchanged, and the new one is simply returned.
 
         Returns
         -------
@@ -645,40 +673,48 @@ def set_discrete_based_on_model(
 
         Notes
         -----
-        1. This method performs in place replacement of the original Function
-        object source.
+        1. This method performs by default in place replacement of the original
+        Function object source. This can be changed by the attribute `mutate_self`.
 
         2. This method is similar to set_discrete, but it uses the domain of a
         model Function to define the domain of the new Function instance.
+
+        3. Currently, this method only supports 1-D and 2-D Functions.
         """
         if not isinstance(model_function.source, np.ndarray):
             raise TypeError("model_function must be a list based Function.")
         if model_function.__dom_dim__ != self.__dom_dim__:
             raise ValueError("model_function must have the same domain dimension.")
 
-        if self.__dom_dim__ == 1:
+        func = deepcopy(self) if not mutate_self else self
+
+        if func.__dom_dim__ == 1:
             xs = model_function.source[:, 0]
-            ys = self.get_value(xs.tolist()) if one_by_one else self.get_value(xs)
-            self.set_source(np.concatenate(([xs], [ys])).transpose())
-        elif self.__dom_dim__ == 2:
+            ys = func.get_value(xs.tolist()) if one_by_one else func.get_value(xs)
+            func.set_source(np.concatenate(([xs], [ys])).transpose())
+        elif func.__dom_dim__ == 2:
             # Create nodes to evaluate function
             xs = model_function.source[:, 0]
             ys = model_function.source[:, 1]
             # Evaluate function at all mesh nodes and convert it to matrix
-            zs = np.array(self.get_value(xs, ys))
-            self.set_source(np.concatenate(([xs], [ys], [zs])).transpose())
+            zs = np.array(func.get_value(xs, ys))
+            func.set_source(np.concatenate(([xs], [ys], [zs])).transpose())
+        else:
+            raise ValueError(
+                "Discretization is only supported for 1-D and 2-D Functions."
+            )
 
         interp = (
-            self.__interpolation__ if keep_self else model_function.__interpolation__
+            func.__interpolation__ if keep_self else model_function.__interpolation__
         )
         extrap = (
-            self.__extrapolation__ if keep_self else model_function.__extrapolation__
+            func.__extrapolation__ if keep_self else model_function.__extrapolation__
         )
 
-        self.set_interpolation(interp)
-        self.set_extrapolation(extrap)
+        func.set_interpolation(interp)
+        func.set_extrapolation(extrap)
 
-        return self
+        return func
 
     def reset(
         self,

tests/unit/test_function.py

@@ -158,3 +158,55 @@ def test_get_value_opt(x, y, z):
     func = Function(source, interpolation="shepard", extrapolation="natural")
     assert isinstance(func.get_value_opt(x, y), float)
     assert np.isclose(func.get_value_opt(x, y), z, atol=1e-6)
+
+
+@pytest.mark.parametrize("samples", [2, 50, 1000])
+def test_set_discrete_mutator(samples):
+    """Tests the set_discrete method of the Function class."""
+    func = Function(lambda x: x**3)
+    discretized_func = func.set_discrete(-10, 10, samples, mutate_self=True)
+
+    assert isinstance(discretized_func, Function)
+    assert isinstance(func, Function)
+    assert discretized_func.source.shape == (samples, 2)
+    assert func.source.shape == (samples, 2)
+
+
+@pytest.mark.parametrize("samples", [2, 50, 1000])
+def test_set_discrete_non_mutator(samples):
+    """Tests the set_discrete method of the Function class.
+    The mutator argument is set to False.
+    """
+    func = Function(lambda x: x**3)
+    discretized_func = func.set_discrete(-10, 10, samples, mutate_self=False)
+
+    assert isinstance(discretized_func, Function)
+    assert isinstance(func, Function)
+    assert discretized_func.source.shape == (samples, 2)
+    assert callable(func.source)
+
+
+def test_set_discrete_based_on_model_mutator(linear_func):
+    """Tests the set_discrete_based_on_model method of the Function class.
+    The mutator argument is set to True.
+    """
+    func = Function(lambda x: x**3)
+    discretized_func = func.set_discrete_based_on_model(linear_func, mutate_self=True)
+
+    assert isinstance(discretized_func, Function)
+    assert isinstance(func, Function)
+    assert discretized_func.source.shape == (4, 2)
+    assert func.source.shape == (4, 2)
+
+
+def test_set_discrete_based_on_model_non_mutator(linear_func):
+    """Tests the set_discrete_based_on_model method of the Function class.
+    The mutator argument is set to False.
+    """
+    func = Function(lambda x: x**3)
+    discretized_func = func.set_discrete_based_on_model(linear_func, mutate_self=False)
+
+    assert isinstance(discretized_func, Function)
+    assert isinstance(func, Function)
+    assert discretized_func.source.shape == (4, 2)
+    assert callable(func.source)