ENH: Callback function for collecting additional data from Monte Carlo sims

docs/notebooks/monte_carlo_analysis/monte_carlo_class_usage.ipynb

@@ -1115,6 +1115,93 @@
     "    type=\"impact\",\n",
     ")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Custom exports using callback functions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We have shown, so far, how to perform to use the `MonteCarlo` class and visualize its results. By default, some variables exported to the output files, such as *apogee* and *x_impact*. The `export_list` argument provides a simplified way for the user to export additional variables listed in the documentation, such as *inclination* and *heading*.  \n",
+    "\n",
+    "There are applications in which you might need to extract more information in the results than the `export_list` argument can handle. To that end, the `MonteCarlo` class has a `export_function` argument which allows you customize further the output of the simulation.\n",
+    "\n",
+    "To exemplify its use, we show how to export the *date* of the environment used in the simulation together with the *average reynolds number* along with the default variables."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We will use the `stochastic_env`, `stochastic_rocket` and `stochastic_flight` objects previously defined, and only change the `MonteCarlo` object. First, we need to define our customized export function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "def custom_export_function(flight):\n",
+    "    reynold_number_list = flight.reynolds_number(flight.time)\n",
+    "    average_reynolds_number = np.mean(reynold_number_list)\n",
+    "    custom_exports = {\n",
+    "        \"average_reynolds_number\": average_reynolds_number,\n",
+    "        \"date\": flight.env.date,\n",
+    "    }\n",
+    "    return custom_exports"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "The `export_function` must be a function which takes a `Flight` object and outputs a dictionary whose keys are variables names to export and their values. Notice how we can compute complex expressions in this function and just export the result. For instance, the *average_reynolds_number* calls the `flight.reynolds_number` method for each value in `flight.time` list and computes the average value using numpy's `mean`. The *date* variable is straightforward.\n",
+    "\n",
+    "After we define the export function, we pass it as an argument to the `MonteCarlo` class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_dispersion = MonteCarlo(\n",
+    "    filename=\"monte_carlo_analysis_outputs/monte_carlo_class_example_customized\",\n",
+    "    environment=stochastic_env,\n",
+    "    rocket=stochastic_rocket,\n",
+    "    flight=stochastic_flight,\n",
+    "    export_function=custom_export_function,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_dispersion.simulate(number_of_simulations=10, append=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_dispersion.prints.all()"
+   ]
   }
  ],
  "metadata": {
@@ -1134,7 +1221,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.11.2"
   }
  },
  "nbformat": 4,

rocketpy/prints/monte_carlo_prints.py

@@ -24,4 +24,7 @@
         print(f"{'Parameter':>25} {'Mean':>15} {'Std. Dev.':>15}")
         print("-" * 60)
         for key, value in self.monte_carlo.processed_results.items():
-            print(f"{key:>25} {value[0]:>15.3f} {value[1]:>15.3f}")
+            try:
+                print(f"{key:>25} {value[0]:>15.3f} {value[1]:>15.3f}")
+            except TypeError:
+                print(f"{key:>25} {str(value[0]):>15} {str(value[1]):>15}")

rocketpy/simulation/monte_carlo.py

@@ -80,7 +80,13 @@
     """
 
     def __init__(
-        self, filename, environment, rocket, flight, export_list=None
+        self,
+        filename,
+        environment,
+        rocket,
+        flight,
+        export_list=None,
+        export_function=None,
     ):  # pylint: disable=too-many-statements
         """
         Initialize a MonteCarlo object.
@@ -104,6 +110,11 @@
             `out_of_rail_stability_margin`, `out_of_rail_time`,
             `out_of_rail_velocity`, `max_mach_number`, `frontal_surface_wind`,
             `lateral_surface_wind`. Default is None.
+        export_function : callable, optional
+            A function which gets called at the end of a simulation to collect
+            additional data to be exported that isn't pre-defined. Takes the
+            Flight object as an argument and returns a dictionary. Default is None.
+
 
         Returns
         -------
@@ -132,6 +143,7 @@
         self._last_print_len = 0  # used to print on the same line
 
         self.export_list = self.__check_export_list(export_list)
+        self.export_function = export_function
 
         try:
             self.import_inputs()
@@ -359,6 +371,22 @@
             for export_item in self.export_list
         }
 
+        if self.export_function is not None:
+            try:
+                additional_exports = self.export_function(flight)
+            except Exception as e:
+                raise ValueError(
+                    "An error was encountered running your custom export function. "
+                    "Check for errors in 'export_function' definition."
+                ) from e
+
+            for key in additional_exports.keys():
+                if key in self.export_list:
+                    raise ValueError(
+                        f"Invalid export function, returns dict which overwrites key, '{key}'"
+                    )
+            results = results | additional_exports
+
         input_file.write(json.dumps(inputs_dict, cls=RocketPyEncoder) + "\n")
         output_file.write(json.dumps(results, cls=RocketPyEncoder) + "\n")
 
@@ -654,9 +682,15 @@
         """
         self.processed_results = {}
         for result, values in self.results.items():
-            mean = np.mean(values)
-            stdev = np.std(values)
-            self.processed_results[result] = (mean, stdev)
+            try:
+                if isinstance(values[0], float):
+                    mean = np.mean(values)
+                    stdev = np.std(values)
+                    self.processed_results[result] = (mean, stdev)
+                else:
+                    self.processed_results[result] = (None, None)
+            except TypeError:
+                self.processed_results[result] = (None, None)
 
     # Import methods
 

tests/integration/test_monte_carlo.py

@@ -111,3 +111,56 @@ def test_monte_carlo_export_ellipses_to_kml(monte_carlo_calisto_pre_loaded):
     )
 
     os.remove("monte_carlo_class_example.kml")
+
+
+@pytest.mark.slow
+def test_monte_carlo_callback(monte_carlo_calisto):
+    """Tests the export_function argument of the MonteCarlo class.
+
+    Parameters
+    ----------
+    monte_carlo_calisto : MonteCarlo
+        The MonteCarlo object, this is a pytest fixture.
+    """
+
+    def valid_export_function(flight):
+        custom_export_dict = {
+            "name": flight.name,
+            "density_t0": flight.env.density(0),
+        }
+        return custom_export_dict
+
+    monte_carlo_calisto.export_function = valid_export_function
+    # NOTE: this is really slow, it runs 10 flight simulations
+    monte_carlo_calisto.simulate(number_of_simulations=10, append=False)
+
+    # tests if print works when we have None in summary
+    monte_carlo_calisto.info()
+
+    # tests if logical errors in export functions raise errors
+    def export_function_with_logical_error(flight):
+        custom_export_dict = {
+            "date": flight.env.date,
+            "density_t0": flight.env.density(0) / "0",
+        }
+        return custom_export_dict
+
+    monte_carlo_calisto.export_function = export_function_with_logical_error
+    # NOTE: this is really slow, it runs 10 flight simulations
+    with pytest.raises(ValueError):
+        monte_carlo_calisto.simulate(number_of_simulations=10, append=False)
+
+    # tests if overwriting default exports raises errors
+    def export_function_with_overwriting_error(flight):
+        custom_export_dict = {
+            "apogee": flight.apogee,
+        }
+        return custom_export_dict
+
+    monte_carlo_calisto.export_function = export_function_with_overwriting_error
+    with pytest.raises(ValueError):
+        monte_carlo_calisto.simulate(number_of_simulations=10, append=False)
+
+    os.remove("monte_carlo_test.errors.txt")
+    os.remove("monte_carlo_test.outputs.txt")
+    os.remove("monte_carlo_test.inputs.txt")