From d3fda4e62fc3fb54f8cea7fe904fc46f19c66a22 Mon Sep 17 00:00:00 2001 From: Peter Verveer Date: Fri, 31 Jan 2025 10:24:17 +0000 Subject: [PATCH] Insert names in everest results using polars --- src/everest/everest_storage.py | 106 +++++++++++++++++++++++---------- 1 file changed, 75 insertions(+), 31 deletions(-) diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py index 30948a66f70..2fab80659b5 100644 --- a/src/everest/everest_storage.py +++ b/src/everest/everest_storage.py @@ -408,24 +408,12 @@ def init(self, everest_config: EverestConfig) -> None: ) def _store_function_results(self, results: FunctionResults) -> _EvaluationResults: - names = { - "variable": self.data.controls["control_name"], - "objective": self.data.objective_functions["objective_name"], - "nonlinear_constraint": ( - self.data.nonlinear_constraints["constraint_name"] - if self.data.nonlinear_constraints is not None - else None - ), - "realization": self.data.realization_weights["realization"], - } - # We could select only objective values, # but we select all to also get the constraint values (if they exist) realization_objectives = polars.from_pandas( results.to_dataframe( "evaluations", select=["objectives", "evaluation_ids"], - names=names, ).reset_index(), ).select( "batch_id", @@ -434,13 +422,20 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult "objectives", "evaluation_ids", ) + realization_objectives = realization_objectives.with_columns( + polars.col("realization").replace_strict( + dict(enumerate(self.data.realization_weights["realization"])) + ), + polars.col("objective").replace_strict( + dict(enumerate(self.data.objective_functions["objective_name"])) + ), + ) if results.functions is not None and results.functions.constraints is not None: realization_constraints = polars.from_pandas( results.to_dataframe( "evaluations", select=["constraints", "evaluation_ids"], - names=names, ).reset_index(), ).select( "batch_id", @@ -449,16 +444,27 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult "nonlinear_constraint", "constraints", ) + realization_constraints = realization_constraints.with_columns( + polars.col("realization").replace_strict( + dict(enumerate(self.data.realization_weights["realization"])) + ), + polars.col("nonlinear_constraint").replace_strict( + dict(enumerate(self.data.nonlinear_constraints["constraint_name"])) + ), + ) realization_constraints = self._rename_ropt_df_columns( realization_constraints ) batch_constraints = polars.from_pandas( - results.to_dataframe( - "functions", select=["constraints"], names=names - ).reset_index() + results.to_dataframe("functions", select=["constraints"]).reset_index() ).select("batch_id", "nonlinear_constraint", "constraints") + batch_constraints = batch_constraints.with_columns( + polars.col("nonlinear_constraint").replace_strict( + dict(enumerate(self.data.nonlinear_constraints["constraint_name"])) + ), + ) batch_constraints = batch_constraints.rename( { @@ -485,13 +491,17 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult results.to_dataframe( "functions", select=["objectives", "weighted_objective"], - names=names, ).reset_index() ).select("batch_id", "objective", "objectives", "weighted_objective") + batch_objectives = batch_objectives.with_columns( + polars.col("objective").replace_strict( + dict(enumerate(self.data.objective_functions["objective_name"])) + ), + ) realization_controls = polars.from_pandas( results.to_dataframe( - "evaluations", select=["variables", "evaluation_ids"], names=names + "evaluations", select=["variables", "evaluation_ids"] ).reset_index() ).select( "batch_id", @@ -500,6 +510,14 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult "variables", "evaluation_ids", ) + realization_controls = realization_controls.with_columns( + polars.col("variable").replace_strict( + dict(enumerate(self.data.controls["control_name"])) + ), + polars.col("realization").replace_strict( + dict(enumerate(self.data.realization_weights["realization"])) + ), + ) realization_controls = self._rename_ropt_df_columns(realization_controls) realization_controls = self._enforce_dtypes(realization_controls) @@ -541,19 +559,8 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult } def _store_gradient_results(self, results: GradientResults) -> _GradientResults: - names = { - "variable": self.data.controls["control_name"], - "objective": self.data.objective_functions["objective_name"], - "nonlinear_constraint": ( - self.data.nonlinear_constraints["constraint_name"] - if self.data.nonlinear_constraints is not None - else None - ), - "realization": self.data.realization_weights["realization"], - } - perturbation_objectives = polars.from_pandas( - results.to_dataframe("evaluations", names=names).reset_index() + results.to_dataframe("evaluations").reset_index() ).select( [ "batch_id", @@ -572,12 +579,30 @@ def _store_gradient_results(self, results: GradientResults) -> _GradientResults: ), ] ) + perturbation_objectives = perturbation_objectives.with_columns( + polars.col("variable").replace_strict( + dict(enumerate(self.data.controls["control_name"])) + ), + polars.col("realization").replace_strict( + dict(enumerate(self.data.realization_weights["realization"])) + ), + polars.col("objective").replace_strict( + dict(enumerate(self.data.objective_functions["objective_name"])) + ), + ) + if results.evaluations.perturbed_constraints is not None: + perturbation_objectives = perturbation_objectives.with_columns( + polars.col("nonlinear_constraint").replace_strict( + dict(enumerate(self.data.nonlinear_constraints["constraint_name"])) + ), + ) + perturbation_objectives = self._rename_ropt_df_columns(perturbation_objectives) perturbation_objectives = self._enforce_dtypes(perturbation_objectives) if results.gradients is not None: batch_objective_gradient = polars.from_pandas( - results.to_dataframe("gradients", names=names).reset_index() + results.to_dataframe("gradients").reset_index() ).select( [ "batch_id", @@ -592,6 +617,25 @@ def _store_gradient_results(self, results: GradientResults) -> _GradientResults: ), ] ) + batch_objective_gradient = batch_objective_gradient.with_columns( + polars.col("variable").replace_strict( + dict(enumerate(self.data.controls["control_name"])) + ), + polars.col("objective").replace_strict( + dict(enumerate(self.data.objective_functions["objective_name"])) + ), + ) + if results.gradients.constraints is not None: + batch_objective_gradient = batch_objective_gradient.with_columns( + polars.col("nonlinear_constraint").replace_strict( + dict( + enumerate( + self.data.nonlinear_constraints["constraint_name"] + ) + ) + ), + ) + batch_objective_gradient = self._rename_ropt_df_columns( batch_objective_gradient )