review: no drop plan_id @ fn results

equinor · Jan 13, 2025 · 2fb4776 · 2fb4776
1 parent fbbc81b
commit 2fb4776
Showing 1 changed file with 73 additions and 52 deletions.
diff --git 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",