diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py index dd246b20301..71d68efba55 100644 --- a/src/everest/everest_storage.py +++ b/src/everest/everest_storage.py @@ -445,62 +445,40 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult "evaluations", select=["objectives", "evaluation_ids"], ).reset_index(), - ).drop("plan_id") - - realization_constraints = polars.from_pandas( - results.to_dataframe( - "evaluations", - select=["constraints", "evaluation_ids"], - ).reset_index(), - ).drop("plan_id") - - realization_constraints = self._rename_ropt_df_columns(realization_constraints) - - batch_objectives = polars.from_pandas( - results.to_dataframe( - "functions", - select=["objectives", "weighted_objective"], - ).reset_index() - ).drop("plan_id") - - realization_controls = polars.from_pandas( - results.to_dataframe( - "evaluations", select=["variables", "evaluation_ids"] - ).reset_index() - ).drop("plan_id") + ).select( + "result_id", + "batch_id", + "realization", + "objective", + "objectives", + "evaluation_ids", + ) - realization_controls = self._rename_ropt_df_columns(realization_controls) - realization_controls = self._enforce_dtypes(realization_controls) + if results.nonlinear_constraints is not None: + realization_constraints = polars.from_pandas( + results.to_dataframe( + "evaluations", + select=["constraints", "evaluation_ids"], + ).reset_index(), + ).select( + "result_id", + "batch_id", + "realization", + "evaluation_ids", + "nonlinear_constraint", + "constraints", + ) - realization_controls = realization_controls.pivot( - on="control_name", - values=["control_value"], - separator=":", - ) + realization_constraints = self._rename_ropt_df_columns( + realization_constraints + ) - try: batch_constraints = polars.from_pandas( results.to_dataframe("nonlinear_constraints").reset_index() - ).drop("plan_id") - batch_constraints = batch_constraints.drop( - [c for c in batch_constraints.columns if c.startswith("scaled_")] + ).select( + "result_id", "batch_id", "nonlinear_constraint", "values", "violations" ) - except AttributeError: - batch_constraints = None - - batch_objectives = self._rename_ropt_df_columns(batch_objectives) - batch_objectives = self._enforce_dtypes(batch_objectives) - - realization_objectives = self._rename_ropt_df_columns(realization_objectives) - realization_objectives = self._enforce_dtypes(realization_objectives) - batch_objectives = batch_objectives.pivot( - on="objective_name", - values=["objective_value"], - separator=":", - ) - - if batch_constraints is not None: batch_constraints = batch_constraints.rename( { "nonlinear_constraint": "constraint_name", @@ -531,9 +509,52 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult realization_constraints = realization_constraints.pivot( values=["constraint_value"], on="constraint_name" ) - batch_objectives = batch_objectives.drop( - [c for c in batch_objectives.columns if "constraint" in c.lower()] - ) + else: + batch_constraints = None + realization_constraints = None + + batch_objectives = polars.from_pandas( + results.to_dataframe( + "functions", + select=["objectives", "weighted_objective"], + ).reset_index() + ).select( + "result_id", "batch_id", "objective", "objectives", "weighted_objective" + ) + + realization_controls = polars.from_pandas( + results.to_dataframe( + "evaluations", select=["variables", "evaluation_ids"] + ).reset_index() + ).select( + "result_id", + "batch_id", + "variable", + "realization", + "variables", + "evaluation_ids", + ) + + realization_controls = self._rename_ropt_df_columns(realization_controls) + realization_controls = self._enforce_dtypes(realization_controls) + + realization_controls = realization_controls.pivot( + on="control_name", + values=["control_value"], + separator=":", + ) + + batch_objectives = self._rename_ropt_df_columns(batch_objectives) + batch_objectives = self._enforce_dtypes(batch_objectives) + + realization_objectives = self._rename_ropt_df_columns(realization_objectives) + realization_objectives = self._enforce_dtypes(realization_objectives) + + batch_objectives = batch_objectives.pivot( + on="objective_name", + values=["objective_value"], + separator=":", + ) realization_objectives = realization_objectives.pivot( values="objective_value",