From 005c29e37d359bea3db76bbc0b5b3011db9d2b11 Mon Sep 17 00:00:00 2001
From: "Yngve S. Kristiansen" <ynk@equinor.com>
Date: Tue, 14 Jan 2025 08:20:53 +0100
Subject: [PATCH] review: no drop plan_id @ grad results

---
 src/everest/everest_storage.py | 58 ++++++++++++++++++++++------------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py
index 35166ab2145..5cd14ef2041 100644
--- 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",