Minor fixups ensembles (#568)

* fixup in chunk sizing and handng over of parameter mappings, fixes a bug which may yield to incorrect predictions if using of chunk size

* allow passing on of umap parameters tothe umap routines

* make user-defined fields of amici return data available to postprocessors

Co-authored-by: Dilan Pathirana <59329744+dilpath@users.noreply.github.com>

pypesto/ensemble/dimension_reduction.py

@@ -16,10 +16,12 @@
 def get_umap_representation_parameters(
         ens: Ensemble,
         n_components: int = 2,
-        normalize_data: bool = False) -> Tuple:
+        normalize_data: bool = False,
+        **kwargs) -> Tuple:
     """
     Compute the representation with reduced dimensionality via umap
     (with a given number of umap components) of the parameter ensemble.
+    Allows to pass on additional keyword arguments to the umap routine.
 
     Parameters
     ==========
@@ -46,18 +48,21 @@ def get_umap_representation_parameters(
     return _get_umap_representation_lowlevel(
         dataset=ens.x_vectors.transpose(),
         n_components=n_components,
-        normalize_data=normalize_data
+        normalize_data=normalize_data,
+        **kwargs
     )
 
 
 def get_umap_representation_predictions(
         ens: Union[Ensemble, EnsemblePrediction],
         prediction_index: int = 0,
         n_components: int = 2,
-        normalize_data: bool = False) -> Tuple:
+        normalize_data: bool = False,
+        **kwargs) -> Tuple:
     """
     Compute the representation with reduced dimensionality via umap
     (with a given number of umap components) of the ensemble predictions.
+    Allows to pass on additional keyword arguments to the umap routine.
 
     Parameters
     ==========
