Skip to content

Commit

Permalink
Insert names in everest results using polars
Browse files Browse the repository at this point in the history
  • Loading branch information
@verveerpj
verveerpj committed Jan 31, 2025
1 parent fac55ca commit d3fda4e
Showing 1 changed file with 75 additions and 31 deletions.
106 changes: 75 additions & 31 deletions src/everest/everest_storage.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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(
{
Expand All @@ -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",
Expand All @@ -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)
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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
)
Expand Down

0 comments on commit d3fda4e

Please sign in to comment.