review: no drop plan_id @ grad results

equinor · Jan 14, 2025 · 005c29e · 005c29e
1 parent 3fc9a64
commit 005c29e
Showing 1 changed file with 38 additions and 20 deletions.
diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py
@@ -574,38 +574,56 @@ def _store_function_results(self, results: FunctionResults) -> _EvaluationResult
             "realization_constraints": realization_constraints,
         }
 
-    def _store_gradient_results(self, results: FunctionResults) -> _GradientResults:
+    def _store_gradient_results(self, results: GradientResults) -> _GradientResults:
         perturbation_objectives = polars.from_pandas(
             results.to_dataframe("evaluations").reset_index()
-        ).drop("plan_id")
-        perturbation_objectives = perturbation_objectives.drop(
-            c for c in perturbation_objectives.columns if c.lower().startswith("scaled")
+        ).select(
+            [
+                "result_id",
+                "batch_id",
+                "variable",
+                "realization",
+                "perturbation",
+                "objective",
+                "variables",
+                "perturbed_variables",
+                "perturbed_objectives",
+                "perturbed_evaluation_ids",
+                *(
+                    ["nonlinear_constraint", "perturbed_constraints"]
+                    if results.evaluations.perturbed_constraints is not None
+                    else []
+                ),
+            ]
         )
+        perturbation_objectives = self._rename_ropt_df_columns(perturbation_objectives)
 
-        try:
-            # ROPT_NOTE: Why is this sometimes None? How can we know if it is
-            # expected to be None?
+        if results.gradients is not None:
             batch_objective_gradient = polars.from_pandas(
                 results.to_dataframe("gradients").reset_index()
-            ).drop("plan_id")
-        except AttributeError:
-            batch_objective_gradient = None
-
-        if batch_objective_gradient is not None:
-            batch_objective_gradient = batch_objective_gradient.drop(
-                c
-                for c in batch_objective_gradient.columns
-                if c.lower().startswith("scaled")
+            ).select(
+                [
+                    "result_id",
+                    "batch_id",
+                    "variable",
+                    "objective",
+                    "weighted_objective",
+                    "objectives",
+                    *(
+                        ["nonlinear_constraint", "constraints"]
+                        if results.gradients.constraints is not None
+                        else []
+                    ),
+                ]
             )
             batch_objective_gradient = self._rename_ropt_df_columns(
                 batch_objective_gradient
             )
             batch_objective_gradient = self._enforce_dtypes(batch_objective_gradient)
+        else:
+            batch_objective_gradient = None
 
-        perturbation_objectives = self._rename_ropt_df_columns(perturbation_objectives)
-        perturbation_objectives = self._rename_ropt_df_columns(perturbation_objectives)
-
-        if "constraint_name" in perturbation_objectives:
+        if results.evaluations.perturbed_constraints is not None:
             perturbation_constraints = (
                 perturbation_objectives[
                     "result_id",