@@ -92,7 +97,8 @@ def get_umap_representation_predictions(
     return _get_umap_representation_lowlevel(
         dataset=dataset,
         n_components=n_components,
-        normalize_data=normalize_data
+        normalize_data=normalize_data,
+        **kwargs
     )
 
 
@@ -197,7 +203,8 @@ def get_pca_representation_predictions(
 def _get_umap_representation_lowlevel(
         dataset: np.ndarray,
         n_components: int = 2,
-        normalize_data: bool = False) -> Tuple:
+        normalize_data: bool = False,
+        **kwargs) -> Tuple:
     """
     Compute the representation with reduced dimensionality via principal
     component analysis (with a given number of principal components) of the
@@ -230,7 +237,7 @@ def _get_umap_representation_lowlevel(
     """
 
     # create a umap object
-    umap_object = umap.UMAP(n_components=n_components)
+    umap_object = umap.UMAP(n_components=n_components, **kwargs)
 
     # normalize data with mean and standard deviation if wanted
     if normalize_data:

pypesto/ensemble/ensemble.py

@@ -294,7 +294,7 @@ def __init__(self,
         # store bounds
         self.lower_bound = np.full((self.n_x,), np.nan)
         if lower_bound is not None:
-            if len(lower_bound) == 1:
+            if np.array(lower_bound).size == 1:
                 self.lower_bound = np.full((x_vectors.shape[0],), lower_bound)
             else:
                 self.lower_bound = lower_bound

pypesto/objective/amici.py

@@ -296,7 +296,8 @@ def check_sensi_orders(self, sensi_orders, mode) -> bool:
     def check_mode(self, mode):
         return mode in [MODE_FUN, MODE_RES]
 
-    def call_unprocessed(self, x, sensi_orders, mode, edatas=None):
+    def call_unprocessed(self, x, sensi_orders, mode, edatas=None,
+                         parameter_mapping=None):
         sensi_order = max(sensi_orders)
 
         x_dct = self.par_arr_to_dct(x)
@@ -309,11 +310,13 @@ def call_unprocessed(self, x, sensi_orders, mode, edatas=None):
 
         if edatas is None:
             edatas = self.edatas
+        if parameter_mapping is None:
+            parameter_mapping = self.parameter_mapping
         ret = self.calculator(
             x_dct=x_dct, sensi_order=sensi_order, mode=mode,
             amici_model=self.amici_model, amici_solver=self.amici_solver,
             edatas=edatas, n_threads=self.n_threads,
-            x_ids=self.x_ids, parameter_mapping=self.parameter_mapping,
+            x_ids=self.x_ids, parameter_mapping=parameter_mapping,
             fim_for_hess=self.fim_for_hess,
         )
 

pypesto/petab/importer.py

@@ -272,6 +272,7 @@ def create_objective(
 
     def create_prediction(self,
                           objective: AmiciObjective = None,
+                          amici_output_fields: Sequence[str] = None,
                           post_processor: Union[Callable, None] = None,
                           post_processor_sensi: Union[Callable, None] = None,
                           post_processor_time: Union[Callable, None] = None,
@@ -285,6 +286,9 @@ def create_prediction(self,
         ----------
         objective:
             An objective object, which will be used to get model simulations
+        amici_output_fields:
+            keys that exist in the return data object from AMICI, which should
+            be available for the post-processors
         post_processor:
             A callable function which applies postprocessing to the simulation
             results. Default are the observables of the amici model.
@@ -343,6 +347,7 @@ def create_prediction(self,
         # wrap around AmiciPredictor
         predictor = AmiciPredictor(
             amici_objective=objective,
+            amici_output_fields=amici_output_fields,
             post_processor=post_processor,
             post_processor_sensi=post_processor_sensi,
             post_processor_time=post_processor_time,

pypesto/prediction/amici_predictor.py

@@ -1,5 +1,6 @@
 import numpy as np
 from typing import Sequence, Union, Callable, Tuple, List
+from copy import deepcopy
 
 from .constants import (MODE_FUN, OBSERVABLE_IDS, TIMEPOINTS, OUTPUT,
                         OUTPUT_SENSI, CSV, H5, AMICI_T, AMICI_X, AMICI_SX,
@@ -25,6 +26,7 @@ class AmiciPredictor:
     """
     def __init__(self,
                  amici_objective: AmiciObjective,
+                 amici_output_fields: Union[Sequence[str], None] = None,
                  post_processor: Union[Callable, None] = None,
                  post_processor_sensi: Union[Callable, None] = None,
                  post_processor_time: Union[Callable, None] = None,
@@ -39,6 +41,9 @@ def __init__(self,
         ----------
         amici_objective:
             An objective object, which will be used to get model simulations
+        amici_output_fields:
+            keys that exist in the return data object from AMICI, which should
+            be available for the post-processors
         post_processor:
             A callable function which applies postprocessing to the simulation
             results and possibly defines different observables than those of
@@ -83,12 +88,23 @@ def __init__(self,
         self.post_processor_time = post_processor_time
         self.condition_ids = condition_ids
 
+        # If the user takes care of everything we can skip default readouts
+        self.skip_default_outputs = False
+        if post_processor is not None and post_processor_sensi is not None \
+                and post_processor_time is not None:
+            self.skip_default_outputs = True
+
         if observable_ids is None:
             self.observable_ids = \
                 amici_objective.amici_model.getObservableIds()
         else:
             self.observable_ids = observable_ids
 
+        if amici_output_fields is None:
+            amici_output_fields = [AMICI_STATUS, AMICI_T, AMICI_X, AMICI_Y,
+                                   AMICI_SX, AMICI_SY]
+        self.amici_output_fields = amici_output_fields
+
     def __call__(
             self,
             x: np.ndarray,
@@ -223,7 +239,8 @@ def _get_outputs(self,
             # call amici
             self._wrap_call_to_amici(
                 amici_outputs=amici_outputs, x=x, sensi_orders=sensi_orders,
-                mode=mode, edatas=self.amici_objective.edatas[ids])
+                parameter_mapping=self.amici_objective.parameter_mapping[ids],
+                edatas=self.amici_objective.edatas[ids], mode=mode)
 
         def _default_output(amici_outputs):
             """
@@ -261,7 +278,8 @@ def _default_output(amici_outputs):
             return timepoints, outputs, outputs_sensi
 
         # Get default output
-        timepoints, outputs, outputs_sensi = _default_output(amici_outputs)
+        if not self.skip_default_outputs:
+            timepoints, outputs, outputs_sensi = _default_output(amici_outputs)
 
         # postprocess (use original Amici outputs)
         if self.post_processor is not None:
@@ -274,19 +292,19 @@ def _default_output(amici_outputs):
         return timepoints, outputs, outputs_sensi
 
     def _wrap_call_to_amici(self, amici_outputs, x, sensi_orders, mode,
-                            edatas):
+                            parameter_mapping, edatas):
         """
         The only purpose of this function is to encapsulate the call to amici:
         This allows to use variable scoping as a mean to clean up the memory
         after calling amici, which is beneficial if large models with large
         datasets are used.
         """
         chunk = self.amici_objective(x=x, sensi_orders=sensi_orders, mode=mode,
+                                     parameter_mapping=parameter_mapping,
                                      edatas=edatas,  return_dict=True)
         for rdata in chunk[RDATAS]:
-            amici_outputs.append({AMICI_STATUS: rdata[AMICI_STATUS],
-                                  AMICI_T: rdata[AMICI_T],
-                                  AMICI_X: rdata[AMICI_X],
-                                  AMICI_SX: rdata[AMICI_SX],
-                                  AMICI_Y: rdata[AMICI_Y],
-                                  AMICI_SY: rdata[AMICI_SY]})
+            amici_outputs.append({
+                output_field: deepcopy(rdata[output_field])
+                for output_field in self.amici_output_fields
+            })
+        del chunk

pypesto/prediction/constants.py

@@ -8,6 +8,7 @@
 
 OBSERVABLE_IDS = 'observable_ids'  # data member in PredictionConditionResult
 PARAMETER_IDS = 'x_names'  # data member in PredictionConditionResult
+CONDITION_IDS = 'condition_ids'
 TIMEPOINTS = 'timepoints'  # data member in PredictionConditionResult
 OUTPUT = 'output'  # field in the return dict of AmiciPredictor
 OUTPUT_SENSI = 'output_sensi'  # field in the return dict of AmiciPredictor

pypesto/prediction/prediction.py

@@ -8,7 +8,7 @@
 import os
 
 from .constants import (OBSERVABLE_IDS, PARAMETER_IDS, TIMEPOINTS, OUTPUT,
-                        OUTPUT_SENSI, TIME, CSV)
+                        OUTPUT_SENSI, TIME, CSV, CONDITION_IDS)
 
 
 class PredictionConditionResult:
@@ -207,6 +207,9 @@ def write_to_h5(self,
             if self.conditions and self.conditions[0].x_names is not None:
                 f.create_dataset(os.path.join(base, PARAMETER_IDS),
                                  data=self.conditions[0].x_names)
+            if self.condition_ids is not None:
+                f.create_dataset(os.path.join(base, CONDITION_IDS),
+                                 data=self.condition_ids)
             for i_cond, cond in enumerate(self.conditions):
                 # each conditions gets a group of its own
                 f.create_group(os.path.join(base, str(i_cond)))