ENH: first implementation of CustomSampler

Author: Lucas-Prates

rocketpy/__init__.py

@@ -44,6 +44,7 @@
 from .sensors import Accelerometer, Barometer, GnssReceiver, Gyroscope
 from .simulation import Flight, MonteCarlo, MultivariateRejectionSampler
 from .stochastic import (
+    CustomSampler,
     StochasticAirBrakes,
     StochasticEllipticalFins,
     StochasticEnvironment,

rocketpy/_encoders.py

@@ -7,8 +7,8 @@
 import numpy as np
 
 from rocketpy.mathutils.function import Function
-from rocketpy.prints.flight_prints import _FlightPrints
 from rocketpy.plots.flight_plots import _FlightPlots
+from rocketpy.prints.flight_prints import _FlightPrints
 
 
 class RocketPyEncoder(json.JSONEncoder):

rocketpy/stochastic/__init__.py

@@ -5,6 +5,7 @@
 associated with each input parameter.
 """
 
+from .custom_sampler import CustomSampler
 from .stochastic_aero_surfaces import (
     StochasticAirBrakes,
     StochasticEllipticalFins,

rocketpy/stochastic/custom_sampler.py

@@ -0,0 +1,38 @@
+"""
+Provides an abstract class so that users can build custom samplers upon
+"""
+
+from abc import ABC, abstractmethod
+
+
+class CustomSampler(ABC):
+    """Abstract subclass for user defined samplers"""
+
+    @abstractmethod
+    def sample(self, n_samples=1):
+        """Generates n samples from the custom distribution
+
+        Parameters
+        ----------
+        n_samples : int, optional
+            Numbers of samples to be generated
+
+        Returns
+        -------
+        sample_list : list
+            A list with n_samples elements, each of which is a valid sample
+        """
+
+    @abstractmethod
+    def reset_seed(self, seed=None):
+        """Resets the seeds of all associated stochastic generators
+
+        Parameters
+        ----------
+        seed : int, optional
+            Seed for the random number generator. The default is None
+
+        Returns
+        -------
+        None
+        """

rocketpy/stochastic/stochastic_model.py

@@ -8,6 +8,7 @@
 import numpy as np
 
 from rocketpy.mathutils.function import Function
+from rocketpy.stochastic.custom_sampler import CustomSampler
 
 from ..tools import get_distribution
 
@@ -96,6 +97,10 @@ def _set_stochastic(self, seed=None):
                             attr_value = self._validate_list(input_name, input_value)
                         elif isinstance(input_value, (int, float)):
                             attr_value = self._validate_scalar(input_name, input_value)
+                        elif isinstance(input_value, CustomSampler):
+                            attr_value = self._validate_custom_sampler(
+                                input_name, input_value, seed
+                            )
                         else:
                             raise AssertionError(
                                 f"'{input_name}' must be a tuple, list, int, or float"
@@ -436,6 +441,33 @@ def _validate_positive_int_list(self, input_name, input_value):
                 isinstance(member, int) and member >= 0 for member in input_value
             ), f"`{input_name}` must be a list of positive integers"
 
+    def _validate_custom_sampler(self, input_name, sampler, seed=None):
+        """
+        Validate a custom sampler.
+
+        Parameters
+        ----------
+        input_name : str
+            Name of the input argument.
+        sampler : CustomSampler object
+            Custom sampler provided by the user
+        seed : int, optional
+            Seed for the random number generator. The default is None
+
+        Raises
+        ------
+        AssertionError
+            If the input is not in a valid format.
+        """
+        try:
+            sampler.reset_seed(seed)
+        except RuntimeError as e:
+            raise RuntimeError(
+                f"An error occurred in the 'reset_seed' of {input_name} CustomSampler"
+            ) from e
+
+        return sampler
+
     def _validate_airfoil(self, airfoil):
         """
         Validate airfoil input.
@@ -490,9 +522,17 @@ def dict_generator(self):
         generated_dict = {}
         for arg, value in self.__dict__.items():
             if isinstance(value, tuple):
-                generated_dict[arg] = value[-1](value[0], value[1])
+                dist_sampler = value[-1]
+                generated_dict[arg] = dist_sampler(value[0], value[1])
             elif isinstance(value, list):
                 generated_dict[arg] = choice(value) if value else value
+            elif isinstance(value, CustomSampler):
+                try:
+                    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"
+                    ) from e
         self.last_rnd_dict = generated_dict
         yield generated_dict
 
@@ -527,6 +567,12 @@ def format_attribute(attr, value):
                         f"{nominal_value:.5f} ± "
                         f"{std_dev:.5f} ({dist_func.__name__})"
                     )
+            elif isinstance(value, CustomSampler):
+                sampler_name = type(value).__name__
+                return (
+                    f"\t{attr.ljust(max_str_length)} "
+                    f"\t{sampler_name.ljust(max_str_length)} "
+                )
             return None
 
         attributes = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
@@ -550,6 +596,9 @@ def format_attribute(attr, value):
         list_attributes = [
             attr for attr, val in items if isinstance(val, list) and len(val) > 1
         ]
+        custom_attributes = [
+            attr for attr, val in items if isinstance(val, CustomSampler)
+        ]
 
         if constant_attributes:
             report.append("\nConstant Attributes:")
@@ -568,5 +617,10 @@ def format_attribute(attr, value):
             report.extend(
                 format_attribute(attr, attributes[attr]) for attr in list_attributes
             )
+        if custom_attributes:
+            report.append("\nStochastic Attributes with Custom user samplers:")
+            report.extend(
+                format_attribute(attr, attributes[attr]) for attr in custom_attributes
+            )
 
         print("\n".join(filter(None, report)))

rocketpy/utilities.py

@@ -1,23 +1,23 @@
 import ast
 import inspect
-import traceback
-import warnings
 import json
 import os
-
-from pathlib import Path
-from importlib.metadata import version
+import traceback
+import warnings
 from datetime import date
+from importlib.metadata import version
+from pathlib import Path
+
 import matplotlib.pyplot as plt
 import numpy as np
 from scipy.integrate import solve_ivp
 
+from ._encoders import RocketPyDecoder, RocketPyEncoder
 from .environment.environment import Environment
 from .mathutils.function import Function
 from .plots.plot_helpers import show_or_save_plot
 from .rocket.aero_surface import TrapezoidalFins
 from .simulation.flight import Flight
-from ._encoders import RocketPyEncoder, RocketPyDecoder
 
 
 def compute_cd_s_from_drop_test(