TST: implementing tests for CustomSampler and making small corrections

Author: Lucas-Prates

.vscode/settings.json

@@ -48,6 +48,7 @@
         "Bigl",
         "Bigr",
         "bijective",
+        "Bivariate",
         "bmatrix",
         "boldsymbol",
         "boxplot",

rocketpy/stochastic/custom_sampler.py

@@ -10,7 +10,7 @@ class CustomSampler(ABC):
 
     @abstractmethod
     def sample(self, n_samples=1):
-        """Generates n samples from the custom distribution
+        """Generates samples from the custom distribution
 
         Parameters
         ----------
@@ -19,7 +19,7 @@ def sample(self, n_samples=1):
 
         Returns
         -------
-        sample_list : list
+        samples_list : list
             A list with n_samples elements, each of which is a valid sample
         """
 

rocketpy/stochastic/stochastic_model.py

@@ -89,7 +89,9 @@ def _set_stochastic(self, seed=None):
                 attr_value = None
                 if input_value is not None:
                     if "factor" in input_name:
-                        attr_value = self._validate_factors(input_name, input_value)
+                        attr_value = self._validate_factors(
+                            input_name, input_value, seed
+                        )
                     elif input_name not in self.exception_list:
                         if isinstance(input_value, tuple):
                             attr_value = self._validate_tuple(input_name, input_value)
@@ -104,6 +106,7 @@ def _set_stochastic(self, seed=None):
                         else:
                             raise AssertionError(
                                 f"'{input_name}' must be a tuple, list, int, or float"
+                                "or a custom sampler"
                             )
                 else:
                     attr_value = [getattr(self.obj, input_name)]
@@ -285,7 +288,7 @@ def _validate_scalar(self, input_name, input_value, getattr=getattr):  # pylint:
             get_distribution("normal", self.__random_number_generator),
         )
 
-    def _validate_factors(self, input_name, input_value):
+    def _validate_factors(self, input_name, input_value, seed):
         """
         Validate factor arguments.
 
@@ -313,8 +316,12 @@ def _validate_factors(self, input_name, input_value):
             return self._validate_tuple_factor(input_name, input_value)
         elif isinstance(input_value, list):
             return self._validate_list_factor(input_name, input_value)
+        elif isinstance(input_value, CustomSampler):
+            return self._validate_custom_sampler(input_name, input_value, seed)
         else:
-            raise AssertionError(f"`{input_name}`: must be either a tuple or list")
+            raise AssertionError(
+                f"`{input_name}`: must be either a tuple or listor a custom sampler"
+            )
 
     def _validate_tuple_factor(self, input_name, factor_tuple):
         """
@@ -463,7 +470,7 @@ def _validate_custom_sampler(self, input_name, sampler, seed=None):
             sampler.reset_seed(seed)
         except RuntimeError as e:
             raise RuntimeError(
-                f"An error occurred in the 'reset_seed' of {input_name} CustomSampler"
+                f"An error occurred in the 'reset_seed' method of {input_name} CustomSampler"
             ) from e
 
         return sampler
@@ -531,7 +538,7 @@ def dict_generator(self):
                     generated_dict[arg] = value.sample(n_samples=1)[0]
                 except RuntimeError as e:
                     raise RuntimeError(
-                        f"An error occurred in the 'sample' of {arg} CustomSampler"
+                        f"An error occurred in the 'sample' method of {arg} CustomSampler"
                     ) from e
         self.last_rnd_dict = generated_dict
         yield generated_dict

test_custom_sampler.ipynb

@@ -18,6 +18,7 @@
     "tests.fixtures.units.numerical_fixtures",
     "tests.fixtures.monte_carlo.monte_carlo_fixtures",
     "tests.fixtures.monte_carlo.stochastic_fixtures",
+    "tests.fixtures.monte_carlo.custom_sampler_fixtures",
     "tests.fixtures.monte_carlo.stochastic_motors_fixtures",
     "tests.fixtures.sensors.sensors_fixtures",
     "tests.fixtures.generic_surfaces.generic_surfaces_fixtures",

tests/conftest.py

@@ -0,0 +1,76 @@
+"""This file contains fixtures of CustomSampler used in stochastic classes."""
+
+import numpy as np
+import pytest
+
+from rocketpy import CustomSampler
+
+
+@pytest.fixture
+def elevation_sampler():
+    """Fixture to create mixture of two gaussian sampler"""
+    means_tuple = [1400, 1500]
+    sd_tuple = [40, 50]
+    prob_tuple = [0.4, 0.6]
+    return TwoGaussianMixture(means_tuple, sd_tuple, prob_tuple)
+
+
+class TwoGaussianMixture(CustomSampler):
+    """Class to sample from a mixture of two Gaussian distributions"""
+
+    def __init__(self, means_tuple, sd_tuple, prob_tuple, seed=None):
+        """Creates a sampler for a mixture of two Gaussian distributions
+
+        Parameters
+        ----------
+        means_tuple : 2-tuple
+            2-Tuple that contains the mean of each normal distribution of the
+            mixture
+        sd_tuple : 2-tuple
+            2-Tuple that contains the sd of each normal distribution of the
+            mixture
+        prob_tuple : 2-tuple
+            2-Tuple that contains the probability of each normal distribution of the
+            mixture. Its entries should be non-negative and sum up to 1.
+        """
+        np.random.default_rng(seed)
+        self.means_tuple = means_tuple
+        self.sd_tuple = sd_tuple
+        self.prob_tuple = prob_tuple
+
+    def sample(self, n_samples=1):
+        """Samples from a mixture of two Gaussian
+
+        Parameters
+        ----------
+        n_samples : int, optional
+            Number of samples, by default 1
+
+        Returns
+        -------
+        samples_list
+            List containing n_samples samples
+        """
+        samples_list = [0] * n_samples
+        mixture_id_list = np.random.binomial(1, self.prob_tuple[0], n_samples)
+        for i, mixture_id in enumerate(mixture_id_list):
+            if mixture_id:
+                samples_list[i] = np.random.normal(
+                    self.means_tuple[0], self.sd_tuple[0]
+                )
+            else:
+                samples_list[i] = np.random.normal(
+                    self.means_tuple[1], self.sd_tuple[1]
+                )
+
+        return samples_list
+
+    def reset_seed(self, seed=None):
+        """Resets all associated random number generators
+
+        Parameters
+        ----------
+        seed : int, optional
+            Seed for the random number generator.
+        """
+        np.random.default_rng(seed)

tests/fixtures/monte_carlo/custom_sampler_fixtures.py

@@ -43,6 +43,36 @@ def stochastic_environment(example_spaceport_env):
     )
 
 
+@pytest.fixture
+def stochastic_environment_custom_sampler(example_spaceport_env, elevation_sampler):
+    """This fixture is used to create a stochastic environment object for the
+    Calisto flight using a custom sampler for the elevation.
+
+    Parameters
+    ----------
+    example_spaceport_env : Environment
+        This is another fixture.
+
+    elevation_sampler: CustomSampler
+        This is another fixture.
+
+    Returns
+    -------
+    StochasticEnvironment
+        The stochastic environment object
+    """
+    return StochasticEnvironment(
+        environment=example_spaceport_env,
+        elevation=elevation_sampler,
+        gravity=None,
+        latitude=None,
+        longitude=None,
+        ensemble_member=None,
+        wind_velocity_x_factor=(1.0, 0.033, "normal"),
+        wind_velocity_y_factor=(1.0, 0.033, "normal"),
+    )
+
+
 @pytest.fixture
 def stochastic_nose_cone(calisto_nose_cone):
     """This fixture is used to create a StochasticNoseCone object for the

tests/fixtures/monte_carlo/stochastic_fixtures.py

@@ -0,0 +1,21 @@
+from rocketpy.environment.environment import Environment
+
+
+def test_create_object(stochastic_environment_custom_sampler):
+    """Test create object method of StochasticEnvironment class.
+
+    This test checks if the create_object method of the StochasticEnvironment
+    class creates a StochasticEnvironment object from the randomly generated
+    input arguments.
+
+    Parameters
+    ----------
+    stochastic_environment : StochasticEnvironment
+        StochasticEnvironment object to be tested.
+
+    Returns
+    -------
+    None
+    """
+    obj = stochastic_environment_custom_sampler.create_object()
+    assert isinstance(obj, Environment)

tests/unit/stochastic/test_custom_sampler.py

@@ -7,6 +7,7 @@
         "stochastic_rail_buttons",
         "stochastic_main_parachute",
         "stochastic_environment",
+        "stochastic_environment_custom_sampler",
         "stochastic_tail",
         "stochastic_calisto",
     ],