247 add convenience function to join extra parameters to the analytic…

…s system quantities output (#248)

* added some properties including direct system quantities access to simulation set

* enhanced test

* Update index.rst

doc/index.rst

@@ -3,7 +3,7 @@ Welcome to RidePy's documentation!
 
 RidePy is a scientific Python library for simulating modern on-demand transit systems such as ridepooling.
 
-To better understand what RidePy is all about probably should start by reading the :doc:`overview section <overview>`. For instructions on installation, contributing and support, head over to the :doc:`*How to...* section <howto>`. And to dive right in and get started with using RidePy for a simple simulation, try the :doc:`tutorial <notebooks>`.
+To better understand what RidePy is all about probably should start by reading the :doc:`overview section <overview>`. For instructions on installation, contributing and support, head over to the :doc:`**How to...** section <howto>`. And to dive right in and get started with using RidePy for a simple simulation, try the :doc:`tutorial <notebooks>`.
 
 In addition to these general guides, we offer a commented :doc:`API reference <reference>`, in which we explain the building blocks and concepts that RidePy is based on in more detail, and how you can leverage them for your specific application. This section is far more technical than the hands-on tutorial, but allows for a deeper understanding.
 

src/ridepy/extras/simulation_set.py

@@ -20,6 +20,7 @@
     Union,
     Sequence,
     Callable,
+    Mapping,
 )
 from pathlib import Path
 
@@ -54,11 +55,13 @@
 logger = logging.getLogger(__name__)
 
 
-def make_file_path(sim_id: str, directory: Path, suffix: str):
+def make_file_path(sim_id: str, directory: Path, suffix: str) -> Path:
     return directory / f"{sim_id}{suffix}"
 
 
-def get_params(directory: Path, sim_id: str, param_path_suffix: str = "_params.json"):
+def get_params(
+    directory: Path, sim_id: str, param_path_suffix: str = "_params.json"
+) -> dict:
     return read_params_json(make_file_path(sim_id, directory, param_path_suffix))
 
 
@@ -568,9 +571,6 @@ def __init__(
             transportation_request_cls = TransportationRequest
             vehicle_state_cls = VehicleState
 
-        fleet_state_cls = SlowSimpleFleetState
-        request_generator_cls = RandomRequestGenerator
-
         self.default_base_params = dict(
             general=dict(
                 n_reqs=100,
@@ -582,11 +582,12 @@ def __init__(
                 seat_capacity=8,
                 transportation_request_cls=transportation_request_cls,
                 vehicle_state_cls=vehicle_state_cls,
-                fleet_state_cls=fleet_state_cls,
+                fleet_state_cls=SlowSimpleFleetState,
             ),
             dispatcher=dict(dispatcher_cls=dispatcher),
             request_generator=dict(
-                request_generator_cls=request_generator_cls,
+                request_generator_cls=RandomRequestGenerator,
+                rate=1,
             ),
         )
 
@@ -626,7 +627,7 @@ def __init__(
 
         self._simulation_ids = None
 
-        self.system_quantities_path = None
+        self._system_quantities_path = None
 
     @property
     def simulation_ids(self) -> list[str]:
@@ -818,7 +819,7 @@ def run_analytics(
         system_quantities_filename
             Filename of the parquet file to store the system quantities in.
         """
-        self.system_quantities_path = self.data_dir / system_quantities_filename
+        self._system_quantities_path = self.data_dir / system_quantities_filename
 
         if not self.simulation_ids:
             warnings.warn(
@@ -845,15 +846,15 @@ def run_analytics(
             else:
                 raise ValueError(f"Got invalid value for {update_existing=}")
 
-            if self.system_quantities_path.exists():
+            if self._system_quantities_path.exists():
                 tasks_if_not_existent -= {"system_quantities"}
                 if check_for_changes:
                     # Currently, we only check for changes in the sense that new
                     # simulations with new ids have been performed and are thus missing from
                     # the system quantities output. For vehicle quantities, stops, and requests,
                     # check_for_changes does not apply, as we compute these in any case, should
                     # they be missing.
-                    sqdf = pd.read_parquet(self.system_quantities_path)
+                    sqdf = pd.read_parquet(self._system_quantities_path)
                     if set(sqdf.index) == set(self.simulation_ids):
                         tasks_if_existent -= {"system_quantities"}
                     del sqdf
@@ -888,4 +889,80 @@ def run_analytics(
             if "system_quantities" in tasks_if_not_existent | tasks_if_existent:
                 system_quantities_df = pd.DataFrame(system_quantities, index=sim_ids)
                 system_quantities_df.rename_axis("simulation_id", inplace=True)
-                system_quantities_df.to_parquet(self.system_quantities_path)
+                system_quantities_df.to_parquet(self._system_quantities_path)
+
+    @property
+    def param_path_suffix(self) -> str:
+        """
+        Get the parameter file suffix.
+        """
+        return self._param_path_suffix
+
+    @property
+    def event_path_suffix(self) -> str:
+        """
+        Get the event file suffix.
+        """
+        return self._event_path_suffix
+
+    @property
+    def system_quantities_path(self) -> Path:
+        """
+        Get path to the parquet file containing the system quantities.
+
+        Returns
+        -------
+        system_quantities_path
+        """
+        if self._system_quantities_path is not None:
+            return self._system_quantities_path
+        else:
+            raise AttributeError("No system quantities path set.")
+
+    def get_system_quantities(
+        self, extra_params: Optional[Mapping] = None
+    ) -> pd.DataFrame:
+        """
+        Return the system quantities, if computed already using `run_analytics`.
+
+        Optionally, join the system quantities dataframe with additional arbitrary
+        parameters used in the simulation.
+
+        Parameters
+        ----------
+        extra_params:
+            Dictionary specifying the desired parameters. The keys are the resulting column
+            names in the returned dataframe, the values are the (nested) keys in the
+            `SimulationSet` params dictionary, joined by ``.``.
+
+            Example:
+
+            .. code-block:: python
+
+                extra_params = {"n": "general.n_reqs"}
+
+        Returns
+        -------
+        system_quantities
+            DataFrame containing the system quantities and optional extra parameters
+
+        """
+
+        try:
+            sqdf = pd.read_parquet(self.system_quantities_path)
+        except FileNotFoundError:
+            raise FileNotFoundError(
+                f"System properties file {self.system_quantities_path} not found."
+            )
+
+        if extra_params:
+            for sim_id in sqdf.index:
+                params = get_params(
+                    self.data_dir, sim_id, param_path_suffix=self.param_path_suffix
+                )
+                for param_col, param_key in extra_params.items():
+                    sqdf.loc[sim_id, param_col] = ft.reduce(
+                        op.getitem, param_key.split("."), params
+                    )
+
+        return sqdf

src/ridepy/util/analytics/system.py

@@ -1,7 +1,7 @@
-from typing import Optional, Any, Dict, Union
-
 import pandas as pd
 
+from typing import Optional, Any, Union
+
 
 def get_system_quantities(
     stops: pd.DataFrame,

test/test_extras.py

@@ -104,6 +104,35 @@ def test_io_simulate(tmp_path):
     pd.testing.assert_frame_equal(requests1, requests2)
 
 
+def test_get_system_quantities(tmp_path):
+    simulation_set = SimulationSet(
+        base_params={"general": {"n_reqs": 5}},
+        zip_params={"general": {"n_vehicles": [10, 20, 30]}},
+        data_dir=tmp_path,
+        debug=True,
+    )
+
+    simulation_set.run()
+
+    with pytest.raises(AttributeError):
+        simulation_set.system_quantities_path
+
+    with pytest.raises(AttributeError):
+        simulation_set.get_system_quantities()
+
+    simulation_set.run_analytics()
+
+    assert simulation_set.system_quantities_path.exists()
+
+    sqdf = simulation_set.get_system_quantities(
+        extra_params={"B": "general.n_vehicles", "n_reqs": "general.n_reqs"}
+    )
+
+    assert len(sqdf) == 3
+    assert np.array_equal(sqdf.B, [10, 20, 30])
+    assert np.array_equal(sqdf.n_reqs, [5, 5, 5])
+
+
 @pytest.mark.parametrize("cython", [True, False])
 def test_io_params(cython, tmp_path):
     param_path = tmp_path / f"params.json"